# rehman047/graph-based-job-load-balancing architecture

> Architecture diagram of the GitHub repository rehman047/graph-based-job-load-balancing, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-11-14). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/rehman047/graph-based-job-load-balancing
- Repository: https://github.com/rehman047/graph-based-job-load-balancing

## Overview

This is a GNN-DDQN batch scheduler for heterogeneous GPU clusters. The main workflow loads Alibaba traces (or reproducible synthetic jobs), trains an agent against a simulated cluster, then evaluates it alongside scheduling baselines and reports metrics. The graph separates workload preparation, simulation, scheduling/learning, and evaluation outputs. Model internals and replay are included as a useful implementation layer. The README documents plots as an output capability, but the sampled training entry point does not call the plotting functions, so no runtime connection to them is asserted.

Main entry points: [`src/train.py`](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/train.py) (Train and evaluate), [`src/model.py`](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py) (DDQN agent), [`src/baselines.py`](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/baselines.py) (Baseline schedulers), [`src/data_loader.py`](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/data_loader.py) (Trace loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workload["Workloads"]
  node_loader["Trace loader<br/>[data_loader.py]"]
end

subgraph group_simulation["Cluster simulation"]
  node_job["Job model<br/>[env.py]"]
  node_cluster["GPU cluster<br/>[env.py]"]
end

subgraph group_scheduling["Scheduling and learning"]
  node_train["Train and evaluate<br/>[train.py]"]
  node_agent["DDQN agent<br/>[model.py]"]
  node_graph_encoder["GNN encoder<br/>[model.py]"]
  node_q_network["Q network<br/>[model.py]"]
  node_replay[("Replay buffer<br/>[model.py]")]
  node_baselines["Baseline schedulers<br/>[baselines.py]"]
end

subgraph group_evaluation["Evaluation and reporting"]
  node_metrics["Scheduling metrics<br/>[metrics.py]"]
  node_plots["Plots and table<br/>[plots.py]"]
  node_results["Evaluation CSVs"]
end

node_trace_source["Alibaba traces"]

node_trace_source -->|"provides traces"| node_loader
node_train -->|"loads workloads"| node_loader
node_train -->|"resets and steps"| node_cluster
node_train -->|"constructs jobs"| node_job
node_train -->|"trains and acts"| node_agent
node_agent -->|"encodes states"| node_graph_encoder
node_agent -->|"computes Q values"| node_q_network
node_train -->|"stores and samples"| node_replay
node_train -->|"evaluates schedulers"| node_baselines
node_baselines -->|"extends DDQN"| node_agent
node_train -->|"scores performance"| node_metrics
node_train -->|"writes evaluation CSV"| node_results
node_train -.-> node_plots

click node_loader "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/data_loader.py"
click node_train "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/train.py"
click node_job "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/env.py"
click node_cluster "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/env.py"
click node_agent "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py"
click node_graph_encoder "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py"
click node_q_network "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py"
click node_replay "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py"
click node_baselines "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/baselines.py"
click node_metrics "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/metrics.py"
click node_plots "https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/plots.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_loader toneBlue
class node_job,node_cluster toneAmber
class node_train,node_agent,node_graph_encoder,node_q_network,node_replay,node_baselines toneMint
class node_metrics,node_plots,node_results toneRose
class node_trace_source toneIndigo
```

## Components

### Workloads

- [Trace loader](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/data_loader.py): `src/data_loader.py`

### Cluster simulation

- [Job model](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/env.py): `src/env.py`
- [GPU cluster](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/env.py): `src/env.py`

### Scheduling and learning

- [Train and evaluate](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/train.py): `src/train.py`
- [DDQN agent](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py): `src/model.py`
- [GNN encoder](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py): `src/model.py`
- [Q network](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py): `src/model.py`
- [Replay buffer](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/model.py): `src/model.py`
- [Baseline schedulers](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/baselines.py): `src/baselines.py`

### Evaluation and reporting

- [Scheduling metrics](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/metrics.py): `src/metrics.py`
- [Plots and table](https://github.com/rehman047/graph-based-job-load-balancing/blob/main/src/plots.py): `src/plots.py`
- **Evaluation CSVs**

### Other components

- **Alibaba traces**

## Connections

- Alibaba traces → Trace loader: provides traces (evidence: `README.md`)
- Train and evaluate → Trace loader: loads workloads (evidence: `src/train.py`)
- Train and evaluate → GPU cluster: resets and steps (evidence: `src/train.py`)
- Train and evaluate → Job model: constructs jobs (evidence: `src/train.py`)
- Train and evaluate → DDQN agent: trains and acts (evidence: `src/train.py`)
- DDQN agent → GNN encoder: encodes states (evidence: `src/model.py`)
- DDQN agent → Q network: computes Q values (evidence: `src/model.py`)
- Train and evaluate → Replay buffer: stores and samples (evidence: `src/train.py`)
- Train and evaluate → Baseline schedulers: evaluates schedulers (evidence: `src/train.py`)
- Baseline schedulers → DDQN agent: extends DDQN (evidence: `src/baselines.py`)
- Train and evaluate → Scheduling metrics: scores performance (evidence: `src/train.py`)
- Train and evaluate → Evaluation CSVs: writes evaluation CSV (evidence: `src/train.py`)
- Train and evaluate → Plots and table (evidence: `src/train.py`)