Skip to content

LLM capacity reservation with BurstGPT workload replay - #1545

Open
lihuaiming2025 wants to merge 1 commit into
openenv:mainfrom
lihuaiming2025:add-llm-capacity-env-20261102
Open

lihuaiming2025 wants to merge 1 commit into
openenv:mainfrom
lihuaiming2025:add-llm-capacity-env-20261102

Conversation

@lihuaiming2025

@lihuaiming2025 lihuaiming2025 commented Nov 2, 2026 •

Copy link
Copy Markdown

Summary

llm_capacity_env is an MCP environment for serving-capacity reservation using real BurstGPT v2.0 workload logs. An agent commits hourly capacity before the next day's demand is revealed and receives reward from explicit reservation and shortfall costs. The environment provides deterministic chronological replay, a client, container deployment, causal baselines, attributed data, structured trajectories, and reproducible evaluation.

Task setting and objective

The task models day-ahead provisioning for an LLM serving workload. Before a day's requests arrive, the agent chooses how many tokens to reserve for each hour and deployment. Reserving more pays for unused capacity; reserving less incurs a larger shortfall cost. The objective is to learn a reservation policy that reduces this asymmetric cost using only demand that has already been revealed.

This provides a forecasting and cost-sensitive resource-allocation task with numerical actions, a deterministic replay environment, and an explicit learning objective. The agent must estimate hour-specific demand from a limited rolling history, distinguish the two deployment streams, and update its decisions as new days arrive. All 48 hourly quantities must be committed before that day's demand is disclosed. Costs are additive across hours; the environment evaluates reservation quality rather than a shared-hardware scheduling problem.

Task and interaction

Each episode contains 14 daily decisions. At each decision, the agent sees the preceding 14 realized days, the next relative log day, the cost coefficients, and its most recent outcome. It submits 24 hourly token reservations for each of the two historical deployment labels, ChatGPT and GPT-4.

The agent has two MCP tools:

Tool Behavior
observe_demand_history() Returns revealed history and costs without advancing time.
reserve_capacity(reservations) Commits all 48 hourly quantities, advances one day, and returns realized demand, cost, and the next observation.

Reservations must be finite, nonnegative values in [0, 1000000]. Invalid submissions do not advance the episode or change its accumulated cost. Reward is emitted once for each successful reservation; observation calls emit zero reward. Completed episodes contain exactly 14 decisions. Reset, split selection, and simulation controls belong to the orchestrator.

A complete episode proceeds as follows:

  1. The orchestrator selects the train/test partition and resets with a seed. The agent receives the first 14 days of history; the same seed and partition reproduce the same start.
  2. The agent reads history, next_log_day, unit_cost_per_1000_tokens, shortfall_multiplier, and its previous outcome. Each history day contains two arrays of 24 hourly token counts.
  3. It calls reserve_capacity with a dictionary containing exactly ChatGPT and GPT-4, each mapped to 24 numeric reservations.
  4. The server scores the committed plan against the recorded next day, returns realized demand and cost, and shifts the history forward one day. The agent can use that outcome for its next decision.
  5. After 14 accepted reservations, the episode terminates. Evaluation compares complete episodes; an unfinished prefix cannot be reported as a successful zero-cost run.

Reward

For each deployment m and hour, reserved tokens r and realized demand d determine:

hourly_cost = unit_cost[m] * (r + 4 * max(d - r, 0)) / 1000
unit_cost = {ChatGPT: 0.002, GPT-4: 0.0375}
daily_reward = -sum(hourly_cost) / 1000

The objective is to minimize total cost over a complete episode, equivalently maximizing accumulated reward. Costs are explicit modeling assumptions. The score is continuous, with no binary acceptance threshold. The rolling 0.75 quantile is a causal baseline motivated by the 4x shortfall coefficient. A perfect-information baseline is separately labeled as a hindsight diagnostic because it reads future demand.

For an illustrative single ChatGPT hour with demand of 150,000 tokens, the accounting is:

Reservation Shortfall Modeled hourly cost
100,000 50,000 0.600
150,000 0 0.300
200,000 0 0.400

The agent chooses before seeing the 150,000-token outcome. This example shows why both systematic overreservation and underreservation can lose to a calibrated forecast. The 1,000,000-token action cap is an input bound, not a modeled physical capacity limit.

Data and evaluation

The workload is an hourly aggregation of Total tokens in BurstGPT v2.0 conversation-service records, covering 231 relative days in two recorded segments (relative days 0–120 and 225–334). There are 11,088 hourly deployment records (231 × 24 × 2). Missing hours within recorded days are zero-filled; the unrecorded gap between segments is preserved. Episodes stay within a segment and require 14 history days followed by 14 target days. Train and test target-day sets are disjoint, with 116 eligible train starts and 35 test starts. The test starts cover 61 unique target dates; adjacent episodes overlap.

Partition Target-day ranges Eligible 14-day starts
Train 14–89 and 239–304 116
Test 90–120 and 305–334 35

The source and test partition are public. During remote interaction, the agent receives only revealed history through the tools; the complete workload stays on the environment server. No request text, completions, or user identifiers are distributed. Runtime does not require a model API key, provider account, or external data service.

The data reconstruction script, BurstGPT v2.0 release attribution, aggregation procedure, source checksums, and aggregate SHA-256 are included in scripts/rebuild_aggregates.py and server/data/PROVENANCE.json. The environment is adapted from S11_system_llm_capacity; the public BurstGPT release is the data source.

Data is distributed under CC-BY-4.0 and environment code under BSD-3-Clause, with the original data license and third-party notices preserved.

Baseline performance and learning evidence

The supplied evaluator runs every eligible start, rather than selecting favorable episodes. On the 35 test starts, the causal baselines produce:

Policy Information used Mean modeled cost per 14-day episode ↓
Zero reservation No demand estimate 2,432.355946
Previous day Latest revealed day's demand 1,293.076713
Fixed 0.75 quantile Initial 14-day history, frozen for the episode 1,304.448138
Rolling 0.75 quantile Latest 14 revealed days at each decision 1,240.313545

The quantile baseline optimizes the empirical hourly loss implied by the stated 4x shortfall coefficient. It is reproducible from the public observation alone. A separately labeled perfect-information diagnostic reserves the realized future demand and gives mean episode cost 608.088987; it is a hindsight lower bound, not an executable causal policy or a promised achievable score.

There is measured improvement beyond the rolling baseline: a trained neural residual policy reduces unique-day test cost by 3.93% across three seeds, with a conditional 95% interval of 0.94–6.84%. In a separate matched-budget program-generation study, Opus reduces unique-day test cost by 6.47% relative to Haiku, with a conditional interval of 0.79–12.79%. Episode means in the table and unique-day improvement statistics use different weightings; the study protocols and uncertainty treatment are specified under Test Plan.

Run a complete episode

From envs/llm_capacity_env, install the locked environment and start the server:

uv sync --python 3.11 --frozen --all-groups --all-extras
CAPACITY_SPLIT=test uv run --frozen uvicorn server.app:app \
  --host 0.0.0.0 --port 8000

In a second shell in the same directory, run the supplied rolling-quantile client:

uv run --frozen python examples/client_demo.py http://localhost:8000

The client resets through the orchestration interface, makes 14 reservation calls, and prints the daily and cumulative modeled cost. Its interaction loop is:

from openenv.core.mcp_client import MCPToolClient
from examples.policies import reservations

with MCPToolClient(base_url="http://localhost:8000").sync() as client:
    client.reset(seed=17)  # Orchestration; not an agent-exposed tool.
    view = client.call_tool("observe_demand_history")
    while not view["done"]:
        plan = reservations(view, policy="rolling_quantile")
        view = client.call_tool("reserve_capacity", reservations=plan)
    print(view["cumulative_modeled_cost"])

For the full baseline table, run uv run --frozen python examples/evaluate.py baseline-results.json. The server records per-session JSONL trajectories containing reset, actions, observations, reward, termination, and decision count. The default output directory is outputs/trajectories; two complete example traces are included in the contribution.

Implementation and documentation

Location Contents
envs/llm_capacity_env/server/ Replay engine, MCP environment, API application, session trajectory recording, and attributed hourly data.
envs/llm_capacity_env/client.py OpenEnv client using the existing transport and MCP types.
envs/llm_capacity_env/examples/ Causal policies, baseline evaluation, client interaction, and trace recording.
envs/llm_capacity_env/ Package metadata, dependency locks, Docker configuration, licenses, reconstruction script, and README.
tests/envs/test_llm_capacity_*.py Environment, MCP, causality, data integrity, reward, termination, isolation, and trajectory checks.
docs/source/ Environment catalog card, navigation entry, and synchronized environment documentation.

Two complete example trajectories and compact runtime/validation records are included with the environment.

Type of Change

  • Bug fix
  • New feature
  • Breaking change
  • Documentation
  • New environment
  • Refactoring

Alignment Checklist

  • I have read .claude/docs/PRINCIPLES.md and this PR aligns with our principles
  • I have checked .claude/docs/INVARIANTS.md and no invariants are violated
  • I have run /pre-submit-pr (or bash .claude/hooks/lint.sh and tests) and addressed all issues

RFC Status

  • Not required (new example environment using existing patterns)
  • RFC exists: #___
  • RFC needed (will create before merge)

The environment uses the existing MCPEnvironment, EnvClient, and create_app interfaces. Agent tools expose observations and decisions; orchestration owns reset and simulation controls.

Test Plan

Core CI and environment checks

The core-test selection and dependency recipe follow the Python 3.11 and 3.12 matrix in .github/workflows/test.yml, using uv 0.9.3. Completed results against the base revision are:

Python Core suite Environment tests
3.11 4,281 passed, 0 failed, 136 skipped, 86 deselected 21 passed
3.12 4,328 passed, 0 failed, 111 skipped, 86 deselected 21 passed

Both versions pass the CI usort/Ruff formatter sequence with no Python diff, Ruff lint, environment formatting, frozen environment-lock validation, documentation synchronization, and whitespace checks. validation-evidence/CI_MATRIX.json records source hashes and results. The 3.11 suite required a full rerun after one AnyIO bridge-cancellation timeout; both attempts are retained.

TextArena's NLTK resources use a local, checksum-verified copy of the official index and corpora because cluster download access is unreliable. The downloader is configured through its index URL; tests, test selection, and NLTK security checks are unchanged. This dependency setup is separate from the environment, which has no NLTK dependency.

The broader contributor lint hook reports 47 existing formatting warnings and two usort warnings on unchanged main. The contribution passes its checks, and the actual CI formatting sequence leaves no Python changes on either tree. The repository-wide checklist item remains unchecked to reflect that broader hook result.

From the repository root, install the dependencies listed in .github/workflows/test.yml, then run:

PYTHONPATH=src:envs uv run pytest tests/ \
  --ignore=tests/envs/test_browsergym_environment.py \
  --ignore=tests/envs/test_dipg_environment.py \
  --ignore=tests/envs/test_websearch_environment.py \
  -m "not integration and not network and not docker" -v --tb=short

PYTHONPATH=src:envs uv run pytest -q \
  tests/envs/test_llm_capacity_mcp.py \
  tests/envs/test_llm_capacity_replay.py \
  tests/envs/test_llm_capacity_trajectory.py

python scripts/sync_env_docs.py --check
git diff --check

The environment's frozen dependency resolution can be checked with:

cd envs/llm_capacity_env
uv sync --frozen --all-groups --all-extras --dry-run --no-install-project

Runtime and provenance

  • Official static manifest validation and all six URL/API checks passed.
  • Actual Docker build and a complete 14-decision network episode passed.
  • Concurrent-session isolation, deterministic reset, invalid/terminal action handling, reward accounting, and web interface checks passed.
  • Reconstruction from all three BurstGPT v2.0 source files reproduced the aggregate hash; wheel packaging retained the client and data.

Container build from the repository root:

docker build -t llm-capacity-env \
  -f envs/llm_capacity_env/server/Dockerfile envs/llm_capacity_env

Policy learning and comparison

The neural study uses a 32-hidden-unit quantile policy with Adam, with cold-start and residual-on-baseline arms. There are three seeds per arm, two learning rates and two regularizers (24 fitting configurations), each trained for 80 epochs. Configuration and checkpoint selection use 82 fit days and 60 validation days from the training partition. The six selected policies are frozen before evaluation over all 35 test starts. This is supervised policy fitting against the exact environment cost.

The matched-budget program study uses Haiku 4.5 and Opus 5.5 gateway aliases, three independently generated programs per model, and three train-feedback submissions per program. Prompts, training starts, feedback rounds, sandbox limits, and token budgets are matched: 32,768 output tokens including 8,192 thinking tokens per response. The train-selected programs are frozen before testing. Requests, responses, selected-program hashes, and all outcomes are retained.

All 502 episodes / 7,028 decisions in the supplementary study passed independent replay, including frozen-weight action reconstruction. Intervals aggregate 61 unique test dates and use seven-day blocks within the two recorded segments. These results measure learnability and policy differences on the recorded workload; overlapping episodes are not treated as independent samples. Study methods and results are documented in envs/llm_capacity_env/README.md. The accompanying quality-evidence-20261102.tar.gz archive contains the frozen programs, weights, requests/responses, trajectories, analysis, and independent replay verifier; reference-model-evidence-20261102.tar.gz retains the earlier model study.

Claude Code Review

N/A


Note

Medium Risk
Large additive package with bundled workload data and new server surface area, but it follows existing OpenEnv MCP patterns and does not modify core auth or shared runtime paths.

Overview
Adds llm_capacity_env, a new OpenEnv MCP environment where agents commit 24 hourly token reservations per deployment (ChatGPT / GPT-4) for 14 days before demand is revealed, minimizing modeled reservation plus 4× shortfall cost on BurstGPT v2.0 hourly aggregates.

The implementation includes a deterministic Replay engine and CapacityEnvironment exposing observe_demand_history and reserve_capacity, a FastAPI app via create_app, optional server-side JSONL trajectories, Docker/lockfiles, causal baselines and trace tooling under examples/, provenance/reconstruction scripts, and shipped validation artifacts plus example JSONL episodes.

Docs register the environment in the catalog and add environments/llm_capacity.md aligned with the Space README (splits, reward math, install/validate/publish).

Reviewed by Cursor Bugbot for commit fc306f0. Bugbot is set up for automated code reviews on this repo. Configure here.

@lihuaiming2025

Copy link
Copy Markdown
Author

Hi @sergiopaniego @burtenshaw, this environment PR is ready for review. The GitHub Actions workflows are awaiting maintainer approval. Could you please approve the workflow runs and help route the environment review? The PR includes the implementation, data provenance, reproducible evaluation materials, and completed Python 3.11/3.12 test results. Thank you!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant