← All posts
RSS
Systems & ML notebook

The scheduler that has to earn its keep

Inside rl-inference-scheduler: a testbed where a reinforcement-learning agent has to beat FIFO, SJF, and priority queues at the unglamorous job of deciding which LLM request runs next — and, so far, honestly reports that it hasn't yet.

Repo github.com/manishklach/rl-inference-scheduler Stack Python, PyTorch Scope DQN → PPO, single-env → 4096-way vectorized
01 — The problem underneath the problem

Every request is a bet on the future

An LLM inference server doesn't process one request at a time so much as referee a crowd of them. Each incoming job carries a prompt length, a priority, a first-token deadline, and a KV-cache footprint that grows as it runs — and the scheduler has to decide, tick by tick, who gets a compute slot and who waits.

The obvious heuristics all trade one thing for another. Run jobs in arrival order and short requests get stuck behind long ones. Run the shortest job first and you starve whichever long request showed up early. Prioritize by importance and you can blow through VRAM. rl-inference-scheduler takes that tension and turns it into an RL environment: a hardware-aware simulator with prefill/decode phases, memory pressure, and SLOs, where DQN and PPO agents are trained to out-schedule fixed policies like FIFO, SJF, Priority, LargestSafeBatch, and MemAwarePriority.

02 — What the agent actually sees and does

A 21-number snapshot, a 12-way decision

At every tick, the environment hands the agent a 21-dimensional, normalized observation — queue length, VRAM headroom, average wait time, how many jobs are close to breaching their TTFT SLO, how many are near their deadline — and asks it to pick one of 12 actions.

Dispatch queue — one simulated tick
incoming
P1
P3
P3
P1
action
action 7 → reserve slot for high-priority
dispatched
held
held
The action space isn't just "who goes next" — it includes batch-size choices (2/4/8, each with its own throughput multiplier), memory-protective dispatch, deferring low-priority jobs, and dropping expired ones before they burn a deadline penalty anyway.

The reward function is where the intent lives: completing a job pays out, high-priority completions pay more, missing a TTFT SLO costs more than a slow-but-on-time job, missing a hard deadline costs more still, and an OOM event is the single worst outcome in the whole function. That shape is deliberate — it's pushing the agent toward the same trade-offs a real inference-serving team argues about in an incident retro.

03 — Four versions in, and the honesty is the point

What the benchmarks actually show

This is the part most repos smooth over. The published results, run across five scenarios and five seeds, don't have RL winning yet — and the README says so directly.

MetricBest performerWhy
Average latencySJFShortest-job-first is hard to beat on raw efficiency
TTFT SLO missesPriority / MemAwarePriorityFast-tracking urgent jobs pays off directly
Adversarial long→short mixFIFO0% miss rate — SJF starves the long jobs that arrived first
DQN (300 episodes)— matches FIFOUnder-trained for a 21-dim / 12-action space
The finding that matters most: no single heuristic dominates across scenarios. SJF's efficiency becomes a liability the moment the workload shifts adversarially. That gap — not a cherry-picked win — is the actual argument for why a learned policy could be worth the complexity.

The project's response to "DQN doesn't beat the baselines yet" wasn't to hide the table — it was to rebuild the training path. v0.3 replaced the single-environment loop with a tensorized environment that runs up to 4,096 simulated servers in parallel and swapped DQN for PPO with GAE. v0.4 added deterministic trace replay and richer tail-latency metrics (p95/p99) so the next round of results can be checked against something closer to production telemetry.

04 — Four releases, in order

How the project actually got here

v0.1

Basic schedulers + DQN

A tick-based server, FIFO/SJF/Priority baselines, and a DQN agent with a replay buffer and target network — the skeleton everything else hangs on.

v0.2

Hardware-aware simulator

Prefill/decode phases, real KV-cache growth, VRAM pressure, batch-efficiency multipliers, and the 21-dim/12-action spaces described above, plus three more baselines and a five-scenario comparison pipeline.

v0.3

Vectorized rollout engine + PPO

State reworked as flat PyTorch tensors so thousands of environments can step together; a PPO agent with an actor-critic head; six equivalence tests proving the fast path matches the reference simulator exactly.

v0.4

Trace replay + RL-wins benchmark

Six synthetic 10k-request traces, deterministic replay through the vectorized env, a six-phase curriculum for PPO training, and distribution-shift evaluation against scenarios the agent never trained on.

25
tests passing across all four versions
4,096
parallel simulated environments per step
12 × 21
action space × observation dimensions
05 — What's still open

The unfinished part is the interesting part

The repo's own limitations table is worth reading as a roadmap rather than a disclaimer. PPO trained on the vectorized engine hasn't yet been shown to beat the heuristics on the richer trace-replay benchmark — that comparison, run with a longer mixed_curriculum schedule, is the next real test of the project's thesis. Past that: no OOM events have been triggered yet, meaning the memory-pressure scenario isn't tight enough to stress the one penalty term (-10) designed to matter most. And everything so far runs against synthetic traces — the CSV trace-replay path is built specifically so real vLLM or Ray Serve logs can be dropped in without touching the environment code.

That's a reasonable place for a systems-RL project to be four versions in: the infrastructure for a fair fight is built, instrumented, and tested. The fight itself is still in progress.

Follow the fight between the agent and the heuristics

Full simulator, scenario configs, benchmark scripts, and the trace-replay pipeline are on GitHub.

→ github.com/manishklach/rl-inference-scheduler