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!

@burtenshaw burtenshaw left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed at fc306f0.

What I ran: the 21 env tests, python scripts/sync_env_docs.py --check, openenv validate envs/llm_capacity_env --level static --skip-build, examples/evaluate.py (it reproduces the baseline table exactly), a build of server/Dockerfile, and live episodes through server.app:app, both in-process and in the built container. scripts/rebuild_aggregates.py on the three public BurstGPT v2.0 release assets reproduces de9db57e9dd8aa1c96f33789e45eaf5fe0b09726160e9554645d2cb8ca58c130.

No observation, tool result, state(), tool listing, HTTP metadata endpoint or error message contains a day's demand before that day is committed. I searched every payload the agent can read before each commit for the target day's hourly arrays, across five policies. reset is not an agent tool, reward is computed in the engine, and calls after the 14th decision return an error with reward 0 and done=True.

Blocking

1. Never committing earns the maximum reward

Every accepted reservation earns a negative reward (server/engine.py:141). Every call that does not advance the day earns 0.0 (server/capacity_environment.py:50). A policy that never commits therefore outscores every policy that does, including the hindsight oracle. Live server, test split, seed 17:

Policy Sum of step rewards done
Hindsight oracle (reads future demand) -0.964335 true
Rolling 0.75 quantile -1.705859 true
50 calls of reserve_capacity(reservations={}) 0.0 false
50 calls of observe_demand_history() 0.0 false

The built container returns the same values. cumulative_modeled_cost stays at 0.0 for the last two rows, so scoring from the tool payload has the same problem. Under a turn limit, a trainer that sums step rewards ranks the policy that never calls reserve_capacity above all others. README.md:31 asks the evaluation harness to require completion, which moves task knowledge out of the environment (PRINCIPLES.md, "Rewards inside environment"). This is also the RFC 008 semantic.floor_gap failure: openenv.yaml:10-12 declares a maximum of 0.0 with floor_margin: 0.1, and doing nothing scores 0.0. Level 3 validation is not implemented yet, so CI does not catch it.

Regression test for tests/envs/test_llm_capacity_mcp.py. It fails at fc306f0 with assert -2.0461811029999994 > 0.0:

from llm_capacity_env.server.capacity_environment import CapacityEnvironment
from llm_capacity_env.server.engine import EPISODE_DAYS, MODELS
from openenv.core.env_server.mcp_types import CallToolAction


def episode_reward(commit):
    env = CapacityEnvironment(split="test")
    try:
        env.reset(seed=17)
        total = 0.0
        for _ in range(EPISODE_DAYS):
            view = env.replay.observe()
            if commit:
                plan = {m: view["history"][-1]["tokens"][m] for m in MODELS}
                action = CallToolAction(
                    tool_name="reserve_capacity", arguments={"reservations": plan}
                )
            else:
                action = CallToolAction(
                    tool_name="observe_demand_history", arguments={}
                )
            total += env.step(action).reward
        return total
    finally:
        env.close()


def test_never_committing_scores_below_a_causal_policy():
    assert episode_reward(commit=True) > episode_reward(commit=False)

Smallest fix: in Replay.reserve, reward each committed day by its savings against reserving nothing.

cost = modeled_cost(plan, demand)
unreserved = modeled_cost({m: [0.0] * 24 for m in MODELS}, demand)
...
"reward": (unreserved - cost) / REWARD_SCALE,

Doing nothing then scores 0, a useful forecast scores above 0, and the oracle stays the maximum. With this change the regression test passes. The four existing tests that assert a negative reward or reward == -cost / 1000 need updating, and the per-step range at openenv.yaml:10 becomes [-0.948, 1.422] under the 500,000-token bound. modeled_cost and cumulative_modeled_cost remain the evaluation metric. Another causal baseline in place of zero reservation works as well.

2. The default server writes every session to disk with no limit

server/app.py:12 sets CAPACITY_TRACE_DIR=outputs/trajectories for every server process. Each reset creates a JSONL file. Each tool call appends about 17 KB, because the tool result is stored three times (content, structured_content and data). In the built image, one 14-decision episode from examples/client_demo.py wrote 268,824 bytes to /app/env/outputs/trajectories. In-process, 20 stateless POST /reset calls created 20 files. Nothing rotates or caps the directory, and the disk_mb: 2048 declared at openenv.yaml:16 fills before 8,000 episodes.

Smallest fix: delete server/app.py:12 so recording stays opt-in through CAPACITY_TRACE_DIR, as it already is for direct library use, and update README.md:189-194.

3. The image installs from a lockfile that CI never reads

server/Dockerfile:3-4 installs requirements.lock. CI validates uv.lock, and no CI job builds this image. The two files already disagree: authlib 1.9.1 vs 1.9.2, openai 3.42.0 vs 3.42.1, platformdirs 4.14.2 vs 4.14.3. requirements.lock:83 and :93 also put pytest and ruff in the runtime image. The root Dockerfile duplicates server/Dockerfile. openenv build reads server/Dockerfile, and openenv push moves it to the root when no root file exists.

Smallest fix: delete requirements.lock and the root Dockerfile, and build from uv.lock as envs/echo_env/server/Dockerfile does.

4. Generated files and submission records are committed

  • outputs/example-traces/ (315,357 bytes) regenerates byte for byte with python examples/record_trace.py <dir>. Both JSONL files and TRACE_RESULT.json match.
  • validation-evidence/ (11 files, 22,301 bytes) records one submission run, including internal paths such as /mnt/cfs_bj_fm/workspace/lihuaiming/... at validation-evidence/SUMMARY.json:8-10. .dockerignore:7 lists validation, so this directory also ships in the image.

Smallest fix: delete both directories, outputs/.gitkeep, and the matching rules in .gitignore and .dockerignore. Keep examples/record_trace.py. The validation record fits in the PR description.

Suggestions

  • README.md:204-279: remove the "Publish" note about asking maintainers and the two study sections. They describe the submission, link archives on a fork's release page and name model gateway aliases. docs/source/environments/llm_capacity.md copies them into the docs site.
  • pyproject.toml:14: the exact pin openenv==0.11.0 blocks runtime fixes. Other envs use a lower bound.
  • pyproject.toml:29-30: no other env carries its own [tool.ruff]. Under the repository config, 7 env files reformat and usort re-sorts examples/client_demo.py and examples/record_trace.py.
  • server/capacity_environment.py:55 and :59: getattr(self.replay, "turn", 0) can be self.replay.turn, which server/engine.py:84 always sets.
  • docs/source/environments.md:10-16: the card sits above Echo. Recent envs append to the end of the grid.
  • Data: day 77 has zero ChatGPT demand in all 24 hours, and test days 101 to 106 have 12 to 20 zero hours per deployment. These look like upstream logging gaps scored as real zero demand, so part of the test cost measures outage prediction. Document them in the README or exclude those days from eligible starts.
  • pyproject.toml:12 and the README front matter declare BSD-3-Clause only, so the wheel metadata and the discovery catalog card report BSD-3-Clause for a package that ships CC-BY-4.0 data. license = "BSD-3-Clause AND CC-BY-4.0" fixes the wheel. An optional discovery.json can add the data license and the two agent tools to the catalog card.
  • .gitattributes:2: binary keeps the CRLF bytes stable but hides future data changes in diffs. -text keeps the bytes and the diff.

Licensing and provenance

BurstGPT is CC-BY-4.0. That license permits sharing adapted material with attribution, a link to the license and a statement of changes. server/data/LICENSE-BurstGPT.txt is byte-identical to the upstream LICENSE. THIRD_PARTY_NOTICES.md, server/data/PROVENANCE.json and the README name the authors, the KDD 2025 paper, the upstream URL and release, and the changes made. The rebuild from the public release assets matches the bundled CSV. Redistributing the CSV is permitted and correctly attributed.

Two points need an answer before merge:

  • LICENSE:3 adds a per-env BSD-3-Clause license with "Copyright (c) 2026, Huaiming Li". CONTRIBUTING.md:31 licenses contributions under the root LICENSE, and no other env ships its own code license. Please remove LICENSE and drop it from license-files at pyproject.toml:9. Keep the data license and THIRD_PARTY_NOTICES.md.
  • server/data/PROVENANCE.json:21-31, README.md:60-61 and README.md:89 say the environment and its cost coefficients are adapted from famou-bench-v3.1/cases/S11_system_llm_capacity. I could not find that benchmark in public. Please confirm that you have the right to contribute this adaptation under the repository license. If the benchmark is public, add it to THIRD_PARTY_NOTICES.md with its license.

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.

2 participants