diff --git a/.claude/docs/REPO_WALKTHROUGH.md b/.claude/docs/REPO_WALKTHROUGH.md index cc87eb77c7..63548b51b1 100644 --- a/.claude/docs/REPO_WALKTHROUGH.md +++ b/.claude/docs/REPO_WALKTHROUGH.md @@ -119,6 +119,7 @@ envs/ │ ├── harbor_env/ # Thin Harbor package and trainer session factory ├── thinkingbox_env/ # Stateful MCP business-workflow benchmark adapter +├── posttrain_env/ # Post-training through brokered, budgeted HF Jobs (RFC 014) ├── coding_env/ # Python code execution environment ├── chat_env/ # Conversational environment ├── textarena_env/ # Text-based games (TextArena) diff --git a/docs/source/_toctree.yml b/docs/source/_toctree.yml index 6a835c0339..e60c30c4cb 100644 --- a/docs/source/_toctree.yml +++ b/docs/source/_toctree.yml @@ -155,6 +155,8 @@ title: Agent World Model - local: environments/pelican_svg title: Pelican SVG + - local: environments/posttrain + title: Post-training - local: environments/sophistry_bench_sprint title: Sophistry Bench Sprint title: Environments diff --git a/docs/source/environments.md b/docs/source/environments.md index 790f1cf2e6..febc884bde 100644 --- a/docs/source/environments.md +++ b/docs/source/environments.md @@ -294,6 +294,13 @@ The OpenEnv community has built a catalog of ready-to-run environments that cove 🤗 HF +
+
Post-training (experimental)
+

posttrain_env lets an agent post-train a model on HF Jobs through a broker that holds the credentials and enforces a per-rollout budget. A grader-only evaluation on a hidden split is the reward. Eval-first prototype for RFC 014; private deployments only, excluded from Spaces releases and HF collections.

+
+ 📄 Docs +
+
Sophistry Bench Sprint

sophistry_bench_sprint_env is a single-turn advocacy reward-hacking environment on QuALITY passages: the policy defends an assigned answer and the reward proxy peaks at 8 <claim> tags, with four weight-0 canaries that detect format hacking.

diff --git a/docs/source/environments/posttrain.md b/docs/source/environments/posttrain.md new file mode 100644 index 0000000000..6949a8915f --- /dev/null +++ b/docs/source/environments/posttrain.md @@ -0,0 +1,262 @@ + +# Post-training Environment (experimental) + +`posttrain_env` is an MCP environment where the agent post-trains a model on +[HF Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs). It is the +experimental, eval-first prototype for [RFC 014: Environment Services +(PR #1508)](https://github.com/openenv/OpenEnv/pull/1508); discussion in +[#1382](https://github.com/openenv/OpenEnv/issues/1382). + +The agent never holds a credential. It submits training scripts through MCP +tools, and the server's broker starts them on HF Jobs with its own token. Each +job's worst-case cost is reserved against a per-rollout budget before it starts +and settled at the real cost when it ends. `dev_evaluate` scores a checkpoint +on a public dev split. `submit` scores it on a hidden split in a grader-only +job, and that score is the reward. + +## Tools + +| Tool | What it does | +|---|---| +| `budget()` | Remaining budget, allowed hardware with its price, jobs and dev evaluations left | +| `jobs_submit(script, flavor, timeout_minutes, dependencies=None, datasets=None, idempotency_key=None)` | Run a Python script as an HF Job, like `hf jobs uv run`, without network access. Reserves `price/min × timeout` first | +| `jobs_status(job_id)` | State, cost once settled, artifact files of a succeeded job, result of a finished evaluation | +| `jobs_wait(job_id, timeout_s=60)` | Wait up to 300 s for a job to finish | +| `jobs_logs(job_id, tail=100)` | Last lines of one of your own jobs' logs | +| `jobs_cancel(job_id)` | Cancel a job; it settles for the minutes it ran | +| `dev_evaluate(job_id)` | Score a succeeded job's artifact on the dev split, as a new job charged to your budget | +| `submit(job_id)` | Final answer: cancels your other jobs, scores the artifact on the hidden split, ends the episode | + +Errors come back as data, with a type the agent can act on: `budget_exhausted`, +`flavor_not_allowed`, `timeout_not_allowed`, `invalid_script`, +`invalid_dependency`, `dataset_not_allowed`, `limit_reached`, `not_ready`, +`unknown_job`, `not_allowed`, `episode_over` or `hf_error`. + +## How a job runs + +Agent scripts run with `HfApi.run_uv_job`, following its usual behaviour: the +script is uploaded to the namespace's private `jobs-artifacts` bucket, and the +job's folder there is mounted read-write at `/data`. Anything the job writes +under `/data/outputs/` is its artifact: + +``` +openenv-artifact:////outputs +``` + +Jobs are started with an explicit flavor, namespace, timeout and labels, and +with no secrets. The broker never passes its own token into a job: a job holding +a "Start and manage Jobs" token could start further jobs that skip the budget. +It also never shells out to `hf jobs uv run`, whose script header can pull the +caller's `HF_TOKEN` into the job. + +### No network in agent jobs + +A job that can reach the internet can download a model already trained for the +task, or the evaluation's own test data. So agent jobs get their inputs as +read-only mounts, which the broker checks and the episode records, and run +without network access: + +- The task's base model is mounted at `/base_model`, at a pinned revision. +- Each dataset in `datasets` is mounted at `/datasets/`, at its current + revision. It must be public and ungated, and not on the task's denylist. The + broker looks it up without its token, so the token can't open private repos. +- Packages from `dependencies` and the script's PEP 723 header are installed + into a fresh venv before the script runs, from PyPI wheels only. Installing a + wheel runs none of its code, so no package code runs while the job still has + network access. URLs and git dependencies are refused. Other settings in the + script's header are ignored. + +Mounts and wheels narrow what a job can fetch, but a public dataset or a PyPI +wheel can still hold weights or test data. Tasks whose test data is public also +need a contamination check after the run. + +HF Jobs has no egress control, and its containers can't create network +namespaces or firewall rules. So the broker sends +[`server/job_launcher.py`](https://github.com/huggingface/OpenEnv/blob/main/envs/posttrain_env/server/job_launcher.py) as the job's script, with the +agent's script and packages embedded. Once it has installed the packages, the +launcher installs a seccomp filter that refuses every socket except `AF_UNIX`, +and io_uring and cross-process memory access; systemd's +`RestrictAddressFamilies=` works the same way. It then closes inherited file +descriptors, switches to a uid no other process in the job uses, and runs the +agent's script with `HF_HUB_OFFLINE=1`. The filter can't be removed and covers every child +process. If any step fails, the script doesn't run. Training on a single GPU, +DataLoader workers and writing `/data/outputs/` work under the filter; +`torch.distributed` over localhost TCP does not. + +### Evaluations + +Evaluations run the task's evaluator script as their own job with that folder +mounted read-only at `/model`, and the task's base model at `/base_model`. The +evaluator prints one `POSTTRAIN_SCORE {...}` line. A missing line, or a score +outside [0, 1], scores 0 as `invalid`. Evaluators take only weights from the +artifact, loaded as safetensors. The tokenizer, chat template, config and +generation settings come from the base model mount, because a chat template is +a program that sees the question and could write the answer into the prompt. Each evaluator pins its +dependencies in its PEP 723 header, with an `exclude-newer` date for the rest. + +## Budget and spend + +Money is tracked in micro-USD from the prices `list_jobs_hardware` reports. +A job settles at its price per minute times its running minutes, rounded up to +whole minutes. Grader evaluations are not charged to the agent. + +Every observation carries the episode's spend, for accounting and for budgets, +concurrency and rollout allocation. It is not a sampling weight (RFC 013), and +an episode that ends unscored still reports it: + +```python +obs.metadata["service_spend"] +# {"version": 1, +# "settled": {"usd": 0.000334, "gpu_s": 0.0, "jobs": 1, "calls": 1}, +# "reserved": {"usd": 0.0}, +# "final": True} +``` + +`jobs` counts agent jobs and `calls` counts dev evaluations. `final` is false +while a job is still running or a cancellation was never confirmed. + +## Episode end + +Through `step(CallToolAction(...))`, an episode ends when the agent calls +`submit`, after `max_steps` tool calls (default 100) or after the task's time limit. +The step and time limits give reward 0. + +Direct `/mcp` calls, including `env.call_tool`, bypass the step and time limits +and do not return the final step observation with reward and `service_spend`. + +If the score job itself fails, times out or is cancelled, the episode ends +unscored: `reward` is `None` and `episode_end.reason` is `score_failed`, since +the fault is the grader's. Trainers should drop such a rollout, not count it as +zero, subject to acceptance of RFC 013 (#1256) and its final unscored-rollout +semantics. Ending the episode cancels every live job. `reset` cancels the previous +episode's jobs first and raises if HF Jobs does not confirm a cancellation, +since that job may still be running and billing. + +## Using it with a trainer + +This environment is eval-first today. Post-training episodes can take hours, +beyond the staleness window of asynchronous GRPO trainers. Use it for evaluation; +training integration needs a service that fits the trainer's staleness window or +a stand-in for the slow service. + +`submit` blocks for up to `score_timeout_s`, reported in reset metadata (2700 s +for `t0_arithmetic_format`). Size the client's tool/message timeout above that +value, with room for cancellation and transport overhead; the examples use +`message_timeout_s=3600`. + +Read `final.reward` and `final.observation.metadata["service_spend"]` from the +final `step` result. For RFC 013 integration, forward one settled amount per +rollout, for example `service_spend["settled"]["usd"]`, as `costs` in the trainer's +chosen unit. Spend is accounting, never a sampling weight. `final: true` means +all jobs have settled, even mid-episode; it does not mean the episode is done. +Use the step result's `done` for that. An orphan keeps spend `final: false` and +may still bill; retain that accounting even when reward is `None`. + +## Tasks + +| Task | What the agent does | Hardware | Budget | +|---|---|---|---| +| `t0_artifact` | Write the episode's nonce to `/data/outputs/answer.txt`. Checks the plumbing | `cpu-basic` | $0.05 | +| `t0_arithmetic_format` | Fine-tune `HuggingFaceTB/SmolLM2-135M-Instruct` to answer two-digit addition as `#### {a+b}`. Reward is exact-match accuracy on a hidden test split | `cpu-upgrade`, `t4-small`, `a10g-small` | $3.00 | + +## Running the server + +```bash +export HF_TOKEN=... # the broker's token, see below +cd envs/posttrain_env +uv run --project . server +``` + +The local entry point binds `127.0.0.1:8000`. The Docker command explicitly +binds `0.0.0.0:8000` so a private host can reach the published port. This +experimental env is excluded from automatic Spaces releases via +`SKIP_HF_DEPLOYMENT`; do not publish it to Spaces or add it to HF collections. + +Configuration, all read by the server only: + +| Variable | Meaning | +|---|---| +| `HF_TOKEN` | Broker token. A fine-grained token with "Start and manage Jobs" plus write access to `/jobs-artifacts`. A Jobs-only token is refused when the script is uploaded | +| `POSTTRAIN_NAMESPACE` | Namespace jobs run and bill in. Defaults to the token's owner | +| `POSTTRAIN_RESOURCE_GROUP_ID` | Resource Group to bill. Its Jobs spend limit caps total spend even if the ledger is wrong | +| `POSTTRAIN_TEST_SEED` | Seed of the hidden test split. Set it for any run whose scores you compare: if unset it is random per process (the server logs a warning), so scores are only comparable within one server process. Keep it secret from the agent | + +Run the server where its token is not shared with the agent, for example a +container you control. Do not deploy it as a public Space: a Space gives no +isolation guarantee for the credential it holds. + +## Client + +```python +from openenv.core.env_server.mcp_types import CallToolAction +from posttrain_env import PostTrainEnv + +with PostTrainEnv( + base_url="http://localhost:8000", message_timeout_s=3600 +).sync() as env: + result = env.reset(task_id="t0_artifact") + print(result.observation.metadata["prompt"]) + result = env.step( + CallToolAction( + tool_name="jobs_submit", + arguments={ + "script": open("train.py").read(), + "flavor": "cpu-basic", + "timeout_minutes": 5, + }, + ) + ) + assert result.observation.error is None + job = result.observation.result["data"] + assert "error" not in job, job + status = "queued" + while status in {"queued", "running"}: + result = env.step( + CallToolAction( + tool_name="jobs_wait", + arguments={ + "job_id": job["job_id"], + "timeout_s": 300, + }, + ) + ) + assert not result.done, result.observation.metadata + status = result.observation.result["data"]["state"] + assert status == "succeeded", result.observation.result + final = env.step( + CallToolAction( + tool_name="submit", + arguments={ + "job_id": job["job_id"], + }, + ) + ) + print(final.done, final.reward) + print(final.observation.metadata["service_spend"]) +``` + +[`examples/posttrain_env_agent.py`](https://github.com/huggingface/OpenEnv/blob/main/examples/posttrain_env_agent.py) +plays a whole episode with an LLM through any OpenAI-compatible endpoint, +Hugging Face Inference Providers by default. The agent's API key is its own; +it never needs the broker's token. + +## Known gaps + +- The server has no authentication of its own: anyone who can reach it can + start episodes and spend its budget. Keep it on localhost or a private network. +- Job folders stay in the `jobs-artifacts` bucket, and the size of what a job + writes there is not capped or charged to the budget. +- `max_steps` and the time limit are checked in `step`; a client calling the + MCP endpoint directly bypasses them, though not the budget. Core enforcement + is tracked in [#1507](https://github.com/openenv/OpenEnv/issues/1507). +- `torch.distributed` over localhost TCP is blocked in agent jobs, so tasks use + single-GPU flavors. + +## Tests + +```bash +PYTHONPATH=src:envs uv run pytest tests/envs/test_posttrain_environment.py -v +``` + +The tests replace HF Jobs with an in-memory fake, so they make no network calls +and spend nothing. diff --git a/envs/posttrain_env/README.md b/envs/posttrain_env/README.md new file mode 100644 index 0000000000..309e24c184 --- /dev/null +++ b/envs/posttrain_env/README.md @@ -0,0 +1,261 @@ +# Post-training Environment (experimental) + +`posttrain_env` is an MCP environment where the agent post-trains a model on +[HF Jobs](https://huggingface.co/docs/huggingface_hub/guides/jobs). It is the +experimental, eval-first prototype for [RFC 014: Environment Services +(PR #1508)](https://github.com/openenv/OpenEnv/pull/1508); discussion in +[#1382](https://github.com/openenv/OpenEnv/issues/1382). + +The agent never holds a credential. It submits training scripts through MCP +tools, and the server's broker starts them on HF Jobs with its own token. Each +job's worst-case cost is reserved against a per-rollout budget before it starts +and settled at the real cost when it ends. `dev_evaluate` scores a checkpoint +on a public dev split. `submit` scores it on a hidden split in a grader-only +job, and that score is the reward. + +## Tools + +| Tool | What it does | +|---|---| +| `budget()` | Remaining budget, allowed hardware with its price, jobs and dev evaluations left | +| `jobs_submit(script, flavor, timeout_minutes, dependencies=None, datasets=None, idempotency_key=None)` | Run a Python script as an HF Job, like `hf jobs uv run`, without network access. Reserves `price/min × timeout` first | +| `jobs_status(job_id)` | State, cost once settled, artifact files of a succeeded job, result of a finished evaluation | +| `jobs_wait(job_id, timeout_s=60)` | Wait up to 300 s for a job to finish | +| `jobs_logs(job_id, tail=100)` | Last lines of one of your own jobs' logs | +| `jobs_cancel(job_id)` | Cancel a job; it settles for the minutes it ran | +| `dev_evaluate(job_id)` | Score a succeeded job's artifact on the dev split, as a new job charged to your budget | +| `submit(job_id)` | Final answer: cancels your other jobs, scores the artifact on the hidden split, ends the episode | + +Errors come back as data, with a type the agent can act on: `budget_exhausted`, +`flavor_not_allowed`, `timeout_not_allowed`, `invalid_script`, +`invalid_dependency`, `dataset_not_allowed`, `limit_reached`, `not_ready`, +`unknown_job`, `not_allowed`, `episode_over` or `hf_error`. + +## How a job runs + +Agent scripts run with `HfApi.run_uv_job`, following its usual behaviour: the +script is uploaded to the namespace's private `jobs-artifacts` bucket, and the +job's folder there is mounted read-write at `/data`. Anything the job writes +under `/data/outputs/` is its artifact: + +``` +openenv-artifact:////outputs +``` + +Jobs are started with an explicit flavor, namespace, timeout and labels, and +with no secrets. The broker never passes its own token into a job: a job holding +a "Start and manage Jobs" token could start further jobs that skip the budget. +It also never shells out to `hf jobs uv run`, whose script header can pull the +caller's `HF_TOKEN` into the job. + +### No network in agent jobs + +A job that can reach the internet can download a model already trained for the +task, or the evaluation's own test data. So agent jobs get their inputs as +read-only mounts, which the broker checks and the episode records, and run +without network access: + +- The task's base model is mounted at `/base_model`, at a pinned revision. +- Each dataset in `datasets` is mounted at `/datasets/`, at its current + revision. It must be public and ungated, and not on the task's denylist. The + broker looks it up without its token, so the token can't open private repos. +- Packages from `dependencies` and the script's PEP 723 header are installed + into a fresh venv before the script runs, from PyPI wheels only. Installing a + wheel runs none of its code, so no package code runs while the job still has + network access. URLs and git dependencies are refused. Other settings in the + script's header are ignored. + +Mounts and wheels narrow what a job can fetch, but a public dataset or a PyPI +wheel can still hold weights or test data. Tasks whose test data is public also +need a contamination check after the run. + +HF Jobs has no egress control, and its containers can't create network +namespaces or firewall rules. So the broker sends +[`server/job_launcher.py`](server/job_launcher.py) as the job's script, with the +agent's script and packages embedded. Once it has installed the packages, the +launcher installs a seccomp filter that refuses every socket except `AF_UNIX`, +and io_uring and cross-process memory access; systemd's +`RestrictAddressFamilies=` works the same way. It then closes inherited file +descriptors, switches to a uid no other process in the job uses, and runs the +agent's script with `HF_HUB_OFFLINE=1`. The filter can't be removed and covers every child +process. If any step fails, the script doesn't run. Training on a single GPU, +DataLoader workers and writing `/data/outputs/` work under the filter; +`torch.distributed` over localhost TCP does not. + +### Evaluations + +Evaluations run the task's evaluator script as their own job with that folder +mounted read-only at `/model`, and the task's base model at `/base_model`. The +evaluator prints one `POSTTRAIN_SCORE {...}` line. A missing line, or a score +outside [0, 1], scores 0 as `invalid`. Evaluators take only weights from the +artifact, loaded as safetensors. The tokenizer, chat template, config and +generation settings come from the base model mount, because a chat template is +a program that sees the question and could write the answer into the prompt. Each evaluator pins its +dependencies in its PEP 723 header, with an `exclude-newer` date for the rest. + +## Budget and spend + +Money is tracked in micro-USD from the prices `list_jobs_hardware` reports. +A job settles at its price per minute times its running minutes, rounded up to +whole minutes. Grader evaluations are not charged to the agent. + +Every observation carries the episode's spend, for accounting and for budgets, +concurrency and rollout allocation. It is not a sampling weight (RFC 013), and +an episode that ends unscored still reports it: + +```python +obs.metadata["service_spend"] +# {"version": 1, +# "settled": {"usd": 0.000334, "gpu_s": 0.0, "jobs": 1, "calls": 1}, +# "reserved": {"usd": 0.0}, +# "final": True} +``` + +`jobs` counts agent jobs and `calls` counts dev evaluations. `final` is false +while a job is still running or a cancellation was never confirmed. + +## Episode end + +Through `step(CallToolAction(...))`, an episode ends when the agent calls +`submit`, after `max_steps` tool calls (default 100) or after the task's time limit. +The step and time limits give reward 0. + +Direct `/mcp` calls, including `env.call_tool`, bypass the step and time limits +and do not return the final step observation with reward and `service_spend`. + +If the score job itself fails, times out or is cancelled, the episode ends +unscored: `reward` is `None` and `episode_end.reason` is `score_failed`, since +the fault is the grader's. Trainers should drop such a rollout, not count it as +zero, subject to acceptance of RFC 013 (#1256) and its final unscored-rollout +semantics. Ending the episode cancels every live job. `reset` cancels the previous +episode's jobs first and raises if HF Jobs does not confirm a cancellation, +since that job may still be running and billing. + +## Using it with a trainer + +This environment is eval-first today. Post-training episodes can take hours, +beyond the staleness window of asynchronous GRPO trainers. Use it for evaluation; +training integration needs a service that fits the trainer's staleness window or +a stand-in for the slow service. + +`submit` blocks for up to `score_timeout_s`, reported in reset metadata (2700 s +for `t0_arithmetic_format`). Size the client's tool/message timeout above that +value, with room for cancellation and transport overhead; the examples use +`message_timeout_s=3600`. + +Read `final.reward` and `final.observation.metadata["service_spend"]` from the +final `step` result. For RFC 013 integration, forward one settled amount per +rollout, for example `service_spend["settled"]["usd"]`, as `costs` in the trainer's +chosen unit. Spend is accounting, never a sampling weight. `final: true` means +all jobs have settled, even mid-episode; it does not mean the episode is done. +Use the step result's `done` for that. An orphan keeps spend `final: false` and +may still bill; retain that accounting even when reward is `None`. + +## Tasks + +| Task | What the agent does | Hardware | Budget | +|---|---|---|---| +| `t0_artifact` | Write the episode's nonce to `/data/outputs/answer.txt`. Checks the plumbing | `cpu-basic` | $0.05 | +| `t0_arithmetic_format` | Fine-tune `HuggingFaceTB/SmolLM2-135M-Instruct` to answer two-digit addition as `#### {a+b}`. Reward is exact-match accuracy on a hidden test split | `cpu-upgrade`, `t4-small`, `a10g-small` | $3.00 | + +## Running the server + +```bash +export HF_TOKEN=... # the broker's token, see below +cd envs/posttrain_env +uv run --project . server +``` + +The local entry point binds `127.0.0.1:8000`. The Docker command explicitly +binds `0.0.0.0:8000` so a private host can reach the published port. This +experimental env is excluded from automatic Spaces releases via +`SKIP_HF_DEPLOYMENT`; do not publish it to Spaces or add it to HF collections. + +Configuration, all read by the server only: + +| Variable | Meaning | +|---|---| +| `HF_TOKEN` | Broker token. A fine-grained token with "Start and manage Jobs" plus write access to `/jobs-artifacts`. A Jobs-only token is refused when the script is uploaded | +| `POSTTRAIN_NAMESPACE` | Namespace jobs run and bill in. Defaults to the token's owner | +| `POSTTRAIN_RESOURCE_GROUP_ID` | Resource Group to bill. Its Jobs spend limit caps total spend even if the ledger is wrong | +| `POSTTRAIN_TEST_SEED` | Seed of the hidden test split. Set it for any run whose scores you compare: if unset it is random per process (the server logs a warning), so scores are only comparable within one server process. Keep it secret from the agent | + +Run the server where its token is not shared with the agent, for example a +container you control. Do not deploy it as a public Space: a Space gives no +isolation guarantee for the credential it holds. + +## Client + +```python +from openenv.core.env_server.mcp_types import CallToolAction +from posttrain_env import PostTrainEnv + +with PostTrainEnv( + base_url="http://localhost:8000", message_timeout_s=3600 +).sync() as env: + result = env.reset(task_id="t0_artifact") + print(result.observation.metadata["prompt"]) + result = env.step( + CallToolAction( + tool_name="jobs_submit", + arguments={ + "script": open("train.py").read(), + "flavor": "cpu-basic", + "timeout_minutes": 5, + }, + ) + ) + assert result.observation.error is None + job = result.observation.result["data"] + assert "error" not in job, job + status = "queued" + while status in {"queued", "running"}: + result = env.step( + CallToolAction( + tool_name="jobs_wait", + arguments={ + "job_id": job["job_id"], + "timeout_s": 300, + }, + ) + ) + assert not result.done, result.observation.metadata + status = result.observation.result["data"]["state"] + assert status == "succeeded", result.observation.result + final = env.step( + CallToolAction( + tool_name="submit", + arguments={ + "job_id": job["job_id"], + }, + ) + ) + print(final.done, final.reward) + print(final.observation.metadata["service_spend"]) +``` + +[`examples/posttrain_env_agent.py`](../../examples/posttrain_env_agent.py) +plays a whole episode with an LLM through any OpenAI-compatible endpoint, +Hugging Face Inference Providers by default. The agent's API key is its own; +it never needs the broker's token. + +## Known gaps + +- The server has no authentication of its own: anyone who can reach it can + start episodes and spend its budget. Keep it on localhost or a private network. +- Job folders stay in the `jobs-artifacts` bucket, and the size of what a job + writes there is not capped or charged to the budget. +- `max_steps` and the time limit are checked in `step`; a client calling the + MCP endpoint directly bypasses them, though not the budget. Core enforcement + is tracked in [#1507](https://github.com/openenv/OpenEnv/issues/1507). +- `torch.distributed` over localhost TCP is blocked in agent jobs, so tasks use + single-GPU flavors. + +## Tests + +```bash +PYTHONPATH=src:envs uv run pytest tests/envs/test_posttrain_environment.py -v +``` + +The tests replace HF Jobs with an in-memory fake, so they make no network calls +and spend nothing. diff --git a/envs/posttrain_env/SKIP_HF_DEPLOYMENT b/envs/posttrain_env/SKIP_HF_DEPLOYMENT new file mode 100644 index 0000000000..6a16aa4e88 --- /dev/null +++ b/envs/posttrain_env/SKIP_HF_DEPLOYMENT @@ -0,0 +1,2 @@ +Experimental credential-holding broker: private deployments only. +Do not release to Spaces or add to HF collections (RFC 014). diff --git a/envs/posttrain_env/__init__.py b/envs/posttrain_env/__init__.py new file mode 100644 index 0000000000..c62dfe96d8 --- /dev/null +++ b/envs/posttrain_env/__init__.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Post-training environment for OpenEnv. + +The agent post-trains a model by submitting scripts to HF Jobs through a broker +that holds the credentials and enforces a per-rollout budget. A grader-only +evaluation of the submitted checkpoint, on a hidden split, is the reward. +""" + +from openenv.core.env_server.mcp_types import CallToolAction, ListToolsAction + +from .client import PostTrainEnv + +__all__ = ["PostTrainEnv", "CallToolAction", "ListToolsAction"] diff --git a/envs/posttrain_env/client.py b/envs/posttrain_env/client.py new file mode 100644 index 0000000000..5f4b2ae02a --- /dev/null +++ b/envs/posttrain_env/client.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Client for the post-training environment.""" + +from openenv.core.mcp_client import MCPToolClient + + +class PostTrainEnv(MCPToolClient): + """ + Client for the post-training environment. + + Dispatch tools with `step(CallToolAction(...))` to count episode steps and + receive the final reward and `service_spend`. Direct `call_tool` bypasses + episode step/time limits and does not return that final observation. + Available tools are: `budget`, `jobs_submit`, + `jobs_status`, `jobs_wait`, `jobs_logs`, `jobs_cancel`, `dev_evaluate` and + `submit`. `submit` ends the episode and can take several minutes, so give + the client a message timeout above the `score_timeout_s` reported by + `reset`, with room for cancellation and transport overhead (2700 s for + `t0_arithmetic_format`; the example uses 3600 s). + + Examples: + + ```python + from openenv.core.env_server.mcp_types import CallToolAction + from posttrain_env import PostTrainEnv + + with PostTrainEnv(base_url="http://localhost:8000", message_timeout_s=3600).sync() as env: + result = env.reset(task_id="t0_artifact") + print(result.observation.metadata["prompt"]) + result = env.step(CallToolAction(tool_name="jobs_submit", arguments={ + "script": script, "flavor": "cpu-basic", "timeout_minutes": 5, + })) + assert result.observation.error is None + job = result.observation.result["data"] + assert "error" not in job, job + status = "queued" + while status in {"queued", "running"}: + result = env.step(CallToolAction(tool_name="jobs_wait", arguments={ + "job_id": job["job_id"], "timeout_s": 300, + })) + assert not result.done, result.observation.metadata + status = result.observation.result["data"]["state"] + assert status == "succeeded", result.observation.result + final = env.step(CallToolAction(tool_name="submit", arguments={ + "job_id": job["job_id"], + })) + print(final.reward, final.observation.metadata["service_spend"]) + ``` + """ + + pass diff --git a/envs/posttrain_env/openenv.yaml b/envs/posttrain_env/openenv.yaml new file mode 100644 index 0000000000..10c8317f99 --- /dev/null +++ b/envs/posttrain_env/openenv.yaml @@ -0,0 +1,42 @@ +spec_version: 1 +name: posttrain_env +runtime: fastapi +app: server.app:app +port: 8000 + +description: > + The agent post-trains a model by submitting scripts to HF Jobs through a + broker that holds the credentials and enforces a per-rollout budget. A + grader-only evaluation of the submitted checkpoint, on a hidden split, is the + reward. Prototype for RFC 014 (Environment Services). +version: "0.1.0" + +# Run this server where its HF token is not shared with the agent: a container +# or host you control. Do not deploy it as a public Space; a Space gives no +# isolation guarantee for the credential it holds (RFC 014, decision 2). +# +# Server configuration, never stored in this file: +# HF_TOKEN broker token: "Start and manage Jobs" plus +# write access to /jobs-artifacts +# POSTTRAIN_NAMESPACE namespace jobs run in (default: token owner) +# POSTTRAIN_RESOURCE_GROUP_ID optional Resource Group with a Jobs spend limit +# POSTTRAIN_TEST_SEED hidden test split seed (default: random per process) + +validation: + reward: + range: [0.0, 1.0] + oracle_tolerance: 0.0 + floor_margin: 0.5 + # The server only brokers; training and evaluation run on HF Jobs. + resources: + cpu: 1.0 + memory_mb: 1024 + disk_mb: 512 + episode_timeout_s: 3600.0 # the policy maximum; t0_artifact fits, longer tasks set their own limit + capabilities: + verifier: + kind: reward_channel + declared_tools: + [budget, jobs_submit, jobs_status, jobs_wait, jobs_logs, jobs_cancel, dev_evaluate, submit] + types: + tags: [post-training, hf-jobs, environment-services] diff --git a/envs/posttrain_env/pyproject.toml b/envs/posttrain_env/pyproject.toml new file mode 100644 index 0000000000..0aff8c897b --- /dev/null +++ b/envs/posttrain_env/pyproject.toml @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: BSD-3-Clause + +[build-system] +requires = ["setuptools>=45", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "openenv-posttrain-env" +version = "0.1.0" +description = "Post-training environment for OpenEnv: train models through brokered, budgeted HF Jobs" +requires-python = ">=3.11" +dependencies = [ + "openenv>=0.9.0", + "fastapi>=0.115.0", + "pydantic>=2.0.0", + "uvicorn>=0.24.0", + # Jobs, job artifact buckets and volume mounts. + "huggingface_hub>=2.3.0", + # Checks the requirements agents pass to their jobs. + "packaging>=23.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=9.0.3", + "pytest-cov>=4.0.0", +] + +[project.scripts] +server = "posttrain_env.server.app:main" + +[tool.setuptools] +include-package-data = true +packages = ["posttrain_env", "posttrain_env.server"] +package-dir = { "posttrain_env" = ".", "posttrain_env.server" = "server" } + +[tool.setuptools.package-data] +"posttrain_env.server" = ["eval_scripts/*.py"] diff --git a/envs/posttrain_env/server/Dockerfile b/envs/posttrain_env/server/Dockerfile new file mode 100644 index 0000000000..12c57b41a5 --- /dev/null +++ b/envs/posttrain_env/server/Dockerfile @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: BSD-3-Clause + +# Multi-stage build using openenv-base +# This Dockerfile is flexible and works for both: +# - In-repo environments (with local src/core) +# - Standalone environments (with openenv from pip) +# The build script (openenv build) handles context detection and sets appropriate build args. + +ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest +FROM ${BASE_IMAGE} AS builder + +WORKDIR /app + +# Build argument to control whether we're building standalone or in-repo +ARG BUILD_MODE=in-repo + +# Copy environment code (always at root of build context) +COPY . /app/env + +# For in-repo builds, openenv is already in the pyproject.toml dependencies +# For standalone builds, openenv will be installed from pip via pyproject.toml +WORKDIR /app/env + +# Install dependencies using uv sync +# First pass: install dependencies without the project (for better caching) +# Second pass: install the project itself +RUN --mount=type=cache,target=/root/.cache/uv \ + if [ -f uv.lock ]; then \ + uv sync --frozen --no-install-project --no-editable; \ + else \ + uv sync --no-install-project --no-editable; \ + fi + +RUN --mount=type=cache,target=/root/.cache/uv \ + if [ -f uv.lock ]; then \ + uv sync --frozen --no-editable; \ + else \ + uv sync --no-editable; \ + fi + +# Final runtime stage +FROM ${BASE_IMAGE} + +WORKDIR /app + +# Copy the virtual environment from builder +COPY --from=builder /app/env/.venv /app/.venv + +# Copy the environment code +COPY --from=builder /app/env /app/env + +# Set PATH to use the virtual environment +ENV PATH="/app/.venv/bin:$PATH" + +# Set PYTHONPATH so imports work correctly +ENV PYTHONPATH="/app/env:$PYTHONPATH" +ENV ENABLE_WEB_INTERFACE=false + +# Health check using Python (more portable than curl/wget) +HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \ + CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1 + +# Run the FastAPI server +# The module path is constructed to work with the /app/env structure +CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"] diff --git a/envs/posttrain_env/server/__init__.py b/envs/posttrain_env/server/__init__.py new file mode 100644 index 0000000000..141022a536 --- /dev/null +++ b/envs/posttrain_env/server/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""Post-training environment server components.""" + +from .posttrain_environment import PostTrainEnvironment + +__all__ = ["PostTrainEnvironment"] diff --git a/envs/posttrain_env/server/app.py b/envs/posttrain_env/server/app.py new file mode 100644 index 0000000000..20ee94bcf3 --- /dev/null +++ b/envs/posttrain_env/server/app.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""FastAPI application for the post-training environment.""" + +import os + +from openenv.core.env_server.http_server import create_app +from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation +from openenv.core.env_server.types import Observation + +from .posttrain_environment import PostTrainEnvironment + +# The class is passed rather than an instance so each WebSocket session gets +# its own environment, broker and budget. +app = create_app( + PostTrainEnvironment, + CallToolAction, + CallToolObservation, + reset_observation_cls=Observation, + env_name="posttrain_env", + max_concurrent_envs=int(os.getenv("MAX_CONCURRENT_ENVS", "8")), +) + + +def main(): + import uvicorn + + uvicorn.run(app, host="127.0.0.1", port=8000) + + +if __name__ == "__main__": + main() diff --git a/envs/posttrain_env/server/broker.py b/envs/posttrain_env/server/broker.py new file mode 100644 index 0000000000..ea1cc4899c --- /dev/null +++ b/envs/posttrain_env/server/broker.py @@ -0,0 +1,848 @@ +# SPDX-License-Identifier: BSD-3-Clause + +"""HF Jobs broker with a reserve-then-settle budget ledger. + +The broker is the only holder of the HF token. Agents never see it: jobs are +started with `run_uv_job` and explicit flavor, namespace, timeout and labels, +and with no secrets. A `job.write` token passed into a job could start further +jobs that skip this ledger, so the broker never passes one in. + +Agent scripts run inside `job_launcher.py`, which installs their dependencies +from wheels and cuts the job's network access before the script starts. Inputs +come as read-only mounts: the task's base model and the public datasets the +agent asks for. + +Money is tracked in integer micro-USD. A job reserves its worst case (the +flavor's per-minute rate times its timeout) before it starts, and settles at +the per-minute rate for the minutes it actually ran. +""" + +import base64 +import json +import math +import re +import tempfile +import threading +import time +import tomllib +from contextlib import suppress +from dataclasses import dataclass +from pathlib import Path +from typing import Any +from uuid import uuid4 + +from huggingface_hub import Volume +from huggingface_hub.errors import HfHubHTTPError, HFValidationError +from huggingface_hub.utils import validate_repo_id +from packaging.requirements import InvalidRequirement, Requirement + +# HF Jobs stage -> job state exposed to the agent. +STATE_FROM_STAGE = { + "SCHEDULING": "queued", + "RUNNING": "running", + "COMPLETED": "succeeded", + "ERROR": "failed", + "CANCELED": "cancelled", + "DELETED": "failed", +} +# States HF Jobs has confirmed; "orphaned" is terminal for the ledger but not confirmed. +CONFIRMED_STATES = {"succeeded", "failed", "cancelled"} +TERMINAL_STATES = CONFIRMED_STATES | {"orphaned"} +TERMINAL_STAGES = {"COMPLETED", "ERROR", "CANCELED", "DELETED"} +SCORE_MARKER = "POSTTRAIN_SCORE " +ARTIFACTS_MOUNT = "/data" +MAX_SCRIPT_BYTES = 200_000 +MAX_DEPENDENCIES = 20 +MAX_DATASETS = 5 +DATASETS_MOUNT = "/datasets" +LAUNCHER = Path(__file__).parent / "job_launcher.py" +# A PEP 723 `script` block, as in the specification's reference regex. +SCRIPT_BLOCK = re.compile(r"(?m)^# /// script$\s(?P(^#(| .*)$\s)+)^# ///$") +MAX_ARTIFACT_FILES = 50 +# Logs can lag a finished job; re-read before declaring the score missing. +MAX_RESULT_READS = 3 +ADMISSION_TIMEOUT_S = 120.0 + + +class BrokerError(Exception): + """A request the broker refuses, with a typed reason for the agent.""" + + def __init__(self, kind: str, message: str): + super().__init__(message) + self.kind = kind + self.message = message + + def as_dict(self) -> dict: + return {"error": {"type": self.kind, "message": self.message}} + + +@dataclass +class Job: + """One HF Job started by the broker, as tracked in the ledger.""" + + job_id: str + role: str # "agent", "dev_eval" or "score" + flavor: str + timeout_minutes: int + reserved_micro_usd: int + bucket: str | None = None + path: str | None = None + source_job_id: str | None = None + state: str = "queued" + reason: str | None = None + running_secs: int | None = None + settled_micro_usd: int = 0 + settled: bool = False + result: dict | None = None + result_reads: int = 0 + + def view(self) -> dict: + view = { + "job_id": self.job_id, + "role": self.role, + "state": self.state, + "flavor": self.flavor, + "timeout_minutes": self.timeout_minutes, + "running_secs": self.running_secs or 0, + } + if self.reason: + view["reason"] = self.reason + if self.source_job_id: + view["evaluates"] = self.source_job_id + if self.settled: + view["cost_usd"] = self.settled_micro_usd / 1e6 + if self.result is not None: + view["result"] = self.result + return view + + +@dataclass +class PendingJob: + """A reserved request whose remote job identity is not yet confirmed.""" + + role: str + reservation: int + flavor: str + timeout_minutes: int + idempotency_key: str | tuple[str] | None + active: bool = True + + +class Broker: + """ + Start, track and cancel HF Jobs for one episode, against a budget. + + Args: + api (`huggingface_hub.HfApi`): + Client holding the broker's token. + namespace (`str`): + Namespace jobs run and are billed in. + episode_id (`str`): + Episode the jobs belong to, added as a label. + budget_usd (`float`): + Budget for agent jobs and dev evaluations. + max_jobs (`int`): + Maximum number of agent jobs. + max_dev_evals (`int`): + Maximum number of dev evaluations. + resource_group_id (`str`, *optional*): + Resource Group to bill, so its Jobs spend limit is a hard backstop. + """ + + def __init__( + self, + api: Any, + *, + namespace: str, + episode_id: str, + budget_usd: float, + max_jobs: int, + max_dev_evals: int, + resource_group_id: str | None = None, + ): + self._api = api + self.namespace = namespace + self.episode_id = episode_id + self.budget_micro_usd = round(budget_usd * 1e6) + self.max_jobs = max_jobs + self.max_dev_evals = max_dev_evals + self.resource_group_id = resource_group_id + self.jobs: dict[str, Job] = {} + self._idempotency: dict[str | tuple[str], str] = {} + self._hardware: dict[str, Any] | None = None + self._admission = threading.Condition(threading.RLock()) + self._pending: dict[str, PendingJob] = {} + self._closing = False + + # -- ledger ------------------------------------------------------------- + + def count(self, role: str) -> int: + """Return how many jobs of `role` this episode has started.""" + with self._admission: + return sum(1 for j in self.jobs.values() if j.role == role) + sum( + 1 for pending in self._pending.values() if pending.role == role + ) + + def _charged(self) -> list[Job]: + return [j for j in self.jobs.values() if j.role in ("agent", "dev_eval")] + + def remaining_micro_usd(self) -> int: + with self._admission: + used = sum( + j.settled_micro_usd if j.settled else j.reserved_micro_usd + for j in self._charged() + ) + sum( + pending.reservation + for pending in self._pending.values() + if pending.role in ("agent", "dev_eval") + ) + return self.budget_micro_usd - used + + def rate_micro_usd_per_minute(self, flavor: str) -> int: + return self._hardware_info(flavor).unit_cost_micro_usd + + def _hardware_info(self, flavor: str) -> Any: + if self._hardware is None: + self._hardware = {h.name: h for h in self._api.list_jobs_hardware()} + if flavor not in self._hardware: + raise BrokerError("flavor_unknown", f"Unknown hardware flavor {flavor!r}") + return self._hardware[flavor] + + def _reserve(self, flavor: str, timeout_minutes: int, charged: bool) -> int: + reservation = self.rate_micro_usd_per_minute(flavor) * timeout_minutes + if charged and reservation > self.remaining_micro_usd(): + raise BrokerError( + "budget_exhausted", + f"Reserving ${reservation / 1e6:.4f} ({flavor} for {timeout_minutes} min) " + f"exceeds the remaining ${self.remaining_micro_usd() / 1e6:.4f}", + ) + return reservation + + def _settle(self, job: Job) -> None: + with self._admission: + if job.running_secs is None: + # A job that ended without a reported duration may still have run; + # charge its full reservation unless it never started. + cost = 0 if job.state == "cancelled" else job.reserved_micro_usd + else: + minutes = math.ceil(job.running_secs / 60) + cost = minutes * self.rate_micro_usd_per_minute(job.flavor) + job.settled_micro_usd = min(cost, job.reserved_micro_usd) + job.settled = True + + def spend(self) -> dict: + """Return the `service_spend` observation entry for this episode.""" + with self._admission: + charged = self._charged() + gpu_s = 0.0 + for job in charged: + accelerator = self._hardware_info(job.flavor).accelerator + if accelerator is not None: + gpu_s += (job.running_secs or 0) * int(accelerator.quantity) + return { + "version": 1, + "settled": { + "usd": sum(j.settled_micro_usd for j in charged if j.settled) / 1e6, + "gpu_s": gpu_s, + "jobs": self.count("agent"), + "calls": self.count("dev_eval"), + }, + "reserved": { + "usd": ( + sum(j.reserved_micro_usd for j in charged if not j.settled) + + sum( + pending.reservation + for pending in self._pending.values() + if pending.role in ("agent", "dev_eval") + ) + ) + / 1e6 + }, + "final": not self._pending + and all(j.state in CONFIRMED_STATES for j in self.jobs.values()), + } + + # -- jobs --------------------------------------------------------------- + + def submit( + self, + script: str, + *, + flavor: str, + timeout_minutes: int, + allowed_flavors: tuple[str, ...], + max_timeout_minutes: int, + dependencies: list[str] | None = None, + datasets: list[str] | None = None, + denied_datasets: tuple[str, ...] = (), + volumes: list[Volume] | None = None, + idempotency_key: str | None = None, + ) -> Job: + """ + Start an agent job after checking limits and reserving its worst case. + + The script runs without network access, inside `job_launcher.py`. Its + dependencies, from `dependencies` and its PEP 723 header, are installed + from wheels before that. Each of `datasets` is mounted read-only at + `/datasets/`, next to `volumes`. + """ + with self._admission: + if idempotency_key and idempotency_key in self._idempotency: + return self.jobs[self._idempotency[idempotency_key]] + if flavor not in allowed_flavors: + raise BrokerError( + "flavor_not_allowed", + f"Flavor {flavor!r} is not one of {list(allowed_flavors)}", + ) + if not 1 <= timeout_minutes <= max_timeout_minutes: + raise BrokerError( + "timeout_not_allowed", + f"timeout_minutes must be in [1, {max_timeout_minutes}]", + ) + if not script.strip() or len(script.encode()) > MAX_SCRIPT_BYTES: + raise BrokerError( + "invalid_script", f"script must be 1 to {MAX_SCRIPT_BYTES} bytes" + ) + requirements = self._requirements(script, dependencies or []) + mounts = self._dataset_volumes(datasets or [], denied_datasets) + return self._start( + self._launcher_script(script, requirements), + role="agent", + flavor=flavor, + timeout_minutes=timeout_minutes, + dependencies=[], + volumes=[*(volumes or []), *mounts], + idempotency_key=idempotency_key, + ) + + def _start( + self, + script: str, + *, + role: str, + flavor: str, + timeout_minutes: int, + idempotency_key: str | tuple[str] | None = None, + **kwargs: Any, + ) -> Job: + # Only admission and ledger publication hold the lock. Provider calls + # may block while other tools inspect or cancel existing jobs. + retry = None + with self._admission: + if idempotency_key: + ready = self._admission.wait_for( + lambda: ( + not any( + p.idempotency_key == idempotency_key and p.active + for p in self._pending.values() + ) + ), + timeout=ADMISSION_TIMEOUT_S, + ) + if not ready: + raise BrokerError( + "creation_pending", + "Remote job creation is unconfirmed; retry later", + ) + if idempotency_key in self._idempotency: + return self.jobs[self._idempotency[idempotency_key]] + retry = next( + ( + request + for request, p in self._pending.items() + if p.idempotency_key == idempotency_key + ), + None, + ) + if retry is None: + if self._closing: + raise BrokerError("episode_over", "The episode is ending") + limit = self.max_jobs if role == "agent" else self.max_dev_evals + if role != "score" and self.count(role) >= limit: + raise BrokerError( + "limit_reached", f"This episode allows {limit} {role} jobs" + ) + reservation = self._reserve( + flavor, timeout_minutes, charged=role != "score" + ) + request = uuid4().hex + self._pending[request] = PendingJob( + role, reservation, flavor, timeout_minutes, idempotency_key + ) + if retry is not None: + try: + job = self._reconcile(retry) + except Exception as error: + raise BrokerError( + "creation_pending", + "Remote job creation is unconfirmed; retry later", + ) from error + if job is not None: + return job + # A successful lookup confirmed no job, or another caller adopted it. + return self._start( + script, + role=role, + flavor=flavor, + timeout_minutes=timeout_minutes, + idempotency_key=idempotency_key, + **kwargs, + ) + try: + job = self._run( + script, + role=role, + flavor=flavor, + timeout_minutes=timeout_minutes, + reservation=reservation, + request=request, + **kwargs, + ) + except Exception: + # A lost response can hide a running job. If lookup also fails, + # retain the request label, reservation and key for later recovery. + job = self._reconcile(request, creating=True) + if job is None: + raise + else: + self._publish(request, job) + return job + + def _publish(self, request: str, job: Job | None) -> None: + with self._admission: + pending = self._pending.pop(request) + if job is not None: + self.jobs[job.job_id] = job + if pending.idempotency_key: + self._idempotency[pending.idempotency_key] = job.job_id + closing = self._closing + self._admission.notify_all() + if job is not None and closing: + # A timed-out cleanup may have returned before creation completed. + # Keep the returned handle and attempt cancellation on this worker. + with suppress(Exception): + self.cancel(job, "episode ended during creation") + + def _reconcile( + self, + request: str, + *, + creating: bool = False, + cleanup: bool = False, + timeout_s: float | None = None, + ) -> Job | None: + with self._admission: + if not creating: + ready = self._admission.wait_for( + lambda: ( + request not in self._pending + or not self._pending[request].active + ), + timeout=ADMISSION_TIMEOUT_S if timeout_s is None else timeout_s, + ) + if not ready: + raise BrokerError( + "creation_pending", "Remote job creation is unconfirmed" + ) + pending = self._pending.get(request) + if pending is None: + return None + if self._closing and not (creating or cleanup): + # Cleanup owns recovery once closing starts, so a retry cannot + # publish a live job after cancel_all snapshots the known jobs. + raise BrokerError( + "creation_pending", "Remote job creation is unconfirmed" + ) + pending.active = True + try: + found = self._api.list_jobs( + labels={"openenv-request": request}, namespace=self.namespace + ) + info = next(iter(found), None) + job = ( + None + if info is None + else self._job_from_info( + info, + role=pending.role, + flavor=pending.flavor, + timeout_minutes=pending.timeout_minutes, + reservation=pending.reservation, + ) + ) + self._publish(request, job) + return job + finally: + with self._admission: + pending.active = False + self._admission.notify_all() + + @staticmethod + def _requirements(script: str, dependencies: list[str]) -> list[str]: + """Merge the script's PEP 723 dependencies with `dependencies`; packages only, no URLs.""" + match = SCRIPT_BLOCK.search(script.replace("\r\n", "\n")) + if match: + content = "".join( + line[2:] if line.startswith("# ") else line[1:] + for line in match.group("content").splitlines(keepends=True) + ) + try: + header = tomllib.loads(content).get("dependencies", []) + except tomllib.TOMLDecodeError as e: + raise BrokerError("invalid_script", f"PEP 723 header: {e}") + if not isinstance(header, list): + raise BrokerError( + "invalid_script", "PEP 723 dependencies must be a list" + ) + dependencies = [*header, *dependencies] + if len(dependencies) > MAX_DEPENDENCIES: + raise BrokerError( + "invalid_dependency", f"At most {MAX_DEPENDENCIES} dependencies" + ) + requirements = [] + for dependency in dependencies: + try: + requirement = Requirement(str(dependency)) + # packaging raises SyntaxError for some malformed markers. + except (InvalidRequirement, SyntaxError): + raise BrokerError( + "invalid_dependency", f"{dependency!r} is not a valid requirement" + ) + if requirement.url: + raise BrokerError( + "invalid_dependency", + f"{dependency!r}: only packages from PyPI, not URLs", + ) + requirements.append(str(requirement)) + return requirements + + def _dataset_volumes( + self, datasets: list[str], denied: tuple[str, ...] + ) -> list[Volume]: + """Mount each public, ungated dataset read-only at its current revision.""" + if len(datasets) > MAX_DATASETS: + raise BrokerError("dataset_not_allowed", f"At most {MAX_DATASETS} datasets") + denied_ids = {d.lower() for d in denied} + volumes = [] + for repo_id in dict.fromkeys(datasets): + try: + validate_repo_id(repo_id) + valid = "/" in repo_id + except HFValidationError: + valid = False + if not valid: + raise BrokerError( + "dataset_not_allowed", "Datasets are Hub repo ids like owner/name" + ) + if repo_id.lower() in denied_ids: + raise BrokerError( + "dataset_not_allowed", f"{repo_id} is not allowed in this task" + ) + # Looked up anonymously: the broker's token must not open private repos. + try: + info = self._api.dataset_info(repo_id, token=False) + except HfHubHTTPError: + info = None + if info is None or info.gated: + raise BrokerError( + "dataset_not_allowed", + f"{repo_id} is not a public, ungated dataset on the Hub", + ) + # A renamed repo's old id redirects to the new one. + if info.id.lower() in denied_ids: + raise BrokerError( + "dataset_not_allowed", f"{repo_id} is not allowed in this task" + ) + volumes.append( + Volume( + type="dataset", + source=info.id, + revision=info.sha, + mount_path=f"{DATASETS_MOUNT}/{repo_id}", + read_only=True, + ) + ) + return volumes + + @staticmethod + def _launcher_script(script: str, requirements: list[str]) -> str: + """Embed `script` and `requirements` in the launcher.""" + encoded = base64.b64encode(script.encode()).decode() + launcher = LAUNCHER.read_text() + for placeholder, value in ( + ('AGENT_SCRIPT = ""', f'AGENT_SCRIPT = "{encoded}"'), + ("REQUIREMENTS: list[str] = []", f"REQUIREMENTS = {requirements!r}"), + ): + launcher = launcher.replace(placeholder, value, 1) + header = ( + '# /// script\n# requires-python = ">=3.10"\n# dependencies = []\n# ///\n' + ) + return header + launcher + + def evaluate( + self, + source: Job, + *, + role: str, + script_path: Path, + flavor: str, + timeout_minutes: int, + env: dict[str, str], + secrets: dict[str, str] | None = None, + volumes: list[Volume] | None = None, + ) -> Job: + """Run an evaluator job on `source`'s artifact, mounted read-only at `/model`, next to `volumes`.""" + self.require_artifact(source) + model = Volume( + type="bucket", + source=source.bucket, + path=f"{source.path}/outputs", + mount_path="/model", + read_only=True, + ) + job = self._start( + script_path.read_text(), + role=role, + flavor=flavor, + timeout_minutes=timeout_minutes, + dependencies=[], + env=env, + secrets=secrets, + volumes=[model, *(volumes or [])], + # Broker lifetime is one episode; the tuple cannot collide with an + # agent-supplied string key. Retain all uncertain-creation guards. + idempotency_key=("score",) if role == "score" else None, + ) + job.source_job_id = source.job_id + return job + + def _run( + self, + script: str, + *, + role: str, + flavor: str, + timeout_minutes: int, + reservation: int, + request: str, + dependencies: list[str], + env: dict[str, str] | None = None, + secrets: dict[str, str] | None = None, + volumes: list[Volume] | None = None, + ) -> Job: + # run_uv_job uploads a local script to the namespace's jobs-artifacts + # bucket and mounts that job's folder read-write at /data, so outputs + # come back without a token inside the job. + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / f"{role}.py" + path.write_text(script) + info = self._api.run_uv_job( + script=str(path), + dependencies=dependencies, + env=env or {}, + secrets=secrets or {}, + flavor=flavor, + timeout=f"{timeout_minutes}m", + labels={ + "openenv-env": "posttrain_env", + "openenv-episode": self.episode_id, + "openenv-role": role, + "openenv-request": request, + }, + volumes=volumes, + namespace=self.namespace, + resource_group_id=self.resource_group_id, + ) + return self._job_from_info( + info, + role=role, + flavor=flavor, + timeout_minutes=timeout_minutes, + reservation=reservation, + ) + + @staticmethod + def _job_from_info(info, *, role, flavor, timeout_minutes, reservation) -> Job: + artifacts = [ + v + for v in info.volumes or [] + if v.type == "bucket" and v.mount_path == ARTIFACTS_MOUNT + ] + job = Job( + job_id=info.id, + role=role, + flavor=flavor, + timeout_minutes=timeout_minutes, + reserved_micro_usd=reservation, + bucket=artifacts[0].source if artifacts else None, + path=artifacts[0].path if artifacts else None, + ) + return job + + def require_artifact(self, job: Job) -> None: + """Raise `not_ready` unless `job` is a succeeded agent job with an artifact.""" + if job.role != "agent" or job.state != "succeeded" or job.bucket is None: + raise BrokerError( + "not_ready", f"Job {job.job_id} has no successful artifact" + ) + + def get(self, job_id: str) -> Job: + if job_id not in self.jobs: + raise BrokerError("unknown_job", f"No job {job_id!r} in this episode") + return self.jobs[job_id] + + def finished(self, job: Job) -> bool: + """True once `job` is settled and, for evaluations, its result is read.""" + return job.settled and (job.role == "agent" or job.result is not None) + + def refresh(self, job: Job) -> Job: + """Update `job` from HF Jobs, settling it and reading its score once it ends.""" + if self.finished(job): + return job + if not job.settled: + info = self._api.inspect_job(job_id=job.job_id, namespace=self.namespace) + stage = info.status.stage + if job.state != "cancelling" or stage in TERMINAL_STAGES: + job.state = STATE_FROM_STAGE.get(stage, job.state) + if info.status.message and job.state in ("failed", "cancelled"): + job.reason = job.reason or info.status.message + if info.durations is not None and info.durations.running_secs is not None: + job.running_secs = info.durations.running_secs + if job.state in TERMINAL_STATES: + self._settle(job) + if job.settled and job.role != "agent": + job.result = self._read_result(job) + return job + + def _read_result(self, job: Job) -> dict | None: + if job.state != "succeeded": + return {"status": job.state, "score": 0.0} + job.result_reads += 1 + lines = self._api.fetch_job_logs( + job_id=job.job_id, namespace=self.namespace, tail=200 + ) + for line in reversed(list(lines)): + if not line.startswith(SCORE_MARKER): + continue + try: + result = json.loads(line[len(SCORE_MARKER) :]) + score = float(result["score"]) + except (ValueError, KeyError, TypeError): + score = -1.0 + if not 0.0 <= score <= 1.0: + return {"status": "error", "score": 0.0, "feedback": "malformed score"} + metrics = result.get("metrics") + return { + "status": str(result.get("status", "ok")), + "score": score, + "metrics": dict(list(metrics.items())[:5]) + if isinstance(metrics, dict) + else {}, + "feedback": str(result.get("feedback", ""))[:500], + } + if job.result_reads < MAX_RESULT_READS: + return None + return { + "status": "error", + "score": 0.0, + "feedback": "evaluator printed no score", + } + + def cancel(self, job: Job, reason: str) -> Job: + if self.refresh(job).state in TERMINAL_STATES: + return job + # The job may have just finished; the next refresh decides. + with suppress(HfHubHTTPError): + self._api.cancel_job(job_id=job.job_id, namespace=self.namespace) + job.state = "cancelling" + job.reason = reason + return self.refresh(job) + + def wait(self, job: Job, timeout_s: float, poll_s: float) -> Job: + """Poll `job` until it is finished or `timeout_s` passes.""" + deadline = time.monotonic() + timeout_s + while not self.finished(self.refresh(job)) and time.monotonic() < deadline: + time.sleep(poll_s) + return job + + def cancel_all(self, reason: str, timeout_s: float, poll_s: float) -> list[str]: + """ + Cancel every live job, re-check earlier orphans, and wait for HF Jobs to confirm. + + Jobs still live after `timeout_s` are marked `orphaned` and settled at + their full reservation, since their real cost is unknown. Orphans are + retried on every later call until HF Jobs confirms they ended. + + Unconfirmed creation requests keep their reservations and are reconciled + again on later calls. They are returned as `request: