GrapheonRL GNN-PPO DAG Scheduler

A trainable, reproducible, standalone DAG scheduler for heterogeneous task scheduling using Graph Neural Networks and Proximal Policy Optimization with HEFT pretraining.

Overview

GrapheonRL DAG + Heterogeneous Resources
              ↓
       HEFT → Expert Trajectories
              ↓
   Behavioral Cloning → Pretrained GNN Actor
              ↓
     PPO → Direct Makespan Optimization
              ↓
         Independent Inference

The Actor learns to schedule tasks on heterogeneous compute nodes independently of HEFT at inference time. The reward function telescopes: sum(r_t) = -M_RL / M_HEFT, directly aligning training with the evaluation metric.

Installation

# Using conda + uv (recommended)
conda create -n gnnrl python=3.11
conda activate gnnrl
pip install -e ".[dev]"

Or with make:

make setup

Requirements

• Python ≥ 3.11
• PyTorch ≥ 2.4
• PyTorch Geometric ≥ 2.5
• See pyproject.toml for full dependency list.

CPU is fully supported. CUDA is auto-detected if available.

Quick Start

# Full pipeline
make all

# Or step by step:
make prepare    # Prepare data and splits
make expert     # Generate HEFT expert trajectories
make train-bc   # Behavioral cloning pretraining
make train-ppo  # PPO training
make evaluate   # Evaluation on test set
make test       # Run tests

Inference

python -m scripts.infer \
  --workflow path/to/workflow.json \
  --system path/to/system.json \
  --checkpoint outputs/checkpoints/ppo_best_actor.pt \
  --output outputs/inference/run_001

The inference script:

• Does NOT run HEFT (except optionally for comparison)
• Does NOT modify model parameters
• Uses model.eval() and torch.inference_mode()
• Produces zero invalid actions

Project Structure

grapheonrl-gnn-ppo/
├── configs/           # YAML configuration files
├── data/              # Raw, processed, splits, expert, cache
├── src/grapheon_scheduler/
│   ├── baselines/     # HEFT, GreedyEFT, Random schedulers
│   ├── data/          # Schema, loader, split, expert dataset
│   ├── env/           # GraphState builder, StaticDAGEnv
│   ├── evaluation/    # Evaluator, metrics, plots
│   ├── models/        # GAT encoder, Actor, Critic, batching
│   ├── scheduling/    # ScheduleEngine, features, validator, export
│   └── training/      # BC, rollout, GAE, PPO trainers
├── scripts/           # CLI entry points
├── tests/             # pytest test suite
└── outputs/           # Checkpoints, logs, evaluation results

Key Design Decisions

1. Single ScheduleEngine: All baselines, expert generation, RL, and evaluation share one scheduling kernel — no duplicate EST/EFT logic.

2. Independent Actor/Critic GNNs: Actor and Critic use separate GATv2 encoders (no parameter sharing), preventing gradient interference.

3. Action Masking: Illegal actions (infeasible task-node pairs, non-ready tasks) are masked at every stage: training, validation, and inference. Illegal logits are set to -inf.

4. No HEFT Dependency at Inference: The model learns from HEFT during BC pretraining but operates independently at inference. HEFT makespan is NOT used as a model input.

5. Static List Scheduling: One DAG = one episode. Exactly num_tasks actions per episode. Non-preemptive serial queue per node.

Architecture

Actor

Task Features + DAG → 3× GATv2Conv (128d, 4 heads) → Task Embeddings
Resource Features → 2× MLP → Resource Embeddings
(Task, Node, Pair) → Pair Scorer MLP (256→128→1) → T×N Logits
Action Mask → Categorical Distribution

Critic

Task Features → Independent 3× GATv2Conv → mean+max pool
Resource Features → Independent MLP → mean+max pool
Global Features → Concat → Value MLP (128→64→1) → V(s)

Configuration

All parameters are in YAML configs:

• configs/default.yaml — Base configuration
• configs/train_bc.yaml — BC training overrides
• configs/train_ppo.yaml — PPO training overrides
• configs/evaluate.yaml — Evaluation overrides

Configs support inheritance via the extends key.

Evaluation Metrics

• mean_ratio = mean(M_RL / M_HEFT) — primary metric
  • < 1: better than HEFT (on average)
  • = 1: equal to HEFT

  • 1: worse than HEFT

• win / tie / loss rates
• Per-variant and per-mode breakdowns

Testing

pytest -q                 # All tests
pytest tests/ -v           # Verbose output

Tests cover: loader, DAG validation, feasibility, timing, communication, HEFT, expert replay, reward, action mask, environment, checkpoint save/load, and integration tests.

Data

Uses the GrapheonRL STG-JSON dataset:

• DOI: https://doi.org/10.5281/zenodo.18927122
• License: CC BY 4.0
• Variants: rnc50, rnc100, rnc300
• Modes: homogeneous, heterogeneous

References

• GrapheonRL Benchmark: https://github.com/AasishKumarSharma/grapheonrl-benchmark
• GOODRL: https://github.com/YifanYang1995/GOODRL
• READYS: https://github.com/nathangrinsztajn/RL_for_dynamic_scheduling
• HEFT: Topcuoglu, Hariri & Wu (2002), IEEE TPDS

License

MIT. See THIRD_PARTY_NOTICES.md for third-party code attributions.