rehman047/graph-based-job-load-balancing

How rehman047/graph-based-job-load-balancing works

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: train.py Train and evaluate · model.py DDQN agent · baselines.py Baseline schedulers · data_loader.py Trace loader

Components 13

Workloads

Trace loader src/data_loader.py

Cluster simulation

Job model src/env.py
GPU cluster src/env.py

Scheduling and learning

Train and evaluate src/train.py
DDQN agent src/model.py
GNN encoder src/model.py
Q network src/model.py
Replay buffer src/model.py
Baseline schedulers src/baselines.py

Evaluation and reporting

Scheduling metrics src/metrics.py
Plots and table src/plots.py
Evaluation CSVs

Other components

Alibaba traces

Connections 13

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

Drawn by GitDiagram on November 14, 2026 from the repository's file tree, README and source files. Read as Markdown ·