One scheduling decision per tick keeps the task compact and inspectable.
GPU Scheduler OpenEnv
A compact OpenEnv benchmark for AI-assisted GPU inference scheduling. An agent observes a small simulated GPU cluster, decides where pending inference-style jobs should run, and is scored on deadline handling, resource fit, allocation, and cost.
Deadline pressure, cost pressure, patience under scarcity, and large-GPU reservation.
Competitive with heuristics on the first two public scenarios with zero invalid actions.
1. Why This Environment Exists
Many agent benchmarks focus on browsing, code generation, or static tool use. Operational systems have a different shape: the state changes over time, each action changes the future decision surface, and reward should come from the environment rather than from a hidden grader. GPU scheduling is a clean version of that problem.
This benchmark is intentionally not a production scheduler. It does not claim to model CUDA kernels, real cluster telemetry, or datacenter-scale orchestration. The question it asks is narrower and more useful for benchmarking: can an agent make good short-horizon scheduling decisions under fit, urgency, allocation, and cost pressure?
Key insight
The benchmark's value is not in simulating a real cluster. It is in creating a clean, scorable sequential decision problem where every tradeoff is explicit and measurable.
2. System Architecture
Definition: OpenEnv
OpenEnv is a framework for building and evaluating agent environments. It provides a standard reset() / step() interface,
an HTTP server layer, and a built-in /web inspection UI. Agents interact through structured actions and receive scalar reward from environment dynamics,
not from a hidden grader.
Scenario ConfigHand-authored scenarios and seeded generation define jobs, GPU tiers, and tick budget.
↓
GPU Scheduling SimulatorValidates actions, advances time, tracks allocations, and computes reward.
Environment WrapperExposes a standard reset/step loop and adds dashboard-oriented fields for inspection.
↓
OpenEnv Server and Web UIProvides /reset, /step, /state, and the interactive console at /web.
Agent ConsumersLocal Python agents, LLM rollout runner, lightweight RL trainer, and heuristic baselines all consume the same interface.
↓
Evaluation ReportsBenchmark tables, generated visuals, and human inspection traces are produced from the same environment loop.
- Scenario definitions create a small GPU cluster and a pending queue of jobs.
- The simulator validates actions, advances time, and computes reward.
- A thin environment wrapper exposes a standard reset() / step() loop.
- That loop is consumed by heuristics, an LLM rollout script, a lightweight learned policy, and the OpenEnv web interface.
The important point is that the reward is environment-derived. Nothing in the benchmark silently grades free-form language output. The agent acts, the environment evolves, and the score follows from those dynamics.
3. The Scheduling Task
Definition: Scenario and Tick
A scenario is one complete run of the scheduling environment, from the initial cluster state until either all jobs are resolved or the tick budget is exhausted.
A tick is a single time step within that run. On each tick, the agent issues exactly one action, the simulator processes it, and the environment transitions to the next state.
Task Shape
- 3 abstract GPU types with different VRAM, speed, and cost
- 6 to 10 pending inference-style jobs
- 10 to 20 scheduling ticks
Action Space
Definition: Action Space
The complete set of actions an agent may issue at any tick. Here it is intentionally minimal: place(job_id, gpu_id) or defer().
- place(job_id, gpu_id) assigns a pending job to a specific GPU.
- defer() takes no placement action this tick and waits for conditions to change.
- Each action also carries a short rationale. Rationales are logged and shown in the UI but are not scored in v1.
Observation Space
Definition: Observation Space
The full state representation the agent receives at each tick before choosing an action.
Base StatePending jobs, running jobs, completed jobs, missed jobs, per-GPU state, cumulative metrics, and action history.
Dashboard Fieldssummary, action_hint, concise GPU status lines, and queue status lines for the OpenEnv surface.
Why It WorksThe same structured observation is usable from heuristics, prompted models, and browser inspection.
Key ConstraintKeeping the action space to two verbs forces the interesting policy work into when and where to place.
Key insight
Keeping the action space to just two verbs, place and defer, forces the policy work into timing and fit rather than parsing a complex command grammar.
4. Reward Design
Definition: Reward Shaping
Reward shaping is the deliberate design of bonus and penalty terms that guide an agent toward desired behavior beyond simple task completion.
The reward function is deliberately simple and environment-grounded. Formally, the total reward accumulated over a scenario is:
R = Σ_j [ r_complete(j) + r_priority(j) ]
- Σ_j p_missed(j)
- Σ_a p_invalid(a)
- Σ_t p_idle(t)
- Σ_t p_cost(t)| Term | Sign | Description |
|---|---|---|
| r_complete(j) | + | Positive reward for successfully completing job j. |
| r_priority(j) | + | Extra bonus scaled to the priority level of job j. |
| p_missed(j) | - | Penalty when job j misses its deadline. |
| p_invalid(a) | - | Penalty for issuing an infeasible placement action. |
| p_idle(t) | - | Penalty for leaving GPU capacity idle while jobs are still pending. |
| p_cost(t) | - | Penalty for running more expensive GPUs than necessary. |
That combination is enough to make policies diverge for the right reasons. A policy that only cares about packing VRAM tightly can still lose if it misses urgent jobs. A policy that always grabs the largest GPU can finish work quickly but overspend. A policy that waits too conservatively can avoid some mistakes but still waste turns and tighten deadlines.
Key insight
The reward function is designed so that no single greedy heuristic can dominate. Each canonical scenario is constructed to expose the failure mode of optimizing one term at the expense of the others.
5. Included Scenarios
The benchmark uses four fixed, hand-authored evaluation scenarios. Each is designed to stress-test a different axis of scheduler competence.
| Scenario | Primary Test | Optimal Strategy |
|---|---|---|
| deadline_crunch | Urgency-aware scheduling under time pressure | Prioritize jobs by deadline, even at higher cost |
| cost_pressure | Balancing throughput against GPU spend | Choose fit-for-purpose GPUs and avoid over-provisioning |
| wait_for_capacity | Patience under short-term resource scarcity | Issue defer() on tick 1 to await cheaper GPU availability |
| reserve_large_gpu | Reservation of scarce resources for the right workload | Resist placing low-value jobs on the only large-VRAM GPU |
Key insight
A policy that scores well on all four scenarios must internalize urgency handling, cost awareness, patience under scarcity, and strategic reservation.
6. Baselines and Results
The current repo includes five reference baselines: FIFO, earliest deadline first, best-fit VRAM, a structured-output gpt-4.1-mini rollout, and the lightweight learned rl_linear_policy.
| Policy | Scenario | Total Reward | Completed | Missed | Invalid | Allocation | Total Cost |
|---|---|---|---|---|---|---|---|
| fifo | deadline_crunch | 29.8 | 5 | 1 | 0 | 0.278 | 19.2 |
| earliest_deadline_first | deadline_crunch | 48.8 | 6 | 0 | 0 | 0.286 | 21.2 |
| best_fit_vram | deadline_crunch | 32.2 | 5 | 1 | 0 | 0.389 | 17.8 |
| fifo | cost_pressure | 51.0 | 7 | 0 | 0 | 0.333 | 29.0 |
| earliest_deadline_first | cost_pressure | 44.7 | 7 | 0 | 0 | 0.212 | 30.8 |
| best_fit_vram | cost_pressure | 34.6 | 6 | 1 | 0 | 0.296 | 24.4 |
| fifo | wait_for_capacity | 36.3 | 5 | 0 | 0 | 0.250 | 18.7 |
| earliest_deadline_first | wait_for_capacity | 35.8 | 5 | 0 | 0 | 0.250 | 18.7 |
| best_fit_vram | wait_for_capacity | 40.4 | 5 | 0 | 0 | 0.250 | 14.6 |
| fifo | reserve_large_gpu | 37.5 | 5 | 0 | 0 | 0.333 | 21.0 |
| earliest_deadline_first | reserve_large_gpu | 33.4 | 5 | 0 | 0 | 0.333 | 24.1 |
| best_fit_vram | reserve_large_gpu | 37.5 | 5 | 0 | 0 | 0.400 | 21.0 |
| gpt-4.1-mini | deadline_crunch | 48.5 | 6 | 0 | 0 | 0.250 | 20.0 |
| gpt-4.1-mini | cost_pressure | 51.5 | 7 | 0 | 0 | 0.333 | 29.0 |
| rl_linear_policy | deadline_crunch | 48.5 | 6 | 0 | 0 | 0.143 | 20.0 |
| rl_linear_policy | cost_pressure | 46.3 | 7 | 0 | 0 | 0.259 | 32.2 |
| rl_linear_policy | wait_for_capacity | 35.8 | 5 | 0 | 0 | 0.200 | 18.7 |
| rl_linear_policy | reserve_large_gpu | 32.0 | 5 | 0 | 0 | 0.267 | 26.0 |
- earliest_deadline_first is strongest on deadline_crunch, which is exactly what that scenario is designed to reward.
- best_fit_vram is strongest on wait_for_capacity, where paying for the large GPU too early is the wrong move.
- gpt-4.1-mini is already competitive with the heuristics and maintains zero invalid actions on the public LLM scenarios.
- The lightweight learned policy remains solid on the original two scenarios, but the defer-sensitive scenarios are harder and expose weaker reservation behavior.
7. Figures
The figures below show the benchmark comparison across policies and scenarios, followed by two scheduler-console views: the environment at reset and the console after one placement action.
8. What the LLM Baseline Shows
The first live model run is useful because it does not fail on schema compliance. It fails on policy sharpness, which is exactly the kind of benchmark signal you want from this environment.
- Strong at fit reasoning and obvious urgency handling.
- Zero invalid actions on the public LLM scenarios.
- Weaker when it should parallelize aggressively instead of deferring conservatively.
9. How to Use the Environment
- Run the browser console at /web.
- Swap in a local policy via examples/custom_agent.py.
- Call the OpenEnv server over /reset, /step, and /state.
- Use green_agent/runner.py as a thin adapter around your own policy or LLM call.