ZonFormer: trained policies and per-trace results

This repository holds the trained policies and the per-trace result files behind the manuscript "ZonFormer: learning cluster scheduling from simulation with a relational pointer transformer" (Parikshieth Harish, Suhani Bhateja, Arun Pandian J). Every number in the manuscript's tables and figures is computed from the CSV files in results/.

ZonFormer picks the next job to start on a batch cluster. It is a pointer transformer over the 128 most schedulable queued jobs, and its job-to-job attention carries a bias that marks which pairs of jobs fit on the free cores together. It was trained on simulated workloads only and evaluated on 17 job logs from the Parallel Workloads Archive, replayed in an event-driven simulator. No job from those 17 logs entered training data, labels or reward.

Results

The score is the geometric mean, over the 17 traces, of a policy's mean bounded slowdown divided by that of F1-BF (the F1 priority function under EASY backfilling). Lower is better. A seed counts as close when its ratio is at most 1.15. Each recipe was trained at seeds 0 to 4 and scored at its final training step, with no checkpoint selection.

Policy	Seeds	Mean ratio	Range over seeds	Seeds close
ZF-BC (behaviour cloning of F1-PC-BF)	5	1.087	1.055 to 1.121	5 of 5
ZF-PD (six-teacher portfolio distillation)	5	1.181	1.023 to 1.402	3 of 5
ZF-RL (anchored PPO, warmed critic)	5	1.745	1.046 to 2.917	2 of 5
ZF-RL cold (anchored PPO, cold critic)	5	5.922	1.062 to 20.057	3 of 5
DeepRM-style agent	3	5.647	3.842 to 8.631	0 of 3
RLScheduler-style agent	3	8.840	8.391 to 9.137	0 of 3
Structured null	10	1.187	1.133 to 1.249	2 of 10
Untrained control (random legal actions)	10	11.681	11.130 to 12.503	0 of 10

The deterministic classical schedulers on the same traces score: F1-PC-BF 0.992, F1-BF 1.000, F1 1.328, EASY with exact runtimes 1.841, EASY 2.649, SJF 3.115, FCFS 17.494.

The two failing ZF-PD seeds fail on one trace, lanl_o2k, with mean bounded slowdowns of 2,791 and 2,552 against F1-BF's 16.75. The worst ZF-RL cold seed left 63% of capacity idle while jobs waited, as a median over the 17 traces. Across the five seed pairs, ZF-RL and ZF-RL cold do not separate (Mann-Whitney U = 9, p = 0.55).

Files

results/
  seed_final/results.csv        five-seed study and classical schedulers (629 rows)
  wide_hetero_v2/results.csv    DeepRM-style and RLScheduler-style agents, nulls, controls (874 rows)
  wide_hetero_v2/g0 ... g5/     the six benchmark shards that results.csv merges
  eval_archived/results.csv     policies trained before the five-seed study (68 rows)
  train_logs/zf-{bc,pd}_s*.csv  per-epoch loss and teacher agreement of the supervised seeds
weights/
  seed_study/                   20 ZonFormer policies: 4 recipes x seeds 0-4
  pre_study/                    4 ZonFormer policies trained with seed 0 before the study
  comparators/                  9 agents: DeepRM-style, RLScheduler-style, F1-DRL x seeds 0-2

Each row of a result file is one policy at one seed on one trace:

Column	Meaning
workload	trace name; the manuscript's Methods map each name to its archive file
policy	policy label (see the table below)
kind	classical, learned or control
seed	training seed for learned, action-sampling seed for control, -1 for classical
bsld	mean bounded slowdown over the trace's jobs, 10-second threshold
idle_q	fraction of capacity left idle while at least one job waited
util	delivered core-seconds over capacity times the span from first arrival to last completion
error	empty on every row; a failed run would carry its exception here

*-RND rows are untrained controls. They take random legal actions and have no weights. F1-DRL-RND is the structured null: it starts F1's top-ranked job whenever it fits and otherwise picks at random among the jobs the teacher's reservation rule admits.

Policies

Each folder in weights/ contains model_seed<k>.zip, a MaskablePPO archive from sb3-contrib 2.9.0, and hparams.json, the configuration the run was trained with.

Folder	Label in the result files	Observation
seed_study/zf-bc-ag_s{0..4}	ZF-BC-AG	version 2
seed_study/zf-pd-ag_s{0..4}	ZF-PD-AG	version 3
seed_study/zf-rl-ag_s{0..4}	ZF-RL-AG	version 3
seed_study/zf-rl-cold-ag_s{0..4}	ZF-RL-COLD-AG	version 3
pre_study/moe_a1-ag_s0	PAPER-ZF-BC	version 2
pre_study/moe_a2p-ag_s0	PAPER-ZF-PD	version 3
pre_study/rl_buff-ag_s0	PAPER-ZF-RL	version 3
pre_study/rl-repl-2m-ag_s0	PAPER-ZF-RL-COLD	version 3
comparators/pure-mlp_s{0..2}	pure-mlp (DeepRM-style)	16-job window
comparators/pure-structured_s{0..2}	pure-structured (RLScheduler-style)	16-job window
comparators/f1-drl_s{0..2}	F1-DRL	16-job window

pre_study/moe_a2p-ag_s0 scores 1.011 and is the starting point of every ZF-RL and ZF-RL cold seed. The -AG suffix marks policies benchmarked with the aging guard: once a visible startable job has waited more than 100 times its user runtime estimate, the guard starts the most overdue such job in place of the policy's choice. The guard is part of the benchmark harness, not of the weights. The comparator agents ran without it.

The ZonFormer archives use features_extractor_kwargs K = 128 visible jobs, H = 48 timeline points, d_model = 256, 6 layers and 8 heads. The action space has 129 choices: start one of the 128 visible jobs, or wait. The observation is a flat vector of 2,693 values in [0, 1].

Loading

policy.pth inside each archive is a plain PyTorch state dict and loads without the project code:

import io
import zipfile

import torch
from huggingface_hub import hf_hub_download

path = hf_hub_download("zParik/zonformer", "weights/seed_study/zf-bc-ag_s0/model_seed0.zip")
with zipfile.ZipFile(path) as archive:
    state_dict = torch.load(io.BytesIO(archive.read("policy.pth")), map_location="cpu", weights_only=True)
print(len(state_dict), next(iter(state_dict)))  # 308 features_extractor.job_proj.weight

MaskablePPO.load on the full archive needs the policy class hpc_sched.models.zonformer.ZonFormerPolicy. That class, the simulator, the workload generator, the schedulers, the training scripts and the benchmark harness are available from the corresponding author on reasonable request.

Evaluation setup

• Traces: 17 logs from the Parallel Workloads Archive, https://www.cs.huji.ac.il/labs/parallel/workload/logs.html, with 60 to 2,755 cores and 1,002 to 53,676 jobs. They come from twelve distinct machines. The logs are not redistributed here.
• Machine: every trace is replayed on the same imposed heterogeneous machine of 16-core nodes, 78% baseline, 19% double-speed and 3% large-memory.
• Metric: bounded slowdown of job j is max(1, (finish − arrival) / max(runtime, 10 s)). A trace's score is the mean over its jobs.
• Training data: generated episodes from a workload generator fitted on ten other archive logs, with machine size randomised from 32 to 4,096 cores.

Limitations

• The 17 traces served as the development benchmark. The teacher, the aging guard, the anchor's minimum strength, the version-3 observation features and the generator's ranges were each settled after runs on them. The results describe a fixed benchmark, not machines the project never touched.
• The generator was fitted partly on two DAS-2 partitions, siblings of three evaluation traces.
• The DeepRM-style and RLScheduler-style agents share the network family of the published systems but are not reimplementations. They were trained with MaskablePPO on an earlier version of the generator and use a 16-job window without backfilling.
• These policies run inside the project's simulator. None has been deployed on a production scheduler.

License

The weights in weights/ are released under the MIT License (LICENSE). The result files in results/ are released under the Creative Commons Attribution 4.0 International License (results/LICENSE).

Citation

@misc{harish2026zonformer,
  title  = {ZonFormer: learning cluster scheduling from simulation with a relational pointer transformer},
  author = {Harish, Parikshieth and Bhateja, Suhani and J, Arun Pandian},
  year   = {2026},
  note   = {Manuscript under review}
}