From f39eb85d197c3ac7a8fe32ed7bf3501c25cc1025 Mon Sep 17 00:00:00 2001 From: luyd Date: Thu, 12 Nov 2026 10:29:41 +0800 Subject: [PATCH] Consolidated release: gateway sharding, hub mode, math training, H-cluster practice - Elastic gateway scaling: static routing + remote-first auto-growth + automatic retirement; validated on 4xH200 - Hub deployment mode: pull-based nodes over HTTP for one-way-firewall clusters; validated end-to-end with real SGLang inference (reward=1) - Unified training algorithms (GRPO, DAPO, NPR, OPD, PRM distillation) on the shared Slime pipeline, validated on math tasks - Controller: deployment preflight, port allocation, startup queue, cleanup admission - Config tooling: options CLI wizard + web console + consistency checks - H-cluster practice: brainbox adapter, rjob provisioning, reverse tunnels Tests: 795 passed / 2 skipped / 3 pre-existing failures. --- .github/workflows/ci.yml | 4 - .gitignore | 2 +- README.md | 40 +- README_zh.md | 30 +- .../p0_p1_config_and_boundaries_zh.md | 81 ++ docs/configuration-wizard.md | 100 +- docs/deployment.md | 33 + docs/math-distill.md | 147 ++ docs/sandbox-rollout.md | 108 +- docs/train-inference-consistency.md | 36 +- docs/web-console.md | 67 +- examples/distill/local-models.yaml | 23 + examples/distill/math.jsonl | 12 + examples/distill/opd.yaml | 24 + examples/distill/prm-opd.yaml | 22 + examples/distill/prm.yaml | 20 + examples/distill/reinforce.yaml | 21 + examples/math/tasks.jsonl | 12 + examples/math/train.yaml | 35 + examples/native/catalog.json | 2 +- examples/native/catalog.jsonl | 2 +- examples/native/claude-code-profile.json | 8 +- examples/native/tasks/hello_world/task.toml | 8 +- .../native/tasks/hello_world/tests/test.sh | 8 +- examples/native/train_qwen_native.yaml | 4 +- harborrl/algorithms/opd.py | 33 + harborrl/backends/slime_v032/advantage.py | 2 +- harborrl/backends/slime_v032/config.py | 4 + harborrl/backends/slime_v032/dapo.py | 28 + harborrl/backends/slime_v032/loss.py | 13 +- harborrl/backends/slime_v032/math_rollout.py | 353 +++++ harborrl/backends/slime_v032/opd.py | 55 + harborrl/backends/slime_v032/postprocess.py | 27 + harborrl/backends/slime_v032/prm.py | 76 + harborrl/backends/slime_v032/rollout.py | 25 +- .../slime_v032/rollout_service_client.py | 12 +- harborrl/cli.py | 4 + harborrl/config/native.py | 16 +- harborrl/consistency/reference.py | 25 + harborrl/consistency/sglang.py | 8 +- harborrl/deployment/controller.py | 672 ++++++--- harborrl/deployment/hub_client.py | 444 ++++++ harborrl/deployment/hub_rollout_executor.py | 222 +++ harborrl/deployment/hub_server.py | 672 +++++++++ harborrl/deployment/hub_synthetic_backend.py | 51 + harborrl/deployment/preflight.py | 18 +- harborrl/deployment/provision.py | 180 +++ harborrl/deployment/schema.py | 280 +++- harborrl/deployment/training_target.py | 1 + harborrl/distill/__init__.py | 1 + harborrl/distill/cli.py | 126 ++ harborrl/distill/config.py | 108 ++ harborrl/distill/contract.py | 71 + harborrl/distill/data.py | 97 ++ harborrl/distill/models.py | 125 ++ harborrl/distill/runner.py | 241 ++++ harborrl/distill/scoring.py | 123 ++ harborrl/export/contract.py | 25 +- harborrl/export/native.py | 4 +- harborrl/export/validate.py | 66 +- harborrl/gateway/messages.py | 24 +- harborrl/gateway/server.py | 31 +- harborrl/gateway/sglang_backend.py | 32 +- harborrl/gateway/shard_service.py | 308 +++++ harborrl/gateway/trace.py | 148 +- harborrl/harnesses/claude_code.py | 9 +- harborrl/math/__init__.py | 1 + harborrl/math/config.py | 193 +++ harborrl/math/dapo.py | 44 + harborrl/math/data.py | 120 ++ harborrl/math/prm.py | 171 +++ harborrl/math/prm_server.py | 111 ++ harborrl/math/scoring.py | 47 + harborrl/math/teacher.py | 93 ++ harborrl/models/qwen.py | 2 +- harborrl/options/catalog.py | 21 +- harborrl/options/presets.py | 3 +- harborrl/options/validation.py | 4 +- harborrl/platform/evidence.py | 50 +- harborrl/platform/native_train.py | 6 +- harborrl/platform/preflight.py | 72 +- harborrl/platform/slime_train.sh | 48 +- harborrl/rollout/harbor_job/audit.py | 120 +- harborrl/rollout/harbor_job/bindings.py | 36 +- harborrl/rollout/harbor_job/collector.py | 17 +- harborrl/rollout/harbor_job/coordinator.py | 66 +- .../harbor_job/environments/opensandbox.py | 66 +- harborrl/rollout/harbor_job/runner.py | 15 +- harborrl/rollout/harbor_job/sandbox_config.py | 31 +- .../rollout/harbor_job/sandbox_transport.py | 27 +- harborrl/rollout/remote_shard.py | 201 +++ harborrl/rollout/runtime.py | 946 ++++++++++++- harborrl/rollout/service.py | 273 +++- harborrl/trajectories/native.py | 14 +- practices/h_cluster/HISTORICAL_VALIDATION.md | 32 + practices/h_cluster/README.md | 55 + practices/h_cluster/VALIDATION.md | 76 + practices/h_cluster/configs/rjob.example.json | 23 + .../h_cluster/configs/tunnel.example.json | 15 + .../src/harborrl_h_cluster/__init__.py | 21 +- .../src/harborrl_h_cluster/environment.py | 2 +- .../h_cluster/src/harborrl_h_cluster/rjob.py | 319 +++++ .../src/harborrl_h_cluster/transport.py | 4 +- .../src/harborrl_h_cluster/tunnel.py | 344 +++++ practices/h_cluster/tests/test_rjob.py | 181 +++ practices/h_cluster/tests/test_tunnel.py | 220 +++ pyproject.toml | 9 + scripts/audit_public_tree.sh | 4 +- scripts/sandbox_sdk_contract.py | 20 + tests/harborrl/test_dapo_training.py | 197 +++ tests/harborrl/test_deployment_targets.py | 31 +- tests/harborrl/test_distill.py | 217 +++ tests/harborrl/test_gateway_auto_shards.py | 1217 +++++++++++++++++ tests/harborrl/test_gateway_shards.py | 356 +++++ tests/harborrl/test_math.py | 113 ++ tests/harborrl/test_math_rollout.py | 145 ++ tests/harborrl/test_native_continuation.py | 135 ++ tests/harborrl/test_native_contracts.py | 10 + tests/harborrl/test_native_example_budget.py | 25 + tests/harborrl/test_native_gateway.py | 273 ++-- tests/harborrl/test_native_lifecycle.py | 134 +- tests/harborrl/test_native_terminal.py | 329 ++++- tests/harborrl/test_native_training.py | 115 +- tests/harborrl/test_opd_training.py | 190 +++ tests/harborrl/test_prm_training.py | 279 ++++ tests/harborrl/test_provision.py | 171 +++ tests/harborrl/test_provision_dispatch.py | 184 +++ tests/harborrl/test_remote_shard.py | 559 ++++++++ tests/harborrl/test_rollout_runtime.py | 28 +- tests/harborrl/test_run_port_allocation.py | 168 +++ .../test_runtime_cleanup_admission.py | 59 + tests/harborrl/test_sandbox.py | 42 + .../test_sandbox_allocation_transport.py | 35 + tests/harborrl/test_sandbox_cleanup.py | 18 + tests/harborrl/test_sandbox_runtime.py | 85 +- tests/harborrl/test_sandbox_startup.py | 195 +++ tests/harborrl/test_service_recovery.py | 224 +++ tests/harborrl/test_service_sampling_lock.py | 79 ++ tests/harborrl/test_service_shutdown.py | 167 +++ tests/harborrl/test_sglang_text_evidence.py | 54 + tests/harborrl/test_slime_v032_exporter.py | 63 +- tests/test_consistency_torch.py | 36 +- 142 files changed, 14776 insertions(+), 889 deletions(-) create mode 100644 docs/architecture/p0_p1_config_and_boundaries_zh.md create mode 100644 docs/math-distill.md create mode 100644 examples/distill/local-models.yaml create mode 100644 examples/distill/math.jsonl create mode 100644 examples/distill/opd.yaml create mode 100644 examples/distill/prm-opd.yaml create mode 100644 examples/distill/prm.yaml create mode 100644 examples/distill/reinforce.yaml create mode 100644 examples/math/tasks.jsonl create mode 100644 examples/math/train.yaml create mode 100644 harborrl/algorithms/opd.py create mode 100644 harborrl/backends/slime_v032/dapo.py create mode 100644 harborrl/backends/slime_v032/math_rollout.py create mode 100644 harborrl/backends/slime_v032/opd.py create mode 100644 harborrl/backends/slime_v032/prm.py create mode 100644 harborrl/consistency/reference.py create mode 100644 harborrl/deployment/hub_client.py create mode 100644 harborrl/deployment/hub_rollout_executor.py create mode 100644 harborrl/deployment/hub_server.py create mode 100644 harborrl/deployment/hub_synthetic_backend.py create mode 100644 harborrl/deployment/provision.py create mode 100644 harborrl/distill/__init__.py create mode 100644 harborrl/distill/cli.py create mode 100644 harborrl/distill/config.py create mode 100644 harborrl/distill/contract.py create mode 100644 harborrl/distill/data.py create mode 100644 harborrl/distill/models.py create mode 100644 harborrl/distill/runner.py create mode 100644 harborrl/distill/scoring.py create mode 100644 harborrl/gateway/shard_service.py create mode 100644 harborrl/math/__init__.py create mode 100644 harborrl/math/config.py create mode 100644 harborrl/math/dapo.py create mode 100644 harborrl/math/data.py create mode 100644 harborrl/math/prm.py create mode 100644 harborrl/math/prm_server.py create mode 100644 harborrl/math/scoring.py create mode 100644 harborrl/math/teacher.py create mode 100644 harborrl/rollout/remote_shard.py create mode 100644 practices/h_cluster/HISTORICAL_VALIDATION.md create mode 100644 practices/h_cluster/VALIDATION.md create mode 100644 practices/h_cluster/configs/rjob.example.json create mode 100644 practices/h_cluster/configs/tunnel.example.json create mode 100644 practices/h_cluster/src/harborrl_h_cluster/rjob.py create mode 100644 practices/h_cluster/src/harborrl_h_cluster/tunnel.py create mode 100644 practices/h_cluster/tests/test_rjob.py create mode 100644 practices/h_cluster/tests/test_tunnel.py create mode 100644 tests/harborrl/test_dapo_training.py create mode 100644 tests/harborrl/test_distill.py create mode 100644 tests/harborrl/test_gateway_auto_shards.py create mode 100644 tests/harborrl/test_gateway_shards.py create mode 100644 tests/harborrl/test_math.py create mode 100644 tests/harborrl/test_math_rollout.py create mode 100644 tests/harborrl/test_native_continuation.py create mode 100644 tests/harborrl/test_native_example_budget.py create mode 100644 tests/harborrl/test_opd_training.py create mode 100644 tests/harborrl/test_prm_training.py create mode 100644 tests/harborrl/test_provision.py create mode 100644 tests/harborrl/test_provision_dispatch.py create mode 100644 tests/harborrl/test_remote_shard.py create mode 100644 tests/harborrl/test_run_port_allocation.py create mode 100644 tests/harborrl/test_runtime_cleanup_admission.py create mode 100644 tests/harborrl/test_sandbox_startup.py create mode 100644 tests/harborrl/test_service_recovery.py create mode 100644 tests/harborrl/test_service_sampling_lock.py create mode 100644 tests/harborrl/test_service_shutdown.py create mode 100644 tests/harborrl/test_sglang_text_evidence.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5b22594a0..c8c90eead 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -26,8 +26,6 @@ jobs: run: python -m pytest - name: Example contract smoke run: bash examples/native/cpu_contract_smoke.sh - - name: New user Hello World - run: harborrl hello --out-dir /tmp/harborrl-hello - name: Training dry run run: python -m harborrl.cli train --config examples/native/train_qwen_native.yaml --dry-run - name: Public tree audit @@ -79,8 +77,6 @@ jobs: python -m venv /tmp/harborrl-wheel /tmp/harborrl-wheel/bin/pip install dist/*.whl /tmp/harborrl-wheel/bin/harborrl --help - cd /tmp - /tmp/harborrl-wheel/bin/harborrl hello --out-dir /tmp/harborrl-wheel-hello - uses: actions/upload-artifact@v4 with: name: harborrl-dist diff --git a/.gitignore b/.gitignore index ec61a7863..6e33a94ee 100644 --- a/.gitignore +++ b/.gitignore @@ -38,8 +38,8 @@ docs/agents/* !docs/agents/environment-check.md !docs/configuration-wizard.md !docs/deployment.md +!docs/math-distill.md !docs/sandbox-rollout.md -!docs/rollout-runtime-lifecycle.md !docs/train-inference-consistency.md !docs/web-console.md practices/**/VALIDATION.md diff --git a/README.md b/README.md index 22efe363e..67d589979 100644 --- a/README.md +++ b/README.md @@ -25,9 +25,11 @@ ## News -- **2026-10-26** Added an [offline configuration wizard](docs/configuration-wizard.md) and a [local Web console](docs/web-console.md) for configuration, validation, dry-runs, training launch, and logs. -- **2026-10-26** Added [controller/GPU deployment configuration](docs/deployment.md), optional [OpenSandbox execution](docs/sandbox-rollout.md), and opt-in [train/inference consistency checks](docs/train-inference-consistency.md). -- **10/2026** 📣📣📣 HarborRL released +- **2026-11-12** Added [gateway horizontal sharding](docs/deployment.md) with static routing, elastic auto-growth (remote-first), and automatic retirement. +- **2026-11-12** Added [Hub deployment mode](docs/deployment.md) for clusters with one-way network paths: nodes pull tasks over HTTP, no cross-machine SSH required. +- **2026-11-12** Added [online math training](docs/math-distill.md) (GRPO, DAPO, process rewards, OPD, PRM distillation) on the shared Slime/Megatron + SGLang pipeline. +- **2026-11-12** Added [configuration wizard](docs/configuration-wizard.md), [web console](docs/web-console.md), and [sandbox rollout](docs/sandbox-rollout.md) with audited cleanup. +- **2026-11-12** 10/2026 📣 HarborRL released ## About @@ -42,18 +44,23 @@ HarborRL supports flexible combinations of **Models × Algorithms × Tasks × Ha ## AI-assisted Configure +For offline option browsing, configuration validation, and interactive or +answer-file setup, see the [configuration CLI guide](docs/configuration-wizard.md). + +For mathematical data validation and scoring utilities, see the [math adapter guide](docs/math-distill.md). + For the offline options catalog, interactive configuration, and answer-file setup, see the [configuration CLI guide](docs/configuration-wizard.md). -Start with `harborrl hello` for a CPU walkthrough requiring no model, GPU or task -JSON. It executes scripted attempts and a real verifier, then exports a GRPO batch -with fixture tokens. Continue with `harborrl options` or `harborrl init --help`. -The wizard detects one shared Local environment, calculates model fingerprints, -and supports incomplete drafts (`:save` / `--resume`). +Start with `harborrl options` or `harborrl init --help`. Run `harborrl web` to open the local browser console for configuration, validation, -drafts, filled Local/Sandbox examples, Hello World, dry-runs, training launch, and logs. See the [web console guide](docs/web-console.md) +dry-runs, training launch, and logs. See the [web console guide](docs/web-console.md) for remote-controller access and lifecycle details. +For the optional math PRM/OPD reference trainer, see +[Math training pipeline](docs/math-distill.md). It includes explicit synthetic +fixtures for testing real optimizer updates and checkpoint restore. + If you want an AI agent to help deploy this repository and verify training, ask it to first read and strictly follow [docs/agents/environment-check.md](docs/agents/environment-check.md). @@ -122,15 +129,10 @@ Expected result: uv executes the installed entry point and tests without activat For native training, install the pinned backend stack in a writable location: ```bash -bash scripts/bootstrap_backends.sh "" -. "/harborrl-backend.env" +bash scripts/bootstrap_backends.sh +. /harborrl-backend.env ``` -The generated file exports `SLIME_DIR`, `MEGATRON_DIR`, and `SGLANG_IMAGE` to child -processes. For files generated by older bootstrap scripts, run `set -a` before -sourcing and `set +a` afterward. See [rollout lifecycle and recovery](docs/rollout-runtime-lifecycle.md) -for cleanup markers and CPU-only regression checks. - ### Train After installation, choose one of the two configuration styles: @@ -169,11 +171,7 @@ task/ Here, `task.toml` declares resources, the agent, and the verifier; `instruction.md` is the agent-facing instruction; `environment/Dockerfile` builds the isolated environment; and `tests/test.sh` verifies the result and writes `reward.json` or `reward.txt`. For training, tasks are locked into a catalog with their task ID, upstream revision, local path, content digest, and reward parsing rules; see [examples/native/catalog.json](examples/native/catalog.json). -The bundled runnable example is `hello_world`. It allows internet access so Harbor can install the pinned Claude Code CLI, reserves 4 GiB each for memory and storage, and writes a 0/1 reward to `/logs/verifier/reward.json`. Its profile allows four turns so the agent can observe a tool result and finish. Offline deployments need the exact CLI version preinstalled and a reachable policy gateway. After changing a task, regenerate its catalog digest. - -The example pins the runtime-validated Claude Code `2.1.141`. The adapter disables prompt caching and cache warmup, compaction, and nonessential traffic, and always disallows `Agent` and `Task` even with a custom tool denylist. CLI `2.1.0` still makes an independent file-path-extraction model request after Bash and is incompatible with the single-context rollout contract. Probe other CLI versions before use; native session audit rejects auxiliary contexts, summaries, and unconsumed model responses. - -Real training tasks are not repackaged by this repository: fetch Harbor-native datasets or tasks from Harbor Hub, and use the [compatibility summary](configs/harbor_hub/public_compatibility_summary.md) to understand validated coverage. After fetching a task root, validate it locally with `harborrl hub inspect` and lock it into your own catalog. +The bundled runnable example is `hello_world`. Real training tasks are not repackaged by this repository: fetch Harbor-native datasets or tasks from Harbor Hub, and use the [compatibility summary](configs/harbor_hub/public_compatibility_summary.md) to understand validated coverage. After fetching a task root, validate it locally with `harborrl hub inspect` and lock it into your own catalog. ## License diff --git a/README_zh.md b/README_zh.md index bce055ca6..3e02f0d1f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -25,9 +25,11 @@ ## News -- **2026-10-26** 新增[离线配置向导](docs/configuration-wizard.md)和[本地 Web 控制台](docs/web-console.md)，支持配置、校验、dry-run、训练启动和日志查看。 -- **2026-10-26** 新增 [controller/GPU 部署配置](docs/deployment.md)、可选的 [OpenSandbox 执行能力](docs/sandbox-rollout.md)，以及按需启用的[训练与推理一致性检查](docs/train-inference-consistency.md)。 -- **10/2026** 📣📣📣 HarborRL released +- **2026-11-12** 新增[网关水平分片](docs/deployment.md)：静态路由、弹性自动生长（远端优先）与自动缩容。 +- **2026-11-12** 新增 [Hub 部署模式](docs/deployment.md)：适配单向网络集群，节点经 HTTP 拉取任务，无需跨机 SSH。 +- **2026-11-12** 新增[在线数学训练](docs/math-distill.md)（GRPO、DAPO、过程奖励、OPD、PRM 蒸馏），统一接入 Slime/Megatron + SGLang 管线。 +- **2026-11-12** 新增[配置向导](docs/configuration-wizard.md)、[Web 控制台](docs/web-console.md)与带审计清理的 [Sandbox 执行](docs/sandbox-rollout.md)。 +- **2026-11-12** 10/2026 📣 HarborRL 发布 ## 项目介绍 @@ -46,12 +48,8 @@ HarborRL 支持 **Models × Algorithms × Tasks × Harnesses (MATH)** 的灵活 [docs/agents/environment-check.md](docs/agents/environment-check.md)。 离线选项目录、交互式配置和 answer-file 配置方式见[配置 CLI 指南](docs/configuration-wizard.md)。 -首次使用可以先运行 `harborrl hello`：无需模型、GPU 或已有 task JSON， -即可执行脚本任务、真实 verifier、0/1 奖励收集和使用示例 token 的 GRPO batch 导出。 -接着使用 `harborrl options` 或 `harborrl init --help`。Local 模式会检测并共用 -一套源码目录与 Python；模型指纹自动计算，未填完的草稿支持 `:save` 保存和 `--resume` 恢复。 -运行 `harborrl web` 可打开浏览器控制台，查看或套用 Local／Sandbox 示例、 -随时保存草稿、一键运行 Hello World，并完成校验、dry-run、训练启动和日志查看； +可以从 `harborrl options` 或 `harborrl init --help` 开始。 +运行 `harborrl web` 可打开本地浏览器控制台，完成配置、校验、dry-run、训练启动和日志查看； 远程访问与生命周期说明见 [Web 控制台指南](docs/web-console.md)。 推荐的实际操作流程是把下面这段提示词交给 agent，并把尖括号内容替换成你的机器信息： @@ -112,14 +110,10 @@ uv run pytest 如需原生训练，请把固定版本的 native 训练栈安装到可写位置： ```bash -bash scripts/bootstrap_backends.sh "" -. "/harborrl-backend.env" +bash scripts/bootstrap_backends.sh +. /harborrl-backend.env ``` -生成的文件会将 `SLIME_DIR`、`MEGATRON_DIR` 和 `SGLANG_IMAGE` 导出给子进程。 -对于旧版 bootstrap 脚本生成的文件，请在 source 前执行 `set -a`，之后执行 `set +a`。 -清理标记的恢复步骤和无 GPU 回归命令见[rollout 生命周期与恢复](docs/rollout-runtime-lifecycle.md)。 - ### 训练 完成安装后，从两种配置方式中选择一种： @@ -158,11 +152,7 @@ task/ 其中 `task.toml` 声明资源、agent 与 verifier，`instruction.md` 是给 agent 的任务说明，`environment/Dockerfile` 构建隔离执行环境，`tests/test.sh` 负责验证并输出 `reward.json` 或 `reward.txt`。训练时，这些任务会被锁定到 catalog 中，记录任务 ID、上游 revision、本地路径、内容 digest 和 reward 解析规则；示例见 [examples/native/catalog.json](examples/native/catalog.json)。 -仓库内置的可直接运行示例是 `hello_world`。示例允许联网，以便 Harbor 安装固定版本的 Claude Code CLI；内存与存储各配置 4 GiB，verifier 将 0/1 奖励写入 `/logs/verifier/reward.json`。profile 允许四轮交互，供 agent 执行工具、读取结果并正常结束。离线部署需要预装完全一致的 CLI 版本，并确保能够访问 policy gateway。修改任务内容后，需要重新生成 catalog 中的内容摘要。 - -示例固定使用经过运行验证的 Claude Code `2.1.141`。adapter 默认关闭 prompt caching 及缓存预热、上下文压缩和非必要流量；自定义工具禁用列表也会保留 `Agent` 和 `Task`。CLI `2.1.0` 在 Bash 执行后仍会发起独立的文件路径提取请求，不符合单上下文 rollout 契约。使用其他 CLI 版本前需要重新验证；原生会话审计会拒绝辅助上下文、摘要和未消费的模型响应。 - -实际训练任务不由本 repo 重新打包，而是从 Harbor Hub 获取符合 Harbor native task 格式的 dataset/task；可先查看[兼容性汇总](configs/harbor_hub/public_compatibility_summary_zh.md)了解已验证覆盖情况。取回 task root 后，先用 `harborrl hub inspect` 做本地静态校验，再锁定进自己的 catalog。 +仓库内置的可直接运行示例是 `hello_world`。实际训练任务不由本 repo 重新打包，而是从 Harbor Hub 获取符合 Harbor native task 格式的 dataset/task；可先查看[兼容性汇总](configs/harbor_hub/public_compatibility_summary_zh.md)了解已验证覆盖情况。取回 task root 后，先用 `harborrl hub inspect` 做本地静态校验，再锁定进自己的 catalog。 ## License diff --git a/docs/architecture/p0_p1_config_and_boundaries_zh.md b/docs/architecture/p0_p1_config_and_boundaries_zh.md new file mode 100644 index 000000000..c82ed09e7 --- /dev/null +++ b/docs/architecture/p0_p1_config_and_boundaries_zh.md @@ -0,0 +1,81 @@ +# P0/P1：配置生效与模块边界 + +本次变更延续已有 GPU 验收的默认训练路径，保持 README、schema-2 和用户命令不变。涉及导出代码的收敛保留原 GRPO 数学公式与权重规则；本轮验证采用 CPU 合同、旧实现数值对照、shell dry-run 和包构建，不代表新增 GPU 实测。 + +## 配置如何生效 + +优先级：YAML 原值 → 按顺序应用 `--set` → 显式 `-model/-harness/-task/-alg` 选择。后端路径、可见 GPU、代理及动态库等部署环境仍由声明的环境变量提供。结构化启动不继承环境中的 `LR`、采样预算、checkpoint 开关或其他未声明训练选项，避免覆盖 YAML。 + +`training.learning_rate` 现在显式映射为 `LR`，进入 Slime `--lr`。旧默认 `1e-6` 保持不变；此前被忽略的非默认 YAML 学习率现在会生效。 + +配置字段到环境的直接映射统一维护于 `harborrl/backends/slime_v032/config.py::CONFIG_ENV`；`field_usage()` 同时记录不通过环境消费的字段，测试要求覆盖全部 schema-2 字段。 + +| 配置 | 实际消费者 | +| --- | --- | +| `execution.backend`、`harness.name` | 当前执行/harness 组合校验及 runner | +| `tasks.catalog` | 静态 task lock 校验；运行时生成锁定 catalog 和 prompt data | +| `harness.profile` | 最终 profile 文件校验；runner 配置 | +| `harbor.python/version/workers` | runner Python、worker runtime probe、worker 调度 | +| `gateway.host/port/advertised_url` | gateway listener 和 worker 访问入口 | +| `gateway.tokenizer_digest/template_digest/audited_raw_logprobs` | 生成证据与训练资格检查 | +| `model.checkpoint/reference/args_file` | `--hf-checkpoint`、`--load`、`--ref-load` 与模型 preset | +| `training.learning_rate/num_rollout/save_interval` | `--lr`、`--num-rollout`、`--save-interval` | +| `training.backend_contract` | 后端 selector 与版本/hooks 检查 | +| `sampling.group_size/groups_per_batch` | `--n-samples-per-prompt`、`--rollout-batch-size` | +| `sampling.max_tokens/max_context` | response/context 预算 | +| `sampling.max_attempts` | 完整组 slot 重试限制 | +| `deployment.*` | GPU 数量、TP、split/colocate 校验和启动参数 | +| `output.root` | run、checkpoint 与指标根目录 | + +严格配置校验与文件检查分离：`validate()` 不读 harness/task 文件，也不执行 live probe；`validate_files()` 校验最终选定的 harness profile；`tasks.catalog.catalog()` 校验任务锁；doctor 执行 live 部署检查。`load_config()` 保留加载后校验 profile 的兼容行为，但在 override 完成后检查文件，因此可以替换已经不存在的旧 profile 路径。 + +## 运行证据 + +每次结构化训练新增： + +- `effective-config.json`：解析后配置、task selectors、组件标识、运行时环境映射、声明的部署信息、doctor 观察到的包版本/后端 commit、源码证据与字段消费者。 +- `effective-command.json`：shell 默认值和模型 preset 展开后的实际可执行 argv，在调用 Slime 前记录。 +- `dry-run-commands/.json`：直接运行 shell dry-run 时的命令证据，允许同目录重复检查不同参数。 + +有效配置快照不转储整个父进程环境，不记录代理 URL 或环境中的 provider 凭据。原有 `launch.json`、源码 manifest、source patch 和任务产物保留。 + +`harborrl train --dry-run` 仍只解析并输出 plan，不启动 doctor 或训练；shell 命令证据是在实际执行 shell launcher 时生成的。若进程在进入 launcher 前失败，最终命令文件可能尚不存在。 + +## 模块职责 + +| 模块 | 职责 | +| --- | --- | +| `config/native.py` | schema-2 解析、类型/约束、路径规则、最终 profile 校验；保留旧 catalog/launch_plan 薄封装 | +| `backends/slime_v032/config.py` | Slime 配置映射与字段消费者清单 | +| `tasks/catalog.py` | native catalog 与 task lock 检查 | +| `tasks/content.py` | task 内容 digest 与共享 TOML reader | +| `platform/preflight.py` | live doctor；task digest 每次 catalog 检查一次 | +| `platform/evidence.py` | 输入准备、源码证据、有效配置与最终命令记录 | +| `platform/native_train.py` | 流程派发、run 目录、子进程、信号转发；旧 doctor/准备函数保留封装 | +| `export/native.py` | canonical immutable GRPO TrainingBatch | +| `rollout/exporters/native.py` | canonical batch 的旧 record 格式视图 | + +模型 tool parser 与 harness 插件接口留给 P2；本阶段不提前扩大支持组合，也不改变 serving、loss、同步屏障或 checkpoint 生命周期。 + +## 历史工具清单与处置 + +| 历史路径 | 当前调用关系 | 处置 | +| --- | --- | --- | +| `data/harbor/__main__.py` | 模块 CLI → inspector/materializer | 保留历史入口，标注 TerminalEnv 兼容用途 | +| `data/harbor/inspector.py` | 历史 CLI、materializer、probe | 保留；共享 digest/TOML 从 `tasks/content.py` 兼容导出 | +| `data/harbor/materializer.py` | 历史 CLI；生成 task.yaml、Compose、wrapper | 保留以支持旧记录复现，native 训练不使用 | +| `data/harbor/probe.py`、`receipt.py` | 旧 TerminalEnv 动态探测和奖励解析 | 保留并标注兼容用途 | +| `data/harbor/native_inspector.py` | native catalog 与 hub inspect | 直接依赖 `tasks/content.py`，不再导入历史 inspector/materializer | +| `rollout/exporters/native.py` | GroupSlots、offline smoke、旧合同调用者 | 保留格式入口，数学及 batch 校验共用 canonical exporter | + +未删除本地环境、缓存、runs 或验收记录。旧接口在确认无复现/外部调用需求后再单独决定退役。 + +## 回归重点 + +1. YAML 原值与 `--set` 的学习率均进入最终 `--lr`；预算、GPU、checkpoint 和输出参数通过命令级检查。 +2. 父环境的训练变量不能覆盖 YAML；有效配置不包含父环境凭据与代理值。 +3. 非法 section、字段类型、override YAML 和 catalog entry 返回可定位的 CLI 错误。 +4. 有效组和零方差组、多 turn、不同轨迹长度的旧导出字段与权重保持一致；缺失证据、不完整组和重复身份继续拒绝。 +5. native task inspection 不导入历史转换路径；原始 verifier 与 task digest 算法保持不变。 + +本次验证结果：102 项 CPU 测试通过；Ruff、shell 语法、示例 smoke、训练 plan dry-run、公开树审计通过。旧/新 exporter 在有方差与零方差、多 turn、不同长度样本上的 JSON 输出逐字一致。使用满足 pyproject 构建要求的隔离环境成功构建 sdist/wheel，并在仓库外安装 wheel 后通过 CLI dry-run 与新增模块导入。 diff --git a/docs/configuration-wizard.md b/docs/configuration-wizard.md index 61a7e93f1..a839a490f 100644 --- a/docs/configuration-wizard.md +++ b/docs/configuration-wizard.md @@ -2,28 +2,8 @@ The configuration tools run with the core package (Python 3.10+ and PyYAML). They enumerate implemented capabilities, check configuration offline, and create -consistent training schema 2 / deployment schema 1 files. Catalog, preview and -validation stay offline. Model preparation explicitly runs the selected Python -locally or over SSH to fingerprint an existing tokenizer. It does not download -models or allocate Sandbox/GPU jobs. - -## First run without a task JSON or GPU - -Install the core package and try the bundled CPU walkthrough: - -```bash -python -m pip install -e . -harborrl hello --out-dir ./hello-first -harborrl web --workspace ./harborrl-web -``` - -Hello World executes two scripted attempts at creating a file, runs the actual -task verifier, collects rewards 0 and 1, assembles native IR and exports a GRPO -batch. It uses fixture token evidence, marked `CONTRACT_READY`; it does not perform -model inference or update GPU weights. All files stay under the new output -directory. The generated `catalog.json` and `claude-code-profile.json` can be reused -when configuring real training. The command also works from an installed wheel -outside a source checkout, on Windows or Linux. +consistent training schema 2 / deployment schema 1 files. They do not install +software, download tasks or models, connect over SSH, or allocate Sandbox/GPU jobs. ## Browse and validate @@ -78,13 +58,12 @@ compatibility is additionally checked by the existing training dry-run. 2. Prepare a pinned Claude Code profile. See [the native examples](../examples/native) for profile structure. It must match the agent and task requirements of your deployment. -3. Stage checkpoint and reference directories on the GPU target. The wizard - calculates both tokenizer fingerprints automatically, using that target's - Python with Transformers installed. No manual SHA-256 entry is needed. In Web, - choose **Prepare model** after setting the checkpoint and environment. The - operation loads tokenizer files offline with `trust_remote_code=False` and calls - the same fingerprint function as the rollout runtime. A missing dependency or - checkpoint produces an actionable error; the draft can still be saved. +3. Stage checkpoint and reference directories on the GPU target. In its existing + training environment, load the actual local tokenizer with the same options as + the runtime and call `harborrl.rollout.runtime.tokenizer_fingerprint(tokenizer)`. + Its return values are `tokenizer_digest, template_digest`. This is an explicit + preparation step, not an operation performed by the wizard. Hashing a filename + or guessing an arbitrary 64-character value does not establish compatibility. 4. Identify controller/GPU/runner Python executables, checkout paths, SSH aliases and shared storage. Service origins must include their declared ports and be reachable by the participating machines. @@ -100,36 +79,13 @@ harborrl init --out-dir ./site-new harborrl init --answers my-answers.yaml --out-dir ./site-new ``` -Interactive mode explains each required field. Local asks for one shared source -directory and Python interpreter, detected from the host running the wizard; -Controller, GPU and rollout all use those values. Storage is configured -separately. Service modes ask for each remote role and its reachable origins. -The wizard prepares model fingerprints automatically, previews both -configurations and requires `yes` to publish. +Interactive mode asks for deployment roles, existing inputs, model fingerprints, +GPU layout and a few sampling settings, then previews both configurations and +requires `yes` to publish. Entering no or EOF cancels. Non-terminal use requires +`--answers`; missing required answers never fall back to prompting. -Progress is saved after every answer to `OUT_DIR.draft.json`. Type `?` for help, -`:skip` to leave a field empty, or `:save` to save and exit. Ctrl+C or EOF after -choosing a deployment also keeps the draft. Resume it later: - -```bash -harborrl init --out-dir ./site-new --resume ./site-new.draft.json -harborrl init --out-dir ./empty-site --save-draft -harborrl init --out-dir ./local-example --example local -harborrl init --out-dir ./sandbox-example --example sandbox -``` - -`--example` saves a filled reference draft using the current host's detected -environment. Replace example paths and hosts before resuming. `--draft-file` -chooses a different progress file. Incomplete drafts never become runnable -bundles. Draft updates check revisions and reject stale writers. - -Non-terminal generation accepts `--answers` or `--resume`. Missing fingerprints -trigger explicit tokenizer preparation; `--prepare-model` recalculates existing -values. Missing required answers never fall back to prompting. Use `--save-draft` -to store incomplete answers without running preparation or final validation. - -Start from one of these reference answer templates and replace example paths -before running: +Start from one of these complete answer templates and replace every `REPLACE_*` +value and example path before running: - [embedded-local](../examples/options/embedded-local.answers.yaml): one machine, local Harbor runner, colocated training/serving by default. @@ -138,18 +94,13 @@ before running: - [service-sandbox-local-runner](../examples/options/service-sandbox.answers.yaml): controller and runner here, SSH GPU target, Sandbox task execution. -Templates leave fingerprints empty for automatic preparation. Web also offers -filled Local and Sandbox recipes with notes, download and apply actions. -GPU allocations illustrate topology; confirm memory requirements for your model -before training. For Local Docker, the gateway advertised origin must be -reachable inside task containers. On Linux, configure the host-gateway mapping -if using `host.docker.internal`, or choose a reachable host address. +Templates deliberately contain invalid digest placeholders. They cannot claim +training readiness without your real inputs. Default one-GPU colocate is a +configuration starting point, not a Qwen model memory-fit recommendation. Answers allow only `answers_version: 1`, `preset`, `training`, `deployment`, and -optional `sandbox_config` and `local_environment`. Local's `local_environment` -contains `source_dir` and `python` and applies to all three roles, taking precedence -over advanced per-role overrides. Overrides use the existing schema field names; -unknown keys fail. Answer values are not evaluated as shell expressions. Gateway fields, +optional `sandbox_config`. Overrides use the existing schema section/field names; +unknown keys fail. No expressions or shell commands are evaluated. Gateway fields, shared run root, GPU budget and worker destinations are derived when only one side is provided. Contradictory explicit values fail. Supply matching origin URLs when changing ports. Paths in training overrides are relative to the **final @@ -187,3 +138,16 @@ harborrl doctor site-new/config.yaml --deployment site-new/deployment.yaml `doctor` performs live environment/connectivity checks. Actual training remains a separate action. Command examples use POSIX shell quoting. + +## Maintenance + +`harborrl.options.build_catalog()` and +`build_config_bundle(answers, output_directory=...)` have no terminal dependency; +the latter returns file contents and validation without publishing them. The +authoritative validators remain in the existing training/deployment modules. +Field descriptions are display metadata, not a second validation engine. + +Dataset metadata is maintained in `harborrl/options/resources/datasets.yaml` and +packaged using importlib resources. Run `python scripts/sync_options_manifest.py` +after changing it; `--check` detects drift in the legacy manifest. No web server, +new algorithm, resource scheduler or H-specific core preset is introduced. diff --git a/docs/deployment.md b/docs/deployment.md index 2f35a8501..bab285b00 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -98,6 +98,39 @@ requires a complete group of at least two slots on the locked policy version, and only then converts trajectories into Slime samples. It is selected by `HARBORRL_ROLLOUT_SERVICE_URL` in the GPU job environment. +## Service shutdown and recovery + +The service holds an OS file lock on its shared run directory until shutdown +completes or its process exits. Registration persists the run manifest; a +different manifest for the same run is rejected even after a restart. Run and +group IDs must each be a single directory component. + +Registration restores durable group outcomes. Repeating a completed request +returns its original result without regenerating trajectories or rewriting +immutable records. Failed and cancelled groups retain their outcomes. The GPU +client still validates all trajectory evidence before accepting a recovered +result. Completion is reported only after the result has been published. + +A group without a durable terminal outcome becomes `unknown`. The service does +not adopt its previous Runner or replay partial work. It writes `unknown.json` +and `native-rollouts/cleanup-required.json`, blocks new groups for that run, and +returns the unknown outcome to polling clients. Inspect and confirm cleanup, +retain the original evidence, then use a fresh run for further sampling. This +recovers results and status; it does not resume an Agent or optimizer. + +`RolloutService.close()` blocks admission, cancels groups and waits for their +asynchronous Runner cleanup, drains generation, then closes the Gateway and +executor loop. Its cleanup budget covers the Runner's 120-second cancellation +deadline and configured Sandbox cleanup deadlines. A generation drain alone +does not wait for tools, verification or environment deletion. + +A cleanup timeout raises an error and leaves cleanup running. Incomplete +shutdown writes a cleanup-required marker; another close can wait for cleanup +without cancelling the same group again. The controller closes its HTTP listener +even when service shutdown fails. Both Docker and Sandbox runtimes reject new +attempts while observed cleanup-required markers remain. Inspect the resources +before removing resolved markers. + ## Operational limits Service mode requires shared storage visible at the configured paths on the diff --git a/docs/math-distill.md b/docs/math-distill.md new file mode 100644 index 000000000..174555126 --- /dev/null +++ b/docs/math-distill.md @@ -0,0 +1,147 @@ +# Math training reference pipeline + +This optional pipeline adds math tasks, result verification, process rewards +and on-policy teacher scoring to the controller-sandbox source tree. Its +contract is `transformers-math-onpolicy-v1`. It uses the existing local/SSH +training targets and immutable record writer. It does not register a Slime +algorithm or change Native IR, GRPO, sandbox rollout or the controller gateway. +Math generation and scoring currently run in the trainer process. + +## Install and run + +Use a Linux source checkout with Python 3.10 or newer: + +```sh +pip install -e '.[distill]' +harborrl distill train examples/distill/opd.yaml --allow-synthetic --dry-run +harborrl distill train examples/distill/opd.yaml --allow-synthetic +harborrl distill evaluate examples/distill/opd.yaml --allow-synthetic \ + --checkpoint runs/math-opd-smoke/checkpoint-000002.pt \ + --output runs/math-opd-evaluation.json +``` + +Training requires a Git source checkout for the source audit. Help, curated +import and dry-run do not import Torch, Transformers, Harbor or sandbox SDKs. +Optional dependencies are loaded only for execution. Dry-run validates the +configuration and tasks; it does not verify device capacity or model loading. + +The bundled tiny Qwen and byte tokenizer are randomly initialized fixtures. +Fake teacher/PRM and tiny models require explicit `--allow-synthetic`. Student +generation, backward, AdamW updates and checkpointing are real even in this +mode. These fixtures cannot demonstrate improved mathematical ability. + +Examples: `reinforce.yaml`, `prm.yaml`, `opd.yaml`, `prm-opd.yaml` and +`local-models.yaml`. All paths resolve relative to the configuration file. +Each run requires a new output directory. Unknown configuration fields fail. + +Set `deployment.device: cuda` for a GPU. Set `deployment.world_size: 2` for +two data-parallel processes; global `training.batch_size` must divide evenly. +`training.microbatch_size` controls accumulation within each rank. Gradients +are averaged across ranks, clipped, and applied once per generated batch. +Metrics are per rank, so aggregate their losses by an equal-rank mean. + +An explicit `deployment.file` can select the existing shared-filesystem SSH +training target. The same absolute source, data, model and run paths must be +accessible to the target; this command does not copy assets or translate +paths. No deployment file means local execution. Cluster addresses and +credentials belong in private deployment files, not the public examples. + +## Data and scoring + +Each task JSONL row has exactly `task_id`, `problem_id`, `prompt`, `reference`, +`verifier`, `split`, `source`, and `license`. The bundled 12 arithmetic tasks +are original MIT fixtures: eight train, two validation and two test. +`verifier` is `exact_rational_v1`. References are bounded integers, decimals +or fractions. Duplicate task IDs, cross-split problem IDs and identical +cross-split prompts are rejected. This is not semantic deduplication. + +The verifier accepts a bare number, a final `####` answer or a final +`\boxed{...}` number. Correct/incorrect/unparseable/unsupported are distinct; +their rewards are +1/-1/0/0. It never executes generated text and does not +claim symbolic algebra support. + +Curated responses can be imported separately: + +```sh +harborrl distill import --tasks tasks.jsonl --samples responses.jsonl \ + --out-dir datasets/math-curated +``` + +Response rows require `task_id`, `response`, `source`, `license`. Import writes +immutable samples, verification sidecars and a manifest. Imported text is +declared evidence with no token probabilities. It is not eligible as an +on-policy batch; the trainer generates fresh student responses from tasks. + +PRM adapters split responses at newline token boundaries. `fake` supplies +deterministic test rewards; `equation` verifies bounded binary rational +equations such as `1 + 1 = 2`. Unsupported steps receive zero. A boundary +token belongs to the preceding step. Each token receives the undiscounted +sum of rewards from its step onward. This is a process-reward integration +point, not a trained neural PRM or a general proof checker. + +## Objectives and admission + +For each valid response token the detached advantage is: + +```text +opd: opd_weight * (teacher_logprob - behavior_logprob) +prm: prm.weight * step_reward_to_go +prm_opd: sum of the two terms above +reinforce: outcome_weight * final_verifier_reward +``` + +`outcome_weight` may additionally contribute to the other methods. The loss +is the negative mean over trajectories of each trajectory's valid-token +mean of `advantage * current_logprob`. OPD implements a sampled-token, +fixed-prefix reverse-KL surrogate, not full-vocabulary KL or the full +sequence-distribution KL gradient including future-prefix costs. + +Sampling uses unmodified raw student probabilities: temperature=1, top_p=1, +no top-k, no special-token filtering, no implicit truncation. Each generated +batch permits one optimizer update. The scorer must score those exact token +IDs at the same causal prefixes. Identity, teacher revision, policy version, +request/response digests, finite probabilities and response masks are checked +before backward. Recomputed actor log probabilities must match the frozen +behavior within the configured tolerance. Teacher parameters remain frozen. + +`hf_local` loads only local, non-remote-code causal models and fast tokenizers. +Teacher/student tokenizer, vocabulary, special tokens and serialization +identity must match exactly. Model vocabulary must equal tokenizer length; +padded-vocabulary models are currently rejected. Prompts use literal text +with tokenizer special tokens, not automatic chat-template rendering. +Use task prompts appropriate to that serialization. A live teacher capability +probe checks first-token, EOS and causal-prefix scoring consistency. + +## Resume and evidence + +To resume, copy the configuration, select a new output directory and raise +`training.steps` to the total desired completed steps: + +```sh +harborrl distill train resume.yaml --allow-synthetic \ + --resume runs/math-opd-smoke/checkpoint-000001.pt +``` + +Only load trusted checkpoints. Restore validates configuration, task digest, +tokenizer identity and parameter digest, and restores optimizer state. +Output/deployment, total steps, microbatch size and save interval may change; +objective and global batch settings must remain compatible. + +Runs contain effective configuration, launch plan, source audit, runtime +versions, optional teacher probe, immutable token batches, per-rank metrics, +checkpoints and completion/failure records. Evaluation reads a compatible +checkpoint on validation/test data and reports verifier accuracy separately. +No training completion flag implies improved model capability. + +## Boundaries + +This is a small reference backend using eager float32 Transformers inference, +serial per-sample generation/scoring and full parameter hashing. It supports +CPU and replicated data parallel GPUs, not tensor/pipeline parallelism or +large-model efficiency claims. Start with small models and contexts. + +Not included: API teacher collection, LLM judges, a trained PRM, SGLang scoring, +Slime/Megatron OPD integration, replay/off-policy training, SFT from imported +samples, multi-turn environments, automatic dataset licensing checks, or +production model-quality evaluation. Keeping these outside this first phase +makes the math contracts and genuine optimizer path independently testable. diff --git a/docs/sandbox-rollout.md b/docs/sandbox-rollout.md index 6fdc9d19a..d31b65cae 100644 --- a/docs/sandbox-rollout.md +++ b/docs/sandbox-rollout.md @@ -43,6 +43,105 @@ The optional `installation` block supports a pinned offline Claude binary: `mode: offline_binary`, `version`, `binary`, and its 64-character `sha256`. Without it, Harbor performs normal agent setup, which can need network access. +For original tasks that define a Dockerfile without a prebuilt image, a binding +can include `task_digests: [SHA256]`. These are HarborRL task tree digests, +including file bytes and modes, rather than registry archive digests. Build the +task's original environment, pin the resulting image in the binding, and keep +the task files unchanged. The runner applies that image to a copy of Harbor's +in-memory environment config. A digest-bound image only accepts the listed +tasks; changed content cannot fall back to the image name. Each task digest +may appear in only one binding. Bindings without `task_digests` retain exact +image-name matching. + +The Claude Code harness profile accepts optional `append_system_prompt` as a +nonempty string. Harbor passes it literally to the pinned CLI's +`--append-system-prompt` option. Use it for explicit harness conventions when +the original task expects a different tool interface, and freeze the same +profile for training and evaluation. The original task files and verifier stay +unchanged; tools execute through the native CLI's declared tool interfaces. +The option is omitted by default. + +Optional `environment_timeout_sec` is a positive integer in the harness profile. +It sets Harbor's outer environment startup budget through the public +`environment_build_timeout_multiplier`, relative to the unchanged task's +`environment.build_timeout_sec` (Harbor's default is 600 seconds). Omitting it +preserves the task's startup budget. This is separate from `setup_timeout_sec`, +which covers agent installation, and from Sandbox SDK `create_timeout_s` and +`request_timeout_s`. For example, use `environment_timeout_sec: 720` to allow +a 600-second SDK creation budget plus runner overhead. Freeze all these +settings before training and evaluation. +The Runner reserves a process-shared startup slot in Harbor's public +`ENVIRONMENT_START` hook before the environment timer starts. Queue waiting +therefore does not consume `environment_timeout_sec`; the coordinator's total +attempt deadline still bounds the wait. The reservation is released after +creation/readiness or failure, allowing ready agents to run concurrently. +Startup lifecycle evidence distinguishes `waiting`, `reserved` and actual +allocation. `sandbox-requests/` records request method, lifecycle operation, +status or exception type, and elapsed time; it excludes raw URLs, credentials +and response bodies. Authentication failures remain fatal and are not retried. + +For pinned Claude Code `2.1.141`, exhausting `max_turns` can produce a valid +`budget_truncated` trajectory. Acceptance requires the matching native CLI +initialization and final `error_max_turns` result, the configured number of +logical turns, a matching session, a final tool-use response, +the completed original verifier result and reward artifact, and confirmed +Sandbox cleanup. The original reward is preserved, including zero; the +trajectory is not retried for its reward. Other agent exceptions, timeouts, +missing responses and unverified CLI versions remain invalid. +The pinned CLI excludes token-limit continuations from its logical turn count. +Each additional response requires a consumed text-only `max_tokens` response +and the immediately following request with the same history, exact assistant +content and pinned native continuation message. Only probed ephemeral cache +placement and string/text-block serialization may differ. The audit records +both response IDs and their trace digests; every sampled token, logprob and +response remains in the trajectory. Arbitrary extra responses are rejected. + +Native session consumption permits the CLI to add omitted optional tool +arguments only when they exactly match defaults in that request's declared +tool schema. The trace records the added defaults, schema digest and both +content digests. Sampled model output, token IDs and logprobs remain unchanged. +Changed arguments, added required fields and undeclared defaults are rejected. +The pinned CLI `2.1.141` also removes trailing ECMAScript whitespace on each +line of `Edit.new_string` and `Write.content` for non-Markdown paths. Session +audit accepts only this exact transformation, with matching profile/session +versions and a declared string field. Line separators, indentation, other +arguments and Markdown content remain exact. The trace records a versioned +`consumption_tool_normalization` proof, schema/value digests and both content +digests. The original sampled response and token/logprob arrays stay intact. +Other CLI input transformations still reject the attempt. + +For pinned Claude Code `2.1.141`, Qwen calls to unavailable tools retain the +generated name, arguments, text, token IDs and logprobs. The CLI returns its +native `No such tool available` observation and can continue the episode. +The gateway accepts that history only with the exact matching error; session +audit independently requires one matching rejection for each unavailable call. +Missing, duplicate or successful results for unavailable tools reject the +attempt. The trace records `unavailable_tool_calls` and `tool_rejection_ref`. +Malformed tool JSON, invalid arguments, thinking output and unrelated API +errors still invalidate the attempt. Rewards come from the original verifier. + +Sandbox creation and readiness failures write `sandbox-startup-failed.json` +with the stage, allocation ID, configured timeouts and exception cause types. +SDK error messages, credentials and signed URLs are excluded. A trial that +fails before creating a native session reports its Harbor exception type +before session auditing. Set `request_timeout_s` to cover slow endpoint +discovery within `create_timeout_s`; increasing the harness setup timeout +alone does not change either limit. + +Set optional `training.seed` to an integer in `0..2147483647` to seed the native +trainer and task order and enable deterministic SGLang inference. Each policy +request derives its sampling seed from that seed, the batch, group, slot, task +digest and input tokens. Retries keep the same seed for the same input, and +concurrent dispatch order does not change it. The actual seed is recorded in +each turn's `serving_input.sampling_params`. External tool output and Sandbox +scheduling can still vary, so this does not guarantee identical complete agent +episodes. Omitting the option keeps the backend defaults. + +Set `training.export_hf: true` to export HF weights under +`hf/rollout-` at each checkpoint save, including the final rollout, for +independent model serving and evaluation. Optimizer checkpoints remain under +`checkpoints/`. The export option defaults to `false`. + ## Concurrency and evidence `max_live_instances` bounds all active instances; each image binding has its @@ -52,9 +151,12 @@ using process-shared locks. Shared-filesystem flock support is required when runners use the same lock directory across hosts. Groups use the same identity, policy barrier, verifier and consumption audit -as Docker. Failed slots cancel their group siblings; submitted independent -groups finish before a partial batch is rejected. Unconfirmed cleanup blocks -new admission. Production export never accepts synthetic evidence. +as Docker. Failed slots cancel their group siblings. Direct Slime rollout +batches cancel and await all remaining groups and their slot cleanup before +closing the runtime and propagating a failure or cancellation. RolloutService +requests wait for submitted independent groups before rejecting a partial +batch. Unconfirmed cleanup blocks new admission. Production export never +accepts synthetic evidence. ## Checks diff --git a/docs/train-inference-consistency.md b/docs/train-inference-consistency.md index f4bd4d2d2..912a7e51d 100644 --- a/docs/train-inference-consistency.md +++ b/docs/train-inference-consistency.md @@ -9,11 +9,20 @@ The three policies are the observed behavior `q`, frozen pre-update trainer drift is `log pθ - log p0`. The Native loss remains unchanged. No TIS, clipping correction, or new training objective is introduced. +## Service sampling lock + +The service client computes the sampling digest from the original group request, +normalizing equivalent JSON probabilities (`1` and `1.0`) the same way as the +executor and Gateway. Each slot descriptor and trajectory identity must match +that requested digest. Native readiness validation then binds every turn's +sampling evidence to the trajectory identity. Internally consistent manifests +from a different distribution are rejected before conversion to training samples. + ## Configuration -Native schema 2 accepts an optional `consistency` mapping. Its default mode is -`off`. Existing configurations remain valid. `audit` records numerical violations; -`enforce` rejects violations before +Both Native schema 2 and the Transformers math configuration accept an optional +`consistency` mapping. Its default mode is `off`. Existing configurations remain +valid. `audit` records numerical violations; `enforce` rejects violations before training or serving admission. Invalid evidence always rejects, including in audit mode. Enforce requires an explicit `max_abs_delta`. Calibrate tolerances for the actual model, precision, parallel layout and kernels. @@ -62,11 +71,28 @@ fraction, bounded worst positions and histogram percentile **upper bounds**. They do not copy raw prompt or response text. TP duplicate counts are excluded. Distributed validation collects rank errors before all ranks reject. +## Serving text and token evidence + +The text passed to the model-output parser must equal the complete visible +decoding of the generation token IDs. Prefix agreement and whitespace +normalization do not prove that the responses match. The only accepted omission +is an exact terminal stop string declared by the request and explicitly reported +as matched by SGLang's stop finish reason. Extra text before or after that stop +is rejected. The original generation IDs and logprobs remain unchanged. + +## Transformers mathematical reference + +The optional gate teacher-forces the frozen batch under `no_grad` before any +backward call and uses the same metrics and report format. The pre-existing +`training.behavior_tolerance` check remains active; audit does not loosen it. +Fake teacher scores never stand in for observed student behavior. + ## Validation and limitations Run `pytest -q tests/test_consistency.py` for the dependency-free checks. With -Torch installed, run `pytest -q tests/test_consistency_torch.py` for two-rank -failure propagation. These tests do not run a separate training loop. +Torch and Transformers installed, run `pytest -q tests/test_consistency_torch.py +tests/harborrl/test_distill.py` for real optimizer rejection, parameter changes, +two-rank failure propagation and the math pipeline tests. GPU acceptance must additionally exercise the locked Megatron / SGLang versions on the target hardware: initial sync, actual rollout tokens, optimizer update, diff --git a/docs/web-console.md b/docs/web-console.md index 885486493..66830acb8 100644 --- a/docs/web-console.md +++ b/docs/web-console.md @@ -6,48 +6,32 @@ harborrl web --workspace ./harborrl-web This opens a browser on an available loopback port. Keep the terminal running. `--port 8765` chooses a fixed port; `--no-browser` prints the URL without opening -it. The workspace stores drafts, saved configuration bundles and execution records. +it. The workspace stores saved configuration bundles and execution records. One console process owns a workspace at a time. The web console uses the core Python dependencies; Node, a frontend build and CDN access are not required. ## Configure and launch -1. New users can first select **Run Hello World**. It needs no task JSON, model - or GPU and opens the Runs view with actual execution logs. This CPU walkthrough - runs scripted attempts and a real verifier, then exports a GRPO batch using - clearly marked fixture tokens; it does not run model inference or training. -2. Open **Local example** or **Sandbox example** to view, download or apply a - filled best-practice recipe. Replace example model paths and service hostnames. - In Local, source directory and Python are detected on the machine running the - console and edited once for Controller, GPU and rollout. Storage is separate. - SSH workers and Sandbox expose the paths belonging to their remote roles. -3. Provide checkpoint/reference directories. **Use bundled Hello World task & - profile** creates valid input files locally, or supply your own pinned catalog - and profile. Select **Prepare model** to calculate both fingerprints using the - real tokenizer and selected GPU Python, over SSH when needed. No manual SHA-256 - entry is shown. Changing the checkpoint or target environment clears the old - fingerprints. See [the CLI guide](configuration-wizard.md) for prerequisites. -4. Choose GPU layout, tensor parallelism, rollout batches and sampling settings. +1. Select Local, SSH workers, or Sandbox. Fill controller/GPU/runner paths and + service origins. Source and interpreter paths refer to the selected machine. + The controller source directory and Python executable select the checkout and + environment used by the launched CLI. +2. Provide existing checkpoint/reference directories, a pinned task catalog, + harness profile and real tokenizer/template fingerprints. The + [configuration CLI guide](configuration-wizard.md) describes preparing these + inputs. The console never downloads or generates model inputs. +3. Choose GPU layout, tensor parallelism, rollout batches and sampling settings. Allocation defaults are configuration values, not a model memory-fit estimate. -5. Review the generated YAML and validation. Missing fields appear in a prominent - list with bold labels, help and links that focus the relevant input. Required - inputs also have inline explanations. **Check local files** additionally +4. Review the generated YAML and validation. **Check local files** additionally checks references on the controller. Deferred GPU, connectivity, model and - Sandbox checks remain explicit. **Create runnable configuration** publishes the - validated bundle. -6. Run **Dry-run** to build the actual CLI launch plan. **Doctor** performs live + Sandbox checks remain explicit. Save the bundle under a new configuration name. +5. Run **Dry-run** to build the actual CLI launch plan. **Doctor** performs live checks after confirmation. **Start training** requires a separate confirmation and invokes the existing `harborrl train` flow, including its preflight. Configuration names are limited to letters, digits, underscores and hyphens. -Use **Save draft** at any step, including an empty or incomplete configuration. -Drafts appear separately from Runnable configurations in the sidebar; reopen one -to continue and save another revision under the same name. They survive server -restarts. Stale revisions from another tab are rejected so work is not silently -overwritten. Drafts cannot start Dry-run, Doctor or training. - -Runnable configurations are immutable. Reopen one to launch it again, or edit it -and publish a new numbered version. The advanced answers editor and YAML/JSON import use +Saved configurations are not overwritten. Reopen one to launch it again, or edit +and save under a new name. The advanced answers editor and YAML/JSON import use the existing `answers_version: 1` contract. Download an individual generated file or export answers from the editor. Advanced remote Sandbox runners can be specified through imported answers; the form preset defaults to a local runner. @@ -64,7 +48,7 @@ the project's supported Linux training environment. ## Runs and shutdown -Only one execution is active per workspace, including Hello World, doctor and dry-run. This +Only one execution is active per workspace, including doctor and dry-run. This prevents duplicate launches within this console; it is not a cluster-wide GPU scheduler. Other CLI processes or different workspaces can still use the same resources, and operators must coordinate those allocations. @@ -95,7 +79,9 @@ After a crash or interrupted shutdown, previously running jobs appear as records and Sandbox cleanup evidence, then use **Confirm cleanup** to acknowledge that external cleanup is complete. This action records operator confirmation; it does not kill a stored PID or automatically relaunch a job. The server never -signals a PID recovered from an old session. +signals a PID recovered from an old session. Job details and **Confirm cleanup** +remain available if a log was never created or cannot be read. The log panel +reports that error explicitly; an unavailable log is not evidence of cleanup. ## Remote controller @@ -124,3 +110,18 @@ no-store headers, and assets are packaged locally. Keep the token private: this is a local single-user console with the controller user's filesystem and training privileges. It is not a public multi-user service or an Internet-facing gateway. Do not expose it through a public reverse proxy. + +## Implementation and verification + +- `harborrl/webui/server.py`: loopback HTTP endpoints, authentication, static + resources and composition of existing catalog/builder/validators. +- `harborrl/webui/jobs.py`: serial CLI launch, snapshots, logs and lifecycle records. +- `harborrl/webui/cli.py`: CLI arguments, workspace lock and browser opening. +- `harborrl/webui/static/`: vanilla browser UI and vendored Lucide assets. +- `tests/harborrl/test_webui.py`: HTTP boundaries, offline configuration, real + dry-run, stand-in training launch/cancellation, persistence and process groups. + +Training schema 2 and deployment schema 1 are unchanged. No alternate training +engine, generic shell endpoint, secret editor or platform-specific preset is +introduced. Live training remains a real user action after reviewing the +configuration; tests use synthetic inputs and CPU-only process stand-ins. diff --git a/examples/distill/local-models.yaml b/examples/distill/local-models.yaml new file mode 100644 index 000000000..aa938e45e --- /dev/null +++ b/examples/distill/local-models.yaml @@ -0,0 +1,23 @@ +# Existing local models only. No download, remote code, or cross-tokenizer alignment. +schema_version: 1 +method: opd +tasks: + path: /srv/math/tasks.jsonl +student: + kind: hf_local + checkpoint: /srv/models/student +teacher: + kind: hf_local + checkpoint: /srv/models/teacher +sampling: + max_new_tokens: 64 + max_context: 512 +training: + steps: 2 + batch_size: 2 + microbatch_size: 1 +deployment: + device: cuda + world_size: 1 +output: + directory: /srv/runs/math-opd-new diff --git a/examples/distill/math.jsonl b/examples/distill/math.jsonl new file mode 100644 index 000000000..65ef09212 --- /dev/null +++ b/examples/distill/math.jsonl @@ -0,0 +1,12 @@ +{"task_id":"train-01","problem_id":"addition-1-1","prompt":"What is 1 + 1? Return only the number.","reference":"2","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-02","problem_id":"addition-2-3","prompt":"What is 2 + 3? Return only the number.","reference":"5","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-03","problem_id":"subtraction-9-4","prompt":"What is 9 - 4? Return only the number.","reference":"5","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-04","problem_id":"multiplication-2-4","prompt":"What is 2 * 4? Return only the number.","reference":"8","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-05","problem_id":"division-6-2","prompt":"What is 6 / 2? Return only the number.","reference":"3","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-06","problem_id":"addition-0-0","prompt":"What is 0 + 0? Return only the number.","reference":"0","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-07","problem_id":"subtraction-7-6","prompt":"What is 7 - 6? Return only the number.","reference":"1","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-08","problem_id":"multiplication-3-2","prompt":"What is 3 * 2? Return only the number.","reference":"6","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"validation-01","problem_id":"addition-6-7","prompt":"What is 6 + 7? Return only the number.","reference":"13","verifier":"exact_rational_v1","split":"validation","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"validation-02","problem_id":"division-5-2","prompt":"What is 5 / 2? Return only the number.","reference":"5/2","verifier":"exact_rational_v1","split":"validation","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"test-01","problem_id":"addition-8-7","prompt":"What is 8 + 7? Return only the number.","reference":"15","verifier":"exact_rational_v1","split":"test","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"test-02","problem_id":"subtraction-20-3","prompt":"What is 20 - 3? Return only the number.","reference":"17","verifier":"exact_rational_v1","split":"test","source":"HarborRL original arithmetic fixture v1","license":"MIT"} diff --git a/examples/distill/opd.yaml b/examples/distill/opd.yaml new file mode 100644 index 000000000..b7be5cfcc --- /dev/null +++ b/examples/distill/opd.yaml @@ -0,0 +1,24 @@ +schema_version: 1 +backend_contract: transformers-math-onpolicy-v1 +method: opd +tasks: + path: math.jsonl +student: + kind: tiny_qwen +teacher: + kind: fake +prm: + kind: none +sampling: + max_new_tokens: 8 + max_context: 128 +training: + steps: 2 + batch_size: 8 + microbatch_size: 2 + learning_rate: 0.0003 +deployment: + device: cpu + world_size: 1 +output: + directory: ../../runs/math-opd-smoke diff --git a/examples/distill/prm-opd.yaml b/examples/distill/prm-opd.yaml new file mode 100644 index 000000000..430c996a7 --- /dev/null +++ b/examples/distill/prm-opd.yaml @@ -0,0 +1,22 @@ +schema_version: 1 +method: prm_opd +tasks: + path: math.jsonl +student: + kind: tiny_qwen +teacher: + kind: fake +prm: + kind: fake + weight: 0.2 +sampling: + max_new_tokens: 8 + max_context: 128 +training: + steps: 2 + batch_size: 8 + outcome_weight: 0.1 +deployment: + device: cpu +output: + directory: ../../runs/math-prm-opd-smoke diff --git a/examples/distill/prm.yaml b/examples/distill/prm.yaml new file mode 100644 index 000000000..6598c8ee6 --- /dev/null +++ b/examples/distill/prm.yaml @@ -0,0 +1,20 @@ +schema_version: 1 +method: prm +tasks: + path: math.jsonl +student: + kind: tiny_qwen +teacher: + kind: none +prm: + kind: fake +sampling: + max_new_tokens: 8 + max_context: 128 +training: + steps: 2 + batch_size: 8 +deployment: + device: cpu +output: + directory: ../../runs/math-prm-smoke diff --git a/examples/distill/reinforce.yaml b/examples/distill/reinforce.yaml new file mode 100644 index 000000000..d228e5665 --- /dev/null +++ b/examples/distill/reinforce.yaml @@ -0,0 +1,21 @@ +schema_version: 1 +method: reinforce +tasks: + path: math.jsonl +student: + kind: tiny_qwen +teacher: + kind: none +prm: + kind: none +sampling: + max_new_tokens: 1 + max_context: 128 +training: + steps: 2 + batch_size: 64 + outcome_weight: 1.0 +deployment: + device: cpu +output: + directory: ../../runs/math-outcome-smoke diff --git a/examples/math/tasks.jsonl b/examples/math/tasks.jsonl new file mode 100644 index 000000000..65ef09212 --- /dev/null +++ b/examples/math/tasks.jsonl @@ -0,0 +1,12 @@ +{"task_id":"train-01","problem_id":"addition-1-1","prompt":"What is 1 + 1? Return only the number.","reference":"2","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-02","problem_id":"addition-2-3","prompt":"What is 2 + 3? Return only the number.","reference":"5","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-03","problem_id":"subtraction-9-4","prompt":"What is 9 - 4? Return only the number.","reference":"5","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-04","problem_id":"multiplication-2-4","prompt":"What is 2 * 4? Return only the number.","reference":"8","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-05","problem_id":"division-6-2","prompt":"What is 6 / 2? Return only the number.","reference":"3","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-06","problem_id":"addition-0-0","prompt":"What is 0 + 0? Return only the number.","reference":"0","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-07","problem_id":"subtraction-7-6","prompt":"What is 7 - 6? Return only the number.","reference":"1","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"train-08","problem_id":"multiplication-3-2","prompt":"What is 3 * 2? Return only the number.","reference":"6","verifier":"exact_rational_v1","split":"train","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"validation-01","problem_id":"addition-6-7","prompt":"What is 6 + 7? Return only the number.","reference":"13","verifier":"exact_rational_v1","split":"validation","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"validation-02","problem_id":"division-5-2","prompt":"What is 5 / 2? Return only the number.","reference":"5/2","verifier":"exact_rational_v1","split":"validation","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"test-01","problem_id":"addition-8-7","prompt":"What is 8 + 7? Return only the number.","reference":"15","verifier":"exact_rational_v1","split":"test","source":"HarborRL original arithmetic fixture v1","license":"MIT"} +{"task_id":"test-02","problem_id":"subtraction-20-3","prompt":"What is 20 - 3? Return only the number.","reference":"17","verifier":"exact_rational_v1","split":"test","source":"HarborRL original arithmetic fixture v1","license":"MIT"} diff --git a/examples/math/train.yaml b/examples/math/train.yaml new file mode 100644 index 000000000..51a6837b4 --- /dev/null +++ b/examples/math/train.yaml @@ -0,0 +1,35 @@ +schema_version: 2 +execution: + backend: text_math +algorithm: + name: grpo +tasks: + path: tasks.jsonl +model: + checkpoint: /models/Qwen3-8B + reference: /models/Qwen3-8B_torch_dist + args_file: qwen3-8B +training: + export_hf: true + num_rollout: 3 + learning_rate: 0.000001 + save_interval: 1 + backend_contract: slime-v0.3.2-native-v1 + resume: null +sampling: + group_size: 4 + groups_per_batch: 8 + max_attempts: 1 + max_tokens: 1024 + max_context: 2048 +deployment: + layout: split + num_gpus: 4 + actor_gpus: 2 + rollout_gpus: 2 + actor_tensor_parallel_size: 2 + rollout_gpus_per_engine: 2 +output: + root: /runs/math +consistency: + mode: "off" diff --git a/examples/native/catalog.json b/examples/native/catalog.json index 70854815e..b8b090ada 100644 --- a/examples/native/catalog.json +++ b/examples/native/catalog.json @@ -12,6 +12,6 @@ "scale": 1.0 }, "revision": "v0.1-example", - "task_digest": "2293cd4682597d9c6c5a4bb35d2ded9978873a8deb48ee07fae4e58121c79691" + "task_digest": "905f496d2852cca5b1e86e9de72395b41c291d2b8f6a80807e0ed24ca19541b3" } ] diff --git a/examples/native/catalog.jsonl b/examples/native/catalog.jsonl index 0afb50238..eba16a31c 100644 --- a/examples/native/catalog.jsonl +++ b/examples/native/catalog.jsonl @@ -1 +1 @@ -{"id": "hello_world", "path": "tasks/hello_world", "revision": "v0.1-example", "reward_profile": {"key": "reward", "offset": 0.0, "raw_range": [0.0, 1.0], "scale": 1.0}, "task_digest": "2293cd4682597d9c6c5a4bb35d2ded9978873a8deb48ee07fae4e58121c79691"} +{"id": "hello_world", "path": "tasks/hello_world", "revision": "v0.1-example", "reward_profile": {"key": "reward", "offset": 0.0, "raw_range": [0.0, 1.0], "scale": 1.0}, "task_digest": "905f496d2852cca5b1e86e9de72395b41c291d2b8f6a80807e0ed24ca19541b3"} diff --git a/examples/native/claude-code-profile.json b/examples/native/claude-code-profile.json index 0434f902a..bbc2804a9 100644 --- a/examples/native/claude-code-profile.json +++ b/examples/native/claude-code-profile.json @@ -1,8 +1,8 @@ { - "agent_timeout_sec": 120, - "cli_version": "2.1.141", - "max_turns": 4, + "agent_timeout_sec": 60, + "cli_version": "2.1.0", + "max_turns": 1, "model": "policy", - "setup_timeout_sec": 300, + "setup_timeout_sec": 60, "verifier_timeout_sec": 60 } diff --git a/examples/native/tasks/hello_world/task.toml b/examples/native/tasks/hello_world/task.toml index 8323405d0..f3534bd83 100644 --- a/examples/native/tasks/hello_world/task.toml +++ b/examples/native/tasks/hello_world/task.toml @@ -3,14 +3,14 @@ version = "1.0" [environment] os = "linux" cpus = 1 -memory_mb = 4096 -storage_mb = 4096 -allow_internet = true +memory_mb = 512 +storage_mb = 256 +allow_internet = false gpus = 0 [agent] name = "claude-code" -max_turns = 4 +max_turns = 1 [verifier] type = "shared" diff --git a/examples/native/tasks/hello_world/tests/test.sh b/examples/native/tasks/hello_world/tests/test.sh index 99dd9cd76..d86585c80 100755 --- a/examples/native/tasks/hello_world/tests/test.sh +++ b/examples/native/tasks/hello_world/tests/test.sh @@ -1,9 +1,5 @@ #!/bin/sh set -eu -reward=0 -if [ "$(cat /hello 2>/dev/null || true)" = "ok" ]; then - reward=1 -fi -mkdir -p /logs/verifier -printf '{"reward":%s}\n' "$reward" >/logs/verifier/reward.json +test "$(cat /hello 2>/dev/null || true)" = "ok" +printf '{"reward":1}\n' >reward.json diff --git a/examples/native/train_qwen_native.yaml b/examples/native/train_qwen_native.yaml index ab8165f7e..5a794c789 100644 --- a/examples/native/train_qwen_native.yaml +++ b/examples/native/train_qwen_native.yaml @@ -40,7 +40,9 @@ sampling: group_size: 2 groups_per_batch: 1 max_attempts: 2 - max_context: 8192 + # Claude Code's system/tool prompt alone can exceed 19K input tokens. + # Reserve room for that prompt and the response; longer tasks may need more. + max_context: 32768 max_tokens: 1024 deployment: diff --git a/harborrl/algorithms/opd.py b/harborrl/algorithms/opd.py new file mode 100644 index 000000000..a9d8df660 --- /dev/null +++ b/harborrl/algorithms/opd.py @@ -0,0 +1,33 @@ +"""Fixed-prefix sampled-token reverse-KL surrogate, not full-vocabulary KL.""" + + +def trajectory_surrogate(current, advantage, mask): + import torch + + if current.shape != advantage.shape or current.shape != mask.shape or current.ndim != 2: + raise ValueError("Expected aligned [trajectory, token] tensors") + if not torch.isfinite(current).all() or not torch.isfinite(advantage).all(): + raise ValueError("Nonfinite loss input") + if not ((mask == 0) | (mask == 1)).all() or (mask.sum(-1) == 0).any(): + raise ValueError("Each trajectory must have a nonempty binary mask") + return -((current * advantage.detach() * mask).sum(-1) / mask.sum(-1)).mean() + + +def sampled_opd_loss(current, behavior, teacher, mask): + if behavior.shape != current.shape or teacher.shape != current.shape: + raise ValueError("OPD tensors must align") + return trajectory_surrogate(current, teacher.detach() - behavior.detach(), mask) + + +def weighted_sampled_opd_loss(current, behavior, teacher, weights, mask): + """Ragged backend sum; the shared scheduler supplies the outer group mean.""" + import torch + + if current.ndim != 1 or any(value.shape != current.shape for value in (behavior, teacher, weights, mask)): + raise ValueError("OPD requires aligned flat token tensors") + if any(not torch.isfinite(value).all() for value in (current, behavior, teacher, weights)): + raise ValueError("Nonfinite OPD loss input") + if (weights < 0).any() or not ((mask == 0) | (mask == 1)).all() or not mask.any(): + raise ValueError("OPD requires nonnegative weights and a nonempty binary mask") + signal = teacher.detach() - behavior.detach() + return -(current * signal * weights.detach() * mask.to(current.dtype)).sum() diff --git a/harborrl/backends/slime_v032/advantage.py b/harborrl/backends/slime_v032/advantage.py index 312480af9..cbf641d76 100644 --- a/harborrl/backends/slime_v032/advantage.py +++ b/harborrl/backends/slime_v032/advantage.py @@ -17,7 +17,7 @@ def _response_values(prototype, value): return torch.full( (len(prototype),), float(value), - dtype=prototype.dtype, + dtype=torch.float32, device=prototype.device, ) return [float(value)] * len(prototype) diff --git a/harborrl/backends/slime_v032/config.py b/harborrl/backends/slime_v032/config.py index cc52b864d..3bba35617 100644 --- a/harborrl/backends/slime_v032/config.py +++ b/harborrl/backends/slime_v032/config.py @@ -91,6 +91,8 @@ def launch_plan(config): 'HARBORRL_CONSISTENCY_MODE': c['consistency']['mode'], } environment['HARBORRL_NATIVE_RUNNER_TRANSPORT'] = c['harbor'].get('runner_transport', 'ssh') + if c['training'].get('seed') is not None: + environment['HARBORRL_TRAINING_SEED'] = c['training']['seed'] if sandbox is not None: environment['HARBORRL_NATIVE_SANDBOX'] = json.dumps(sandbox, sort_keys=True) return {'config': config, 'math': selection.asdict(), @@ -117,6 +119,8 @@ def field_usage(): "harness.name": ["schema:claude_code", "runner:Claude Code"], "harbor.version": ["preflight:worker runtime version"], "training.backend_contract": ["environment:HARBORRL_NATIVE_SLIME_CONTRACT", "preflight:backend hooks"], + "training.seed": ["environment:HARBORRL_TRAINING_SEED"], + "training.export_hf": ["run_environment:SAVE_HF"], "harbor.workers": ["environment:HARBORRL_NATIVE_RUNNER_WORKERS", "preflight:worker probe"], "harbor.runner_transport": ["environment:HARBORRL_NATIVE_RUNNER_TRANSPORT"], "harbor.environment": ["environment:HARBORRL_NATIVE_SANDBOX", "runner:environment"], diff --git a/harborrl/backends/slime_v032/dapo.py b/harborrl/backends/slime_v032/dapo.py new file mode 100644 index 000000000..34139aa0c --- /dev/null +++ b/harborrl/backends/slime_v032/dapo.py @@ -0,0 +1,28 @@ +"""DAPO token sums; Megatron supplies the global valid-token denominator.""" +from harborrl.backends.slime_v032.loss import clipped_loss + + +def loss_function(args, batch, logits, unused_reducer): + import torch + from slime.backends.megatron_utils.loss import get_log_probs_and_entropy + + if not args.calculate_per_token_loss or args.context_parallel_size != 1: + raise ValueError('DAPO requires global per-token normalization and CP=1') + _, output = get_log_probs_and_entropy( + logits, args=args, unconcat_tokens=batch['unconcat_tokens'], + total_lengths=batch['total_lengths'], response_lengths=batch['response_lengths'], + with_entropy=False) + current = torch.cat(output['log_probs']) + old, advantages, masks = (torch.cat(batch[key]).detach() for key in ( + 'rollout_log_probs', 'advantages', 'loss_masks')) + # The pinned Slime denominator clamps each sample's token count to >=1. + # This DAPO variant retains every response token, including truncations. + if not (masks == 1).all(): + raise ValueError('DAPO requires all response tokens, including truncations') + loss = clipped_loss(current, old, advantages, torch.ones_like(current), + args.eps_clip, masks, epsilon_high=args.eps_clip_high) + ratio = (current.detach() - old).exp() + clipped = ((advantages > 0) & (ratio > 1 + args.eps_clip_high)) | ( + (advantages < 0) & (ratio < 1 - args.eps_clip)) + return loss, {'dapo_pg_loss': loss.detach(), 'dapo_token_count': masks.sum(), + 'dapo_clipped_tokens': clipped.sum()} diff --git a/harborrl/backends/slime_v032/loss.py b/harborrl/backends/slime_v032/loss.py index 502c146eb..461320226 100644 --- a/harborrl/backends/slime_v032/loss.py +++ b/harborrl/backends/slime_v032/loss.py @@ -6,9 +6,9 @@ from harborrl.export.validate import token_weight from harborrl.trajectories.native import finite -# Slime v0.3.2 rescales custom loss by microbatch, step rollout count, and DP -# size outside this hook. Keep the correction explicit while GPU acceptance -# determines whether the official scaling already preserves Native semantics. +# The pinned Slime scheduler counts distinct rollout ids, one per group. +# Its outer microbatch/DP scaling therefore averages these group-weighted +# sums over groups, with no additional group-size correction. LOSS_SCALE_CORRECTION = 1.0 @@ -40,7 +40,7 @@ def clipped_pg_loss_pure(new_logprobs, old_logprobs, advantages, weights, epsilo return total -def clipped_loss(new_logprobs, old_logprobs, advantages, weights, epsilon=0.2, loss_masks=None): +def clipped_loss(new_logprobs, old_logprobs, advantages, weights, epsilon=0.2, loss_masks=None, epsilon_high=None): """Global sum of weighted PPO losses; usable on CPU for gradient contracts.""" import torch if not (new_logprobs.shape == old_logprobs.shape == advantages.shape == weights.shape): @@ -49,10 +49,13 @@ def clipped_loss(new_logprobs, old_logprobs, advantages, weights, epsilon=0.2, l raise ValueError("native loss mask must match token weights") if not 0 < epsilon < float("inf"): raise ValueError("native clipping epsilon must be finite and positive") + epsilon_high = epsilon if epsilon_high is None else epsilon_high + if not 0 < epsilon_high < float('inf'): + raise ValueError('native upper clipping epsilon must be finite and positive') if any(not torch.isfinite(tensor).all() for tensor in (new_logprobs, old_logprobs, advantages, weights)) or (weights < 0).any(): raise ValueError("nonfinite native loss input or negative weight") ratio = (new_logprobs - old_logprobs).exp() - loss = torch.maximum(-advantages * ratio, -advantages * ratio.clamp(1 - epsilon, 1 + epsilon)) + loss = torch.maximum(-advantages * ratio, -advantages * ratio.clamp(1 - epsilon, 1 + epsilon_high)) if loss_masks is not None: weights = weights * loss_masks.to(device=weights.device, dtype=weights.dtype) return (loss * weights).sum() diff --git a/harborrl/backends/slime_v032/math_rollout.py b/harborrl/backends/slime_v032/math_rollout.py new file mode 100644 index 000000000..aa1d18713 --- /dev/null +++ b/harborrl/backends/slime_v032/math_rollout.py @@ -0,0 +1,353 @@ +"""Single-turn math evidence adapter for the shared Slime training loop.""" +from __future__ import annotations + +import asyncio +import json +import math +import os +from dataclasses import asdict +from pathlib import Path +from urllib import request + +from harborrl.backends.slime_v032.converter import STANDARD_FIELDS +from harborrl.export.contract import ( + DAPO_REDUCTION, + OPD_REDUCTION, + REDUCTION, + SCHEMA_VERSION, + ProcessRewards, + TokenSpan, + TrainingBatch, + TrainingGroup, + TrajectoryTrainingUnit, +) +from harborrl.export.validate import validate_batch +from harborrl.math import dapo +from harborrl.math import prm as prm_api +from harborrl.math import teacher as teacher_api +from harborrl.math.data import verify +from harborrl.trajectories.native import digest, finite, publish + + +def _request(url, payload=None): + opener = request.build_opener(request.ProxyHandler({})) + req = request.Request(url, data=None if payload is None else json.dumps(payload).encode(), + headers={'Content-Type': 'application/json'}) + with opener.open(req, timeout=600) as response: + return json.load(response) + + +def _pool(endpoint): + from harborrl.backends.slime_v032.rollout_service_client import fetch_pool_states + states = fetch_pool_states(endpoint) + versions = {state['weight_version'] for state in states} + if len(versions) != 1 or any(not state['healthy'] for state in states): + raise ValueError('math rollout requires one healthy synchronized policy pool') + return states, next(iter(versions)) + + +def build_batch(args, samples): + """Rebuild neutral groups from exact serving evidence, never Harbor IR.""" + from collections import defaultdict + algorithm = os.environ.get('HARBORRL_MATH_ALGORITHM', 'grpo') + if algorithm not in ('grpo', 'opd', 'dapo'): + raise ValueError('Unknown math algorithm') + teacher = teacher_api.load(os.environ['HARBORRL_OPD_TEACHER']) if algorithm == 'opd' else None + process = prm_api.load(os.environ['HARBORRL_PROCESS_REWARD']) if os.environ.get('HARBORRL_PROCESS_REWARD') else None + if process is not None: + from slime.utils.processing_utils import load_tokenizer + if algorithm != 'grpo': + raise ValueError('PRM currently requires GRPO') + tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + settings = dapo.environment(args.rollout_max_response_len) if algorithm == 'dapo' else None + groups = defaultdict(list) + for sample in samples: + evidence = sample.metadata['math_evidence'] + if evidence.get('algorithm', 'grpo') != algorithm: + raise ValueError('Math sample algorithm changed') + if (sample.tokens != evidence['tokens'] or sample.rollout_log_probs != evidence['logprobs'] + or sample.response != evidence['response'] + or sample.response_length != evidence['response_length'] + or sample.loss_mask != [1] * evidence['response_length']): + raise ValueError('math sample differs from serving evidence') + task = sample.metadata['math_task'] + if (task['split'] != 'train' or digest(task) != sample.metadata['task_digest'] + or evidence['task_digest'] != sample.metadata['task_digest']): + raise ValueError('math task identity or train split changed') + result = verify(sample.response, task['reference'], task['verifier']) + if result != evidence['verification']: + raise ValueError('math verifier evidence changed') + if settings is not None: + expected = dapo.reward(result['reward'], sample.response_length, + args.rollout_max_response_len, settings) + if evidence.get('dapo') != expected or sample.rollout_id != evidence['rollout_id']: + raise ValueError('DAPO reward or rollout identity changed') + elif 'dapo' in evidence: + raise ValueError('Only DAPO may consume shaped rewards') + if teacher is not None: + scores = evidence.get('teacher', {}) + if (scores.get('identity') != teacher['identity'] + or scores.get('weight_version') != teacher['weight_version'] + or scores.get('tokens_digest') != digest(sample.tokens) + or scores.get('logprobs') is None + or getattr(sample, 'teacher_log_probs', None) != scores['logprobs']): + raise ValueError('OPD teacher identity or exact-token evidence changed') + elif 'teacher' in evidence or getattr(sample, 'teacher_log_probs', None) is not None: + raise ValueError('Only OPD may consume teacher evidence') + if process is not None: + if tokenizer.decode(sample.tokens[-sample.response_length:], skip_special_tokens=True) != sample.response: + raise ValueError('PRM response text differs from exact student tokens') + prm_api.validate_receipt(process, evidence.get('process_reward', {}), task['prompt'], + sample.tokens[-sample.response_length:], tokenizer) + elif 'process_reward' in evidence: + raise ValueError('Process rewards require a frozen PRM manifest') + groups[evidence['group_key']].append((sample, evidence, result)) + output, versions = [], set() + for key, rows in groups.items(): + if len(rows) != args.n_samples_per_prompt: + raise ValueError('incomplete math group') + if len({row[1]['task_digest'] for row in rows}) != 1: + raise ValueError('math group mixes tasks') + if settings is not None and len({row[2]['reward'] for row in rows}) == 1: + raise ValueError('DAPO must filter zero-variance outcome groups') + rewards = [row[1]['dapo']['reward'] if settings is not None else row[2]['reward'] for row in rows] + if process is not None: + rewards = [row[1]['process_reward']['token_returns'][0] for row in rows] + mean = math.fsum(rewards) / len(rows) + std = math.sqrt(math.fsum((reward - mean) ** 2 for reward in rewards) / len(rows)) + units = [] + for (sample, evidence, result), reward in zip(rows, rewards, strict=True): + response_length = evidence['response_length'] + prompt_length = len(sample.tokens) - response_length + version = evidence['policy_version'] + versions.add(version) + span = TokenSpan(0, sample.tokens, prompt_length, response_length, + [False] * prompt_length + [True] * response_length, + sample.rollout_log_probs, version, + sample.teacher_log_probs if teacher is not None else None, + ProcessRewards(process['identity'], evidence['process_reward']['step_ends'], + evidence['process_reward']['probabilities'], + evidence['process_reward']['token_returns']) if process else None) + advantage = 0.0 if teacher is not None else reward - mean + if settings is not None: + advantage = dapo.advantage(reward, mean, std, len(rows)) + units.append(TrajectoryTrainingUnit(evidence['trajectory_id'], evidence['rollout_id'], key, + reward, advantage, + response_length, (span,))) + output.append(TrainingGroup(key, len(rows), mean, std, tuple(units))) + reduction = {'grpo': REDUCTION, 'opd': OPD_REDUCTION, 'dapo': DAPO_REDUCTION}[algorithm] + batch = TrainingBatch(SCHEMA_VERSION, tuple(output), reduction, + frozenset(versions), algorithm, teacher['identity'] if teacher else None, + process['identity'] if process else None) + validate_batch(batch) + return batch + + +async def _generate(args, rollout_id, groups): + from slime.utils.processing_utils import load_tokenizer + tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True) + endpoint = f'http://{args.sglang_router_ip}:{args.sglang_router_port}' + before, version = await asyncio.to_thread(_pool, endpoint) + algorithm = os.environ.get('HARBORRL_MATH_ALGORITHM', 'grpo') + settings = dapo.environment(args.rollout_max_response_len) if algorithm == 'dapo' else None + teacher = teacher_api.load(os.environ['HARBORRL_OPD_TEACHER']) if algorithm == 'opd' else None + process = prm_api.load(os.environ['HARBORRL_PROCESS_REWARD']) if os.environ.get('HARBORRL_PROCESS_REWARD') else None + if process is not None: + if algorithm != 'grpo' or process['endpoint'].rstrip('/') == endpoint: + raise ValueError('PRM requires GRPO and a separate frozen scorer endpoint') + process_before = await asyncio.to_thread(prm_api.check, process) + if teacher is not None: + if teacher['endpoint'].rstrip('/') == endpoint: + raise ValueError('OPD teacher must use a separate frozen serving endpoint') + teacher_before = await asyncio.to_thread(teacher_api.check, teacher) + semaphore = asyncio.Semaphore(max(1, args.sglang_server_concurrency)) + + async def one(sample, group_key, group_rollout_id): + async with semaphore: + task = sample.metadata['math_task'] + if task['split'] != 'train' or digest(task) != sample.metadata['task_digest']: + raise ValueError('math generation requires an unchanged train task') + prompt = tokenizer.encode(sample.prompt, add_special_tokens=False) + if len(prompt) + args.rollout_max_response_len > args.rollout_max_context_len: + raise ValueError('math prompt exceeds the context budget') + payload = {'input_ids': prompt, 'return_logprob': True, 'stream': False, + 'sampling_params': {'temperature': 1.0, 'top_p': 1.0, 'top_k': -1, + 'max_new_tokens': args.rollout_max_response_len, + 'skip_special_tokens': True, 'no_stop_trim': True}} + if os.environ.get('HARBORRL_MATH_SEED'): + payload['sampling_params']['sampling_seed'] = (args.rollout_seed + sample.index) % 2**31 + raw = await asyncio.to_thread(_request, endpoint + '/generate', payload) + meta = raw['meta_info'] + if str(meta.get('weight_version', '')) != version: + raise ValueError('math generation served a stale or missing policy version') + rows = meta.get('output_token_logprobs') + if not rows or any(len(row) < 2 or not finite(row[0]) or row[0] > 0 + or type(row[1]) is not int or row[1] < 0 for row in rows): + raise ValueError('math generation requires exact token logprobs') + response_ids = [row[1] for row in rows] + text = tokenizer.decode(response_ids, skip_special_tokens=True) + if text != raw['text']: + raise ValueError('math serving text differs from output tokens') + finish = meta['finish_reason']['type'] + if finish not in ('stop', 'length'): + raise ValueError('unsupported math finish reason') + sample.tokens = prompt + response_ids + sample.response = text + sample.response_length = len(response_ids) + sample.rollout_log_probs = [row[0] for row in rows] + sample.loss_mask = [1] * len(response_ids) + sample.rollout_id = group_rollout_id + sample.status = sample.Status.TRUNCATED if finish == 'length' else sample.Status.COMPLETED + sample.weight_versions = [version] + result = verify(text, task['reference'], task['verifier']) + sample.reward = {'score': result['reward']} + evidence = {'tokens': sample.tokens, 'response': text, 'response_length': len(response_ids), + 'logprobs': sample.rollout_log_probs, 'policy_version': version, + 'group_key': group_key, 'trajectory_id': f'{group_key}:{sample.index}', + 'rollout_id': group_rollout_id, 'task_digest': digest(task), + 'verification': result, 'truncated': finish == 'length', + 'serving_input': payload, 'serving_output_digest': digest(raw)} + sample.metadata['math_evidence'] = evidence + evidence['algorithm'] = algorithm + if settings is not None: + evidence['dapo'] = dapo.reward(result['reward'], sample.response_length, + args.rollout_max_response_len, settings) + sample.reward = {'score': evidence['dapo']['reward']} + if teacher is not None: + scores = await asyncio.to_thread(teacher_api.score_response, teacher, prompt, response_ids) + sample.teacher_log_probs = scores['logprobs'] + evidence['teacher'] = scores + if process is not None: + scores = await asyncio.to_thread(prm_api.score_response, process, task['prompt'], response_ids, tokenizer) + evidence['process_reward'] = scores + sample.reward = {'score': scores['token_returns'][0]} + return sample + + pending = [] + for group_index, group in enumerate(groups): + if len({sample.metadata['task_digest'] for sample in group}) != 1: + raise ValueError('math prompt group mixes tasks') + if settings is not None: + source_ids = {sample.group_index for sample in group} + if len(source_ids) != 1 or type(group[0].group_index) is not int or group[0].group_index < 0: + raise ValueError('DAPO requires persistent data-source group indices') + group_index = group[0].group_index + group_key = digest({'rollout': rollout_id, 'group': group_index, + 'task': group[0].metadata['task_digest'], 'policy': version}) + group_rollout_id = group_index if settings is not None else rollout_id * args.rollout_batch_size + group_index + pending.extend(one(sample, group_key, group_rollout_id) for sample in group) + samples = await asyncio.gather(*pending) + after, after_version = await asyncio.to_thread(_pool, endpoint) + if before != after or after_version != version: + raise ValueError('policy pool changed during synchronous math rollout') + if teacher is not None and await asyncio.to_thread(teacher_api.check, teacher) != teacher_before: + raise ValueError('OPD teacher changed during scoring') + if process is not None and await asyncio.to_thread(prm_api.check, process) != process_before: + raise ValueError('PRM changed during scoring') + return samples + + +def generate_rollout(args, rollout_id, data_source, evaluation=False): + from slime.rollout.base_types import RolloutFnTrainOutput + if evaluation: + raise ValueError('math evaluation must use its independent locked split') + if (args.rollout_temperature, args.rollout_top_p, args.rollout_top_k) != (1.0, 1.0, -1): + raise ValueError('math training requires unmodified raw-model sampling') + root = Path(os.environ['RUN_DIR']) / 'math' / f'{rollout_id:06d}' + algorithm = os.environ.get('HARBORRL_MATH_ALGORITHM', 'grpo') + settings = dapo.environment(args.rollout_max_response_len) if algorithm == 'dapo' else None + samples, versions = [], set() + selection = {'status': 'sampling', 'accepted': [], 'rejected': [], 'generation_batches': 0} + limit = settings['max_refill_batches'] if settings is not None else 1 + for attempt in range(limit): + remaining = args.rollout_batch_size - len(samples) // args.n_samples_per_prompt + groups = [list(group) for group in data_source.get_samples(remaining)] + if len(groups) != remaining or any(len(group) != args.n_samples_per_prompt for group in groups): + raise ValueError('incomplete math rollout batch') + generated = asyncio.run(_generate(args, rollout_id, groups)) + if len(generated) != remaining * args.n_samples_per_prompt: + raise ValueError('incomplete math generation batch') + versions.update(sample.metadata['math_evidence']['policy_version'] for sample in generated) + if len(versions) != 1: + raise ValueError('policy changed during math refill sampling') + if settings is None: + samples = generated + break + publish(root / f'sampling-evidence-{attempt:03d}.json', [sample.metadata for sample in generated]) + for start in range(0, len(generated), args.n_samples_per_prompt): + group = generated[start:start + args.n_samples_per_prompt] + keys = {sample.metadata['math_evidence']['group_key'] for sample in group} + if len(keys) != 1: + raise ValueError('DAPO generation mixed prompt groups') + key = next(iter(keys)) + if key in selection['accepted'] or key in selection['rejected']: + raise ValueError('DAPO generated a duplicate prompt group') + outcomes = [sample.metadata['math_evidence']['verification']['reward'] for sample in group] + if len(set(outcomes)) == 1: + selection['rejected'].append(key) + else: + selection['accepted'].append(key) + samples.extend(group) + selection['generation_batches'] = attempt + 1 + complete = len(samples) == args.rollout_batch_size * args.n_samples_per_prompt + selection['status'] = 'complete' if complete else 'sampling' + publish(root / f'sampling-selection-{attempt:03d}.json', selection) + if complete: + publish(root / 'sampling.json', selection) + break + else: + selection['status'] = 'exhausted' + publish(root / 'sampling.json', selection) + raise ValueError(f'DAPO refill exhausted after {limit} batches: ' + f'{len(selection["accepted"])} accepted groups; see {root}') + batch = build_batch(args, samples) + publish(root / 'evidence.json', [sample.metadata for sample in samples]) + publish(root / 'batch.json', {**asdict(batch), 'policy_versions': sorted(batch.policy_versions)}) + statuses = [sample.metadata['math_evidence']['verification']['status'] for sample in samples] + metrics = {f'math_{status}': float(statuses.count(status)) + for status in ('correct', 'incorrect', 'unparseable', 'unsupported')} + metrics.update(math_group_count=float(len(batch.groups)), + math_nonzero_group_count=float(sum(group.reward_std > 0 for group in batch.groups)), + math_response_tokens=float(sum(sample.response_length for sample in samples))) + if batch.process_reward_identity is not None: + scores = [sample.metadata['math_evidence']['process_reward'] for sample in samples] + probabilities = [value for receipt in scores for value in receipt['probabilities']] + metrics.update(prm_step_count=float(len(probabilities)), + prm_probability_mean=math.fsum(probabilities) / len(probabilities), + prm_total_reward_mean=math.fsum(row['token_returns'][0] for row in scores) / len(scores), + prm_truncated_responses=float(sum(sample.metadata['math_evidence']['truncated'] + for sample in samples))) + if settings is not None: + penalties = [sample.metadata['math_evidence']['dapo']['overlong_penalty'] for sample in samples] + metrics.update(dapo_generated_groups=float(len(selection['accepted']) + len(selection['rejected'])), + dapo_rejected_groups=float(len(selection['rejected'])), + dapo_generation_batches=float(selection['generation_batches']), + dapo_overlong_responses=float(sum(penalty < 0 for penalty in penalties)), + dapo_overlong_penalty_sum=math.fsum(penalties), + dapo_truncated_responses=float(sum(sample.metadata['math_evidence']['truncated'] + for sample in samples))) + return RolloutFnTrainOutput(samples=samples, metrics=metrics) + + +def convert_samples_to_train_data(args, samples): + samples = [item for sample in samples for item in (sample if isinstance(sample, (list, tuple)) else [sample])] + batch = build_batch(args, samples) + units = {unit.trajectory_id: unit for group in batch.groups for unit in group.trajectories} + groups = {group.group_key: group for group in batch.groups} + data = {key: [] for key in STANDARD_FIELDS} + if batch.algorithm == 'opd': + data['teacher_log_probs'] = [] + for sample in samples: + evidence = sample.metadata['math_evidence'] + unit = units[evidence['trajectory_id']] + values = (sample.tokens, sample.response_length, unit.advantage, evidence['verification']['reward'], + int(evidence['truncated']), sample.index, unit.rollout_id, + sample.loss_mask, sample.rollout_log_probs, unit.token_count) + if batch.process_reward_identity is not None: + from harborrl.backends.slime_v032.prm import reward_payload + values = (*values[:2], reward_payload(unit, groups[unit.group_key], sample.index), *values[3:]) + for key, value in zip(STANDARD_FIELDS, values, strict=True): + data[key].append(value) + if batch.algorithm == 'opd': + data['teacher_log_probs'].append(list(unit.turns[0].teacher_logprobs)) + return data diff --git a/harborrl/backends/slime_v032/opd.py b/harborrl/backends/slime_v032/opd.py new file mode 100644 index 000000000..398719ff5 --- /dev/null +++ b/harborrl/backends/slime_v032/opd.py @@ -0,0 +1,55 @@ +"""Sampled-token OPD hooks for the shared Slime/Megatron optimizer loop.""" +from __future__ import annotations + +from harborrl.algorithms.opd import weighted_sampled_opd_loss +from harborrl.export.validate import token_weight + + +def compute_advantages_and_returns(args, rollout_data): + import torch + + teacher = rollout_data.get('teacher_log_probs') + behavior = rollout_data.get('rollout_log_probs') + masks = rollout_data.get('loss_masks') + if teacher is None or behavior is None or masks is None or not len(teacher): + raise ValueError('OPD requires teacher and behavior token logprobs and masks') + if not len(teacher) == len(behavior) == len(masks): + raise ValueError('OPD sample counts differ') + advantages = [] + for target, old, mask in zip(teacher, behavior, masks, strict=True): + if target.shape != old.shape or old.shape != mask.shape or old.ndim != 1: + raise ValueError('OPD teacher, behavior and mask token positions differ') + if not torch.isfinite(target).all() or not torch.isfinite(old).all(): + raise ValueError('Nonfinite OPD token logprob') + advantages.append((target - old).detach()) + rollout_data['advantages'] = advantages + rollout_data['returns'] = advantages + + +def loss_function(args, batch, logits, unused_reducer): + import torch + from slime.backends.megatron_utils.loss import get_log_probs_and_entropy + + _, output = get_log_probs_and_entropy( + logits, args=args, unconcat_tokens=batch['unconcat_tokens'], + total_lengths=batch['total_lengths'], response_lengths=batch['response_lengths'], + with_entropy=False, + ) + weights = torch.cat([ + torch.full((int(length),), token_weight(int(args.n_samples_per_prompt), int(count)), + dtype=torch.float32, device=logits.device) + for length, count in zip(batch['response_lengths'], batch['rollout_mask_sums'], strict=True) + ]) + current = torch.cat(output['log_probs']) + behavior = torch.cat(batch['rollout_log_probs']) + teacher = torch.cat(batch['teacher_log_probs']) + mask = torch.cat(batch['loss_masks']) + loss = weighted_sampled_opd_loss(current, behavior, teacher, weights, mask) + valid_weights = weights * mask.to(weights.dtype) + return loss, { + 'opd_surrogate_loss': loss.detach(), + 'opd_sampled_reverse_kl': ((behavior - teacher).detach() * valid_weights).sum(), + 'opd_teacher_behavior_abs': ((teacher - behavior).detach().abs() * valid_weights).sum(), + 'opd_weight_sum': valid_weights.sum().detach(), + 'opd_token_count': mask.sum().detach(), + } diff --git a/harborrl/backends/slime_v032/postprocess.py b/harborrl/backends/slime_v032/postprocess.py index 7d4c89bb3..ab8ab1df8 100644 --- a/harborrl/backends/slime_v032/postprocess.py +++ b/harborrl/backends/slime_v032/postprocess.py @@ -2,6 +2,7 @@ from __future__ import annotations import math +import os REQUIRED_FIELDS = ( "tokens", @@ -46,9 +47,28 @@ def _validate_rollout_data(args, rollout_id, rollout_data): raise ValueError("native rollout data fields have mismatched sample counts") if count < 1: raise ValueError("native rollout data is empty") + teacher = rollout_data.get('teacher_log_probs') + algorithm = os.environ.get('HARBORRL_MATH_ALGORITHM', 'grpo') + if algorithm == 'opd': + if teacher is None or len(teacher) != count: + raise ValueError('OPD requires teacher scores for every sample') + elif teacher is not None: + raise ValueError('GRPO must not consume teacher scores') group_size = int(getattr(args, "n_samples_per_prompt", 0)) if group_size < 2: raise ValueError("native training requires a complete group") + if os.environ.get('HARBORRL_PROCESS_REWARD'): + import torch + if algorithm != 'grpo' or teacher is not None: + raise ValueError('PRM requires GRPO without teacher scores') + for field in ('advantages', 'returns'): + values = rollout_data.get(field) + if values is None or len(values) != count: + raise ValueError('PRM requires computed token advantages and returns') + for index, value in enumerate(values): + if (value.ndim != 1 or value.numel() != rollout_data['response_lengths'][index] + or not torch.isfinite(value).all() or value.requires_grad): + raise ValueError('PRM advantages and returns must be aligned finite detached token tensors') group_weights = {} for index in range(count): group_id = _integer_scalar(rollout_data["rollout_ids"][index]) @@ -63,6 +83,13 @@ def _validate_rollout_data(args, rollout_id, rollout_data): raise ValueError("invalid native response length") if len(mask) != response_length or len(logprobs) != response_length: raise ValueError("native response field lengths differ") + if any(flag not in (0, 1) for flag in mask): + raise ValueError('Native response masks must be binary') + if algorithm == 'dapo' and (not all(mask) or mask_sum != response_length): + raise ValueError('DAPO requires all response tokens and exact per-sample token counts') + if teacher is not None and (len(teacher[index]) != response_length or any( + not _finite_scalar(value) or value > 0 for value in teacher[index])): + raise ValueError('OPD teacher token scores are invalid or misaligned') if any(not _finite_scalar(value) for value in logprobs): raise ValueError("nonfinite native rollout logprob") if (not _finite_scalar(rollout_data["rewards"][index]) diff --git a/harborrl/backends/slime_v032/prm.py b/harborrl/backends/slime_v032/prm.py new file mode 100644 index 000000000..4eebd1d1c --- /dev/null +++ b/harborrl/backends/slime_v032/prm.py @@ -0,0 +1,76 @@ +"""Process reward transport through Slime's standard per-sample rewards field.""" +from __future__ import annotations + +import math +import os + +from harborrl.math import prm +from harborrl.trajectories.native import digest, finite + + +def reward_payload(unit, group, sample_index): + span = unit.turns[0] + scores = span.process_rewards + return { + 'kind': prm.ADAPTER, 'identity': scores.identity, + 'tokens_digest': digest(list(span.tokens)), 'sample_index': sample_index, + 'rollout_id': unit.rollout_id, 'trajectory_id': unit.trajectory_id, + 'group_key': group.group_key, 'group_size': group.group_size, + 'group_totals': {member.trajectory_id: member.reward for member in group.trajectories}, + 'step_ends': list(scores.step_ends), 'probabilities': list(scores.probabilities), + 'token_returns': list(scores.token_returns), 'baseline': group.reward_mean, + } + + +def compute_advantages_and_returns(args, rollout_data): + import torch + + frozen = prm.load(os.environ['HARBORRL_PROCESS_REWARD']) + fields = ('rewards', 'tokens', 'response_lengths', 'sample_indices', 'rollout_ids', 'loss_masks', + 'rollout_log_probs', 'rollout_mask_sums') + if (any(key not in rollout_data for key in fields) or not rollout_data['rewards'] + or len({len(rollout_data[key]) for key in fields}) != 1 + or rollout_data.get('teacher_log_probs') is not None): + raise ValueError('PRM requires complete aligned sample fields without teacher scores') + advantages, returns, summaries = [], [], [] + group_records = {} + for index, payload in enumerate(rollout_data['rewards']): + if not isinstance(payload, dict) or payload.get('kind') != prm.ADAPTER: + raise ValueError('PRM requires structured process rewards before advantage computation') + tokens = rollout_data['tokens'][index] + tokens = tokens.tolist() if hasattr(tokens, 'tolist') else list(tokens) + length = rollout_data['response_lengths'][index] + old = rollout_data['rollout_log_probs'][index] + mask = rollout_data['loss_masks'][index] + totals = payload.get('group_totals', {}) + if (payload.get('identity') != frozen['identity'] or payload.get('tokens_digest') != digest(tokens) + or payload.get('sample_index') != rollout_data['sample_indices'][index] + or payload.get('rollout_id') != rollout_data['rollout_ids'][index] + or payload.get('group_size') != args.n_samples_per_prompt + or not isinstance(totals, dict) or len(totals) != args.n_samples_per_prompt + or payload.get('trajectory_id') not in totals or any(not finite(value) for value in totals.values())): + raise ValueError('PRM identity, sample binding or complete-group baseline changed') + values = prm.token_returns(payload.get('step_ends', []), payload.get('probabilities', []), length) + baseline = math.fsum(totals.values()) / len(totals) + if (payload.get('token_returns') != values or payload.get('baseline') != baseline + or not math.isclose(totals[payload['trajectory_id']], values[0], rel_tol=1e-9, abs_tol=1e-9)): + raise ValueError('PRM reward-to-go or complete-group baseline changed') + key = payload.get('group_key') + if not isinstance(key, str) or not key: + raise ValueError('PRM group key must be nonempty') + record = (payload['rollout_id'], totals) + if key in group_records and group_records[key] != record: + raise ValueError('PRM local group members disagree on their complete group') + group_records[key] = record + if (old.ndim != 1 or mask.shape != old.shape or old.numel() != length + or not torch.isfinite(old).all() or not torch.all(mask == 1) + or rollout_data['rollout_mask_sums'][index] != length): + raise ValueError('PRM requires every exact response token and a matching mask') + target = torch.tensor(values, dtype=torch.float32, device=old.device) + returns.append(target) + advantages.append((target - baseline).detach()) + summaries.append(values[0] - baseline) + rollout_data['advantages'] = advantages + rollout_data['returns'] = returns + # Slime's subsequent metric logger expects scalar rewards, not receipts. + rollout_data['rewards'] = summaries diff --git a/harborrl/backends/slime_v032/rollout.py b/harborrl/backends/slime_v032/rollout.py index c8b6ffd9d..822e360a1 100644 --- a/harborrl/backends/slime_v032/rollout.py +++ b/harborrl/backends/slime_v032/rollout.py @@ -16,10 +16,10 @@ from harborrl.export.native import export_training_batch from harborrl.rollout.runtime import ( RolloutRuntimeConfig, + cancel_and_wait, + close_runtime, current_policy_version, digest, - gather_or_cancel, - run_embedded, run_group, runtime, task_entry, @@ -197,8 +197,23 @@ async def run_native_groups(args, groups, policy_version, rollout_id): ir_groups = await run_service_groups(args, groups, rollout_id) assert_locked_policy_versions(ir_groups, policy_version) return ir_groups - results = await gather_or_cancel( - generate_group(args, group, params, rollout_id=rollout_id) for group in groups) + pending = [asyncio.create_task( + generate_group(args, group, params, rollout_id=rollout_id)) for group in groups] + batch = asyncio.gather(*pending) + try: + # Keep parent cancellation from being forwarded twice during cleanup. + results = await asyncio.shield(batch) + except BaseException: + # Groups own their slot cleanup; finish it before shutting down the gateway. + await cancel_and_wait(pending, batch) + # Close the cached runtime here unless an embedded run owns the loop: + # run_embedded's failure handler closes it after the loop unwinds, so + # the original error is what propagates. + from harborrl.rollout import runtime as _runtime + loop = getattr(_runtime, "_EMBEDDED_LOOP", None) + if not (loop is not None and loop.is_running()): + close_runtime() + raise ir_groups = [] for group in results: trajectories = [] @@ -306,7 +321,7 @@ def generate_rollout( if runner is None: async def runner(args_, groups_, policy_version): return await run_native_groups(args_, groups_, policy_version, rollout_id) - ir_groups = run_embedded(runner(args, groups, None)) + ir_groups = asyncio.run(runner(args, groups, None)) policy_version = lock_policy_version(args) assert_locked_policy_versions(ir_groups, policy_version) batch = export_training_batch( diff --git a/harborrl/backends/slime_v032/rollout_service_client.py b/harborrl/backends/slime_v032/rollout_service_client.py index 4fc73e492..161fb24f2 100644 --- a/harborrl/backends/slime_v032/rollout_service_client.py +++ b/harborrl/backends/slime_v032/rollout_service_client.py @@ -13,7 +13,12 @@ from urllib import error, parse, request from harborrl.rollout.runtime import task_entry -from harborrl.trajectories.native import digest, read_json, require_ready +from harborrl.trajectories.native import ( + canonical_sampling, + digest, + read_json, + require_ready, +) SERVICE_SCHEMA = "harborrl-rollout-service-v1" @@ -185,8 +190,11 @@ def load_group_result(run_id, group_request, result, *, allow_synthetic=False): if (not isinstance(slots, list) or len(slots) != len(expected) or len(slots) < 2 or [slot.get("slot_id") for slot in slots] != expected): raise ServiceRolloutError("group result is incomplete or reordered") + sampling_sha = digest(canonical_sampling(group_request["sampling"])) trajectories = [] for descriptor, slot in zip(group_request["slots"], slots): + if slot.get("sampling_digest") != sampling_sha: + raise ServiceRolloutError(f"slot {descriptor['slot_id']} sampling digest differs from the request") path = Path(slot["manifest_path"]) try: payload = read_json(path) @@ -195,6 +203,8 @@ def load_group_result(run_id, group_request, result, *, allow_synthetic=False): if digest(payload) != slot.get("manifest_sha256"): raise ServiceRolloutError(f"slot {descriptor['slot_id']} manifest digest mismatch") identity = payload.get("identity", {}) + if identity.get("sampling_digest") != sampling_sha: + raise ServiceRolloutError(f"slot {descriptor['slot_id']} manifest sampling differs from the request") if (identity.get("policy_version") != group_request["policy_version"] or identity.get("task_digest") != group_request["task"]["task_digest"] or identity.get("run_id") != run_id diff --git a/harborrl/cli.py b/harborrl/cli.py index 59af49567..a39370444 100644 --- a/harborrl/cli.py +++ b/harborrl/cli.py @@ -17,6 +17,10 @@ def main(argv=None): from harborrl.webui.cli import main as web_main return web_main(argv[1:]) + if argv and argv[0] == "distill": + from harborrl.distill.cli import main as distill_main + + return distill_main(argv[1:]) if argv and argv[0] in {"options", "validate", "init"}: from harborrl.options.cli import main as options_main diff --git a/harborrl/config/native.py b/harborrl/config/native.py index 50b3ddb86..43a3a2808 100644 --- a/harborrl/config/native.py +++ b/harborrl/config/native.py @@ -15,12 +15,12 @@ 'harbor': {'python', 'version', 'workers', 'runner_transport', 'environment'}, 'gateway': {'host', 'port', 'advertised_url', 'tokenizer_digest', 'template_digest', 'audited_raw_logprobs'}, 'model': {'checkpoint', 'reference', 'args_file'}, - 'training': {'num_rollout', 'learning_rate', 'save_interval', 'backend_contract'}, + 'training': {'num_rollout', 'learning_rate', 'save_interval', 'backend_contract', 'seed', 'export_hf'}, 'sampling': {'group_size', 'groups_per_batch', 'max_attempts', 'max_tokens', 'max_context'}, 'deployment': {'layout', 'num_gpus', 'actor_gpus', 'rollout_gpus', 'actor_tensor_parallel_size', 'rollout_gpus_per_engine'}, 'output': {'root'}, } -OPTIONAL = {'harbor': {'runner_transport', 'environment'}} +OPTIONAL = {'harbor': {'runner_transport', 'environment'}, 'training': {'seed', 'export_hf'}} class SelectedConfig(dict): @@ -40,6 +40,9 @@ def load_config(path, overrides=None, task_ids=None): raise TypeError('configuration must be a mapping') if config.get('schema_version') != 2: raise ValueError('native configuration requires schema_version: 2') + if config.get('execution', {}).get('backend') == 'text_math': + from harborrl.math.config import load_config as load_math_config + return load_math_config(config, path, overrides, task_ids) validate(config, path) for override in overrides or (): key, separator, value = override.partition('=') @@ -73,6 +76,7 @@ def validate(config, path): raise TypeError('training must be a mapping') config['training'].setdefault( 'backend_contract', 'slime-v0.3.2-native-v1') + config['training'].setdefault('export_hf', False) for section, fields in FIELDS.items(): actual = set(config[section]) if isinstance(config[section], dict) else set() if not isinstance(config[section], dict) or not fields - OPTIONAL.get(section, set()) <= actual <= fields: @@ -97,6 +101,11 @@ def validate(config, path): raise ValueError(f"training.backend_contract must be one of {sorted(CONFIG_CONTRACTS)}") if not finite(config['training']['learning_rate']) or config['training']['learning_rate'] <= 0: raise ValueError('training.learning_rate must be positive and finite') + seed = config['training'].get('seed') + if seed is not None and (type(seed) is not int or not 0 <= seed < 2**31): + raise ValueError('training.seed must be null or an integer in 0..2147483647') + if type(config['training']['export_hf']) is not bool: + raise ValueError('training.export_hf must be a boolean') g, d, s = config['gateway'], config['deployment'], config['sampling'] if type(g['audited_raw_logprobs']) is not bool or not 0 < g['port'] < 65536: raise ValueError('gateway.audited_raw_logprobs must be boolean and gateway.port must be in 1..65535') @@ -177,5 +186,8 @@ def catalog(config): def launch_plan(config): """Compatibility entry point for backend launch mapping.""" + if config['execution']['backend'] == 'text_math': + from harborrl.math.config import launch_plan as build_math_plan + return build_math_plan(config) from harborrl.backends.slime_v032.config import launch_plan as build_plan return build_plan(config) diff --git a/harborrl/consistency/reference.py b/harborrl/consistency/reference.py new file mode 100644 index 000000000..4ce4890a0 --- /dev/null +++ b/harborrl/consistency/reference.py @@ -0,0 +1,25 @@ +"""Frozen pre-update reference check, performed before any backward call.""" +from pathlib import Path + +from .report import assess, distributed_gate + + +def check_reference(model, rows, config, tokenizer): + import torch + import torch.distributed as dist + + from harborrl.distill.models import response_logprobs + + settings = config['consistency'] + pairs = [] + try: + with torch.no_grad(): + for row in rows: + p0 = response_logprobs(model, [row['prompt_ids']], [row['response_ids']], tokenizer.pad_token_id)[0] + pairs.append((p0.detach().float().cpu().tolist(), row['behavior_logprobs'], row['loss_mask'])) + part = assess(pairs, settings, metadata={'backend': 'transformers', 'rank': dist.get_rank() if dist.is_initialized() else 0}) + except (ValueError, TypeError, KeyError, RuntimeError) as exc: + part = {'statistics': None, 'error': str(exc), 'metadata': {'backend': 'transformers'}} + step = rows[0]['sample_id'].split('-')[0] if rows else 'empty' + return distributed_gate(part, settings, Path(config['output']['directory']) / 'consistency' / f'batch-{step}.json', + phase='pre-update-backend-gap') diff --git a/harborrl/consistency/sglang.py b/harborrl/consistency/sglang.py index d9a5ac654..e70b161e9 100644 --- a/harborrl/consistency/sglang.py +++ b/harborrl/consistency/sglang.py @@ -18,22 +18,22 @@ def http_json(url, payload=None, timeout=120): return json.load(response) -def score(endpoint, prompt, response, version, *, transport=http_json): +def score(endpoint, prompt, response, version, *, transport=http_json, label="Canary"): tokens = prompt + response payload = {"input_ids": tokens, "sampling_params": {"temperature": 1.0, "top_p": 1.0, "top_k": -1, "max_new_tokens": 0}, "return_logprob": True, "logprob_start_len": 0} output = transport(endpoint.rstrip("/") + "/generate", payload) meta = output.get("meta_info", {}) if str(meta.get("weight_version", "")) != str(version): - raise ValueError("Canary generation-time weight version mismatch") + raise ValueError(f"{label} generation-time weight version mismatch") rows = meta.get("input_token_logprobs") if (not isinstance(rows, list) or len(rows) != len(tokens) or any(not isinstance(row, (tuple, list)) or len(row) < 2 for row in rows) or [row[1] for row in rows] != tokens): - raise ValueError("Canary input logprobs must preserve every exact token and position") + raise ValueError(f"{label} input logprobs must preserve every exact token and position") values = [row[0] for row in rows[len(prompt):]] if any(not finite(value) or value > 0 for value in values): - raise ValueError("Invalid canary raw logprob") + raise ValueError(f"Invalid {label} raw logprob") return values diff --git a/harborrl/deployment/controller.py b/harborrl/deployment/controller.py index 0ee3700e3..4cb937a64 100644 --- a/harborrl/deployment/controller.py +++ b/harborrl/deployment/controller.py @@ -1,207 +1,465 @@ -"""Service-mode dispatch: the manager-side controller flow (P1 + P3). - -An independent CPU manager runs the same preflight graph, materializes the -run on shared storage, starts the RolloutService control plane, launches the -existing Slime job on the GPU target over SSH, and waits for the shared -exit-code record. Rollout execution itself moves to the manager only through -the service; training stays on the GPU. -""" -from __future__ import annotations - -import contextlib -import json -import os -import signal -import sys -import threading -from pathlib import Path - -from harborrl.config import ROOT -from harborrl.deployment import remote -from harborrl.deployment.schema import destination -from harborrl.deployment.training_target import SshTrainingTarget -from harborrl.platform.evidence import ( - _materialize, - _source_audit, - effective_config, - new_run_root, - run_environment, -) -from harborrl.rollout.service import RolloutService, make_control_server -from harborrl.trajectories.native import publish - -TOKEN_MODE = 0o600 - - -def check_gpu_control(deployment): - manager, gpu = deployment['manager'], deployment['gpu_target'] - timeout = deployment['connectivity']['http_timeout_s'] - result = remote.run_ssh(gpu['ssh'] or {'host': 'localhost'}, - f"curl --noproxy '*' -fsS --max-time {timeout} " + remote.quote(manager['control_url'] + '/healthz'), - timeout=timeout + 15) - if result.returncode != 0: - raise RuntimeError('GPU cannot reach manager control service') - payload = json.loads(result.stdout) - if payload != {'ok': True, 'service': 'harborrl-rollout-service-v1'}: - raise RuntimeError('GPU reached an unexpected control service') - return {'gpu_to_controller': True, 'control_url': manager['control_url']} - - -def _ensure_token(token_file): - """Load the run token, creating a fresh one on first use.""" - path = Path(token_file) - if path.is_file(): - token = path.read_text().strip() - if token: - return token - import secrets - - token = secrets.token_urlsafe(32) - path.parent.mkdir(parents=True, exist_ok=True) - fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, TOKEN_MODE) - with os.fdopen(fd, "w") as stream: - stream.write(token + "\n") - return token - - -def _rollout_environment(plan, deployment, environment): - """Environment overlay that points the GPU job at the manager service.""" - manager = deployment["manager"] - workers = deployment["rollout"]["workers"] - overlay = dict(environment) - overlay.update({ - "HARBORRL_ROLLOUT_MODE": "service", - "HARBORRL_ROLLOUT_SERVICE_URL": manager["control_url"], - "HARBORRL_ROLLOUT_SERVICE_TOKEN_FILE": manager['run_token_file'], - "TRAIN_PYTHON": deployment['gpu_target']['train_python'], - "HARBORRL_NATIVE_GATEWAY_HOST": manager["bind"], - "HARBORRL_NATIVE_GATEWAY_PORT": str(manager["gateway_port"]), - "HARBORRL_NATIVE_GATEWAY_URL": manager["gateway_url"], - "HARBORRL_NATIVE_RUNNER_WORKERS": json.dumps( - [destination(worker["ssh"]) for worker in workers]), - "HARBORRL_NATIVE_RUNNER_PYTHON": (plan['config']['harbor']['python'] - if deployment['rollout']['provider'] == 'sandbox' - else workers[0]['runner_python']), - }) - if deployment['rollout']['provider'] == 'sandbox': - if 'HARBORRL_NATIVE_SANDBOX' not in overlay: - raise ValueError('sandbox deployment requires sandbox training config') - transport = plan['config']['harbor'].get('runner_transport', 'ssh') - if (transport == 'local' and workers) or (transport == 'ssh' and not workers): - raise ValueError('sandbox runner transport and deployment workers disagree') - overlay['HARBORRL_NATIVE_RUNNER_TRANSPORT'] = transport - return overlay - - -def _runtime_factory(deployment, environment, *, concurrency_limit, group_size, tokenizer=None): - """Build the per-run rollout runtime when the GPU job announces serving.""" - - def factory(run_id, router_url): - from harborrl.rollout.runtime import RolloutRuntime, RolloutRuntimeConfig - - env = {**os.environ, **environment} - config = RolloutRuntimeConfig.from_environment( - env, endpoint=router_url, concurrency_limit=concurrency_limit, - group_size=group_size, - gateway_host=deployment["manager"]["bind"], - gateway_port=deployment["manager"]["gateway_port"], - gateway_url=deployment["manager"]["gateway_url"]) - resolved = tokenizer - if resolved is None: - from transformers import AutoTokenizer - - resolved = AutoTokenizer.from_pretrained( - env["HF_CKPT"], trust_remote_code=True) - return RolloutRuntime(config, resolved) - - return factory - - -def dispatch_service(plan, deployment, *, checks): - manager, gpu, storage = deployment["manager"], deployment["gpu_target"], deployment["storage"] - config = plan["config"] - root = new_run_root(storage["run_root"]) - root.mkdir(parents=True, mode=0o700) - run_id = root.name - - catalog_path, prompt_path = _materialize(plan, root) - source = _source_audit(root) - environment = run_environment(plan, root, catalog_path, prompt_path) - token = _ensure_token(manager["run_token_file"]) - environment = _rollout_environment(plan, deployment, environment) - - sampling = config["sampling"] - concurrency = sampling["group_size"] * sampling["groups_per_batch"] - service = RolloutService( - run_root=root, - runtime_factory=_runtime_factory( - deployment, environment, - concurrency_limit=concurrency, group_size=sampling["group_size"]), - token_file=manager["run_token_file"]) - control = make_control_server(service, host=manager["bind"], - port=manager["control_port"]) - control_thread = None - target = None - exit_code = 1 - try: - service.register_run(run_id, { - "run_id": run_id, "deployment_source": plan["deployment"]["source"], - "source_commit": source["commit"], - "config_digest": plan["deployment"]["digest"], - "group_size": sampling["group_size"], - "workers": [worker["id"] for worker in deployment["rollout"]["workers"]], - }, token=token) - control_thread = threading.Thread( - target=control.serve_forever, name="harborrl-control-plane", daemon=True) - control_thread.start() - publish(root / 'control-readiness.json', check_gpu_control(deployment)) - - publish(root / "launch.json", { - "plan": plan, "run_dir": str(root), "source": source, - "materialized_catalog": str(catalog_path), "prompt_data": str(prompt_path), - "control_url": manager["control_url"]}) - publish(root / "effective-config.json", - effective_config(plan, environment, dict(environment), source, checks)) - - target = SshTrainingTarget( - gpu["ssh"] or {"host": "localhost"}, source_dir=gpu["source_dir"], run_dir=root, - environment_file=gpu.get("environment_file"), - command_timeout=deployment["connectivity"]["command_timeout_s"] + 30, - ssh_timeout=deployment["connectivity"]["ssh_timeout_s"]) - # The launcher executes inside the GPU checkout; manager-side absolute - # paths are rewritten so shared-storage artifacts keep one spelling. - target.write_launcher( - ["bash", str(ROOT / "harborrl/platform/slime_train.sh")], environment, - path_rewrites=[(str(manager["source_dir"]), str(gpu["source_dir"]))]) - target.launch() - - def forward_signal(signum, frame): - if target is not None: - target.cancel() - - previous = {sig: signal.signal(sig, forward_signal) - for sig in (signal.SIGINT, signal.SIGTERM)} - try: - exit_code = target.wait() - finally: - for sig, handler in previous.items(): - signal.signal(sig, handler) - except BaseException: - import traceback - - traceback.print_exc() - if target is not None: - # Best-effort teardown; the run record keeps the failure evidence. - with contextlib.suppress(Exception): - target.cancel() - raise - finally: - with contextlib.suppress(Exception): - service.drain_run(run_id) - service.close() - if control_thread is not None: - control.shutdown() - control_thread.join(5) - control.server_close() - print(json.dumps({"run_dir": str(root), "exit_code": exit_code}), file=sys.stderr) - return exit_code +"""Service-mode dispatch: the manager-side controller flow (P1 + P3). + +An independent CPU manager runs the same preflight graph, materializes the +run on shared storage, starts the RolloutService control plane, launches the +existing Slime job on the GPU target over SSH, and waits for the shared +exit-code record. Rollout execution itself moves to the manager only through +the service; training stays on the GPU. +Hub mode extends service mode by allowing the manager to accept multiple runs +as a long-running daemon. Jobs are submitted via HTTP API and queued for +execution on available GPU resources. +""" +from __future__ import annotations + +import contextlib +import json +import os +import signal +import socket +import sys +import threading +from pathlib import Path +from urllib.parse import urlsplit + +from harborrl.config import ROOT +from harborrl.deployment import remote +from harborrl.deployment.schema import destination +from harborrl.deployment.training_target import SshTrainingTarget +from harborrl.platform.evidence import ( + _materialize, + _source_audit, + effective_config, + new_run_root, + run_environment, +) +from harborrl.rollout.service import RolloutService, make_control_server +from harborrl.trajectories.native import publish + +TOKEN_MODE = 0o600 + +PORT_SLOTS = 16 +_ALLOC_LOCK = threading.Lock() +_ALLOCATED = {} # port -> run_id, across control and gateway labels + + +def _allocate_port(bind, base, run_id): + """Pick the first bindable port from ``base`` within PORT_SLOTS candidates. + + The probe closes the socket before the real server binds, so a race can + still fail the later bind loudly; allocation never silently shares a port. + Candidates already held by another label of the same run are skipped so + control and gateway ranges may safely overlap. + """ + with _ALLOC_LOCK: + for candidate in range(base, base + PORT_SLOTS): + if candidate in _ALLOCATED: + continue + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as probe: + probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + probe.bind((bind, candidate)) + except OSError: + continue + _ALLOCATED[candidate] = run_id + return candidate + raise RuntimeError(f"no free manager port within {PORT_SLOTS} slots from {base}") + + +def _rewrite_port(url, port): + parts = urlsplit(url) + host = f"[{parts.hostname}]" if ":" in parts.hostname else parts.hostname + return f"{parts.scheme}://{host}:{port}" + + +def _with_allocated_ports(deployment, run_id): + """Give one dispatch exclusive manager ports; concurrent runs then coexist. + + The deployment's declared ports act as the base of a small allocation + range: every consumer (control server, rollout runtime, GPU environment + overlay, evidence records) reads the patched manager mapping. Every + LOCAL gateway shard claims its own exclusive port: the declared + gateway_port is the first local shard's base and undeclared shard ports + default to consecutive slots, so sibling runs may never land on the same + gateway. Remote shards bind nothing here; their origins pass through. + """ + manager = deployment["manager"] + control = _allocate_port(manager["bind"], manager["control_port"], run_id) + declared = manager.get("gateway_shards") or [{"name": "gateway"}] + allocated = [] + for index, spec in enumerate(declared): + if spec.get("mode", "local") == "remote": + allocated.append(None) + continue + base = spec.get("port", manager["gateway_port"] + index) + allocated.append(_allocate_port(manager["bind"], base, run_id)) + local_ports = [port for port in allocated if port is not None] + patched = dict(manager) + if local_ports: + patched.update({ + "control_port": control, "control_url": _rewrite_port(manager["control_url"], control), + "gateway_port": local_ports[0], "gateway_url": _rewrite_port( + manager["gateway_url"], local_ports[0]), + }) + else: + patched.update({ + "control_port": control, + "control_url": _rewrite_port(manager["control_url"], control), + }) + if manager.get("gateway_shards"): + patched["gateway_shards"] = [ + {**spec, "port": port} if port is not None else dict(spec) + for spec, port in zip(manager["gateway_shards"], allocated)] + return {**deployment, "manager": patched} + + +def _release_run_ports(run_id): + with _ALLOC_LOCK: + for port, owner in list(_ALLOCATED.items()): + if owner == run_id: + del _ALLOCATED[port] + + +def check_gpu_control(deployment): + manager, gpu = deployment['manager'], deployment['gpu_target'] + timeout = deployment['connectivity']['http_timeout_s'] + result = remote.run_ssh(gpu['ssh'] or {'host': 'localhost'}, + f"curl --noproxy '*' -fsS --max-time {timeout} " + remote.quote(manager['control_url'] + '/healthz'), + timeout=timeout + 15) + if result.returncode != 0: + raise RuntimeError('GPU cannot reach manager control service') + payload = json.loads(result.stdout) + if payload != {'ok': True, 'service': 'harborrl-rollout-service-v1'}: + raise RuntimeError('GPU reached an unexpected control service') + return {'gpu_to_controller': True, 'control_url': manager['control_url']} + + +def check_gateway_shards(deployment): + """Preflight every remote shard before the GPU job can depend on it. + + Each remote process must answer as its declared shard name with the run + token; a missing or misnamed shard fails dispatch loudly instead of + failing attempts mid-training. + """ + manager = deployment['manager'] + specs = [spec for spec in manager.get('gateway_shards') or [] + if spec.get('mode', 'local') == 'remote'] + if not specs: + return None + from harborrl.rollout.remote_shard import RemoteShard, read_shard_token + from harborrl.rollout.runtime import GatewayShardSpec + + token = read_shard_token(manager['run_token_file']) + checks = [] + for spec in specs: + shard = RemoteShard( + GatewayShardSpec(spec['name'], urlsplit(spec['url']).port, spec['url'], + mode='remote', admin_url=spec['admin_url']), + token=token, trace_dir='unused') + shard.start() # Raises on unreachable or wrong shard identity. + checks.append({'name': spec['name'], 'admin_url': spec['admin_url'], + 'gateway_url': spec['url'], 'ok': True}) + return checks + + +def _ensure_token(token_file): + """Load the run token, creating a fresh one on first use.""" + path = Path(token_file) + if path.is_file(): + token = path.read_text().strip() + if token: + return token + import secrets + + token = secrets.token_urlsafe(32) + path.parent.mkdir(parents=True, exist_ok=True) + fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, TOKEN_MODE) + with os.fdopen(fd, "w") as stream: + stream.write(token + "\n") + return token + + +def _rollout_environment(plan, deployment, environment): + """Environment overlay that points the GPU job at the manager service.""" + manager = deployment["manager"] + workers = deployment["rollout"]["workers"] + overlay = dict(environment) + overlay.update({ + "HARBORRL_ROLLOUT_MODE": "service", + "HARBORRL_ROLLOUT_SERVICE_URL": manager["control_url"], + "HARBORRL_ROLLOUT_SERVICE_TOKEN_FILE": manager['run_token_file'], + "TRAIN_PYTHON": deployment['gpu_target']['train_python'], + "HARBORRL_NATIVE_GATEWAY_HOST": manager["bind"], + "HARBORRL_NATIVE_GATEWAY_PORT": str(manager["gateway_port"]), + "HARBORRL_NATIVE_GATEWAY_URL": manager["gateway_url"], + "HARBORRL_NATIVE_RUNNER_WORKERS": json.dumps( + [destination(worker["ssh"]) for worker in workers]), + "HARBORRL_NATIVE_RUNNER_PYTHON": (plan['config']['harbor']['python'] + if deployment['rollout']['provider'] == 'sandbox' + else workers[0]['runner_python']), + }) + if deployment['rollout']['provider'] == 'sandbox': + if 'HARBORRL_NATIVE_SANDBOX' not in overlay: + raise ValueError('sandbox deployment requires sandbox training config') + transport = plan['config']['harbor'].get('runner_transport', 'ssh') + if (transport == 'local' and workers) or (transport == 'ssh' and not workers): + raise ValueError('sandbox runner transport and deployment workers disagree') + overlay['HARBORRL_NATIVE_RUNNER_TRANSPORT'] = transport + shards = manager.get("gateway_shards") + if shards: + entries = [] + for spec in shards: + if spec.get("mode", "local") == "remote": + entries.append({"name": spec["name"], "mode": "remote", + "url": spec["url"], "admin_url": spec["admin_url"]}) + else: + entries.append({"name": spec["name"], "port": spec["port"]}) + overlay["HARBORRL_NATIVE_GATEWAY_SHARDS"] = json.dumps(entries) + if any(spec.get("mode") == "remote" for spec in shards): + # The remote processes read the same shared-storage token file. + overlay["HARBORRL_NATIVE_SHARD_TOKEN_FILE"] = manager["run_token_file"] + worker_shards = manager.get("worker_shards") + if worker_shards: + # The deployment speaks worker ids; the runtime routes by SSH + # destination (or the 'local' transport marker). + destinations = {worker["id"]: destination(worker["ssh"]) for worker in workers} + overlay["HARBORRL_NATIVE_GATEWAY_WORKER_SHARDS"] = json.dumps( + {destinations.get(key, key): shard + for key, shard in worker_shards.items()}, sort_keys=True) + routing = manager.get("gateway_routing") + if routing is not None: + # Declared even without a shard table: an explicit routing opts in or + # out of elastic growth; undeclared leaves the runtime's auto default. + overlay["HARBORRL_NATIVE_GATEWAY_ROUTING"] = routing + if manager.get("sandboxes_per_shard") is not None: + overlay["HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD"] = str( + manager["sandboxes_per_shard"]) + gateway_hosts = manager.get("gateway_hosts") + if gateway_hosts: + # Auto growth fills this pool first and only spawns locally when no + # candidate answers; remotely spawned shards authenticate on the same + # shared-storage run token as statically declared ones. + overlay["HARBORRL_NATIVE_GATEWAY_HOSTS"] = json.dumps(gateway_hosts) + overlay["HARBORRL_NATIVE_SHARD_TOKEN_FILE"] = manager["run_token_file"] + if manager.get("shard_python"): + # Growth hosts run the serving stack, which need not share the + # runner's sandbox-framework interpreter. + overlay["HARBORRL_NATIVE_SHARD_PYTHON"] = manager["shard_python"] + return overlay + + +def _runtime_factory(deployment, environment, *, concurrency_limit, group_size, tokenizer=None): + """Build the per-run rollout runtime when the GPU job announces serving.""" + + def factory(run_id, router_url): + from harborrl.rollout.runtime import RolloutRuntime, RolloutRuntimeConfig + + env = {**os.environ, **environment} + config = RolloutRuntimeConfig.from_environment( + env, endpoint=router_url, concurrency_limit=concurrency_limit, + group_size=group_size, + gateway_host=deployment["manager"]["bind"], + gateway_port=deployment["manager"]["gateway_port"], + gateway_url=deployment["manager"]["gateway_url"]) + resolved = tokenizer + if resolved is None: + from transformers import AutoTokenizer + + resolved = AutoTokenizer.from_pretrained( + env["HF_CKPT"], trust_remote_code=True) + return RolloutRuntime(config, resolved) + + return factory + + + + +def check_hub_readiness(deployment): + """Preflight check: verify hub manager is reachable and healthy.""" + import requests + + manager = deployment["manager"] + hub_url = manager.get("hub_url") + timeout = deployment["connectivity"]["http_timeout_s"] + + if not hub_url: + raise RuntimeError("hub mode requires manager.hub_url") + + token_file = manager.get("run_token_file") + hub_token = "" + if token_file: + try: + hub_token = open(token_file).read().strip() + except OSError as exc: + raise RuntimeError(f"hub mode cannot read manager.run_token_file {token_file}: {exc}") + + try: + response = requests.get( + f"{hub_url}/hub/status", + timeout=timeout, + headers={"X-Api-Key": hub_token}, + ) + response.raise_for_status() + status = response.json() + + # The hub server identifies itself with schema_version, not a + # service field; keep the check loose so both spellings pass. + identity = status.get("service") or status.get("schema_version") + if not str(identity or "").startswith("harborrl-hub"): + raise RuntimeError(f"Unexpected service at {hub_url}: {status}") + + return { + "hub_reachable": True, + "hub_url": hub_url, + "hub_status": status, + } + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Hub manager not reachable at {hub_url}: {e}") + + +def dispatch_hub(plan, deployment, *, checks): + """Hub mode: submit job to long-running manager, return immediately. + + In hub mode, the manager is a persistent daemon that accepts multiple + training runs. This function submits the run to the hub HTTP API + and returns control immediately rather than waiting for completion. + """ + import requests + + manager = deployment["manager"] + hub_url = manager.get("hub_url") + + if not hub_url: + raise RuntimeError("hub mode requires manager.hub_url to be configured") + + try: + token_file = manager.get("run_token_file") + hub_token = open(token_file).read().strip() if token_file else "" + response = requests.post( + f"{hub_url}/hub/runs", + json={"plan": plan, "deployment": deployment}, + timeout=deployment["connectivity"]["http_timeout_s"], + headers={"Content-Type": "application/json", "X-Api-Key": hub_token}, + ) + response.raise_for_status() + result = response.json() + + run_id = result.get("run_id") + status_url = f"{hub_url}/hub/runs/{run_id}" + + print(json.dumps({ + "mode": "hub", + "status": "submitted", + "run_id": run_id, + "status_url": status_url, + "hub_url": hub_url, + }), file=sys.stderr) + + return 0 # Success - job submitted + + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Failed to submit run to hub at {hub_url}: {e}") + + +def dispatch_service(plan, deployment, *, checks): + manager, gpu, storage = deployment["manager"], deployment["gpu_target"], deployment["storage"] + + # Hub mode: submit to long-running manager instead of direct execution + mode = manager.get("mode") + if mode == "hub": + return dispatch_hub(plan, deployment, checks=checks) + + # Standard service mode execution + config = plan["config"] + root = new_run_root(storage["run_root"]) + root.mkdir(parents=True, mode=0o700) + run_id = root.name + # Concurrent dispatches on one manager must not collide on the declared + # ports; patch every consumer's view before the environment is rendered. + deployment = _with_allocated_ports(deployment, run_id) + manager = deployment["manager"] + + catalog_path, prompt_path = _materialize(plan, root) + source = _source_audit(root) + environment = run_environment(plan, root, catalog_path, prompt_path) + token = _ensure_token(manager["run_token_file"]) + environment = _rollout_environment(plan, deployment, environment) + + sampling = config["sampling"] + concurrency = sampling["group_size"] * sampling["groups_per_batch"] + service = RolloutService( + run_root=root, + runtime_factory=_runtime_factory( + deployment, environment, + concurrency_limit=concurrency, group_size=sampling["group_size"]), + token_file=manager["run_token_file"]) + control = make_control_server(service, host=manager["bind"], + port=manager["control_port"]) + control_thread = None + target = None + exit_code = 1 + try: + service.register_run(run_id, { + "run_id": run_id, "deployment_source": plan["deployment"]["source"], + "source_commit": source["commit"], + "config_digest": plan["deployment"]["digest"], + "group_size": sampling["group_size"], + "workers": [worker["id"] for worker in deployment["rollout"]["workers"]], + }, token=token) + control_thread = threading.Thread( + target=control.serve_forever, name="harborrl-control-plane", daemon=True) + control_thread.start() + publish(root / 'control-readiness.json', check_gpu_control(deployment)) + shard_checks = check_gateway_shards(deployment) + if shard_checks: + publish(root / 'gateway-shards-readiness.json', shard_checks) + + publish(root / "launch.json", { + "plan": plan, "run_dir": str(root), "source": source, + "materialized_catalog": str(catalog_path), "prompt_data": str(prompt_path), + "control_url": manager["control_url"]}) + publish(root / "effective-config.json", + effective_config(plan, environment, dict(environment), source, checks)) + + target = SshTrainingTarget( + gpu["ssh"] or {"host": "localhost"}, source_dir=gpu["source_dir"], run_dir=root, + environment_file=gpu.get("environment_file"), + command_timeout=deployment["connectivity"]["command_timeout_s"] + 30, + ssh_timeout=deployment["connectivity"]["ssh_timeout_s"]) + # The launcher executes inside the GPU checkout; manager-side absolute + # paths are rewritten so shared-storage artifacts keep one spelling. + target.write_launcher( + ["bash", str(ROOT / "harborrl/platform/slime_train.sh")], environment, + path_rewrites=[(str(manager["source_dir"]), str(gpu["source_dir"]))]) + target.launch() + + def forward_signal(signum, frame): + if target is not None: + target.cancel() + + previous = {sig: signal.signal(sig, forward_signal) + for sig in (signal.SIGINT, signal.SIGTERM)} + try: + exit_code = target.wait() + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + except BaseException: + import traceback + + traceback.print_exc() + if target is not None: + # Best-effort teardown; the run record keeps the failure evidence. + with contextlib.suppress(Exception): + target.cancel() + raise + finally: + with contextlib.suppress(Exception): + service.drain_run(run_id) + service.close() + if control_thread is not None: + control.shutdown() + control_thread.join(5) + control.server_close() + _release_run_ports(run_id) + print(json.dumps({"run_dir": str(root), "exit_code": exit_code}), file=sys.stderr) + return exit_code diff --git a/harborrl/deployment/hub_client.py b/harborrl/deployment/hub_client.py new file mode 100644 index 000000000..295d812de --- /dev/null +++ b/harborrl/deployment/hub_client.py @@ -0,0 +1,444 @@ +"""Hub client pullers for GPU training nodes and rollout worker nodes. + +GPU Puller: runs on GPU nodes, pulls training tasks from hub manager. +Rollout Puller: runs on worker nodes, pulls rollout tasks from hub manager. + +Both implement exponential backoff retry. The contract below mirrors +harborrl/deployment/hub_server.py exactly: + + GET /pull/gpu_task header X-GPU-ID -> 200 {task_id, run_id, plan, deployment} | 204 + GET /pull/rollout_task header X-Worker-ID -> 200 {task_id, run_id, group_id, request} | 204 + POST /report/gpu_status {gpu_id, task_id, status: {success, result|error}} + POST /report/rollout_result {worker_id, task_id, result} + +Authentication is a single header: X-Api-Key (never combined with Authorization). +""" +from __future__ import annotations + +import json +import os +import signal +import subprocess +import sys +import threading +import time +from dataclasses import dataclass +from pathlib import Path +from urllib.parse import urljoin +from urllib.request import Request, urlopen +from urllib.error import HTTPError, URLError + +# Exponential backoff parameters +INITIAL_BACKOFF = 1.0 +MAX_BACKOFF = 60.0 +BACKOFF_MULTIPLIER = 2.0 +POLL_INTERVAL = 5.0 + +# Timeout settings +TASK_PULL_TIMEOUT = 30 +STATUS_REPORT_TIMEOUT = 15 + + +@dataclass +class PullerConfig: + """Common configuration for hub pullers.""" + hub_url: str + node_id: str + node_type: str # 'gpu' or 'rollout' + token: str + working_dir: str = None + log_dir: str = None + # rollout pullers need somewhere to send model traffic; the hub holds no + # gateway of its own (see hub_server docstring), so the worker brings it. + gateway_url: str = None + runner_python: str = None + source_dir: str = None + profile_path: str = None + task_path: str = None + gateway_host: str = None + + +class HubClientError(Exception): + """Base exception for hub client errors.""" + + +class TaskExecutionError(HubClientError): + """Raised when task execution fails.""" + + +def _http_request(url, data=None, *, token=None, timeout=30, method=None, headers=None): + """Make an HTTP request carrying exactly one authentication header.""" + hdrs = {'Content-Type': 'application/json'} + if headers: + hdrs.update(headers) + if token: + # Server rejects requests presenting both x-api-key and Authorization. + hdrs['X-Api-Key'] = token + + body = json.dumps(data).encode('utf-8') if data is not None else None + req = Request(url, data=body, headers=hdrs, method=method) + + try: + with urlopen(req, timeout=timeout) as response: + content = response.read().decode('utf-8') + return json.loads(content) if content else {} + except HTTPError as exc: + if exc.code == 204: + return {} + error_body = exc.read().decode('utf-8', errors='replace') + try: + error_data = json.loads(error_body) + message = error_data.get('error', error_body) + except json.JSONDecodeError: + message = error_body + raise HubClientError(f"HTTP {exc.code}: {message}") from exc + except URLError as exc: + raise HubClientError(f"Network error: {exc.reason}") from exc + + +class BasePuller: + """Base class for hub client pullers with common retry logic.""" + + def __init__(self, config: PullerConfig): + self.config = config + self.running = False + self.current_task = None + self.task_lock = threading.RLock() + self._backoff = INITIAL_BACKOFF + + def start(self): + """Run the task polling loop until stopped.""" + self.running = True + label = self.config.node_type.upper() + print(f"[{label}] Puller started: {self.config.node_id}") + print(f"[{label}] Hub URL: {self.config.hub_url}") + + try: + self._main_loop() + except KeyboardInterrupt: + print(f"\n[{label}] Shutting down...") + finally: + self.running = False + + def _main_loop(self): + """Main task polling loop with exponential backoff.""" + label = self.config.node_type.upper() + while self.running: + try: + task = self._pull_task() + if task: + self._backoff = INITIAL_BACKOFF # Reset backoff on success + self._execute_task(task) + else: + # No task available, back off + time.sleep(self._backoff) + self._backoff = min(self._backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF) + except KeyboardInterrupt: + raise + except Exception as exc: # noqa: BLE001 - loop must survive task errors + print(f"[{label}] Error in main loop: {exc}", flush=True) + time.sleep(self._backoff) + self._backoff = min(self._backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF) + + def _pull_task(self): + """Pull next available task from the hub. Implemented per node type.""" + raise NotImplementedError + + def _execute_task(self, task): + """Execute a task. Implemented per node type.""" + raise NotImplementedError + + def _stream_process(self, cmd, *, cwd, env, task_id, label): + """Run cmd, stream its output, and return the exit code.""" + log_file = None + if self.config.log_dir: + log_path = Path(self.config.log_dir) / f'{task_id}.log' + log_path.parent.mkdir(parents=True, exist_ok=True) + log_file = log_path.open('w') + + process = subprocess.Popen( + cmd, cwd=cwd, env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + text=True, bufsize=1) + try: + for line in process.stdout: + print(f"[{label}:{task_id}] {line}", end='', flush=True) + if log_file: + log_file.write(line) + log_file.flush() + return process.wait() + finally: + if log_file: + log_file.close() + + +class GpuPuller(BasePuller): + """GPU training node puller for executing training jobs.""" + + def __init__(self, config: PullerConfig): + super().__init__(config) + if config.node_type != 'gpu': + raise ValueError("GpuPuller requires node_type='gpu'") + + def _pull_task(self): + """GET /pull/gpu_task with the X-GPU-ID header.""" + url = urljoin(self.config.hub_url, '/pull/gpu_task') + try: + response = _http_request( + url, token=self.config.token, timeout=TASK_PULL_TIMEOUT, + method='GET', headers={'X-GPU-ID': self.config.node_id}) + except HubClientError as exc: + if '404' in str(exc) or '503' in str(exc): + return None + raise + if response and response.get('task_id'): + return response + return None + + def _report(self, task_id, *, success, result=None, error=None): + """POST /report/gpu_status.""" + status = {'success': bool(success)} + if success: + status['result'] = result + else: + status['error'] = error or 'unknown error' + url = urljoin(self.config.hub_url, '/report/gpu_status') + payload = { + 'gpu_id': self.config.node_id, + 'task_id': task_id, + 'status': status, + } + try: + _http_request(url, payload, token=self.config.token, + timeout=STATUS_REPORT_TIMEOUT, method='POST') + except HubClientError as exc: + print(f"[GPU] Status report failed: {exc}", flush=True) + + def _execute_task(self, task): + """Execute one training task on this GPU node.""" + task_id = task['task_id'] + run_id = task.get('run_id') + print(f"[GPU] Starting training task: {task_id} (run: {run_id})", flush=True) + + with self.task_lock: + self.current_task = task_id + + try: + deployment = task.get('deployment') or {} + plan = task.get('plan') or {} + gpu_target = deployment.get('gpu_target') or {} + + working_dir = Path(self.config.working_dir or '/tmp') / f'task-{task_id}' + working_dir.mkdir(parents=True, exist_ok=True) + deployment_file = working_dir / 'deployment.json' + plan_file = working_dir / 'plan.json' + deployment_file.write_text(json.dumps(deployment, indent=2)) + plan_file.write_text(json.dumps(plan, indent=2)) + + env = os.environ.copy() + if task.get('environment'): + env.update(task['environment']) + env['HARBORRL_DEPLOYMENT_FILE'] = str(deployment_file) + env['HARBORRL_PLAN_FILE'] = str(plan_file) + + train_python = gpu_target.get('train_python') or sys.executable + source_dir = gpu_target.get('source_dir') or self.config.source_dir or os.getcwd() + train_module = task.get('train_module', 'harborrl.platform.native_train') + + cmd = [train_python, '-m', train_module, str(plan_file)] + print(f"[GPU] Executing: {' '.join(cmd)}", flush=True) + print(f"[GPU] Working directory: {source_dir}", flush=True) + + returncode = self._stream_process( + cmd, cwd=source_dir, env=env, task_id=task_id, label='GPU') + + if returncode == 0: + print(f"[GPU] Task completed successfully: {task_id}", flush=True) + self._report(task_id, success=True, result={'exit_code': 0}) + else: + print(f"[GPU] Task failed with exit code {returncode}: {task_id}", flush=True) + self._report(task_id, success=False, + error=f'Training exited with code {returncode}') + except Exception as exc: # noqa: BLE001 - report every failure to the hub + print(f"[GPU] Task execution error: {exc}", flush=True) + self._report(task_id, success=False, error=str(exc)) + raise TaskExecutionError(f"Training task failed: {exc}") from exc + finally: + with self.task_lock: + self.current_task = None + + +class RolloutPuller(BasePuller): + """Rollout worker node puller for executing Harbor trials. + + The hub server carries no gateway and no RolloutRuntime, so execution + happens here: the worker spawns the existing runner module and speaks its + stdin/stdout JSON protocol, using the gateway this node was pointed at. + """ + + def __init__(self, config: PullerConfig): + super().__init__(config) + if config.node_type != 'rollout': + raise ValueError("RolloutPuller requires node_type='rollout'") + + def _pull_task(self): + """GET /pull/rollout_task with the X-Worker-ID header.""" + url = urljoin(self.config.hub_url, '/pull/rollout_task') + try: + response = _http_request( + url, token=self.config.token, timeout=TASK_PULL_TIMEOUT, + method='GET', headers={'X-Worker-ID': self.config.node_id}) + except HubClientError as exc: + if '404' in str(exc) or '503' in str(exc): + return None + raise + if response and response.get('task_id'): + return response + return None + + def _report(self, task_id, *, success, result=None, error=None): + """POST /report/rollout_result.""" + payload_result = {'success': bool(success)} + if success: + payload_result['result'] = result + else: + payload_result['error'] = error or 'unknown error' + url = urljoin(self.config.hub_url, '/report/rollout_result') + payload = { + 'worker_id': self.config.node_id, + 'task_id': task_id, + 'result': payload_result, + } + try: + _http_request(url, payload, token=self.config.token, + timeout=STATUS_REPORT_TIMEOUT, method='POST') + except HubClientError as exc: + print(f"[ROLLOUT] Result report failed: {exc}", flush=True) + + def _execute_task(self, task): + """Execute one rollout task by hosting a gateway and driving the runner.""" + from harborrl.deployment.hub_rollout_executor import RolloutExecutor + + task_id = task['task_id'] + group_request = dict(task.get('request') or {}) + group_request.setdefault('task_id', task_id) + group_request.setdefault('run_id', task.get('run_id')) + group_request.setdefault('group_id', task.get('group_id')) + print(f"[ROLLOUT] Starting rollout task: {task_id}", flush=True) + + with self.task_lock: + self.current_task = task_id + + try: + task_root = Path(self.config.working_dir or '/tmp') / f'task-{task_id}' + task_root.mkdir(parents=True, exist_ok=True) + + profile_path = self.config.profile_path + if not profile_path: + raise TaskExecutionError('rollout puller requires --profile') + profile = json.loads(Path(profile_path).read_text()) + + task_path = group_request.get('task_path') or self.config.task_path + if not task_path: + raise TaskExecutionError('rollout task carries no task_path') + + executor = RolloutExecutor( + run_root=task_root, + task_path=task_path, + profile=profile, + runner_python=self.config.runner_python or sys.executable, + source_dir=self.config.source_dir or os.getcwd(), + gateway_host=self.config.gateway_host, + log=lambda m: print(m, flush=True)) + result = executor.execute(group_request) + + self._report(task_id, success=bool(result.get('success')), + result=result, + error=None if result.get('success') else 'runner reported a failed terminal') + except Exception as exc: # noqa: BLE001 - report every failure to the hub + print(f"[ROLLOUT] Task execution error: {exc}", flush=True) + self._report(task_id, success=False, error=str(exc)) + raise TaskExecutionError(f"Rollout task failed: {exc}") from exc + finally: + with self.task_lock: + self.current_task = None + + +def main(): + """CLI entry point for pullers.""" + import argparse + + parser = argparse.ArgumentParser(description='HarborRL Hub Client Puller') + parser.add_argument('--hub-url', required=True, + help='Hub manager URL (e.g., http://manager:49300)') + parser.add_argument('--node-id', required=True, + help='Unique node identifier') + parser.add_argument('--node-type', required=True, choices=['gpu', 'rollout'], + help='Node type: gpu or rollout') + parser.add_argument('--token-file', required=True, + help='Path to hub authentication token file') + parser.add_argument('--working-dir', + help='Working directory for task execution') + parser.add_argument('--log-dir', + help='Directory for task execution logs') + parser.add_argument('--gateway-url', + help='Gateway origin for rollout tasks (required for rollout node)') + parser.add_argument('--runner-python', + help='Interpreter that runs harborrl.rollout.harbor_job.runner') + parser.add_argument('--source-dir', + help='Checkout that contains the harborrl package') + parser.add_argument('--profile', + help='Claude-code profile JSON (required for rollout node)') + parser.add_argument('--task-path', + help='Default Harbor task path used when the task omits one') + parser.add_argument('--gateway-host', + help='Address the sandbox uses to reach this worker gateway') + + args = parser.parse_args() + + token_path = Path(args.token_file) + if not token_path.is_file(): + print(f"Error: Token file not found: {args.token_file}", file=sys.stderr) + sys.exit(1) + + token = token_path.read_text().strip() + if not token: + print(f"Error: Token file is empty: {args.token_file}", file=sys.stderr) + sys.exit(1) + + config = PullerConfig( + hub_url=args.hub_url, + node_id=args.node_id, + node_type=args.node_type, + token=token, + working_dir=args.working_dir, + log_dir=args.log_dir, + gateway_url=args.gateway_url, + runner_python=args.runner_python, + source_dir=args.source_dir, + profile_path=args.profile, + task_path=args.task_path, + gateway_host=args.gateway_host, + ) + + if args.node_type == 'gpu': + puller = GpuPuller(config) + else: + puller = RolloutPuller(config) + + def signal_handler(signum, frame): + print(f"\nReceived signal {signum}, shutting down...") + puller.running = False + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + try: + puller.start() + except Exception as exc: # noqa: BLE001 - surface fatal startup errors + print(f"Fatal error: {exc}", file=sys.stderr) + sys.exit(1) + + +if __name__ == '__main__': + main() diff --git a/harborrl/deployment/hub_rollout_executor.py b/harborrl/deployment/hub_rollout_executor.py new file mode 100644 index 000000000..abde1855d --- /dev/null +++ b/harborrl/deployment/hub_rollout_executor.py @@ -0,0 +1,222 @@ +"""Rollout task executor for hub-mode workers. + +The hub server deliberately holds no Gateway and no RolloutRuntime, so a +pulled rollout task is executed here on the worker: + + 1. start a local Gateway + TraceRegistry for this task, + 2. register the attempt to obtain its credential, + 3. hand the runner a request whose gateway_url points at that gateway, + 4. speak the runner's stdin/stdout protocol (prepared -> run -> terminal). + +The gateway must be reachable from inside the sandbox, so the advertised +origin is composed from --gateway-host (defaults to the local IP) and the +ephemeral port the server actually bound. +""" +from __future__ import annotations + +import asyncio +import json +import secrets +import socket +import threading +from pathlib import Path + +from harborrl.gateway.server import Gateway, make_server +from harborrl.gateway.trace import TraceRegistry +from harborrl.trajectories.native import Identity, RewardProfile, digest, publish + + +def local_ip(): + """Best-effort address the sandbox can dial back to.""" + try: + with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as probe: + probe.connect(("10.255.255.255", 1)) + return probe.getsockname()[0] + except OSError: + return "127.0.0.1" + + +def sampling_digest_for(max_output_tokens, stop=None): + """Digest the same sampling structure the gateway's convert() produces. + + The gateway compares this digest against the converted request, so the + recipe here must match messages.convert() exactly: temperature=1, + top_p=1, top_k=-1 are enforced for the initial profile, max_new_tokens + is the profile output budget, and stop is the caller's sequence list + (empty when unspecified). + """ + return digest({ + # floats on purpose: messages.convert() defaults temperature/top_p to + # 1.0, and digest() hashes JSON text, so 1 and 1.0 hash differently. + "temperature": 1.0, + "top_p": 1.0, + "top_k": -1, + "max_new_tokens": max_output_tokens, + "stop": list(stop or []), + }) + + +class RolloutExecutor: + """Run one rollout group request against a locally hosted gateway.""" + + def __init__(self, *, run_root, task_path, profile, runner_python, source_dir, + gateway_host=None, tracing=True, log=print): + self.run_root = Path(run_root) + self.task_path = Path(task_path) + self.profile = profile + self.runner_python = runner_python + self.source_dir = str(source_dir) + self.gateway_host = gateway_host or local_ip() + self.tracing = tracing + self.log = log + + def execute(self, group_request): + """Execute one group request; returns a result dict for the hub.""" + task_id = group_request.get('task_id') or 'rollout_' + secrets.token_hex(8) + run_id = group_request.get('run_id', 'adhoc') + group_id = group_request.get('group_id', task_id) + slot_id = str(group_request.get('slot_id', '0')) + + attempt_id = f'{task_id}-{slot_id}' + attempt_root = self.run_root / attempt_id + attempt_root.mkdir(parents=True, exist_ok=True) + + trace_root = attempt_root / 'traces' + registry = TraceRegistry(trace_root) + + model_tag = group_request.get('model_tag', 'qwen3-8B') + from harborrl.math_registry import MATH_REGISTRY + model = MATH_REGISTRY.model_for_preset(model_tag) + + backend = group_request.get('backend') + if backend is None and group_request.get('synthetic'): + # Transport-only verification: generate without a serving pool. + from .hub_synthetic_backend import SyntheticServingBackend + backend = SyntheticServingBackend( + context=group_request.get('max_context', 32768)) + if backend is None: + # A live rollout needs a serving pool; the hub does not carry one, + # so the caller must supply an endpoint for real generation. + endpoint = group_request.get('serving_endpoint') + if not endpoint: + raise ValueError( + 'rollout execution requires serving_endpoint (or an injected backend)') + # The tokenizer cannot cross JSON, so load it locally from the + # checkpoint and verify it still matches the run's lock digests. + from transformers import AutoTokenizer + ckpt = group_request.get('hf_ckpt') or self.hf_ckpt + if not ckpt: + raise ValueError('live rollout requires hf_ckpt to load the tokenizer') + tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True) + from harborrl.rollout.runtime import tokenizer_fingerprint + tok_digest, tpl_digest = tokenizer_fingerprint(tokenizer) + expect_tok = group_request.get('tokenizer_digest') + expect_tpl = group_request.get('template_digest') + if expect_tok and tok_digest != expect_tok: + raise ValueError(f'tokenizer digest mismatch: {tok_digest} != {expect_tok}') + if expect_tpl and tpl_digest != expect_tpl: + raise ValueError(f'template digest mismatch: {tpl_digest} != {expect_tpl}') + from harborrl.gateway.sglang_backend import SGLangBackend + backend = SGLangBackend( + endpoint, tokenizer, model.adapter, + tokenizer_digest=tok_digest, template_digest=tpl_digest, + max_context=group_request.get('max_context', 32768)) + + gateway = Gateway(registry, backend, model=self.profile.get('model', 'policy'), + max_output_tokens=group_request.get('max_output_tokens', 8192)) + server = make_server(gateway, host='0.0.0.0', port=0) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True, + name=f'hub-gateway-{task_id}') + thread.start() + + advertised = f'http://{self.gateway_host}:{port}' + self.log(f'[ROLLOUT] gateway up at {advertised} (attempt {attempt_id})') + + try: + identity = Identity( + run_id=run_id, + batch_id=group_request.get('batch_id', 'batch'), + group_id=group_id, + slot_id=slot_id, + attempt_id=attempt_id, + trajectory_id=group_request.get('trajectory_id', f'traj-{slot_id}'), + task_digest=group_request.get('task_digest', ''), + policy_version=group_request.get('policy_version', '1'), + harness_digest=digest(self.profile), + reward_digest=RewardProfile().digest, + sampling_digest=group_request.get('sampling_digest') + or sampling_digest_for(group_request.get('max_output_tokens', 8192), + group_request.get('stop')), + ) + # register() mints the attempt credential; the runner must hand + # exactly this token to the sandbox or the gateway rejects it. + attempt_credential = registry.register(identity, task_id) + + request = { + 'identity': identity.to_dict(), + 'task_path': str(self.task_path), + 'profile': self.profile, + 'gateway_url': advertised, + 'credential': attempt_credential, + 'output': str(attempt_root / 'runner'), + 'max_output_tokens': group_request.get('max_output_tokens', 8192), + } + + terminal = asyncio.run( + self._drive_runner(request, attempt_root)) + publish(attempt_root / 'terminal.json', terminal) + return {'success': terminal.get('type') == 'terminal', 'terminal': terminal, + 'attempt_root': str(attempt_root)} + finally: + server.shutdown() + server.server_close() + thread.join(5) + + async def _drive_runner(self, request, attempt_root): + """Speak the runner stdin/stdout protocol to completion.""" + env = {**__import__('os').environ} + env['PYTHONPATH'] = self.source_dir + process = await asyncio.create_subprocess_exec( + self.runner_python, '-m', 'harborrl.rollout.harbor_job.runner', + cwd=self.source_dir, env=env, + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE) + + process.stdin.write(json.dumps(request).encode() + b'\n') + await process.stdin.drain() + + terminal = None + try: + while True: + line = await process.stdout.readline() + if not line: + break + try: + event = json.loads(line) + except ValueError: + self.log(f'[ROLLOUT] non-JSON line: {line[:200]!r}') + continue + kind = event.get('type') + self.log(f'[ROLLOUT] runner event: {kind}') + if kind == 'prepared': + process.stdin.write(b'{"command":"run"}\n') + await process.stdin.drain() + elif kind in ('terminal', 'cancelled'): + terminal = event + break + finally: + with __import__('contextlib').suppress(Exception): + await asyncio.wait_for(process.wait(), timeout=60) + stderr = b'' + if process.stderr is not None: + with __import__('contextlib').suppress(Exception): + stderr = await process.stderr.read() + if stderr: + publish(attempt_root / 'runner.stderr.txt', + stderr.decode('utf-8', 'replace')[-8000:]) + + if terminal is None: + raise RuntimeError('runner exited before emitting a terminal record') + return terminal diff --git a/harborrl/deployment/hub_server.py b/harborrl/deployment/hub_server.py new file mode 100644 index 000000000..4efd41ee1 --- /dev/null +++ b/harborrl/deployment/hub_server.py @@ -0,0 +1,672 @@ +"""Hub Server: Centralized manager for multi-run GPU and rollout task scheduling. + +This implements the hub mode control plane, where a single long-running manager +service accepts and queues multiple training runs, managing shared GPU resources +and rollout workers. Unlike service mode (1 manager → 1 run), hub mode supports +multi-tenancy (1 manager → N runs). + +Architecture: +- Pull-based task distribution for GPU training and rollout workers +- Shared GPU pool with FIFO/priority scheduling +- Run queue with configurable concurrency limits +- HTTP API for task submission, status queries, and worker pulls +""" +from __future__ import annotations + +import json +import secrets +import threading +import time +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Dict, List, Optional +from urllib.parse import urlsplit + +HUB_SCHEMA = "harborrl-hub-service-v1" + + +class HubError(Exception): + """Maps to HTTP status codes with machine-readable error codes.""" + + def __init__(self, status: int, code: str, message: str): + super().__init__(message) + self.status = status + self.code = code + + +@dataclass +class GpuResource: + """Represents a single GPU node in the resource pool.""" + id: str + ssh: dict + source_dir: str + train_python: str + expected_gpus: int + status: str = "idle" # idle | busy | reserved + current_run: Optional[str] = None + current_task: Optional[str] = None + assigned_at: Optional[float] = None + + +@dataclass +class RolloutTask: + """A rollout task ready to be pulled by workers.""" + task_id: str + run_id: str + group_id: str + request: dict + status: str = "queued" # queued | assigned | running | completed | failed + worker_id: Optional[str] = None + created_at: float = field(default_factory=time.time) + assigned_at: Optional[float] = None + completed_at: Optional[float] = None + result: Optional[dict] = None + error: Optional[str] = None + + +@dataclass +class GpuTask: + """A GPU training task ready to be pulled by GPU workers.""" + task_id: str + run_id: str + plan: dict + deployment: dict + status: str = "queued" # queued | assigned | running | completed | failed + gpu_id: Optional[str] = None + created_at: float = field(default_factory=time.time) + assigned_at: Optional[float] = None + completed_at: Optional[float] = None + result: Optional[dict] = None + error: Optional[str] = None + + +@dataclass +class RunInfo: + """Tracking information for a submitted training run.""" + run_id: str + plan: dict + deployment: dict + status: str = "pending" # pending | scheduled | running | completed | failed + gpu_task_id: Optional[str] = None + rollout_tasks: List[str] = field(default_factory=list) + created_at: float = field(default_factory=time.time) + scheduled_at: Optional[float] = None + completed_at: Optional[float] = None + error: Optional[str] = None + + +class HubManager: + """Core hub manager coordinating GPU pool and task queues.""" + + def __init__(self, *, hub_config: dict, run_root: str, max_concurrent_runs: int = 4, + run_queue_size: int = 20, token_file: Optional[str] = None): + self.hub_config = hub_config + self.run_root = Path(run_root).resolve() + self.max_concurrent_runs = max_concurrent_runs + self.run_queue_size = run_queue_size + + # Authentication + self.hub_token = None + if token_file is not None: + token = Path(token_file).read_text().strip() + if not token: + raise HubError(500, "invalid_token", "hub token file is empty") + self.hub_token = token + + # State tracking + self.runs: Dict[str, RunInfo] = {} + self.gpu_pool: Dict[str, GpuResource] = {} + self.gpu_tasks: Dict[str, GpuTask] = {} + self.rollout_tasks: Dict[str, RolloutTask] = {} + self.lock = threading.RLock() + + # Initialize GPU pool from config + if "gpu_pool" in hub_config: + for gpu_spec in hub_config["gpu_pool"]: + gpu = GpuResource( + id=gpu_spec["id"], + ssh=gpu_spec.get("ssh", {}), + source_dir=gpu_spec["source_dir"], + train_python=gpu_spec["train_python"], + expected_gpus=gpu_spec.get("expected_gpus", 1) + ) + self.gpu_pool[gpu.id] = gpu + + # Start scheduler thread + self.scheduler_running = True + self.scheduler_thread = threading.Thread( + target=self._scheduler_loop, name="hub-scheduler", daemon=True + ) + self.scheduler_thread.start() + + def close(self): + """Shutdown the hub manager.""" + self.scheduler_running = False + self.scheduler_thread.join(timeout=5) + + def _authenticate(self, token: str): + """Validate hub access token.""" + if self.hub_token is not None and token != self.hub_token: + raise HubError(401, "unauthorized", "invalid hub token") + + def _scheduler_loop(self): + """Background scheduler for GPU resource allocation.""" + while self.scheduler_running: + try: + self._schedule_pending_runs() + time.sleep(1.0) + except Exception as e: + # Log but don't crash the scheduler + print(f"Scheduler error: {e}") + + def _schedule_pending_runs(self): + """Allocate GPU resources to pending runs if available.""" + with self.lock: + # Count currently running + running = sum(1 for r in self.runs.values() if r.status == "running") + if running >= self.max_concurrent_runs: + return + + # Find pending runs (FIFO order) + pending = [ + (r.created_at, run_id, r) + for run_id, r in self.runs.items() + if r.status == "pending" + ] + pending.sort() + + # Try to schedule each pending run + for _, run_id, run_info in pending: + if running >= self.max_concurrent_runs: + break + + # Find idle GPU + idle_gpu = next( + (g for g in self.gpu_pool.values() if g.status == "idle"), + None + ) + + if idle_gpu is None: + continue # No GPU available + + # Create GPU task + task_id = f"gpu_{secrets.token_hex(8)}" + gpu_task = GpuTask( + task_id=task_id, + run_id=run_id, + plan=run_info.plan, + deployment=run_info.deployment + ) + self.gpu_tasks[task_id] = gpu_task + + # Reserve GPU and update run + idle_gpu.status = "reserved" + idle_gpu.current_run = run_id + idle_gpu.current_task = task_id + run_info.status = "scheduled" + run_info.gpu_task_id = task_id + run_info.scheduled_at = time.time() + running += 1 + + def submit_run(self, plan: dict, deployment: dict, *, token: str) -> dict: + """Submit a new training run to the hub.""" + self._authenticate(token) + + with self.lock: + # Check queue capacity + if len(self.runs) >= self.run_queue_size: + raise HubError(503, "queue_full", "hub run queue is at capacity") + + # Generate run ID + run_id = f"run_{secrets.token_hex(12)}" + + # Create run info + run_info = RunInfo( + run_id=run_id, + plan=plan, + deployment=deployment + ) + self.runs[run_id] = run_info + + return { + "schema_version": HUB_SCHEMA, + "run_id": run_id, + "status": "submitted", + "queue_position": sum( + 1 for r in self.runs.values() if r.status == "pending" + ) + } + + def query_run(self, run_id: str, *, token: str) -> dict: + """Query the status of a submitted run.""" + self._authenticate(token) + + with self.lock: + run_info = self.runs.get(run_id) + if run_info is None: + raise HubError(404, "not_found", f"unknown run: {run_id}") + + result = { + "schema_version": HUB_SCHEMA, + "run_id": run_id, + "status": run_info.status, + "created_at": run_info.created_at + } + + if run_info.scheduled_at: + result["scheduled_at"] = run_info.scheduled_at + if run_info.completed_at: + result["completed_at"] = run_info.completed_at + if run_info.error: + result["error"] = run_info.error + if run_info.gpu_task_id: + result["gpu_task_id"] = run_info.gpu_task_id + + return result + + def cancel_run(self, run_id: str, *, token: str) -> dict: + """Cancel a pending or running run.""" + self._authenticate(token) + + with self.lock: + run_info = self.runs.get(run_id) + if run_info is None: + raise HubError(404, "not_found", f"unknown run: {run_id}") + + if run_info.status in ("completed", "failed"): + raise HubError(400, "invalid_state", "cannot cancel completed run") + + # Release GPU if allocated + if run_info.gpu_task_id: + gpu_task = self.gpu_tasks.get(run_info.gpu_task_id) + if gpu_task and gpu_task.gpu_id: + gpu = self.gpu_pool.get(gpu_task.gpu_id) + if gpu: + gpu.status = "idle" + gpu.current_run = None + gpu.current_task = None + gpu_task.status = "failed" + gpu_task.error = "cancelled" + + run_info.status = "failed" + run_info.error = "cancelled by user" + run_info.completed_at = time.time() + + return {"run_id": run_id, "status": "cancelled"} + + def hub_status(self, *, token: str) -> dict: + """Query overall hub status and resource utilization.""" + self._authenticate(token) + + with self.lock: + gpu_status = { + "total": len(self.gpu_pool), + "idle": sum(1 for g in self.gpu_pool.values() if g.status == "idle"), + "reserved": sum(1 for g in self.gpu_pool.values() if g.status == "reserved"), + "busy": sum(1 for g in self.gpu_pool.values() if g.status == "busy") + } + + run_status = { + "total": len(self.runs), + "pending": sum(1 for r in self.runs.values() if r.status == "pending"), + "scheduled": sum(1 for r in self.runs.values() if r.status == "scheduled"), + "running": sum(1 for r in self.runs.values() if r.status == "running"), + "completed": sum(1 for r in self.runs.values() if r.status == "completed"), + "failed": sum(1 for r in self.runs.values() if r.status == "failed") + } + + return { + "schema_version": HUB_SCHEMA, + "gpu_pool": gpu_status, + "runs": run_status, + "max_concurrent_runs": self.max_concurrent_runs, + "queue_capacity": self.run_queue_size + } + + def pull_gpu_task(self, gpu_id: str) -> Optional[dict]: + """GPU worker pulls next available task.""" + with self.lock: + # Find reserved task for this GPU + for task_id, task in self.gpu_tasks.items(): + if task.status == "queued": + gpu = self.gpu_pool.get(gpu_id) + if gpu and gpu.current_task == task_id: + # Assign task + task.status = "assigned" + task.gpu_id = gpu_id + task.assigned_at = time.time() + gpu.status = "busy" + + # Update run status + run_info = self.runs.get(task.run_id) + if run_info: + run_info.status = "running" + + return { + "task_id": task_id, + "run_id": task.run_id, + "plan": task.plan, + "deployment": task.deployment + } + + return None + + def report_gpu_status(self, gpu_id: str, task_id: str, status: dict) -> dict: + """GPU worker reports task completion or failure.""" + with self.lock: + task = self.gpu_tasks.get(task_id) + if task is None or task.gpu_id != gpu_id: + raise HubError(404, "not_found", "unknown or mismatched task") + + # Update task + if status.get("success"): + task.status = "completed" + task.result = status.get("result") + else: + task.status = "failed" + task.error = status.get("error", "unknown error") + task.completed_at = time.time() + + # Release GPU + gpu = self.gpu_pool.get(gpu_id) + if gpu: + gpu.status = "idle" + gpu.current_run = None + gpu.current_task = None + gpu.assigned_at = None + + # Update run + run_info = self.runs.get(task.run_id) + if run_info: + if task.status == "completed": + run_info.status = "completed" + else: + run_info.status = "failed" + run_info.error = task.error + run_info.completed_at = time.time() + + return {"task_id": task_id, "acknowledged": True} + + def submit_rollout_task(self, run_id: str, group_request: dict) -> dict: + """Submit a rollout group task for worker processing.""" + with self.lock: + task_id = f"rollout_{secrets.token_hex(8)}" + task = RolloutTask( + task_id=task_id, + run_id=run_id, + group_id=group_request.get("group_id", task_id), + request=group_request + ) + self.rollout_tasks[task_id] = task + + # Link to run + run_info = self.runs.get(run_id) + if run_info: + run_info.rollout_tasks.append(task_id) + + return {"task_id": task_id, "status": "queued"} + + def pull_rollout_task(self, worker_id: str) -> Optional[dict]: + """Rollout worker pulls next available task.""" + with self.lock: + # Find oldest queued task + queued = [ + (task.created_at, task_id, task) + for task_id, task in self.rollout_tasks.items() + if task.status == "queued" + ] + + if not queued: + return None + + queued.sort() + _, task_id, task = queued[0] + + # Assign to worker + task.status = "assigned" + task.worker_id = worker_id + task.assigned_at = time.time() + + return { + "task_id": task_id, + "run_id": task.run_id, + "group_id": task.group_id, + "request": task.request + } + + def report_rollout_result(self, worker_id: str, task_id: str, result: dict) -> dict: + """Rollout worker reports task completion.""" + with self.lock: + task = self.rollout_tasks.get(task_id) + if task is None or task.worker_id != worker_id: + raise HubError(404, "not_found", "unknown or mismatched task") + + # Update task + if result.get("success"): + task.status = "completed" + task.result = result + else: + task.status = "failed" + task.error = result.get("error", "unknown error") + task.completed_at = time.time() + + return {"task_id": task_id, "acknowledged": True} + + +def make_hub_server(hub_manager: HubManager, host: str = "0.0.0.0", port: int = 0, + *, max_body_bytes: int = 16 * 1024 * 1024): + """Create HTTP server for hub control plane.""" + + class HubHandler(BaseHTTPRequestHandler): + def setup(self): + super().setup() + self.connection.settimeout(120) + + def log_message(self, fmt, *args): + # Keep one line per request: without it the hub log carried only + # the startup banner and scheduling was untraceable. + print("HUB %s %s %s" % (self.command, self.path, fmt % args if args else ""), + flush=True) + + def _parse_token(self): + """Extract authentication token from headers.""" + # Accept either spelling; the clients only ever send one, but a + # proxy that adds Authorization must not turn a good request into + # a 401. X-Api-Key wins when both appear. + token = self.headers.get("x-api-key", "") or self.headers.get("X-Api-Key", "") + if not token: + auth = self.headers.get("Authorization", "") + if auth.startswith("Bearer "): + token = auth[7:] + return token + + def _read_json_body(self): + """Read and parse JSON request body.""" + if self.headers.get("Transfer-Encoding"): + raise HubError(400, "invalid_request", "chunked bodies not supported") + size = int(self.headers.get("Content-Length", "0")) + if size <= 0 or size > max_body_bytes: + raise HubError(400, "invalid_request", "invalid body size") + return json.loads(self.rfile.read(size)) + + def reply(self, status: int, payload: dict): + """Send JSON response.""" + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + try: + path = urlsplit(self.path).path + token = self._parse_token() + + if path == "/hub/status": + result = hub_manager.hub_status(token=token) + return self.reply(200, result) + + elif path.startswith("/hub/runs/"): + run_id = path.split("/")[-1] + result = hub_manager.query_run(run_id, token=token) + return self.reply(200, result) + + elif path == "/pull/gpu_task": + gpu_id = self.headers.get("X-GPU-ID", "") + task = hub_manager.pull_gpu_task(gpu_id) + if task is None: + return self.reply(204, {}) # No content + return self.reply(200, task) + + elif path == "/pull/rollout_task": + worker_id = self.headers.get("X-Worker-ID", "") + task = hub_manager.pull_rollout_task(worker_id) + if task is None: + return self.reply(204, {}) + return self.reply(200, task) + + else: + return self.reply(404, { + "error": {"type": "not_found", "message": "endpoint not found"} + }) + + except HubError as e: + return self.reply(e.status, { + "error": {"type": e.code, "message": str(e)} + }) + except Exception as e: + return self.reply(500, { + "error": {"type": "internal_error", "message": str(e)} + }) + + def do_POST(self): + try: + path = urlsplit(self.path).path + token = self._parse_token() + + if path == "/hub/runs": + body = self._read_json_body() + result = hub_manager.submit_run( + body.get("plan", {}), + body.get("deployment", {}), + token=token + ) + return self.reply(201, result) + + elif path == "/report/gpu_status": + body = self._read_json_body() + result = hub_manager.report_gpu_status( + body["gpu_id"], + body["task_id"], + body["status"] + ) + return self.reply(200, result) + + elif path == "/report/rollout_result": + body = self._read_json_body() + result = hub_manager.report_rollout_result( + body["worker_id"], + body["task_id"], + body["result"] + ) + return self.reply(200, result) + + elif path == "/hub/rollout_tasks": + body = self._read_json_body() + result = hub_manager.submit_rollout_task( + body["run_id"], + body["request"] + ) + return self.reply(201, result) + + else: + return self.reply(404, { + "error": {"type": "not_found", "message": "endpoint not found"} + }) + + except HubError as e: + return self.reply(e.status, { + "error": {"type": e.code, "message": str(e)} + }) + except Exception as e: + return self.reply(500, { + "error": {"type": "internal_error", "message": str(e)} + }) + + def do_DELETE(self): + try: + path = urlsplit(self.path).path + token = self._parse_token() + + if path.startswith("/hub/runs/"): + run_id = path.split("/")[-1] + result = hub_manager.cancel_run(run_id, token=token) + return self.reply(200, result) + + else: + return self.reply(404, { + "error": {"type": "not_found", "message": "endpoint not found"} + }) + + except HubError as e: + return self.reply(e.status, { + "error": {"type": e.code, "message": str(e)} + }) + except Exception as e: + return self.reply(500, { + "error": {"type": "internal_error", "message": str(e)} + }) + + class HubServer(ThreadingHTTPServer): + daemon_threads = True + request_queue_size = 128 + + server = HubServer((host, port), HubHandler) + return server + + +if __name__ == "__main__": + import sys + import signal + + # Simple CLI for testing + if len(sys.argv) < 2: + print("Usage: python -m harborrl.deployment.hub_server ") + sys.exit(1) + + import yaml + config_path = Path(sys.argv[1]) + with config_path.open() as f: + config = yaml.safe_load(f) + + manager_config = config.get("manager", {}) + + hub = HubManager( + hub_config=manager_config, + run_root=config.get("storage", {}).get("run_root", "/tmp/harborrl"), + max_concurrent_runs=manager_config.get("max_concurrent_runs", 4), + run_queue_size=manager_config.get("run_queue_size", 20), + token_file=manager_config.get("run_token_file") + ) + + server = make_hub_server( + hub, + host=manager_config.get("bind", "0.0.0.0"), + port=manager_config.get("hub_port", 49300) + ) + + print(f"Hub server listening on {server.server_address}") + + def shutdown_handler(signum, frame): + print("\nShutting down hub server...") + hub.close() + server.shutdown() + sys.exit(0) + + signal.signal(signal.SIGINT, shutdown_handler) + signal.signal(signal.SIGTERM, shutdown_handler) + + server.serve_forever() + diff --git a/harborrl/deployment/hub_synthetic_backend.py b/harborrl/deployment/hub_synthetic_backend.py new file mode 100644 index 000000000..b340fa9b4 --- /dev/null +++ b/harborrl/deployment/hub_synthetic_backend.py @@ -0,0 +1,51 @@ +"""Synthetic serving backend for transport-only rollout verification. + +This backend never touches a model: it exists so the hub -> worker -> gateway +-> runner path can be exercised end to end without a GPU serving pool. Its +output is marked synthetic, and trajectory assembly downgrades such evidence +to EVAL_ONLY, so it can never be mistaken for real RL evidence. +""" +from __future__ import annotations + +from harborrl.gateway.messages import ProtocolError + + +class SyntheticServingBackend: + """Minimal backend implementing the two calls Gateway makes.""" + + def __init__(self, *, context=32768): + self.max_context = context + self.endpoint = "synthetic://local" + + def prepare(self, converted): + ids = list(converted.get("input_ids") or [1, 2, 3]) + if len(ids) + converted["sampling"]["max_new_tokens"] > self.max_context: + raise ProtocolError("context budget exceeded") + return ids + + def count_tokens(self, converted): + return len(self.prepare(converted)) + + def weight_version(self, endpoint=None): + return "synthetic" + + def pool_states(self): + return [{"endpoint": self.endpoint, "healthy": True, + "weight_version": "synthetic"}] + + def generate(self, converted, identity, response_id): + return { + "input_ids": [1, 2], + "output_ids": [3, 4], + "logprobs": [-0.2, -0.3], + "logprob_semantics": "raw_model", + "policy_version": identity.policy_version, + "tokenizer_digest": "synthetic-tokenizer", + "template_digest": "synthetic-template", + "engine_id": "synthetic", + "serving_input": {"input_ids": [1, 2]}, + "sampling": converted["sampling"], + "finish_reason": "end_turn", + "evidence_kind": "synthetic", + "content": [{"type": "text", "text": "synthetic rollout output"}], + } diff --git a/harborrl/deployment/preflight.py b/harborrl/deployment/preflight.py index 94cb6fd6f..72448333c 100644 --- a/harborrl/deployment/preflight.py +++ b/harborrl/deployment/preflight.py @@ -104,11 +104,25 @@ def test(expression): return run, test -def _gpu_checks(plan, deployment): +def gpu_checks(plan, deployment): + """Probes that need a live GPU machine; the dispatcher reruns them. + + Plan-time preflight cannot probe a machine that provisioning has not + created yet, so it defers; the dispatcher calls this again once ``up`` + has produced one, before anything is launched. + """ gpu, connectivity = deployment["gpu_target"], deployment["connectivity"] config = plan["config"] machine = gpu["id"] checks = [] + if gpu["transport"] == "ssh" and gpu.get("ssh") is None: + # A provisioned target has no machine until the dispatch obtains one: + # probing here would test nothing, and a green check would be a lie. + checks.append(_check( + "gpu_target", machine, "gpu_target.provision", True, + "deferred: gpu_target.provision obtains the machine at dispatch, " + "which then reruns these checks before launching")) + return checks if gpu["transport"] == "ssh": spec = gpu["ssh"] started = time.monotonic() @@ -297,7 +311,7 @@ def preflight(plan, deployment, *, live=True): True, "dry-run: parsed and validated, remote probes skipped")] checks = [] checks.extend(_manager_checks(deployment)) - checks.extend(_gpu_checks(plan, deployment)) + checks.extend(gpu_checks(plan, deployment)) if deployment["rollout"]["provider"] == "ssh_docker": checks.extend(_worker_checks(plan, deployment)) checks.extend(_control_loop_checks(deployment)) diff --git a/harborrl/deployment/provision.py b/harborrl/deployment/provision.py new file mode 100644 index 000000000..6260b9605 --- /dev/null +++ b/harborrl/deployment/provision.py @@ -0,0 +1,180 @@ +"""Generic provisioning: obtain and release the GPU machine a dispatch needs. + +A deployment may declare ``gpu_target.provision`` instead of a fixed +``gpu_target.ssh``. The core owns the lifecycle — ``up`` runs before any +connectivity check, ``down`` runs in the dispatcher's ``finally`` so a failed, +cancelled or crashed run still releases the machine — while the commands and +the machine family they speak to stay site-specific: an optional practice +package implements them, the core never learns those names. + +Both commands are opaque argv lists. ``up`` prints one JSON object naming the +machine it produced; ``down`` receives that same object as JSON on stdin. +Nothing secret belongs in either direction: the record carries an ssh +destination, an optional port and identity-file *path*, and an opaque handle +the site may need to release the machine. Commands inherit the dispatcher's +environment, so a site that needs credentials keeps them in its own +environment or a mode-600 file — never in the deployment config. +""" +from __future__ import annotations + +import json +import os +import subprocess + +from harborrl.trajectories.native import publish + +# What one provision record may name. Anything else is a site's own business +# and must stay inside its opaque handle. +RECORD_FIELDS = {"host", "user", "port", "identity", "handle"} +# Error text is diagnostic, not evidence: keep it short and never echo a +# command's whole output into a record that outlives the run. +ERROR_TAIL = 400 + + +class ProvisionError(RuntimeError): + """The provisioning contract was not honoured.""" + + +def _run(argv, *, stdin_json=None, timeout=None): + return subprocess.run( + list(argv), + input=None if stdin_json is None else json.dumps(stdin_json), + capture_output=True, text=True, errors="replace", timeout=timeout, check=False, + env=dict(os.environ)) + + +def _tail(text): + text = (text or "").strip() + return text[-ERROR_TAIL:] + + +def _valid_record(record, label): + """Validate one candidate record; raises ProvisionError when unusable.""" + if not isinstance(record, dict): + raise ProvisionError(f"{label} record must be a JSON object") + unknown = sorted(set(record) - RECORD_FIELDS) + if unknown: + raise ProvisionError( + f"{label} record names unsupported fields {unknown}; keep site data in 'handle'") + host = record.get("host") + if not isinstance(host, str) or not host.strip() or any( + character.isspace() for character in host) or host.startswith("-"): + raise ProvisionError(f"{label} record requires an ssh destination host") + user = record.get("user") + if user is not None and (not isinstance(user, str) or not user.strip() + or user.startswith("-")): + raise ProvisionError(f"{label} record user must not look like a flag") + port = record.get("port") + if port is not None and (type(port) is not int or not 0 < port < 65536): + raise ProvisionError(f"{label} record port must be an integer in 1..65535") + identity = record.get("identity") + if identity is not None and (not isinstance(identity, str) or not identity.strip()): + raise ProvisionError(f"{label} record identity must be a path") + handle = record.get("handle") + if handle is not None and not isinstance(handle, str): + raise ProvisionError(f"{label} record handle must be a string") + return record + + +def _record(stdout, label): + """Parse the JSON object a provision command reports. + + A site submission tool is chatty and keeps logging after it prints the + machine, so the record is not necessarily the last line — a trailing log + entry may itself be JSON with a ``host`` key, and taking it would dispatch + onto the wrong machine. Scan backwards for the last line that validates as + a complete record; a stream with no such line fails closed rather than + guessing. + """ + lines = [line for line in (stdout or "").splitlines() if line.strip()] + if not lines: + raise ProvisionError(f"{label} reported no record on stdout") + first_error = None + for line in reversed(lines): + try: + candidate = json.loads(line) + except ValueError: + continue + try: + return _valid_record(candidate, label) + except ProvisionError as exc: + # Keep looking: a chatty trailing log object is not the record. + first_error = first_error or exc + raise first_error or ProvisionError(f"{label} did not report a JSON record") + + +def ssh_spec(record): + """The normalized ssh spec a provision record licenses.""" + spec = {"host": record["host"]} + for key in ("user", "port", "identity"): + if record.get(key) is not None: + spec[key] = record[key] + return spec + + +def up(provision, *, evidence_path=None): + """Run ``up`` and return the normalized machine record. + + The record is published to ``evidence_path`` (when given) before dispatch + proceeds, so the machine that a run actually used — and the handle needed + to release it — outlives the process that obtained it. + + A publish failure releases the machine before raising. Nothing downstream + would ever learn the handle, so leaving it running would bill an + allocation that no run directory can account for. + """ + try: + result = _run(provision["up"], timeout=provision["timeout_s"]) + except subprocess.TimeoutExpired as exc: + raise ProvisionError( + f"gpu_target.provision.up exceeded {provision['timeout_s']}s") from exc + except OSError as exc: + raise ProvisionError(f"gpu_target.provision.up could not start: {exc}") from exc + if result.returncode != 0: + raise ProvisionError( + f"gpu_target.provision.up failed with {result.returncode}: " + f"{_tail(result.stderr)}") + record = _record(result.stdout, "gpu_target.provision.up") + if evidence_path is not None: + try: + publish(evidence_path, {"action": "up", "record": record}) + except Exception as exc: # noqa: BLE001 — any evidence failure leaks otherwise + outcome, detail = down(provision, record) + raise ProvisionError( + f"cannot record the provisioned machine {record.get('handle')!r} " + f"at {evidence_path}: {exc}; release reported {outcome} ({detail})" + ) from exc + return record + + +def down(provision, record, *, evidence_path=None): + """Run ``down`` for a machine this dispatch obtained. + + Returns ``"released"``, ``"undeclared"`` (no down command: the machine + lives out its own declared lifetime) or ``"failed"``; the release record + is published either way so a leaked machine is visible, never silent. + """ + command = provision.get("down") + if command is None: + outcome, detail = "undeclared", "no gpu_target.provision.down declared" + result = None + else: + try: + result = _run(command, stdin_json=record, timeout=provision["timeout_s"]) + outcome = "released" if result.returncode == 0 else "failed" + detail = (f"exit {result.returncode}: {_tail(result.stderr)}" + if result.returncode else "released") + except subprocess.TimeoutExpired: + outcome, detail = "failed", f"exceeded {provision['timeout_s']}s" + except OSError as exc: + outcome, detail = "failed", f"could not start: {exc}" + if evidence_path is not None: + payload = {"action": "down", "outcome": outcome, "detail": detail, + "record": record} + try: + publish(evidence_path, payload) + except ValueError: + # Release evidence is advisory; the machine's fate is what the + # outcome reports, not whether the record could be written. + pass + return outcome, detail diff --git a/harborrl/deployment/schema.py b/harborrl/deployment/schema.py index 3cc7bf7c8..ec179be1e 100644 --- a/harborrl/deployment/schema.py +++ b/harborrl/deployment/schema.py @@ -9,6 +9,7 @@ from __future__ import annotations import re +import shlex import sys from pathlib import Path from urllib.parse import urlsplit @@ -23,10 +24,23 @@ SLUG = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]*") LOOPBACK_HOSTS = {"localhost", "127.0.0.1", "::1", "0.0.0.0", "[::1]"} -MANAGER_FIELDS = {"id", "mode", "bind", "control_port", "control_url", "gateway_port", - "gateway_url", "run_token_file", "source_dir", "controller_python"} +MANAGER_REQUIRED = {"id", "mode", "bind", "control_port", "control_url", "gateway_port", + "gateway_url", "run_token_file", "source_dir", "controller_python"} +MANAGER_OPTIONAL = {"gateway_shards", "gateway_routing", "sandboxes_per_shard", + "worker_shards", "gateway_hosts", "shard_python"} +MANAGER_FIELDS = MANAGER_REQUIRED | MANAGER_OPTIONAL +SHARD_FIELDS = {"name", "port", "mode", "url", "admin_url"} +SHARD_MODES = ("local", "remote") +SHARD_NAME = re.compile(r"[a-z0-9][a-z0-9-]{0,31}") +GATEWAY_ROUTINGS = ("modulo", "worker_affinity", "auto") +DEFAULT_SANDBOXES_PER_SHARD = 30 SSH_FIELDS = {"host", "user", "port", "identity"} -GPU_FIELDS = {"id", "transport", "ssh", "source_dir", "train_python", "environment_file", "expected_gpus"} +GPU_REQUIRED = {"id", "transport", "ssh", "source_dir", "train_python", "environment_file", + "expected_gpus"} +GPU_OPTIONAL = {"provision"} +GPU_FIELDS = GPU_REQUIRED | GPU_OPTIONAL +PROVISION_FIELDS = {"up", "down", "timeout_s"} +DEFAULT_PROVISION_TIMEOUT_S = 900 STORAGE_FIELDS = {"mode", "run_root", "manager_writable", "gpu_writable", "worker_writable"} WORKER_FIELDS = {"id", "ssh", "source_dir", "runner_python", "docker", "task_roots"} ROLLOUT_FIELDS = {"provider", "workers"} @@ -123,9 +137,88 @@ def destination(spec): return f"{spec['user']}@{spec['host']}" if spec.get("user") else spec["host"] +def _growth_host(value, label): + """Validate one growth-pool entry, preserving its declared spelling. + + A plain string is an ssh destination that inherits every transport option + from the invoking user's ssh config — the historical form, kept byte for + byte so legacy digests do not move. A mapping narrows those options for + an ephemeral machine whose port and identity are known per run rather + than in the config file. + """ + if isinstance(value, str): + if (not value.strip() or any(character.isspace() for character in value) + or value.startswith("-") or value.count("@") > 1): + raise DeploymentError(f"{label} must be an ssh destination (user@host)") + return value + if not isinstance(value, dict) or set(value) - SSH_FIELDS or "host" not in value: + raise DeploymentError( + f"{label} requires at most {sorted(SSH_FIELDS)} and a host") + return _ssh_spec(value, label) + + +def _growth_key(entry): + """The identity used for duplicate detection across both entry forms.""" + return destination(entry) if isinstance(entry, dict) else entry + + +def _gateway_shards(value, control_port, *, service): + """Normalize manager.gateway_shards; absent keeps the legacy single gateway.""" + if value is None: + return None + if not isinstance(value, list) or not value: + raise DeploymentError("manager.gateway_shards must be a nonempty list") + shards, names, ports = [], set(), set() + for entry in value: + if not isinstance(entry, dict) or set(entry) - SHARD_FIELDS or "name" not in entry: + raise DeploymentError( + f"manager.gateway_shards entries require at most {sorted(SHARD_FIELDS)} and a name") + name = entry["name"] + if not isinstance(name, str) or not SHARD_NAME.fullmatch(name): + raise DeploymentError("manager.gateway_shards names are lowercase slugs") + mode = entry.get("mode", "local") + if mode not in SHARD_MODES: + raise DeploymentError(f"manager.gateway_shards mode must be one of {SHARD_MODES}") + shard = {"name": name} + if mode == "remote": + if not service: + raise DeploymentError( + "remote gateway shards require a service-mode manager") + if "port" in entry: + raise DeploymentError("a remote gateway shard takes its port from url") + shard["mode"] = "remote" + # A remote shard may bind on a hub loopback when reached through a + # reverse tunnel: the forward onto 127.0.0.1:N is the real route. + shard["url"] = _plain_origin(entry.get("url"), "manager.gateway_shards url", + allow_loopback=True) + shard["admin_url"] = _plain_origin(entry.get("admin_url"), + "manager.gateway_shards admin_url", + allow_loopback=True) + else: + if "url" in entry or "admin_url" in entry: + raise DeploymentError("a local gateway shard binds a port and derives its url") + port = entry.get("port") + if port is not None: + port = _positive_int(port, "manager.gateway_shards port") + if port == control_port: + raise DeploymentError("manager.gateway_shards port collides with control_port") + if port in ports: + raise DeploymentError(f"duplicate gateway shard port: {port}") + ports.add(port) + shard["port"] = port + if name in names: + raise DeploymentError(f"duplicate gateway shard name: {name}") + names.add(name) + shards.append(shard) + return shards + + def _manager(config): - if not isinstance(config, dict) or set(config) != MANAGER_FIELDS: - raise DeploymentError(f"manager requires exactly {sorted(MANAGER_FIELDS)}") + if not isinstance(config, dict) or set(config) - MANAGER_FIELDS: + raise DeploymentError(f"manager supports only {sorted(MANAGER_FIELDS)}") + missing = MANAGER_REQUIRED - set(config) + if missing: + raise DeploymentError(f"manager requires {sorted(missing)}") mode = config["mode"] if mode not in ("embedded", "service"): raise DeploymentError("manager.mode must be embedded or service") @@ -140,7 +233,65 @@ def _manager(config): port=control_port, allow_loopback=not service) gateway_url = _plain_origin(config["gateway_url"], "manager.gateway_url", port=gateway_port, allow_loopback=not service) - return { + routing = config.get("gateway_routing", "modulo") + if routing not in GATEWAY_ROUTINGS: + raise DeploymentError(f"manager.gateway_routing must be one of {GATEWAY_ROUTINGS}") + shards = _gateway_shards(config.get("gateway_shards"), control_port, service=service) + # A static table and elastic growth may coexist: declared shards exist + # from startup (a long-lived host served through a tunnel), and auto + # routing grows beyond them when concurrency passes the per-shard budget. + sandboxes_per_shard = config.get("sandboxes_per_shard", DEFAULT_SANDBOXES_PER_SHARD) + if type(sandboxes_per_shard) is not int or sandboxes_per_shard < 1: + raise DeploymentError( + "manager.sandboxes_per_shard must be a positive integer (default " + f"{DEFAULT_SANDBOXES_PER_SHARD} sandboxes per gateway)") + worker_shards = config.get("worker_shards") + if worker_shards is not None: + if shards is None: + raise DeploymentError("manager.worker_shards requires manager.gateway_shards") + if (not isinstance(worker_shards, dict) + or not all(isinstance(key, str) and key.strip() + and isinstance(value, str) and value.strip() + for key, value in worker_shards.items())): + raise DeploymentError( + "manager.worker_shards maps worker ids (or 'local') to gateway shard names") + gateway_hosts = config.get("gateway_hosts") + if gateway_hosts is not None: + if not service: + raise DeploymentError( + "manager.gateway_hosts grows remote shards on a service-mode manager") + if not isinstance(gateway_hosts, list) or not gateway_hosts: + raise DeploymentError( + "manager.gateway_hosts is a nonempty list of ssh destinations (user@host)") + normalized_hosts = [] + for index, entry in enumerate(gateway_hosts): + if (isinstance(entry, str) and entry.strip() + and not any(character.isspace() for character in entry)): + normalized_hosts.append(entry) + elif isinstance(entry, dict): + normalized_hosts.append( + _growth_host(entry, f"manager.gateway_hosts[{index}]")) + else: + # Historical wording: a malformed string entry is still a + # missing destination, not a new option form. + raise DeploymentError( + "manager.gateway_hosts is a nonempty list of ssh destinations (user@host)") + gateway_hosts = normalized_hosts + if len({_growth_key(entry) for entry in gateway_hosts}) != len(gateway_hosts): + raise DeploymentError("manager.gateway_hosts entries must be unique") + if config.get("gateway_routing") in ("modulo", "worker_affinity"): + # An undeclared routing keeps the runtime's elastic auto default; + # only an explicit static routing refuses the growth pool. + raise DeploymentError( + "manager.gateway_hosts grows gateways under auto routing") + shard_python = config.get("shard_python") + if shard_python is not None: + if gateway_hosts is None: + raise DeploymentError( + "manager.shard_python interprets gateway_hosts growth; " + "declare the pool first") + shard_python = _absolute_path(shard_python, "manager.shard_python") + manager = { "id": _slug(config["id"], "manager.id"), "mode": mode, "bind": _nonempty_str(config["bind"], "manager.bind"), "control_port": control_port, "control_url": control_url, @@ -149,17 +300,96 @@ def _manager(config): "source_dir": _absolute_path(config["source_dir"], "manager.source_dir"), "controller_python": _absolute_path(config["controller_python"], "manager.controller_python"), } + # Shard keys appear only when sharding is declared, so legacy deployments + # keep their exact historical shape (and digest). An explicitly declared + # routing without shards opts in or out of auto growth; undeclared means + # the runtime's elastic default (one gateway until sandboxes_per_shard). + if shards is not None: + manager["gateway_shards"] = shards + # A declared table pins the historical static default, unless the + # growth pool asks auto routing to expand beyond it. + manager["gateway_routing"] = ( + "auto" if gateway_hosts is not None and "gateway_routing" not in config + else routing) + if worker_shards is not None: + manager["worker_shards"] = dict(worker_shards) + elif gateway_hosts is not None: + # The pool is meaningless without growth; never let the runtime see a + # static routing that it would reject after the GPU job is submitted. + manager["gateway_routing"] = "auto" + elif "gateway_routing" in config: + manager["gateway_routing"] = routing + if "sandboxes_per_shard" in config: + manager["sandboxes_per_shard"] = sandboxes_per_shard + if gateway_hosts is not None: + # Auto growth fills this pool before falling back to the manager + # machine; the entry spelling is the URL host runners will use. + manager["gateway_hosts"] = list(gateway_hosts) + if shard_python is not None: + manager["shard_python"] = shard_python + return manager + + +def _command(value, label): + """A provision command: an argv list, or a string split without a shell.""" + if isinstance(value, str): + parts = shlex.split(value) + if not parts: + raise DeploymentError(f"{label} must be a nonempty command") + return parts + if (not isinstance(value, list) or not value + or not all(isinstance(part, str) and part for part in value)): + raise DeploymentError(f"{label} must be a nonempty argv list of strings") + return list(value) + + +def _provision(value): + """Normalize gpu_target.provision: how a dispatch obtains its GPU machine. + + The commands are opaque to the core: ``up`` prints one JSON object naming + the machine it produced, ``down`` releases it. Nothing site-specific is + interpreted here, and no secret belongs in either command. + """ + if (not isinstance(value, dict) or set(value) - PROVISION_FIELDS + or "up" not in value): + raise DeploymentError( + f"gpu_target.provision supports {sorted(PROVISION_FIELDS)} and requires up") + provision = {"up": _command(value["up"], "gpu_target.provision.up")} + down = value.get("down") + if down is not None: + provision["down"] = _command(down, "gpu_target.provision.down") + provision["timeout_s"] = _positive_int( + value.get("timeout_s", DEFAULT_PROVISION_TIMEOUT_S), + "gpu_target.provision.timeout_s") + return provision def _gpu_target(config): - if not isinstance(config, dict) or set(config) != GPU_FIELDS: - raise DeploymentError(f"gpu_target requires exactly {sorted(GPU_FIELDS)}") + if (not isinstance(config, dict) or set(config) - GPU_FIELDS + or GPU_REQUIRED - set(config)): + raise DeploymentError( + f"gpu_target requires {sorted(GPU_REQUIRED)} and supports at most " + f"{sorted(GPU_FIELDS)}") transport = config["transport"] if transport not in TRANSPORTS: raise DeploymentError(f"gpu_target.transport must be one of {TRANSPORTS}") + provision = config.get("provision") + if provision is not None: + # A provisioned machine is reached over ssh; a local target already + # exists and has nothing to obtain or release. + if transport != "ssh": + raise DeploymentError("gpu_target.provision requires ssh transport") + if config["ssh"] is not None: + raise DeploymentError( + "gpu_target.provision supplies the machine; declare no gpu_target.ssh") + provision = _provision(provision) ssh = config["ssh"] if transport == "ssh": - ssh = _ssh_spec(ssh, "gpu_target.ssh") + if ssh is None and provision is None: + raise DeploymentError( + "gpu_target.ssh is required unless gpu_target.provision supplies the machine") + if ssh is not None: + ssh = _ssh_spec(ssh, "gpu_target.ssh") elif ssh is not None: raise DeploymentError("gpu_target.ssh must be null for local transport") target = { @@ -171,6 +401,8 @@ def _gpu_target(config): environment_file = config["environment_file"] if environment_file is not None: target["environment_file"] = _absolute_path(environment_file, "gpu_target.environment_file") + if provision is not None: + target["provision"] = provision return target @@ -235,6 +467,29 @@ def _connectivity(config): } +def _check_gateway_shard_routing(manager, raw_manager, rollout): + """Static routing must be total before dispatch, not discovered mid-run.""" + declared_routing = raw_manager.get("gateway_routing", "modulo") + shards = manager.get("gateway_shards") + worker_shards = manager.get("worker_shards") or {} + if declared_routing == "worker_affinity" and not shards: + raise DeploymentError("worker_affinity routing requires manager.gateway_shards") + worker_ids = {worker["id"] for worker in rollout["workers"]} + if declared_routing == "worker_affinity": + missing = sorted(worker_ids - set(worker_shards)) + if missing: + raise DeploymentError( + f"worker_affinity requires a shard for every rollout worker; missing {missing}") + if not worker_shards: + return + unknown_keys = sorted(set(worker_shards) - worker_ids - {"local"}) + if unknown_keys: + raise DeploymentError(f"manager.worker_shards names unknown workers {unknown_keys}") + unknown_shards = sorted(set(worker_shards.values()) - {shard["name"] for shard in shards}) + if unknown_shards: + raise DeploymentError(f"manager.worker_shards names unknown shards {unknown_shards}") + + def validate(config): """Strictly validate and normalize a deployment mapping.""" if not isinstance(config, dict): @@ -245,12 +500,15 @@ def validate(config): + ", ".join(sorted(DEPLOYMENT_SECTIONS))) if config["schema_version"] != DEPLOYMENT_SCHEMA_VERSION or type(config["schema_version"]) is not int: raise DeploymentError(f"deployment schema_version must be {DEPLOYMENT_SCHEMA_VERSION}") + manager = _manager(config["manager"]) + rollout = _rollout(config["rollout"]) + _check_gateway_shard_routing(manager, config["manager"], rollout) return { "schema_version": DEPLOYMENT_SCHEMA_VERSION, - "manager": _manager(config["manager"]), + "manager": manager, "gpu_target": _gpu_target(config["gpu_target"]), "storage": _storage(config["storage"]), - "rollout": _rollout(config["rollout"]), + "rollout": rollout, "connectivity": _connectivity(config["connectivity"]), } diff --git a/harborrl/deployment/training_target.py b/harborrl/deployment/training_target.py index 4de5d2468..8df8168cb 100644 --- a/harborrl/deployment/training_target.py +++ b/harborrl/deployment/training_target.py @@ -201,6 +201,7 @@ def cancel(self, *, grace_seconds=30): def local_target(plan, *, environment): """Build the legacy in-process target from a resolved plan.""" allowed = {"PATH", "HOME", "USER", "LANG", "LC_ALL", "TMPDIR", "LD_LIBRARY_PATH", + "LD_PRELOAD", "CUDA_HOME", "NVTE_CUDA_INCLUDE_DIR", "CUDA_VISIBLE_DEVICES", "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY", "http_proxy", "https_proxy", "no_proxy", "PYTHONPATH", "SLIME_DIR", "MEGATRON_DIR", "SGLANG_IMAGE"} diff --git a/harborrl/distill/__init__.py b/harborrl/distill/__init__.py new file mode 100644 index 000000000..821550746 --- /dev/null +++ b/harborrl/distill/__init__.py @@ -0,0 +1 @@ +"""Independent single-turn math records and reference on-policy training.""" diff --git a/harborrl/distill/cli.py b/harborrl/distill/cli.py new file mode 100644 index 000000000..6789064b1 --- /dev/null +++ b/harborrl/distill/cli.py @@ -0,0 +1,126 @@ +"""Math-only data import, on-policy reference training and held-out evaluation.""" +import argparse +import json +import os +import sys +from pathlib import Path + +from harborrl.deployment.schema import embedded_local, load_deployment +from harborrl.deployment.training_target import LocalTrainingTarget, SshTrainingTarget +from harborrl.platform.evidence import _source_audit +from harborrl.trajectories.native import digest, publish + +from .config import load +from .data import import_samples, load_tasks + + +def plan(config, *, allow_synthetic=False, resume=None): + tasks = load_tasks(config["tasks"]["path"]) + if not any(task["split"] == "train" for task in tasks): + raise ValueError("No training split") + root = Path(config["output"]["directory"]) + deployment, source = (load_deployment(config["deployment"]["file"]) if config["deployment"]["file"] + else (embedded_local(), "embedded-local")) + gpu = deployment["gpu_target"] + if config["deployment"]["device"] == "cuda" and source != "embedded-local" and gpu["expected_gpus"] < config["deployment"]["world_size"]: + raise ValueError("GPU target budget below requested world size") + python = gpu["train_python"] if source != "embedded-local" else sys.executable + command = [python] + if config["deployment"]["world_size"] > 1: + command += ["-m", "torch.distributed.run", "--standalone", f"--nproc-per-node={config['deployment']['world_size']}", "--module", "harborrl.distill.runner"] + else: + command += ["-m", "harborrl.distill.runner"] + command += ["--config", str(root / "config.json")] + if allow_synthetic: + command.append("--allow-synthetic") + if resume: + command += ["--resume", str(Path(resume).resolve())] + return {"backend_contract": config["backend_contract"], "method": config["method"], "command": command, + "config": config, "config_digest": digest(config), "tasks_digest": digest(tasks), + "deployment_source": source, "gpu_target": gpu, "connectivity": deployment["connectivity"], + "scope": "reference math pipeline; live dependencies, model compatibility and capacity remain unchecked"} + + +def launch(value): + root = Path(value["config"]["output"]["directory"]) + if root.exists(): + raise ValueError("Run output must be a new directory; resume into a different output") + root.mkdir(parents=True) + publish(root / "config.json", value["config"]) + publish(root / "plan.json", value) + _source_audit(root) + gpu = value["gpu_target"] + environment = {"PYTHONUNBUFFERED": "1", "HF_HUB_OFFLINE": "1", "TRANSFORMERS_OFFLINE": "1", + "TOKENIZERS_PARALLELISM": "false"} + if gpu["transport"] == "ssh": + target = SshTrainingTarget(gpu["ssh"], source_dir=gpu["source_dir"], run_dir=root, + environment_file=gpu.get("environment_file"), + command_timeout=value["connectivity"]["command_timeout_s"] + 30, + ssh_timeout=value["connectivity"]["ssh_timeout_s"]) + target.write_launcher(value["command"], {**environment, "PYTHONPATH": gpu["source_dir"]}) + import signal + previous = {sig: signal.getsignal(sig) for sig in (signal.SIGINT, signal.SIGTERM)} + def cancel(_signum, _frame): + target.cancel() + raise KeyboardInterrupt + try: + for sig in previous: + signal.signal(sig, cancel) + target.launch() + return target.wait() + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + from harborrl.config import ROOT + env = {**os.environ, **environment} + env["PYTHONPATH"] = os.pathsep.join(filter(None, [str(ROOT), env.get("PYTHONPATH")])) + target = LocalTrainingTarget(value["command"], cwd=ROOT, env=env) + target.launch() + try: + return target.wait() + finally: + if target.child.poll() is None: + target.cancel() + target.child.wait(timeout=30) + + +def main(argv=None): + parser = argparse.ArgumentParser(prog="harborrl distill", description=__doc__) + commands = parser.add_subparsers(dest="command", required=True) + imported = commands.add_parser("import", help="Import curated responses into immutable sample/verification records") + imported.add_argument("--tasks", required=True) + imported.add_argument("--samples", required=True) + imported.add_argument("--out-dir", required=True) + for name in ("train", "evaluate"): + command = commands.add_parser(name) + command.add_argument("config") + command.add_argument("--allow-synthetic", action="store_true") + if name == "train": + command.add_argument("--dry-run", action="store_true") + command.add_argument("--resume", help="Trusted checkpoint; use a new output directory") + else: + command.add_argument("--checkpoint", required=True) + command.add_argument("--output", required=True) + args = parser.parse_args(argv) + try: + if args.command == "import": + print(json.dumps(import_samples(args.tasks, args.samples, args.out_dir), indent=2)) + return 0 + config = load(args.config, allow_synthetic=args.allow_synthetic) + if args.command == "evaluate": + from .runner import evaluate + result = evaluate(config, args.checkpoint) + publish(args.output, result) + print(json.dumps(result, indent=2)) + return 0 + value = plan(config, allow_synthetic=args.allow_synthetic, resume=args.resume) + if args.dry_run: + print(json.dumps(value, indent=2)) + return 0 + return launch(value) + except (ValueError, TypeError, KeyError, OSError) as exc: + parser.error(str(exc)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harborrl/distill/config.py b/harborrl/distill/config.py new file mode 100644 index 000000000..5098bc11f --- /dev/null +++ b/harborrl/distill/config.py @@ -0,0 +1,108 @@ +"""Independent distill schema; the legacy GRPO schema is unchanged.""" +from copy import deepcopy +from pathlib import Path + +import yaml + +from harborrl.consistency.config import DEFAULTS as CONSISTENCY_DEFAULTS +from harborrl.consistency.config import validate as validate_consistency +from harborrl.trajectories.native import finite + +CONTRACT = "transformers-math-onpolicy-v1" +METHODS = ("reinforce", "prm", "opd", "prm_opd") +DEFAULTS = { + "consistency": CONSISTENCY_DEFAULTS, + "schema_version": 1, "backend_contract": CONTRACT, "method": "opd", + "tasks": {"path": None, "train_split": "train", "eval_split": "test"}, + "student": {"kind": "tiny_qwen", "checkpoint": None, "hidden_size": 32, "layers": 1, "heads": 2}, + "teacher": {"kind": "fake", "checkpoint": None}, + "prm": {"kind": "none", "weight": 1.0}, + "sampling": {"max_new_tokens": 8, "max_context": 256, "temperature": 1.0, "top_p": 1.0}, + "training": {"steps": 2, "batch_size": 8, "microbatch_size": 2, "learning_rate": 0.0003, + "seed": 17, "max_grad_norm": 1.0, "opd_weight": 1.0, "outcome_weight": 0.0, + "save_interval": 1, "behavior_tolerance": 0.0001}, + "deployment": {"file": None, "device": "cpu", "world_size": 1}, + "output": {"directory": None}, +} + + +def merge(base, override, prefix=""): + if not isinstance(override, dict): + raise TypeError(f"{prefix or 'config'} must be a mapping") + result = deepcopy(base) + for key, value in override.items(): + if key not in base: + raise ValueError(f"Unknown field: {prefix}{key}") + result[key] = merge(base[key], value, f"{prefix}{key}.") if isinstance(base[key], dict) else value + return result + + +def validate(raw, path, *, allow_synthetic=False): + config = merge(DEFAULTS, raw) + config['consistency'] = validate_consistency(config['consistency']) + if type(config["schema_version"]) is not int or config["schema_version"] != 1 or config["backend_contract"] != CONTRACT: + raise ValueError("distill requires schema 1 / transformers-math-onpolicy-v1") + if config["method"] not in METHODS: + raise ValueError(f"method must be one of {METHODS}") + student, teacher, prm = config["student"], config["teacher"], config["prm"] + if student["kind"] not in {"tiny_qwen", "hf_local"} or teacher["kind"] not in {"none", "fake", "tiny_qwen", "hf_local"}: + raise ValueError("Unsupported student or teacher adapter") + if prm["kind"] not in {"none", "fake", "equation"}: + raise ValueError("Unsupported PRM adapter") + if not allow_synthetic and (student["kind"] == "tiny_qwen" or teacher["kind"] in {"fake", "tiny_qwen"} or prm["kind"] == "fake"): + raise ValueError("Tiny/fake adapters require explicit --allow-synthetic") + opd = config["method"] in {"opd", "prm_opd"} + process = config["method"] in {"prm", "prm_opd"} + if opd == (teacher["kind"] == "none") or process == (prm["kind"] == "none"): + raise ValueError("Selected method requires exactly its teacher/PRM adapters; disable unused adapters") + if teacher["kind"] == "tiny_qwen" and student["kind"] != "tiny_qwen": + raise ValueError("Tiny teacher requires the same tiny tokenizer") + for section, fields in (("student", ("hidden_size", "layers", "heads")), + ("sampling", ("max_new_tokens", "max_context")), + ("training", ("steps", "batch_size", "microbatch_size", "save_interval")), + ("deployment", ("world_size",))): + for field in fields: + if type(config[section][field]) is not int or config[section][field] <= 0: + raise ValueError(f"{section}.{field} must be a positive integer") + if type(config["training"]["seed"]) is not int or config["training"]["seed"] < 0: + raise ValueError("training.seed must be a nonnegative integer") + for section, fields in (("training", ("learning_rate", "max_grad_norm", "behavior_tolerance")),): + for field in fields: + if not finite(config[section][field]) or config[section][field] <= 0: + raise ValueError(f"{section}.{field} must be positive and finite") + for section, field in (("training", "opd_weight"), ("training", "outcome_weight"), ("prm", "weight")): + if not finite(config[section][field]) or config[section][field] < 0: + raise ValueError(f"{section}.{field} must be nonnegative and finite") + if (opd and config["training"]["opd_weight"] == 0) or (process and prm["weight"] == 0): + raise ValueError("The selected method's loss weight must be positive") + if config["method"] == "reinforce" and config["training"]["outcome_weight"] == 0: + raise ValueError("reinforce requires positive outcome_weight") + if config["sampling"]["temperature"] != 1.0 or config["sampling"]["top_p"] != 1.0: + raise ValueError("Only unmodified raw-policy sampling (temperature=top_p=1) is supported") + if config["sampling"]["max_new_tokens"] >= config["sampling"]["max_context"]: + raise ValueError("Response limit must be below context limit") + if student["hidden_size"] % student["heads"] or (student["hidden_size"] // student["heads"]) % 2: + raise ValueError("Tiny Qwen head size must be an even integer") + if config["deployment"]["device"] not in {"cpu", "cuda"}: + raise ValueError("deployment.device must be cpu or cuda") + if config["training"]["batch_size"] % config["deployment"]["world_size"]: + raise ValueError("Global batch must divide evenly among data-parallel ranks") + for key in ("train_split", "eval_split"): + if config["tasks"][key] not in {"train", "validation", "test"}: + raise ValueError("Unknown task split") + if config["tasks"]["train_split"] != "train" or config["tasks"]["eval_split"] == "train": + raise ValueError("Training uses train; evaluation must be held out") + for section, field, required in (("tasks", "path", True), ("output", "directory", True), + ("deployment", "file", False), ("student", "checkpoint", student["kind"] == "hf_local"), + ("teacher", "checkpoint", teacher["kind"] == "hf_local")): + value = config[section][field] + if value is None and not required: + continue + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"{section}.{field} requires a path") + config[section][field] = str((Path(path).resolve().parent / Path(value).expanduser()).resolve()) + return config + + +def load(path, *, allow_synthetic=False): + return validate(yaml.safe_load(Path(path).read_text(encoding="utf-8")), path, allow_synthetic=allow_synthetic) diff --git a/harborrl/distill/contract.py b/harborrl/distill/contract.py new file mode 100644 index 000000000..a7d4853bd --- /dev/null +++ b/harborrl/distill/contract.py @@ -0,0 +1,71 @@ +"""Fail-closed on-policy batch admission independent of Native RL_READY.""" +from harborrl.trajectories.native import digest, finite + +ROW_FIELDS = {"schema_version", "sample_id", "task_id", "task_digest", "split", "student_behavior_version", + "student_evidence", "teacher_revision", "teacher_evidence", "identity", "prompt_ids", "response_ids", + "behavior_logprobs", "teacher_logprobs", "loss_mask", "sampling", "stop_reason", "score_request_digest", + "score_response_digest", "response", "verification", "prm", "digest"} + + +def score_request(row): + return {key: row[key] for key in ("prompt_ids", "response_ids", "identity", "teacher_revision", "student_behavior_version")} + + +def validate_row(row, *, version, identity, context, allow_synthetic=False, require_teacher=True): + if not isinstance(row, dict) or set(row) != ROW_FIELDS or row["schema_version"] != "math-onpolicy-batch-v1": + raise ValueError("Invalid on-policy batch row fields") + if digest({k: v for k, v in row.items() if k != "digest"}) != row["digest"]: + raise ValueError("Batch content digest mismatch") + if row["student_behavior_version"] != version or row["identity"] != identity: + raise ValueError("Stale student version or tokenizer/template/vocabulary mismatch") + if row["student_evidence"] != "observed" or row["split"] != "train": + raise ValueError("Only observed student generation on the train split is admitted") + if require_teacher and (not row["teacher_revision"] or row["teacher_evidence"] not in {"observed", "synthetic"}): + raise ValueError("Missing locked teacher scoring evidence") + if row["teacher_evidence"] == "synthetic" and not allow_synthetic: + raise ValueError("Synthetic teacher is test-only") + if row["prm"] and row["prm"]["evidence_kind"] == "synthetic" and not allow_synthetic: + raise ValueError("Synthetic PRM is test-only") + if row["sampling"] != {"temperature": 1.0, "top_p": 1.0, "raw_logprobs": True}: + raise ValueError("Batch must use raw on-policy sampling") + p, r = row["prompt_ids"], row["response_ids"] + if not p or not r or len(p) + len(r) > context: + raise ValueError("Empty or over-context sequence") + if any(type(token) is not int or not 0 <= token < identity["vocab_size"] for token in p + r): + raise ValueError("Invalid token ID") + for key in ("behavior_logprobs", "teacher_logprobs", "loss_mask"): + if len(row[key]) != len(r) and not (key == "teacher_logprobs" and not require_teacher and row[key] == []): + raise ValueError("Token/probability/mask length mismatch") + for key in ("behavior_logprobs", "teacher_logprobs"): + if any(not finite(value) or value > 0 for value in row[key]): + raise ValueError("Nonfinite or positive raw log probability") + if row["loss_mask"] != [1] * len(r): + raise ValueError("Only actual single-turn generated tokens enter this batch") + verification = row["verification"] + rewards = {"correct": 1.0, "incorrect": -1.0, "unparseable": 0.0, "unsupported": 0.0} + if (not isinstance(verification, dict) or verification.get("status") not in rewards + or verification.get("reward") != rewards[verification["status"]]): + raise ValueError("Invalid verification reward") + prm = row["prm"] + if prm is not None: + values = prm.get("token_returns", []) + if len(values) != len(r) or any(not finite(value) for value in values): + raise ValueError("Invalid PRM token returns") + cursor, cumulative = len(r), 0.0 + for step in reversed(prm.get("steps", [])): + if (type(step.get("start")) is not int or type(step.get("end")) is not int + or not 0 <= step["start"] < step["end"] == cursor or not finite(step.get("score"))): + raise ValueError("Invalid PRM step boundary or score") + cumulative += step["score"] + if values[step["start"]:cursor] != [cumulative] * (cursor - step["start"]): + raise ValueError("PRM returns disagree with step rewards") + cursor = step["start"] + if cursor: + raise ValueError("PRM steps do not cover the response") + if row["score_request_digest"] != digest(score_request(row)): + raise ValueError("Scorer request mismatch") + if row["score_response_digest"] != digest({"request_digest": row["score_request_digest"], "logprobs": row["teacher_logprobs"]}): + raise ValueError("Scorer response mismatch") + if row["stop_reason"] not in {"eos", "length"}: + raise ValueError("Unknown stop reason") + return row diff --git a/harborrl/distill/data.py b/harborrl/distill/data.py new file mode 100644 index 000000000..187643008 --- /dev/null +++ b/harborrl/distill/data.py @@ -0,0 +1,97 @@ +"""Locked math tasks and bounded exact-rational verification; never eval text.""" +import json +import re +from fractions import Fraction +from pathlib import Path + +from harborrl.trajectories.native import digest, publish + +NUMBER = r"[+-]?(?:\d{1,32}(?:\.\d{1,16})?|\d{1,32}/[1-9]\d{0,31})" +TASK_FIELDS = {"task_id", "problem_id", "prompt", "reference", "verifier", "split", "source", "license"} + + +def load_tasks(path): + rows, identities, problems, prompts = [], set(), {}, {} + source = Path(path) + if source.stat().st_size > 16 * 1024 * 1024: + raise ValueError("Task file exceeds the 16 MiB reference-backend limit") + for index, line in enumerate(source.read_text(encoding="utf-8").splitlines(), 1): + row = json.loads(line) + if not isinstance(row, dict) or set(row) != TASK_FIELDS: + raise ValueError(f"Task line {index} requires exactly {sorted(TASK_FIELDS)}") + if any(not isinstance(value, str) or not value.strip() for value in row.values()): + raise ValueError("Task fields must be nonempty strings") + if row["split"] not in {"train", "validation", "test"} or row["verifier"] != "exact_rational_v1": + raise ValueError("Only split-locked exact_rational_v1 tasks are supported") + if len(row["prompt"]) > 8192 or not re.fullmatch(NUMBER, row["reference"]): + raise ValueError("Unsupported prompt length or rational reference") + if row["task_id"] in identities: + raise ValueError("Duplicate task_id") + identities.add(row["task_id"]) + for key, table in ((row["problem_id"], problems), (digest(row["prompt"]), prompts)): + if key in table and table[key] != row["split"]: + raise ValueError("Problem or prompt leakage across splits") + table[key] = row["split"] + rows.append(row) + if not rows: + raise ValueError("Empty math task source") + return rows + + +def verify(response, reference): + record = {"verifier": "exact_rational_v1", "input_digest": digest({"response": response, "reference": reference}), + "status": "unparseable", "reward": 0.0, "parsed": None} + if not isinstance(response, str) or len(response) > 16384: + record["status"] = "unsupported" + return record + text = response.strip() + if "####" in text: + text = text.rsplit("####", 1)[1].strip() + elif "\\boxed{" in text: + match = re.search(r"\\boxed\{([^{}]{1,100})\}\s*$", text) + text = match[1] if match else "" + if not re.fullmatch(NUMBER, text): + return record + correct = Fraction(text) == Fraction(reference) + record.update(status="correct" if correct else "incorrect", reward=1.0 if correct else -1.0, parsed=text) + return record + + +def import_samples(tasks_path, samples_path, output): + tasks = {row["task_id"]: row for row in load_tasks(tasks_path)} + root = Path(output) + if root.exists(): + raise ValueError("Dataset output must be a new directory") + if Path(samples_path).stat().st_size > 16 * 1024 * 1024: + raise ValueError("Sample file exceeds the 16 MiB reference-backend limit") + raw = Path(samples_path).read_bytes() + rows = [json.loads(line) for line in raw.decode("utf-8").splitlines()] + if not rows: + raise ValueError("Empty sample source") + prepared = [] + identities = set() + for row in rows: + if not isinstance(row, dict) or set(row) != {"task_id", "response", "source", "license"}: + raise ValueError("Imported row requires task_id, response, source and license") + if any(not isinstance(value, str) or not value.strip() for value in row.values()): + raise ValueError("Sample values must be nonempty strings") + task = tasks[row["task_id"]] + record = {"schema_version": "math-sample-v1", "task_id": task["task_id"], "problem_id": task["problem_id"], + "split": task["split"], "prompt": task["prompt"], "response": row["response"], + "source_kind": "curated", "evidence_kind": "declared", "status": "collected", + "source": row["source"], "license": row["license"], "source_digest": digest(raw.hex()), + "task_digest": digest(task), "parent_digest": None, "token_evidence": None} + record["sample_id"] = digest(record) + if record["sample_id"] in identities: + raise ValueError("Duplicate imported sample") + identities.add(record["sample_id"]) + prepared.append((record, verify(row["response"], task["reference"]))) + root.mkdir(parents=True) + for sample, verification in prepared: + publish(root / "samples" / f"{sample['sample_id']}.json", sample) + publish(root / "verification" / f"{sample['sample_id']}.json", {**verification, "sample_digest": digest(sample)}) + manifest = {"schema_version": "math-dataset-v1", "tasks_digest": digest(list(tasks.values())), + "samples": [{"sample_id": s["sample_id"], "digest": digest(s), "verification_digest": digest({**v, "sample_digest": digest(s)})} + for s, v in prepared], "training_eligibility": "not_an_on_policy_batch"} + publish(root / "manifest.json", manifest) + return manifest diff --git a/harborrl/distill/models.py b/harborrl/distill/models.py new file mode 100644 index 000000000..954ae0df8 --- /dev/null +++ b/harborrl/distill/models.py @@ -0,0 +1,125 @@ +"""Small exact-token Transformers adapter; imported only by training/evaluation.""" +import hashlib + +import torch +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + Qwen2Config, + Qwen2ForCausalLM, +) + +from harborrl.trajectories.native import digest + + +class ByteTokenizer: + """Test-only byte tokenizer with explicit BOS/EOS/PAD and no hidden template.""" + pad_token_id, bos_token_id, eos_token_id = 0, 1, 2 + + def encode(self, text): + return [1, *[value + 3 for value in text.encode("utf-8")]] + + def decode(self, tokens): + return bytes(token - 3 for token in tokens if token >= 3).decode("utf-8", errors="replace") + + def identity(self): + return {"tokenizer_digest": digest("test-byte-v1"), "template_digest": digest("literal-prompt-with-bos-v1"), + "vocabulary_digest": digest(list(range(259))), "vocab_size": 259} + + +class LocalTokenizer: + def __init__(self, checkpoint): + self.tokenizer = AutoTokenizer.from_pretrained(checkpoint, local_files_only=True, trust_remote_code=False) + self.eos_token_id = self.tokenizer.eos_token_id + self.pad_token_id = self.tokenizer.pad_token_id + if self.pad_token_id is None: + self.pad_token_id = self.eos_token_id + if self.eos_token_id is None or self.pad_token_id is None: + raise ValueError("Tokenizer requires explicit EOS and PAD IDs") + + def encode(self, text): + return self.tokenizer.encode(text, add_special_tokens=True) + + def decode(self, tokens): + return self.tokenizer.decode(tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False) + + def identity(self): + tok = self.tokenizer + if not tok.is_fast: + raise ValueError("Exact identity requires a serializable fast tokenizer") + return {"tokenizer_digest": digest(tok.backend_tokenizer.to_str()), + "template_digest": digest({"serialization": "literal-prompt-add-special-tokens-v1", + "chat_template": tok.chat_template, "special_tokens": tok.special_tokens_map}), + "vocabulary_digest": digest(sorted(tok.get_vocab().items())), "vocab_size": len(tok)} + + +def load_student(config, device, *, teacher=False): + spec = config["teacher"] if teacher else config["student"] + torch.manual_seed(config["training"]["seed"] + int(teacher)) + if spec["kind"] == "tiny_qwen": + dimensions = config["student"] + model = Qwen2ForCausalLM(Qwen2Config(vocab_size=259, hidden_size=dimensions["hidden_size"], + intermediate_size=dimensions["hidden_size"] * 2, num_hidden_layers=dimensions["layers"], + num_attention_heads=dimensions["heads"], num_key_value_heads=dimensions["heads"], + max_position_embeddings=config["sampling"]["max_context"], attention_dropout=0.0, + bos_token_id=1, eos_token_id=2, pad_token_id=0, use_sliding_window=False, + _attn_implementation="eager")) + tokenizer = ByteTokenizer() + else: + tokenizer = LocalTokenizer(spec["checkpoint"]) + model = AutoModelForCausalLM.from_pretrained(spec["checkpoint"], local_files_only=True, + trust_remote_code=False, torch_dtype=torch.float32, + attn_implementation="eager") + if model.config.vocab_size != tokenizer.identity()["vocab_size"]: + raise ValueError("Model vocabulary and tokenizer must match exactly (no padded-vocab inference)") + if config["sampling"]["max_context"] > model.config.max_position_embeddings: + raise ValueError("Configured context exceeds model capacity") + model.to(device).eval() + if teacher: + model.requires_grad_(False) + return model, tokenizer + + +def parameter_digest(model): + value = hashlib.sha256() + for name, tensor in model.state_dict().items(): + value.update(name.encode()) + value.update(tensor.detach().cpu().contiguous().numpy().tobytes()) + return value.hexdigest() + + +def response_logprobs(model, prompts, responses, pad_id): + """Causal shift: response[0] is scored by the final prompt position.""" + device = next(model.parameters()).device + sequences = [p + r for p, r in zip(prompts, responses, strict=True)] + maximum = max(map(len, sequences)) + ids = torch.full((len(sequences), maximum), pad_id, dtype=torch.long, device=device) + attention = torch.zeros_like(ids) + for index, tokens in enumerate(sequences): + ids[index, :len(tokens)] = torch.tensor(tokens, device=device) + attention[index, :len(tokens)] = 1 + logits = model(input_ids=ids, attention_mask=attention, use_cache=False).logits.float() + results = [] + for index, (prompt, response) in enumerate(zip(prompts, responses, strict=True)): + logprobs = logits[index, len(prompt) - 1:len(prompt) + len(response) - 1].log_softmax(-1) + results.append(logprobs.gather(-1, torch.tensor(response, device=device).unsqueeze(-1)).squeeze(-1)) + return results + + +@torch.no_grad() +def generate(model, tokenizer, prompt_ids, maximum, context, seed): + if not prompt_ids or len(prompt_ids) + maximum > context: + raise ValueError("Prompt plus response budget exceeds context; no implicit truncation") + device = next(model.parameters()).device + generator = torch.Generator(device=device).manual_seed(seed) + tokens, probabilities = [], [] + for _ in range(maximum): + ids = torch.tensor([prompt_ids + tokens], device=device) + logits = model(input_ids=ids, attention_mask=torch.ones_like(ids), use_cache=False).logits[0, -1].float() + logprobs = logits.log_softmax(-1) + token = int(torch.multinomial(logprobs.exp(), 1, generator=generator)) + tokens.append(token) + probabilities.append(float(logprobs[token])) + if token == tokenizer.eos_token_id: + return tokens, probabilities, "eos" + return tokens, probabilities, "length" diff --git a/harborrl/distill/runner.py b/harborrl/distill/runner.py new file mode 100644 index 000000000..89cb48d1a --- /dev/null +++ b/harborrl/distill/runner.py @@ -0,0 +1,241 @@ +"""Bounded reference trainer: raw student rollout, scoring, one optimizer update.""" +import argparse +import json +import os +import time +from copy import deepcopy +from pathlib import Path + +import torch +import torch.distributed as dist + +from harborrl.algorithms.opd import trajectory_surrogate +from harborrl.trajectories.native import digest, publish + +from .config import load +from .contract import score_request, validate_row +from .data import load_tasks, verify +from .models import generate, load_student, parameter_digest, response_logprobs +from .scoring import FakeTeacher, ModelTeacher, probe_teacher, process_rewards + + +def training_signature(config, tasks_digest): + value = deepcopy(config) + for key in ("steps", "microbatch_size", "save_interval"): + value["training"].pop(key) + value.pop("output") + value.pop("deployment") + return digest({"config": value, "tasks_digest": tasks_digest}) + + +def atomic_checkpoint(path, value): + temporary = path.with_suffix(".pending") + torch.save(value, temporary) + temporary.replace(path) + + +def pad(values, length, device): + return torch.stack([torch.cat((value, value.new_zeros(length - len(value)))) if isinstance(value, torch.Tensor) + else torch.tensor([*value, *([0] * (length - len(value)))], device=device) for value in values]) + + +def optimize(model, optimizer, rows, config, tokenizer, world): + """Equal-size rank shards; mean per trajectory, then mean over global batch.""" + device = next(model.parameters()).device + local_count = len(rows) + optimizer.zero_grad(set_to_none=True) + loss_sum, maximum_error = 0.0, 0.0 + micro = config["training"]["microbatch_size"] + if config.get('consistency', {}).get('mode', 'off') != 'off': + from harborrl.consistency.reference import check_reference + check_reference(model, rows, config, tokenizer) + for offset in range(0, local_count, micro): + chunk = rows[offset:offset + micro] + current_rows = response_logprobs(model, [r["prompt_ids"] for r in chunk], [r["response_ids"] for r in chunk], tokenizer.pad_token_id) + maximum = max(map(len, current_rows)) + current = pad(current_rows, maximum, device) + behavior = pad([r["behavior_logprobs"] for r in chunk], maximum, device) + mask = pad([r["loss_mask"] for r in chunk], maximum, device) + error = float(((current.detach() - behavior).abs() * mask).max()) + maximum_error = max(maximum_error, error) + if error > config["training"]["behavior_tolerance"]: + raise ValueError(f"Actor no longer matches frozen behavior: {error}") + advantage = torch.zeros_like(current) + if config["method"] in {"opd", "prm_opd"}: + teacher = pad([r["teacher_logprobs"] for r in chunk], maximum, device) + advantage += config["training"]["opd_weight"] * (teacher - behavior) + if config["method"] in {"prm", "prm_opd"}: + advantage += config["prm"]["weight"] * pad([r["prm"]["token_returns"] for r in chunk], maximum, device) + rewards = torch.tensor([r["verification"]["reward"] for r in chunk], device=device).unsqueeze(-1) + advantage += config["training"]["outcome_weight"] * rewards + loss = trajectory_surrogate(current, advantage, mask) * (len(chunk) / local_count) + loss.backward() + loss_sum += float(loss.detach()) + if world > 1: + for parameter in model.parameters(): + if parameter.grad is not None: + dist.all_reduce(parameter.grad) + parameter.grad.div_(world) + norm = torch.nn.utils.clip_grad_norm_(model.parameters(), config["training"]["max_grad_norm"], error_if_nonfinite=True) + optimizer.step() + return {"loss": loss_sum, "gradient_norm": float(norm), "max_behavior_error": maximum_error} + + +def train(config, *, allow_synthetic=False, resume=None): + torch.set_num_threads(1) + rank, world, local_rank = (int(os.environ.get(key, default)) for key, default in + (("RANK", "0"), ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0"))) + if world != config["deployment"]["world_size"]: + raise ValueError("Launch world size differs from the declared data-parallel layout") + cuda = config["deployment"]["device"] == "cuda" + device = torch.device(f"cuda:{local_rank}" if cuda else "cpu") + if cuda: + torch.cuda.set_device(device) + if world > 1: + from datetime import timedelta + dist.init_process_group("nccl" if cuda else "gloo", timeout=timedelta(seconds=120)) + root = Path(config["output"]["directory"]) + try: + tasks = load_tasks(config["tasks"]["path"]) + selected = [task for task in tasks if task["split"] == "train"] + if not selected: + raise ValueError("No train tasks") + signature = training_signature(config, digest(tasks)) + model, tokenizer = load_student(config, device) + identity = tokenizer.identity() + optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"], weight_decay=0) + start = 0 + if resume: + checkpoint = torch.load(resume, map_location=device, weights_only=True) + if checkpoint["signature"] != signature or checkpoint["identity"] != identity: + raise ValueError("Checkpoint config/data/tokenizer identity mismatch") + model.load_state_dict(checkpoint["model"]) + optimizer.load_state_dict(checkpoint["optimizer"]) + start = checkpoint["step"] + if parameter_digest(model) != checkpoint["parameter_digest"]: + raise ValueError("Checkpoint parameter digest mismatch") + if start >= config["training"]["steps"]: + raise ValueError("training.steps must exceed resumed completed steps") + initial = parameter_digest(model) + teacher = None + if config["teacher"]["kind"] != "none": + teacher = (FakeTeacher(identity, config["sampling"]["max_context"]) if config["teacher"]["kind"] == "fake" + else ModelTeacher(config, device, identity)) + capabilities = probe_teacher(teacher, tokenizer) + if rank == 0: + publish(root / "teacher-probe.json", capabilities) + if rank == 0: + publish(root / "runtime.json", {"torch": str(torch.__version__), "transformers": __import__("transformers").__version__, + "device": str(device), "world_size": world, "identity": identity, "parameter_count": sum(p.numel() for p in model.parameters()), + "initial_parameter_digest": initial, "resume": str(resume) if resume else None, + "synthetic_allowed": allow_synthetic, "method": config["method"], "signature": signature, + "teacher_evidence": teacher.evidence_kind if teacher else None, "tasks_digest": digest(tasks)}) + metrics = [] + local_count = config["training"]["batch_size"] // world + for step in range(start, config["training"]["steps"]): + begin = time.monotonic() + version = parameter_digest(model) + rows = [] + for index in range(rank * local_count, (rank + 1) * local_count): + task = selected[(step * config["training"]["batch_size"] + index) % len(selected)] + prompt = tokenizer.encode(task["prompt"]) + tokens, behavior, stop = generate(model, tokenizer, prompt, config["sampling"]["max_new_tokens"], + config["sampling"]["max_context"], config["training"]["seed"] + step * config["training"]["batch_size"] + index) + text = tokenizer.decode(tokens) + teacher_scores = teacher.score_tokens(prompt, tokens, identity) if teacher else [] + row = {"schema_version": "math-onpolicy-batch-v1", "sample_id": f"{step:06d}-{index:06d}", + "task_id": task["task_id"], "task_digest": digest(task), "split": "train", + "student_behavior_version": version, "student_evidence": "observed", + "teacher_revision": teacher.revision if teacher else None, + "teacher_evidence": teacher.evidence_kind if teacher else None, + "identity": identity, "prompt_ids": prompt, "response_ids": tokens, + "behavior_logprobs": behavior, "teacher_logprobs": teacher_scores, "loss_mask": [1] * len(tokens), + "sampling": {"temperature": 1.0, "top_p": 1.0, "raw_logprobs": True}, "stop_reason": stop, + "response": text, "verification": verify(text, task["reference"]), + "prm": process_rewards(config["prm"]["kind"], tokens, tokenizer) if config["prm"]["kind"] != "none" else None} + row["score_request_digest"] = digest(score_request(row)) + row["score_response_digest"] = digest({"request_digest": row["score_request_digest"], "logprobs": teacher_scores}) + row["digest"] = digest(row) + validate_row(row, version=version, identity=identity, context=config["sampling"]["max_context"], + allow_synthetic=allow_synthetic, require_teacher=teacher is not None) + publish(root / "batches" / f"{row['sample_id']}.json", row) + rows.append(row) + if world > 1: + dist.barrier() + result = optimize(model, optimizer, rows, config, tokenizer, world) + after = parameter_digest(model) + if teacher and not teacher.unchanged(): + raise ValueError("Teacher changed or received gradients") + result.update(step=step + 1, rank=rank, student_behavior_version=version, updated_version=after, + parameter_changed=after != version, teacher_unchanged=True, + tokens=sum(len(r["response_ids"]) for r in rows), seconds=time.monotonic() - begin, + verification_counts={status: sum(r["verification"]["status"] == status for r in rows) + for status in ("correct", "incorrect", "unparseable", "unsupported")}) + publish(root / "metrics" / f"{step + 1:06d}-rank{rank}.json", result) + metrics.append(result) + if rank == 0: + print(json.dumps(result), flush=True) + if (step + 1) % config["training"]["save_interval"] == 0 or step + 1 == config["training"]["steps"]: + atomic_checkpoint(root / f"checkpoint-{step + 1:06d}.pt", {"model": model.state_dict(), + "optimizer": optimizer.state_dict(), "step": step + 1, "identity": identity, + "signature": signature, "parameter_digest": after}) + if world > 1: + versions = [None] * world + dist.all_gather_object(versions, after) + if len(set(versions)) != 1: + raise ValueError("Data-parallel ranks diverged") + if rank == 0: + summary = {"status": "completed", "completed_steps": config["training"]["steps"], "resumed_step": start, + "initial_parameter_digest": initial, "final_parameter_digest": parameter_digest(model), + "parameter_changed": initial != parameter_digest(model), "teacher_unchanged": True, + "world_size": world, "method": config["method"], "capability_improvement": "not_evaluated", + "synthetic_allowed": allow_synthetic, "nonzero_gradient_steps": sum(m["gradient_norm"] > 0 for m in metrics)} + publish(root / "summary.json", summary) + print(json.dumps(summary), flush=True) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + + +def evaluate(config, checkpoint, *, split=None): + torch.set_num_threads(1) + tasks = load_tasks(config["tasks"]["path"]) + model, tokenizer = load_student(config, torch.device(config["deployment"]["device"])) + state = torch.load(checkpoint, map_location=next(model.parameters()).device, weights_only=True) + if state["signature"] != training_signature(config, digest(tasks)) or state["identity"] != tokenizer.identity(): + raise ValueError("Evaluation checkpoint config/data/tokenizer mismatch") + model.load_state_dict(state["model"]) + if parameter_digest(model) != state["parameter_digest"]: + raise ValueError("Checkpoint digest mismatch") + selected = [task for task in tasks if task["split"] == (split or config["tasks"]["eval_split"])] + if not selected: + raise ValueError("Empty held-out split") + rows = [] + for index, task in enumerate(selected): + tokens, _, _ = generate(model, tokenizer, tokenizer.encode(task["prompt"]), config["sampling"]["max_new_tokens"], + config["sampling"]["max_context"], config["training"]["seed"] + index) + response = tokenizer.decode(tokens) + rows.append({"task_id": task["task_id"], "response": response, "verification": verify(response, task["reference"])}) + return {"schema_version": "math-evaluation-v1", "checkpoint_digest": state["parameter_digest"], + "split": split or config["tasks"]["eval_split"], "samples": rows, + "accuracy": sum(r["verification"]["status"] == "correct" for r in rows) / len(rows), + "scope": "pipeline smoke; not evidence of capability improvement"} + + +def main(argv=None): + parser = argparse.ArgumentParser() + parser.add_argument("--config", required=True) + parser.add_argument("--allow-synthetic", action="store_true") + parser.add_argument("--resume") + args = parser.parse_args(argv) + config = load(args.config, allow_synthetic=args.allow_synthetic) + try: + train(config, allow_synthetic=args.allow_synthetic, resume=args.resume) + except BaseException as exc: + publish(Path(config["output"]["directory"]) / f"failure-rank{os.environ.get('RANK', '0')}.json", + {"status": "failed", "type": type(exc).__name__, "message": str(exc)}) + raise + + +if __name__ == "__main__": + main() diff --git a/harborrl/distill/scoring.py b/harborrl/distill/scoring.py new file mode 100644 index 000000000..d92f2bb56 --- /dev/null +++ b/harborrl/distill/scoring.py @@ -0,0 +1,123 @@ +"""Exact-prefix teacher and process-reward adapters, with explicit fake evidence.""" +import operator +import re +from fractions import Fraction + +from harborrl.distill.data import NUMBER +from harborrl.trajectories.native import digest + + +class FakeTeacher: + """Known normalized test distribution; never presented as model inference.""" + revision = digest("fake-prefix-categorical-v1") + evidence_kind = "synthetic" + + def __init__(self, identity, context): + self.identity, self.context = identity, context + + def capabilities(self): + return {"exact_token_scoring": True, "raw_logprob_semantics": True, "identity": self.identity, + "max_context": self.context, "revision": self.revision, "evidence_kind": self.evidence_kind} + + def score_tokens(self, prompt, response, identity): + import torch + if identity != self.identity or len(prompt) + len(response) > self.context: + raise ValueError("Fake scorer identity/context mismatch") + logits = torch.linspace(-1, 1, identity["vocab_size"], dtype=torch.float64) + base = logits.log_softmax(0) + prefix, result = list(prompt), [] + for token in response: + shift = sum(prefix) % identity["vocab_size"] + result.append(float(base[(token - shift) % identity["vocab_size"]])) + prefix.append(token) + return result + + def unchanged(self): + return True + + +class ModelTeacher: + evidence_kind = "observed" + + def __init__(self, config, device, identity): + from .models import load_student, parameter_digest + self.model, self.tokenizer = load_student(config, device, teacher=True) + self.identity = self.tokenizer.identity() + if self.identity != identity: + raise ValueError("Teacher and student tokenizer/template/vocabulary must match exactly") + self.context = config["sampling"]["max_context"] + self.revision = parameter_digest(self.model) + + def capabilities(self): + return {"exact_token_scoring": True, "raw_logprob_semantics": True, "identity": self.identity, + "max_context": self.context, "revision": self.revision, "evidence_kind": self.evidence_kind} + + def score_tokens(self, prompt, response, identity): + import torch + + from .models import response_logprobs + if identity != self.identity or len(prompt) + len(response) > self.context: + raise ValueError("Teacher identity/context mismatch") + with torch.no_grad(): + return response_logprobs(self.model, [prompt], [response], self.tokenizer.pad_token_id)[0].cpu().tolist() + + def unchanged(self): + from .models import parameter_digest + return all(parameter.grad is None for parameter in self.model.parameters()) and parameter_digest(self.model) == self.revision + + +def probe_teacher(teacher, tokenizer): + """Exercise the actual adapter at the first-token boundary, EOS and longer prefixes.""" + import math + capabilities = teacher.capabilities() + prompt = tokenizer.encode("1 + 1 =") + response = tokenizer.encode("2\n")[1:] + [tokenizer.eos_token_id] + for tokens in ([tokenizer.eos_token_id], response): + scores = teacher.score_tokens(prompt, tokens, tokenizer.identity()) + if len(scores) != len(tokens) or any(not math.isfinite(v) or v > 0 for v in scores): + raise ValueError("Teacher scoring probe failed") + for index in range(len(tokens)): + separate = teacher.score_tokens(prompt + tokens[:index], [tokens[index]], tokenizer.identity())[0] + if abs(separate - scores[index]) > 1e-4: + raise ValueError("Teacher causal-prefix probe failed") + return {**capabilities, "probe": "passed", "request_digest": digest({"prompt": prompt, "response": response})} + + +def process_rewards(kind, response_ids, tokenizer): + """Each newline terminates a step; a boundary token belongs to that step.""" + if kind not in {"fake", "equation"} or not response_ids or len(response_ids) > 16384: + raise ValueError("Invalid or oversized process-reward request") + ends, previous = [], 0 + for index in range(len(response_ids)): + count = tokenizer.decode(response_ids[:index + 1]).count("\n") + if count > previous: + ends.append(index + 1) + previous = count + if not ends or ends[-1] != len(response_ids): + ends.append(len(response_ids)) + steps, start = [], 0 + for end in ends: + text = tokenizer.decode(response_ids[start:end]).strip() + if kind == "fake": + score = (int(digest(response_ids[start:end])[:8], 16) % 9 + 1) / 10 + status = "synthetic" + else: + match = re.fullmatch(rf"\s*({NUMBER})\s*([+*/-])\s*({NUMBER})\s*=\s*({NUMBER})\s*", text) + score, status = 0.0, "unsupported" + if match and len(text) <= 1024: + a, b, expected = Fraction(match[1]), Fraction(match[3]), Fraction(match[4]) + try: + actual = {"+": operator.add, "-": operator.sub, "*": operator.mul, "/": operator.truediv}[match[2]](a, b) + score = 1.0 if actual == expected else -1.0 + status = "correct" if actual == expected else "incorrect" + except ZeroDivisionError: + status = "unsupported" + steps.append({"start": start, "end": end, "text": text, "score": score, "status": status}) + start = end + values = [0.0] * len(response_ids) + cumulative = 0.0 + for step in reversed(steps): + cumulative += step["score"] + values[step["start"]:step["end"]] = [cumulative] * (step["end"] - step["start"]) + return {"adapter": f"{kind}-prm-v1", "evidence_kind": "synthetic" if kind == "fake" else "observed", + "steps": steps, "token_returns": values, "aggregation": "undiscounted-step-reward-to-go"} diff --git a/harborrl/export/contract.py b/harborrl/export/contract.py index d64ab8829..3b627ec24 100644 --- a/harborrl/export/contract.py +++ b/harborrl/export/contract.py @@ -10,6 +10,8 @@ SCHEMA_VERSION = "harborrl-training-batch-v1" REDUCTION = "trajectory_weighted_clipped_pg_v1" +OPD_REDUCTION = "trajectory_weighted_sampled_opd_v1" +DAPO_REDUCTION = "global_token_mean_asymmetric_clipped_pg_v1" def _as_tuple(value, label): @@ -18,6 +20,20 @@ def _as_tuple(value, label): return tuple(value) +@dataclass(frozen=True) +class ProcessRewards: + """Observed PRM probabilities and their exact response-token partition.""" + + identity: str + step_ends: tuple[int, ...] + probabilities: tuple[float, ...] + token_returns: tuple[float, ...] + + def __post_init__(self): + for name in ('step_ends', 'probabilities', 'token_returns'): + object.__setattr__(self, name, _as_tuple(getattr(self, name), name)) + + @dataclass(frozen=True) class TokenSpan: """One policy turn expanded from a Native trajectory.""" @@ -29,16 +45,20 @@ class TokenSpan: loss_mask: tuple[bool, ...] old_logprobs: tuple[float, ...] policy_version: str + teacher_logprobs: tuple[float, ...] | None = None + process_rewards: ProcessRewards | None = None def __post_init__(self): object.__setattr__(self, "tokens", _as_tuple(self.tokens, "tokens")) object.__setattr__(self, "loss_mask", _as_tuple(self.loss_mask, "loss_mask")) object.__setattr__(self, "old_logprobs", _as_tuple(self.old_logprobs, "old_logprobs")) + if self.teacher_logprobs is not None: + object.__setattr__(self, "teacher_logprobs", _as_tuple(self.teacher_logprobs, "teacher_logprobs")) @dataclass(frozen=True) class TrajectoryTrainingUnit: - """A complete trajectory with its group-normalized training semantics.""" + """A complete trajectory; OPD uses token scores and a zero scalar advantage.""" trajectory_id: str rollout_id: int @@ -74,6 +94,9 @@ class TrainingBatch: groups: tuple[TrainingGroup, ...] reduction: str policy_versions: frozenset[str] + algorithm: str = "grpo" + teacher_identity: str | None = None + process_reward_identity: str | None = None def __post_init__(self): object.__setattr__(self, "groups", _as_tuple(self.groups, "groups")) diff --git a/harborrl/export/native.py b/harborrl/export/native.py index 0062e9e2d..96519a3f6 100644 --- a/harborrl/export/native.py +++ b/harborrl/export/native.py @@ -16,7 +16,7 @@ def group_statistics(rewards, *, epsilon=1e-6): - """Population mean/std and GRPO advantages for one complete group.""" + """Population statistics and centered GRPO advantages without std scaling.""" values = list(rewards) if len(values) < 2: raise ValueError("a complete group of at least two trajectories is required") @@ -28,7 +28,7 @@ def group_statistics(rewards, *, epsilon=1e-6): std = math.sqrt(math.fsum((value - mean) ** 2 for value in values) / len(values)) if not finite(mean) or not finite(std): raise ValueError("nonfinite group statistics") - advantages = [0.0 if std == 0 else (value - mean) / (std + epsilon) for value in values] + advantages = [value - mean for value in values] return mean, std, advantages diff --git a/harborrl/export/validate.py b/harborrl/export/validate.py index c9519bd61..1de575fc4 100644 --- a/harborrl/export/validate.py +++ b/harborrl/export/validate.py @@ -4,6 +4,8 @@ import math from harborrl.export.contract import ( + DAPO_REDUCTION, + OPD_REDUCTION, REDUCTION, SCHEMA_VERSION, TokenSpan, @@ -58,6 +60,16 @@ def validate_token_span(span: TokenSpan): if any(not finite(value) or value > 0 for value in span.old_logprobs): raise ValueError("old logprobs must be finite and nonpositive") _nonempty_str(span.policy_version, "policy_version") + if span.process_rewards is not None: + from harborrl.export.contract import ProcessRewards + from harborrl.math.prm import token_returns + scores = span.process_rewards + if not isinstance(scores, ProcessRewards): + raise ValueError('process_rewards must be a ProcessRewards contract') + _nonempty_str(scores.identity, 'process reward identity') + expected = token_returns(scores.step_ends, scores.probabilities, span.response_length) + if tuple(expected) != scores.token_returns: + raise ValueError('PRM token returns differ from step reward-to-go') def validate_trajectory(unit: TrajectoryTrainingUnit): @@ -83,7 +95,9 @@ def validate_trajectory(unit: TrajectoryTrainingUnit): raise ValueError("trajectory token count differs from its turns") -def validate_group(group: TrainingGroup, *, epsilon=1e-6): +def validate_group(group: TrainingGroup, *, epsilon=1e-6, algorithm="grpo"): + if algorithm not in ('grpo', 'opd', 'dapo'): + raise ValueError('unsupported group algorithm') _nonempty_str(group.group_key, "group_key") _plain_int(group.group_size, "group_size") if group.group_size < 2 or len(group.trajectories) != group.group_size: @@ -105,10 +119,33 @@ def validate_group(group: TrainingGroup, *, epsilon=1e-6): raise ValueError("group reward mean differs from its trajectories") if not math.isclose(std, group.reward_std, rel_tol=1e-9, abs_tol=1e-9): raise ValueError("group reward std differs from its trajectories") - if group.reward_std == 0 and any(unit.advantage != 0.0 for unit in group.trajectories): + if algorithm == "grpo" and group.reward_std == 0 and any(unit.advantage != 0.0 for unit in group.trajectories): raise ValueError("zero-std groups must produce zero advantages") for unit in group.trajectories: validate_trajectory(unit) + if algorithm in ('grpo', 'dapo'): + expected = unit.reward - mean + if algorithm == 'dapo': + from harborrl.math.dapo import advantage + if std == 0: + raise ValueError('DAPO requires a nonzero reward std') + expected = advantage(unit.reward, mean, std, group.group_size) + if not math.isclose(unit.advantage, expected, rel_tol=1e-9, abs_tol=1e-9): + if algorithm == 'dapo': + raise ValueError('DAPO advantage must use complete-group sample std plus 1e-6') + raise ValueError("GRPO advantage must be centered without std normalization") + if any(span.teacher_logprobs is not None for span in unit.turns): + raise ValueError("Outcome policy gradients must not consume teacher scores") + else: + if unit.advantage != 0.0: + raise ValueError("OPD requires token scores, not a scalar reward advantage") + for span in unit.turns: + values = span.teacher_logprobs + if (values is None or len(values) != span.response_length + or any(not finite(value) or value > 0 for value in values)): + raise ValueError("OPD requires aligned finite teacher token logprobs") + if algorithm == 'dapo': + continue weight_sum = math.fsum( token_weight(group.group_size, unit.token_count) for span in unit.turns for _ in range(span.response_length) ) @@ -119,8 +156,17 @@ def validate_group(group: TrainingGroup, *, epsilon=1e-6): def validate_batch(batch: TrainingBatch, *, epsilon=1e-6): if batch.schema_version != SCHEMA_VERSION: raise ValueError(f"unsupported training batch schema: {batch.schema_version}") - if batch.reduction != REDUCTION: + reductions = {"grpo": REDUCTION, "opd": OPD_REDUCTION, "dapo": DAPO_REDUCTION} + if batch.algorithm not in reductions or batch.reduction != reductions[batch.algorithm]: raise ValueError(f"unsupported training reduction: {batch.reduction}") + if batch.algorithm == "opd": + _nonempty_str(batch.teacher_identity, "teacher_identity") + elif batch.teacher_identity is not None: + raise ValueError("Only OPD may declare a teacher") + if batch.process_reward_identity is not None: + _nonempty_str(batch.process_reward_identity, 'process_reward_identity') + if batch.algorithm != 'grpo': + raise ValueError('Process rewards currently require GRPO') if not batch.groups: raise ValueError("native training batch requires at least one group") if len(batch.policy_versions) != 1: @@ -132,7 +178,16 @@ def validate_batch(batch: TrainingBatch, *, epsilon=1e-6): if len(set(rollout_ids)) != len(rollout_ids): raise ValueError("each native group needs a distinct rollout id") for group in batch.groups: - validate_group(group, epsilon=epsilon) + validate_group(group, epsilon=epsilon, algorithm=batch.algorithm) + for unit in group.trajectories: + scores = [span.process_rewards for span in unit.turns] + if batch.process_reward_identity is None: + if any(value is not None for value in scores): + raise ValueError('PRM evidence requires a batch process_reward_identity') + elif (len(scores) != 1 or scores[0] is None + or scores[0].identity != batch.process_reward_identity + or not math.isclose(unit.reward, scores[0].token_returns[0], rel_tol=1e-9, abs_tol=1e-9)): + raise ValueError('PRM batch requires one aligned process reward source and matching total reward') versions = {span.policy_version for unit in group.trajectories for span in unit.turns} if versions != batch.policy_versions: raise ValueError("training batch mixes policy versions") @@ -140,6 +195,9 @@ def validate_batch(batch: TrainingBatch, *, epsilon=1e-6): def batch_weight_sum(batch: TrainingBatch): """Sum of trajectory weights; each complete group contributes exactly one.""" + if batch.algorithm == 'dapo': + validate_batch(batch) + return 1.0 return math.fsum( token_weight(group.group_size, unit.token_count) for group in batch.groups diff --git a/harborrl/gateway/messages.py b/harborrl/gateway/messages.py index 08578c2a3..3c3d4dadd 100644 --- a/harborrl/gateway/messages.py +++ b/harborrl/gateway/messages.py @@ -9,6 +9,11 @@ class ProtocolError(ValueError): pass +def unavailable_tool_error(name): + """Exact native Claude CLI observation for an unavailable tool.""" + return f"Error: No such tool available: {name}" + + def keys(value, allowed, label): if not isinstance(value, dict) or set(value) - set(allowed): raise ProtocolError(f"unsupported fields in {label}") @@ -80,7 +85,7 @@ def convert(payload, *, model, max_output_tokens=8192, count_only=False): history = payload.get("messages") if not isinstance(history, list) or not history: raise ProtocolError("nonempty messages required") - pending, seen = set(), set() + pending, seen = {}, set() for item in history: keys(item, ("role", "content"), "message") role = item.get("role") @@ -101,10 +106,12 @@ def convert(payload, *, model, max_output_tokens=8192, count_only=False): elif kind == "tool_use" and role == "assistant": keys(block, ("type", "id", "name", "input", "cache_control"), "tool_use") call_id = block.get("id") - if not isinstance(call_id, str) or not call_id or call_id in seen or block.get("name") not in names or not isinstance(block.get("input"), dict): + name = block.get("name") + if (not isinstance(call_id, str) or not call_id or call_id in seen + or not isinstance(name, str) or not name or not isinstance(block.get("input"), dict)): raise ProtocolError("invalid tool_use identity or input") seen.add(call_id) - pending.add(call_id) + pending[call_id] = name calls.append({"id": call_id, "type": "function", "function": { "name": block["name"], "arguments": json.dumps(block["input"], ensure_ascii=False, allow_nan=False)}}) elif kind == "tool_result" and role == "user": @@ -112,9 +119,12 @@ def convert(payload, *, model, max_output_tokens=8192, count_only=False): call_id = block.get("tool_use_id") if call_id not in pending or type(block.get("is_error", False)) is not bool: raise ProtocolError("orphan or duplicate tool_result") - pending.remove(call_id) + name = pending.pop(call_id) # Preserve error flag as part of model-visible tool observation. content = text_blocks(block.get("content", "")) + if name not in names and (block.get("is_error") is not True + or content != unavailable_tool_error(name)): + raise ProtocolError("undeclared tool requires its exact native rejection") if block.get("is_error"): content = "[tool_error]\n" + content results.append({"role": "tool", "tool_call_id": call_id, "content": content}) @@ -134,7 +144,7 @@ def convert(payload, *, model, max_output_tokens=8192, count_only=False): "max_new_tokens": max_tokens, "stop": stop})} -def parse_qwen_output(text, *, response_id, tools): +def parse_qwen_output(text, *, response_id, tools, allow_unknown_tools=False): """Deterministic Qwen tool markers; malformed model output is never repaired.""" if "" in text or "" in text: raise ProtocolError("thinking output is not enabled by this profile") @@ -164,8 +174,10 @@ def parse_qwen_output(text, *, response_id, tools): raise ProtocolError("finish marker must be the final output") return blocks canonical_name = names.get(generated_name.casefold()) if isinstance(generated_name, str) else None - if canonical_name is None: + if canonical_name is None and (not allow_unknown_tools or not isinstance(generated_name, str) or not generated_name): raise ProtocolError(f"unknown generated tool: {generated_name!r}") + if canonical_name is None: + canonical_name = generated_name # Preserve the invalid action for native CLI rejection. if not isinstance(call.get("arguments"), dict): raise ProtocolError("generated tool arguments are not an object") blocks.append({"type": "tool_use", "id": f"toolu_{response_id}_{index}", "name": canonical_name, "input": call["arguments"]}) diff --git a/harborrl/gateway/server.py b/harborrl/gateway/server.py index fd011c505..db07a3ed2 100644 --- a/harborrl/gateway/server.py +++ b/harborrl/gateway/server.py @@ -1,6 +1,7 @@ """Bounded stdlib HTTP transport. All policy requests require attempt credentials.""" import json import secrets +import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import urlsplit @@ -46,7 +47,17 @@ def count_tokens(self, token, payload): return {"input_tokens": self.backend.count_tokens(converted)} -def make_server(gateway, host="127.0.0.1", port=0, *, max_body_bytes=8 * 1024 * 1024): +def make_server(gateway, host="127.0.0.1", port=0, *, max_body_bytes=8 * 1024 * 1024, + max_inflight=None): + """Serve one gateway. ``max_inflight`` bounds concurrent policy requests. + + Beyond the bound the server answers 503 before any attempt state is + touched, so an overloaded gateway never poisons a seal: no generation + began and a client retry is safe. ``None`` keeps the historic unbounded + behavior for tests and embedded single-attempt use. + """ + inflight = threading.BoundedSemaphore(max_inflight) if max_inflight else None + class Handler(BaseHTTPRequestHandler): def setup(self): super().setup() @@ -69,6 +80,16 @@ def reply(self, status, payload): self.wfile.write(body) def do_POST(self): + if inflight is not None and not inflight.acquire(blocking=False): + return self.reply(503, {"type": "error", "error": {"type": "overloaded_error", + "message": "gateway in-flight request limit reached"}}) + try: + return self._handle_post() + finally: + if inflight is not None: + inflight.release() + + def _handle_post(self): aid = message = None try: path = urlsplit(self.path).path @@ -112,6 +133,10 @@ def do_POST(self): # Do not echo backend errors or request contents (credentials may appear there). self.reply(status, {"type": "error", "error": {"type": "authentication_error" if status == 401 else "invalid_request_error" if status == 400 else "api_error", "message": f"{type(exc).__name__}: {exc}"}}) - server = ThreadingHTTPServer((host, port), Handler) - server.daemon_threads = True + class Server(ThreadingHTTPServer): + # One thread per connection; the kernel backlog bounds queued peers. + daemon_threads = True + request_queue_size = 128 + + server = Server((host, port), Handler) return server diff --git a/harborrl/gateway/sglang_backend.py b/harborrl/gateway/sglang_backend.py index d2efe5e63..0471ef8f2 100644 --- a/harborrl/gateway/sglang_backend.py +++ b/harborrl/gateway/sglang_backend.py @@ -15,13 +15,16 @@ class SGLangBackend: def __init__(self, endpoint, tokenizer, model_adapter, *, tokenizer_digest, template_digest, timeout=120, max_context=32768, - audited_raw_logprobs=False): + audited_raw_logprobs=False, sampling_seed=None): + if sampling_seed is not None and (type(sampling_seed) is not int or not 0 <= sampling_seed < 2**31): + raise ValueError('sampling_seed must be null or an integer in 0..2147483647') self.endpoint = endpoint.rstrip("/") self.tokenizer = tokenizer self.model_adapter = model_adapter self.tokenizer_digest, self.template_digest = tokenizer_digest, template_digest self.timeout, self.max_context = timeout, max_context self.audited_raw_logprobs = audited_raw_logprobs + self.sampling_seed = sampling_seed self.opener = request.build_opener(request.ProxyHandler({})) def prepare(self, converted): @@ -65,7 +68,14 @@ def pool_states(self): def generate(self, converted, identity, response_id): ids = self.prepare(converted) - payload = {"input_ids": ids, "sampling_params": converted["sampling"], + sampling = dict(converted['sampling']) + if self.sampling_seed is not None: + sampling['sampling_seed'] = int(digest({ + 'seed': self.sampling_seed, 'batch': identity.batch_id, + 'group': identity.group_id, 'slot': identity.slot_id, + 'task': identity.task_digest, 'input_ids': ids, + }), 16) % 2**31 + payload = {"input_ids": ids, "sampling_params": sampling, "return_logprob": True, "stream": False} req = request.Request(self.endpoint + "/generate", data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}) @@ -91,18 +101,23 @@ def decode(self, output, payload, converted, identity, response_id): if reported_ids is not None and reported_ids != output_ids: raise ValueError("serving output token IDs disagree with logprob evidence") decoded = self.tokenizer.decode(output_ids, skip_special_tokens=True) - if not decoded.startswith(text) and not text.startswith(decoded.rstrip("\n")): - raise ValueError("serving text disagrees with output token IDs") finish = meta.get("finish_reason", {}) - if finish.get("type") not in ("stop", "length"): + if not isinstance(finish, dict) or finish.get("type") not in ("stop", "length"): raise ValueError("unsupported serving finish reason") + # SGLang may omit the explicitly matched terminal stop string from text. + # Nothing else may be dropped, appended, or whitespace-normalized. + matched = finish.get("matched") + stop_sequence = matched if (finish["type"] == "stop" and isinstance(matched, str) + and matched in converted["sampling"]["stop"]) else None + if decoded != text and not (stop_sequence and decoded == text + stop_sequence): + raise ValueError("serving text disagrees with output token IDs") blocks = self.model_adapter.parse_output( text, response_id=response_id, tools=converted["tools"] ) has_tools = any(b["type"] == "tool_use" for b in blocks) - stop_sequence = finish.get("matched") if isinstance(finish.get("matched"), str) else None - if stop_sequence not in converted["sampling"]["stop"]: - stop_sequence = None + declared = {tool['function']['name'] for tool in converted['tools']} + unavailable = [{'id': b['id'], 'name': b['name']} for b in blocks + if b['type'] == 'tool_use' and b['name'] not in declared] eos_id = getattr(self.tokenizer, "eos_token_id", None) eos = finish["type"] == "stop" and finish.get("matched") == eos_id reason = "max_tokens" if finish["type"] == "length" else ( @@ -114,4 +129,5 @@ def decode(self, output, payload, converted, identity, response_id): "serving_input": payload, "evidence_kind": "serving", "content": blocks, "finish_reason": reason, "stop_sequence": stop_sequence, "eos": eos, "sampling": converted["sampling"], + "model_output_text": text, "unavailable_tool_calls": unavailable, "serving_output_digest": digest(output)} diff --git a/harborrl/gateway/shard_service.py b/harborrl/gateway/shard_service.py new file mode 100644 index 000000000..34308162a --- /dev/null +++ b/harborrl/gateway/shard_service.py @@ -0,0 +1,308 @@ +"""The remote gateway shard process: one registry, one gateway, one admin plane. + +Usually launched on another machine that shares the run directory: + + RUN_ID=... RUN_DIR=/mnt/shared-storage-user/.../runs/Run-... \\ + HARBORRL_SHARD_NAME=gw1 \\ + HARBORRL_SHARD_GATEWAY_HOST=0.0.0.0 HARBORRL_SHARD_GATEWAY_PORT=49300 \\ + HARBORRL_SHARD_GATEWAY_URL=http://10.98.250.180:49300 \\ + HARBORRL_SHARD_ADMIN_HOST=0.0.0.0 HARBORRL_SHARD_ADMIN_PORT=49301 \\ + HARBORRL_SHARD_TOKEN_FILE=/mnt/.../run-token \\ + HARBORRL_SHARD_ENDPOINT=http://router:30000 \\ + HARBORRL_NATIVE_PROFILE=/mnt/.../profile.json HF_CKPT=/mnt/.../ckpt \\ + python -m harborrl.gateway.shard_service + +A fixed gateway port requires its declared URL (the spelling must exist +before the process starts); an ephemeral bind — port 0, no URL — instead +reports the resolved gateway and admin ports on stdout as one ready line, +which is how an auto-growing manager launches this service over ssh. + +Traces land in RUN_DIR/native-traces/shard-{name} on the shared storage, so +the GPU side digests them exactly like local-shard traces. The admin plane +carries only manager-side registry operations, run-token authenticated; +generation traffic hits the gateway port directly from runners. +""" +from __future__ import annotations + +import hmac +import json +import signal +import threading +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path + +from harborrl.trajectories.native import Identity + +SHARD_SERVICE = "harborrl-gateway-shard-v1" +ADMIN_MAX_BODY_BYTES = 64 * 1024 * 1024 + + +@dataclass(frozen=True) +class ShardServiceConfig: + """Deployment-independent description of one shard process.""" + + run_id: str + run_root: str + shard_name: str + gateway_host: str + gateway_port: int + gateway_url: str + admin_host: str + admin_port: int + token_file: str + endpoint: str + profile_path: str + model_tag: str + max_output_tokens: int + max_context: int + expected_tokenizer_digest: str + expected_template_digest: str + audited_raw_logprobs: bool + + @property + def trace_dir(self): + # The same spelling the manager-side runtime derives from RUN_DIR. + return Path(self.run_root) / "native-traces" / f"shard-{self.shard_name}" + + @staticmethod + def _required(env, name): + value = env.get(name) + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"gateway shard service requires {name}") + return value + + @classmethod + def from_environment(cls, env): + gateway_port = int(cls._required(env, "HARBORRL_SHARD_GATEWAY_PORT")) + if gateway_port == 0: + # An ephemeral gateway bind reports its port in the ready line; + # whoever launched this process composes and advertises the URL. + gateway_url = "" + else: + gateway_url = cls._required(env, "HARBORRL_SHARD_GATEWAY_URL") + return cls( + run_id=cls._required(env, "RUN_ID"), + run_root=cls._required(env, "RUN_DIR"), + shard_name=cls._required(env, "HARBORRL_SHARD_NAME"), + gateway_host=env.get("HARBORRL_SHARD_GATEWAY_HOST", "0.0.0.0"), + gateway_port=gateway_port, + gateway_url=gateway_url, + admin_host=env.get("HARBORRL_SHARD_ADMIN_HOST", "0.0.0.0"), + admin_port=int(cls._required(env, "HARBORRL_SHARD_ADMIN_PORT")), + token_file=cls._required(env, "HARBORRL_SHARD_TOKEN_FILE"), + endpoint=cls._required(env, "HARBORRL_SHARD_ENDPOINT").rstrip("/"), + profile_path=cls._required(env, "HARBORRL_NATIVE_PROFILE"), + model_tag=env.get("MODEL_TAG", env.get("MODEL_ARGS_FILE", "qwen3-8B")), + max_output_tokens=int(env.get("ROLLOUT_MAX_RESPONSE_LEN", "8192")), + max_context=int(env.get("ROLLOUT_MAX_CONTEXT_LEN", "32768")), + expected_tokenizer_digest=cls._required(env, "HARBORRL_NATIVE_TOKENIZER_DIGEST"), + expected_template_digest=cls._required(env, "HARBORRL_NATIVE_TEMPLATE_DIGEST"), + audited_raw_logprobs=env.get("HARBORRL_NATIVE_AUDITED_LOGPROBS") == "1", + ) + + +def _seal(registry, body): + turns, seal = registry.seal(body["attempt_id"]) + return {"turns": turns, "seal": seal} + + +def _registry_actions(): + """The manager-side registry surface this process serves.""" + return { + "register": lambda registry, body: {"credential": registry.register( + Identity(**body["identity"]), body["trial_id"], + credential=body.get("credential"))}, + "close": lambda registry, body: registry.close(body["attempt_id"]), + "wait_delivered": lambda registry, body: registry.wait_delivered( + body["attempt_id"], float(body["timeout"])), + "unconsumed_responses": lambda registry, body: registry.unconsumed_responses( + body["attempt_id"]), + "consume": lambda registry, body: registry.consume( + body["attempt_id"], body["response_id"], body["content"], + body["evidence_ref"], cli_version=body.get("cli_version")), + "audit_tool_rejections": lambda registry, body: registry.audit_tool_rejections( + body["attempt_id"], body["results"], body["evidence_ref"]), + "seal": _seal, + "close_admission": lambda registry, body: registry.close_admission(), + "drain": lambda registry, body: registry.drain(float(body["timeout"])), + "publish_version": lambda registry, body: registry.publish_version( + body["states"], body["expected"]), + } + + +def make_admin_server(registry, *, shard_name, token, host="127.0.0.1", port=0, + on_shutdown=None): + """Run-token-authenticated admin plane for one shard process. + + Every registry failure crosses the channel as data ({"ok": false, + "error": {type, message}}) so the manager re-raises the exact exception + the local registry would have raised; only transport failures stay + transport errors. + """ + actions = _registry_actions() + + class Handler(BaseHTTPRequestHandler): + def setup(self): + super().setup() + self.connection.settimeout(135) + + def log_message(self, *args): + pass + + def reply(self, status, payload): + body = json.dumps(payload).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def authorized(self): + given = self.headers.get("x-api-key", "") + return bool(given) and hmac.compare_digest(given, token) + + def do_GET(self): + if self.path == "/healthz": + return self.reply(200, {"ok": True, "service": SHARD_SERVICE, + "shard": shard_name}) + if self.path == "/status": + if not self.authorized(): + return self.reply(401, {"error": "invalid shard token"}) + return self.reply(200, {"shard": shard_name, "version": registry.version, + "accepting": registry.accepting}) + self.reply(404, {"error": "unsupported endpoint"}) + + def do_POST(self): + if not self.authorized(): + return self.reply(401, {"error": "invalid shard token"}) + if self.path == "/shutdown": + self.reply(200, {"ok": True, "shard": shard_name}) + if on_shutdown is not None: + threading.Thread(target=on_shutdown, daemon=True).start() + return + if not self.path.startswith("/registry/"): + return self.reply(404, {"error": "unsupported endpoint"}) + action = self.path[len("/registry/"):] + handler = actions.get(action) + if handler is None: + return self.reply(200, {"ok": False, "error": { + "type": "ValueError", + "message": f"unknown registry action {action!r}"}}) + try: + size = int(self.headers.get("Content-Length", "0")) + if size < 0 or size > ADMIN_MAX_BODY_BYTES: + raise ValueError("invalid request body size") + # Payload-free actions (close_admission) post an empty body. + body = json.loads(self.rfile.read(size)) if size else {} + result = handler(registry, body) + except Exception as exc: # noqa: BLE001 — registry failures cross as data + return self.reply(200, {"ok": False, "error": { + "type": type(exc).__name__, "message": str(exc)}}) + self.reply(200, {"ok": True, "result": result}) + + server = ThreadingHTTPServer((host, port), Handler) + server.daemon_threads = True + return server + + +def _ready_record(config, *, gateway_port, admin_port): + """The one stdout line announcing this shard's resolved endpoints. + + An ephemeral gateway bind (port 0) has no declared URL, so the record + carries the bound port for the launcher to compose the real URL from; + a fixed port repeats the declared URL and port for a single spelling + everywhere. + """ + return {"shard": config.shard_name, "service": SHARD_SERVICE, + "gateway": config.gateway_url or f"http://localhost:{gateway_port}", + "gateway_port": gateway_port, + "admin_port": admin_port, + "trace_dir": str(config.trace_dir)} + + +def serve_shard(config, tokenizer, *, backend=None): + """Run one shard process until SIGTERM, SIGINT or admin shutdown.""" + from harborrl.gateway.server import Gateway, make_server + from harborrl.gateway.sglang_backend import SGLangBackend + from harborrl.gateway.trace import TraceRegistry + from harborrl.math_registry import MATH_REGISTRY + from harborrl.rollout.runtime import tokenizer_fingerprint + from harborrl.trajectories.native import read_json + + tokenizer_digest, template_digest = tokenizer_fingerprint(tokenizer) + if (tokenizer_digest != config.expected_tokenizer_digest + or template_digest != config.expected_template_digest): + raise ValueError("live tokenizer/template digest differs from the native lock") + profile = read_json(config.profile_path) + registry = TraceRegistry(config.trace_dir) + model = MATH_REGISTRY.model_for_preset(config.model_tag) + resolved = backend or SGLangBackend( + config.endpoint, tokenizer, model.adapter, + tokenizer_digest=tokenizer_digest, template_digest=template_digest, + max_context=config.max_context, + audited_raw_logprobs=config.audited_raw_logprobs) + gateway = Gateway(registry, resolved, model=profile["model"], + max_output_tokens=config.max_output_tokens) + token = Path(config.token_file).read_text().strip() + if not token: + raise ValueError(f"shard token file is empty: {config.token_file}") + gateway_server = make_server(gateway, host=config.gateway_host, + port=config.gateway_port) + gateway_thread = threading.Thread( + target=gateway_server.serve_forever, + name=f"harborrl-gateway-{config.shard_name}", daemon=True) + admin = None + stopped = [] + + def on_shutdown(): + if stopped: + return + stopped.append(True) + gateway_server.shutdown() + if admin is not None: + admin.shutdown() + + admin = make_admin_server(registry, shard_name=config.shard_name, token=token, + host=config.admin_host, port=config.admin_port, + on_shutdown=on_shutdown) + + def handle_signal(signum, frame): + # shutdown() must never run on the serve_forever thread. + threading.Thread(target=on_shutdown, daemon=True).start() + + previous = {sig: signal.signal(sig, handle_signal) + for sig in (signal.SIGTERM, signal.SIGINT)} + gateway_thread.start() + admin_port = admin.server_address[1] + print(json.dumps(_ready_record(config, + gateway_port=gateway_server.server_address[1], + admin_port=admin_port)), flush=True) + try: + admin.serve_forever() + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + gateway_thread.join(5) + gateway_server.server_close() + admin.server_close() + print(json.dumps({"shard": config.shard_name, "stopped": True}), flush=True) + + +def main(): + import os + + config = ShardServiceConfig.from_environment(dict(os.environ)) + + def load_tokenizer(): + from transformers import AutoTokenizer + + return AutoTokenizer.from_pretrained(os.environ["HF_CKPT"], + trust_remote_code=True) + + serve_shard(config, load_tokenizer()) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/harborrl/gateway/trace.py b/harborrl/gateway/trace.py index 5f3ba0b12..cd06ce9f5 100644 --- a/harborrl/gateway/trace.py +++ b/harborrl/gateway/trace.py @@ -1,12 +1,89 @@ """Attempt-scoped credentials, request draining and immutable trace sealing.""" import hashlib +import re import secrets import threading import time from copy import deepcopy from pathlib import Path -from harborrl.trajectories.native import digest, publish, read_json +from harborrl.trajectories.native import digest, publish + + +def consumed_tool_defaults(turn, content): + """Validate only omitted optional defaults declared by the request's tool schema.""" + original = turn['response']['content'] + if not isinstance(content, list) or len(content) != len(original): + raise ValueError('consumed response differs or delivery is unknown') + expected = deepcopy(original) + changes = [] + tools = turn.get('client_request', {}).get('tools', []) + for index, (served, consumed) in enumerate(zip(expected, content)): + if digest(served) == digest(consumed): + continue + if (served.get('type') != 'tool_use' or not isinstance(consumed, dict) + or not isinstance(served.get('input'), dict) or not isinstance(consumed.get('input'), dict)): + raise ValueError('consumed response differs or delivery is unknown') + matches = [tool for tool in tools if tool.get('name') == served.get('name')] + if len(matches) != 1: + raise ValueError('consumed tool defaults require a unique declared schema') + schema = matches[0].get('input_schema', {}) + added = {} + for key in consumed['input'].keys() - served['input'].keys(): + prop = schema.get('properties', {}).get(key, {}) + if ('default' not in prop or key in schema.get('required', []) + or digest(consumed['input'][key]) != digest(prop['default'])): + raise ValueError('consumed tool parameter is not a declared optional default') + served['input'][key] = deepcopy(prop['default']) + added[key] = prop['default'] + if digest(served) != digest(consumed): + raise ValueError('consumed response differs or delivery is unknown') + changes.append({'block': index, 'tool': served['name'], 'defaults': added, + 'schema_digest': digest(schema)}) + return {'kind': 'declared_tool_defaults_v1', 'changes': changes, + 'served_content_digest': digest(original), 'consumed_content_digest': digest(content)} + + +def consumed_cli_tool_inputs(turn, content, cli_version): + """Prove pinned native trailing-whitespace removal, then schema defaults.""" + original = turn['response']['content'] + if cli_version != '2.1.141' or not isinstance(content, list) or len(content) != len(original): + raise ValueError('consumed response differs or delivery is unknown') + expected = deepcopy(turn) + changes = [] + # ECMAScript \s, rather than Python's broader Unicode whitespace class. + whitespace = '\t\n\v\f\r \u00a0\u1680\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000\ufeff' + for index, (served, consumed) in enumerate(zip(expected['response']['content'], content)): + if digest(served) == digest(consumed): + continue + if (served.get('type') != 'tool_use' or not isinstance(consumed, dict) + or not isinstance(served.get('input'), dict) or not isinstance(consumed.get('input'), dict)): + continue # The strict defaults check below rejects unrelated changes. + field = {'Edit': 'new_string', 'Write': 'content'}.get(served.get('name')) + inputs = served['input'] + path = inputs.get('file_path') + tools = [tool for tool in turn.get('client_request', {}).get('tools', []) + if tool.get('name') == served.get('name')] + if (not field or not isinstance(path, str) or re.search(r'\.(md|mdx)$', path, re.IGNORECASE) + or len(tools) != 1 or not isinstance(inputs.get(field), str) + or tools[0].get('input_schema', {}).get('properties', {}).get(field, {}).get('type') != 'string'): + continue + pieces = re.split(r'(\r\n|\n|\r)', inputs[field]) + normalized = ''.join(piece.rstrip(whitespace) if i % 2 == 0 else piece + for i, piece in enumerate(pieces)) + if normalized != inputs[field] and consumed['input'].get(field) == normalized: + changes.append({'block': index, 'tool': served['name'], 'field': field, + 'rule': 'non_markdown_line_trailing_ecmascript_whitespace_v1', + 'schema_digest': digest(tools[0]['input_schema']), + 'served_value_digest': digest(inputs[field]), + 'consumed_value_digest': digest(normalized)}) + inputs[field] = normalized + if not changes: + raise ValueError('consumed response differs or delivery is unknown') + defaults = consumed_tool_defaults(expected, content) + return {'kind': 'claude_code_2.1.141_tool_inputs_v1', 'changes': changes, + 'declared_defaults': defaults, 'served_content_digest': digest(original), + 'consumed_content_digest': digest(content)} class TraceRegistry: @@ -19,11 +96,6 @@ def __init__(self, root): self.accepting = True self.version = None - def _turns(self, state): - if state.get("sealed"): - return read_json(state["root"] / "trace.json")["turns"] - return state["turns"] - def register(self, identity, trial_id, *, credential=None): with self.lock: if not self.accepting or (self.version is not None and self.version != identity.policy_version): @@ -109,10 +181,31 @@ def wait_delivered(self, aid, timeout): raise TimeoutError("policy delivery acknowledgement still pending") self.changed.wait(remaining) if any(turn.get("delivery") not in ("sent", "consumed_confirmed") - for turn in self._turns(state)): + for turn in state["turns"]): raise ValueError("policy response delivery was not confirmed") - def consume(self, aid, response_id, content, evidence_ref): + def audit_tool_rejections(self, aid, results, evidence_ref): + """Unavailable model actions are valid only after an exact native rejection.""" + from .messages import unavailable_tool_error + + with self.lock: + state = self.attempts[aid] + if state['sealed']: + raise ValueError('trace already sealed') + for turn in state['turns']: + declared = {tool['name'] for tool in turn['client_request'].get('tools', [])} + unavailable = [block for block in turn['response']['content'] + if block['type'] == 'tool_use' and block['name'] not in declared] + for block in unavailable: + matches = [result for result in results if result.get('tool_use_id') == block['id']] + if (len(matches) != 1 or matches[0].get('is_error') is not True + or matches[0].get('content') != unavailable_tool_error(block['name'])): + state['errors'].append('unconfirmed_tool_rejection') + raise ValueError('unavailable tool lacks its exact native rejection') + if unavailable: + turn['tool_rejection_ref'] = evidence_ref + + def consume(self, aid, response_id, content, evidence_ref, *, cli_version=None): """Called only by trusted native-session audit, never by an HTTP policy client.""" with self.lock: state = self.attempts[aid] @@ -122,18 +215,28 @@ def consume(self, aid, response_id, content, evidence_ref): if len(matches) != 1 or not evidence_ref: raise ValueError("unknown consumed response") turn = matches[0] - if turn["response"]["content"] != content or turn["delivery"] not in ("sent", "consumed_confirmed"): + if turn["delivery"] not in ("sent", "consumed_confirmed"): raise ValueError("consumed response differs or delivery is unknown") - publish(state["root"] / "consumption" / f"{response_id}.json", {"response_id": response_id, "content_digest": digest(content), "evidence_ref": evidence_ref}) + normalization = {} + if digest(turn['response']['content']) != digest(content): + try: + normalization = {'tool_defaults': consumed_tool_defaults(turn, content)} + except ValueError: + normalization = {'tool_normalization': consumed_cli_tool_inputs(turn, content, cli_version)} + publish(state["root"] / "consumption" / f"{response_id}.json", { + "response_id": response_id, "content_digest": digest(content), "evidence_ref": evidence_ref, + **normalization}) + if normalization: + for kind, proof in normalization.items(): + turn['consumption_' + kind] = proof turn.update(delivery="consumed_confirmed", consumption_ref=evidence_ref) def unconsumed_responses(self, aid): with self.lock: state = self.attempts[aid] - turns = self._turns(state) return [{"response_id": turn["response_id"], "content": deepcopy(turn["response"]["content"]), - "finish_reason": turn.get("finish_reason"), "is_final": index == len(turns) - 1} - for index, turn in enumerate(turns) if turn.get("delivery") != "consumed_confirmed"] + "finish_reason": turn.get("finish_reason"), "is_final": index == len(state["turns"]) - 1} + for index, turn in enumerate(state["turns"]) if turn.get("delivery") != "consumed_confirmed"] def close(self, aid): with self.lock: @@ -144,9 +247,6 @@ def close(self, aid): def seal(self, aid): with self.lock: state = self.attempts[aid] - if state["sealed"]: - record = read_json(state["root"] / "trace.json") - return record["turns"], record["seal"] if not state["closed"] or state["pending"]: raise ValueError("close admission and drain requests before sealing") turns = deepcopy(state["turns"]) @@ -158,17 +258,23 @@ def seal(self, aid): "turns_digest": digest(turns), "errors": sorted(set(errors))} publish(state["root"] / "trace.json", {"turns": turns, "seal": seal}) state["sealed"] = True - # Keep only admission/identity bookkeeping after durable publication. - state["turns"].clear() - state["errors"].clear() return turns, seal - def drain(self, timeout): - deadline = time.monotonic() + timeout + def close_admission(self): + """Reject new attempts immediately; drain() still waits for pending ones. + + Split from drain so a cross-shard barrier can close admission on every + shard before waiting on any of them. + """ with self.changed: self.accepting = False for aid in self.attempts: self.close(aid) + + def drain(self, timeout): + deadline = time.monotonic() + timeout + self.close_admission() + with self.changed: while any(s["pending"] for s in self.attempts.values()): remaining = deadline - time.monotonic() if remaining <= 0: diff --git a/harborrl/harnesses/claude_code.py b/harborrl/harnesses/claude_code.py index fb5c27519..2ebd03f19 100644 --- a/harborrl/harnesses/claude_code.py +++ b/harborrl/harnesses/claude_code.py @@ -5,7 +5,7 @@ import json from pathlib import Path -from harborrl.rollout.harbor_job.audit import audit_session +from harborrl.rollout.harbor_job.audit import audit_session, audit_termination from harborrl.rollout.harbor_job.bindings import ( claude_config, sanitized_runner_env, @@ -24,8 +24,11 @@ def agent_config(self, profile, gateway_url, credential, max_output_tokens=8192) def runner_environment(self, environment): return sanitized_runner_env(environment) - def audit_session(self, path, registry, attempt_id): - return audit_session(path, registry, attempt_id) + def audit_session(self, path, registry, attempt_id, *, cli_version=None): + return audit_session(path, registry, attempt_id, cli_version=cli_version) + + def audit_termination(self, path, result_path, profile, audit, *, turns): + return audit_termination(path, result_path, profile, audit, turns=turns) def audit_empty_final(self, path, registry, attempt_id, audit): """Claude 2.1.141 omits empty final messages from persisted sessions.""" diff --git a/harborrl/math/__init__.py b/harborrl/math/__init__.py new file mode 100644 index 000000000..b4f90dec2 --- /dev/null +++ b/harborrl/math/__init__.py @@ -0,0 +1 @@ +"""Reusable mathematical data and reward adapters for the training backend.""" diff --git a/harborrl/math/config.py b/harborrl/math/config.py new file mode 100644 index 000000000..766eb2a7f --- /dev/null +++ b/harborrl/math/config.py @@ -0,0 +1,193 @@ +"""Plain math input to the shared Slime training launcher.""" +import json +import re +import sys +from pathlib import Path + +import yaml + +from harborrl.backends.slime_v032.config import CONFIG_ENV +from harborrl.backends.slime_v032.versions import BACKEND_CONTRACT +from harborrl.config import ROOT +from harborrl.config.native import FIELDS +from harborrl.consistency.config import validate as validate_consistency +from harborrl.math.data import load_tasks +from harborrl.trajectories.native import finite + +MATH_FIELDS = {key: FIELDS[key] for key in ( + 'execution', 'model', 'training', 'sampling', 'deployment', 'output')} +MATH_FIELDS = {**MATH_FIELDS, 'tasks': {'path'}, 'consistency': FIELDS['consistency']} + + +def load_config(config, path, overrides=None, task_ids=None): + if task_ids is not None: + raise ValueError('text_math uses the locked train split, not Harbor task selectors') + config['consistency'] = validate_consistency(config.get('consistency')) + if not isinstance(config.get('training'), dict): + raise TypeError('training must be a mapping') + config['training'].setdefault('backend_contract', BACKEND_CONTRACT) + config['training'].setdefault('resume', None) + config['training'].setdefault('export_hf', False) + config['training'].setdefault('seed', None) + config.setdefault('algorithm', {'name': 'grpo'}) + fields = {**MATH_FIELDS, 'training': MATH_FIELDS['training'] | {'resume', 'export_hf', 'seed'}, + 'algorithm': {'name'}} + if 'teacher' in config: + fields['teacher'] = {'endpoint', 'checkpoint', 'weight_version'} + if 'dapo' in config: + from harborrl.math.dapo import FIELDS as DAPO_FIELDS + fields['dapo'] = DAPO_FIELDS + if 'process_reward' in config: + from harborrl.math.prm import FIELDS as PRM_FIELDS + fields['process_reward'] = PRM_FIELDS + if set(config) != set(fields) | {'schema_version'}: + raise ValueError('text_math requires exactly the documented schema 2 sections') + for section, keys in fields.items(): + if not isinstance(config[section], dict) or set(config[section]) != keys: + raise ValueError(f'{section} requires exactly {sorted(keys)}') + for override in overrides or (): + key, sep, value = override.partition('=') + section, _, field = key.partition('.') + if not sep or field not in fields.get(section, ()): + raise ValueError(f'unknown text_math override: {key}') + config[section][field] = yaml.safe_load(value) + config['consistency'] = validate_consistency(config['consistency']) + algorithm = config['algorithm']['name'] + if algorithm not in ('grpo', 'opd', 'dapo'): + raise ValueError('text_math algorithm.name must be grpo, opd or dapo') + if (algorithm == 'opd') != ('teacher' in config): + raise ValueError('Only OPD requires a teacher section') + if (algorithm == 'dapo') != ('dapo' in config): + raise ValueError('Only DAPO requires a dapo section') + if 'process_reward' in config: + from harborrl.math.prm import validate_config + if algorithm != 'grpo': + raise ValueError('process_reward currently requires algorithm.name: grpo') + validate_config(config['process_reward']) + checkpoint = Path(config['process_reward']['checkpoint']).expanduser() + config['process_reward']['checkpoint'] = str((Path(path).parent / checkpoint).resolve()) + if algorithm == 'opd': + from harborrl.math.teacher import validate_config + validate_config(config['teacher']) + checkpoint = Path(config['teacher']['checkpoint']).expanduser() + config['teacher']['checkpoint'] = str((Path(path).parent / checkpoint).resolve()) + if type(config['training']['export_hf']) is not bool: + raise ValueError('training.export_hf must be a boolean') + seed = config['training']['seed'] + if seed is not None and (type(seed) is not int or not 0 <= seed < 2**31): + raise ValueError('training.seed must be null or an integer in 0..2147483647') + if config['execution']['backend'] != 'text_math': + raise ValueError('execution.backend must be text_math') + if config['training']['backend_contract'] != BACKEND_CONTRACT: + raise ValueError('text_math requires the pinned Slime backend') + for section, field in [('tasks', 'path'), ('model', 'checkpoint'), + ('model', 'reference'), ('output', 'root'), ('training', 'resume')]: + value = config[section][field] + if section == 'training' and value is None: + continue + if not isinstance(value, str) or not value.strip(): + raise ValueError(f'{section}.{field} must be a path') + config[section][field] = str((Path(path).parent / Path(value).expanduser()).resolve()) + for section, keys in [('sampling', FIELDS['sampling']), + ('deployment', FIELDS['deployment'] - {'layout'}), + ('training', {'num_rollout', 'save_interval'})]: + for key in keys: + if type(config[section][key]) is not int or config[section][key] <= 0: + raise ValueError(f'{section}.{key} must be a positive integer') + rate = config['training']['learning_rate'] + if not finite(rate) or rate <= 0: + raise ValueError('training.learning_rate must be positive and finite') + sampling, deployment = config['sampling'], config['deployment'] + if algorithm == 'dapo': + from harborrl.math.dapo import validate + validate(config['dapo'], sampling['max_tokens']) + if sampling['group_size'] < 2 or sampling['max_tokens'] >= sampling['max_context']: + raise ValueError('text_math requires complete groups and max_tokens below max_context') + if sampling['max_attempts'] != 1: + raise ValueError('text_math requires max_attempts: 1; DAPO refills use max_refill_batches') + if deployment['layout'] not in ('split', 'colocate'): + raise ValueError('layout must be split or colocate') + if deployment['layout'] == 'split' and deployment['actor_gpus'] + deployment['rollout_gpus'] > deployment['num_gpus']: + raise ValueError('GPU budget exceeded') + if deployment['layout'] == 'colocate' and not deployment['actor_gpus'] == deployment['rollout_gpus'] == deployment['num_gpus']: + raise ValueError('colocate GPU counts must match') + if (deployment['actor_gpus'] % deployment['actor_tensor_parallel_size'] + or deployment['rollout_gpus'] % deployment['rollout_gpus_per_engine']): + raise ValueError('GPU counts must be divisible by tensor parallel sizes') + if not isinstance(config['model']['args_file'], str) or not re.fullmatch(r'[A-Za-z0-9_-]+', config['model']['args_file']): + raise ValueError('model.args_file must name a Slime model preset') + if config['consistency']['mode'] != 'off': + raise ValueError('text_math currently requires consistency.mode: off; Harbor sidecars are not math evidence') + rows = load_tasks(config['tasks']['path']) + if not any(row['split'] == 'train' for row in rows): + raise ValueError('text_math requires a nonempty train split') + return config + + +def launch_plan(config): + environment = { + name: config[section][field] + for name, (section, field) in CONFIG_ENV.items() + if section in config and field in config[section] + } + environment.update({ + 'TRAIN_PYTHON': sys.executable, + 'HARBORRL_EXECUTION_BACKEND': 'text_math', + 'HARBORRL_MATH_ALGORITHM': config['algorithm']['name'], + 'HARBORRL_NATIVE_SLIME_CONTRACT': 'slime-v032', + 'HARBORRL_CONSISTENCY': json.dumps(config['consistency'], sort_keys=True), + 'HARBORRL_CONSISTENCY_MODE': 'off', + 'SLIME_RAY_PLACEMENT_GPU_PROBE': '1', + 'HARBORRL_SKIP_GLOBAL_CLEANUP': '1', + 'PYTORCH_CUDA_ALLOC_CONF': 'max_split_size_mb:2048', + }) + if config['training']['resume']: + environment['RESUME_LOAD'] = config['training']['resume'] + if config['training']['seed'] is not None: + environment['HARBORRL_MATH_SEED'] = config['training']['seed'] + if config['algorithm']['name'] == 'dapo': + environment['HARBORRL_DAPO'] = json.dumps(config['dapo'], sort_keys=True) + return { + 'config': config, + 'math': {'algorithm': {'id': config['algorithm']['name']}, 'harness': {'id': 'text_math'}}, + 'environment': {key: str(value) for key, value in environment.items()}, + 'training_command': ['bash', str(ROOT / 'harborrl/platform/slime_train.sh')], + 'command': [sys.executable, '-m', 'harborrl.platform.native_train'], + 'required_services': ['GPU Ray / SGLang'], + 'layout': {'mode': config['deployment']['layout'], 'offload': config['deployment']['layout'] == 'colocate'}, + } + + +def materialize(plan, root): + from slime.utils.processing_utils import load_tokenizer + + from harborrl.trajectories.native import digest, publish + + config = plan['config'] + rows = load_tasks(config['tasks']['path']) + tokenizer = load_tokenizer(config['model']['checkpoint'], trust_remote_code=True) + normalized = root / 'config' + normalized.mkdir(parents=True) + if config['algorithm']['name'] == 'opd': + from harborrl.math.teacher import check, freeze + teacher = freeze(config['teacher'], config['model']['checkpoint']) + check(teacher) + publish(normalized / 'teacher.json', teacher) + if 'process_reward' in config: + from harborrl.math import prm + scorer = prm.freeze(config['process_reward']) + prm.check(scorer) + publish(normalized / 'process_reward.json', scorer) + catalog = normalized / 'math_tasks.json' + publish(catalog, {'tasks': rows, 'digest': digest(rows)}) + prompts = normalized / 'math_train.jsonl' + output = [] + for row in rows: + if row['split'] != 'train': + continue + prompt = tokenizer.apply_chat_template( + [{'role': 'user', 'content': row['prompt']}], tokenize=False, + add_generation_prompt=True, enable_thinking=False) + output.append({'task': prompt, 'metadata': {'math_task': row, 'task_digest': digest(row)}}) + prompts.write_text(''.join(json.dumps(row, ensure_ascii=False) + '\n' for row in output), encoding='utf-8') + return catalog, prompts diff --git a/harborrl/math/dapo.py b/harborrl/math/dapo.py new file mode 100644 index 000000000..91890901b --- /dev/null +++ b/harborrl/math/dapo.py @@ -0,0 +1,44 @@ +"""Explicit DAPO math settings and complete-group reward semantics.""" +import json +import math +import os + +from harborrl.trajectories.native import finite + +FIELDS = {'clip_low', 'clip_high', 'overlong_buffer_tokens', 'overlong_penalty', 'max_refill_batches'} +ADVANTAGE_EPSILON = 1e-6 + + +def validate(config, max_tokens): + if not isinstance(config, dict) or set(config) != FIELDS: + raise ValueError(f'dapo requires exactly {sorted(FIELDS)}') + low, high = config['clip_low'], config['clip_high'] + if not finite(low) or not finite(high) or not 0 < low < 1 or high <= low: + raise ValueError('DAPO requires 0 < clip_low < 1 and clip_high > clip_low') + buffer = config['overlong_buffer_tokens'] + if type(buffer) is not int or not 0 < buffer < max_tokens: + raise ValueError('DAPO overlong_buffer_tokens must be between zero and max_tokens') + penalty = config['overlong_penalty'] + if not finite(penalty) or not 0 < penalty <= 1: + raise ValueError('DAPO overlong_penalty must be in (0, 1]') + if type(config['max_refill_batches']) is not int or config['max_refill_batches'] < 1: + raise ValueError('DAPO max_refill_batches must be a positive integer') + return config + + +def environment(max_tokens): + return validate(json.loads(os.environ['HARBORRL_DAPO']), max_tokens) + + +def reward(outcome_reward, length, max_tokens, config): + if outcome_reward not in (-1., 0., 1.) or type(length) is not int or not 0 < length <= max_tokens: + raise ValueError('DAPO requires a verifier outcome of -1, 0 or 1 and a response within max_tokens') + buffer = config['overlong_buffer_tokens'] + penalty = -max(0, length - (max_tokens - buffer)) / buffer * config['overlong_penalty'] + return {'reward': outcome_reward + penalty, 'overlong_penalty': penalty} + + +def advantage(reward, mean, population_std, group_size): + """Use the sample std convention in the published verl DAPO recipe.""" + sample_std = population_std * math.sqrt(group_size / (group_size - 1)) + return (reward - mean) / (sample_std + ADVANTAGE_EPSILON) diff --git a/harborrl/math/data.py b/harborrl/math/data.py new file mode 100644 index 000000000..044a29258 --- /dev/null +++ b/harborrl/math/data.py @@ -0,0 +1,120 @@ +"""Locked math tasks and bounded exact-rational verification; never eval text.""" +import json +import re +from fractions import Fraction +from pathlib import Path + +from harborrl.trajectories.native import digest, publish + +NUMBER = r"[+-]?(?:\d{1,32}(?:\.\d{1,16})?|\d{1,32}/[1-9]\d{0,31})" +GSM_NUMBER = r"[+-]?(?:\d{1,3}(?:,\d{3}){1,10}(?:\.\d{1,16})?|\d{1,32}(?:\.\d{1,16})?|\d{1,32}/[1-9]\d{0,31})" +VERIFIERS = {"exact_rational_v1", "gsm8k_numeric_v1"} +TASK_FIELDS = {"task_id", "problem_id", "prompt", "reference", "verifier", "split", "source", "license"} + + +def load_tasks(path): + rows, identities, problems, prompts = [], set(), {}, {} + source = Path(path) + if source.stat().st_size > 16 * 1024 * 1024: + raise ValueError("Task file exceeds the 16 MiB input limit") + for index, line in enumerate(source.read_text(encoding="utf-8").splitlines(), 1): + row = json.loads(line) + if not isinstance(row, dict) or set(row) != TASK_FIELDS: + raise ValueError(f"Task line {index} requires exactly {sorted(TASK_FIELDS)}") + if any(not isinstance(value, str) or not value.strip() for value in row.values()): + raise ValueError("Task fields must be nonempty strings") + if row["split"] not in {"train", "validation", "test"} or row["verifier"] not in VERIFIERS: + raise ValueError("Only split-locked tasks with a supported numeric verifier are supported") + if len(row["prompt"]) > 8192 or not re.fullmatch(NUMBER, row["reference"]): + raise ValueError("Unsupported prompt length or rational reference") + if row["task_id"] in identities: + raise ValueError("Duplicate task_id") + identities.add(row["task_id"]) + for key, table in ((row["problem_id"], problems), (digest(row["prompt"]), prompts)): + if key in table and table[key] != row["split"]: + raise ValueError("Problem or prompt leakage across splits") + table[key] = row["split"] + rows.append(row) + if not rows: + raise ValueError("Empty math task source") + return rows + + +def _gsm_answer(response): + candidates = [] + for pattern in (r"(?m)^[ \t]*####[ \t]+([^\n]{1,100})$", r"\\boxed\{([^{}]{1,100})\}"): + for match in re.finditer(pattern, response): + text = match[1].strip() + for wrapper in ("**", "$"): + if text.startswith(wrapper) and text.endswith(wrapper) and len(text) >= 2 * len(wrapper): + text = text[len(wrapper):-len(wrapper)].strip() + if re.fullmatch(GSM_NUMBER, text): + candidates.append((match.start(), text.replace(",", ""))) + if candidates: + return max(candidates)[1] + text = response.strip() + return text if re.fullmatch(NUMBER, text) else "" + + +def verify(response, reference, verifier="exact_rational_v1"): + if verifier not in VERIFIERS: + raise ValueError(f"Unsupported numeric verifier: {verifier}") + record = {"verifier": verifier, "input_digest": digest({"response": response, "reference": reference}), + "status": "unparseable", "reward": 0.0, "parsed": None} + if not isinstance(response, str) or len(response) > 16384: + record["status"] = "unsupported" + return record + text = response.strip() + if verifier == "gsm8k_numeric_v1": + text = _gsm_answer(text) + elif "####" in text: + text = text.rsplit("####", 1)[1].strip() + elif "\\boxed{" in text: + match = re.search(r"\\boxed\{([^{}]{1,100})\}\s*$", text) + text = match[1] if match else "" + if not re.fullmatch(NUMBER, text): + return record + correct = Fraction(text) == Fraction(reference) + incorrect_reward = 0.0 if verifier == "gsm8k_numeric_v1" else -1.0 + record.update(status="correct" if correct else "incorrect", reward=1.0 if correct else incorrect_reward, parsed=text) + return record + + +def import_samples(tasks_path, samples_path, output): + tasks = {row["task_id"]: row for row in load_tasks(tasks_path)} + root = Path(output) + if root.exists(): + raise ValueError("Dataset output must be a new directory") + if Path(samples_path).stat().st_size > 16 * 1024 * 1024: + raise ValueError("Sample file exceeds the 16 MiB input limit") + raw = Path(samples_path).read_bytes() + rows = [json.loads(line) for line in raw.decode("utf-8").splitlines()] + if not rows: + raise ValueError("Empty sample source") + prepared = [] + identities = set() + for row in rows: + if not isinstance(row, dict) or set(row) != {"task_id", "response", "source", "license"}: + raise ValueError("Imported row requires task_id, response, source and license") + if any(not isinstance(value, str) or not value.strip() for value in row.values()): + raise ValueError("Sample values must be nonempty strings") + task = tasks[row["task_id"]] + record = {"schema_version": "math-sample-v1", "task_id": task["task_id"], "problem_id": task["problem_id"], + "split": task["split"], "prompt": task["prompt"], "response": row["response"], + "source_kind": "curated", "evidence_kind": "declared", "status": "collected", + "source": row["source"], "license": row["license"], "source_digest": digest(raw.hex()), + "task_digest": digest(task), "parent_digest": None, "token_evidence": None} + record["sample_id"] = digest(record) + if record["sample_id"] in identities: + raise ValueError("Duplicate imported sample") + identities.add(record["sample_id"]) + prepared.append((record, verify(row["response"], task["reference"], task["verifier"]))) + root.mkdir(parents=True) + for sample, verification in prepared: + publish(root / "samples" / f"{sample['sample_id']}.json", sample) + publish(root / "verification" / f"{sample['sample_id']}.json", {**verification, "sample_digest": digest(sample)}) + manifest = {"schema_version": "math-dataset-v1", "tasks_digest": digest(list(tasks.values())), + "samples": [{"sample_id": s["sample_id"], "digest": digest(s), "verification_digest": digest({**v, "sample_digest": digest(s)})} + for s, v in prepared], "training_eligibility": "not_an_on_policy_batch"} + publish(root / "manifest.json", manifest) + return manifest diff --git a/harborrl/math/prm.py b/harborrl/math/prm.py new file mode 100644 index 000000000..34067cfa5 --- /dev/null +++ b/harborrl/math/prm.py @@ -0,0 +1,171 @@ +"""Frozen Qwen math PRM identity, step alignment and scoring receipts.""" +from __future__ import annotations + +import hashlib +import math +import re +from pathlib import Path +from urllib.parse import urlsplit + +from harborrl.consistency.sglang import http_json +from harborrl.trajectories.native import digest, finite, read_json + +ADAPTER = 'qwen-math-prm-v1' +MARKER = '' +SYSTEM = 'Please reason step by step, and put your final answer within \\boxed{}.' +FIELDS = {'endpoint', 'checkpoint', 'revision'} + + +def validate_config(config): + if not isinstance(config, dict) or set(config) != FIELDS: + raise ValueError('process_reward requires endpoint, checkpoint and revision') + if any(not isinstance(value, str) or not value.strip() for value in config.values()): + raise ValueError('process_reward fields must be nonempty strings') + url = urlsplit(config['endpoint']) + if (url.scheme not in ('http', 'https') or not url.hostname or url.username or url.password + or url.query or url.fragment or url.path not in ('', '/')): + raise ValueError('process_reward.endpoint must be a direct HTTP(S) origin') + if not re.fullmatch(r'[0-9a-f]{40}', config['revision']): + raise ValueError('process_reward.revision must be an immutable 40-character commit') + + +def freeze(config): + """Pin local weights, tokenizer and executable model sources by content.""" + validate_config(config) + root = Path(config['checkpoint']).resolve() + model = read_json(root / 'config.json') + if (model.get('architectures') != ['Qwen2ForProcessRewardModel'] + or model.get('auto_map') != { + 'AutoConfig': 'configuration_qwen2_rm.Qwen2RMConfig', + 'AutoModel': 'modeling_qwen2_rm.Qwen2ForProcessRewardModel'}): + raise ValueError('process_reward requires the Qwen2 math process classification model') + names = {'config.json', 'tokenizer.json', 'tokenizer_config.json', + 'configuration_qwen2_rm.py', 'modeling_qwen2_rm.py'} + index = root / 'model.safetensors.index.json' + if index.is_file(): + weights = read_json(index).get('weight_map', {}) + if not isinstance(weights, dict) or not weights: + raise ValueError('PRM requires a nonempty safetensors index') + names.update(weights.values()) + names.add(index.name) + else: + names.add('model.safetensors') + names.update(name for name in ('special_tokens_map.json', 'added_tokens.json', 'vocab.json', 'merges.txt') + if (root / name).is_file()) + files = {} + for name in sorted(names): + if not isinstance(name, str) or Path(name).name != name or not (root / name).is_file(): + raise ValueError('PRM requires a complete local checkpoint and model sources') + value = hashlib.sha256() + with (root / name).open('rb') as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b''): + value.update(chunk) + files[name] = value.hexdigest() + frozen = {**config, 'checkpoint': str(root), 'endpoint': config['endpoint'].rstrip('/'), + 'adapter': ADAPTER, 'files': files} + # Endpoint is transport configuration, not part of the model's identity. + return {**frozen, 'identity': digest({key: value for key, value in frozen.items() if key != 'endpoint'})} + + +def load(path): + frozen = read_json(path) + expected = digest({key: value for key, value in frozen.items() if key not in ('identity', 'endpoint')}) + if frozen.get('identity') != expected or frozen.get('adapter') != ADAPTER: + raise ValueError('PRM manifest identity changed') + validate_config({key: frozen[key] for key in FIELDS}) + return frozen + + +def check(frozen, *, transport=http_json): + info = transport(frozen['endpoint'] + '/model_info') + if any(info.get(key) != frozen[key] for key in ('identity', 'revision', 'adapter', 'checkpoint')): + raise ValueError('PRM scorer identity changed') + if info.get('frozen') is not True or info.get('is_generation') is not False: + raise ValueError('PRM must be a frozen inference-only scorer') + return {key: info[key] for key in ('identity', 'revision', 'adapter', 'checkpoint', 'frozen')} + + +def split_steps(response_ids, tokenizer): + """Double-newline steps; an indivisible boundary token stays in its step.""" + if (not response_ids or len(response_ids) > 16384 + or any(type(token) is not int or token < 0 for token in response_ids)): + raise ValueError('PRM requires a nonempty bounded response token sequence') + text = tokenizer.decode(response_ids, skip_special_tokens=True) + ends, steps, char_start, seen = [], [], 0, 0 + for index in range(1, len(response_ids) + 1): + prefix = tokenizer.decode(response_ids[:index], skip_special_tokens=True) + # Byte-level tokenizers can temporarily end in an incomplete code point. + if not text.startswith(prefix): + continue + count = prefix.count('\n\n') + if count < seen: + raise ValueError('PRM tokenizer changed an established step boundary') + if count > seen and prefix[char_start:].strip(): + ends.append(index) + steps.append(prefix[char_start:]) + char_start = len(prefix) + seen = count + tail = text[char_start:] + if tail.strip(): + ends.append(len(response_ids)) + steps.append(tail) + elif steps: + ends[-1] = len(response_ids) + steps[-1] += tail + else: + raise ValueError('PRM cannot score a response with no visible reasoning') + if ''.join(steps) != text: + raise ValueError('PRM step alignment changed the visible response') + return ends, steps + + +def token_returns(ends, probabilities, response_length): + """Signed rewards 2*p-1, undiscounted reward-to-go within each token span.""" + if (type(response_length) is not int or response_length < 1 or not ends + or len(ends) != len(probabilities) or ends[-1] != response_length + or any(type(end) is not int for end in ends) + or any(left >= right for left, right in zip((0, *ends[:-1]), ends, strict=True))): + raise ValueError('PRM step ends must partition every response token exactly once') + if any(not finite(value) or not 0 <= value <= 1 for value in probabilities): + raise ValueError('PRM probabilities must be finite and in [0, 1]') + rewards = [2 * value - 1 for value in probabilities] + values = [0.] * response_length + starts = (0, *ends[:-1]) + for index, (start, end) in enumerate(zip(starts, ends, strict=True)): + values[start:end] = [math.fsum(rewards[index:])] * (end - start) + return values + + +def request_payload(identity, query, steps): + if (not isinstance(query, str) or not query.strip() or not isinstance(steps, (tuple, list)) + or not steps or any(not isinstance(step, str) or not step.strip() for step in steps) + or any(MARKER in value for value in (query, *steps))): + raise ValueError('PRM query and nonempty steps must not contain the reward marker') + return {'identity': identity, 'query': query, 'steps': list(steps)} + + +def validate_receipt(frozen, receipt, query, response_ids, tokenizer): + ends, steps = split_steps(response_ids, tokenizer) + payload = request_payload(frozen['identity'], query, steps) + if (receipt.get('identity') != frozen['identity'] + or receipt.get('request_digest') != digest(payload) + or receipt.get('response_tokens_digest') != digest(response_ids) + or receipt.get('step_ends') != ends or receipt.get('steps') != steps): + raise ValueError('PRM scorer, request or exact response-token evidence changed') + values = token_returns(ends, receipt.get('probabilities', []), len(response_ids)) + if receipt.get('token_returns') != values: + raise ValueError('PRM reward-to-go evidence changed') + if (not isinstance(receipt.get('scorer_input_digest'), str) + or not re.fullmatch(r'[0-9a-f]{64}', receipt['scorer_input_digest'])): + raise ValueError('PRM scoring input digest is missing') + return values + + +def score_response(frozen, query, response_ids, tokenizer, *, transport=http_json): + ends, steps = split_steps(response_ids, tokenizer) + payload = request_payload(frozen['identity'], query, steps) + result = transport(frozen['endpoint'] + '/score', payload) + receipt = {**result, 'response_tokens_digest': digest(response_ids), 'step_ends': ends, 'steps': steps, + 'token_returns': token_returns(ends, result.get('probabilities', []), len(response_ids))} + validate_receipt(frozen, receipt, query, response_ids, tokenizer) + return receipt diff --git a/harborrl/math/prm_server.py b/harborrl/math/prm_server.py new file mode 100644 index 000000000..7840f0697 --- /dev/null +++ b/harborrl/math/prm_server.py @@ -0,0 +1,111 @@ +"""Inference-only local HTTP service for a pinned Qwen math process model.""" +from __future__ import annotations + +import argparse +import json +from http.server import BaseHTTPRequestHandler, HTTPServer + +from harborrl.math import prm +from harborrl.trajectories.native import digest + + +class Scorer: + def __init__(self, config, device): + import torch + from transformers import AutoModel, AutoTokenizer + + self.frozen = prm.freeze(config) + self.tokenizer = AutoTokenizer.from_pretrained(config['checkpoint'], trust_remote_code=True, + local_files_only=True) + self.model = AutoModel.from_pretrained( + config['checkpoint'], torch_dtype=torch.bfloat16, trust_remote_code=True, + local_files_only=True, attn_implementation='sdpa').to(device).eval() + self.model.requires_grad_(False) + self.marker_id = self.tokenizer.encode(prm.MARKER, add_special_tokens=False) + if len(self.marker_id) != 1: + raise ValueError('PRM marker must encode to one token') + if prm.freeze(config) != self.frozen: + raise ValueError('PRM checkpoint changed while loading the scorer') + + def info(self): + frozen = not self.model.training and all(not parameter.requires_grad for parameter in self.model.parameters()) + return {key: self.frozen[key] for key in ('identity', 'revision', 'adapter', 'checkpoint')} | { + 'frozen': frozen, 'is_generation': False, + 'max_context': self.model.config.max_position_embeddings} + + def score(self, payload): + import torch + + if not isinstance(payload, dict) or set(payload) != {'identity', 'query', 'steps'}: + raise ValueError('PRM score requires identity, query and steps') + expected = prm.request_payload(payload['identity'], payload['query'], payload['steps']) + if expected['identity'] != self.frozen['identity'] or not self.info()['frozen']: + raise ValueError('PRM scoring requires the frozen manifest identity') + messages = [{'role': 'system', 'content': prm.SYSTEM}, + {'role': 'user', 'content': payload['query']}, + {'role': 'assistant', 'content': prm.MARKER.join(payload['steps']) + prm.MARKER}] + text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False) + ids = self.tokenizer.encode(text, add_special_tokens=False) + if len(ids) > self.model.config.max_position_embeddings: + raise ValueError('PRM scoring context exceeded; truncation is not permitted') + positions = [index for index, token in enumerate(ids) if token == self.marker_id[0]] + if len(positions) != len(payload['steps']): + raise ValueError('PRM marker positions differ from the submitted steps') + with torch.inference_mode(): + inputs = torch.tensor([ids], dtype=torch.long, device=self.model.device) + logits = self.model(input_ids=inputs, use_cache=False).logits + if logits.shape != (1, len(ids), 2) or not torch.isfinite(logits).all(): + raise ValueError('PRM must return finite two-class token logits') + values = logits[0, positions].float().softmax(-1)[:, 1].cpu().tolist() + return {'identity': self.frozen['identity'], 'request_digest': digest(expected), + 'scorer_input_digest': digest(ids), 'probabilities': values} + + +def serve(scorer, host, port): + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def reply(self, status, payload): + content = json.dumps(payload, allow_nan=False).encode() + self.send_response(status) + self.send_header('Content-Type', 'application/json') + self.send_header('Content-Length', str(len(content))) + self.end_headers() + self.wfile.write(content) + + def do_GET(self): + self.reply(200, scorer.info()) if self.path == '/model_info' else self.reply(404, {'error': 'not found'}) + + def do_POST(self): + if self.path != '/score': + return self.reply(404, {'error': 'not found'}) + try: + length = int(self.headers.get('Content-Length', '0')) + if not 0 < length <= 1024 * 1024: + raise ValueError('PRM request size must be between 1 byte and 1 MiB') + result = scorer.score(json.loads(self.rfile.read(length))) + except (ValueError, TypeError, KeyError) as exc: + return self.reply(400, {'error': str(exc)}) + self.reply(200, result) + + with HTTPServer((host, port), Handler) as server: + print(json.dumps({'ready': True, **scorer.info()}), flush=True) + server.serve_forever() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('--checkpoint', required=True) + parser.add_argument('--revision', required=True) + parser.add_argument('--host', default='127.0.0.1') + parser.add_argument('--port', type=int, default=30091) + parser.add_argument('--device', default='cuda:0') + args = parser.parse_args() + config = {'checkpoint': args.checkpoint, 'revision': args.revision, + 'endpoint': f'http://127.0.0.1:{args.port}'} + serve(Scorer(config, args.device), args.host, args.port) + + +if __name__ == '__main__': + main() diff --git a/harborrl/math/scoring.py b/harborrl/math/scoring.py new file mode 100644 index 000000000..c6307c841 --- /dev/null +++ b/harborrl/math/scoring.py @@ -0,0 +1,47 @@ +"""Token-aligned arithmetic process rewards; no training loop.""" +import operator +import re +from fractions import Fraction + +from harborrl.math.data import NUMBER +from harborrl.trajectories.native import digest + + +def process_rewards(kind, response_ids, tokenizer): + """Each newline terminates a step; a boundary token belongs to that step.""" + if kind not in {"fake", "equation"} or not response_ids or len(response_ids) > 16384: + raise ValueError("Invalid or oversized process-reward request") + ends, previous = [], 0 + for index in range(len(response_ids)): + count = tokenizer.decode(response_ids[:index + 1]).count("\n") + if count > previous: + ends.append(index + 1) + previous = count + if not ends or ends[-1] != len(response_ids): + ends.append(len(response_ids)) + steps, start = [], 0 + for end in ends: + text = tokenizer.decode(response_ids[start:end]).strip() + if kind == "fake": + score = (int(digest(response_ids[start:end])[:8], 16) % 9 + 1) / 10 + status = "synthetic" + else: + match = re.fullmatch(rf"\s*({NUMBER})\s*([+*/-])\s*({NUMBER})\s*=\s*({NUMBER})\s*", text) + score, status = 0.0, "unsupported" + if match and len(text) <= 1024: + a, b, expected = Fraction(match[1]), Fraction(match[3]), Fraction(match[4]) + try: + actual = {"+": operator.add, "-": operator.sub, "*": operator.mul, "/": operator.truediv}[match[2]](a, b) + score = 1.0 if actual == expected else -1.0 + status = "correct" if actual == expected else "incorrect" + except ZeroDivisionError: + status = "unsupported" + steps.append({"start": start, "end": end, "text": text, "score": score, "status": status}) + start = end + values = [0.0] * len(response_ids) + cumulative = 0.0 + for step in reversed(steps): + cumulative += step["score"] + values[step["start"]:step["end"]] = [cumulative] * (step["end"] - step["start"]) + return {"adapter": f"{kind}-prm-v1", "evidence_kind": "synthetic" if kind == "fake" else "observed", + "steps": steps, "token_returns": values, "aggregation": "undiscounted-step-reward-to-go"} diff --git a/harborrl/math/teacher.py b/harborrl/math/teacher.py new file mode 100644 index 000000000..c957cb980 --- /dev/null +++ b/harborrl/math/teacher.py @@ -0,0 +1,93 @@ +"""Frozen, local-checkpoint SGLang teacher with exact student-token scoring.""" +from __future__ import annotations + +import hashlib +from pathlib import Path +from urllib.parse import urlsplit + +from harborrl.consistency.sglang import http_json, score +from harborrl.trajectories.native import digest, read_json + +TOKENIZER_FILES = ( + 'tokenizer.json', 'tokenizer_config.json', 'special_tokens_map.json', + 'added_tokens.json', 'vocab.json', 'merges.txt', +) + + +def validate_config(config): + if not isinstance(config, dict) or set(config) != {'endpoint', 'checkpoint', 'weight_version'}: + raise ValueError('teacher requires endpoint, checkpoint and weight_version') + if any(not isinstance(value, str) or not value.strip() for value in config.values()): + raise ValueError('teacher fields must be nonempty strings') + url = urlsplit(config['endpoint']) + if (url.scheme not in ('http', 'https') or not url.hostname or url.username or url.password + or url.query or url.fragment or url.path not in ('', '/')): + raise ValueError('teacher.endpoint must be a direct SGLang server HTTP(S) origin') + + +def _sha256(path): + result = hashlib.sha256() + with path.open('rb') as stream: + for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b''): + result.update(chunk) + return result.hexdigest() + + +def _tokenizer_files(root): + if any(not (root / name).is_file() for name in TOKENIZER_FILES[:2]): + raise ValueError('OPD requires tokenizer.json and tokenizer_config.json') + return {name: _sha256(root / name) for name in TOKENIZER_FILES if (root / name).is_file()} + + +def freeze(config, student_checkpoint): + """Hash exactly the teacher safetensors checkpoint and compatible tokenizer.""" + validate_config(config) + root = Path(config['checkpoint']).resolve() + tokenizer = _tokenizer_files(root) + if tokenizer != _tokenizer_files(Path(student_checkpoint)): + raise ValueError('OPD teacher and student tokenizer files must match exactly') + index = root / 'model.safetensors.index.json' + names = {'config.json'} + if index.is_file(): + weight_map = read_json(index).get('weight_map', {}) + if not weight_map: + raise ValueError('Empty teacher safetensors weight map') + names.update(weight_map.values()) + names.add(index.name) + else: + names.add('model.safetensors') + files = {} + for name in sorted(names): + if Path(name).name != name or not (root / name).is_file(): + raise ValueError('Teacher requires a complete local safetensors checkpoint') + files[name] = _sha256(root / name) + frozen = {**config, 'checkpoint': str(root), 'endpoint': config['endpoint'].rstrip('/'), + 'files': files, 'tokenizer_files': tokenizer} + return {**frozen, 'identity': digest(frozen)} + + +def load(path): + frozen = read_json(path) + if frozen.get('identity') != digest({key: value for key, value in frozen.items() if key != 'identity'}): + raise ValueError('OPD teacher manifest identity changed') + return frozen + + +def check(frozen, *, transport=http_json): + """Check the trusted server's model, tokenizer and explicit weight version.""" + info = transport(frozen['endpoint'] + '/model_info') + checkpoint = Path(frozen['checkpoint']).resolve() + if (Path(info.get('model_path', '')).resolve() != checkpoint + or Path(info.get('tokenizer_path') or info.get('model_path', '')).resolve() != checkpoint + or str(info.get('weight_version', '')) != frozen['weight_version'] + or info.get('is_generation') is not True): + raise ValueError('OPD teacher model, tokenizer or weight version changed') + return {'checkpoint': str(checkpoint), 'weight_version': frozen['weight_version'], + 'identity': frozen['identity']} + + +def score_response(frozen, prompt, response, *, transport=http_json): + values = score(frozen['endpoint'], prompt, response, frozen['weight_version'], + transport=transport, label='OPD teacher') + return {'identity': frozen['identity'], 'weight_version': frozen['weight_version'], + 'logprobs': values, 'tokens_digest': digest(prompt + response)} diff --git a/harborrl/models/qwen.py b/harborrl/models/qwen.py index cbd096dad..bd2fb063c 100644 --- a/harborrl/models/qwen.py +++ b/harborrl/models/qwen.py @@ -36,4 +36,4 @@ def render_chat(self, tokenizer, converted): return token_ids def parse_output(self, text, *, response_id, tools): - return parse_qwen_output(text, response_id=response_id, tools=tools) + return parse_qwen_output(text, response_id=response_id, tools=tools, allow_unknown_tools=True) diff --git a/harborrl/options/catalog.py b/harborrl/options/catalog.py index 76ede0794..cebe93cae 100644 --- a/harborrl/options/catalog.py +++ b/harborrl/options/catalog.py @@ -38,6 +38,8 @@ "model.args_file": ("string", "Supported Slime preset name; currently the qwen prefix"), "training.num_rollout": ("integer", "Positive number of rollout batches"), "training.learning_rate": ("number", "Positive finite learning rate"), + "training.seed": ("number_or_null", "Sampling seed shared by every group; null keeps sampling nondeterministic"), + "training.export_hf": ("boolean", "Also export an HF checkpoint next to the training checkpoints"), "training.save_interval": ("integer", "Positive checkpoint interval in rollout batches"), "training.backend_contract": ("string", "Pinned backend/algorithm contract"), "sampling.group_size": ("integer", "Samples per prompt, at least two for GRPO"), @@ -61,8 +63,22 @@ "rollout.workers[].ssh": schema.SSH_FIELDS, "connectivity": schema.CONNECTIVITY_FIELDS, } DEPLOYMENT_FIELDS = { - "id": ("string", "Short role identifier"), "mode": ("string", "Manager mode or shared storage mode"), + "id": ("string", "Short role identifier"), + "mode": ("string", "Manager mode or shared storage mode"), "bind": ("string", "Manager listener address"), + "gateway_shards": ("array", "Static gateway shard table; omit to let auto routing own the shard set"), + "gateway_routing": ("string", "Shard routing rule: modulo, worker_affinity, or auto"), + "sandboxes_per_shard": ("integer", "Concurrent sandboxes per gateway before auto routing grows one"), + "worker_shards": ("object", "Worker id (or local) to gateway shard name; worker_affinity only"), + "gateway_hosts": ("array", "Ssh destinations (user@host) auto growth fills first; empty grows locally"), + "shard_python": ("path", "Interpreter on gateway_hosts for remotely grown shards; default runner_python"), + "hub_port": ("integer", "Hub control-plane listener port; hub mode only"), + "hub_url": ("string", "Hub control-plane origin reachable from every puller node"), + "hub_pull_interval": ("integer", "Seconds between puller polls in hub mode"), + "gpu_pool": ("array", "Declared GPU nodes the hub schedules runs onto"), + "max_concurrent_runs": ("integer", "Hub mode: how many runs may execute at once"), + "run_queue_size": ("integer", "Hub mode: pending run queue capacity"), + "provision": ("object", "up/down commands that obtain and release the GPU machine; site-specific"), "control_port": ("integer", "Control listener port; different from gateway_port"), "control_url": ("string", "Control origin reachable from the GPU target"), "gateway_port": ("integer", "Gateway listener port"), @@ -85,7 +101,8 @@ "docker": ("boolean", "Worker Docker requirement"), "task_roots": ("array", "Absolute task roots visible on this worker"), "host": ("string", "SSH alias/destination, not an automatically verified host"), - "user": ("string", "Optional SSH username"), "port": ("integer", "Optional SSH port"), + "user": ("string", "Optional SSH username"), + "port": ("integer", "Optional SSH port"), "identity": ("path", "Optional identity-file reference; no private-key contents"), "command_timeout_s": ("integer", "Positive remote command timeout"), "ssh_timeout_s": ("integer", "Positive SSH connection timeout"), diff --git a/harborrl/options/presets.py b/harborrl/options/presets.py index e6e2d6842..f52fdc365 100644 --- a/harborrl/options/presets.py +++ b/harborrl/options/presets.py @@ -53,7 +53,8 @@ def defaults(preset, output_directory): "tokenizer_digest": None, "template_digest": None, "audited_raw_logprobs": True}, "model": {"checkpoint": None, "reference": None, "args_file": "qwen3-8B"}, "training": {"num_rollout": 2, "learning_rate": 1e-6, "save_interval": 1, - "backend_contract": "slime-v0.3.2-native-v1"}, + "backend_contract": "slime-v0.3.2-native-v1", + "seed": None, "export_hf": False}, "sampling": {"group_size": 2, "groups_per_batch": 1, "max_attempts": 1, "max_tokens": 2048, "max_context": 32768}, "deployment": {"layout": "colocate", "num_gpus": 1, "actor_gpus": 1, diff --git a/harborrl/options/validation.py b/harborrl/options/validation.py index 53465d83e..d0fb76224 100644 --- a/harborrl/options/validation.py +++ b/harborrl/options/validation.py @@ -60,7 +60,9 @@ def require(condition, path, message): destinations = [schema.destination(worker["ssh"]) for worker in workers] require(harbor["workers"] == destinations, "training.harbor.workers", "Must match deployment worker destinations in order") - require(harbor["python"] == workers[0]["runner_python"], "training.harbor.python", + import os + require(os.path.normpath(harbor["python"]) == os.path.normpath(workers[0]["runner_python"]), + "training.harbor.python", "Must match the deployment's first runner_python") return checks diff --git a/harborrl/platform/evidence.py b/harborrl/platform/evidence.py index d988cf86f..c9f49797e 100644 --- a/harborrl/platform/evidence.py +++ b/harborrl/platform/evidence.py @@ -25,12 +25,23 @@ def run_environment(plan, root, catalog_path, prompt_path): "CKPT_ROOT": str(root / "checkpoints"), "WANDB_DIR": str(root / "metrics/wandb"), "HARBORRL_NATIVE_CATALOG": str(catalog_path), "ROLLOUT_PROMPT_DATA": str(prompt_path), - "HARBORRL_NATIVE_PROFILE": plan["config"]["harness"]["profile"], + **({"HARBORRL_NATIVE_PROFILE": plan["config"]["harness"]["profile"]} + if "harness" in plan["config"] else {}), "HARBOR_IR_ROOT": str(root / "ir"), + **({"SAVE_HF": str(root / "hf" / "rollout-{rollout_id}")} + if plan['config']['training'].get('export_hf') else {}), + **({'HARBORRL_OPD_TEACHER': str(root / 'config' / 'teacher.json')} + if plan['math']['algorithm']['id'] == 'opd' else {}), + **({'HARBORRL_PROCESS_REWARD': str(root / 'config' / 'process_reward.json') + if 'process_reward' in plan['config'] else ''} + if plan['config']['execution']['backend'] == 'text_math' else {}), } def _materialize(plan, root): + if plan['config']['execution']['backend'] == 'text_math': + from harborrl.math.config import materialize + return materialize(plan, root) from harborrl.math_registry import select_math config = plan["config"] @@ -95,12 +106,47 @@ def effective_config(plan, environment, child_env, source, checks): "environment": environment, "deployment_environment": {key: child_env[key] for key in DEPLOYMENT_EVIDENCE_KEYS if key in child_env}, "training_command": plan["training_command"], + "update_semantics": ({ + "objective": "sampled_token_fixed_prefix_reverse_kl_surrogate", + "advantage": "detached_teacher_logprob_minus_detached_behavior_logprob", + "reward_mixing": False, + "clipping": False, + "std_normalization": False, + "reduction": "trajectory_weighted_sampled_opd_v1", + "teacher_manifest": "config/teacher.json", + } if plan['math']['algorithm']['id'] == 'opd' else { + 'objective': 'asymmetric_clipped_policy_gradient', + 'advantage': 'shaped_reward_minus_complete_group_mean_div_sample_std_plus_1e-6', + 'std_normalization': True, + 'reduction': 'global_token_mean_asymmetric_clipped_pg_v1', + 'dynamic_sampling': 'reject_zero_variance_outcome_groups_and_refill_complete_groups', + 'truncation': 'retain_all_response_tokens_with_soft_overlong_penalty', + **config['dapo'], + } if plan['math']['algorithm']['id'] == 'dapo' else { + 'objective': 'trajectory_weighted_clipped_process_policy_gradient', + 'reward_source': 'qwen-math-prm-v1', + 'step_reward': '2 * positive_probability - 1', + 'advantage': 'undiscounted_step_reward_to_go_minus_complete_group_mean_total_reward', + 'step_boundary': 'double_newline_at_student_token_boundary; boundary_token_belongs_to_previous_step', + 'std_normalization': False, + 'outcome_reward_mixing': False, + 'truncation': 'score_observed_tail; retain_all_response_tokens; no_unobserved_suffix_reward', + 'reduction': 'trajectory_weighted_clipped_pg_v1', + 'scorer_manifest': 'config/process_reward.json', + } if 'process_reward' in config else { + "advantage": "reward_minus_complete_group_mean", + "std_normalization": False, + "reduction": "trajectory_weighted_clipped_pg_v1", + }), "expanded_command_file": "effective-command.json", "source": source, "runtime_versions": [check for check in checks if "version" in check or check.get("service") in ("slime commit", "megatron commit", "sglang image")], - "field_usage": field_usage(), + "field_usage": ({f'{section}.{key}': ['text_math configuration'] + for section, values in config.items() if isinstance(values, dict) + for key in values} + if config['execution']['backend'] == 'text_math' else field_usage()), "precedence": "YAML < --set in order < explicit MATH selectors; declared deployment environment only", } diff --git a/harborrl/platform/native_train.py b/harborrl/platform/native_train.py index 0c0eedba2..87d8bba72 100644 --- a/harborrl/platform/native_train.py +++ b/harborrl/platform/native_train.py @@ -43,8 +43,10 @@ def dispatch(plan, *, doctor_only=False): explicit = plan.get("deployment", {}).get("source", "embedded-local") != "embedded-local" print(json.dumps({"checks": checks, "scope": ("deployment topology, " if explicit else "") - + "native config, original tasks, isolated Harbor, training " - "dependencies and GPU budget; protocol/token semantics " + + ("math config, locked tasks, training " if + plan['config']['execution']['backend'] == 'text_math' else + "native config, original tasks, isolated Harbor, training ") + + "dependencies and GPU budget; protocol/token semantics " "require the deferred live probe"}, indent=2)) return int(any(not check["ok"] for check in checks)) diff --git a/harborrl/platform/preflight.py b/harborrl/platform/preflight.py index 3a35fe88f..a0be564ff 100644 --- a/harborrl/platform/preflight.py +++ b/harborrl/platform/preflight.py @@ -24,6 +24,40 @@ def _deployment(plan): def doctor(plan): + if plan['config']['execution']['backend'] == 'text_math': + return _math_doctor(plan) + return _doctor(plan) + + +def _math_doctor(plan): + """Use the shared backend probes without invoking Harbor task workers.""" + config = plan['config'] + from harborrl.math.data import load_tasks + checks = [] + try: + rows = load_tasks(config['tasks']['path']) + checks.append({'service': 'locked math tasks', 'ok': True, 'count': len(rows)}) + except (ValueError, OSError, TypeError) as exc: + checks.append({'service': 'locked math tasks', 'ok': False, 'detail': str(exc)}) + if config['algorithm']['name'] == 'opd': + from harborrl.math.teacher import check + try: + check({**config['teacher'], 'identity': 'preflight'}) + checks.append({'service': 'OPD teacher identity', 'ok': True, + 'detail': 'live model/version checked; checkpoint and tokenizer hashes freeze at launch'}) + except (ValueError, OSError, TypeError) as exc: + checks.append({'service': 'OPD teacher identity', 'ok': False, 'detail': str(exc)}) + if 'process_reward' in config: + from harborrl.math import prm + try: + prm.check(prm.freeze(config['process_reward'])) + checks.append({'service': 'frozen PRM identity', 'ok': True}) + except (ValueError, OSError, TypeError, KeyError) as exc: + checks.append({'service': 'frozen PRM identity', 'ok': False, 'detail': str(exc)}) + return checks + _doctor(plan, math_only=True) + + +def _doctor(plan, *, math_only=False): config = plan["config"] checks = [] deployment = _deployment(plan) @@ -45,10 +79,11 @@ def doctor(plan): "by slime_train.sh at GPU job startup on the GPU machine", }) - runner = Path(config["harbor"]["python"]) - checks.append({"path": "harbor.python", "ok": runner.is_file()}) + runner = Path(config.get("harbor", {}).get("python", '/nonexistent')) + if not math_only: + checks.append({"path": "harbor.python", "ok": runner.is_file()}) from harborrl.rollout.harbor_job.coordinator import runner_command - if legacy_workers and runner.is_file(): + if not math_only and legacy_workers and runner.is_file(): workers = ([None] if config['harbor'].get('runner_transport') == 'local' else config['harbor']['workers']) for worker in workers: @@ -87,10 +122,16 @@ def doctor(plan): checks.append({'service': 'sandbox offline runtime', 'worker': worker, 'ok': False, 'detail': type(exc).__name__}) - for key, path in (("tasks.catalog", config["tasks"]["catalog"]), - ("harness.profile", config["harness"]["profile"]), - ("model.checkpoint", config["model"]["checkpoint"]), - ("model.reference", config["model"]["reference"])): + paths = [("model.checkpoint", config["model"]["checkpoint"]), + ("model.reference", config["model"]["reference"])] + if not math_only: + paths.extend([("tasks.catalog", config["tasks"]["catalog"]), + ("harness.profile", config["harness"]["profile"])]) + if config['training'].get('resume'): + paths.append(('training.resume', config['training']['resume'])) + if math_only and config['algorithm']['name'] == 'opd': + paths.append(('teacher.checkpoint', config['teacher']['checkpoint'])) + for key, path in paths: checks.append({"path": key, "ok": Path(path).exists()}) model_config = Path(config["model"]["checkpoint"]) / "config.json" @@ -108,7 +149,7 @@ def doctor(plan): from harborrl.tasks.catalog import catalog try: - entries = catalog(config) + entries = [] if math_only else catalog(config) for entry in entries: # catalog() already checks the complete task tree against its lock. checks.append({"task": entry["id"], "ok": True, "detail": []}) @@ -200,7 +241,18 @@ def doctor(plan): if config["training"]["backend_contract"] != "slime-legacy": from harborrl.backends.slime_v032.launcher import hook_modules - for hook_name, module_name in hook_modules().items(): + modules = hook_modules() + if math_only: + modules.update(rollout='harborrl.backends.slime_v032.math_rollout', + converter='harborrl.backends.slime_v032.math_rollout') + if config['algorithm']['name'] == 'opd': + modules.update(advantage='harborrl.backends.slime_v032.opd', + loss='harborrl.backends.slime_v032.opd') + elif config['algorithm']['name'] == 'dapo': + modules.update(loss='harborrl.backends.slime_v032.dapo') + elif 'process_reward' in config: + modules.update(advantage='harborrl.backends.slime_v032.prm') + for hook_name, module_name in modules.items(): try: importlib.import_module(module_name) checks.append({"service": "training backend hook", "hook": hook_name, "ok": True}) @@ -214,7 +266,7 @@ def doctor(plan): "service": "raw-model logprob audit", "ok": True, "detail": ("enabled by lock; live serving semantic probe remains a GPU gate" - if config["gateway"]["audited_raw_logprobs"] else + if math_only or config["gateway"]["audited_raw_logprobs"] else "disabled; generated evidence remains EVAL_ONLY and training will fail closed"), }) return checks diff --git a/harborrl/platform/slime_train.sh b/harborrl/platform/slime_train.sh index db00a0948..2ea9ed238 100644 --- a/harborrl/platform/slime_train.sh +++ b/harborrl/platform/slime_train.sh @@ -65,17 +65,19 @@ CKPT_ARGS=( --megatron-to-hf-mode bridge --save "${CKPT_ROOT}/${RUN_ID:-run}" --save-interval "${SAVE_INTERVAL}" - --no-save-optim --max-ckpt-keep "${MAX_CKPT_KEEP:-2}" --checkpoint-min-free-gb "${CHECKPOINT_MIN_FREE_GB:-128}" --checkpoint-expected-gb "${CHECKPOINT_EXPECTED_GB:-0}" --checkpoint-space-margin-ratio "${CHECKPOINT_SPACE_MARGIN_RATIO:-1.15}" ) if [[ -n "${RESUME_LOAD:-}" ]]; then - CKPT_ARGS+=(--load "${RESUME_LOAD}" --no-load-optim) + CKPT_ARGS+=(--load "${RESUME_LOAD}" --use-checkpoint-opt-param-scheduler) else CKPT_ARGS+=(--load "${HF_CKPT}") fi +if [[ -n "${SAVE_HF:-}" ]]; then + CKPT_ARGS+=(--save-hf "${SAVE_HF}") +fi DATA_ARGS=( --prompt-data "${ROLLOUT_PROMPT_DATA}" @@ -146,9 +148,6 @@ PARALLEL_ARGS=( --attention-backend flash --no-gradient-accumulation-fusion ) -if [[ "${HARBORRL_GPU_LAYOUT:-split}" == "colocate" ]]; then - PARALLEL_ARGS+=(--colocate) -fi CUSTOM_ARGS=( --rollout-function-path harborrl.backends.slime_v032.rollout.generate_rollout @@ -158,6 +157,45 @@ CUSTOM_ARGS=( --custom-loss-function-path harborrl.backends.slime_v032.loss.loss_function --rollout-data-postprocess-path harborrl.backends.slime_v032.postprocess.rollout_data_postprocess ) +TRAINING_SEED="${HARBORRL_TRAINING_SEED:-${HARBORRL_MATH_SEED:-}}" +if [[ -n "${TRAINING_SEED}" ]]; then + [[ "${TRAINING_SEED}" =~ ^[0-9]+$ && ${#TRAINING_SEED} -le 10 ]] || die "training seed must be an integer in 0..2147483647" + (( 10#${TRAINING_SEED} <= 2147483647 )) || die "training seed must be an integer in 0..2147483647" + DATA_ARGS+=(--seed "${TRAINING_SEED}" --rollout-seed "${TRAINING_SEED}" + --sglang-enable-deterministic-inference) +fi +if [[ "${HARBORRL_EXECUTION_BACKEND:-harbor_job}" == "text_math" ]]; then + ADVANTAGE_MODULE=harborrl.backends.slime_v032.advantage + LOSS_MODULE=harborrl.backends.slime_v032.loss + case "${HARBORRL_MATH_ALGORITHM:-grpo}" in + grpo) ;; + dapo) + : "${HARBORRL_DAPO:?DAPO requires explicit settings}" + LOSS_MODULE=harborrl.backends.slime_v032.dapo + mapfile -t DAPO_CLIP < <("${TRAIN_PYTHON:-python3}" -c \ + 'import json,os; c=json.loads(os.environ["HARBORRL_DAPO"]); print(c["clip_low"]); print(c["clip_high"])') + OPT_ARGS+=(--calculate-per-token-loss --eps-clip "${DAPO_CLIP[0]}" --eps-clip-high "${DAPO_CLIP[1]}") + ;; + opd) + : "${HARBORRL_OPD_TEACHER:?OPD requires a frozen teacher manifest}" + ADVANTAGE_MODULE=harborrl.backends.slime_v032.opd + LOSS_MODULE=harborrl.backends.slime_v032.opd + ;; + *) die "Unsupported math algorithm: ${HARBORRL_MATH_ALGORITHM}" ;; + esac + if [[ -n "${HARBORRL_PROCESS_REWARD:-}" ]]; then + [[ "${HARBORRL_MATH_ALGORITHM:-grpo}" == "grpo" ]] || die "PRM currently requires GRPO" + ADVANTAGE_MODULE=harborrl.backends.slime_v032.prm + fi + CUSTOM_ARGS=( + --rollout-function-path harborrl.backends.slime_v032.math_rollout.generate_rollout + --custom-convert-samples-to-train-data-path harborrl.backends.slime_v032.math_rollout.convert_samples_to_train_data + --custom-advantage-function-path "${ADVANTAGE_MODULE}.compute_advantages_and_returns" + --loss-type custom_loss + --custom-loss-function-path "${LOSS_MODULE}.loss_function" + --rollout-data-postprocess-path harborrl.backends.slime_v032.postprocess.rollout_data_postprocess + ) +fi if [[ -f "${RUN_DIR}/config/rollout_config.yaml" ]]; then CUSTOM_ARGS+=(--custom-config-path "${RUN_DIR}/config/rollout_config.yaml") fi diff --git a/harborrl/rollout/harbor_job/audit.py b/harborrl/rollout/harbor_job/audit.py index bc93f75a3..693cd42e1 100644 --- a/harborrl/rollout/harbor_job/audit.py +++ b/harborrl/rollout/harbor_job/audit.py @@ -11,13 +11,121 @@ import json from pathlib import Path +from harborrl.trajectories.native import digest, read_json -def audit_session(path, registry, attempt_id): +_CONTINUATION = ('Output token limit hit. Resume directly — no apology, no recap of what you were doing. ' + 'Pick up mid-thought if that is where the cut happened. Break remaining work into smaller pieces.') + + +def _message_content(messages): + """Ignore only probed cache placement and string/text-block serialization.""" + result = [] + for message in messages: + if not isinstance(message, dict): + raise ValueError('turn budget continuation lacks native message content') # noqa: TRY004 - audit rejects malformed evidence + content = message.get('content') + if isinstance(content, str): + content = [{'type': 'text', 'text': content}] + if not isinstance(content, list) or any(not isinstance(block, dict) for block in content): + raise ValueError('turn budget continuation lacks native message content') + blocks = [] + for block in content: + cache = block.get('cache_control') + if 'cache_control' in block and cache != {'type': 'ephemeral'}: + raise ValueError('turn budget continuation has unprobed cache metadata') + blocks.append({key: value for key, value in block.items() if key != 'cache_control'}) + result.append({**message, 'content': blocks}) + return result + + +def _budget_continuations(turns, budget): + """Every extra response must be a consumed length stop with a linked retry.""" + if not turns or turns[-1].get('finish_reason') != 'tool_use': + raise ValueError('turn budget differs from sealed policy trace') + proofs = [] + for index, turn in enumerate(turns): + if turn.get('finish_reason') == 'tool_use': + continue + response = turn.get('response', {}) + if (turn.get('finish_reason') != 'max_tokens' or response.get('stop_reason') != 'max_tokens' + or index + 1 == len(turns)): + raise ValueError('turn budget includes an unverified extra response') + content = response.get('content') + if (not isinstance(content, list) or not content + or any(not isinstance(block, dict) or block.get('type') != 'text' for block in content)): + raise ValueError('turn budget continuation requires a text-only length stop') + following = turns[index + 1] + before = turn.get('client_request', {}).get('messages') + after = following.get('client_request', {}).get('messages') + if not isinstance(before, list) or not before or not isinstance(after, list): + raise ValueError('turn budget continuation lacks request history') + expected = [*before, {'role': 'assistant', 'content': content}, + {'role': 'user', 'content': [{'type': 'text', 'text': _CONTINUATION}]}] + if _message_content(after) != _message_content(expected): + raise ValueError('turn budget continuation differs from the consumed length stop') + for item in (turn, following): + for key in ('client_request', 'response'): + if item.get(key + '_digest') != digest(item.get(key)): + raise ValueError('turn budget continuation has conflicting trace digests') + proofs.append({'response_id': turn['response_id'], 'next_response_id': following['response_id'], + 'response_digest': turn['response_digest'], + 'continuation_request_digest': following['client_request_digest']}) + if len(turns) - len(proofs) != budget: + raise ValueError('turn budget differs from sealed policy trace') + return proofs + + +def audit_termination(path, result_path, profile, audit, *, turns): + """Recognize the probed CLI turn-budget exit from matching native evidence.""" + result_path = Path(result_path) + result = read_json(result_path) + exception = result.get('exception_info') + if not exception: + return None + if (profile['cli_version'] != '2.1.141' + or exception.get('exception_type') != 'NonZeroAgentExitCodeError'): + raise ValueError('exceptional trial requires a verified termination profile') + path = Path(path) + raw = path.read_bytes() + events = [json.loads(line) for line in raw.splitlines() if line.strip()] + results = [event for event in events if event.get('type') == 'result'] + inits = [event for event in events if event.get('type') == 'system' and event.get('subtype') == 'init'] + if len(results) != 1 or events[-1] != results[0] or len(inits) != 1: + raise ValueError('turn budget requires one native init and final result') + final, initial = results[0], inits[0] + budget = profile['max_turns'] + expected = {'subtype': 'error_max_turns', 'terminal_reason': 'max_turns', + 'stop_reason': 'tool_use', 'session_id': audit['session_id'], + 'errors': [f'Reached maximum number of turns ({budget})']} + if (any(final.get(key) != value for key, value in expected.items()) + or final.get('is_error') is not True or final.get('permission_denials') != [] + or type(final.get('num_turns')) is not int or final['num_turns'] != budget + 1 + or initial.get('session_id') != audit['session_id'] + or initial.get('claude_code_version') != profile['cli_version']): + raise ValueError('turn budget lacks matching CLI termination evidence') + responses = [event for event in events if event.get('type') == 'assistant'] + ids = list(dict.fromkeys(event.get('message', {}).get('id') for event in responses)) + if (ids != audit['consumed_response_ids'] or ids != [turn['response_id'] for turn in turns] + or any(event.get('session_id') != audit['session_id'] for event in responses) + or not ids or not all(isinstance(value, str) and value for value in ids)): + raise ValueError('turn budget differs from consumed native responses') + continuations = _budget_continuations(turns, budget) + return {'profile': 'claude_code_2.1.141_max_turns_v1', 'reason': 'max_turns', + 'session_id': audit['session_id'], 'max_turns': budget, + 'policy_response_count': len(turns), 'logical_turns': budget, 'continuations': continuations, + 'consumed_response_ids': ids, 'session_evidence_ref': audit['evidence_ref'], + 'source_result_digest': hashlib.sha256(result_path.read_bytes()).hexdigest(), + 'cli_stream_ref': str(path), 'cli_stream_digest': hashlib.sha256(raw).hexdigest()} + + +def audit_session(path, registry, attempt_id, *, cli_version=None): path = Path(path) raw = path.read_bytes() evidence = f"{path}#sha256={hashlib.sha256(raw).hexdigest()}" sessions = set() + tool_results = [] merged = {} # message id -> merged content blocks, in first-seen order + versions = {} for line in raw.splitlines(): if not line.strip(): continue @@ -26,12 +134,18 @@ def audit_session(path, registry, attempt_id): raise ValueError("auxiliary generation or compact is not enabled") if event.get("sessionId"): sessions.add(event["sessionId"]) + if event.get('type') == 'user': + content = event.get('message', {}).get('content') + if isinstance(content, list): + tool_results.extend(block for block in content + if isinstance(block, dict) and block.get('type') == 'tool_result') if event.get("type") != "assistant": continue message = event.get("message", {}) if message.get("role") != "assistant" or not isinstance(message.get("content"), list) or not message.get("id"): raise ValueError("assistant consumption lacks ID or content") rid, content = message["id"], message["content"] + versions.setdefault(rid, set()).add(event.get('version')) if rid not in merged: merged[rid] = list(content) continue @@ -46,6 +160,8 @@ def audit_session(path, registry, attempt_id): if len(sessions) != 1 or not records: raise ValueError("expected one complete native session") # Validate the complete file before mutating the registry; seal catches omitted responses. + registry.audit_tool_rejections(attempt_id, tool_results, evidence) for rid, content in records: - registry.consume(attempt_id, rid, content, evidence) + verified_version = cli_version if versions[rid] == {cli_version} else None + registry.consume(attempt_id, rid, content, evidence, cli_version=verified_version) return {"session_id": next(iter(sessions)), "consumed_response_ids": list(merged), "evidence_ref": evidence} diff --git a/harborrl/rollout/harbor_job/bindings.py b/harborrl/rollout/harbor_job/bindings.py index d0da81703..e9ed2f5bd 100644 --- a/harborrl/rollout/harbor_job/bindings.py +++ b/harborrl/rollout/harbor_job/bindings.py @@ -1,25 +1,34 @@ """Harbor 0.23 Claude Code configuration. No import of Harbor in the parent.""" +import math import re from urllib.parse import urlsplit def claude_config(profile, gateway_url, credential, max_output_tokens=8192): required = {"cli_version", "model", "max_turns", "agent_timeout_sec", "setup_timeout_sec", "verifier_timeout_sec"} - allowed = required | {"disallowed_tools"} + optional = {"disallowed_tools", "append_system_prompt", "environment_timeout_sec"} + allowed = required | optional if not required.issubset(profile) or set(profile) - allowed: - raise ValueError(f"Claude profile requires {sorted(required)} and optionally disallowed_tools") + raise ValueError(f"Claude profile requires {sorted(required)} and optionally {sorted(optional)}") if not isinstance(profile["cli_version"], str) or not re.fullmatch(r"\d+\.\d+\.\d+(?:-[A-Za-z0-9.]+)?", profile["cli_version"]): raise ValueError("exact Claude CLI version required; latest is forbidden") if not isinstance(profile["model"], str) or not re.fullmatch(r"[A-Za-z0-9_.:/-]+", profile["model"]): raise ValueError("explicit policy model alias required") - for key in required - {"model", "cli_version"}: + for key in (required - {"model", "cli_version"}) | ({'environment_timeout_sec'} & profile.keys()): if type(profile[key]) is not int or profile[key] <= 0: raise ValueError(f"positive integer required: {key}") disallowed_tools = profile.get("disallowed_tools", "Agent,Task") if not isinstance(disallowed_tools, str) or not re.fullmatch( r"[A-Za-z][A-Za-z0-9_]*(?:,[A-Za-z][A-Za-z0-9_]*)*", disallowed_tools): raise ValueError("disallowed tools must be a comma-separated identifier list") - disallowed_tools = ",".join(dict.fromkeys(["Agent", "Task", *disallowed_tools.split(",")])) + kwargs = {"version": profile["cli_version"], "max_turns": profile["max_turns"], + "max_thinking_tokens": 0, "disallowed_tools": disallowed_tools, + "permission_mode": "bypassPermissions"} + if "append_system_prompt" in profile: + prompt = profile["append_system_prompt"] + if not isinstance(prompt, str) or not prompt.strip() or "\x00" in prompt: + raise ValueError("append_system_prompt must be a nonempty string without NUL") + kwargs["append_system_prompt"] = prompt url = urlsplit(gateway_url) if url.scheme not in ("http", "https") or not url.hostname or url.username or url.password or url.query or url.fragment or url.path not in ("", "/"): raise ValueError("Anthropic base URL must be a plain origin (no /v1 suffix)") @@ -32,9 +41,7 @@ def claude_config(profile, gateway_url, credential, max_output_tokens=8192): return {"name": "claude-code", "model_name": profile["model"], "override_timeout_sec": profile["agent_timeout_sec"], "override_setup_timeout_sec": profile["setup_timeout_sec"], - "kwargs": {"version": profile["cli_version"], "max_turns": profile["max_turns"], - "max_thinking_tokens": 0, "disallowed_tools": disallowed_tools, - "permission_mode": "bypassPermissions"}, + "kwargs": kwargs, "env": {**client_controls(), "ANTHROPIC_API_KEY": credential, "ANTHROPIC_MODEL": profile["model"], "ANTHROPIC_BASE_URL": gateway_url.rstrip("/"), @@ -42,13 +49,24 @@ def claude_config(profile, gateway_url, credential, max_output_tokens=8192): "CLAUDE_CODE_MAX_OUTPUT_TOKENS": str(max_output_tokens)}} +def environment_timeouts(profile, resources): + """Use Harbor's public multiplier without changing the original task config.""" + if 'environment_timeout_sec' not in profile: + return {} + budget = profile['environment_timeout_sec'] + base = resources.get('build_timeout_sec', 600.0) # Harbor 0.23 default. + if type(budget) is not int or budget <= 0: + raise ValueError('positive integer required: environment_timeout_sec') + if type(base) not in (int, float) or not math.isfinite(base) or base <= 0: + raise ValueError('positive finite task environment build timeout required') + return {'environment_build_timeout_multiplier': budget / base} + + def client_controls(): """Pass controls into the Agent environment, not only the Runner host.""" - # Disabling prompt caching also prevents cache-warming model requests. return {"CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1", "CLAUDE_CODE_DISABLE_ADAPTIVE_THINKING": "1", "CLAUDE_CODE_DISABLE_AUTO_MEMORY": "1", - "DISABLE_PROMPT_CACHING": "1", "DISABLE_COMPACT": "1", "DISABLE_AUTOUPDATER": "1", "DISABLE_AUTO_COMPACT": "1"} diff --git a/harborrl/rollout/harbor_job/collector.py b/harborrl/rollout/harbor_job/collector.py index 8dfaeda89..c71313dfe 100644 --- a/harborrl/rollout/harbor_job/collector.py +++ b/harborrl/rollout/harbor_job/collector.py @@ -5,7 +5,8 @@ from harborrl.trajectories.native import publish, read_json -def collect(identity, trial_id, result_path, verifier_dir, profile, *, evidence_kind="harbor"): +def collect(identity, trial_id, result_path, verifier_dir, profile, *, evidence_kind="harbor", + termination_evidence=None): if profile.digest != identity.reward_digest: raise ValueError("reward profile differs from attempt manifest") result_path, verifier_dir = Path(result_path), Path(verifier_dir) @@ -15,8 +16,15 @@ def collect(identity, trial_id, result_path, verifier_dir, profile, *, evidence_ raise ValueError("missing terminal result or mismatched trial") if result.get("step_results") or result.get("verifier_environment_mode") != "shared": raise ValueError("only single-step shared verifier is enabled") + result_digest = hashlib.sha256(result_path.read_bytes()).hexdigest() if result.get("exception_info"): - raise ValueError("exceptional trial requires a separately verified termination profile") + if (not isinstance(termination_evidence, dict) + or termination_evidence.get('profile') != 'claude_code_2.1.141_max_turns_v1' + or termination_evidence.get('source_result_digest') != result_digest + or result['exception_info'].get('exception_type') != 'NonZeroAgentExitCodeError'): + raise ValueError("exceptional trial requires a separately verified termination profile") + elif termination_evidence is not None: + raise ValueError('termination evidence contradicts normal trial result') rewards = (result.get("verifier_result") or {}).get("rewards") artifact = verifier_dir / "reward.json" if artifact.exists(): @@ -33,9 +41,10 @@ def collect(identity, trial_id, result_path, verifier_dir, profile, *, evidence_ "evaluation_trial_id": trial_id, "source_trial_id": trial_id, "verifier_environment_mode": "shared", "evaluation_status": "valid", "evidence_kind": evidence_kind, "source_result_ref": str(result_path), - "source_result_digest": hashlib.sha256(result_path.read_bytes()).hexdigest(), + "source_result_digest": result_digest, "source_artifact_ref": str(artifact), - "source_artifact_digest": hashlib.sha256(artifact.read_bytes()).hexdigest()} + "source_artifact_digest": hashlib.sha256(artifact.read_bytes()).hexdigest(), + **({'termination_evidence': termination_evidence} if termination_evidence else {})} def commit_receipt(path, *args, **kwargs): diff --git a/harborrl/rollout/harbor_job/coordinator.py b/harborrl/rollout/harbor_job/coordinator.py index 7316d5bba..ad6131087 100644 --- a/harborrl/rollout/harbor_job/coordinator.py +++ b/harborrl/rollout/harbor_job/coordinator.py @@ -82,27 +82,23 @@ async def run_attempt(identity, *, task_path, profile, gateway_url, runner_pytho credential = secrets.token_urlsafe(32) project = str(Path(__file__).resolve().parents[3]) env = CLAUDE_CODE_HARNESS.adapter.runner_environment(os.environ) - env["PYTHONPATH"] = project + env["PYTHONPATH"] = os.pathsep.join(filter(None, (project, env.get('PYTHONPATH')))) stderr = (root / "runner.stderr").open("wb") try: command = runner_command(runner_python, project, runner_host) process = await asyncio.create_subprocess_exec(*command, stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=stderr, env=env) - except BaseException: + except BaseException as exc: stderr.close() - publish(root / "launch-failed.json", {"identity": identity.to_dict(), "reason": "Runner process did not start"}) + # A cancelled group kills sibling attempts while their runner process + # is still starting; that is not a launch failure and must not read + # as one in the shared evidence. + reason = "Runner launch cancelled" if isinstance(exc, asyncio.CancelledError) \ + else "Runner process did not start" + publish(root / "launch-failed.json", {"identity": identity.to_dict(), "reason": reason}) raise - registered, trial_id, lifecycle_timeout, sealed = False, None, False, False + registered, trial_id, lifecycle_timeout = False, None, False messages = [] - - def mark_cleanup(reason): - marker = root / "cleanup-required.json" - if not marker.exists(): - publish(marker, { - "identity": identity.to_dict(), "trial_id": trial_id, - "runner_host": runner_host, "reason": reason, - }) - try: req = {"identity": identity.to_dict(), "task_path": str(Path(task_path).resolve()), "profile": profile, "gateway_url": gateway_url, "credential": credential, "output": str(root / "runner"), @@ -110,9 +106,13 @@ def mark_cleanup(reason): if sandbox is not None: from harborrl.trajectories.native import digest req.update(sandbox=sandbox, runtime_digest=digest(sandbox)) - timeout = (sandbox['create_timeout_s'] + sandbox['cleanup_timeout_s'] * 2 + timeout = (max(sandbox['create_timeout_s'], profile.get('environment_timeout_sec', 0)) + + sandbox['cleanup_timeout_s'] * 2 + sum(profile[k] for k in ('setup_timeout_sec', 'agent_timeout_sec', 'verifier_timeout_sec'))) cleanup_timeout = sandbox['cleanup_timeout_s'] + 15 + elif 'environment_timeout_sec' in profile: + timeout = max(timeout, profile['environment_timeout_sec'] + cleanup_timeout * 2 + + sum(profile[k] for k in ('setup_timeout_sec', 'agent_timeout_sec', 'verifier_timeout_sec'))) process.stdin.write(json.dumps(req).encode() + b"\n") await process.stdin.drain() @@ -151,7 +151,7 @@ async def lifecycle(): lifecycle_path = root / 'runner' / 'sandbox-lifecycle.json' lifecycle = read_json(lifecycle_path) if lifecycle_path.is_file() else {} if lifecycle.get('state') != 'deleted' or lifecycle.get('runtime_digest') != req['runtime_digest']: - mark_cleanup("Sandbox deletion unconfirmed") + publish(root / 'cleanup-required.json', {'identity': identity.to_dict(), 'reason': 'Sandbox deletion unconfirmed'}) raise RuntimeError('sandbox deletion unconfirmed') trial_dir = Path(terminal["trial_dir"]).resolve() result_path = Path(terminal["result_path"]).resolve() @@ -159,16 +159,26 @@ async def lifecycle(): raise ValueError("Runner result path is outside this attempt") sessions = list((trial_dir / "agent" / "sessions" / "projects").rglob("*.jsonl")) if len(sessions) != 1: - raise ValueError("expected exactly one native Claude session; auxiliary sessions require review") + from harborrl.trajectories.native import read_json + exception = read_json(result_path).get('exception_info') + if not sessions and exception: + error_type = exception.get('exception_type', 'unknown') + raise RuntimeError(f'native trial failed before session audit: {error_type}; inspect Harbor result and Sandbox lifecycle') + raise ValueError(f"expected exactly one native Claude session, found {len(sessions)}; auxiliary sessions require review") registry.wait_delivered(identity.attempt_id, 30.0) - audit = CLAUDE_CODE_HARNESS.adapter.audit_session(sessions[0], registry, identity.attempt_id) - receipt = collect(identity, trial_id, result_path, trial_dir / "verifier", reward_profile) + audit = CLAUDE_CODE_HARNESS.adapter.audit_session( + sessions[0], registry, identity.attempt_id, cli_version=profile['cli_version']) CLAUDE_CODE_HARNESS.adapter.audit_empty_final( trial_dir / "agent" / "claude-code.txt", registry, identity.attempt_id, audit) turns, seal = registry.seal(identity.attempt_id) - sealed = True + termination_evidence = CLAUDE_CODE_HARNESS.adapter.audit_termination( + trial_dir / 'agent/claude-code.txt', result_path, profile, audit, turns=turns) + receipt = collect(identity, trial_id, result_path, trial_dir / "verifier", reward_profile, + termination_evidence=termination_evidence) publish(root / "evaluation.json", receipt) - ir = assemble(identity, trial_id, turns, seal, receipt) + ir = assemble(identity, trial_id, turns, seal, receipt, + termination='budget_truncated' if termination_evidence else 'completed', + termination_evidence=termination_evidence) publish(root / "trajectory.v2.json", ir) return ir except BaseException as exc: @@ -176,8 +186,9 @@ async def lifecycle(): publish(root / "failure.json", {"identity": identity.to_dict(), "error_type": type(exc).__name__}) raise finally: - if (root / "runner" / "cleanup-required.json").exists(): - mark_cleanup("Inspect runner/cleanup-required.json") + if (sandbox is not None and (root / 'runner' / 'cleanup-required.json').exists() + and not (root / 'cleanup-required.json').exists()): + publish(root / 'cleanup-required.json', {'identity': identity.to_dict(), 'reason': 'Inspect runner/cleanup-required.json'}) if registered: registry.close(identity.attempt_id) if process.returncode is None: @@ -191,10 +202,13 @@ async def lifecycle(): except ProcessLookupError: pass await process.wait() - mark_cleanup("Runner cleanup deadline exceeded; inspect Harbor resources before retry") - if process.returncode != 0 or lifecycle_timeout: - mark_cleanup("Runner exited abnormally; inspect Harbor resources before retry") - if registered and not sealed: + if not (root / "cleanup-required.json").exists(): + publish(root / "cleanup-required.json", {"identity": identity.to_dict(), "trial_id": trial_id, + "reason": "Runner cleanup deadline exceeded; inspect Harbor resources before retry"}) + if (process.returncode != 0 or lifecycle_timeout) and not (root / "cleanup-required.json").exists(): + publish(root / "cleanup-required.json", {"identity": identity.to_dict(), "trial_id": trial_id, + "reason": "Runner exited abnormally; inspect Harbor resources before retry"}) + if registered: try: registry.seal(identity.attempt_id) except ValueError: diff --git a/harborrl/rollout/harbor_job/environments/opensandbox.py b/harborrl/rollout/harbor_job/environments/opensandbox.py index 829f9dadc..436b913a5 100644 --- a/harborrl/rollout/harbor_job/environments/opensandbox.py +++ b/harborrl/rollout/harbor_job/environments/opensandbox.py @@ -41,6 +41,11 @@ async def startup_slot(directory, count): class ManagedOpenSandboxEnvironment(OpenSandboxEnvironment): + def _validate_definition(self): + # The deployment binding selects the built image without rewriting task files. + self.task_env_config = self.task_env_config.model_copy(update={'docker_image': self.binding['image']}) + super()._validate_definition() + @classmethod def validate_config(cls, cfg): if cfg['options'] or any(b['options'] for b in cfg['bindings']): @@ -67,7 +72,8 @@ def startup_transport(self, allocated): from harborrl.rollout.harbor_job.sandbox_transport import ( SandboxStartupTransport, ) - return SandboxStartupTransport(self.cfg['domain'].rstrip('/') + '/v1', allocated) + return SandboxStartupTransport(self.cfg['domain'].rstrip('/') + '/v1', allocated, + on_request=self.record_request) def __init__(self, *args, sandbox_config, binding, evidence_dir, startup_dir, gateway_url=None, **kwargs): self.validate_config(sandbox_config) @@ -76,6 +82,8 @@ def __init__(self, *args, sandbox_config, binding, evidence_dir, startup_dir, ga self.startup_dir = startup_dir self.gateway_url = gateway_url self._startup_transport = None + self._startup_reservation = None + self._startup_requests = 0 self.record = {'runtime_digest': digest(sandbox_config), 'image': binding['image'], 'sandbox_id': None, 'state': 'pending', 'events': []} super().__init__(*args, domain=self.cfg['domain'], api_key=read_credential(self.cfg), @@ -99,9 +107,51 @@ def cleanup_required(self, reason): 'sandbox_id': self.record['sandbox_id'], 'image': self.binding['image'], 'reason': reason, 'runtime_digest': self.record['runtime_digest']}) + def record_request(self, record): + self._startup_requests += 1 + publish(self.evidence / 'sandbox-requests' / f'{self._startup_requests:04d}.json', record) + + async def reserve_startup(self): + """The Runner reserves before Harbor starts its environment timeout.""" + if self._startup_reservation is not None: + raise RuntimeError('Sandbox startup slot already reserved') + self.record_state('waiting') + reservation = startup_slot(self.startup_dir, self.cfg['startup_concurrency']) + await reservation.__aenter__() + self._startup_reservation = reservation + try: + self.record_state('reserved') + except BaseException: + await self.release_startup() + raise + + async def release_startup(self): + reservation, self._startup_reservation = self._startup_reservation, None + if reservation is not None: + await reservation.__aexit__(None, None, None) + + def startup_failed(self, stage, exc): + """Keep exception classes/status only; SDK messages may contain credentials.""" + causes, seen = [], set() + while exc is not None and id(exc) not in seen: + seen.add(id(exc)) + row = {'error_type': type(exc).__name__} + status = getattr(exc, 'status_code', None) + if type(status) is int: + row['http_status'] = status + causes.append(row) + exc = exc.__cause__ or exc.__context__ + publish(self.evidence / 'sandbox-startup-failed.json', { + 'stage': stage, 'sandbox_id': self.record['sandbox_id'], 'causes': causes, + 'create_timeout_s': self.cfg['create_timeout_s'], + 'request_timeout_s': self.cfg['request_timeout_s']}) + async def _create_sandbox(self, sdk): - await self.validate_binding() - async with startup_slot(self.startup_dir, self.cfg['startup_concurrency']): + # Direct environment callers retain the same process-shared limit. + if self._startup_reservation is None: + await self.reserve_startup() + try: + await self.validate_binding() self.record_state('creating') def allocated(sandbox_id): self.record['sandbox_id'] = sandbox_id @@ -118,7 +168,8 @@ def allocated(sandbox_id): resource=self._build_resource(), env=dict(self._persistent_env), network_policy=self._build_network_policy(sdk), skip_health_check=True), self.cfg['create_timeout_s']) - except BaseException: # noqa: BLE001 - persist unknown allocation on cancellation too + except BaseException as exc: # noqa: BLE001 - persist unknown allocation on cancellation too + self.startup_failed('create', exc) if self.record['sandbox_id']: await self._cleanup_allocated() else: @@ -133,9 +184,12 @@ def allocated(sandbox_id): await asyncio.wait_for(self._wait_until_ready(sandbox), self.cfg['create_timeout_s']) self.record_state('ready') return sandbox - except BaseException: # noqa: BLE001 - cleanup before redacting SDK errors + except BaseException as exc: # noqa: BLE001 - cleanup before redacting SDK errors + self.startup_failed('readiness', exc) await self.stop(delete=True) raise RuntimeError('Sandbox readiness failed') from None + finally: + await self.release_startup() async def _cleanup_allocated(self): """The SDK may already have killed an instance whose endpoints failed.""" @@ -163,6 +217,8 @@ async def start(self, force_build): if self._sandbox is not None: await self.stop(delete=True) raise + finally: + await self.release_startup() async def stop(self, delete=True): if self._sandbox is None: diff --git a/harborrl/rollout/harbor_job/runner.py b/harborrl/rollout/harbor_job/runner.py index 7aa23fe4a..19cfc59ee 100644 --- a/harborrl/rollout/harbor_job/runner.py +++ b/harborrl/rollout/harbor_job/runner.py @@ -15,6 +15,7 @@ from pathlib import Path from harborrl.math_registry import CLAUDE_CODE_HARNESS, NATIVE_TASK_PROVIDER +from harborrl.rollout.harbor_job.bindings import environment_timeouts from harborrl.trajectories.native import Identity, digest, publish @@ -81,6 +82,7 @@ async def execute(request, emit, read_command): "trial_name": digest(identity.to_dict()), "trials_dir": str(root / "harbor"), "agent": agent, "environment": environment, + **environment_timeouts(request['profile'], inspection['resources']), "verifier": {"disable": False, "override_timeout_sec": request["profile"]["verifier_timeout_sec"]}}) trial = await Trial.create(config) trial_id = str(trial.id) @@ -90,6 +92,9 @@ async def execute(request, emit, read_command): raise ValueError("expected run handshake before starting trial") async def event_hook(event): + if sandbox is not None and event.event == TrialEvent.ENVIRONMENT_START: + # Harbor awaits this public hook before starting its environment timer. + await trial.agent_environment.reserve_startup() record = {"type": "event", "event": event.event.value, "trial_id": str(event.trial_id), "timestamp": event.timestamp.isoformat()} # Each event is an independent atomic record; never serialize config credentials. @@ -98,7 +103,15 @@ async def event_hook(event): for event in TrialEvent: trial.add_hook(event, event_hook) - running = asyncio.create_task(trial.run()) + + async def run(): + try: + return await trial.run() + finally: + if sandbox is not None: + await trial.agent_environment.release_startup() + + running = asyncio.create_task(run()) loop = asyncio.get_running_loop() for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, running.cancel) diff --git a/harborrl/rollout/harbor_job/sandbox_config.py b/harborrl/rollout/harbor_job/sandbox_config.py index 063cce41e..1d52b7909 100644 --- a/harborrl/rollout/harbor_job/sandbox_config.py +++ b/harborrl/rollout/harbor_job/sandbox_config.py @@ -9,7 +9,7 @@ import yaml -from harborrl.tasks.content import tomllib +from harborrl.tasks.content import tomllib, tree_digest def _check_options(value): @@ -58,8 +58,9 @@ def load_sandbox(path): if not isinstance(bindings, list) or not bindings: raise ValueError('sandbox image bindings required') images = set() + task_digests = set() for binding in bindings: - if not isinstance(binding, dict) or set(binding) - {'image', 'max_live_instances', 'options'} or not {'image', 'max_live_instances'} <= set(binding): + if not isinstance(binding, dict) or set(binding) - {'image', 'max_live_instances', 'options', 'task_digests'} or not {'image', 'max_live_instances'} <= set(binding): raise ValueError('invalid sandbox binding fields') binding.setdefault('options', {}) _check_options(binding['options']) @@ -68,6 +69,15 @@ def load_sandbox(path): raise ValueError('invalid or duplicate sandbox image binding') if type(binding['max_live_instances']) is not int or binding['max_live_instances'] <= 0: raise ValueError('binding max_live_instances must be positive') + if 'task_digests' in binding: + digests = binding['task_digests'] + if (not isinstance(digests, list) or not digests or any( + not isinstance(value, str) or not re.fullmatch(r'[0-9a-f]{64}', value) + for value in digests)): + raise ValueError('binding task_digests must be a nonempty list of SHA256 digests') + if len(set(digests)) != len(digests) or task_digests.intersection(digests): + raise ValueError('duplicate sandbox task digest binding') + task_digests.update(digests) images.add(binding['image']) install = cfg['installation'] if install is None: @@ -84,10 +94,15 @@ def load_sandbox(path): def task_binding(cfg, task_path): + if any(binding.get('task_digests') for binding in cfg['bindings']): + task_digest = tree_digest(Path(task_path)) + for binding in cfg['bindings']: + if task_digest in binding.get('task_digests', ()): + return binding spec = tomllib.loads((Path(task_path) / 'task.toml').read_text()) image = spec.get('environment', {}).get('docker_image') for binding in cfg['bindings']: - if image == binding['image']: + if image == binding['image'] and not binding.get('task_digests'): return binding raise ValueError('sandbox task requires a prebuilt image with an exact binding') @@ -120,12 +135,7 @@ class SandboxAdmission: def __init__(self, cfg): self.total = asyncio.Semaphore(cfg['max_live_instances']) self.bindings = {b['image']: asyncio.Semaphore(b['max_live_instances']) for b in cfg['bindings']} - self.cleanup_markers = set() - - @property - def blocked(self): - self.cleanup_markers = {path for path in self.cleanup_markers if path.exists()} - return bool(self.cleanup_markers) + self.blocked = False @asynccontextmanager async def lease(self, binding, root): @@ -136,4 +146,5 @@ async def lease(self, binding, root): try: yield finally: - self.cleanup_markers.update(Path(root).rglob('cleanup-required.json')) + if (Path(root) / 'cleanup-required.json').exists(): + self.blocked = True diff --git a/harborrl/rollout/harbor_job/sandbox_transport.py b/harborrl/rollout/harbor_job/sandbox_transport.py index 385ad5701..6fa98b36e 100644 --- a/harborrl/rollout/harbor_job/sandbox_transport.py +++ b/harborrl/rollout/harbor_job/sandbox_transport.py @@ -1,18 +1,41 @@ """Observe SDK allocation IDs without retrying creation requests.""" import re +import time import httpx class SandboxStartupTransport(httpx.AsyncBaseTransport): - def __init__(self, base_url, on_created, *, transport=None): + def __init__(self, base_url, on_created, *, transport=None, on_request=None): self.base_url = base_url.rstrip('/') self.on_created = on_created self.transport = transport or httpx.AsyncHTTPTransport() + self.on_request = on_request + + def record_request(self, request, started, *, status=None, error=None): + """Persist only lifecycle route templates, never URLs or SDK messages.""" + if self.on_request is None: + return + url = str(request.url).split('?', 1)[0] + operation = 'other' + for route, name in ((r'/sandboxes', 'sandboxes'), + (r'/sandboxes/[A-Za-z0-9_-]+', 'sandbox'), + (r'/sandboxes/[A-Za-z0-9_-]+/endpoints/\d+', 'endpoint')): + if re.fullmatch(re.escape(self.base_url) + route, url): + operation = name + break + self.on_request({'method': request.method, 'operation': operation, 'http_status': status, + 'error_type': error, 'elapsed_sec': time.monotonic() - started, 'time': time.time()}) async def handle_async_request(self, request): url = str(request.url).split('?', 1)[0] - response = await self.transport.handle_async_request(request) + started = time.monotonic() + try: + response = await self.transport.handle_async_request(request) + except BaseException as exc: + self.record_request(request, started, error=type(exc).__name__) + raise + self.record_request(request, started, status=response.status_code) if request.method == 'POST' and url == self.base_url + '/sandboxes' and response.is_success: await response.aread() try: diff --git a/harborrl/rollout/remote_shard.py b/harborrl/rollout/remote_shard.py new file mode 100644 index 000000000..523f92057 --- /dev/null +++ b/harborrl/rollout/remote_shard.py @@ -0,0 +1,201 @@ +"""Remote gateway shards: the local shard surface, served over admin HTTP. + +A remote shard is an independent process (usually on another machine) that +owns one TraceRegistry and one gateway server. Generation traffic never +crosses this module: runners talk to the shard's gateway URL directly with +their attempt credentials, exactly as with a local shard. Only the +manager-side registry surface — admission epochs, registration, delivery +waits, session audits, sealing — crosses the admin channel, authenticated +by the run token. + +Fail-closed throughout: a transport failure surfaces as OSError so the +scheduler fails exactly this shard's attempts, and an unreachable shard +reads as a closed admission (version unknown, not accepting) until the +next activation republishes a consistent epoch. A restarted-but-reachable +shard instead reports no version — the facade surfaces that disagreement +as an unknown epoch, so the service rejects new groups until the next +activation — while its fresh registry still accepts and enforces +settle-before-republish for the attempts it owns. +""" +from __future__ import annotations + +import contextlib +import json +from pathlib import Path +from urllib import error as urlerror +from urllib.parse import urlsplit +from urllib.request import ProxyHandler, Request, build_opener + +# Manager<->shard traffic is direct LAN HTTP; never route it through a proxy. +_OPENER = build_opener(ProxyHandler({})) +_REMOTE_ERRORS = {"ValueError": ValueError, "TimeoutError": TimeoutError, + "PermissionError": PermissionError, "KeyError": KeyError, + "TypeError": TypeError, "RuntimeError": RuntimeError} + + +class RemoteShardError(OSError): + """The remote shard answered with an unmapped failure.""" + + +def plain_origin(url, label): + """Validate and normalize a plain http(s) origin.""" + parts = urlsplit(url if isinstance(url, str) else "") + if (parts.scheme not in ("http", "https") or not parts.hostname + or parts.username or parts.password or parts.query or parts.fragment + or parts.path not in ("", "/")): + raise ValueError(f"{label} must be a plain http(s) origin") + return f"{parts.scheme}://{parts.netloc}".rstrip("/") + + +def read_shard_token(path): + """Load the run token a remote shard and its manager must share.""" + token = Path(path).read_text().strip() + if not token: + raise ValueError(f"gateway shard token file is empty: {path}") + return token + + +class RemoteRegistry: + """The manager-side TraceRegistry surface of one remote shard. + + ``version`` and ``accepting`` are live reads, not a cache: admission + decisions never trust manager-local memory about a process that can + crash and restart between two group submissions. + """ + + def __init__(self, admin_url, token, *, timeout=30.0, status_timeout=5.0): + self.admin_url = plain_origin(admin_url, "shard admin_url") + if not isinstance(token, str) or not token: + raise ValueError("remote shard registry requires the run token") + self.token = token + self.timeout = float(timeout) + self.status_timeout = float(status_timeout) + + # -- transport --------------------------------------------------------- + + def request(self, method, path, payload=None, *, timeout=None): + data = None if payload is None else json.dumps(payload).encode() + headers = {"x-api-key": self.token} + if data is not None: + headers["Content-Type"] = "application/json" + request = Request(f"{self.admin_url}{path}", data=data, + headers=headers, method=method) + try: + with _OPENER.open(request, timeout=timeout or self.timeout) as response: + body = response.read() + except urlerror.HTTPError as exc: + if exc.code == 401: + raise PermissionError("remote shard rejected the run token") from exc + raise + return json.loads(body) if body else None + + def _call(self, action, payload=None, *, timeout=None): + body = self.request("POST", f"/registry/{action}", payload, timeout=timeout) + if isinstance(body, dict) and body.get("ok"): + return body.get("result") + error = (body or {}).get("error") or {} + raise _REMOTE_ERRORS.get(error.get("type"), RemoteShardError)( + error.get("message") or "remote shard call failed") + + def status(self): + return self.request("GET", "/status", timeout=self.status_timeout) or {} + + # -- admission epoch --------------------------------------------------- + + @property + def version(self): + try: + return self.status().get("version") + except (OSError, ValueError): + return None + + @property + def accepting(self): + try: + return bool(self.status().get("accepting")) + except (OSError, ValueError): + return False + + def close_admission(self): + self._call("close_admission") + + def drain(self, timeout): + timeout = float(timeout) + self._call("drain", {"timeout": timeout}, timeout=timeout + 15) + + def publish_version(self, states, expected): + self._call("publish_version", {"states": states, "expected": str(expected)}) + + # -- per-attempt surface ----------------------------------------------- + + def register(self, identity, trial_id, *, credential=None): + result = self._call("register", {"identity": identity.to_dict(), + "trial_id": trial_id, "credential": credential}) + return (result or {}).get("credential") + + def close(self, attempt_id): + self._call("close", {"attempt_id": attempt_id}) + + def wait_delivered(self, attempt_id, timeout): + timeout = float(timeout) + self._call("wait_delivered", {"attempt_id": attempt_id, "timeout": timeout}, + timeout=timeout + 15) + + def unconsumed_responses(self, attempt_id): + return self._call("unconsumed_responses", {"attempt_id": attempt_id}) or [] + + def consume(self, attempt_id, response_id, content, evidence_ref, *, cli_version=None): + self._call("consume", {"attempt_id": attempt_id, "response_id": response_id, + "content": content, "evidence_ref": evidence_ref, + "cli_version": cli_version}) + + def audit_tool_rejections(self, attempt_id, results, evidence_ref): + self._call("audit_tool_rejections", {"attempt_id": attempt_id, "results": results, + "evidence_ref": evidence_ref}) + + def seal(self, attempt_id): + result = self._call("seal", {"attempt_id": attempt_id}) + return result["turns"], result["seal"] + + +class RemoteShard: + """One gateway shard served by an independent process. + + Presents the GatewayShard surface — ``name``, ``url``, ``port``, + ``registry``, ``start``/``close`` — so scheduling and the ShardAdmission + barrier never learn which form a shard takes. ``start()`` proves the + process is reachable and answers as this shard; ``close()`` asks it to + shut down, best-effort, because the shard process owns its lifecycle. + """ + + def __init__(self, spec, *, token, trace_dir, timeout=30.0): + self.spec = spec + self.trace_dir = Path(trace_dir) + self.registry = RemoteRegistry(spec.admin_url, token, timeout=timeout) + + @property + def name(self): + return self.spec.name + + @property + def url(self): + return self.spec.url + + @property + def port(self): + return urlsplit(self.spec.url).port + + @property + def admin_url(self): + return self.spec.admin_url + + def start(self): + status = self.registry.status() + if status.get("shard") != self.spec.name: + raise RuntimeError( + f"gateway shard {self.spec.name!r} at {self.spec.admin_url} " + f"answers as {status.get('shard')!r}") + + def close(self): + with contextlib.suppress(OSError, ValueError): + self.registry.request("POST", "/shutdown", timeout=5.0) diff --git a/harborrl/rollout/runtime.py b/harborrl/rollout/runtime.py index 9288ebef8..5e3107800 100644 --- a/harborrl/rollout/runtime.py +++ b/harborrl/rollout/runtime.py @@ -10,23 +10,113 @@ import asyncio import atexit -import json import logging +import json +import re import secrets +import shlex import threading -from contextlib import nullcontext -from dataclasses import dataclass +import time +from contextlib import nullcontext, suppress +from dataclasses import dataclass, field from pathlib import Path +from urllib.parse import urlsplit from harborrl.models.qwen import normalize_chat_template_ids from harborrl.rollout.harbor_job.coordinator import GroupSlots, run_attempt -from harborrl.trajectories.native import Identity, RewardProfile, digest, read_json +from harborrl.rollout.remote_shard import RemoteShard, read_shard_token +from harborrl.trajectories.native import ( + Identity, + RewardProfile, + digest, + encode, + publish, + read_json, +) TASK_ENTRY_KEYS = {"id", "revision", "path", "task_digest", "reward_profile"} +SHARD_NAME = re.compile(r"[a-z0-9][a-z0-9-]{0,31}") +GATEWAY_ROUTINGS = ("modulo", "worker_affinity", "auto") +DEFAULT_SANDBOXES_PER_SHARD = 30 +LOCAL_RUNNER_KEY = "local" +SHARD_SERVICE_MODULE = "harborrl.gateway.shard_service" +# One remote shard must report its ready line within this budget: ssh connect +# plus tokenizer load on the remote host. A candidate that misses it is +# abandoned so growth can try the next host instead of hanging a slot. +REMOTE_SHARD_START_TIMEOUT = 120.0 + + +@dataclass(frozen=True) +class GatewayShardSpec: + """One gateway's static placement: name, exclusive port, advertised URL. + + ``mode`` is "local" (this process binds the port) or "remote" (an + independent shard process serves ``url``; the manager reaches its + registry through ``admin_url``). + """ + + name: str + port: int + url: str + mode: str = "local" + admin_url: str = None + + +def _port_url(url, port): + """Rewrite only the port of a plain origin.""" + parts = urlsplit(url) + host = f"[{parts.hostname}]" if ":" in (parts.hostname or "") else parts.hostname + return f"{parts.scheme}://{host}:{port}" + + +GROWTH_HOST_FIELDS = {"host", "user", "port", "identity"} + + +def _growth_host(entry): + """Normalize one growth-pool entry into (destination, port, identity, url_host). + + A plain string is an ssh destination (``user@host``) that inherits every + transport option from the invoking user's SSH config. An object narrows + those options for this host only — an ephemeral machine whose key and + port are known per run, not in ``~/.ssh/config``. The url host is the + bare host: the SSH port is the transport's, never the shard's gateway + port, which the shard process binds and reports on its ready line. + """ + if isinstance(entry, str): + if not entry.strip() or any(character.isspace() for character in entry): + raise ValueError('gateway_hosts lists ssh destinations (user@host)') + return entry, None, None, entry.split("@")[-1] + if not isinstance(entry, dict) or set(entry) - GROWTH_HOST_FIELDS: + raise ValueError( + f'gateway_hosts entries are ssh destinations or {sorted(GROWTH_HOST_FIELDS)}') + host = entry.get("host") + if (not isinstance(host, str) or not host.strip() + or any(character.isspace() for character in host) + or host.startswith("-") or host.count("@") > 1): + raise ValueError('gateway_hosts entries require an ssh destination host') + user, port, identity = entry.get("user"), entry.get("port"), entry.get("identity") + if user is not None and (not isinstance(user, str) or not user.strip() + or user.startswith("-")): + raise ValueError('gateway_hosts entry user must not look like a flag') + if port is not None and (type(port) is not int or not 0 < port < 65536): + raise ValueError('gateway_hosts entry port must be an integer in 1..65535') + if identity is not None and (not isinstance(identity, str) or not identity.strip()): + raise ValueError('gateway_hosts entry identity must be a nonempty path') + destination = f"{user}@{host}" if user else host + return destination, port, identity, host.split("@")[-1] + + +def _plain_origin(url, label): + parts = urlsplit(url if isinstance(url, str) else "") + if (parts.scheme not in ("http", "https") or not parts.hostname + or parts.username or parts.password or parts.query or parts.fragment + or parts.path not in ("", "/")): + raise ValueError(f"{label} must be a plain http(s) origin") + return f"{parts.scheme}://{parts.netloc}".rstrip("/") class RuntimeLifecycleError(RuntimeError): - """An async resource was used outside its owning runtime lifecycle.""" + """Embedded rollout loop/thread-affinity or lifecycle misuse.""" class CleanupRequiredError(RuntimeError): @@ -77,6 +167,89 @@ class RolloutRuntimeConfig: endpoint: str sandbox: dict = None runner_transport: str = 'ssh' + sampling_seed: int | None = None + shards: tuple = () + gateway_routing: str = 'modulo' + sandboxes_per_shard: int = DEFAULT_SANDBOXES_PER_SHARD + worker_shards: dict = field(default_factory=dict) + shard_token_file: str = None + gateway_hosts: tuple = () + hf_ckpt: str = None + shard_python: str = None + + def __post_init__(self): + if self.sampling_seed is not None and ( + type(self.sampling_seed) is not int or not 0 <= self.sampling_seed < 2**31): + raise ValueError('sampling_seed must be null or an integer in 0..2147483647') + if self.gateway_routing not in GATEWAY_ROUTINGS: + raise ValueError(f'gateway_routing must be one of {GATEWAY_ROUTINGS}') + if (type(self.sandboxes_per_shard) is not int + or self.sandboxes_per_shard < 1): + raise ValueError('sandboxes_per_shard must be a positive integer') + if self.gateway_routing == 'auto' and self.gateway_hosts: + # Growth needs the shared run token and a tokenizer checkpoint the + # spawned process loads itself. + if not self.shard_token_file: + raise ValueError( + 'gateway_hosts requires HARBORRL_NATIVE_SHARD_TOKEN_FILE') + if not self.hf_ckpt: + raise ValueError( + 'gateway_hosts requires HF_CKPT so remote shards load the ' + 'run tokenizer') + if self.gateway_hosts and self.gateway_routing != 'auto': + # Remote growth is an auto-routing concern: static routing has no + # mid-run growth for the pool to serve. + raise ValueError('gateway_hosts grows gateways under auto routing only') + if not isinstance(self.gateway_hosts, tuple): + raise ValueError('gateway_hosts lists ssh destinations (user@host)') + # Validated here so a malformed pool fails at startup; the runtime + # normalizes and rotates the same entries when it grows a shard. + for entry in self.gateway_hosts: + _growth_host(entry) + if (not isinstance(self.worker_shards, dict) + or not all(isinstance(key, str) and isinstance(value, str) + for key, value in self.worker_shards.items())): + raise ValueError('worker_shards maps runner destinations to gateway shard names') + if self.worker_shards and self.gateway_routing == 'auto': + # auto leases from the live shard set and never consults the + # affinity map, so accepting one here would silently route by + # lease order while the operator believes affinity is in force. + raise ValueError( + 'worker_shards is static affinity routing; auto routing ignores it. ' + 'Declare gateway_routing: worker_affinity, or drop the mapping') + if self.shards: + names = [spec.name for spec in self.shards] + ports = [spec.port for spec in self.shards] + if any(not SHARD_NAME.fullmatch(name) for name in names): + raise ValueError('gateway shard names are lowercase slugs') + if len(set(names)) != len(names): + raise ValueError('gateway shard names must be unique') + if (any(type(port) is not int or not 0 < port < 65536 for port in ports) + or len(set(ports)) != len(ports)): + raise ValueError('gateway shard ports must be unique integers in 1..65535') + if any(spec.mode not in ('local', 'remote') for spec in self.shards): + raise ValueError('gateway shard mode is local or remote') + if any(spec.mode == 'remote' and not spec.admin_url for spec in self.shards): + raise ValueError('a remote gateway shard requires its admin_url') + if any(spec.mode == 'remote' for spec in self.shards) and not self.shard_token_file: + raise ValueError('remote gateway shards require HARBORRL_NATIVE_SHARD_TOKEN_FILE') + if self.gateway_routing == 'worker_affinity': + # Static routing must be total before dispatch, never + # discovered mid-run: every runner names exactly one shard. + required = set(self.workers) + if self.runner_transport == 'local': + required.add(LOCAL_RUNNER_KEY) + allowed = set(self.workers) | {LOCAL_RUNNER_KEY} + unknown = sorted(set(self.worker_shards) - allowed) + if unknown: + raise ValueError(f'worker_shards names unknown runners {unknown}') + missing = sorted(required - set(self.worker_shards)) + if missing: + raise ValueError('worker_affinity routes every runner to a declared ' + f'gateway shard; missing {missing}') + unsharded = sorted(set(self.worker_shards.values()) - set(names)) + if unsharded: + raise ValueError(f'worker_shards names unknown gateway shards {unsharded}') @staticmethod def _required(env, name): @@ -101,14 +274,50 @@ def from_environment(cls, env, *, endpoint, concurrency_limit, group_size, overrides = json.loads(env.get("HARBORRL_NATIVE_TASK_WORKER_OVERRIDES") or "{}") except ValueError as exc: raise ValueError("invalid HARBORRL_NATIVE_TASK_WORKER_OVERRIDES") from exc + try: + worker_shards = json.loads(env.get("HARBORRL_NATIVE_GATEWAY_WORKER_SHARDS") or "{}") + except ValueError as exc: + raise ValueError("invalid HARBORRL_NATIVE_GATEWAY_WORKER_SHARDS") from exc + try: + gateway_hosts = json.loads(env.get("HARBORRL_NATIVE_GATEWAY_HOSTS") or "[]") + except ValueError as exc: + raise ValueError("invalid HARBORRL_NATIVE_GATEWAY_HOSTS") from exc + if (not isinstance(gateway_hosts, list) + or not all(isinstance(host, (str, dict)) for host in gateway_hosts)): + # Entries are ssh destinations or option mappings; __post_init__ + # validates each one's fields. + raise ValueError( + "HARBORRL_NATIVE_GATEWAY_HOSTS is a list of ssh destinations or host mappings") + gateway_port = int(gateway_port or cls._required(env, "HARBORRL_NATIVE_GATEWAY_PORT")) + gateway_url = gateway_url or cls._required(env, "HARBORRL_NATIVE_GATEWAY_URL") + shards = cls._gateway_shards(env, gateway_port, gateway_url) + try: + sandboxes_per_shard = int(env.get( + "HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD", + str(DEFAULT_SANDBOXES_PER_SHARD))) + except ValueError as exc: + raise ValueError("invalid HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD") from exc + # Without a static table the default is elastic: exactly one gateway + # until the live sandbox count passes sandboxes_per_shard, then one + # more. A declared table keeps the historical static default. + declared_routing = env.get("HARBORRL_NATIVE_GATEWAY_ROUTING") + declared_routing = declared_routing.strip() if isinstance(declared_routing, str) else "" return cls( run_id=cls._required(env, "RUN_ID"), run_root=cls._required(env, "RUN_DIR"), catalog_path=cls._required(env, "HARBORRL_NATIVE_CATALOG"), profile_path=cls._required(env, "HARBORRL_NATIVE_PROFILE"), gateway_host=gateway_host or env.get("HARBORRL_NATIVE_GATEWAY_HOST", "0.0.0.0"), - gateway_port=int(gateway_port or cls._required(env, "HARBORRL_NATIVE_GATEWAY_PORT")), - gateway_url=gateway_url or cls._required(env, "HARBORRL_NATIVE_GATEWAY_URL"), + gateway_port=gateway_port, + gateway_url=gateway_url, + shards=shards, + gateway_routing=declared_routing or ("auto" if not shards else "modulo"), + sandboxes_per_shard=sandboxes_per_shard, + worker_shards=dict(worker_shards), + shard_token_file=env.get("HARBORRL_NATIVE_SHARD_TOKEN_FILE") or None, + gateway_hosts=tuple(gateway_hosts), + hf_ckpt=env.get("HF_CKPT") or None, + shard_python=env.get("HARBORRL_NATIVE_SHARD_PYTHON") or None, runner_python=runner_python, workers=tuple(workers), task_worker_overrides=dict(overrides), @@ -125,8 +334,57 @@ def from_environment(cls, env, *, endpoint, concurrency_limit, group_size, endpoint=endpoint.rstrip("/"), sandbox=json.loads(env.get('HARBORRL_NATIVE_SANDBOX', 'null')), runner_transport=env.get('HARBORRL_NATIVE_RUNNER_TRANSPORT', 'ssh'), + sampling_seed=(int(env['HARBORRL_TRAINING_SEED']) + if 'HARBORRL_TRAINING_SEED' in env else None), ) + @staticmethod + def _gateway_shards(env, gateway_port, gateway_url): + """Parse the static shard table; an absent table means one elastic gateway. + + Undeclared local shard ports default to consecutive slots after the + declared gateway port; every local shard URL is the gateway URL with + its port rewritten, so shard placement stays one spelling everywhere. + A remote shard names the origins of its independent process instead. + An absent table routes "auto": the implicit first gateway serves until + the live sandbox count passes sandboxes_per_shard, then the runtime + spawns the next local shard itself. + """ + raw = env.get("HARBORRL_NATIVE_GATEWAY_SHARDS") + if raw is None or not raw.strip(): + return () + try: + entries = json.loads(raw) + except ValueError as exc: + raise ValueError("invalid HARBORRL_NATIVE_GATEWAY_SHARDS") from exc + if not isinstance(entries, list) or not entries: + raise ValueError("HARBORRL_NATIVE_GATEWAY_SHARDS must be a nonempty shard list") + specs = [] + for index, entry in enumerate(entries): + if (not isinstance(entry, dict) + or set(entry) - {"name", "port", "mode", "url", "admin_url"} + or not isinstance(entry.get("name"), str)): + raise ValueError("each gateway shard is {name, port?, mode?, url?, admin_url?}") + mode = entry.get("mode", "local") + if mode not in ("local", "remote"): + raise ValueError("gateway shard mode is local or remote") + if mode == "remote": + if "port" in entry: + raise ValueError("a remote gateway shard takes its port from url") + url = _plain_origin(entry.get("url"), "remote gateway shard url") + admin_url = _plain_origin(entry.get("admin_url"), + "remote gateway shard admin_url") + specs.append(GatewayShardSpec(name=entry["name"], port=urlsplit(url).port, + url=url, mode="remote", admin_url=admin_url)) + continue + if "url" in entry or "admin_url" in entry: + raise ValueError("a local gateway shard binds a port and derives its url") + port = entry.get("port") + port = gateway_port + index if port is None else int(port) + specs.append(GatewayShardSpec(name=entry["name"], port=port, + url=_port_url(gateway_url, port))) + return tuple(specs) + def environment(self): """Render the environment mapping this config was built from.""" return { @@ -150,6 +408,25 @@ def environment(self): "MODEL_TAG": self.model_tag, "HARBORRL_NATIVE_SANDBOX": json.dumps(self.sandbox), "HARBORRL_NATIVE_RUNNER_TRANSPORT": self.runner_transport, + "HARBORRL_NATIVE_GATEWAY_ROUTING": self.gateway_routing, + "HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD": str(self.sandboxes_per_shard), + "HARBORRL_NATIVE_GATEWAY_WORKER_SHARDS": json.dumps( + self.worker_shards, sort_keys=True), + **({'HARBORRL_NATIVE_GATEWAY_SHARDS': json.dumps([ + {'name': spec.name, 'port': spec.port} if spec.mode == 'local' else { + 'name': spec.name, 'mode': spec.mode, 'url': spec.url, + 'admin_url': spec.admin_url} + for spec in self.shards])} + if self.shards else {}), + **({'HARBORRL_NATIVE_SHARD_TOKEN_FILE': self.shard_token_file} + if self.shard_token_file else {}), + **({'HARBORRL_NATIVE_GATEWAY_HOSTS': json.dumps(list(self.gateway_hosts))} + if self.gateway_hosts else {}), + **({'HF_CKPT': self.hf_ckpt} if self.hf_ckpt else {}), + **({'HARBORRL_NATIVE_SHARD_PYTHON': self.shard_python} + if self.shard_python else {}), + **({'HARBORRL_TRAINING_SEED': str(self.sampling_seed)} + if self.sampling_seed is not None else {}), } @@ -167,8 +444,230 @@ def tokenizer_fingerprint(tokenizer): return tokenizer_digest, template +def _shard_service_environment(config, name): + """The environment one remotely spawned shard process starts under. + + Everything the shard service validates comes from the manager's own + config, so a spawned shard is locked to the same run, digests and + serving pool as the gateways that started locally. Both ports bind + ephemeral — the OS owns exclusivity across every machine that might + host a shard — and the manager reads the resolved ports from the + process' ready line before advertising the shard anywhere. + """ + return { + "RUN_ID": config.run_id, + "RUN_DIR": str(config.run_root), + "HARBORRL_SHARD_NAME": name, + "HARBORRL_SHARD_GATEWAY_HOST": "0.0.0.0", + "HARBORRL_SHARD_GATEWAY_PORT": "0", + "HARBORRL_SHARD_ADMIN_HOST": "0.0.0.0", + "HARBORRL_SHARD_ADMIN_PORT": "0", + "HARBORRL_SHARD_TOKEN_FILE": config.shard_token_file, + "HARBORRL_SHARD_ENDPOINT": config.endpoint, + "HARBORRL_NATIVE_PROFILE": config.profile_path, + "MODEL_TAG": config.model_tag, + "ROLLOUT_MAX_RESPONSE_LEN": str(config.max_output_tokens), + "ROLLOUT_MAX_CONTEXT_LEN": str(config.max_context), + "HARBORRL_NATIVE_TOKENIZER_DIGEST": config.expected_tokenizer_digest, + "HARBORRL_NATIVE_TEMPLATE_DIGEST": config.expected_template_digest, + "HARBORRL_NATIVE_AUDITED_LOGPROBS": str(int(config.audited_raw_logprobs)), + "HF_CKPT": config.hf_ckpt, + } + + +def _shard_service_command(host, environment, *, python, project, port=None, identity=None): + """The ssh command that starts one shard process on a remote host. + + ssh forwards no environment, so the shard service's env contract crosses + as explicit ``env`` assignments ahead of the interpreter; the same + BatchMode/ClearAllForwardings options the runner transport uses keep the + launch non-interactive and proxy-free. ``port`` and ``identity`` come + from a growth-host mapping: an ephemeral machine often listens on a + non-standard ssh port behind a key that is not in the invoking user's + ssh config. + """ + assignments = " ".join(f"{key}={shlex.quote(value)}" + for key, value in sorted(environment.items())) + remote = (f"cd {shlex.quote(project)} && exec env {assignments} " + f"{shlex.quote(python)} -m {SHARD_SERVICE_MODULE}") + argv = ["ssh", "-o", "BatchMode=yes", "-o", "ClearAllForwardings=yes", + "-o", "StrictHostKeyChecking=accept-new"] + if port is not None: + argv += ["-p", str(port)] + if identity is not None: + argv += ["-i", str(identity)] + argv += [host, remote] + return argv + + +class GatewayShard: + """One local gateway instance: exclusive port, trace registry and trace tree. + + A remote shard (a separate gateway process, Phase 2) presents the same + surface — ``name``, ``url``, per-attempt admission, ``start``/``close`` — + so scheduling and the version barrier never learn which form a shard + takes. Within one process the sandbox admission budget stays shared: + shards split admission and tracing, not the worker quota. + """ + + def __init__(self, spec, *, registry, gateway, server, trace_dir): + self.spec = spec + self.registry = registry + self.gateway = gateway + self.server = server + self.trace_dir = Path(trace_dir) + self.thread = threading.Thread( + target=server.serve_forever, name=f"harborrl-gateway-{spec.name}", daemon=True) + self._started = False + + @property + def name(self): + return self.spec.name + + @property + def url(self): + return self.spec.url + + @property + def port(self): + return self.spec.port + + def start(self): + self._started = True + self.thread.start() + + def close(self): + # shutdown() before serve_forever() starts would deadlock the thread. + if self._started: + self.server.shutdown() + if self.thread.is_alive(): + self.thread.join(5) + self.server.server_close() + + +class SpawnedRemoteShard(RemoteShard): + """A remote shard this runtime launched, plus the ssh transport owning it. + + ``close`` asks the process to stop through the admin plane — it unwinds + cleanly and flushes its traces — and only kills the ssh transport when it + outlives a grace period. Killing first would hang up the remote + mid-shutdown; the transport is nevertheless always reaped so the manager + never leaks child processes. + """ + + SHUTDOWN_GRACE = 15.0 + + def __init__(self, spec, *, token, trace_dir, process, timeout=30.0): + super().__init__(spec, token=token, trace_dir=trace_dir, timeout=timeout) + self.process = process + + def close(self): + super().close() + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return # No loop left to reap on; the shutdown request went out. + loop.create_task(self._reap()) + + async def _reap(self): + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(self.process.wait(), self.SHUTDOWN_GRACE) + if self.process.returncode is None: + with suppress(ProcessLookupError): + self.process.kill() + with suppress(Exception): + await self.process.wait() + + +class ShardAdmission: + """The one policy epoch shared by every gateway shard. + + ``drain`` closes admission on ALL shards before waiting on any of them, + and ``publish_version`` opens the next epoch only after every shard + drained. Per-attempt operations (register, begin, finish, seal) stay on + the shard that owns the attempt; this facade is the only surface the + rollout service and the version barrier ever touch. Shard disagreement + (a partially completed activation) reads as a closed admission until the + next activation republishes a consistent epoch. + + Auto routing grows the shard set mid-run: a runtime passes its live + shard list here, and ``join`` attaches a freshly spawned shard onto the + current epoch atomically so the facade never observes disagreement. + """ + + def __init__(self, shards): + # Hold a caller's list by reference so mid-run growth stays visible; + # any other iterable is copied into a private list. + self._shards = shards if isinstance(shards, list) else list(shards) + self._lock = threading.Lock() + self._epoch = None # (states, version) of the last published epoch + + @property + def shards(self): + with self._lock: + return tuple(self._shards) + + @property + def version(self): + versions = {shard.registry.version for shard in self.shards} + return versions.pop() if len(versions) == 1 else None + + @property + def accepting(self): + return all(shard.registry.accepting for shard in self.shards) + + def drain(self, timeout): + deadline = time.monotonic() + timeout + shards = self.shards + for shard in shards: + shard.registry.close_admission() + for shard in shards: + shard.registry.drain(max(0.0, deadline - time.monotonic())) + + def publish_version(self, states, expected): + with self._lock: + self._epoch = (states, str(expected)) + for shard in self.shards: + shard.registry.publish_version(states, expected) + + def join(self, shard): + """Attach a mid-run shard onto the current epoch, then admit it. + + A fresh registry starts (version None, accepting True); joining with + that version would read as shard disagreement and fail-close group + admission until the next activation. Replaying the last published + epoch inside the lock keeps one version across every shard; before + the first activation there is no epoch yet and every shard + legitimately starts unversioned. + + The replay is the facade's own close-then-publish sequence scoped to + the new shard: ``publish_version`` refuses to open a barrier that was + never closed, and a fresh registry opens accepting, so its admission + closes first exactly like a drain would. + """ + with self._lock: + epoch = self._epoch + if epoch is not None: + states, version = epoch + shard.registry.close_admission() + shard.registry.publish_version(states, version) + self._shards.append(shard) + + def leave(self, shard): + """Detach a retiring shard from the epoch under the same lock. + + The mirror of ``join``: the facade stops spanning the shard, so a + concurrent drain or publish no longer waits on or versions a gateway + that is closing. Retirement happens only after the shard's attempts + all settled, so its registry state is already sealed history. + """ + with self._lock: + if shard in self._shards: + self._shards.remove(shard) + + class RolloutRuntime: - """One gateway, one trace registry and one rollout root per run.""" + """One rollout root and one or more local gateway shards per run.""" def __init__(self, config, tokenizer, *, backend=None): from harborrl.gateway.server import Gateway, make_server @@ -182,10 +681,9 @@ def __init__(self, config, tokenizer, *, backend=None): self.profile = read_json(config.profile_path) self.rollout_root = self.run_root / "native-rollouts" self.rollout_root.mkdir(parents=True, exist_ok=True) + self._loop = None self.cleanup = CleanupAdmission(self.rollout_root) self.cleanup.check() - self._loop = None - self.registry = TraceRegistry(self.run_root / "native-traces") self.semaphore = asyncio.Semaphore(config.concurrency_limit) if config.sandbox: from harborrl.rollout.harbor_job.sandbox_config import SandboxAdmission @@ -195,17 +693,83 @@ def __init__(self, config, tokenizer, *, backend=None): or template_digest != config.expected_template_digest): raise ValueError("live tokenizer/template digest differs from the native lock") model = MATH_REGISTRY.model_for_preset(config.model_tag) + # One shared serving backend: shards split admission and tracing, + # never the inference pool. self.backend = backend or SGLangBackend( config.endpoint, tokenizer, model.adapter, tokenizer_digest=tokenizer_digest, template_digest=template_digest, max_context=config.max_context, audited_raw_logprobs=config.audited_raw_logprobs, + sampling_seed=config.sampling_seed, ) - gateway = Gateway(self.registry, self.backend, model=self.profile["model"], - max_output_tokens=config.max_output_tokens) - self.server = make_server(gateway, host=config.gateway_host, - port=config.gateway_port) - self._start_gateway() + specs = config.shards or (GatewayShardSpec( + name="gateway", port=config.gateway_port, url=config.gateway_url),) + trace_root = self.run_root / "native-traces" + # An explicit shard table gives every shard its own immutable trace + # subtree; the implicit first gateway keeps the historical layout. + shards = [] + for spec in specs: + trace_dir = trace_root / f"shard-{spec.name}" if config.shards else trace_root + if spec.mode == "remote": + # An independent process owns the registry and the gateway; + # this side only holds the admin-channel proxy. Resolved + # through the module attribute so deployments can substitute + # the surface (tests pin it this way). + from harborrl.rollout.remote_shard import RemoteShard + shards.append(RemoteShard(spec, token=read_shard_token(config.shard_token_file), + trace_dir=trace_dir)) + continue + registry = TraceRegistry(trace_dir) + gateway = Gateway(registry, self.backend, model=self.profile["model"], + max_output_tokens=config.max_output_tokens) + # Each attempt holds at most one policy generation plus occasional + # count_tokens calls; double the scheduler budget with a floor of 4. + server = make_server(gateway, host=config.gateway_host, port=spec.port, + max_inflight=max(4, 2 * config.concurrency_limit)) + shards.append(GatewayShard(spec, registry=registry, gateway=gateway, + server=server, trace_dir=trace_dir)) + # The live shard set: auto routing appends spawned gateways here and + # the admission facade holds the same list, so the one epoch always + # spans every shard that exists. + self._shards = shards + # Auto routing occupancy lives on the event loop that dispatches + # attempts; nothing outside the loop touches it. + self._auto_inflight = {} + # Auto routing names spawned gateways monotonically so a retired + # name's trace subtree is never reused by a later regrowth. A static + # table may coexist with the growth pool, so the counter also skips + # every name the table already claims. + self._static_names = {spec.name for spec in config.shards} + self._shard_seq = 0 + self._auto_spawned = set() + # Growth serializes under this lock: concurrent slots at capacity + # must grow exactly once, and the fill re-check happens under it. + self._lease_lock = asyncio.Lock() + # Remote growth rotates through the declared pool so consecutive + # spawns spread across machines instead of piling onto the first. + self._host_seq = 0 + # Remotely spawned shards share the run token with statically + # declared ones; loading it here fails loudly at startup instead of + # on the first mid-run growth attempt. + self._shard_token = (read_shard_token(config.shard_token_file) + if config.gateway_hosts else None) + # The service-facing admission epoch spans every shard. + self.registry = ShardAdmission(self._shards) + if config.gateway_routing == "auto": + # Growth legitimately advances the record, so auto routing keeps + # it replaceable; static tables keep the immutable one-shot form. + self._write_shards_evidence() + else: + publish(self.run_root / "gateway-shards.json", self._shards_payload()) + try: + for shard in self._shards: + shard.start() + except BaseException: + # Never leak local servers because a later shard failed to start. + for shard in self._shards: + with suppress(Exception): + shard.close() + raise def check_event_loop(self): loop = asyncio.get_running_loop() @@ -214,17 +778,255 @@ def check_event_loop(self): elif self._loop is not loop: raise RuntimeLifecycleError("rollout runtime belongs to a different event loop") - def _start_gateway(self): - self.thread = threading.Thread( - target=self.server.serve_forever, name="harborrl-native-gateway", daemon=True - ) - self.thread.start() + @property + def shards(self): + """The live shard set; auto routing appends spawned gateways mid-run.""" + return tuple(self._shards) + + def _shards_payload(self): + return { + "gateway_routing": self.config.gateway_routing, + **({"sandboxes_per_shard": self.config.sandboxes_per_shard} + if self.config.gateway_routing == "auto" else {}), + "shards": [ + {"name": shard.name, "port": shard.port, "url": shard.url, + "mode": shard.spec.mode, + **({"admin_url": shard.spec.admin_url} + if shard.spec.mode == "remote" else {}), + "trace_dir": str(shard.trace_dir)} + for shard in self.shards]} + + def _write_shards_evidence(self): + """Rewrite the shard table evidence after auto routing grows it.""" + path = self.run_root / "gateway-shards.json" + temp = self.run_root / "gateway-shards.tmp" + temp.write_bytes(encode(self._shards_payload()) + b"\n") + temp.replace(path) + + async def lease_shard(self): + """Pick the gateway for the next running sandbox attempt. + + Auto routing keeps at most ``sandboxes_per_shard`` concurrent sandbox + attempts per gateway: the (n+1)-th concurrent attempt grows the fleet + instead of piling onto a full one — on a declared remote host first, + locally only when no candidate answers. Fail-closed by construction: + a spawn failure fails this attempt loudly, never a silent fallback + onto an overloaded gateway. + """ + for shard in self.shards: + if self._auto_inflight.get(shard.name, 0) < self.config.sandboxes_per_shard: + break + else: + async with self._lease_lock: + # Slots that waited on the lock while a sibling spawned must + # re-check: capacity may have appeared, and growth that + # started under the lock must not repeat. + for shard in self.shards: + if (self._auto_inflight.get(shard.name, 0) + < self.config.sandboxes_per_shard): + break + else: + shard = await self._spawn_shard() + self._auto_inflight[shard.name] = self._auto_inflight.get(shard.name, 0) + 1 + return shard + + def release_shard(self, shard): + remaining = self._auto_inflight.get(shard.name, 0) - 1 + if remaining > 0: + self._auto_inflight[shard.name] = remaining + else: + self._auto_inflight.pop(shard.name, None) + + def _candidate_hosts(self): + """The declared growth pool, rotated so spawns spread across it. + + Returns normalized ``(destination, port, identity, url_host)`` + tuples; a plain-string entry inherits its transport options from the + invoking user's ssh config. + """ + hosts = self.config.gateway_hosts + if not hosts: + return () + rotation = self._host_seq % len(hosts) + self._host_seq += 1 + rotated = hosts[rotation:] + hosts[:rotation] + return tuple(_growth_host(entry) for entry in rotated) + + async def _spawn_shard(self): + """Grow by one gateway: declared remote hosts first, local last.""" + # One growth consumes exactly one name: candidates that never + # materialized leave no trace subtree, so they reuse this name + # instead of burning the sequence. A declared static table may name + # some gateways already, so the sequence skips those names. + self._shard_seq += 1 + name = f"gw{self._shard_seq}" + while name in self._static_names: + self._shard_seq += 1 + name = f"gw{self._shard_seq}" + attempts = [] + for destination, port, identity, url_host in self._candidate_hosts(): + try: + return await self._spawn_remote_shard( + destination, name, port=port, identity=identity, url_host=url_host) + except Exception as exc: # noqa: BLE001 — next candidate, then local + attempts.append({ + "host": destination, + **({"port": port} if port is not None else {}), + **({"identity": identity} if identity is not None else {}), + "error": f"{type(exc).__name__}: {exc}"}) + if attempts: + # Observable, not silent: every remote candidate failed, so this + # gateway grows on the manager machine instead. + publish(self.run_root / "gateway-spawn-fallback.json", + {"fallback": "local", "attempts": attempts}) + return self._spawn_local_shard(name) + + def _shard_launch_python(self): + """The interpreter growth hosts run the shard service under. + + The shard process needs the serving stack (transformers, the harborrl + gateway) rather than the runner's sandbox framework, so deployments + with split environments point this at their serving interpreter; + undeclared keeps the runner's, which uniform deployments share. + """ + return self.config.shard_python or self.config.runner_python + + async def _start_shard_process(self, host, environment, *, port=None, identity=None): + """Launch the shard process over ssh and read its ready line. + + Returns the parsed ready record — shard name, resolved gateway and + admin ports — or raises; the caller abandons the candidate on any + failure and tries the next host. + """ + command = _shard_service_command( + host, environment, python=self._shard_launch_python(), + project=str(Path(__file__).resolve().parents[2]), + port=port, identity=identity) + process = await asyncio.create_subprocess_exec( + *command, stdin=asyncio.subprocess.DEVNULL, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.DEVNULL) + try: + line = await asyncio.wait_for( + process.stdout.readline(), REMOTE_SHARD_START_TIMEOUT) + record = json.loads(line) + if record.get("shard") is None or "gateway_port" not in record: + raise ValueError("no shard service ready line") + return process, record + except BaseException: + if process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + with suppress(Exception): + await process.wait() + raise + + async def _spawn_remote_shard(self, host, name, *, port=None, identity=None, + url_host=None): + """Start one shard process on a remote host and join it to the epoch.""" + from harborrl.gateway.shard_service import SHARD_SERVICE + + trace_dir = self.run_root / "native-traces" / f"shard-{name}" + process, ready = await self._start_shard_process( + host, _shard_service_environment(self.config, name), + port=port, identity=identity) + shard = None + try: + if ready.get("shard") != name or ready.get("service") != SHARD_SERVICE: + raise RuntimeError(f"gateway host {host} answered as another service") + # The shard is advertised at the host both planes reach it by, + # which is the destination's host part (user@ stripped) unless + # the caller names another spelling — e.g. a tunnel endpoint that + # the manager's side forwards to the same process. + if url_host is None: + url_host = host.split("@")[-1] + gateway_port, admin_port = int(ready["gateway_port"]), int(ready["admin_port"]) + spec = GatewayShardSpec( + name=name, port=gateway_port, url=f"http://{url_host}:{gateway_port}", + mode="remote", admin_url=f"http://{url_host}:{admin_port}") + shard = SpawnedRemoteShard(spec, token=self._shard_token, + trace_dir=trace_dir, process=process) + # start() proves the process answers as this shard over the admin + # plane — ssh alone only proved the process launched. + shard.start() + self.registry.join(shard) + except BaseException: + if shard is not None: + with suppress(Exception): + shard.close() + elif process.returncode is None: + # The shard object never formed; kill the transport directly + # so a rejected candidate leaves no process behind. + with suppress(ProcessLookupError): + process.kill() + with suppress(Exception): + await process.wait() + raise + self._auto_spawned.add(name) + self._write_shards_evidence() + return shard + + def _spawn_local_shard(self, name): + """Start one more local gateway and attach it to the live epoch.""" + from harborrl.gateway.server import Gateway, make_server + from harborrl.gateway.trace import TraceRegistry + + trace_dir = self.run_root / "native-traces" / f"shard-{name}" + registry = TraceRegistry(trace_dir) + gateway = Gateway(registry, self.backend, model=self.profile["model"], + max_output_tokens=self.config.max_output_tokens) + # The OS owns exclusivity: an ephemeral bind cannot collide with a + # sibling run's port allocation, and the advertised URL carries the + # port the socket actually took. + server = make_server(gateway, host=self.config.gateway_host, port=0, + max_inflight=max(4, 2 * self.config.concurrency_limit)) + port = server.server_address[1] + spec = GatewayShardSpec(name=name, port=port, + url=_port_url(self.config.gateway_url, port)) + shard = GatewayShard(spec, registry=registry, gateway=gateway, + server=server, trace_dir=trace_dir) + try: + shard.start() + self.registry.join(shard) + except BaseException: + with suppress(Exception): + shard.close() + raise + self._auto_spawned.add(name) + self._write_shards_evidence() + return shard + + def retire_idle_shards(self): + """Retire auto-spawned gateways once their sandboxes all settled. + + The mirror of growth: when a group settles and the running sandbox + count no longer justifies a gateway, it detaches from the admission + epoch, closes, and the evidence record shrinks back. Only idle + auto-spawned shards retire — the implicit first gateway never does, + and a busy shard (or anything behind it, since leases fill in order) + blocks the tail. The group's own outcome is already decided by the + time this runs, so a close failure is contained instead of replacing + the settlement result. + """ + retired = False + while len(self._shards) > 1: + shard = self._shards[-1] + if (shard.name not in self._auto_spawned + or self._auto_inflight.get(shard.name, 0)): + break + self._shards.pop() + self._auto_spawned.discard(shard.name) + self._auto_inflight.pop(shard.name, None) + self.registry.leave(shard) + with suppress(Exception): + shard.close() + retired = True + if retired: + self._write_shards_evidence() def close(self): - self.server.shutdown() - if getattr(self.thread, "is_alive", lambda: False)(): - self.thread.join(5) - self.server.server_close() + for shard in self.shards: + shard.close() def manifest_path(self, identity): """Deterministic shared-storage path of one trajectory manifest.""" @@ -392,18 +1194,35 @@ async def run_slot(config, runtime_, group_index, slot, rollout_id, entry, return await _run_slot(config, runtime_, slots, slot, entry, reward_profile) +def route_shard(config, shards, *, runner_host, slot, attempt_index): + """Static per-attempt shard routing: fixed at dispatch, never balanced. + + Fail-closed by construction — an unmapped runner or an unknown shard + raises instead of falling back to another shard's admission; a dead + shard fails exactly its own attempts. + """ + if len(shards) == 1: + return shards[0] + if config.gateway_routing == "worker_affinity": + key = runner_host if runner_host is not None else LOCAL_RUNNER_KEY + name = config.worker_shards.get(key) + if name is None: + raise ValueError(f"worker_affinity has no gateway shard for runner {key!r}") + for shard in shards: + if shard.name == name: + return shard + raise ValueError(f"worker_affinity names unknown gateway shard {name!r}") + return shards[(slot + attempt_index) % len(shards)] + + async def _run_slot(config, runtime_, slots, slot, entry, reward_profile): """Execute one slot with the group's shared identity and retry bookkeeping.""" workers = config.workers overrides = config.task_worker_overrides - check_loop = getattr(runtime_, "check_event_loop", None) - if check_loop is not None: - check_loop() if not hasattr(runtime_, "cleanup"): runtime_.cleanup = CleanupAdmission(runtime_.rollout_root) cleanup = runtime_.cleanup cleanup.check() - last_error = None for attempt_index in range(slots.max_attempts): identity = slots.start(str(slot)) runner_host = overrides.get(entry["id"]) @@ -413,6 +1232,11 @@ async def _run_slot(config, runtime_, slots, slot, entry, reward_profile): if not workers: raise ValueError("rollout runtime requires at least one worker destination") runner_host = workers[(slot + attempt_index) % len(workers)] + auto = config.gateway_routing == "auto" + shard = None + if not auto: + shard = route_shard(config, runtime_.shards, runner_host=runner_host, + slot=slot, attempt_index=attempt_index) root = runtime_.rollout_root / identity.group_id / identity.attempt_id admission = nullcontext() if config.sandbox: @@ -421,46 +1245,57 @@ async def _run_slot(config, runtime_, slots, slot, entry, reward_profile): try: async with admission, runtime_.semaphore: cleanup.check() + if auto: + # Auto routing leases a gateway once this attempt actually + # runs: the (n+1)-th concurrent sandbox, not a queued one, + # is what pulls up the next gateway. Remote growth waits + # on an ssh subprocess, so the lease itself awaits. + shard = await runtime_.lease_shard() try: ir = await run_attempt( identity, task_path=entry["path"], profile=runtime_.profile, - gateway_url=config.gateway_url, + gateway_url=shard.url, runner_python=config.runner_python, runner_host=runner_host, root=root, - registry=runtime_.registry, + registry=shard.registry, reward_profile=reward_profile, max_output_tokens=config.max_output_tokens, **({'sandbox': config.sandbox} if config.sandbox else {}), ) finally: + if auto: + runtime_.release_shard(shard) cleanup.observe(root) - cleanup.check() - except (CleanupRequiredError, RuntimeLifecycleError) as exc: + cleanup.check() + except CleanupRequiredError as exc: slots.finish(str(slot), error=f"{type(exc).__name__}: {exc}") raise except (OSError, ValueError, RuntimeError, asyncio.TimeoutError) as exc: slots.finish(str(slot), error=f"{type(exc).__name__}: {exc}") cleanup.check() - last_error = exc + if config.sandbox and ((root / 'cleanup-required.json').exists() or runtime_.sandbox_admission.blocked): + raise RuntimeError('sandbox cleanup unresolved; automatic retry blocked') from exc continue slots.finish(str(slot), ir) return ir - raise RuntimeError(f"native slot {slot} exhausted its retry budget") from last_error + raise RuntimeError(f"native slot {slot} exhausted its retry budget") -async def gather_or_cancel(coroutines): - """Return ordered results, awaiting sibling cleanup before propagating failure.""" - tasks = [asyncio.create_task(coroutine) for coroutine in coroutines] - try: - return await asyncio.gather(*tasks) - except BaseException: - for task in tasks: - task.cancel() - await asyncio.gather(*tasks, return_exceptions=True) - raise +async def cancel_and_wait(tasks, result): + """Cancel owned tasks once and finish cleanup despite further cancellation.""" + for task in tasks: + task.cancel() + settled = asyncio.gather(result, *tasks, return_exceptions=True) + while True: + try: + await asyncio.shield(settled) + return + except asyncio.CancelledError: + # The caller re-raises its original failure after cleanup completes. + continue async def run_group(config, runtime_, *, entry, group_index, rollout_id, @@ -470,6 +1305,23 @@ async def run_group(config, runtime_, *, entry, group_index, rollout_id, slots, reward_profile = _group_slots( config, runtime_, group_index, rollout_id, entry, sampling_digest, policy_version, size) - return await gather_or_cancel( - _run_slot(config, runtime_, slots, slot, entry, reward_profile) - for slot in range(size)) + tasks = [asyncio.create_task( + _run_slot(config, runtime_, slots, slot, entry, reward_profile)) + for slot in range(size) + ] + result = asyncio.gather(*tasks) + try: + # Cancel children once here, so gather cannot interrupt slower cleanup. + return list(await asyncio.shield(result)) + except BaseException: + await cancel_and_wait(tasks, result) + raise + finally: + # Settlement (success or failure) is the shrink trigger: every lease + # of this group is released by now, so an auto-spawned gateway the + # running sandbox count no longer justifies goes away before the + # next group can lease it. + if getattr(config, "gateway_routing", "auto") == "auto": + retire = getattr(runtime_, "retire_idle_shards", None) + if retire is not None: + retire() diff --git a/harborrl/rollout/service.py b/harborrl/rollout/service.py index 81d2ed273..fc9b8d6ae 100644 --- a/harborrl/rollout/service.py +++ b/harborrl/rollout/service.py @@ -16,9 +16,11 @@ from __future__ import annotations import asyncio +import concurrent.futures import hashlib import hmac import json +import os import secrets import threading import time @@ -60,6 +62,8 @@ def validate_group_request(request): for key in ("group_id", "policy_version"): if not isinstance(request[key], str) or not request[key].strip(): raise ServiceError(400, "invalid_request", f"{key} must be a nonempty string") + if request["group_id"] in (".", "..") or any(c in request["group_id"] for c in ("/", "\\")): + raise ServiceError(400, "invalid_request", "group_id must be one directory component") for key in ("rollout_id", "group_index"): if type(request[key]) is not int or request[key] < 0: raise ServiceError(400, "invalid_request", f"{key} must be a nonnegative integer") @@ -166,6 +170,8 @@ class GroupState: error: str = None created_at: float = field(default_factory=time.time) future: object = None + task: object = None + cancel_requested: bool = False def summary(self, *, include_result=True): payload = {"group_id": self.request["group_id"], "status": self.status, @@ -196,7 +202,9 @@ class RolloutService: def __init__(self, *, run_root, executor=None, runtime_factory=None, token_file=None, group_wait_timeout=3600.0, poll_interval=0.05, - allow_synthetic=False): + allow_synthetic=False, cleanup_timeout=None): + if cleanup_timeout is not None and (not finite(cleanup_timeout) or cleanup_timeout <= 0): + raise ValueError("cleanup_timeout must be positive and finite") self.run_root = Path(run_root).resolve() self.executor = executor or default_executor self.runtime_factory = runtime_factory @@ -211,25 +219,112 @@ def __init__(self, *, run_root, executor=None, runtime_factory=None, self.poll_interval = poll_interval self.runs = {} self.lock = threading.RLock() + self._close_lock = threading.Lock() + self._closing = self._closed = False + self.cleanup_timeout = cleanup_timeout + self._store_lock = self._lock_run_directory() self.loop = asyncio.new_event_loop() self.thread = threading.Thread(target=self._serve_loop, name="harborrl-rollout-service", daemon=True) - self.thread.start() + try: + self.thread.start() + except RuntimeError: + self.loop.close() + self._store_lock.close() + raise + + def _lock_run_directory(self): + directory = self.run_root / "rollout-service" + directory.mkdir(parents=True, exist_ok=True) + stream = (directory / "service.lock").open("a+b") + try: + if os.name == "posix": + import fcntl + fcntl.flock(stream.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + else: + import msvcrt + stream.seek(0) + if not stream.read(1): + stream.write(b"0") + stream.flush() + stream.seek(0) + msvcrt.locking(stream.fileno(), msvcrt.LK_NBLCK, 1) + except OSError as exc: + stream.close() + raise ServiceError(409, "run_directory_busy", "another service owns this run directory") from exc + return stream def _serve_loop(self): asyncio.set_event_loop(self.loop) self.loop.run_forever() - def close(self): + def _cleanup_budget(self): + if self.cleanup_timeout is not None: + return self.cleanup_timeout + # Runner cancellation allows 120s, or the configured Sandbox deletion + # deadline plus 15s. The service must not cut that lifecycle to 30s. + budgets = [120.0] with self.lock: - states = [state for state in self.runs.values() if state.runtime is not None] + states = list(self.runs.values()) for state in states: + sandbox = getattr(state.config, "sandbox", None) + if sandbox: + budgets.append(float(sandbox["cleanup_timeout_s"]) + 15) + return max(budgets) + 30 + + def _ensure_open(self): + if self._closing: + raise ServiceError(503, "service_closing", "rollout service is closing") + + def _shutdown_failed(self, error): + marker = self.run_root / "native-rollouts" / "cleanup-required.json" + if not marker.exists(): + publish(marker, {"reason": "RolloutService shutdown incomplete", + "error": str(error), + "groups": {state.run_id: {key: group.status for key, group in state.groups.items()} + for state in self.runs.values()}}) + + async def _shutdown_loop_resources(self): + await self.loop.shutdown_asyncgens() + await self.loop.shutdown_default_executor() + + def close(self): + if threading.current_thread() is self.thread: + raise RuntimeError("close the rollout service outside its executor loop") + with self._close_lock: + if self._closed: + return + with self.lock: + self._closing = True + states = list(self.runs.values()) + for state in states: + state.drained = True + groups = [group for state in states for group in state.groups.values()] try: - state.runtime.close() - except AttributeError: - pass - self.loop.call_soon_threadsafe(self.loop.stop) - self.thread.join(5) + self._cancel_groups(groups) + for state in states: + if state.runtime is None: + continue + registry = getattr(state.runtime, "registry", None) + if registry is not None: + registry.drain(float(getattr(state.config, "drain_timeout", 30))) + close = getattr(state.runtime, "close", None) + if close is not None: + close() + future = asyncio.run_coroutine_threadsafe(self._shutdown_loop_resources(), self.loop) + future.result(timeout=self._cleanup_budget()) + self.loop.call_soon_threadsafe(self.loop.stop) + self.thread.join(5) + if self.thread.is_alive(): + raise ServiceError(409, "shutdown_timeout", "rollout executor thread has not stopped") + self.loop.close() + self._store_lock.close() + self._closed = True + except Exception as exc: + # Keep the executor alive to finish pending cleanup; never claim + # that a deadline or a cancelled wrapper Future proves deletion. + self._shutdown_failed(exc) + raise # -- run lifecycle ---------------------------------------------------- @@ -237,14 +332,16 @@ def _token_digest(self, token): return hashlib.sha256(token.encode()).hexdigest() def register_run(self, run_id, manifest, *, token=None): - if not isinstance(run_id, str) or not run_id.strip(): - raise ServiceError(400, "invalid_request", "run id required") + if (not isinstance(run_id, str) or not run_id.strip() or run_id in (".", "..") + or any(c in run_id for c in ("/", "\\"))): + raise ServiceError(400, "invalid_request", "run id must be one directory component") if not isinstance(manifest, dict) or not manifest: raise ServiceError(400, "invalid_request", "manifest must be a nonempty mapping") if manifest.get("run_id", run_id) != run_id: raise ServiceError(400, "invalid_request", "manifest run id must match the path") manifest_sha = digest(manifest) with self.lock: + self._ensure_open() existing = self.runs.get(run_id) if existing is not None: if existing.manifest_sha != manifest_sha: @@ -253,14 +350,21 @@ def register_run(self, run_id, manifest, *, token=None): return {"run_id": run_id, "state": "registered", "existed": True} if self.token is not None and token != self.token: raise ServiceError(401, "unauthorized", "invalid manager run token") + manifest_path = self.run_root / "rollout-service/runs" / run_id / "manifest.json" + existed = manifest_path.exists() + if existed and digest(read_json(manifest_path)) != manifest_sha: + raise ServiceError(409, "conflict", "persisted run manifest differs from registration") + groups = self._restore_groups(run_id) + publish(manifest_path, manifest) generated = token is None if generated: token = secrets.token_urlsafe(32) self.runs[run_id] = RunState(run_id=run_id, manifest=manifest, manifest_sha=manifest_sha, - token_digest=self._token_digest(token)) - return {"run_id": run_id, "state": "registered", "existed": False, - **({"run_token": token} if generated else {})} + token_digest=self._token_digest(token), groups=groups, + drained=any(group.status == "unknown" for group in groups.values())) + return {"run_id": run_id, "state": "registered", "existed": existed, + "recovered_groups": len(groups), **({"run_token": token} if generated else {})} def authenticate(self, run_id, token): with self.lock: @@ -281,6 +385,7 @@ def _run(self, run_id): def attach_runtime(self, run_id, config, runtime): """Bind the rollout runtime that executes this run's groups.""" with self.lock: + self._ensure_open() state = self._run(run_id) state.config, state.runtime = config, runtime return {"run_id": run_id, "state": "attached"} @@ -292,6 +397,7 @@ def announce_serving(self, run_id, router_url): or parts.fragment or parts.path not in ("", "/")): raise ServiceError(400, "invalid_request", "router_url must be a plain origin") with self.lock: + self._ensure_open() state = self._run(run_id) state.serving = f"{parts.scheme}://{parts.netloc}".rstrip("/") if state.runtime is None and self.runtime_factory is not None: @@ -310,6 +416,7 @@ def activate_policy(self, run_id, states, expected): from harborrl.platform.policy_pool import validate_pool with self.lock: + self._ensure_open() state = self._run(run_id) if state.runtime is None: raise ServiceError(409, "admission_closed", "run has no serving runtime yet") @@ -324,16 +431,18 @@ def activate_policy(self, run_id, states, expected): states = live_states except (RuntimeError, TypeError, ValueError) as exc: raise ServiceError(409, "pool_invalid", str(exc)) from exc - registry = state.runtime.registry - registry.drain(float(getattr(state.config, "drain_timeout", 30))) - registry.publish_version(states, version) - # Mirror policy_pool.publish_pool into the shared run directory so the - # GPU side can lock the same version without manager-local state; the - # record is replaceable because activation legitimately advances. - payload = {"weight_version": version, "engines": states, "timestamp": time.time()} - temp = self.run_root / "policy-pool.tmp" - temp.write_text(json.dumps(payload) + "\n") - temp.replace(self.run_root / "policy-pool.json") + with self.lock: + # Pool probes above can race a shutdown. Do not reopen generation + # admission after close has frozen the service. + self._ensure_open() + registry = state.runtime.registry + registry.drain(float(getattr(state.config, "drain_timeout", 30))) + registry.publish_version(states, version) + # Mirror the admitted version into shared storage. + payload = {"weight_version": version, "engines": states, "timestamp": time.time()} + temp = self.run_root / "policy-pool.tmp" + temp.write_text(json.dumps(payload) + "\n") + temp.replace(self.run_root / "policy-pool.json") return {"run_id": run_id, "policy_version": version, "admission": "open"} # -- group lifecycle ---------------------------------------------------- @@ -341,6 +450,70 @@ def activate_policy(self, run_id, states, expected): def _group_dir(self, run_id, group_id): return self.run_root / GROUP_DIR / run_id / group_id + def _restore_groups(self, run_id): + """Restore durable outcomes, never resume an unowned Runner or resample.""" + groups = {} + directory = self.run_root / GROUP_DIR / run_id + for root in sorted(directory.glob("*")): + if not root.is_dir(): + continue + request_path = root / "request.json" + if not request_path.exists(): + if any((root / name).exists() for name in ("running.json", "result.json", "failed.json", "cancelled.json", "unknown.json")): + raise ServiceError(500, "invalid_persisted_group", f"missing request: {root.name}") + continue + try: + request = validate_group_request(read_json(request_path)) + if request["group_id"] != root.name: + raise ValueError("group directory differs from request") + sha = request_digest(request) + terminal = [(name, status) for name, status in ( + ("result.json", "complete"), ("failed.json", "failed"), ("cancelled.json", "cancelled")) + if (root / name).exists()] + if len(terminal) > 1: + raise ValueError("conflicting terminal group records") + if terminal: + name, status = terminal[0] + record = read_json(root / name) + else: + status = "unknown" + record_path = root / "unknown.json" if (root / "unknown.json").exists() else root / "running.json" + record = read_json(record_path) if record_path.exists() else { + "group_id": root.name, "request_digest": sha, "created_at": time.time()} + if not isinstance(record, dict): + raise TypeError("persisted group record must be a mapping") + if record.get("group_id") != root.name or record.get("request_digest") != sha: + raise ValueError("persisted group does not match its request") + if terminal and record.get("status") != status: + raise ValueError("terminal status differs from its record") + if status == "complete" and ( + record.get("run_id") != run_id or record.get("schema_version") != SERVICE_SCHEMA + or record.get("policy_version") != request["policy_version"] + or [slot["slot_id"] for slot in record["slots"]] != [slot["slot_id"] for slot in request["slots"]] + ): + raise ValueError("persisted result identity differs from its request") + created_at = record.get("created_at") + if created_at is None and (root / "running.json").exists(): + running = read_json(root / "running.json") + if (not isinstance(running, dict) or running.get("request_digest") != sha + or running.get("group_id") != root.name): + raise ValueError("running record differs from its request") + created_at = running["created_at"] + group = GroupState(request, sha, status=status, + result=record if status == "complete" else None, + error=record.get("error"), created_at=created_at or 0) + if status == "unknown": + group.error = "Service restarted before durable group completion; confirm Runner cleanup and use a new run" + publish(root / "unknown.json", group.summary(include_result=False)) + marker = self.run_root / "native-rollouts/cleanup-required.json" + if not marker.exists(): + publish(marker, {"run_id": run_id, "group_id": root.name, + "reason": "Interrupted group has unknown Runner cleanup status"}) + groups[root.name] = group + except (OSError, ValueError, KeyError, TypeError) as exc: + raise ServiceError(500, "invalid_persisted_group", f"cannot restore {root.name}: {exc}") from exc + return groups + def submit_group(self, run_id, request, *, wait=True): request = validate_group_request(request) sha = request_digest(request) @@ -353,6 +526,7 @@ def submit_group(self, run_id, request, *, wait=True): "group id already submitted with different content") group = existing else: + self._ensure_open() if state.drained: raise ServiceError(409, "admission_closed", "run is draining") if state.runtime is None or state.config is None: @@ -387,23 +561,27 @@ def _wait_group(self, run_id, group): async def _execute(self, state, group): # Immutable state records: each transition gets its own file, and # publish() refuses to rewrite history. + group.task = asyncio.current_task() request = group.request directory = self._group_dir(state.run_id, request["group_id"]) - publish(directory / "request.json", request) - group.status = "running" - publish(directory / "running.json", group.summary(include_result=False)) try: + publish(directory / "request.json", request) + if group.cancel_requested: + raise asyncio.CancelledError + group.status = "running" + publish(directory / "running.json", group.summary(include_result=False)) trajectories = await self.executor(state.config, state.runtime, request) export = export_group( trajectories, len(trajectories), allow_synthetic=self.allow_synthetic) export_path = directory / "group-export.json" publish(export_path, export) - group.result = build_group_result( + result = build_group_result( state.run_id, request, trajectories, runtime=state.runtime, export={"path": str(export_path), "sha256": digest(export)}) + publish(directory / "result.json", result) + group.result = result group.status = "complete" - publish(directory / "result.json", group.result) except asyncio.CancelledError: group.status = "cancelled" publish(directory / "cancelled.json", group.summary(include_result=False)) @@ -428,6 +606,8 @@ def cancel_group(self, run_id, group_id): group = state.groups.get(group_id) if group is None: raise ServiceError(404, "not_found", f"unknown group: {group_id}") + if group.status == "unknown": + raise ServiceError(409, "cleanup_required", "interrupted group has no owned Task; inspect Runner cleanup") self._cancel_groups([group]) return group.summary() @@ -450,19 +630,40 @@ def drain_run(self, run_id): def cancel_run(self, run_id): with self.lock: state = self._run(run_id) + state.drained = True groups = list(state.groups.values()) + if any(group.status == "unknown" for group in groups): + raise ServiceError(409, "cleanup_required", "interrupted groups require external cleanup confirmation") self._cancel_groups(groups) return {"run_id": run_id, "state": "cancelled", "groups": {group.request["group_id"]: group.status for group in groups}} - def _cancel_groups(self, groups): + async def _cancel_and_wait(self, groups): for group in groups: - if group.status in ("queued", "running") and group.future is not None: - group.future.cancel() - deadline = time.monotonic() + 30 - while any(group.status in ("queued", "running") for group in groups) \ - and time.monotonic() < deadline: - time.sleep(self.poll_interval) + if not group.cancel_requested and group.future is not None and not group.future.done(): + group.cancel_requested = True + if group.task is not None: + group.task.cancel() + # Do not cancel the run_coroutine_threadsafe Future: that wrapper can + # become done before the actual Task's asynchronous finally completes. + pending = [asyncio.wrap_future(group.future) for group in groups if group.future is not None] + results = await asyncio.gather(*pending, return_exceptions=True) + for result in results: + if isinstance(result, BaseException) and not isinstance(result, asyncio.CancelledError): + raise ServiceError(500, "cleanup_failed", f"group completion failed: {result}") + + def _cancel_groups(self, groups): + if not any(group.future is not None for group in groups): + return + if threading.current_thread() is self.thread: + raise RuntimeError("cancel rollout groups outside their executor loop") + joined = asyncio.run_coroutine_threadsafe(self._cancel_and_wait(groups), self.loop) + try: + joined.result(timeout=self._cleanup_budget()) + except concurrent.futures.TimeoutError as exc: + # Leave the waiter and group cleanup running; a second cancel would + # interrupt their finally blocks. Retrying close may join them later. + raise ServiceError(409, "cleanup_timeout", "rollout groups are still cleaning up") from exc def make_control_server(service, host="127.0.0.1", port=0, *, max_body_bytes=8 * 1024 * 1024): diff --git a/harborrl/trajectories/native.py b/harborrl/trajectories/native.py index eaf9ed32e..230e1fd89 100644 --- a/harborrl/trajectories/native.py +++ b/harborrl/trajectories/native.py @@ -167,6 +167,10 @@ def validate_turn(turn, identity, trial_id): for key in ("client_request", "serving_input", "response"): if key not in turn or turn.get(key + "_digest") != digest(turn[key]): errors.append(key + "_digest") + declared = {tool['name'] for tool in turn.get('client_request', {}).get('tools', [])} + if any(block.get('type') == 'tool_use' and block.get('name') not in declared + for block in turn.get('response', {}).get('content', [])) and not turn.get('tool_rejection_ref'): + errors.append('unconfirmed_tool_rejection') if turn.get("finish_reason") not in ("end_turn", "tool_use", "max_tokens", "stop_sequence"): errors.append("finish_reason") return errors @@ -210,8 +214,14 @@ def validate_v2(ir): errors.append("mixed_tokenizer_template") if ir.get("termination") not in ("completed", "budget_truncated"): errors.append("termination") - if ir.get("termination") == "budget_truncated" and not ir.get("termination_evidence"): - errors.append("truncation_evidence") + if ir.get("termination") == "budget_truncated": + evidence = ir.get('termination_evidence') + if (not isinstance(evidence, dict) or evidence != receipt.get('termination_evidence') + or evidence.get('source_result_digest') != receipt.get('source_result_digest') + or evidence.get('consumed_response_ids') != [t['response_id'] for t in turns]): + errors.append("truncation_evidence") + elif ir.get('termination_evidence') is not None or receipt.get('termination_evidence') is not None: + errors.append('unexpected_termination_evidence') except (KeyError, TypeError, ValueError, AttributeError) as exc: errors.append(f"malformed_contract:{type(exc).__name__}") return sorted(set(errors)) diff --git a/practices/h_cluster/HISTORICAL_VALIDATION.md b/practices/h_cluster/HISTORICAL_VALIDATION.md new file mode 100644 index 000000000..664de6e68 --- /dev/null +++ b/practices/h_cluster/HISTORICAL_VALIDATION.md @@ -0,0 +1,32 @@ +# Historical platform validation, 2026-10-20 + +This predates the controller integration and practice package extraction. +See README.md for the subsequent relay results and current validation limits. + +Project identity was verified through authenticated Sandbox and registry APIs. +The platform AccessKey can log into the primary registry. The Sandbox registry +uses separate credentials obtained through the platform's Harbor v2 account API; +the same AccessKey is not itself the Sandbox registry username/password. + +A dedicated task image was pushed successfully and a dedicated Environment was +created with capacity 2 and prewarm 0. A GPU-local runner completed real Sandbox +creation/readiness, positive and negative command exit-code checks, binary file +upload/download, and acknowledged deletion. These checks used the existing +prepared humanevalfix task image, not a newly built GPU environment image. + +Gateway reachability is blocked. The GPU-local HTTP probe responds with HTTP 200 +on loopback. Sandbox connections to two GPU interface addresses timed out. The +Environment initially ignored egressEnabled during creation. A subsequent PATCH +requires networkPolicy; setting default deny with explicit GPU /32 allow rules +succeeded, but the connection still timed out. No training was launched. + +One allocation timed out during SDK endpoint initialization. SDK cleanup was +followed by an explicit API check: the instance returned 404 and the Sandbox list +was empty. Subsequent controlled validation succeeded through command/files but +again failed Gateway reachability, then acknowledged instance deletion. + +Next gate: platform networking must provide a Sandbox-reachable GPU listener or +resolve the current route/policy restriction. Do not claim RL_READY or training +acceptance until a live group passes the existing session and token audits. +Deployment-specific IDs, addresses and detailed probe evidence are maintained +outside the public source tree in the validation directory on shared storage. diff --git a/practices/h_cluster/README.md b/practices/h_cluster/README.md index 36fcd7cd3..db669f3f5 100644 --- a/practices/h_cluster/README.md +++ b/practices/h_cluster/README.md @@ -78,6 +78,61 @@ both before reusing capacity. Never delete templates or instances belonging to other runs. The framework records unresolved cleanup rather than assuming it succeeded. It does not automatically replace platform Environment templates. +## GPU jobs and outside hosts + +Two optional adapters cover machines the core deliberately does not know +about. Both are site adapters; the core only pins their contracts. + +**Provisioned GPU jobs.** A deployment may declare `gpu_target.provision` +instead of a fixed `gpu_target.ssh`: + +```yaml +gpu_target: + id: trainer + transport: ssh + ssh: null + provision: + up: ["python", "-m", "harborrl_h_cluster.rjob", "up", "--config", "/srv/harborrl/site/rjob.json"] + down: ["python", "-m", "harborrl_h_cluster.rjob", "down", "--config", "/srv/harborrl/site/rjob.json"] +``` + +The core runs `up` before anything dials the machine, reruns the deferred GPU +probes against the machine it obtained, and runs `down` in its `finally` path — +a failed launch, a cancelled wait or a crashed run all still release it. `up` +prints the machine record as its last stdout line: `host`, optional `user`, +`port`, `identity` and an opaque `handle`; any other key is rejected. `down` +receives the same record on stdin. Nothing secret belongs in either direction. + +`rjob up` refuses to submit while a previous pod record exists under its +`workdir`: a second submit silently spends a second allocation. Release test +pods with `down`; the release outcome is published even when it fails, so a +leaked pod is visible instead of assumed clean. + +**A host the cluster cannot dial.** When a machine has no inbound route but can +reach the hub, it serves shard ports through reverse forwards: + +```bash +python -m harborrl_h_cluster.tunnel install --config /srv/harborrl/site/tunnel.json +``` + +It prints two commands as JSON, because one paste cannot do both ends: the +host's key is not authorized yet, so it cannot reach the hub until the hub +side runs. Run `hub` on the hub (it authorizes the public key, idempotently), +then run `host` in the outside host's interactive session (it pulls the +package over the same one-way path, writes the site config, and starts the +supervised daemon). Both are idempotent and carry no private key material. + +The network portal swallows non-interactive input, so that paste is the only +installation channel; the platform may wipe the host's `/root` and take the +installed key with it, in which case the daemon's state file names the failure +and the pair must be run again. Declare the shard in +`manager.gateway_shards` pointing at the hub's `127.0.0.1:`, and keep +`gateway_hosts` for elastic growth; the two may coexist. + +Such a host has no shared storage. `push` copies declared directories back to +the hub over the same ssh direction with bounded retries, and never exposes a +partially written tree. + ## Configuration migration Experimental flat Brainbox configs are intentionally not accepted as standard diff --git a/practices/h_cluster/VALIDATION.md b/practices/h_cluster/VALIDATION.md new file mode 100644 index 000000000..c5e525328 --- /dev/null +++ b/practices/h_cluster/VALIDATION.md @@ -0,0 +1,76 @@ +# Merged Controller and Sandbox validation + +Validation date: 2026-10-22. This is an infrastructure acceptance record, not a +model benchmark. Private addresses, credentials, task artifacts and platform IDs +are intentionally excluded. + +## Setup + +- One developer host runs Controller, RolloutService, Gateway and Harbor Runner. +- The optional H-cluster adapter creates real instances through OpenSandbox SDK + 0.1.9. Harbor 0.23.0 runs a pinned Claude Code 2.1.141 binary in each instance. +- Sandbox model requests use the developer mailbox relay. The developer reaches + GPU SGLang; Sandbox does not need a direct GPU route. +- Eight H200 GPUs run Qwen3-8B: four actor GPUs with TP4 and four serving GPUs + with TP4. All tests retain strict serving and consumption evidence checks. + +## Observed results + +| Test | Running instance peak | Actual agent execution peak | Result | +|---|---:|---:|---| +| Serial baseline | 1 | 1 | Four real trajectories, two training iterations and two checkpoints | +| SDK command capacity probe | 8 ready | Not an agent test | Eight concurrent ten-second commands completed in 11.037 seconds | +| Four-dataset stress | 4 | 2 | Batch rejected; four valid trajectories, three failed groups | +| Eight-dataset stress | 6 | 3 | Batch rejected; four valid trajectories, six failed groups | +| Three-dataset regression | 6 | 6 | Batch rejected after an agent API error | +| Two-dataset regression before actor fix | 6 | 6 | Six valid trajectories; actor incorrectly rejected two groups' weights | +| Identical regression after actor fix | 6 | 6 | First complete batch trained and saved a checkpoint; second batch hit an unknown generated tool | + +The two-dataset regression uses AARR and an existing HumanEvalFix task, three +samples per group, two groups per batch, a total instance limit of six, three +instances per Environment, startup concurrency four, and a 32-turn CLI limit. +These tasks were selected after the broad stress tests to isolate infrastructure +behavior. They are not an unbiased dataset sample. Each test starts from the +reference model; successful subsets from rejected batches are never repackaged +as new training batches. + +The SDK path demonstrably supports three simultaneous instances from one +Environment in this deployment. This does not override per-Environment capacity, +project quotas or platform scheduling. Pending allocations can take minutes. + +## Core fixes + +- `9a84641a`: port preflight handles immediate HTTP-server restarts after + TIME_WAIT while rejecting an existing listener. Real socket tests cover both. +- `dc208dc9`: actor postprocessing validates each group's local weight budget + using its unique rollout ID, rather than limiting a multi-group batch to one. + Partial actor shards remain valid; excess weight in one group cannot be hidden + by another group's smaller weight. Regression tests cover both properties. + +After these fixes, all 217 CPU tests passed, including the optional practice +tests. Changed Python files passed Ruff and the public-tree audit passed. + +## Limits + +Stable eight-agent mixed-dataset training has not passed. Observed failures +include Pending/startup errors, missing offline verifier dependencies, CLI turn +limits, unknown generated tools, and CLI tool-argument normalization that fails +exact consumption auditing. The relay is not a general internet proxy. + +In the second batch of the final regression, the model generated an undeclared +`finish` tool. Gateway returned a protocol error and the trial was rejected even +though its verifier reward was positive. This must not be bypassed or converted +into a synthetic successful trajectory. + +Successful training batches had uniform rewards within each group and zero +gradient norm. Optimizer execution, checkpoint creation and serving policy +version changes establish workflow execution, not learning improvement or +nonzero-gradient scaling correctness. The current `native_group_count` training +metric rounds microbatch weights and must not be used to count complete groups; +use RolloutService group results and audited trajectories instead. + +To repeat this test, configure private site files using the practice examples, +pin task/image/binary digests, start with a serial complete-group test, then +increase binding and global concurrency within actual platform limits. Record +agent start/end intervals separately from instance states. Retain rejected +groups and verify deletion of every allocated instance after completion. diff --git a/practices/h_cluster/configs/rjob.example.json b/practices/h_cluster/configs/rjob.example.json new file mode 100644 index 000000000..79d00cb11 --- /dev/null +++ b/practices/h_cluster/configs/rjob.example.json @@ -0,0 +1,23 @@ +{ + "_comment": "Replace every placeholder. Referenced by gpu_target.provision.up/down in a deployment config. No credential belongs in this file.", + "name_prefix": "harborrl-train", + "namespace": "REPLACE-project-namespace", + "charged_group": "REPLACE-charged-group", + "private_machine": "group", + "image": "REPLACE-registry/image:tag", + "image_pull_policy": "IfNotPresent", + "mounts": [ + "REPLACE-storage-scheme://REPLACE-volume/REPLACE-export:REPLACE-mount-point" + ], + "gpu": 4, + "cpu": 64, + "memory": 400000, + "priority": 1, + "workdir": "/REPLACE/shared/path/for/this/run", + "identity": "/REPLACE/path/to/manager_id_ed25519", + "ssh_user": "root", + "ssh_port": 22, + "timeout_s": 1800, + "poll_interval_s": 15, + "alive_seconds": 21600 +} diff --git a/practices/h_cluster/configs/tunnel.example.json b/practices/h_cluster/configs/tunnel.example.json new file mode 100644 index 000000000..334ed657e --- /dev/null +++ b/practices/h_cluster/configs/tunnel.example.json @@ -0,0 +1,15 @@ +{ + "_comment": "Replace every placeholder. Run `python -m harborrl_h_cluster.tunnel install --config this-file` and paste the printed line into the outside host's interactive portal session once.", + "hub": "REPLACE-user@REPLACE-manager-host", + "identity": "/REPLACE/path/to/outside_host_id_ed25519", + "state_path": "/REPLACE/path/tunnel-state.json", + "mappings": [ + { "local_port": 5010, "hub_bind": "127.0.0.1", "hub_port": 15010 }, + { "local_port": 5011, "hub_bind": "127.0.0.1", "hub_port": 15011 } + ], + "push": [ + "/REPLACE/local/traces:REPLACE/remote/incoming/traces" + ], + "backoff_start": 1.0, + "backoff_cap": 60.0 +} diff --git a/practices/h_cluster/src/harborrl_h_cluster/__init__.py b/practices/h_cluster/src/harborrl_h_cluster/__init__.py index e8e8a73b2..c2d7a3fb7 100644 --- a/practices/h_cluster/src/harborrl_h_cluster/__init__.py +++ b/practices/h_cluster/src/harborrl_h_cluster/__init__.py @@ -1 +1,20 @@ -"""Optional H-cluster deployment adapters; not imported by HarborRL core.""" +"""Optional H-cluster deployment adapters; not imported by HarborRL core. + +Modules in this package, each usable alone: + +``environment`` OpenSandbox platform environment adapter (site Environment + IDs, mounts, command port) used by the sandbox runtime. +``transport`` bounded retries for platform endpoint discovery; allocation + POSTs are never retried. +``relay`` manager-side relay pump for sandboxes that may only reach the + platform proxy. +``relay_server`` the loopback mailbox the relay and a sandbox share. +``reverse_relay`` manager-initiated variant of the same relay. +``rjob`` ``gpu_target.provision`` implementation: submit and release a + GPU pod, print the machine record the core consumes. +``tunnel`` reverse ssh tunnel so a host the cluster cannot dial can still + serve gateway shard ports, plus the artifact push back. + +Nothing here is imported at package import time; each adapter is loaded by the +component that needs it. +""" diff --git a/practices/h_cluster/src/harborrl_h_cluster/environment.py b/practices/h_cluster/src/harborrl_h_cluster/environment.py index 6b0b68527..9639c1515 100644 --- a/practices/h_cluster/src/harborrl_h_cluster/environment.py +++ b/practices/h_cluster/src/harborrl_h_cluster/environment.py @@ -96,7 +96,7 @@ def api_params(self): def startup_transport(self, allocated): return SandboxStartupTransport(self.cfg['domain'].rstrip('/') + '/v1', allocated, - timeout=self.cfg['create_timeout_s']) + timeout=self.cfg['create_timeout_s'], on_request=self.record_request) def _build_resource(self): resource = super()._build_resource() diff --git a/practices/h_cluster/src/harborrl_h_cluster/rjob.py b/practices/h_cluster/src/harborrl_h_cluster/rjob.py new file mode 100644 index 000000000..7836a2248 --- /dev/null +++ b/practices/h_cluster/src/harborrl_h_cluster/rjob.py @@ -0,0 +1,319 @@ +"""rjob site adapter: obtain and release a GPU pod through the brainpp CLI. + +The core speaks one provisioning contract (``harborrl.deployment.provision``): +``up`` prints one JSON record naming the machine it produced, ``down`` +receives that record on stdin and releases it. This module implements that +contract for the H cluster, whose GPU jobs are submitted with ``rjob``. + +Two facts about ``rjob`` shape the implementation. It has no ``exec`` action, +so the pod's command must be a script written *before* submission — the +script bootstraps sshd, publishes the pod's address, then sleeps so the pod +does not exit underneath us. And a submitted pod is billed until it is +deleted, so ``up`` refuses to submit while a previous handle is still +recorded: a double submit quietly burns a second allocation. +""" +from __future__ import annotations + +import argparse +import json +import os +import re +import shlex +import subprocess +import sys +import time +from pathlib import Path + +DEFAULTS = { + "namespace": "ailab-datafoundation", + "charged_group": "datafoundation_gpu_pool", + "private_machine": "group", + "image_pull_policy": "IfNotPresent", + "mounts": [], + "gpu": 4, + "cpu": 64, + "memory": 400000, + "priority": 1, + "ssh_user": "root", + "ssh_port": 22, + "timeout_s": 1800, + "poll_interval_s": 15, + "alive_seconds": 21600, +} +REQUIRED = ("name_prefix", "image", "workdir", "identity") +MOUNT_TARGET = re.compile(r"^[A-Za-z0-9_./:@-]+$") + +# The submission tool and the clock are indirections so the tests can drive +# the whole lifecycle without a cluster or a real wait. +RUNNER = subprocess.run +CLOCK = time + + +class RjobError(RuntimeError): + """A submission, discovery or release step did not behave as declared.""" + + +def load_config(path): + """Read and validate one site config; never accept a half-known file.""" + try: + raw = json.loads(Path(path).read_text()) + except OSError as exc: + raise RjobError(f"cannot read rjob config {path}: {exc}") from exc + except ValueError as exc: + raise RjobError(f"rjob config {path} is not JSON: {exc}") from exc + if not isinstance(raw, dict): + raise RjobError("rjob config must be a JSON object") + unknown = sorted(set(raw) - set(DEFAULTS) - set(REQUIRED)) + if unknown: + raise RjobError(f"rjob config names unsupported fields {unknown}") + config = {**DEFAULTS, **raw} + missing = [key for key in REQUIRED if not config.get(key)] + if missing: + raise RjobError(f"rjob config requires {missing}") + for key in ("name_prefix", "image", "identity"): + if not isinstance(config[key], str) or not config[key].strip(): + raise RjobError(f"rjob config {key} must be a nonempty string") + workdir = Path(config["workdir"]) + if not workdir.is_absolute(): + # The pod mounts this directory and reads the bootstrap from it, so a + # relative path would resolve differently on the two sides. + raise RjobError("rjob config workdir must be an absolute path") + config["workdir"] = str(workdir) + config["identity"] = str(Path(config["identity"])) + mounts = config.get("mounts") or [] + if not isinstance(mounts, list) or not all( + isinstance(mount, str) and MOUNT_TARGET.match(mount) for mount in mounts): + raise RjobError("rjob config mounts is a list of local:target paths") + config["mounts"] = list(mounts) + for key in ("gpu", "cpu", "memory", "priority", "ssh_port", "timeout_s"): + if type(config[key]) is not int or config[key] <= 0: + raise RjobError(f"rjob config {key} must be a positive integer") + if type(config["poll_interval_s"]) not in (int, float) or config["poll_interval_s"] <= 0: + raise RjobError("rjob config poll_interval_s must be a positive number") + if not isinstance(config["alive_seconds"], int) or config["alive_seconds"] <= 0: + raise RjobError("rjob config alive_seconds must be a positive integer") + return config + + +def _run(argv, *, timeout=None, stdin=None): + try: + return RUNNER(argv, capture_output=True, text=True, + input=stdin, timeout=timeout, check=False, + env=dict(os.environ)) + except OSError as exc: + raise RjobError(f"cannot run {argv[0]}: {exc}") from exc + except subprocess.TimeoutExpired as exc: + raise RjobError(f"{' '.join(argv[:2])} exceeded {timeout}s") from exc + + +def _public_key(identity): + """Derive the pod-side authorized key from the manager's private key. + + Only the public half ever enters the generated script; the private key + stays on the manager, where the core's ssh calls read it by path. + """ + result = _run(["ssh-keygen", "-y", "-f", identity], timeout=30) + if result.returncode != 0 or not result.stdout.strip(): + raise RjobError( + f"cannot derive a public key from {identity}: {result.stderr.strip()[-200:]}") + return result.stdout.strip() + + +def bootstrap_script(config, public_key): + """The pod command: bring sshd up, publish the address, then stay alive. + + ``rjob`` cannot execute anything inside a running pod, so this script is + the only channel from the site into the pod. It is idempotent — re-running + it appends nothing twice and starts no second sshd — because a pod may be + re-entered by hand while the acceptance run is being diagnosed. + """ + workdir = shlex.quote(config["workdir"]) + pubkey = shlex.quote(public_key) + return f"""#!/usr/bin/env bash +set -uo pipefail +mkdir -p /root/.ssh && chmod 700 /root/.ssh +touch /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys +grep -qxF {pubkey} /root/.ssh/authorized_keys 2>/dev/null || echo {pubkey} >> /root/.ssh/authorized_keys +if ! pgrep -x sshd >/dev/null 2>&1; then + ssh-keygen -A >/dev/null 2>&1 + /usr/sbin/sshd +fi +printf '{{"hostname":"%s","ips":"%s"}}\\n' "$(hostname)" "$(hostname -I)" > {workdir}/pod-net.json +sleep {config['alive_seconds']} +""" + + +def submit_argv(config, script_path): + argv = ["rjob", "submit", f"--name={config['name_prefix']}", + "--enable-sshd", + f"--gpu={config['gpu']}", f"--cpu={config['cpu']}", + f"--memory={config['memory']}", "-P", str(config["priority"]), + f"--charged-group={config['charged_group']}", + f"--private-machine={config['private_machine']}", + f"--namespace={config['namespace']}", + f"--image={config['image']}", + f"--image-pull-policy={config['image_pull_policy']}"] + for mount in config["mounts"]: + argv.append(f"--mount={mount}") + argv += ["--", "bash", str(script_path)] + return argv + + +CREATED = re.compile(r"created rjob_name:\s*(\S+)") +# A replica line reads "|- replica : STATE, "; the trailing node +# name means the state is not at end-of-line, and a summary line may instead +# end in a human state word such as "Inqueue" — only the uppercase token is +# the machine state these commands are asked about. +STATE = re.compile(r":\s*([A-Z][A-Z_]*)\s*(?:,|$)") + + +def _job_name(stdout): + matches = CREATED.findall(stdout or "") + if not matches: + raise RjobError("rjob submit did not report the created job name") + return matches[-1] + + +def job_state(name, config): + """The replica state of one job, or None when the job is gone. + + Three outcomes, and they must not be confused: a live job answers with a + replica line; a deleted job answers successfully with no replica line at + all; a broken probe exits nonzero. Treating the last as "gone" would + release the record for a pod that is still billed, so it raises instead. + """ + result = _run(["rjob", "get", name, f"--namespace={config['namespace']}"], timeout=120) + if result.returncode != 0: + raise RjobError( + f"rjob get {name} failed with {result.returncode}: " + f"{result.stderr.strip()[-300:]}") + states = [match for match in STATE.findall(result.stdout or "") if match] + # No replica line means the job no longer exists; that is a legitimate + # answer, not an unknown state to wait on forever. + return states[-1] if states else None + + +def pod_address(config): + """The address record the bootstrap wrote into the shared workdir.""" + path = Path(config["workdir"]) / "pod-net.json" + try: + record = json.loads(path.read_text()) + except (OSError, ValueError): + return None + ips = str(record.get("ips") or "").split() + # The pod network address is the first global-looking IPv4; loopback and + # the docker bridge are never reachable from the manager. + for candidate in ips: + if re.fullmatch(r"\d{1,3}(?:\.\d{1,3}){3}", candidate) and not candidate.startswith( + ("127.", "172.17.")): + return candidate + return None + + +def _ssh_ok(config, host, *, timeout=30): + argv = ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=accept-new", + "-o", "ClearAllForwardings=yes", "-o", "ConnectTimeout=10", + "-i", config["identity"], "-p", str(config["ssh_port"]), + f"{config['ssh_user']}@{host}", "echo ok"] + result = _run(argv, timeout=timeout) + return result.returncode == 0 and result.stdout.strip() == "ok" + + +def _wait(predicate, *, timeout, interval, description): + deadline = CLOCK.monotonic() + timeout + last = None + while CLOCK.monotonic() < deadline: + last = predicate() + if last: + return last + CLOCK.sleep(interval) + raise RjobError(f"timed out after {timeout}s waiting for {description}") + + +def up(config, *, record_path=None): + """Submit one pod and return the provision record naming it.""" + handle_path = Path(config["workdir"]) / "rjob-handle.json" + net_path = Path(config["workdir"]) / "pod-net.json" + if net_path.exists() or handle_path.exists(): + # A leftover record means an earlier submission may still be billed; + # refuse instead of silently stacking a second pod on the quota. + raise RjobError( + f"a previous pod record exists under {config['workdir']}; " + "release it with down or remove the stale files deliberately") + workdir = Path(config["workdir"]) + workdir.mkdir(parents=True, exist_ok=True) + script = workdir / "bootstrap.sh" + script.write_text(bootstrap_script(config, _public_key(config["identity"]))) + script.chmod(0o755) + + result = _run(submit_argv(config, script), timeout=600) + if result.returncode != 0: + raise RjobError( + f"rjob submit failed with {result.returncode}: {result.stderr.strip()[-400:]}") + name = _job_name(result.stdout) + print(f"submitted {name}", file=sys.stderr) + handle_path.write_text(json.dumps({"handle": name}) + "\n") + + _wait(lambda: job_state(name, config) in ("ACTIVE", "RUNNING"), + timeout=config["timeout_s"], interval=config["poll_interval_s"], + description=f"job {name} to become active") + host = _wait(lambda: pod_address(config), timeout=config["timeout_s"], + interval=config["poll_interval_s"], + description=f"pod {name} to publish its address") + _wait(lambda: _ssh_ok(config, host), timeout=config["timeout_s"], + interval=config["poll_interval_s"], + description=f"non-interactive ssh to {host}") + + record = {"host": host, "user": config["ssh_user"], + "port": config["ssh_port"], "identity": config["identity"], + "handle": name} + if record_path is not None: + Path(record_path).write_text(json.dumps(record) + "\n") + return record + + +def down(config, record): + """Delete the recorded pod; exit nonzero if it cannot be released.""" + handle = (record or {}).get("handle") + if not isinstance(handle, str) or not handle.strip(): + raise RjobError("release record carries no handle; refusing to guess a job name") + result = _run(["rjob", "delete", handle, f"--namespace={config['namespace']}"], + timeout=600) + if result.returncode != 0: + raise RjobError( + f"rjob delete {handle} failed with {result.returncode}: " + f"{result.stderr.strip()[-400:]}") + _wait(lambda: job_state(handle, config) is None, timeout=config["timeout_s"], + interval=config["poll_interval_s"], description=f"job {handle} to disappear") + for name in ("rjob-handle.json", "pod-net.json"): + with __import__("contextlib").suppress(OSError): + (Path(config["workdir"]) / name).unlink() + return handle + + +def main(argv=None): + parser = argparse.ArgumentParser(prog="harborrl_h_cluster.rjob", + description="GPU pod provisioning through rjob") + sub = parser.add_subparsers(dest="action", required=True) + for action in ("up", "down"): + cmd = sub.add_parser(action) + cmd.add_argument("--config", required=True) + args = parser.parse_args(argv) + config = load_config(args.config) + try: + if args.action == "up": + record = up(config) + print(json.dumps(record)) + return 0 + payload = sys.stdin.read().strip() + if not payload: + raise RjobError("down reads the provision record as JSON on stdin") + down(config, json.loads(payload)) + return 0 + except (RjobError, ValueError) as exc: + print(f"rjob practice: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/practices/h_cluster/src/harborrl_h_cluster/transport.py b/practices/h_cluster/src/harborrl_h_cluster/transport.py index b72151681..f3f8c65be 100644 --- a/practices/h_cluster/src/harborrl_h_cluster/transport.py +++ b/practices/h_cluster/src/harborrl_h_cluster/transport.py @@ -9,8 +9,8 @@ class SandboxStartupTransport(AllocationTransport): - def __init__(self, base_url, on_created, *, timeout, transport=None, retry_delay=2): - super().__init__(base_url, on_created, transport=transport) + def __init__(self, base_url, on_created, *, timeout, transport=None, retry_delay=2, on_request=None): + super().__init__(base_url, on_created, transport=transport, on_request=on_request) self.deadline = time.monotonic() + timeout self.retry_delay = retry_delay diff --git a/practices/h_cluster/src/harborrl_h_cluster/tunnel.py b/practices/h_cluster/src/harborrl_h_cluster/tunnel.py new file mode 100644 index 000000000..79e9df312 --- /dev/null +++ b/practices/h_cluster/src/harborrl_h_cluster/tunnel.py @@ -0,0 +1,344 @@ +"""Reverse-tunnel adapter: serve shard ports from a host the cluster cannot dial. + +The H cluster gives every machine an inbound port only inside the cluster; a +host outside it that is entered through an interactive web portal has no +inbound route at all. The one direction that works is the outside host +initiating ssh to the hub. This module exploits exactly that: the outside +host opens ssh connections carrying ``-R`` forwards, so the hub binds the +ports, and the in-cluster manager dials ``127.0.0.1:`` as if the +shard lived next door. + +Three properties drive the design. The tunnel is supervised, because the jump +gateway drops idle connections and a silent death would look like a shard +crash. The state file records why the last connection ended, because the +platform may wipe the outside host's ``/root`` and take the installed keys +with it — an operator needs to see that, not infer it. And artifacts are +pushed back over the same ssh direction, because the outside host has no +shared storage: anything written there dies with the host. +""" +from __future__ import annotations + +import argparse +import json +import os +import shlex +import signal +import subprocess +import sys +import time +from pathlib import Path + +SSH_OPTIONS = ( + "-N", + "-o", "ExitOnForwardFailure=yes", + "-o", "BatchMode=yes", + "-o", "ServerAliveInterval=15", + "-o", "ServerAliveCountMax=3", + "-o", "StrictHostKeyChecking=accept-new", + "-o", "ClearAllForwardings=yes", +) +BACKOFF_START = 1.0 +BACKOFF_CAP = 60.0 + +RUNNER = subprocess.run +POPEN = subprocess.Popen +CLOCK = time + + +class TunnelError(RuntimeError): + """The tunnel configuration or a tunnel step did not behave as declared.""" + + +def load_config(path): + """Read and validate one site config; a bad mapping fails before connecting.""" + try: + raw = json.loads(Path(path).read_text()) + except OSError as exc: + raise TunnelError(f"cannot read tunnel config {path}: {exc}") from exc + except ValueError as exc: + raise TunnelError(f"tunnel config {path} is not JSON: {exc}") from exc + if not isinstance(raw, dict): + raise TunnelError("tunnel config must be a JSON object") + unknown = sorted(set(raw) - {"hub", "identity", "mappings", "state_path", + "push", "backoff_cap", "backoff_start"}) + if unknown: + raise TunnelError(f"tunnel config names unsupported fields {unknown}") + hub = raw.get("hub") + if not isinstance(hub, str) or not hub.strip() or hub.startswith("-"): + raise TunnelError("tunnel config hub must be an ssh destination") + identity = raw.get("identity") + if not isinstance(identity, str) or not identity.strip(): + raise TunnelError("tunnel config identity must be a key path") + mappings = raw.get("mappings") + if not isinstance(mappings, list) or not mappings: + raise TunnelError("tunnel config mappings must be a nonempty list") + normalized, hub_ports = [], set() + for index, entry in enumerate(mappings): + label = f"tunnel config mappings[{index}]" + if not isinstance(entry, dict) or set(entry) != {"local_port", "hub_bind", "hub_port"}: + raise TunnelError(f"{label} requires exactly local_port, hub_bind and hub_port") + for key in ("local_port", "hub_port"): + if type(entry[key]) is not int or not 0 < entry[key] < 65536: + raise TunnelError(f"{label}.{key} must be an integer in 1..65535") + bind = entry["hub_bind"] + if not isinstance(bind, str) or not bind.strip(): + raise TunnelError(f"{label}.hub_bind must be an address to bind on the hub") + if entry["hub_port"] in hub_ports: + # Two forwards onto one hub port: ssh silently keeps the first, + # so the second shard would never be reachable. + raise TunnelError(f"duplicate hub port {entry['hub_port']}") + hub_ports.add(entry["hub_port"]) + normalized.append({"local_port": entry["local_port"], "hub_bind": bind, + "hub_port": entry["hub_port"]}) + state_path = raw.get("state_path") + if not isinstance(state_path, str) or not state_path.strip(): + raise TunnelError("tunnel config state_path must be a path") + push = raw.get("push") or [] + if not isinstance(push, list) or not all(isinstance(item, str) and item.strip() for item in push): + raise TunnelError("tunnel config push is a list of local:remote directories") + config = { + "hub": hub.strip(), "identity": identity.strip(), "mappings": normalized, + "state_path": state_path, "push": list(push), + "backoff_start": float(raw.get("backoff_start", BACKOFF_START)), + "backoff_cap": float(raw.get("backoff_cap", BACKOFF_CAP)), + } + if config["backoff_start"] <= 0 or config["backoff_cap"] < config["backoff_start"]: + raise TunnelError("tunnel config backoff must start positive and not exceed the cap") + return config + + +def forward_argv(config): + """The exact ssh argv: one -R per mapping, stable option order. + + Each forward binds ``hub_bind:hub_port`` on the hub and connects it to + ``127.0.0.1:local_port`` on this side. The local side is always loopback: + binding a shard process on a routable address would expose it to the + outside network for no gain. + """ + argv = ["ssh", *SSH_OPTIONS, "-i", config["identity"]] + for mapping in config["mappings"]: + argv += ["-R", f"{mapping['hub_bind']}:{mapping['hub_port']}:" + f"127.0.0.1:{mapping['local_port']}"] + argv.append(config["hub"]) + return argv + + +def write_state(path, config, *, pid, attempts, last_error=None, connected=False): + """Record what the supervisor is doing; the operator reads this, not a log.""" + payload = { + "pid": pid, "connected": bool(connected), "attempts": int(attempts), + "last_error": last_error, "updated": CLOCK.strftime("%Y-%m-%dT%H:%M:%S"), + "mappings": list(config["mappings"]), "hub": config["hub"], + } + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + temp = target.with_suffix(".tmp") + temp.write_text(json.dumps(payload, sort_keys=True) + "\n") + temp.replace(target) + + +def daemon(config, *, stop=None): + """Supervise the tunnel: run ssh, wait, reconnect with capped backoff. + + Returns the number of completed connection attempts. A caller passes a + ``threading.Event`` (or any object with ``is_set``) to stop the loop from + another thread; SIGINT and SIGTERM do the same from the process side. + """ + stop = stop or _StopFlag() + previous = {} + for signame in ("SIGINT", "SIGTERM"): + sig = getattr(signal, signame, None) + if sig is None: + continue + previous[sig] = signal.getsignal(sig) + signal.signal(sig, lambda *_args: stop.set()) + delay, attempts = config["backoff_start"], 0 + try: + while not stop.is_set(): + attempts += 1 + write_state(config["state_path"], config, pid=os.getpid(), + attempts=attempts, connected=True) + child = None + error = None + try: + child = POPEN(forward_argv(config), + stdout=subprocess.DEVNULL, stderr=subprocess.PIPE) + while not stop.is_set(): + if child.poll() is not None: + break + CLOCK.sleep(0.2) + if stop.is_set(): + break + stderr = b"" + try: + stderr = child.stderr.read() if child.stderr else b"" + except (OSError, ValueError): + stderr = b"" + error = (stderr.decode(errors="replace").strip()[-300:] + or f"ssh exited {child.returncode}") + except OSError as exc: + error = f"cannot start ssh: {exc}" + finally: + if child is not None and child.poll() is None: + child.terminate() + try: + child.wait(timeout=10) + except subprocess.TimeoutExpired: + child.kill() + if stop.is_set(): + break + write_state(config["state_path"], config, pid=os.getpid(), + attempts=attempts, last_error=error, connected=False) + print(f"tunnel: {error}; retrying in {delay:.0f}s", file=sys.stderr) + # Sleep in small slices so a stop request is honoured promptly + # instead of waiting out a full backoff interval. + waited = 0.0 + while waited < delay and not stop.is_set(): + CLOCK.sleep(min(0.2, delay - waited)) + waited += 0.2 + delay = min(delay * 2, config["backoff_cap"]) + finally: + for sig, handler in previous.items(): + signal.signal(sig, handler) + return attempts + + +class _StopFlag: + """A minimal stand-in for threading.Event, used when no stop is passed.""" + + def __init__(self): + self._set = False + + def set(self): + self._set = True + + def is_set(self): + return self._set + + +def install_commands(config, *, root="/opt/harborrl-tunnel", package_dir=None, + public_key=None): + """The two commands that install this adapter across the only working path. + + The outside host has neither this package nor a route to anything but the + hub, and its key is not authorized yet — so one paste cannot do both ends. + ``hub`` runs on the hub (authorize the key, stage nothing else); ``host`` + runs on the outside host (pull the package over the one direction that + works, write the config, start the supervised daemon). Both are + idempotent: re-running adds no duplicate key and starts no second daemon. + + ``public_key`` exists because of the same chicken-and-egg: the keypair is + generated on the outside host (it is the only side that needs the private + half), so the hub cannot read it from disk. The operator copies the public + half out of the host's portal session and passes it here. + """ + public_key = public_key or _public_key(config["identity"]) + site = json.dumps({ + "hub": config["hub"], "identity": config["identity"], + "mappings": config["mappings"], "state_path": config["state_path"], + "push": config["push"], + }, separators=(",", ":")) + source = str(Path(package_dir or Path(__file__).resolve().parent.parent)) + quoted_root = shlex.quote(root) + hub = ( + "mkdir -p ~/.ssh && chmod 700 ~/.ssh " + "&& touch ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys " + # Idempotent by content: a blind append leaves duplicates that look + # like several hosts sharing one key. + f"&& (grep -qxF {shlex.quote(public_key)} ~/.ssh/authorized_keys " + f"|| echo {shlex.quote(public_key)} >> ~/.ssh/authorized_keys) " + "&& echo hub-ready" + ) + host = ( + f"mkdir -p {quoted_root}/lib " + # The package arrives over the same direction the tunnel uses; there + # is no shared storage here to read it from. + f"&& ssh -o BatchMode=yes -i {shlex.quote(config['identity'])} " + f"{shlex.quote(config['hub'])} " + f"'tar -C {shlex.quote(source)} -cf - harborrl_h_cluster' " + f"| tar -C {quoted_root}/lib -xf - " + f"&& printf '%s\\n' {shlex.quote(site)} > {quoted_root}/tunnel.json " + f"&& if [ -f {quoted_root}/tunnel.pid ] " + f"&& kill -0 \"$(cat {quoted_root}/tunnel.pid)\" 2>/dev/null; then echo tunnel-already-running; " + f"else PYTHONPATH={quoted_root}/lib nohup python3 -m harborrl_h_cluster.tunnel daemon " + f"--config {quoted_root}/tunnel.json >> {quoted_root}/tunnel.log 2>&1 " + f"& echo $! > {quoted_root}/tunnel.pid; echo tunnel-started; fi" + ) + return {"hub": hub, "host": host} + + +def _public_key(identity): + """The public half of a key file; only this ever leaves the host.""" + result = RUNNER(["ssh-keygen", "-y", "-f", identity], + capture_output=True, text=True, timeout=30, check=False) + if result.returncode != 0 or not (result.stdout or "").strip(): + raise TunnelError( + f"cannot derive a public key from {identity}: " + f"{(result.stderr or '').strip()[-200:]}") + return result.stdout.strip() + + +def push(config, *, once=False): + """Copy declared directories back to the hub without exposing partial files. + + No shared storage exists on the outside host, so traces written there must + arrive before the host is recycled. ``rsync`` writes through a temporary + name and renames, so a reader never sees a half-written tree; without + rsync, ``tar`` streams over ssh and unpacks atomically in a temp dir. + """ + results = [] + for index, entry in enumerate(config["push"]): + local, _, remote = entry.partition(":") + if not remote: + raise TunnelError(f"push entry {index} must be a local:remote pair") + argv = ["rsync", "-a", "--timeout=120", + "-e", f"ssh -o BatchMode=yes -i {config['identity']}", + str(local), f"{config['hub']}:{remote}"] + try: + result = RUNNER(argv, capture_output=True, text=True, + timeout=1800, check=False, env=dict(os.environ)) + results.append({"entry": entry, "returncode": result.returncode, + "stderr": (result.stderr or "").strip()[-200:]}) + except (OSError, subprocess.TimeoutExpired) as exc: + results.append({"entry": entry, "returncode": None, "stderr": str(exc)}) + return results + + +def main(argv=None): + parser = argparse.ArgumentParser(prog="harborrl_h_cluster.tunnel", + description="Reverse ssh tunnel for an outside host") + sub = parser.add_subparsers(dest="action", required=True) + for action in ("daemon", "install", "push"): + cmd = sub.add_parser(action) + cmd.add_argument("--config", required=True) + if action == "push": + cmd.add_argument("--once", action="store_true") + if action == "install": + cmd.add_argument("--root", default="/opt/harborrl-tunnel") + # The keypair lives on the outside host; the hub runs install + # without it and is handed the public half from the portal. + cmd.add_argument("--public-key", default=None) + args = parser.parse_args(argv) + config = load_config(args.config) + try: + if args.action == "install": + # Both ends are printed: the key is not authorized yet, so the + # outside host cannot reach the hub until the hub command runs. + # Neither command carries private key material. + commands = install_commands(config, root=args.root, + public_key=args.public_key) + print(json.dumps(commands)) + return 0 + if args.action == "daemon": + daemon(config) + return 0 + results = push(config, once=args.once) + print(json.dumps(results)) + return 0 if all(item["returncode"] == 0 for item in results) else 1 + except TunnelError as exc: + print(f"tunnel practice: {exc}", file=sys.stderr) + return 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/practices/h_cluster/tests/test_rjob.py b/practices/h_cluster/tests/test_rjob.py new file mode 100644 index 000000000..c1c902414 --- /dev/null +++ b/practices/h_cluster/tests/test_rjob.py @@ -0,0 +1,181 @@ +"""rjob provisioning adapter: contract, polling and quota safety. + +No cluster is contacted: the submission runner, the clock and the address +record are all faked, so these tests pin the decisions — what a record may +contain, when a submit is refused, and what happens when a pod never answers. +""" +from __future__ import annotations + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from harborrl_h_cluster import rjob + + +def write_config(tmp_path, **overrides): + values = { + "name_prefix": "harborrl-test", + "image": "registry.example/ml-base:1", + "workdir": str(tmp_path / "work"), + "identity": str(tmp_path / "id_ed25519"), + "mounts": ["gpfs://gpfs1/data:/mnt/data"], + "timeout_s": 30, + # The clock is faked in these tests; the interval only has to be legal. + "poll_interval_s": 0.001, + } + values.update(overrides) + path = tmp_path / "rjob.json" + path.write_text(json.dumps(values)) + return path + + +def fake_runtime(monkeypatch, *, states, submissions=None, ssh_ok=True, clock=None): + """Replace the runner and clock; record every argv the adapter builds.""" + calls = [] + state_iter = iter(states) + + def run(argv, **kwargs): + calls.append(list(argv)) + head = argv[:2] + if head == ["rjob", "submit"]: + if submissions is not None: + submissions.append(list(argv)) + return SimpleNamespace(returncode=0, stdout="created rjob_name: job-42\n", + stderr="") + if head == ["rjob", "get"]: + state = next(state_iter, None) + if state is None: + # A deleted job answers successfully with no replica line. + return SimpleNamespace(returncode=0, stdout="", stderr="") + return SimpleNamespace( + returncode=0, stdout=f" |- replica job-42-x: {state}, node-1\n", stderr="") + if head == ["rjob", "delete"]: + return SimpleNamespace(returncode=0, stdout="", stderr="") + if head[0] == "ssh": + return SimpleNamespace(returncode=0 if ssh_ok else 255, + stdout="ok\n" if ssh_ok else "", stderr="denied") + if head[0] == "ssh-keygen": + return SimpleNamespace(returncode=0, stdout="ssh-ed25519 AAAA test\n", stderr="") + raise AssertionError(f"unexpected command {argv}") + + monkeypatch.setattr(rjob, "RUNNER", run) + monkeypatch.setattr(rjob, "_public_key", lambda identity: "ssh-ed25519 AAAA test") + # A clock that advances with every sleep: a frozen clock would make the + # deadline unreachable and hang the poll loop forever. + now = [0.0] + monkeypatch.setattr(rjob.CLOCK, "monotonic", lambda: now[0]) + + def sleep(seconds): + now[0] += max(float(seconds), 1.0) + + monkeypatch.setattr(rjob.CLOCK, "sleep", sleep) + return calls + + +def test_config_requires_the_submission_fields(tmp_path): + path = write_config(tmp_path) + path.write_text(json.dumps({"name_prefix": "x"})) + with pytest.raises(rjob.RjobError, match="requires"): + rjob.load_config(path) + + +def test_config_rejects_a_relative_workdir(tmp_path): + with pytest.raises(rjob.RjobError, match="absolute path"): + rjob.load_config(write_config(tmp_path, workdir="relative/work")) + + +def test_config_rejects_unknown_fields(tmp_path): + path = write_config(tmp_path) + raw = json.loads(path.read_text()) + raw["gpuu"] = 8 + path.write_text(json.dumps(raw)) + with pytest.raises(rjob.RjobError, match="unsupported fields"): + rjob.load_config(path) + + +def test_submit_argv_carries_the_quota_and_namespace_flags(tmp_path): + config = rjob.load_config(write_config(tmp_path)) + argv = rjob.submit_argv(config, "/tmp/work/bootstrap.sh") + assert argv[:2] == ["rjob", "submit"] + assert "--charged-group=datafoundation_gpu_pool" in argv + assert "--namespace=ailab-datafoundation" in argv + assert "--enable-sshd" in argv + assert "--mount=gpfs://gpfs1/data:/mnt/data" in argv + # The bootstrap script is the pod command: rjob has no exec action. + assert argv[-3:] == ["--", "bash", "/tmp/work/bootstrap.sh"] + + +def test_bootstrap_script_is_idempotent_and_publishes_the_address(tmp_path): + config = rjob.load_config(write_config(tmp_path)) + script = rjob.bootstrap_script(config, "ssh-ed25519 AAAA test") + # Appending twice would leave a duplicate key, so the guard is by content. + assert "grep -qxF" in script and "authorized_keys" in script + assert "pgrep -x sshd" in script + assert "pod-net.json" in script + assert "ssh-ed25519 AAAA test" in script + + +def test_up_returns_the_record_after_state_address_and_ssh(tmp_path, monkeypatch): + config = rjob.load_config(write_config(tmp_path)) + Path(config["workdir"]).mkdir(parents=True) + fake_runtime(monkeypatch, states=["STARTING", "ACTIVE"]) + writes = iter([None, "100.98.93.42"]) + monkeypatch.setattr(rjob, "pod_address", lambda cfg: next(writes, None)) + record = rjob.up(config) + assert record == {"host": "100.98.93.42", "user": "root", "port": 22, + "identity": config["identity"], "handle": "job-42"} + + +def test_up_refuses_a_second_submission_before_spending_quota(tmp_path, monkeypatch): + config = rjob.load_config(write_config(tmp_path)) + workdir = Path(config["workdir"]) + workdir.mkdir(parents=True) + (workdir / "rjob-handle.json").write_text('{"handle": "old-job"}\n') + calls = fake_runtime(monkeypatch, states=[]) + with pytest.raises(rjob.RjobError, match="previous pod record"): + rjob.up(config) + assert not [call for call in calls if call[:2] == ["rjob", "submit"]] + + +def test_up_fails_closed_when_the_pod_never_publishes_an_address(tmp_path, monkeypatch): + config = rjob.load_config(write_config(tmp_path, timeout_s=1)) + Path(config["workdir"]).mkdir(parents=True) + fake_runtime(monkeypatch, states=["ACTIVE"] * 50) + monkeypatch.setattr(rjob, "pod_address", lambda cfg: None) + with pytest.raises(rjob.RjobError, match="publish its address"): + rjob.up(config) + + +def test_up_fails_closed_when_ssh_is_refused(tmp_path, monkeypatch): + config = rjob.load_config(write_config(tmp_path, timeout_s=1)) + Path(config["workdir"]).mkdir(parents=True) + fake_runtime(monkeypatch, states=["ACTIVE"] * 50, ssh_ok=False) + monkeypatch.setattr(rjob, "pod_address", + lambda cfg: "100.98.93.42") + with pytest.raises(rjob.RjobError, match="non-interactive ssh"): + rjob.up(config) + + +def test_down_deletes_only_the_recorded_handle(tmp_path, monkeypatch): + config = rjob.load_config(write_config(tmp_path)) + Path(config["workdir"]).mkdir(parents=True) + calls = fake_runtime(monkeypatch, states=[None]) + assert rjob.down(config, {"handle": "job-42"}) == "job-42" + deletes = [call for call in calls if call[:2] == ["rjob", "delete"]] + assert deletes and deletes[0][2] == "job-42" + + +def test_down_refuses_a_record_without_a_handle(tmp_path, monkeypatch): + config = rjob.load_config(write_config(tmp_path)) + fake_runtime(monkeypatch, states=[]) + with pytest.raises(rjob.RjobError, match="no handle"): + rjob.down(config, {"host": "1.2.3.4"}) + + +def test_cli_prints_one_machine_readable_line(tmp_path, monkeypatch, capsys): + monkeypatch.setattr(rjob, "up", lambda cfg: {"host": "h", "handle": "job-1"}) + assert rjob.main(["up", "--config", str(write_config(tmp_path))]) == 0 + assert json.loads(capsys.readouterr().out) == {"host": "h", "handle": "job-1"} diff --git a/practices/h_cluster/tests/test_tunnel.py b/practices/h_cluster/tests/test_tunnel.py new file mode 100644 index 000000000..72bb8fff9 --- /dev/null +++ b/practices/h_cluster/tests/test_tunnel.py @@ -0,0 +1,220 @@ +"""Reverse tunnel adapter: mapping, supervision and pasteable installation. + +Nothing here opens a socket: ssh runs through a fake subprocess layer, so +these tests pin the decisions — the exact -R strings, what a reconnect does +after a failure, and what the operator's one-line install contains. +""" +from __future__ import annotations + +import json +import threading + +import pytest + +from harborrl_h_cluster import tunnel + + +def write_config(tmp_path, **overrides): + values = { + "hub": "manager.example", + "identity": str(tmp_path / "id_ed25519"), + "mappings": [ + {"local_port": 5010, "hub_bind": "127.0.0.1", "hub_port": 15010}, + {"local_port": 5011, "hub_bind": "127.0.0.1", "hub_port": 15011}, + ], + "state_path": str(tmp_path / "state.json"), + "push": [f"{tmp_path}/traces:~/incoming/traces"], + } + values.update(overrides) + path = tmp_path / "tunnel.json" + path.write_text(json.dumps(values)) + return path + + +def test_config_defaults_and_absolute_path_preservation(tmp_path): + config = tunnel.load_config(write_config(tmp_path)) + assert len(config["mappings"]) == 2 + assert config["backoff_start"] == tunnel.BACKOFF_START + assert config["backoff_cap"] == tunnel.BACKOFF_CAP + + +@pytest.mark.parametrize("mutation, message", [ + ({"hub": ""}, "ssh destination"), + ({"hub": "-oProxyCommand=x"}, "ssh destination"), + ({"identity": ""}, "key path"), + ({"mappings": []}, "nonempty list"), + ({"mappings": [{"local_port": 0, "hub_bind": "127.0.0.1", "hub_port": 1}]}, + "local_port must be an integer"), + ({"mappings": [{"local_port": 5010, "hub_bind": "", "hub_port": 1}]}, "hub_bind"), + ({"mappings": [{"local_port": 5010, "hub_bind": "127.0.0.1"}]}, "requires exactly"), + ({"mappings": [ + {"local_port": 1, "hub_bind": "127.0.0.1", "hub_port": 9}, + {"local_port": 2, "hub_bind": "127.0.0.1", "hub_port": 9}]}, "duplicate hub port"), + ({"state_path": ""}, "state_path"), + ({"backoff_start": 0}, "backoff"), + ({"backoff_cap": 0.1}, "backoff"), + ({"push": [""]}, "local:remote"), +]) +def test_config_validation_fails_closed(tmp_path, mutation, message): + with pytest.raises(tunnel.TunnelError, match=message): + tunnel.load_config(write_config(tmp_path, **mutation)) + + +def test_forward_argv_binds_loopback_and_orders_one_forward_per_mapping(tmp_path): + config = tunnel.load_config(write_config(tmp_path)) + argv = tunnel.forward_argv(config) + assert argv[0] == "ssh" + forwards = [argv[index + 1] for index, part in enumerate(argv) if part == "-R"] + assert forwards == ["127.0.0.1:15010:127.0.0.1:5010", + "127.0.0.1:15011:127.0.0.1:5011"] + # The destination is last and the local side is never a routable address. + assert argv[-1] == "manager.example" + assert all(forward.split(":")[2] == "127.0.0.1" for forward in forwards) + assert "-N" in argv and "ExitOnForwardFailure=yes" in argv + + +def test_daemon_reconnects_with_backoff_and_records_the_error(tmp_path, monkeypatch): + config = tunnel.load_config(write_config(tmp_path, backoff_start=0.001, + backoff_cap=0.004)) + sleeps = [] + monkeypatch.setattr(tunnel.CLOCK, "sleep", lambda seconds: sleeps.append(seconds)) + errors, attempts = [], [] + monkeypatch.setattr(tunnel, "POPEN", lambda argv, **kw: (_ for _ in ()).throw( + OSError("host key verification failed"))) + real_write = tunnel.write_state + + def capture(path, cfg, **kwargs): + # The supervisor writes twice per attempt: once when connecting (no + # error yet) and once when the connection ends, carrying the reason. + if kwargs.get("last_error") is not None: + errors.append(kwargs["last_error"]) + attempts.append(kwargs.get("attempts")) + return real_write(path, cfg, **kwargs) + + monkeypatch.setattr(tunnel, "write_state", capture) + stopped = threading.Event() + + def sleep_then_stop(seconds): + sleeps.append(seconds) + if len(sleeps) >= 3: + stopped.set() + + monkeypatch.setattr(tunnel.CLOCK, "sleep", sleep_then_stop) + attempts_made = tunnel.daemon(config, stop=stopped) + assert attempts_made == 3 + # Every recorded failure names the cause, not just "ssh died". + assert all("host key verification failed" in item for item in errors) + state = json.loads((tmp_path / "state.json").read_text()) + assert state["connected"] is False + assert state["mappings"] == config["mappings"] + + +def test_daemon_retries_a_nonzero_ssh_exit_with_capped_backoff(tmp_path, monkeypatch): + config = tunnel.load_config(write_config(tmp_path, backoff_start=0.001, + backoff_cap=0.002)) + delays, rounds = [], {"n": 0} + + class Child: + returncode = None + stderr = None + + def poll(self): + return 7 + + def wait(self, timeout=None): + return 7 + + def terminate(self): + pass + + def kill(self): + pass + + class Buf: + def read(self): + return b"remote port forwarding failed\n" + + Child.stderr = Buf() + monkeypatch.setattr(tunnel, "POPEN", lambda argv, **kw: Child()) + stopped = threading.Event() + + def fake_sleep(seconds): + delays.append(seconds) + rounds["n"] += 1 + if rounds["n"] >= 4: + stopped.set() + + monkeypatch.setattr(tunnel.CLOCK, "sleep", fake_sleep) + tunnel.daemon(config, stop=stopped) + # Backoff doubles from the start value and stops growing at the cap. + assert delays[0] == 0.001 + assert max(delays) <= 0.002 + + +def test_install_prints_both_ends_and_never_the_private_key(tmp_path): + config = tunnel.load_config(write_config(tmp_path)) + local = tmp_path / "id_ed25519" + local.write_text("PRIVATE-KEY-MATERIAL") + (tmp_path / "id_ed25519.pub").write_text("ssh-ed25519 AAAA test\n") + + class Run: + returncode = 0 + stdout = "ssh-ed25519 AAAA test\n" + stderr = "" + + monkey = tunnel.RUNNER + tunnel.RUNNER = lambda *args, **kwargs: Run() + try: + commands = tunnel.install_commands( + config, root="/opt/harborrl-tunnel", + public_key="ssh-ed25519 AAAA test") + finally: + tunnel.RUNNER = monkey + + # The key is not authorized yet, so one paste cannot do both ends: the + # hub command authorizes it, the host command can only follow. + assert set(commands) == {"hub", "host"} + hub, host = commands["hub"], commands["host"] + assert "authorized_keys" in hub and "grep -qxF" in hub # idempotent append + assert "tunnel.json" in host and "harborrl_h_cluster.tunnel daemon" in host + # The package travels over the same one-way path: the host has no route + # to anything but the hub, and no shared storage to read it from. + assert "tar -C" in host and "harborrl_h_cluster" in host + # The mappings travel inside the site config the host command writes. + assert json.dumps(config["mappings"], separators=(",", ":")) in host + # Neither command may carry the private key itself. + for command in commands.values(): + assert "PRIVATE-KEY-MATERIAL" not in command + + +def test_state_writes_are_atomic_and_carry_the_reason(tmp_path): + config = tunnel.load_config(write_config(tmp_path)) + tunnel.write_state(config["state_path"], config, pid=1234, attempts=2, + last_error="connection reset", connected=False) + payload = json.loads((tmp_path / "state.json").read_text()) + assert payload["pid"] == 1234 and payload["attempts"] == 2 + assert payload["last_error"] == "connection reset" + assert payload["connected"] is False + assert not (tmp_path / "state.tmp").exists() + + +def test_push_uses_rsync_with_a_bounded_timeout(tmp_path, monkeypatch): + config = tunnel.load_config(write_config(tmp_path)) + calls = [] + + def run(argv, **kwargs): + calls.append((argv, kwargs)) + return type("R", (), {"returncode": 0, "stderr": ""})() + + monkeypatch.setattr(tunnel, "RUNNER", run) + results = tunnel.push(config, once=True) + assert results[0]["returncode"] == 0 + argv, kwargs = calls[0] + assert argv[0] == "rsync" and "-a" in argv + assert kwargs["timeout"] == 1800 + + +def test_push_reports_a_malformed_entry_instead_of_guessing(tmp_path): + config = tunnel.load_config(write_config(tmp_path, push=["/only/local"])) + with pytest.raises(tunnel.TunnelError, match="local:remote"): + tunnel.push(config, once=True) diff --git a/pyproject.toml b/pyproject.toml index 1cbdca1ad..3c9068768 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,15 @@ dependencies = [ harborrl = "harborrl.cli:main" [project.optional-dependencies] +distill = [ + "torch>=2.5,<3", + "transformers>=4.49,<5", +] + +prm = [ + "torch>=2.4", + "transformers==4.46.3", +] sandbox = [ "harbor==0.23.0; python_version >= '3.12'", "opensandbox==0.1.9; python_version >= '3.12'", diff --git a/scripts/audit_public_tree.sh b/scripts/audit_public_tree.sh index 7d6d5dfd5..917f58819 100755 --- a/scripts/audit_public_tree.sh +++ b/scripts/audit_public_tree.sh @@ -19,7 +19,7 @@ public_documents = { "configs/harbor_hub/public_compatibility_summary.md", "configs/harbor_hub/public_compatibility_summary_zh.md", "docs/agents/environment-check.md", "docs/configuration-wizard.md", - "docs/deployment.md", "docs/rollout-runtime-lifecycle.md", "docs/sandbox-rollout.md", + "docs/deployment.md", "docs/math-distill.md", "docs/sandbox-rollout.md", "docs/train-inference-consistency.md", "docs/web-console.md", "practices/h_cluster/README.md", "harborrl/webui/static/lucide-LICENSE.txt", } @@ -98,7 +98,7 @@ for path in files: for label, pattern in internal_patterns.items(): for match in pattern.finditer(data): value = match.group(0) - if b"127.0.0.1" in value or b"localhost" in value or value == b"host.docker.internal": + if b"127.0.0.1" in value or b"localhost" in value: continue raise SystemExit(f"{path} contains {label}: {value.decode(errors='replace')}") diff --git a/scripts/sandbox_sdk_contract.py b/scripts/sandbox_sdk_contract.py index 9957d6ee7..ce368b578 100644 --- a/scripts/sandbox_sdk_contract.py +++ b/scripts/sandbox_sdk_contract.py @@ -6,6 +6,7 @@ from pathlib import Path from types import SimpleNamespace +from harbor.models.task.config import EnvironmentConfig from harbor.models.trial.config import TrialConfig from opensandbox import Sandbox @@ -20,6 +21,25 @@ assert 'import_path' in TrialConfig.model_fields['environment'].annotation.model_fields assert not inspect.isabstract(OfflineClaudeCode) +with tempfile.TemporaryDirectory() as directory: + task_dir = Path(directory) + (task_dir / 'Dockerfile').write_text('FROM example/original\n') + config = EnvironmentConfig() + env = object.__new__(ManagedOpenSandboxEnvironment) + env.environment_dir = task_dir + env.task_env_config = config + env.binding = {'image': 'example/built@sha256:' + 'a' * 64} + env._validate_definition() + assert config.docker_image is None + assert env.task_env_config.docker_image == env.binding['image'] + (task_dir / 'docker-compose.yaml').write_text('services: {}\n') + try: + env._validate_definition() + except ValueError as exc: + assert 'docker-compose' in str(exc) + else: + raise AssertionError('compose tasks must still be rejected') + async def cleanup_contract(): with tempfile.TemporaryDirectory() as directory: diff --git a/tests/harborrl/test_dapo_training.py b/tests/harborrl/test_dapo_training.py new file mode 100644 index 000000000..f8003f1e3 --- /dev/null +++ b/tests/harborrl/test_dapo_training.py @@ -0,0 +1,197 @@ +"""DAPO contracts for complete-group sampling and global token gradients.""" +import json +import math +import os +import sys +from copy import deepcopy +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import yaml +from test_math_rollout import samples + +from harborrl.backends.slime_v032 import math_rollout +from harborrl.backends.slime_v032.postprocess import rollout_data_postprocess +from harborrl.config import ROOT, launch_plan, load_config +from harborrl.export.contract import DAPO_REDUCTION +from harborrl.export.validate import validate_batch +from harborrl.math import dapo +from harborrl.math.data import verify +from harborrl.platform.evidence import effective_config + +SETTINGS = {'clip_low': .2, 'clip_high': .28, 'overlong_buffer_tokens': 2, + 'overlong_penalty': 1., 'max_refill_batches': 3} + + +def dapo_samples(monkeypatch, lengths=(1, 3)): + monkeypatch.setenv('HARBORRL_MATH_ALGORITHM', 'dapo') + monkeypatch.setenv('HARBORRL_DAPO', json.dumps(SETTINGS)) + rows = samples() + for sample, length in zip(rows, lengths, strict=True): + sample.tokens = [1, 2] + [3] * length + sample.response_length = length + sample.loss_mask = [1] * length + sample.rollout_log_probs = [-.5] * length + sample.rollout_id = 0 + evidence = sample.metadata['math_evidence'] + evidence.update(algorithm='dapo', tokens=sample.tokens, response_length=length, + logprobs=sample.rollout_log_probs, + dapo=dapo.reward(evidence['verification']['reward'], length, 4, SETTINGS)) + return rows + + +def test_dapo_shaping_advantages_and_backend_contract(monkeypatch): + rows = dapo_samples(monkeypatch) + args = SimpleNamespace(n_samples_per_prompt=2, rollout_max_response_len=4) + batch = math_rollout.build_batch(args, rows) + assert batch.reduction == DAPO_REDUCTION + group = batch.groups[0] + assert [unit.reward for unit in group.trajectories] == [1., -1.5] + assert group.reward_std == 1.25 + expected = 1.25 / (math.sqrt(3.125) + 1e-6) + assert [unit.advantage for unit in group.trajectories] == pytest.approx([expected, -expected]) + data = math_rollout.convert_samples_to_train_data(args, rows) + assert data['raw_reward'] == [1., -1.] + assert data['rollout_mask_sums'] == [1, 3] + rollout_data_postprocess(args, 0, data) + data['loss_masks'][1][-1] = 0 + with pytest.raises(ValueError, match='all response tokens'): + rollout_data_postprocess(args, 0, data) + wrong = replace(group.trajectories[0], advantage=1.) + with pytest.raises(ValueError, match='sample std'): + validate_batch(replace(batch, groups=(replace(group, trajectories=(wrong, group.trajectories[1])),))) + rows[0].metadata['math_evidence']['dapo']['reward'] = 2. + with pytest.raises(ValueError, match='DAPO reward'): + math_rollout.build_batch(args, rows) + + +@pytest.mark.parametrize('length,expected', [(1, 0.), (2, 0.), (3, -.5), (4, -1.)]) +def test_soft_overlong_boundary(length, expected): + assert dapo.reward(1., length, 4, SETTINGS) == {'reward': 1. + expected, 'overlong_penalty': expected} + + +@pytest.mark.parametrize('outcome', [-1., 0., 1.]) +def test_dapo_accepts_every_verifier_outcome(outcome): + assert dapo.reward(outcome, 1, 4, SETTINGS) == {'reward': outcome, 'overlong_penalty': 0.} + with pytest.raises(ValueError, match='verifier outcome'): + dapo.reward(.5, 1, 4, SETTINGS) + + +def test_dapo_config_and_effective_semantics(tmp_path): + config = yaml.safe_load((ROOT / 'examples/math/train.yaml').read_text(encoding='utf-8')) + config['tasks']['path'] = str(ROOT / 'examples/math/tasks.jsonl') + config['algorithm'] = {'name': 'dapo'} + path = tmp_path / 'train.yaml' + path.write_text(yaml.safe_dump(config), encoding='utf-8') + with pytest.raises(ValueError, match='dapo section'): + load_config(path) + config['dapo'] = SETTINGS + path.write_text(yaml.safe_dump(config), encoding='utf-8') + plan = launch_plan(load_config(path)) + assert json.loads(plan['environment']['HARBORRL_DAPO']) == SETTINGS + semantics = effective_config(plan, plan['environment'], {}, {}, [])['update_semantics'] + assert semantics['std_normalization'] is True + assert semantics['reduction'] == DAPO_REDUCTION + for override in ('dapo.clip_low=1', 'dapo.clip_high=0.1', 'dapo.max_refill_batches=0', + 'dapo.overlong_penalty=2', 'dapo.overlong_buffer_tokens=0', 'algorithm.name=grpo'): + with pytest.raises(ValueError): + load_config(path, [override]) + + +@pytest.mark.parametrize('mode', ['success', 'exhausted', 'version', 'duplicate']) +@pytest.mark.skipif(os.name == 'nt', reason='production immutable evidence requires POSIX directory fsync') +def test_dynamic_sampling_complete_groups_and_failure_evidence(tmp_path, monkeypatch, mode): + prototype = dapo_samples(monkeypatch, (1, 1)) + monkeypatch.setenv('RUN_DIR', str(tmp_path)) + monkeypatch.setitem(sys.modules, 'slime.rollout.base_types', SimpleNamespace(RolloutFnTrainOutput=SimpleNamespace)) + args = SimpleNamespace(n_samples_per_prompt=2, rollout_batch_size=2, rollout_max_response_len=4, + rollout_temperature=1., rollout_top_p=1., rollout_top_k=-1) + + class Source: + offset = 0 + + def get_samples(self, count): + groups = [] + for index in range(self.offset, self.offset + count): + rows = deepcopy(prototype) + for sample in rows: + sample.group_index = index + groups.append(rows) + self.offset += count + return groups + + async def generate(args, rollout_id, groups): + result = [] + for group in groups: + for index, sample in enumerate(group): + group_id = sample.group_index + if group_id == 0 or mode == 'exhausted': + sample.response = '999999' + task = sample.metadata['math_task'] + outcome = verify(sample.response, task['reference']) + key = f'group-{group_id if mode != "duplicate" else 0}' + evidence = sample.metadata['math_evidence'] + evidence.update(response=sample.response, verification=outcome, group_key=key, + rollout_id=group_id, trajectory_id=f'{key}:{index}', + policy_version='2' if mode == 'version' and group_id > 1 else '1', + dapo=dapo.reward(outcome['reward'], 1, 4, SETTINGS)) + sample.rollout_id = group_id + sample.index = group_id * 2 + index + result.append(sample) + return result + + monkeypatch.setattr(math_rollout, '_generate', generate) + source = Source() + if mode == 'success': + output = math_rollout.generate_rollout(args, 0, source) + assert source.offset == 3 + assert len(output.samples) == 4 + assert {sample.group_index for sample in output.samples} == {1, 2} + assert output.metrics['dapo_rejected_groups'] == 1. + assert output.metrics['dapo_generation_batches'] == 2. + selection = json.loads((tmp_path / 'math/000000/sampling.json').read_text()) + assert selection['status'] == 'complete' + assert len(selection['accepted']) == 2 and len(selection['rejected']) == 1 + else: + with pytest.raises(ValueError, match={'exhausted': 'refill exhausted', 'version': 'policy changed', + 'duplicate': 'duplicate'}[mode]): + math_rollout.generate_rollout(args, 0, source) + if mode == 'exhausted': + selection = json.loads((tmp_path / 'math/000000/sampling.json').read_text()) + assert selection['status'] == 'exhausted' + assert len(selection['rejected']) == 6 + assert (tmp_path / 'math/000000/sampling-evidence-000.json').is_file() + + +def test_dapo_loss_partition_invariance_and_asymmetric_gradient(monkeypatch): + torch = pytest.importorskip('torch') + from harborrl.backends.slime_v032.dapo import loss_function + + ratios = torch.tensor([.7, .85, 1.22, 1.3, .7, .85, 1.22, 1.3], dtype=torch.float64) + current = (ratios.log() - 1).requires_grad_() + old = torch.full((8,), -1., dtype=torch.float64, requires_grad=True) + advantages = torch.tensor([1.] * 4 + [-1.] * 4, dtype=torch.float64, requires_grad=True) + args = SimpleNamespace(calculate_per_token_loss=True, context_parallel_size=1, + eps_clip=.2, eps_clip_high=.28) + + def loss(indices): + selected = current[indices] + batch = {'unconcat_tokens': [], 'total_lengths': [len(indices) + 1], + 'response_lengths': [len(indices)], 'rollout_log_probs': [old[indices]], + 'advantages': [advantages[indices]], 'loss_masks': [torch.ones(len(indices))]} + monkeypatch.setitem(sys.modules, 'slime.backends.megatron_utils.loss', SimpleNamespace( + get_log_probs_and_entropy=lambda *a, **k: (None, {'log_probs': [selected]}))) + return loss_function(args, batch, selected, None)[0] + + full = loss(list(range(8))) / 8 + # Uneven DP shards and microbatches must sum before the global token division. + partitioned = (loss([0, 3, 5]) + loss([2]) + loss([1, 4, 6, 7])) / 8 + assert partitioned.item() == pytest.approx(full.item()) + expected = torch.maximum(-advantages.detach() * ratios, + -advantages.detach() * ratios.clamp(.8, 1.28)).mean() + assert full.item() == pytest.approx(expected.item()) + partitioned.backward() + assert current.grad.tolist() == pytest.approx([-.7 / 8, -.85 / 8, -1.22 / 8, 0., + 0., .85 / 8, 1.22 / 8, 1.3 / 8]) + assert old.grad is None and advantages.grad is None diff --git a/tests/harborrl/test_deployment_targets.py b/tests/harborrl/test_deployment_targets.py index ef609d9ad..08db522b9 100644 --- a/tests/harborrl/test_deployment_targets.py +++ b/tests/harborrl/test_deployment_targets.py @@ -19,6 +19,19 @@ def completed(returncode=0, stdout="", stderr=""): return subprocess.CompletedProcess([], returncode, stdout, stderr) +def test_local_target_preserves_cuda_runtime_environment(monkeypatch): + runtime = { + "CUDA_HOME": "/opt/cuda", + "LD_LIBRARY_PATH": "/opt/cuda/lib64", + "LD_PRELOAD": "/opt/cuda/lib64/libnvrtc.so.13", + "NVTE_CUDA_INCLUDE_DIR": "/opt/cuda/include", + } + for key, value in runtime.items(): + monkeypatch.setenv(key, value) + target = training_target.local_target({"training_command": ["true"]}, environment={}) + assert {key: target.env[key] for key in runtime} == runtime + + def test_local_target_runs_and_forwards_exit_code(tmp_path): target = LocalTrainingTarget(["bash", "-c", "exit 3"], cwd=str(tmp_path), env={"PATH": "/bin"}) target.launch() @@ -142,7 +155,8 @@ def test_ssh_command_assembles_noninteractive_ssh(): assert argv[-2:] == ["root@gpu", "true"] -def test_service_dispatch_launches_gpu_job_through_controller(tmp_path, monkeypatch, capsys): +@pytest.mark.parametrize("shutdown_error", [False, True]) +def test_service_dispatch_launches_gpu_job_through_controller(tmp_path, monkeypatch, capsys, shutdown_error): """P1+P3 manager flow: materialize, register the run, launch, wait, report.""" import json as json_module from pathlib import Path @@ -185,10 +199,10 @@ def serve_forever(self): pass def shutdown(self): - pass + launched["control_shutdown"] = True def server_close(self): - pass + launched["control_closed"] = True class FakeTarget: def __init__(self, spec, *, source_dir, run_dir, **kwargs): @@ -214,8 +228,19 @@ def cancel(self): monkeypatch.setattr(controller_module, "SshTrainingTarget", FakeTarget) monkeypatch.setattr(controller_module, 'check_gpu_control', lambda _: {'gpu_to_controller': True}) + if shutdown_error: + original_close = controller_module.RolloutService.close + def failed_close(service): + original_close(service) + raise RuntimeError("cleanup failure") + monkeypatch.setattr(controller_module.RolloutService, "close", failed_close) + with pytest.raises(RuntimeError, match="cleanup failure"): + controller_module.dispatch_service(plan, deployment, checks=[]) + assert launched["control_shutdown"] and launched["control_closed"] + return exit_code = controller_module.dispatch_service(plan, deployment, checks=[]) assert exit_code == 0 + assert launched["control_shutdown"] and launched["control_closed"] run_dir = Path(launched["run_dir"]) assert run_dir.is_dir() environment = launched["environment"] diff --git a/tests/harborrl/test_distill.py b/tests/harborrl/test_distill.py new file mode 100644 index 000000000..280fbc83c --- /dev/null +++ b/tests/harborrl/test_distill.py @@ -0,0 +1,217 @@ +"""Math data boundaries, sampled OPD mathematics and reference-trainer contracts.""" +import json +import subprocess +import sys +from copy import deepcopy +from pathlib import Path + +import pytest + +from harborrl.config import ROOT +from harborrl.distill.config import validate +from harborrl.distill.data import import_samples, load_tasks, verify +from harborrl.trajectories.native import digest + + +def config(tmp_path, **overrides): + raw = {"tasks": {"path": str(ROOT / "examples/distill/math.jsonl")}, + "output": {"directory": str(tmp_path / "run")}, "training": {"steps": 1, "batch_size": 2, "microbatch_size": 1}, + "sampling": {"max_new_tokens": 3, "max_context": 128}, **overrides} + return validate(raw, tmp_path / "config.yaml", allow_synthetic=True) + + +@pytest.mark.parametrize("response,reference,status", [("2", "2", "correct"), ("#### 1/2", "0.5", "correct"), + (r"reasoning \boxed{3}", "3", "correct"), ("2", "3", "incorrect"), ("2 or 3", "3", "unparseable"), + ("__import__('os').system('echo unsafe')", "0", "unparseable"), ("1/0", "1", "unparseable"), ("x" * 20000, "0", "unsupported")]) +def test_rational_verifier(response, reference, status): + assert verify(response, reference)["status"] == status + + +def test_split_lock_and_curated_import(tmp_path): + source = ROOT / "examples/distill/math.jsonl" + tasks = load_tasks(source) + duplicate = {**tasks[0], "task_id": "different", "split": "test"} + leaked = tmp_path / "leaked.jsonl" + leaked.write_text("\n".join(json.dumps(row) for row in [tasks[0], duplicate])) + with pytest.raises(ValueError, match="leakage"): + load_tasks(leaked) + samples = tmp_path / "curated.jsonl" + samples.write_text(json.dumps({"task_id": tasks[0]["task_id"], "response": "2", "source": "human-reviewed", "license": "MIT"})) + output = tmp_path / "dataset" + manifest = import_samples(source, samples, output) + sample = json.loads(next((output / "samples").glob("*.json")).read_text()) + assert sample["evidence_kind"] == "declared" and sample["token_evidence"] is None + assert manifest["training_eligibility"] == "not_an_on_policy_batch" + with pytest.raises(ValueError, match="new directory"): + import_samples(source, samples, output) + + +@pytest.mark.parametrize("mutation", ["fake", "unknown", "temperature", "method", "group", "split"]) +def test_config_rejections(tmp_path, mutation): + value = config(tmp_path) + if mutation == "fake": + with pytest.raises(ValueError, match="allow-synthetic"): + validate(value, tmp_path / "config") + return + if mutation == "unknown": + value["training"]["epochs"] = 2 + elif mutation == "temperature": + value["sampling"]["temperature"] = 0.7 + elif mutation == "method": + value["teacher"]["kind"] = "none" + elif mutation == "group": + value["training"]["group_size"] = 1 + else: + value["tasks"]["eval_split"] = "train" + with pytest.raises(ValueError): + validate(value, tmp_path / "config", allow_synthetic=True) + + +def test_offline_help_and_plan_without_torch(tmp_path): + script = ''' +import builtins, json, sys +original = builtins.__import__ +def guarded(name, *args, **kwargs): + assert name.split('.')[0] not in {'torch','transformers','harbor','sglang','opensandbox'} + return original(name,*args,**kwargs) +builtins.__import__ = guarded +from harborrl.cli import main +assert main(['distill','train',sys.argv[1],'--allow-synthetic','--dry-run']) == 0 +''' + path = tmp_path / "config.json" + path.write_text(json.dumps(config(tmp_path))) + result = subprocess.run([sys.executable, "-c", script, str(path)], cwd=ROOT, capture_output=True, text=True, check=False) + assert result.returncode == 0, result.stderr + assert json.loads(result.stdout)["backend_contract"] == "transformers-math-onpolicy-v1" + assert not (tmp_path / "run").exists() + + +def test_opd_expectation_gradient_and_detach(): + torch = pytest.importorskip("torch") + from harborrl.algorithms.opd import sampled_opd_loss + logits = torch.tensor([-.5, .2, .8], dtype=torch.float64, requires_grad=True) + teacher = torch.tensor([-.2, .4, -.1], dtype=torch.float64, requires_grad=True).log_softmax(0) + logq = logits.log_softmax(0) + exact = (logq.exp() * (logq - teacher.detach())).sum() + expected = torch.autograd.grad(exact, logits, retain_graph=True)[0] + surrogate = sum(logq[i].detach().exp() * sampled_opd_loss(logq[i].reshape(1, 1), + logq[i].reshape(1, 1), teacher[i].reshape(1, 1), torch.ones(1, 1)) for i in range(3)) + observed = torch.autograd.grad(surrogate, logits, retain_graph=True)[0] + torch.testing.assert_close(expected, observed) + assert observed.norm() > 0 + assert torch.autograd.grad(surrogate, teacher, allow_unused=True)[0] is None + + +def test_loss_mask_and_microbatch_reduction(): + torch = pytest.importorskip("torch") + from harborrl.algorithms.opd import trajectory_surrogate + current = torch.tensor([[-1., -2., 0.], [-3., 0., 0.], [-2., -3., -4.]], requires_grad=True) + advantages = torch.tensor([[1., 2., 9.], [3., 8., 8.], [1., -1., 2.]], requires_grad=True) + mask = torch.tensor([[1, 1, 0], [1, 0, 0], [1, 1, 1]]) + full = trajectory_surrogate(current, advantages, mask) + split = trajectory_surrogate(current[:2], advantages[:2], mask[:2]) * 2 / 3 + trajectory_surrogate(current[2:], advantages[2:], mask[2:]) / 3 + torch.testing.assert_close(full, split) + full.backward() + assert current.grad[0, 2] == 0 and advantages.grad is None + + +def test_prm_step_alignment_and_equation(): + pytest.importorskip("transformers") + from harborrl.distill.models import ByteTokenizer + from harborrl.distill.scoring import process_rewards + tokenizer = ByteTokenizer() + tokens = tokenizer.encode("1 + 1 = 2\n2 * 3 = 7")[1:] + scored = process_rewards("equation", tokens, tokenizer) + assert [step["status"] for step in scored["steps"]] == ["correct", "incorrect"] + assert len(scored["token_returns"]) == len(tokens) + assert scored["token_returns"][0] == 0 and scored["token_returns"][-1] == -1 + fake = process_rewards("fake", tokens, tokenizer) + assert fake["evidence_kind"] == "synthetic" + + +def test_model_shift_padding_and_teacher_probe(tmp_path): + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + from harborrl.distill.models import load_student, response_logprobs + from harborrl.distill.scoring import ModelTeacher, probe_teacher + value = config(tmp_path, teacher={"kind": "tiny_qwen"}) + model, tokenizer = load_student(value, torch.device("cpu")) + prompt, response = [1, 40, 50], [52, 2] + scored = response_logprobs(model, [prompt, [1]], [response, [2]], tokenizer.pad_token_id) + logits = model(torch.tensor([prompt + response])).logits[0].log_softmax(-1) + torch.testing.assert_close(scored[0][0], logits[2, 52]) + torch.testing.assert_close(scored[0][1], logits[3, 2]) + teacher = ModelTeacher(value, torch.device("cpu"), tokenizer.identity()) + assert probe_teacher(teacher, tokenizer)["probe"] == "passed" + assert teacher.unchanged() + with pytest.raises(ValueError, match="mismatch"): + teacher.score_tokens(prompt, response, {**tokenizer.identity(), "template_digest": "bad"}) + + +def test_training_resume_and_admission(tmp_path): + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + from harborrl.distill.contract import validate_row + from harborrl.distill.runner import train + value = config(tmp_path) + root = Path(value["output"]["directory"]) + root.mkdir() + train(value, allow_synthetic=True) + summary = json.loads((root / "summary.json").read_text()) + assert summary["parameter_changed"] and summary["nonzero_gradient_steps"] == 1 + row = json.loads(next((root / "batches").glob("*.json")).read_text()) + kwargs = {"version": row["student_behavior_version"], "identity": row["identity"], "context": 128} + with pytest.raises(ValueError, match="Synthetic"): + validate_row(row, **kwargs) + with pytest.raises(ValueError, match="Stale"): + validate_row(row, **{**kwargs, "version": "old"}, allow_synthetic=True) + tampered = deepcopy(row) + tampered["behavior_logprobs"][0] = float("nan") + with pytest.raises(ValueError): + validate_row(tampered, **kwargs, allow_synthetic=True) + for key, replacement in (("teacher_logprobs", [0.5] * len(row["response_ids"])), + ("loss_mask", [0] * len(row["response_ids"])), + ("verification", {"status": "correct", "reward": -1.0})): + altered = deepcopy(row) + altered[key] = replacement + altered["score_response_digest"] = digest({"request_digest": altered["score_request_digest"], + "logprobs": altered["teacher_logprobs"]}) + altered["digest"] = digest({k: v for k, v in altered.items() if k != "digest"}) + with pytest.raises(ValueError): + validate_row(altered, **kwargs, allow_synthetic=True) + resume = root / "checkpoint-000001.pt" + resumed = deepcopy(value) + resumed["output"]["directory"] = str(tmp_path / "resumed") + resumed["training"]["steps"] = 2 + Path(resumed["output"]["directory"]).mkdir() + train(resumed, allow_synthetic=True, resume=resume) + continuous = deepcopy(resumed) + continuous["output"]["directory"] = str(tmp_path / "continuous") + Path(continuous["output"]["directory"]).mkdir() + train(continuous, allow_synthetic=True) + expected = torch.load(Path(continuous["output"]["directory"]) / "checkpoint-000002.pt", weights_only=True) + actual = torch.load(Path(resumed["output"]["directory"]) / "checkpoint-000002.pt", weights_only=True) + assert expected["parameter_digest"] == actual["parameter_digest"] + + +def test_prm_failure_precedes_optimizer(tmp_path, monkeypatch): + pytest.importorskip("torch") + pytest.importorskip("transformers") + from harborrl.distill import runner + value = config(tmp_path, method="prm", teacher={"kind": "none"}, prm={"kind": "fake"}) + root = Path(value["output"]["directory"]) + root.mkdir() + def invalid(*args): + return {"evidence_kind": "synthetic", "steps": [], "token_returns": []} + monkeypatch.setattr(runner, "process_rewards", invalid) + with pytest.raises(ValueError, match="PRM"): + runner.train(value, allow_synthetic=True) + assert not list(root.glob("checkpoint-*.pt")) + assert not (root / "summary.json").exists() + + +def test_deployment_requires_explicit_opt_in(tmp_path, monkeypatch): + from harborrl.distill.cli import plan + monkeypatch.chdir(tmp_path) + (tmp_path / "deployment.yaml").write_text("invalid deployment") + assert plan(config(tmp_path))["deployment_source"] == "embedded-local" diff --git a/tests/harborrl/test_gateway_auto_shards.py b/tests/harborrl/test_gateway_auto_shards.py new file mode 100644 index 000000000..d57823cbc --- /dev/null +++ b/tests/harborrl/test_gateway_auto_shards.py @@ -0,0 +1,1217 @@ +"""CPU contract evidence for elastic auto-sharded gateways. + +Default is exactly one gateway. Auto routing leases one gateway per RUNNING +sandbox attempt, keeps at most ``sandboxes_per_shard`` concurrent attempts +per gateway (default 30, configurable), and grows the fleet when capacity is +exhausted — on a declared remote host first (a shard process launched over +ssh), locally only when no candidate answers. Either way the new shard joins +the live admission epoch atomically, so group admission never observes shard +disagreement. Group settlement retires idle spawned gateways symmetrically +(tail-only, never the initial gateway, names never reused). Everything fails +closed: bad capacity values, a static table combined with auto routing, and +leaked leases are all loud. +""" +from __future__ import annotations + +import asyncio +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +def shard_environment(tmp_path, **extra): + environment = { + "RUN_ID": "run-1", "RUN_DIR": str(tmp_path), + "HARBORRL_NATIVE_CATALOG": str(tmp_path / "catalog.json"), + "HARBORRL_NATIVE_PROFILE": str(tmp_path / "profile.json"), + "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, + "HARBORRL_NATIVE_TEMPLATE_DIGEST": "1" * 64, + "HARBORRL_NATIVE_GATEWAY_HOST": "127.0.0.1", + "HARBORRL_NATIVE_GATEWAY_PORT": "49184", + "HARBORRL_NATIVE_GATEWAY_URL": "http://127.0.0.1:49184", + "HARBORRL_NATIVE_RUNNER_PYTHON": "/python", + "HARBORRL_NATIVE_RUNNER_WORKERS": json.dumps(["root@worker-a"]), + "ROLLOUT_MAX_RESPONSE_LEN": "128", + "ROLLOUT_MAX_CONTEXT_LEN": "1024", + } + environment.update(extra) + return environment + + +def shard_config(tmp_path, monkeypatch, **extra): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment(tmp_path, **extra) + for key, value in environment.items(): + monkeypatch.setenv(key, value) + return RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + + +def build_auto_runtime(tmp_path, monkeypatch, *, sandboxes_per_shard, concurrency=2, + gateway_hosts=(), shard_python=None, shards_json=None): + """One real auto-routing runtime over fake server/registry/backend parts.""" + from harborrl.gateway import server as gateway_server + from harborrl.gateway import sglang_backend, trace + from harborrl.rollout import runtime as rollout_runtime + + (tmp_path / "catalog.json").write_text("[]") + (tmp_path / "profile.json").write_text(json.dumps({"model": "policy"})) + tokenizer = SimpleNamespace( + apply_chat_template=lambda *args, **kwargs: [1, 2], + get_vocab=lambda: {"token": 1}, + ) + tokenizer_digest, template_digest = rollout_runtime.tokenizer_fingerprint(tokenizer) + extra = { + "HARBORRL_NATIVE_TOKENIZER_DIGEST": tokenizer_digest, + "HARBORRL_NATIVE_TEMPLATE_DIGEST": template_digest, + "HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD": str(sandboxes_per_shard), + } + if shards_json: + # A declared table coexisting with the growth pool: remote entries are + # reached through the patched RemoteShard, never a real connection. + # A table alone pins routing to modulo, so the coexistence case must + # declare auto explicitly (the schema default does this for a pool). + extra["HARBORRL_NATIVE_GATEWAY_SHARDS"] = shards_json + if gateway_hosts: + extra["HARBORRL_NATIVE_GATEWAY_ROUTING"] = "auto" + if gateway_hosts: + # Remote growth needs the shared run token and a tokenizer checkpoint + # the spawned process loads itself. + (tmp_path / "run-token").write_text("shared-run-token") + extra.update( + HARBORRL_NATIVE_GATEWAY_HOSTS=json.dumps(list(gateway_hosts)), + HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(tmp_path / "run-token"), + HF_CKPT="/shared/ckpt", + ) + if shard_python: + extra["HARBORRL_NATIVE_SHARD_PYTHON"] = shard_python + environment = shard_environment(tmp_path, **extra) + for key, value in environment.items(): + monkeypatch.setenv(key, value) + + state = SimpleNamespace(servers=[], registries=[], closed=[]) + + class FakeGateway: + def __init__(self, registry, backend, *, model, max_output_tokens): + pass + + class FakeServer: + sequence = 0 + + def __init__(self): + FakeServer.sequence += 1 + self.server_address = ("127.0.0.1", 40000 + FakeServer.sequence) + + def serve_forever(self): + pass + + def shutdown(self): + pass + + def server_close(self): + state.closed.append(self.server_address[1]) + + class FakeTraceRegistry: + def __init__(self, path): + self.path = str(path) + self.version = None + self.accepting = True + state.registries.append(self) + + def publish_version(self, states, expected): + self.version, self.accepting = str(expected), True + + def close_admission(self): + self.accepting = False + + def drain(self, timeout): + return None + + class FakeBackend: + def __init__(self, *args, **kwargs): + pass + + def fake_make_server(gateway, host, port, *, max_inflight=None): + server = FakeServer() + state.servers.append((host, port, server.server_address[1], max_inflight)) + return server + + monkeypatch.setattr(gateway_server, "Gateway", FakeGateway) + monkeypatch.setattr(gateway_server, "make_server", fake_make_server) + monkeypatch.setattr(sglang_backend, "SGLangBackend", FakeBackend) + monkeypatch.setattr(trace, "TraceRegistry", FakeTraceRegistry) + if shards_json: + # A declared remote shard would otherwise construct the real + # RemoteShard and dial its admin origin. The static path imports it + # lazily inside __init__, so the seam is the source module's + # attribute, not the runtime module's re-export. + from harborrl.rollout import remote_shard as remote_module + + monkeypatch.setattr(remote_module, "RemoteShard", FakeRemoteShard) + + config = rollout_runtime.RolloutRuntimeConfig.from_environment( + environment, endpoint="http://127.0.0.1:30000", + concurrency_limit=concurrency, group_size=2) + runtime_ = rollout_runtime.RolloutRuntime(config, tokenizer) + return runtime_, state + + +def group_entry(tmp_path): + return {"id": "task-0", "revision": 1, "path": str(tmp_path / "task.md"), + "task_digest": "d" * 64, "reward_profile": {}} + + +class PermissiveSlots: + """Slot bookkeeping that accepts fake IRs. + + The real GroupSlots runs ``require_ready`` on every result and no fake + attempt can produce a valid trajectory.v2 record; the lease behavior + under test only needs start/finish bookkeeping. + """ + + max_attempts = 2 + + def __init__(self, identities, *, max_attempts=2): + self.identities = {identity.slot_id: identity for identity in identities} + + def start(self, slot): + return self.identities[slot] + + def finish(self, slot, ir=None, error=None): + assert (ir is None) != (error is None) + + +def lease(runtime_): + """lease_shard is async: remote growth waits on an ssh subprocess.""" + return asyncio.run(runtime_.lease_shard()) + + +class FakeRemoteShard: + """A SpawnedRemoteShard stand-in carrying a REAL TraceRegistry. + + The join replay must cross publish_version's closed-barrier precondition + — fake registries cannot see it (see the regression test below) — so the + remote-spawn tests keep the real registry behind the fake transport. + """ + + def __init__(self, spec, *, token, trace_dir, process=None, timeout=30.0): + from harborrl.gateway.trace import TraceRegistry + + self.spec = spec + self.trace_dir = Path(trace_dir) + self.registry = TraceRegistry(trace_dir) + self.process = process + self.started = False + self.closed = False + + @property + def name(self): + return self.spec.name + + @property + def url(self): + return self.spec.url + + @property + def port(self): + return self.spec.port + + def start(self): + self.started = True + + def close(self): + self.closed = True + + +def fake_remote_launch(monkeypatch, runtime_module, *, delay=0.0, failures=0): + """Patch the remote transport; record every launch attempt. + + The first ``failures`` launches raise, exercising the candidate loop and + the local fallback; every launch resolves distinct ports like a real + ready line would. + """ + from harborrl.gateway.shard_service import SHARD_SERVICE + + launched = [] + + class FakeProcess: + returncode = 0 + + async def start_process(host, environment, *, port=None, identity=None): + entry = {"host": host, "name": environment["HARBORRL_SHARD_NAME"]} + if port is not None: + entry["port"] = port + if identity is not None: + entry["identity"] = identity + launched.append(entry) + if len(launched) <= failures: + raise OSError(f"connection timed out: {host}") + if delay: + await asyncio.sleep(delay) + port = 41000 + len(launched) + return (FakeProcess(), {"shard": environment["HARBORRL_SHARD_NAME"], + "service": SHARD_SERVICE, + "gateway_port": port, "admin_port": port + 1}) + + monkeypatch.setattr(runtime_module, "SpawnedRemoteShard", FakeRemoteShard) + return start_process, launched + + +# -- config ------------------------------------------------------------------- + + +def test_absent_table_defaults_to_one_auto_gateway_of_thirty(tmp_path, monkeypatch): + config = shard_config(tmp_path, monkeypatch) + assert config.shards == () + assert config.gateway_routing == "auto" + assert config.sandboxes_per_shard == 30 + rendered = config.environment() + assert rendered["HARBORRL_NATIVE_GATEWAY_ROUTING"] == "auto" + assert rendered["HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD"] == "30" + + +def test_sandboxes_per_shard_is_configurable_and_round_trips(tmp_path, monkeypatch): + config = shard_config( + tmp_path, monkeypatch, HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD="7") + assert config.sandboxes_per_shard == 7 + assert config.environment()["HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD"] == "7" + + +@pytest.mark.parametrize("value", ["0", "-3", "many"]) +def test_sandboxes_per_shard_fails_closed(tmp_path, monkeypatch, value): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment( + tmp_path, HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD=value) + with pytest.raises(ValueError, match="(?i)sandboxes_per_shard"): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", + concurrency_limit=2, group_size=2) + + +def test_auto_routing_accepts_a_static_table_with_a_growth_pool(tmp_path, monkeypatch): + """A long-lived declared shard is the floor growth expands beyond. + + The tunnel-host topology: one statically declared remote shard exists at + startup, and auto routing still grows extra gateways on the pool when + concurrency passes the per-shard budget. + """ + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment( + tmp_path, + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([ + {"name": "static-1", "mode": "remote", + "url": "http://tunnel-host:5010", "admin_url": "http://tunnel-host:5011"}]), + HARBORRL_NATIVE_GATEWAY_HOSTS=json.dumps(["gpu-a"]), + HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(tmp_path / "t"), + HF_CKPT="/ckpt", + HARBORRL_NATIVE_GATEWAY_ROUTING="auto") + (tmp_path / "t").write_text("token") + config = RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", + concurrency_limit=2, group_size=2) + assert config.gateway_routing == "auto" + assert [spec.name for spec in config.shards] == ["static-1"] + assert config.gateway_hosts == ("gpu-a",) + + +def test_growth_names_skip_names_a_static_table_claims(tmp_path, monkeypatch): + """A regrown shard must never reuse a declared static name.""" + from harborrl.rollout import runtime as rollout_runtime + + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, gateway_hosts=("gpu-a",), + shards_json=json.dumps([{"name": "gw1", "mode": "remote", + "url": "http://declared:5010", + "admin_url": "http://declared:5011"}])) + start_process, launched = fake_remote_launch(monkeypatch, rollout_runtime) + monkeypatch.setattr(runtime_, "_start_shard_process", start_process) + try: + lease(runtime_) # fills the declared gw1 + states = [{"endpoint": "http://router:30000", "healthy": True, + "weight_version": "2"}] + runtime_.registry.publish_version(states, "2") + grown = lease(runtime_) + # gw1 is taken by the table, so growth starts at gw2. + assert grown.name == "gw2" + assert launched == [{"host": "gpu-a", "name": "gw2"}] + finally: + runtime_.close() + + +def test_explicit_modulo_without_a_table_opts_out_of_auto(tmp_path, monkeypatch): + config = shard_config( + tmp_path, monkeypatch, HARBORRL_NATIVE_GATEWAY_ROUTING="modulo") + assert config.gateway_routing == "modulo" + + +def test_declared_table_keeps_the_static_modulo_default(tmp_path, monkeypatch): + config = shard_config( + tmp_path, monkeypatch, + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([{"name": "gw0"}, {"name": "gw1"}])) + assert config.gateway_routing == "modulo" + + +def test_gateway_hosts_round_trip_through_the_environment(tmp_path, monkeypatch): + config = shard_config( + tmp_path, monkeypatch, + HARBORRL_NATIVE_GATEWAY_HOSTS=json.dumps(["root@gpu-a", "gpu-b"]), + HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(tmp_path / "run-token"), + HF_CKPT="/shared/ckpt") + assert config.gateway_hosts == ("root@gpu-a", "gpu-b") + rendered = config.environment() + assert rendered["HARBORRL_NATIVE_GATEWAY_HOSTS"] == '["root@gpu-a", "gpu-b"]' + assert rendered["HF_CKPT"] == "/shared/ckpt" + assert "HARBORRL_NATIVE_SHARD_PYTHON" not in rendered + + +def test_shard_python_overrides_the_growth_interpreter(tmp_path, monkeypatch): + """Growth hosts run the serving stack, not the runner's sandbox framework.""" + runtime_, _ = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, + gateway_hosts=("gpu-a",), shard_python="/serving/python") + try: + assert runtime_._shard_launch_python() == "/serving/python" + assert (runtime_.config.environment()["HARBORRL_NATIVE_SHARD_PYTHON"] + == "/serving/python") + finally: + runtime_.close() + + runtime_, _ = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, gateway_hosts=("gpu-a",)) + try: + # Undeclared keeps the runner's interpreter (uniform deployments). + assert runtime_._shard_launch_python() == runtime_.config.runner_python + assert "HARBORRL_NATIVE_SHARD_PYTHON" not in runtime_.config.environment() + finally: + runtime_.close() + + +@pytest.mark.parametrize("overrides, message", [ + # Remote growth shares the run token and needs a tokenizer checkpoint the + # spawned process loads itself. + ({"HARBORRL_NATIVE_GATEWAY_HOSTS": json.dumps(["gpu-a"]), + "HF_CKPT": "/ckpt"}, "SHARD_TOKEN_FILE"), + ({"HARBORRL_NATIVE_GATEWAY_HOSTS": json.dumps(["gpu-a"]), + "HARBORRL_NATIVE_SHARD_TOKEN_FILE": "t"}, "HF_CKPT"), + ({"HARBORRL_NATIVE_GATEWAY_HOSTS": json.dumps(["gpu-a"]), + "HARBORRL_NATIVE_SHARD_TOKEN_FILE": "t", "HF_CKPT": "/ckpt", + "HARBORRL_NATIVE_GATEWAY_ROUTING": "modulo"}, "auto routing only"), +]) +def test_gateway_hosts_fail_closed(tmp_path, monkeypatch, overrides, message): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment(tmp_path, **overrides) + with pytest.raises(ValueError, match=message): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", + concurrency_limit=2, group_size=2) + + +@pytest.mark.parametrize("value", ["gpu-a", '["", "gpu-a"]', '{"host": "a"}', "[[]]"]) +def test_gateway_hosts_env_garbage_fails_closed(tmp_path, monkeypatch, value): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment( + tmp_path, HARBORRL_NATIVE_GATEWAY_HOSTS=value, + HARBORRL_NATIVE_SHARD_TOKEN_FILE="t", HF_CKPT="/ckpt") + with pytest.raises(ValueError, match="(?i)gateway_hosts|ssh destinations"): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", + concurrency_limit=2, group_size=2) + + +# -- lease, spawn, join ------------------------------------------------------- + + +def test_leases_fill_one_gateway_then_spawn_the_next(tmp_path, monkeypatch): + runtime_, state = build_auto_runtime(tmp_path, monkeypatch, sandboxes_per_shard=2) + try: + first = lease(runtime_) + second = lease(runtime_) + assert first is second is runtime_.shards[0] + assert first.name == "gateway" + # Capacity n reached: the (n+1)-th concurrent sandbox spawns gw1 on an + # ephemeral port with its own trace subtree. + third = lease(runtime_) + assert third.name == "gw1" + assert third.port != first.port and third.url != first.url + assert third.url == f"http://127.0.0.1:{third.port}" + assert third.trace_dir == runtime_.run_root / "native-traces" / "shard-gw1" + assert state.servers[1][1] == 0 # the OS owns spawned-shard exclusivity + # The evidence record grows with the live shard set. + evidence = json.loads((runtime_.run_root / "gateway-shards.json").read_text()) + assert evidence["gateway_routing"] == "auto" + assert evidence["sandboxes_per_shard"] == 2 + assert [(s["name"], s["port"]) for s in evidence["shards"]] == [ + ("gateway", 49184), ("gw1", third.port)] + # A released lease reuses freed capacity instead of spawning again. + runtime_.release_shard(third) + again = lease(runtime_) + assert again is third + assert len(state.servers) == 2 + finally: + runtime_.close() + assert len(state.closed) == 2 + + +def test_spawned_shard_joins_the_live_policy_epoch(tmp_path, monkeypatch): + runtime_, state = build_auto_runtime(tmp_path, monkeypatch, sandboxes_per_shard=1) + try: + lease(runtime_) + states = [{"endpoint": "http://router:30000", "healthy": True, + "weight_version": "2"}] + runtime_.registry.publish_version(states, "2") + # The (n+1)-th lease spawns gw1 straight onto the live epoch: joining + # with version None would read as shard disagreement and fail-close + # group admission until the next activation. + spawned = lease(runtime_) + assert spawned.name == "gw1" + assert spawned.registry.version == "2" and spawned.registry.accepting + assert runtime_.registry.version == "2" and runtime_.registry.accepting + # The epoch barrier spans the late shard too. + runtime_.registry.drain(0.05) + assert all(registry.accepting is False for registry in state.registries) + finally: + runtime_.close() + + +def test_spawned_shard_before_the_first_epoch_stays_unversioned(tmp_path, monkeypatch): + runtime_, _ = build_auto_runtime(tmp_path, monkeypatch, sandboxes_per_shard=1) + try: + lease(runtime_) + spawned = lease(runtime_) + assert spawned.registry.version is None + assert runtime_.registry.version is None and runtime_.registry.accepting + finally: + runtime_.close() + + +def test_join_replays_a_published_epoch_onto_a_real_registry(tmp_path): + """Regression: the replay must satisfy publish_version's barrier. + + A fresh TraceRegistry opens accepting and publish_version refuses to + open a barrier that was never closed; replaying an epoch without closing + first raised "all attempts must be sealed behind a closed admission + barrier" and every auto-grown shard failed to spawn mid-run (the fake + registries used elsewhere cannot see this precondition). + """ + from harborrl.gateway.trace import TraceRegistry + from harborrl.rollout.runtime import ShardAdmission + + shards = [] + admission = ShardAdmission(shards) + shards.append(SimpleNamespace(registry=TraceRegistry(tmp_path / "gw0"))) + states = [{"endpoint": "http://router:30000", "healthy": True, + "weight_version": "2"}] + admission.drain(0.05) + admission.publish_version(states, "2") + + grown = SimpleNamespace(registry=TraceRegistry(tmp_path / "gw1")) + admission.join(grown) + + # The late shard lands on the live epoch, open for admission. + assert grown.registry.version == "2" and grown.registry.accepting + assert admission.version == "2" and admission.accepting + assert len(admission.shards) == 2 + # The epoch barrier spans the late shard too. + admission.drain(0.05) + assert not grown.registry.accepting + + +# -- remote-first growth ------------------------------------------------------ + + +def test_growth_prefers_declared_remote_hosts(tmp_path, monkeypatch): + """Growth fills the declared pool before ever binding a local port.""" + from harborrl.rollout import runtime as rollout_runtime + + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, gateway_hosts=("root@gpu-a",)) + start_process, launched = fake_remote_launch(monkeypatch, rollout_runtime) + monkeypatch.setattr(runtime_, "_start_shard_process", start_process) + try: + first = lease(runtime_) + states = [{"endpoint": "http://router:30000", "healthy": True, + "weight_version": "2"}] + runtime_.registry.publish_version(states, "2") + grown = lease(runtime_) + assert first.name == "gateway" + assert launched == [{"host": "root@gpu-a", "name": "gw1"}] + assert grown.name == "gw1" and grown.spec.mode == "remote" + # The ssh destination spelling becomes the URL host (user@ stripped), + # with the ports the ready line resolved. + assert grown.url == "http://gpu-a:41001" + assert grown.spec.admin_url == "http://gpu-a:41002" + assert grown.trace_dir == runtime_.run_root / "native-traces" / "shard-gw1" + assert grown.started and not grown.closed + # The grown shard joined the live epoch over its real registry — the + # replay crossed publish_version's barrier precondition. + assert grown.registry.version == "2" and grown.registry.accepting + assert runtime_.registry.version == "2" + # Growth bound no local port and stayed quiet about fallbacks. + assert len(state.servers) == 1 + assert not (runtime_.run_root / "gateway-spawn-fallback.json").exists() + evidence = json.loads((runtime_.run_root / "gateway-shards.json").read_text()) + assert [(s["name"], s["mode"]) for s in evidence["shards"]] == [ + ("gateway", "local"), ("gw1", "remote")] + assert evidence["shards"][1]["admin_url"] == "http://gpu-a:41002" + finally: + runtime_.close() + + +def test_remote_candidates_rotate_and_fall_back_to_local(tmp_path, monkeypatch): + """Every declared host is tried in rotation; only then does growth bind locally.""" + from harborrl.rollout import runtime as rollout_runtime + + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, + gateway_hosts=("gpu-a", "gpu-b")) + start_process, launched = fake_remote_launch(monkeypatch, rollout_runtime, + failures=2) + monkeypatch.setattr(runtime_, "_start_shard_process", start_process) + try: + # Entries normalize to (destination, ssh port, identity, url host); + # a plain string inherits its transport options from ssh config. + assert runtime_._candidate_hosts() == ( + ("gpu-a", None, None, "gpu-a"), ("gpu-b", None, None, "gpu-b")) + assert runtime_._candidate_hosts() == ( + ("gpu-b", None, None, "gpu-b"), ("gpu-a", None, None, "gpu-a")) + lease(runtime_) + # Both remote candidates time out: growth falls back to the manager + # machine and records why. + grown = lease(runtime_) + assert grown.name == "gw1" and grown.spec.mode == "local" + assert len(state.servers) == 2 # a local ephemeral port was bound + record = json.loads( + (runtime_.run_root / "gateway-spawn-fallback.json").read_text()) + assert record["fallback"] == "local" + assert [attempt["host"] for attempt in record["attempts"]] == ["gpu-a", "gpu-b"] + assert all("OSError" in attempt["error"] for attempt in record["attempts"]) + finally: + runtime_.close() + + +def test_a_spawn_in_flight_does_not_double_grow(tmp_path, monkeypatch): + """Concurrent leases at capacity grow exactly once. + + With the initial gateway full (n=2 held), two more leases arrive while + a remote spawn is in flight: the second waiter re-checks capacity under + the lease lock after that spawn finishes, so it fills the fresh shard + instead of pulling up another one. + """ + from harborrl.rollout import runtime as rollout_runtime + + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=2, concurrency=4, + gateway_hosts=("gpu-a",)) + start_process, launched = fake_remote_launch(monkeypatch, rollout_runtime, + delay=0.05) + monkeypatch.setattr(runtime_, "_start_shard_process", start_process) + + async def scenario(): + await runtime_.lease_shard() + await runtime_.lease_shard() # the initial gateway is at capacity + third, fourth = await asyncio.gather( + runtime_.lease_shard(), runtime_.lease_shard()) + return third, fourth + + try: + third, fourth = asyncio.run(scenario()) + assert launched == [{"host": "gpu-a", "name": "gw1"}] + assert third.name == fourth.name == "gw1" + assert len(state.servers) == 1 # no local growth happened either + assert runtime_._auto_inflight == {"gateway": 2, "gw1": 2} + finally: + runtime_.close() + + +def test_spawned_remote_shard_close_reaps_the_transport(tmp_path): + """close() shuts down via the admin plane, then kills a hung transport.""" + from harborrl.rollout.remote_shard import RemoteShard + from harborrl.rollout.runtime import GatewayShardSpec, SpawnedRemoteShard + + class FakeProcess: + def __init__(self): + self.returncode = None + self.killed = False + + async def wait(self): + if self.killed: + return self.returncode + await asyncio.sleep(3600) + + def kill(self): + self.killed = True + self.returncode = 1 + + spec = GatewayShardSpec(name="gw1", port=41001, url="http://gpu-a:41001", + mode="remote", admin_url="http://127.0.0.1:1") + process = FakeProcess() + shard = SpawnedRemoteShard(spec, token="tok", trace_dir=tmp_path, + process=process) + shard.SHUTDOWN_GRACE = 0.05 + assert isinstance(shard, RemoteShard) # presents the remote-shard surface + + async def scenario(): + shard.close() + await asyncio.sleep(0.2) # the reap task must run past its grace + + asyncio.run(scenario()) + assert process.killed and process.returncode == 1 + + +def test_spawned_remote_shard_close_waits_for_a_clean_exit(tmp_path): + """A transport that exits after the admin shutdown is reaped, never killed.""" + from harborrl.rollout.runtime import GatewayShardSpec, SpawnedRemoteShard + + class FakeProcess: + def __init__(self): + self.returncode = None + self.killed = False + + async def wait(self): + await asyncio.sleep(0.02) + self.returncode = 0 + return 0 + + def kill(self): + self.killed = True + + spec = GatewayShardSpec(name="gw1", port=41001, url="http://gpu-a:41001", + mode="remote", admin_url="http://127.0.0.1:1") + process = FakeProcess() + shard = SpawnedRemoteShard(spec, token="tok", trace_dir=tmp_path, + process=process) + shard.SHUTDOWN_GRACE = 5.0 + + async def scenario(): + shard.close() + await asyncio.sleep(0.3) + + asyncio.run(scenario()) + assert not process.killed and process.returncode == 0 + + +def test_remote_shard_spawn_environment_satisfies_the_service_contract( + tmp_path, monkeypatch): + """The inline ssh environment is exactly what ShardServiceConfig parses. + + A missing key would only surface as a remote process failing to start; + parsing the built environment here pins the contract to compile time. + """ + from harborrl.gateway.shard_service import ShardServiceConfig + from harborrl.rollout.runtime import _shard_service_command, _shard_service_environment + + config = shard_config( + tmp_path, monkeypatch, + HARBORRL_NATIVE_GATEWAY_HOSTS=json.dumps(["root@gpu-a"]), + HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(tmp_path / "run-token"), + HF_CKPT="/shared/ckpt") + environment = _shard_service_environment(config, "gw1") + parsed = ShardServiceConfig.from_environment(environment) + # Both ports bind ephemeral; the manager reads the resolved ports from + # the ready line and composes the URL itself. + assert parsed.gateway_port == 0 and parsed.gateway_url == "" + assert parsed.admin_port == 0 + assert parsed.run_id == config.run_id and parsed.shard_name == "gw1" + assert parsed.run_root == config.run_root + assert parsed.token_file == config.shard_token_file + assert parsed.endpoint == config.endpoint + assert parsed.profile_path == config.profile_path + assert parsed.model_tag == config.model_tag + assert parsed.max_output_tokens == config.max_output_tokens + assert parsed.max_context == config.max_context + assert parsed.expected_tokenizer_digest == config.expected_tokenizer_digest + assert parsed.expected_template_digest == config.expected_template_digest + assert parsed.audited_raw_logprobs == config.audited_raw_logprobs + + command = _shard_service_command( + "root@gpu-a", environment, python="/train/python", project="/project") + assert command[0] == "ssh" and command[-2] == "root@gpu-a" + assert {"BatchMode=yes", "ClearAllForwardings=yes", + "StrictHostKeyChecking=accept-new"} <= set(command) + remote = command[-1] + assert remote.startswith("cd /project && exec env ") + assert remote.endswith("/train/python -m harborrl.gateway.shard_service") + assert "RUN_ID=run-1" in remote and "HARBORRL_SHARD_NAME=gw1" in remote + # ssh forwards no environment: every value crosses shell-quoted. + quoted = _shard_service_command( + "gpu-a", {"K": "a b"}, python="/p", project="/pr")[-1] + assert "K='a b'" in quoted + + +# -- group dispatch ----------------------------------------------------------- + + +def test_group_dispatch_leases_and_releases_gateways(tmp_path, monkeypatch): + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, concurrency=2) + from harborrl.rollout import runtime as rollout_runtime + + seen = [] + + async def fake_attempt(identity, **kwargs): + await asyncio.sleep(0) # overlap the two slots so both leases are live + seen.append(kwargs["gateway_url"]) + return {"identity": identity.to_dict()} + + monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) + monkeypatch.setattr(rollout_runtime, "run_attempt", fake_attempt) + try: + result = asyncio.run(rollout_runtime.run_group( + runtime_.config, runtime_, entry=group_entry(tmp_path), group_index=0, + rollout_id=0, sampling_digest="s" * 64, policy_version="1", + group_size=2)) + assert len(result) == 2 + # n=1 and two concurrent slots: each running sandbox gets its own + # gateway, and the second one was pulled up on demand. + assert len(set(seen)) == 2 + assert len(state.servers) == 2 + # Every lease was released once the group settled. + assert runtime_._auto_inflight == {} + finally: + runtime_.close() + + +def test_failed_attempt_releases_its_gateway_lease(tmp_path, monkeypatch): + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, concurrency=1) + from harborrl.rollout import runtime as rollout_runtime + + calls = [] + + async def flaky_attempt(identity, **kwargs): + calls.append(identity.slot_id) + if identity.slot_id == "0" and calls.count("0") == 1: + raise ValueError("sandbox exploded") + await asyncio.sleep(0) + return {"identity": identity.to_dict()} + + monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) + monkeypatch.setattr(rollout_runtime, "run_attempt", flaky_attempt) + try: + result = asyncio.run(rollout_runtime.run_group( + runtime_.config, runtime_, entry=group_entry(tmp_path), group_index=0, + rollout_id=0, sampling_digest="s" * 64, policy_version="1", + group_size=2)) + assert len(result) == 2 + # Slot 0 failed once and its retry succeeded; the semaphore may wake + # slot 1 between the two attempts, so the order stays unspecified. + assert sorted(calls) == ["0", "0", "1"] + # A released lease is reusable: no leaked occupancy ever spawns a + # gateway that the running sandbox count does not justify. + assert len(state.servers) == 1 + assert runtime_._auto_inflight == {} + finally: + runtime_.close() + + +def test_auto_group_below_capacity_stays_on_one_gateway(tmp_path, monkeypatch): + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=30, concurrency=2) + from harborrl.rollout import runtime as rollout_runtime + + async def fake_attempt(identity, **kwargs): + await asyncio.sleep(0) + return {"identity": identity.to_dict()} + + monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) + monkeypatch.setattr(rollout_runtime, "run_attempt", fake_attempt) + try: + result = asyncio.run(rollout_runtime.run_group( + runtime_.config, runtime_, entry=group_entry(tmp_path), group_index=0, + rollout_id=0, sampling_digest="s" * 64, policy_version="1", + group_size=2)) + assert len(result) == 2 + assert len(state.servers) == 1 + assert runtime_._auto_inflight == {} + finally: + runtime_.close() + + +# -- auto shrink --------------------------------------------------------------- + + +def test_settlement_retires_idle_spawned_gateways(tmp_path, monkeypatch): + """Group settlement is the shrink trigger. + + n=1 with two concurrent slots pulls up gw1 mid-group; once the group + settles and every lease is released, the running sandbox count no longer + justifies gw1 — it detaches, closes, and the evidence record shrinks + back to exactly the initial gateway. + """ + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, concurrency=2) + from harborrl.rollout import runtime as rollout_runtime + + async def fake_attempt(identity, **kwargs): + await asyncio.sleep(0) # overlap both slots so gw1 is pulled up + return {"identity": identity.to_dict()} + + monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) + monkeypatch.setattr(rollout_runtime, "run_attempt", fake_attempt) + try: + result = asyncio.run(rollout_runtime.run_group( + runtime_.config, runtime_, entry=group_entry(tmp_path), group_index=0, + rollout_id=0, sampling_digest="s" * 64, policy_version="1", + group_size=2)) + assert len(result) == 2 + assert len(state.servers) == 2 # gw1 was pulled up mid-group + # Settlement retired it: one live shard, one closed server, and the + # evidence record shrank back with the live set. + assert [shard.name for shard in runtime_.shards] == ["gateway"] + assert len(state.closed) == 1 + evidence = json.loads((runtime_.run_root / "gateway-shards.json").read_text()) + assert [s["name"] for s in evidence["shards"]] == ["gateway"] + assert runtime_._auto_spawned == set() + assert runtime_._auto_inflight == {} + finally: + runtime_.close() + assert len(state.closed) == 2 # the initial gateway closes with the runtime + + +def test_failed_group_settlement_still_retires_spawned_gateways(tmp_path, monkeypatch): + """A failed group settles too: shrink runs on the failure path as well.""" + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, concurrency=2) + from harborrl.rollout import runtime as rollout_runtime + + async def broken_attempt(identity, **kwargs): + await asyncio.sleep(0) # overlap so gw1 spawns before both fail + raise ValueError("sandbox exploded") + + monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) + monkeypatch.setattr(rollout_runtime, "run_attempt", broken_attempt) + try: + with pytest.raises(RuntimeError, match="exhausted"): + asyncio.run(rollout_runtime.run_group( + runtime_.config, runtime_, entry=group_entry(tmp_path), group_index=0, + rollout_id=0, sampling_digest="s" * 64, policy_version="1", + group_size=2)) + assert [shard.name for shard in runtime_.shards] == ["gateway"] + assert runtime_._auto_inflight == {} + finally: + runtime_.close() + + +def test_retired_shard_names_are_never_reused(tmp_path, monkeypatch): + """A retired gateway's trace subtree belongs to its own history. + + Retirement advances the name sequence, so regrowth is a fresh gw2 with a + fresh trace subtree instead of colliding with retired gw1's evidence. + """ + runtime_, state = build_auto_runtime(tmp_path, monkeypatch, sandboxes_per_shard=1) + try: + first = lease(runtime_) + second = lease(runtime_) + assert first.name == "gateway" and second.name == "gw1" + runtime_.release_shard(first) + runtime_.release_shard(second) + runtime_.retire_idle_shards() + assert [shard.name for shard in runtime_.shards] == ["gateway"] + + lease(runtime_) + regrown = lease(runtime_) + assert regrown.name == "gw2" + assert regrown.trace_dir == runtime_.run_root / "native-traces" / "shard-gw2" + assert regrown is not second + finally: + runtime_.close() + + +def test_busy_tail_blocks_retirement(tmp_path, monkeypatch): + """Only idle shards retire, and only from the tail. + + Leases fill in order, so anything behind a busy shard stays: here gw1 + still runs a sandbox while gw0 is idle — the busy tail blocks, and the + idle initial gateway is never a retirement candidate anyway. + """ + runtime_, _ = build_auto_runtime(tmp_path, monkeypatch, sandboxes_per_shard=1) + try: + first = lease(runtime_) + second = lease(runtime_) + runtime_.release_shard(first) # gw0 idle, gw1 still busy + runtime_.retire_idle_shards() + assert [shard.name for shard in runtime_.shards] == ["gateway", "gw1"] + + runtime_.release_shard(second) + runtime_.retire_idle_shards() + assert [shard.name for shard in runtime_.shards] == ["gateway"] + finally: + runtime_.close() + + +def test_leave_detaches_a_shard_from_the_admission_facade(): + """leave is join's mirror under the same lock, and idempotent.""" + from harborrl.rollout.runtime import ShardAdmission + + admission = ShardAdmission([]) + first, second = SimpleNamespace(name="gw0"), SimpleNamespace(name="gw1") + admission.join(first) + admission.join(second) + assert admission.shards == (first, second) + admission.leave(second) + assert admission.shards == (first,) + admission.leave(second) # a settled retry must not raise on a gone shard + assert admission.shards == (first,) + + +# -- deployment schema and env overlay ---------------------------------------- + + +def test_deployment_accepts_auto_routing_with_capacity(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_routing": "auto", "sandboxes_per_shard": 20})) + manager = config["manager"] + assert manager["gateway_routing"] == "auto" + assert manager["sandboxes_per_shard"] == 20 + + +def test_deployment_auto_routing_accepts_a_shard_table(tmp_path): + """Auto routing with a declared floor is now a supported topology. + + The declared table is fixed at startup; auto routing grows beyond it on + the growth pool. Declaring the routing explicitly must simply stick. + """ + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_routing": "auto", + "gateway_shards": [{"name": "gw0"}]})) + assert config["manager"]["gateway_routing"] == "auto" + assert [shard["name"] for shard in config["manager"]["gateway_shards"]] == ["gw0"] + + +@pytest.mark.parametrize("value", [0, -3, "many", 2.5]) +def test_deployment_sandboxes_per_shard_fails_closed(tmp_path, value): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import DeploymentError, validate + + with pytest.raises(DeploymentError, match="sandboxes_per_shard"): + validate(deployment_mapping(tmp_path, manager={"sandboxes_per_shard": value})) + + +def test_deployment_legacy_manager_shape_stays_unchanged(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + legacy = validate(deployment_mapping(tmp_path))["manager"] + assert not ({"gateway_shards", "gateway_routing", "sandboxes_per_shard", + "worker_shards"} & set(legacy)) + + +def _auto_overlay(tmp_path, manager_over): + from test_deployment_config import deployment_mapping + + from harborrl.deployment import controller as dispatch + + mapped = deployment_mapping(tmp_path, manager=manager_over) + mapped["gpu_target"] = {"train_python": "/train/python"} + mapped["rollout"] = { + "provider": "ssh_docker", + "workers": [{"id": "worker-a", "ssh": {"host": "worker-a"}, + "runner_python": "/runner/python"}]} + return dispatch._rollout_environment({"config": {"harbor": {}}}, mapped, {}) + + +def test_rollout_environment_carries_auto_routing_and_capacity(tmp_path): + overlay = _auto_overlay( + tmp_path, {"gateway_routing": "auto", "sandboxes_per_shard": 20}) + assert overlay["HARBORRL_NATIVE_GATEWAY_ROUTING"] == "auto" + assert overlay["HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD"] == "20" + assert "HARBORRL_NATIVE_GATEWAY_SHARDS" not in overlay + + +def test_rollout_environment_leaves_undeclared_routing_to_the_default(tmp_path): + overlay = _auto_overlay(tmp_path, {}) + assert "HARBORRL_NATIVE_GATEWAY_ROUTING" not in overlay + assert "HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD" not in overlay + + +def test_rollout_environment_carries_an_explicit_static_opt_out(tmp_path): + overlay = _auto_overlay(tmp_path, {"gateway_routing": "modulo"}) + assert overlay["HARBORRL_NATIVE_GATEWAY_ROUTING"] == "modulo" + assert "HARBORRL_NATIVE_GATEWAY_SANDBOXES_PER_SHARD" not in overlay + + +# -- remote growth through the deployment and dispatch layers ------------------ + + +def test_settlement_retires_a_remotely_spawned_gateway(tmp_path, monkeypatch): + """Shrink is placement-agnostic: a remote shard retires like a local one.""" + from harborrl.rollout import runtime as rollout_runtime + + runtime_, state = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, concurrency=2, + gateway_hosts=("gpu-a",)) + start_process, launched = fake_remote_launch(monkeypatch, rollout_runtime) + monkeypatch.setattr(runtime_, "_start_shard_process", start_process) + + async def fake_attempt(identity, **kwargs): + await asyncio.sleep(0) # overlap both slots so gw1 is pulled up + return {"identity": identity.to_dict()} + + monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) + monkeypatch.setattr(rollout_runtime, "run_attempt", fake_attempt) + try: + result = asyncio.run(rollout_runtime.run_group( + runtime_.config, runtime_, entry=group_entry(tmp_path), group_index=0, + rollout_id=0, sampling_digest="s" * 64, policy_version="1", + group_size=2)) + assert len(result) == 2 + assert launched == [{"host": "gpu-a", "name": "gw1"}] + assert len(state.servers) == 1 # growth stayed remote throughout + # Settlement detached and closed the remote shard like a local one. + assert [shard.name for shard in runtime_.shards] == ["gateway"] + evidence = json.loads((runtime_.run_root / "gateway-shards.json").read_text()) + assert [s["name"] for s in evidence["shards"]] == ["gateway"] + assert runtime_._auto_spawned == set() + assert runtime_._auto_inflight == {} + finally: + runtime_.close() + + +def test_retire_detaches_and_closes_a_remotely_spawned_gateway(tmp_path, monkeypatch): + """Retirement leaves the facade and closes the remote shard in place.""" + from harborrl.rollout import runtime as rollout_runtime + + runtime_, _ = build_auto_runtime( + tmp_path, monkeypatch, sandboxes_per_shard=1, gateway_hosts=("gpu-a",)) + start_process, _ = fake_remote_launch(monkeypatch, rollout_runtime) + monkeypatch.setattr(runtime_, "_start_shard_process", start_process) + try: + first = lease(runtime_) + grown = lease(runtime_) + assert grown.name == "gw1" + runtime_.release_shard(first) + runtime_.release_shard(grown) + runtime_.retire_idle_shards() + assert [shard.name for shard in runtime_.shards] == ["gateway"] + assert grown.closed + assert runtime_.registry.shards == (first,) + finally: + runtime_.close() + + +def test_deployment_accepts_gateway_hosts(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_hosts": ["root@gpu-a", "gpu-b"]})) + assert config["manager"]["gateway_hosts"] == ["root@gpu-a", "gpu-b"] + # Without a declared table the pool itself activates auto routing. + assert config["manager"]["gateway_routing"] == "auto" + assert "gateway_shards" not in config["manager"] + + +def test_deployment_accepts_a_static_table_with_the_growth_pool(tmp_path): + """The tunnel floor plus elastic growth, declared in one deployment.""" + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_hosts": ["root@gpu-a"], + "gateway_shards": [ + {"name": "static-1", "mode": "remote", + "url": "http://tunnel-host:5010", "admin_url": "http://tunnel-host:5011"}]})) + assert [shard["name"] for shard in config["manager"]["gateway_shards"]] == ["static-1"] + assert config["manager"]["gateway_hosts"] == ["root@gpu-a"] + # Undeclared routing with a pool means growth beyond the declared floor. + assert config["manager"]["gateway_routing"] == "auto" + + +def test_deployment_accepts_growth_host_mappings(tmp_path): + """An ephemeral host carries its own key and port instead of ssh config.""" + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_hosts": [ + {"host": "10.0.0.7", "user": "root", "port": 2222, + "identity": str(tmp_path / "id_ed25519")}, + "root@gpu-b"]})) + assert config["manager"]["gateway_hosts"] == [ + {"host": "10.0.0.7", "user": "root", "port": 2222, + "identity": str(tmp_path / "id_ed25519")}, + "root@gpu-b"] + + +def test_deployment_accepts_shard_python_with_the_growth_pool(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import DeploymentError, validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_hosts": ["gpu-a"], + "shard_python": "/serving/venv/bin/python"})) + assert config["manager"]["shard_python"] == "/serving/venv/bin/python" + # The interpreter exists only for growth; without a pool it is a typo. + with pytest.raises(DeploymentError, match="gateway_hosts"): + validate(deployment_mapping(tmp_path, manager={ + "shard_python": "/serving/venv/bin/python"})) + + +@pytest.mark.parametrize("manager_over, message", [ + ({"gateway_hosts": ["root@gpu-a", "root@gpu-a"]}, "unique"), + ({"gateway_hosts": []}, "nonempty list"), + ({"gateway_hosts": "root@gpu-a"}, "nonempty list"), + ({"gateway_hosts": [""]}, "nonempty list"), + ({"gateway_hosts": ["gpu a"]}, "nonempty list"), + ({"gateway_hosts": ["gpu-a"], "gateway_routing": "modulo"}, + "auto routing"), +]) +def test_deployment_gateway_hosts_fail_closed(tmp_path, manager_over, message): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import DeploymentError, validate + + with pytest.raises(DeploymentError, match=message): + validate(deployment_mapping(tmp_path, manager=manager_over)) + + +def test_deployment_gateway_hosts_require_a_service_mode_manager(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import DeploymentError, validate + + values = deployment_mapping(tmp_path, manager={"gateway_hosts": ["gpu-a"]}) + values["manager"]["mode"] = "embedded" + with pytest.raises(DeploymentError, match="service-mode"): + validate(values) + + +def test_rollout_environment_carries_gateway_hosts(tmp_path): + overlay = _auto_overlay(tmp_path, {"gateway_hosts": ["root@gpu-a"]}) + assert overlay["HARBORRL_NATIVE_GATEWAY_HOSTS"] == '["root@gpu-a"]' + # Remotely spawned shards authenticate on the shared run token. + assert overlay["HARBORRL_NATIVE_SHARD_TOKEN_FILE"] + assert "HARBORRL_NATIVE_SHARD_PYTHON" not in overlay + + +def test_rollout_environment_carries_shard_python(tmp_path): + overlay = _auto_overlay(tmp_path, { + "gateway_hosts": ["root@gpu-a"], + "shard_python": "/serving/venv/bin/python"}) + assert overlay["HARBORRL_NATIVE_SHARD_PYTHON"] == "/serving/venv/bin/python" + + +def test_rollout_environment_omits_gateway_hosts_when_undeclared(tmp_path): + overlay = _auto_overlay(tmp_path, {}) + assert "HARBORRL_NATIVE_GATEWAY_HOSTS" not in overlay diff --git a/tests/harborrl/test_gateway_shards.py b/tests/harborrl/test_gateway_shards.py new file mode 100644 index 000000000..bc58d6aa2 --- /dev/null +++ b/tests/harborrl/test_gateway_shards.py @@ -0,0 +1,356 @@ +"""CPU contract evidence for static gateway sharding (same-process form). + +These tests exercise the shard table, static routing and the cross-shard +admission barrier without SSH, GPU or a live rollout service: shard ports +and trace subtrees stay exclusive, routing never falls back to another +shard's admission, and one pending generation on ANY shard blocks the next +policy epoch on ALL of them. +""" +from __future__ import annotations + +import json +from types import SimpleNamespace + +import pytest + + +def shard_environment(tmp_path, **extra): + environment = { + "RUN_ID": "run-1", "RUN_DIR": str(tmp_path), + "HARBORRL_NATIVE_CATALOG": str(tmp_path / "catalog.json"), + "HARBORRL_NATIVE_PROFILE": str(tmp_path / "profile.json"), + "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, + "HARBORRL_NATIVE_TEMPLATE_DIGEST": "1" * 64, + "HARBORRL_NATIVE_GATEWAY_HOST": "127.0.0.1", + "HARBORRL_NATIVE_GATEWAY_PORT": "49184", + "HARBORRL_NATIVE_GATEWAY_URL": "http://127.0.0.1:49184", + "HARBORRL_NATIVE_RUNNER_PYTHON": "/python", + "HARBORRL_NATIVE_RUNNER_WORKERS": json.dumps(["root@worker-a"]), + "ROLLOUT_MAX_RESPONSE_LEN": "128", + "ROLLOUT_MAX_CONTEXT_LEN": "1024", + } + environment.update(extra) + return environment + + +def shard_config(tmp_path, monkeypatch, **extra): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment(tmp_path, **extra) + for key, value in environment.items(): + monkeypatch.setenv(key, value) + return RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + + +def shard_identity(version, slot): + from harborrl.trajectories.native import Identity + + return Identity(run_id="run-1", batch_id="rollout-0", group_id="group-0", + slot_id=slot, attempt_id=f"attempt-{version}-{slot}", + trajectory_id=f"trajectory-{version}-{slot}", + task_digest="d" * 64, policy_version=version, + harness_digest="h" * 64, reward_digest="r" * 64, + sampling_digest="s" * 64) + + +# -- shard table -------------------------------------------------------------- + + +def test_config_parses_the_static_shard_table(tmp_path, monkeypatch): + from harborrl.rollout.runtime import GatewayShardSpec + + # An absent table keeps the single implicit gateway, now routed "auto": + # one gateway until the live sandbox count passes sandboxes_per_shard. + legacy = shard_config(tmp_path, monkeypatch) + assert legacy.shards == () and legacy.gateway_routing == "auto" + + config = shard_config(tmp_path, monkeypatch, HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps( + [{"name": "gw0"}, {"name": "gw1", "port": 49300}])) + assert config.shards == ( + GatewayShardSpec(name="gw0", port=49184, url="http://127.0.0.1:49184"), + GatewayShardSpec(name="gw1", port=49300, url="http://127.0.0.1:49300")) + # The rendered environment round-trips the resolved ports. + assert json.loads(config.environment()["HARBORRL_NATIVE_GATEWAY_SHARDS"]) == [ + {"name": "gw0", "port": 49184}, {"name": "gw1", "port": 49300}] + + +@pytest.mark.parametrize("shards, message", [ + (json.dumps([{"name": "gw0"}, {"name": "gw0"}]), "unique"), + (json.dumps([{"name": "gw0", "port": 49184}, {"name": "gw1", "port": 49184}]), + "unique integers"), + (json.dumps([{"name": "GW0"}]), "lowercase slugs"), + (json.dumps([{"name": "gw0", "extra": 1}]), "name, port"), + (json.dumps([]), "nonempty"), + ("{not-json", "invalid HARBORRL_NATIVE_GATEWAY_SHARDS"), +]) +def test_shard_table_fails_closed(tmp_path, monkeypatch, shards, message): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment(tmp_path, HARBORRL_NATIVE_GATEWAY_SHARDS=shards) + with pytest.raises(ValueError, match=message): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + + +@pytest.mark.parametrize("routing, worker_shards, message", [ + ("worker_affinity", {}, "missing"), + ("worker_affinity", {"root@worker-a": "gwX"}, "unknown gateway shards"), + ("worker_affinity", {"root@other": "gw0"}, "unknown runners"), + ("round_robin", {}, "gateway_routing"), +]) +def test_worker_affinity_routing_table_fails_closed( + tmp_path, monkeypatch, routing, worker_shards, message): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment( + tmp_path, + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([{"name": "gw0"}, {"name": "gw1"}]), + HARBORRL_NATIVE_GATEWAY_ROUTING=routing, + HARBORRL_NATIVE_GATEWAY_WORKER_SHARDS=json.dumps(worker_shards)) + with pytest.raises(ValueError, match=message): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + + +# -- runtime construction ----------------------------------------------------- + + +def test_runtime_splits_ports_and_trace_trees_across_shards(tmp_path, monkeypatch): + from harborrl.gateway import server as gateway_server + from harborrl.gateway import sglang_backend, trace + from harborrl.rollout import runtime as rollout_runtime + + (tmp_path / "catalog.json").write_text("[]") + (tmp_path / "profile.json").write_text(json.dumps({"model": "policy"})) + tokenizer = SimpleNamespace( + apply_chat_template=lambda *args, **kwargs: [1, 2], + get_vocab=lambda: {"token": 1}, + ) + tokenizer_digest, template_digest = rollout_runtime.tokenizer_fingerprint(tokenizer) + environment = shard_environment( + tmp_path, + HARBORRL_NATIVE_TOKENIZER_DIGEST=tokenizer_digest, + HARBORRL_NATIVE_TEMPLATE_DIGEST=template_digest, + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([{"name": "gw0"}, {"name": "gw1", "port": 49300}])) + for key, value in environment.items(): + monkeypatch.setenv(key, value) + + registries, servers, gateways = [], [], [] + + class FakeGateway: + def __init__(self, registry, backend, *, model, max_output_tokens): + gateways.append((registry, backend)) + + class FakeServer: + def serve_forever(self): + pass + + def shutdown(self): + pass + + def server_close(self): + pass + + class FakeTraceRegistry: + version = None + accepting = True + + def __init__(self, path): + registries.append(path) + + class FakeBackend: + def __init__(self, *args, **kwargs): + pass + + def weight_version(self): + return "7" + + def fake_make_server(gateway, host, port, *, max_inflight=None): + servers.append((host, port, max_inflight)) + return FakeServer() + + monkeypatch.setattr(gateway_server, "Gateway", FakeGateway) + monkeypatch.setattr(gateway_server, "make_server", fake_make_server) + monkeypatch.setattr(sglang_backend, "SGLangBackend", FakeBackend) + monkeypatch.setattr(trace, "TraceRegistry", FakeTraceRegistry) + + config = rollout_runtime.RolloutRuntimeConfig.from_environment( + environment, endpoint="http://127.0.0.1:30000", concurrency_limit=2, group_size=2) + runtime_ = rollout_runtime.RolloutRuntime(config, tokenizer) + try: + assert [shard.name for shard in runtime_.shards] == ["gw0", "gw1"] + # Exclusive ports and per-shard trace subtrees. + assert servers == [("127.0.0.1", 49184, 4), ("127.0.0.1", 49300, 4)] + assert registries == [runtime_.run_root / "native-traces" / "shard-gw0", + runtime_.run_root / "native-traces" / "shard-gw1"] + # One shared serving backend; registries stay per shard. + assert gateways[0][1] is gateways[1][1] is runtime_.backend + assert gateways[0][0] is not gateways[1][0] + # The service-facing admission epoch spans every shard. + assert isinstance(runtime_.registry, rollout_runtime.ShardAdmission) + assert runtime_.registry.shards == runtime_.shards + assert runtime_.registry.version is None and runtime_.registry.accepting + evidence = json.loads((runtime_.run_root / "gateway-shards.json").read_text()) + assert [(s["name"], s["port"], s["url"]) for s in evidence["shards"]] == [ + ("gw0", 49184, "http://127.0.0.1:49184"), + ("gw1", 49300, "http://127.0.0.1:49300")] + assert evidence["gateway_routing"] == "modulo" + finally: + runtime_.close() + + +# -- static routing ----------------------------------------------------------- + + +def test_modulo_routing_spreads_attempts_like_worker_rotation(): + from harborrl.rollout.runtime import route_shard + + config = SimpleNamespace(gateway_routing="modulo", worker_shards={}) + shards = [SimpleNamespace(name=f"gw{index}") for index in range(3)] + picked = [route_shard(config, shards, runner_host=None, slot=slot, + attempt_index=attempt).name + for slot in range(4) for attempt in range(2)] + assert picked == ["gw0", "gw1", "gw1", "gw2", "gw2", "gw0", "gw0", "gw1"] + + +def test_worker_affinity_routing_is_static_and_fails_closed(): + from harborrl.rollout.runtime import route_shard + + config = SimpleNamespace(gateway_routing="worker_affinity", + worker_shards={"root@worker-a": "gw0", "local": "gw1"}) + gw0, gw1 = SimpleNamespace(name="gw0"), SimpleNamespace(name="gw1") + shards = (gw0, gw1) + assert route_shard(config, shards, runner_host="root@worker-a", + slot=0, attempt_index=0) is gw0 + # Retries of the same slot keep their worker's shard. + assert route_shard(config, shards, runner_host="root@worker-a", + slot=5, attempt_index=1) is gw0 + # The local runner transport pins to its own shard. + assert route_shard(config, shards, runner_host=None, slot=1, attempt_index=0) is gw1 + unmapped = SimpleNamespace(gateway_routing="worker_affinity", + worker_shards={"local": "gw1"}) + with pytest.raises(ValueError, match="no gateway shard for runner"): + route_shard(unmapped, shards, runner_host="root@worker-b", slot=0, attempt_index=0) + unknown = SimpleNamespace(gateway_routing="worker_affinity", + worker_shards={"root@worker-a": "gwX"}) + with pytest.raises(ValueError, match="unknown gateway shard"): + route_shard(unknown, shards, runner_host="root@worker-a", slot=0, attempt_index=0) + # A single shard needs no routing table. + assert route_shard(unmapped, shards[:1], runner_host="anywhere", + slot=0, attempt_index=0) is gw0 + + +# -- cross-shard version barrier ---------------------------------------------- + + +def test_pending_generation_on_any_shard_blocks_the_whole_epoch(tmp_path): + from harborrl.gateway.trace import TraceRegistry + from harborrl.rollout.runtime import ShardAdmission + + registries = [TraceRegistry(tmp_path / f"shard-{index}") for index in range(2)] + admission = ShardAdmission([SimpleNamespace(registry=r) for r in registries]) + first = registries[0].register(shard_identity("1", "0"), "trial-0") + second = registries[1].register(shard_identity("1", "1"), "trial-1") + aid0, request0, *_ = registries[0].begin(first) + registries[0].finish(aid0, request0, error="interrupted_by_drain") + aid1, request1, *_ = registries[1].begin(second) + turn = {"request_id": request1, "response_id": "resp-1", + "response": {"content": [{"type": "text", "text": "pending"}]}} + # One in-flight generation on shard 1 keeps admission closed everywhere. + with pytest.raises(TimeoutError): + admission.drain(0.05) + assert not admission.accepting + assert admission.version == "1" + # Settle, seal and reopen: the next epoch lands on every shard at once. + registries[1].finish(aid1, request1, turn) + registries[1].delivery(aid1, "resp-1", "sent") + registries[1].consume(aid1, "resp-1", turn["response"]["content"], "session#1") + for registry, aid in ((registries[0], aid0), (registries[1], aid1)): + registry.close(aid) + registry.seal(aid) + states = [{"endpoint": "http://router:30000", "healthy": True, "weight_version": "2"}] + admission.drain(1.0) + admission.publish_version(states, "2") + assert admission.version == "2" and admission.accepting + for index, registry in enumerate(registries): + assert registry.version == "2" and registry.accepting + # The old epoch never admits new attempts on any shard. + with pytest.raises(ValueError, match="version mismatch"): + registry.register(shard_identity("1", f"old-{index}"), f"trial-old-{index}") + + +def test_service_activation_spans_every_shard(tmp_path): + from harborrl.gateway.trace import TraceRegistry + from harborrl.rollout.runtime import ShardAdmission + from harborrl.rollout.service import RolloutService + + registries = [TraceRegistry(tmp_path / f"shard-{index}") for index in range(2)] + states = [{"endpoint": "http://router:30000", "healthy": True, "weight_version": "2"}] + + class Pool: + @staticmethod + def pool_states(): + return [dict(item) for item in states] + + runtime_ = SimpleNamespace( + registry=ShardAdmission([SimpleNamespace(registry=r) for r in registries]), + backend=Pool()) + service = RolloutService(run_root=tmp_path) + try: + service.register_run("run-1", {"run_id": "run-1"}) + service.attach_runtime("run-1", SimpleNamespace(drain_timeout=1.0), runtime_) + result = service.activate_policy("run-1", states, "2") + assert result["admission"] == "open" + for registry in registries: + assert registry.version == "2" and registry.accepting + assert json.loads( + (tmp_path / "policy-pool.json").read_text())["weight_version"] == "2" + finally: + service.close() + + +# -- deployment schema -------------------------------------------------------- + + +def test_deployment_accepts_a_static_shard_table(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_shards": [{"name": "gw0"}, {"name": "gw1", "port": 49300}], + "gateway_routing": "worker_affinity", + "worker_shards": {"worker-a": "gw0", "worker-b": "gw1"}})) + manager = config["manager"] + assert manager["gateway_shards"] == [{"name": "gw0"}, {"name": "gw1", "port": 49300}] + assert manager["gateway_routing"] == "worker_affinity" + assert manager["worker_shards"] == {"worker-a": "gw0", "worker-b": "gw1"} + # Legacy deployments keep their exact historical shape. + legacy = validate(deployment_mapping(tmp_path))["manager"] + assert not ({"gateway_shards", "gateway_routing", "worker_shards"} & set(legacy)) + + +@pytest.mark.parametrize("manager_over, message", [ + ({"gateway_shards": []}, "nonempty"), + ({"gateway_shards": [{"name": "gw0"}, {"name": "gw0"}]}, "duplicate"), + ({"gateway_shards": [{"name": "gw0", "extra": 1}]}, "name"), + ({"gateway_shards": [{"name": "GW0"}]}, "lowercase slugs"), + ({"gateway_shards": [{"name": "gw0", "port": 49200}]}, "control_port"), + ({"gateway_shards": [{"name": "gw0", "port": 1}, {"name": "gw1", "port": 1}]}, "duplicate"), + ({"gateway_routing": "round_robin"}, "gateway_routing"), + ({"worker_shards": {"worker-a": "gw0"}}, "requires manager.gateway_shards"), + ({"gateway_shards": [{"name": "gw0"}], + "worker_shards": {"worker-a": "gwX"}}, "unknown shards"), + ({"gateway_shards": [{"name": "gw0"}], + "worker_shards": {"worker-c": "gw0"}}, "unknown workers"), + ({"gateway_shards": [{"name": "gw0"}], "gateway_routing": "worker_affinity"}, "missing"), + ({"gateway_routing": "worker_affinity"}, "requires manager.gateway_shards"), +]) +def test_deployment_shard_routing_fails_closed(tmp_path, manager_over, message): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import DeploymentError, validate + + with pytest.raises(DeploymentError, match=message): + validate(deployment_mapping(tmp_path, manager=manager_over)) diff --git a/tests/harborrl/test_math.py b/tests/harborrl/test_math.py new file mode 100644 index 000000000..97c68e521 --- /dev/null +++ b/tests/harborrl/test_math.py @@ -0,0 +1,113 @@ +"""Math data boundaries, sampled OPD mathematics and process-reward alignment.""" +import json + +import pytest + +from harborrl.config import ROOT +from harborrl.math.data import import_samples, load_tasks, verify + + +@pytest.mark.parametrize("response,reference,status", [("2", "2", "correct"), ("#### 1/2", "0.5", "correct"), + (r"reasoning \boxed{3}", "3", "correct"), ("2", "3", "incorrect"), ("2 or 3", "3", "unparseable"), + ("__import__('os').system('echo unsafe')", "0", "unparseable"), ("1/0", "1", "unparseable"), ("x" * 20000, "0", "unsupported")]) +def test_rational_verifier(response, reference, status): + assert verify(response, reference)["status"] == status + + +@pytest.mark.parametrize("response,reference,status", [ + (r"#### Final Answer: $\boxed{57,500}$", "57500", "correct"), + ("work\n#### **42**\nMore explanation.", "42", "correct"), + (r"\boxed{14} years ago.", "14", "correct"), + (r"\boxed{12} corrected to \boxed{13}", "13", "correct"), + (r"\boxed{1/2}", "0.5", "correct"), + (r"\boxed{12,34}", "1234", "unparseable"), + (r"\boxed{1/0}", "1", "unparseable"), + ("The answer could be 2 or 3", "3", "unparseable"), + (r"\boxed{2}", "3", "incorrect"), +]) +def test_gsm8k_numeric_verifier(response, reference, status): + result = verify(response, reference, "gsm8k_numeric_v1") + assert result["status"] == status + assert result["verifier"] == "gsm8k_numeric_v1" + assert result["reward"] == (1.0 if status == "correct" else 0.0) + + +def test_verifier_version_is_locked_in_task_import(tmp_path): + row = {**load_tasks(ROOT / "examples/math/tasks.jsonl")[0], "verifier": "gsm8k_numeric_v1"} + tasks = tmp_path / "tasks.jsonl" + tasks.write_text(json.dumps(row)) + samples = tmp_path / "samples.jsonl" + samples.write_text(json.dumps({"task_id": row["task_id"], "response": r"\boxed{999} done", + "source": "fixture", "license": "MIT"})) + output = tmp_path / "out" + import_samples(tasks, samples, output) + result = json.loads(next((output / "verification").glob("*.json")).read_text()) + assert result["verifier"] == "gsm8k_numeric_v1" and result["reward"] == 0.0 + assert verify(r"\boxed{2} done", "2")["status"] == "unparseable" + with pytest.raises(ValueError, match="Unsupported numeric verifier"): + verify("2", "2", "unknown") + + +def test_split_lock_and_curated_import(tmp_path): + source = ROOT / "examples/math/tasks.jsonl" + tasks = load_tasks(source) + duplicate = {**tasks[0], "task_id": "different", "split": "test"} + leaked = tmp_path / "leaked.jsonl" + leaked.write_text("\n".join(json.dumps(row) for row in [tasks[0], duplicate])) + with pytest.raises(ValueError, match="leakage"): + load_tasks(leaked) + samples = tmp_path / "curated.jsonl" + samples.write_text(json.dumps({"task_id": tasks[0]["task_id"], "response": "2", "source": "human-reviewed", "license": "MIT"})) + output = tmp_path / "dataset" + manifest = import_samples(source, samples, output) + sample = json.loads(next((output / "samples").glob("*.json")).read_text()) + assert sample["evidence_kind"] == "declared" and sample["token_evidence"] is None + assert manifest["training_eligibility"] == "not_an_on_policy_batch" + with pytest.raises(ValueError, match="new directory"): + import_samples(source, samples, output) + + +def test_opd_expectation_gradient_and_detach(): + torch = pytest.importorskip("torch") + from harborrl.algorithms.opd import sampled_opd_loss + logits = torch.tensor([-.5, .2, .8], dtype=torch.float64, requires_grad=True) + teacher = torch.tensor([-.2, .4, -.1], dtype=torch.float64, requires_grad=True).log_softmax(0) + logq = logits.log_softmax(0) + exact = (logq.exp() * (logq - teacher.detach())).sum() + expected = torch.autograd.grad(exact, logits, retain_graph=True)[0] + surrogate = sum(logq[i].detach().exp() * sampled_opd_loss(logq[i].reshape(1, 1), + logq[i].reshape(1, 1), teacher[i].reshape(1, 1), torch.ones(1, 1)) for i in range(3)) + observed = torch.autograd.grad(surrogate, logits, retain_graph=True)[0] + torch.testing.assert_close(expected, observed) + assert observed.norm() > 0 + assert torch.autograd.grad(surrogate, teacher, allow_unused=True)[0] is None + + +def test_loss_mask_and_microbatch_reduction(): + torch = pytest.importorskip("torch") + from harborrl.algorithms.opd import trajectory_surrogate + current = torch.tensor([[-1., -2., 0.], [-3., 0., 0.], [-2., -3., -4.]], requires_grad=True) + advantages = torch.tensor([[1., 2., 9.], [3., 8., 8.], [1., -1., 2.]], requires_grad=True) + mask = torch.tensor([[1, 1, 0], [1, 0, 0], [1, 1, 1]]) + full = trajectory_surrogate(current, advantages, mask) + split = trajectory_surrogate(current[:2], advantages[:2], mask[:2]) * 2 / 3 + trajectory_surrogate(current[2:], advantages[2:], mask[2:]) / 3 + torch.testing.assert_close(full, split) + full.backward() + assert current.grad[0, 2] == 0 and advantages.grad is None + + +def test_prm_step_alignment_and_equation(): + from harborrl.math.scoring import process_rewards + + class Tokenizer: + def decode(self, tokens): + return bytes(tokens).decode("ascii") + + tokenizer = Tokenizer() + tokens = list(b"1 + 1 = 2\n2 * 3 = 7") + scored = process_rewards("equation", tokens, tokenizer) + assert [step["status"] for step in scored["steps"]] == ["correct", "incorrect"] + assert len(scored["token_returns"]) == len(tokens) + assert scored["token_returns"][0] == 0 and scored["token_returns"][-1] == -1 + fake = process_rewards("fake", tokens, tokenizer) + assert fake["evidence_kind"] == "synthetic" diff --git a/tests/harborrl/test_math_rollout.py b/tests/harborrl/test_math_rollout.py new file mode 100644 index 000000000..3874a2a08 --- /dev/null +++ b/tests/harborrl/test_math_rollout.py @@ -0,0 +1,145 @@ +"""Math evidence validation at the production backend boundary.""" +import asyncio +import sys +from copy import deepcopy +from types import SimpleNamespace + +import pytest + +from harborrl.backends.slime_v032 import math_rollout +from harborrl.backends.slime_v032.math_rollout import ( + build_batch, + convert_samples_to_train_data, +) +from harborrl.config import ROOT, launch_plan, load_config +from harborrl.math.data import load_tasks, verify +from harborrl.trajectories.native import digest + + +def samples(): + task = load_tasks(ROOT / 'examples/math/tasks.jsonl')[0] + output = [] + for index, response in enumerate([task['reference'], '999999']): + evidence = { + 'tokens': [1, 2, 3], 'logprobs': [-0.5], 'response': response, + 'response_length': 1, 'group_key': 'group', 'policy_version': '1', + 'trajectory_id': f'trajectory-{index}', 'rollout_id': 0, + 'task_digest': digest(task), 'verification': verify(response, task['reference']), + 'truncated': False, + } + output.append(SimpleNamespace( + tokens=[1, 2, 3], rollout_log_probs=[-0.5], response=response, response_length=1, + loss_mask=[1], index=index, + metadata={'math_task': task, 'task_digest': digest(task), 'math_evidence': evidence})) + return output + + +def test_math_neutral_batch_and_standard_fields(): + args = SimpleNamespace(n_samples_per_prompt=2) + batch = build_batch(args, samples()) + assert batch.policy_versions == {'1'} + assert [unit.advantage for unit in batch.groups[0].trajectories] == [1., -1.] + data = convert_samples_to_train_data(args, samples()) + assert data['rewards'] == [1., -1.] + assert data['loss_masks'] == [[1], [1]] + assert data['rollout_mask_sums'] == [1, 1] + + +@pytest.mark.parametrize('mutation', ['split', 'version', 'incomplete', 'tokens', 'reward']) +def test_math_rejects_invalid_evidence(mutation): + rows = deepcopy(samples()) + if mutation == 'split': + rows[0].metadata['math_task']['split'] = 'test' + elif mutation == 'version': + rows[0].metadata['math_evidence']['policy_version'] = '2' + elif mutation == 'incomplete': + rows.pop() + elif mutation == 'tokens': + rows[0].tokens = [2, 2, 3] + else: + rows[0].metadata['math_evidence']['verification']['reward'] = 9. + with pytest.raises(ValueError): + build_batch(SimpleNamespace(n_samples_per_prompt=2), rows) + + +def test_math_config_uses_shared_launcher_without_harbor(): + config = load_config(ROOT / 'examples/math/train.yaml') + assert 'harbor' not in config + plan = launch_plan(config) + assert plan['training_command'] == ['bash', str(ROOT / 'harborrl/platform/slime_train.sh')] + assert plan['environment']['HARBORRL_EXECUTION_BACKEND'] == 'text_math' + with pytest.raises(ValueError, match='max_attempts'): + load_config(ROOT / 'examples/math/train.yaml', ['sampling.max_attempts=2']) + + +def test_math_checkpoint_export_and_resume(tmp_path): + from harborrl.platform.evidence import run_environment + config = load_config(ROOT / 'examples/math/train.yaml', [f'training.resume={tmp_path.as_posix()}']) + plan = launch_plan(config) + assert plan['environment']['RESUME_LOAD'] == str(tmp_path.resolve()) + environment = run_environment(plan, tmp_path, tmp_path / 'tasks.json', tmp_path / 'prompts.jsonl') + assert environment['SAVE_HF'] == str(tmp_path / 'hf' / 'rollout-{rollout_id}') + with pytest.raises(ValueError, match='boolean'): + load_config(ROOT / 'examples/math/train.yaml', ['training.export_hf=1']) + + +def test_math_seed_is_explicit_and_validated(): + path = ROOT / 'examples/math/train.yaml' + assert 'HARBORRL_MATH_SEED' not in launch_plan(load_config(path))['environment'] + for seed in (0, 42, 2147483647): + plan = launch_plan(load_config(path, [f'training.seed={seed}'])) + assert plan['environment']['HARBORRL_MATH_SEED'] == str(seed) + for value in ('true', '-1', '2147483648', '1.5'): + with pytest.raises(ValueError, match='training.seed'): + load_config(path, [f'training.seed={value}']) + + +def test_math_batch_uses_task_verifier(): + rows = deepcopy(samples()) + for sample in rows: + task = sample.metadata['math_task'] + task['verifier'] = 'gsm8k_numeric_v1' + sample.metadata['task_digest'] = digest(task) + evidence = sample.metadata['math_evidence'] + evidence['task_digest'] = digest(task) + evidence['verification'] = verify(sample.response, task['reference'], task['verifier']) + batch = build_batch(SimpleNamespace(n_samples_per_prompt=2), rows) + assert [unit.advantage for unit in batch.groups[0].trajectories] == [0.5, -0.5] + + +@pytest.mark.parametrize('seed', [None, 0, 2147483647]) +def test_generation_seed_tracks_persistent_sample_index(monkeypatch, seed): + requests = [] + tokenizer = SimpleNamespace(encode=lambda *_args, **_kwargs: [10], + decode=lambda *_args, **_kwargs: '2') + monkeypatch.setitem(sys.modules, 'slime.utils.processing_utils', SimpleNamespace( + load_tokenizer=lambda *_args, **_kwargs: tokenizer)) + monkeypatch.delenv('HARBORRL_MATH_ALGORITHM', raising=False) + monkeypatch.delenv('HARBORRL_PROCESS_REWARD', raising=False) + monkeypatch.delenv('HARBORRL_MATH_SEED', raising=False) + if seed is not None: + monkeypatch.setenv('HARBORRL_MATH_SEED', str(seed)) + monkeypatch.setattr(math_rollout, '_pool', lambda _endpoint: ([], '0')) + + def request(_url, payload): + requests.append(deepcopy(payload)) + return {'text': '2', 'meta_info': {'weight_version': '0', + 'output_token_logprobs': [[-0.5, 20]], 'finish_reason': {'type': 'stop'}}} + + monkeypatch.setattr(math_rollout, '_request', request) + task = load_tasks(ROOT / 'examples/math/tasks.jsonl')[0] + indices = [37, 38] + group = [SimpleNamespace(index=index, prompt=task['prompt'], + metadata={'math_task': task, 'task_digest': digest(task)}, + Status=SimpleNamespace(COMPLETED='completed', TRUNCATED='truncated')) for index in indices] + args = SimpleNamespace(hf_checkpoint='unused', sglang_router_ip='localhost', sglang_router_port=1, + sglang_server_concurrency=1, rollout_max_response_len=10, + rollout_max_context_len=20, rollout_batch_size=1, rollout_seed=seed) + generated = asyncio.run(math_rollout._generate(args, 8, [group])) + assert len(generated) == 2 + assert [sample.metadata['math_evidence']['serving_input'] for sample in generated] == requests + if seed is None: + assert all('sampling_seed' not in payload['sampling_params'] for payload in requests) + else: + assert [payload['sampling_params']['sampling_seed'] for payload in requests] == [ + (seed + index) % 2**31 for index in indices] diff --git a/tests/harborrl/test_native_continuation.py b/tests/harborrl/test_native_continuation.py new file mode 100644 index 000000000..02dd5e8f3 --- /dev/null +++ b/tests/harborrl/test_native_continuation.py @@ -0,0 +1,135 @@ +"""Pinned CLI logical turns require exact consumed length-stop continuations.""" +import json +from copy import deepcopy + +import pytest + +from harborrl.rollout.harbor_job.audit import audit_termination +from harborrl.trajectories.native import digest + +CONTINUATION = ('Output token limit hit. Resume directly — no apology, no recap of what you were doing. ' + 'Pick up mid-thought if that is where the cut happened. Break remaining work into smaller pieces.') + + +def fixture(lengths): + messages = [{'role': 'user', 'content': [{'type': 'text', 'text': 'Inspect the file.'}]}] + turns = [] + for index in range(2 + len(lengths)): + length = index in lengths + content = ([{'type': 'text', 'text': 'Still inspecting'}] if length else + [{'type': 'tool_use', 'id': f'tool-{index}', 'name': 'Read', 'input': {'file_path': '/code.py'}}]) + turn = {'response_id': f'r{index}', 'finish_reason': 'max_tokens' if length else 'tool_use', + 'client_request': {'messages': deepcopy(messages)}, + 'response': {'content': content, 'stop_reason': 'max_tokens' if length else 'tool_use'}} + turns.append(turn) + messages.append({'role': 'assistant', 'content': deepcopy(content)}) + messages.append({'role': 'user', 'content': [{'type': 'text', 'text': CONTINUATION}] if length else + [{'type': 'tool_result', 'tool_use_id': f'tool-{index}', 'content': 'value = 1'}]}) + return turns + + +def audit(tmp_path, turns): + for turn in turns: + for key in ('client_request', 'response'): + turn.setdefault(key + '_digest', digest(turn[key])) + session = {'session_id': 'session', 'consumed_response_ids': [turn['response_id'] for turn in turns], + 'evidence_ref': 'native-session#sha256=diagnostic'} + events = [{'type': 'system', 'subtype': 'init', 'session_id': 'session', 'claude_code_version': '2.1.141'}] + events.extend({'type': 'assistant', 'session_id': 'session', 'message': {'id': turn['response_id']}} + for turn in turns) + events.append({'type': 'result', 'subtype': 'error_max_turns', 'is_error': True, + 'terminal_reason': 'max_turns', 'stop_reason': 'tool_use', 'num_turns': 3, + 'session_id': 'session', 'permission_denials': [], 'errors': ['Reached maximum number of turns (2)']}) + path, result = tmp_path / 'stream.jsonl', tmp_path / 'result.json' + path.write_text('\n'.join(map(json.dumps, events))) + result.write_text(json.dumps({'exception_info': {'exception_type': 'NonZeroAgentExitCodeError'}})) + return audit_termination(path, result, {'cli_version': '2.1.141', 'max_turns': 2}, session, turns=turns) + + +@pytest.mark.parametrize('lengths', [set(), {0}, {1}, {0, 1}, {0, 2}]) +def test_logical_turns_preserve_every_consumed_response(tmp_path, lengths): + turns = fixture(lengths) + # The pinned CLI moves ephemeral cache metadata and serializes older text as strings. + if lengths == {0, 1}: + turns[2]['client_request']['messages'][2]['content'] = CONTINUATION + for turn in turns: + turn['client_request']['messages'][-1]['content'][-1]['cache_control'] = {'type': 'ephemeral'} + snapshot = deepcopy(turns) + evidence = audit(tmp_path, turns) + assert evidence['logical_turns'] == evidence['max_turns'] == 2 + assert evidence['policy_response_count'] == len(turns) == 2 + len(lengths) + assert evidence['consumed_response_ids'] == [turn['response_id'] for turn in turns] + assert len(evidence['continuations']) == len(lengths) + assert [{key: value for key, value in turn.items() if not key.endswith('_digest')} + for turn in turns] == snapshot + for proof in evidence['continuations']: + index = int(proof['response_id'][1:]) + assert proof['next_response_id'] == turns[index + 1]['response_id'] + assert proof['continuation_request_digest'] == digest(turns[index + 1]['client_request']) + + +@pytest.mark.parametrize('mutation', [ + 'arbitrary_extra', 'missing_continuation', 'changed_continuation', 'continuation_role', + 'changed_history', 'missing_history', 'extra_history', 'changed_assistant', 'assistant_role', + 'extra_field', 'extra_block', 'unknown_cache', 'null_cache', 'malformed_message', + 'malformed_block', 'tool_length', 'empty_length', 'stop_reason', + 'end_turn', 'final_length', 'missing_logical_turn', 'extra_logical_turn', + 'request_digest', 'response_digest', 'next_digest', 'reordered', +]) +def test_unlinked_or_unprobed_extra_response_rejected(tmp_path, mutation): + turns = fixture({0}) + following = turns[1]['client_request']['messages'] + if mutation == 'arbitrary_extra': + turns[0]['finish_reason'] = 'tool_use' + elif mutation == 'missing_continuation': + following.pop() + elif mutation == 'changed_continuation': + following[-1]['content'][0]['text'] += ' modified' + elif mutation == 'continuation_role': + following[-1]['role'] = 'assistant' + elif mutation == 'changed_history': + following[0]['content'][0]['text'] += ' modified' + elif mutation == 'missing_history': + turns[0]['client_request'].pop('messages') + elif mutation == 'extra_history': + following.insert(0, {'role': 'user', 'content': 'injected'}) + elif mutation == 'changed_assistant': + following[-2]['content'][0]['text'] += ' modified' + elif mutation == 'assistant_role': + following[-2]['role'] = 'user' + elif mutation == 'extra_field': + following[-1]['injected'] = True + elif mutation == 'extra_block': + following[-1]['content'].append({'type': 'text', 'text': 'injected'}) + elif mutation == 'unknown_cache': + following[0]['content'][0]['cache_control'] = {'type': 'ephemeral', 'injected': True} + elif mutation == 'null_cache': + following[0]['content'][0]['cache_control'] = None + elif mutation == 'malformed_message': + following[0] = None + elif mutation == 'malformed_block': + turns[0]['response']['content'] = [None] + elif mutation == 'tool_length': + turns[0]['response']['content'] = turns[1]['response']['content'] + elif mutation == 'empty_length': + turns[0]['response']['content'] = [] + elif mutation == 'stop_reason': + turns[0]['response']['stop_reason'] = 'tool_use' + elif mutation == 'end_turn': + turns[0]['finish_reason'] = 'end_turn' + elif mutation == 'final_length': + turns[-1]['finish_reason'] = 'max_tokens' + elif mutation == 'missing_logical_turn': + turns.pop() + elif mutation == 'extra_logical_turn': + turns.append({**deepcopy(turns[-1]), 'response_id': 'extra'}) + elif mutation == 'request_digest': + turns[0]['client_request_digest'] = 'conflicting' + elif mutation == 'response_digest': + turns[0]['response_digest'] = 'conflicting' + elif mutation == 'next_digest': + turns[1]['client_request_digest'] = 'conflicting' + else: + turns.reverse() + with pytest.raises(ValueError): + audit(tmp_path, turns) diff --git a/tests/harborrl/test_native_contracts.py b/tests/harborrl/test_native_contracts.py index b63851386..1767261f8 100644 --- a/tests/harborrl/test_native_contracts.py +++ b/tests/harborrl/test_native_contracts.py @@ -73,6 +73,16 @@ def test_roundtrip_and_synthetic_gate(tmp_path): publish(path, {**ir, "readiness": "RL_READY"}) +def test_unavailable_tool_cannot_enter_training_without_native_rejection(tmp_path): + ir = make_ir(tmp_path) + turn = ir['turns'][0] + turn['response']['content'] = [{'type': 'tool_use', 'id': 'u', 'name': 'Check', 'input': {}}] + turn['response_digest'] = digest(turn['response']) + ir['trace_seal']['turns_digest'] = digest(ir['turns']) + with pytest.raises(ValueError, match='unconfirmed_tool_rejection'): + require_ready(ir, allow_synthetic=True) + + @pytest.mark.parametrize("reward", [True, False, None, "1", float("nan"), float("inf"), -1, 2]) def test_invalid_reward(reward): with pytest.raises(ValueError): diff --git a/tests/harborrl/test_native_example_budget.py b/tests/harborrl/test_native_example_budget.py new file mode 100644 index 000000000..4fbe498f4 --- /dev/null +++ b/tests/harborrl/test_native_example_budget.py @@ -0,0 +1,25 @@ +"""Keep the native training example compatible with the CLI prompt budget.""" +from types import SimpleNamespace + +import pytest +import yaml + +from harborrl.config import ROOT +from harborrl.gateway.messages import ProtocolError +from harborrl.gateway.sglang_backend import SGLangBackend + + +def test_native_example_reserves_room_for_cli_system_and_tools(): + config = yaml.safe_load((ROOT / "examples/native/train_qwen_native.yaml").read_text()) + sampling = config["sampling"] + # A controlled CPU rendering floor based on the pinned CLI's measured prompt. + # This checks the example budget, not tokenizer/model/GPU compatibility. + model = SimpleNamespace(render_chat=lambda *_: [1] * 20000) + backend = SGLangBackend("http://unused", None, model, + tokenizer_digest="fixture", template_digest="fixture", + max_context=sampling["max_context"]) + converted = {"sampling": {"max_new_tokens": sampling["max_tokens"]}} + assert len(backend.prepare(converted)) == 20000 + backend.max_context = 8192 + with pytest.raises(ProtocolError, match="context budget exceeded"): + backend.prepare(converted) diff --git a/tests/harborrl/test_native_gateway.py b/tests/harborrl/test_native_gateway.py index 0055fe515..90d6c38a8 100644 --- a/tests/harborrl/test_native_gateway.py +++ b/tests/harborrl/test_native_gateway.py @@ -43,8 +43,6 @@ def setup(tmp_path): @pytest.mark.parametrize("change", [{"model": "external"}, {"thinking": {"type": "enabled"}}, {"temperature": .7}, {"top_p": .9}, {"top_k": 1}, {"max_tokens": True}, {"stream": 1}, - {"temperature": True}, {"top_p": True}, {"temperature": "1"}, {"top_p": "1"}, - {"temperature": float("nan")}, {"top_p": float("inf")}, {"messages": []}, {"messages": [{"role": "assistant", "content": [{"type": "image"}]}]}, {"tool_choice": {"type": "any"}}]) def test_unsupported_request(change): @@ -112,21 +110,13 @@ def test_model_output_never_repaired(text): parse_qwen_output(text, response_id="m", tools=[]) -@pytest.mark.parametrize("delivery", ["sent", "delivery_unknown"]) -def test_trace_requires_actual_consumption_and_is_immutable(tmp_path, delivery): +def test_trace_requires_actual_consumption_and_is_immutable(tmp_path): registry, _ident, token, gateway = setup(tmp_path) aid, message = gateway.generate(token, payload()) - registry.delivery(aid, message["id"], delivery) - unconsumed = registry.unconsumed_responses(aid) + registry.delivery(aid, message["id"], "sent") registry.close(aid) _turns, seal = registry.seal(aid) assert seal["errors"] == ["unconfirmed_consumption"] - assert registry.unconsumed_responses(aid) == unconsumed - if delivery == "delivery_unknown": - with pytest.raises(ValueError, match="delivery was not confirmed"): - registry.wait_delivered(aid, .1) - else: - registry.wait_delivered(aid, .1) with pytest.raises(ValueError, match="sealed"): registry.consume(aid, message["id"], message["content"], "session") with pytest.raises(PermissionError): @@ -147,67 +137,6 @@ def test_session_audit_covers_final_response(tmp_path): assert turns[0]["delivery"] == "consumed_confirmed" -def test_sealed_trace_releases_payload_and_remains_idempotent(tmp_path): - registry = TraceRegistry(tmp_path / "traces") - gateway = Gateway(registry, Backend(), model="policy") - for version in range(1, 5): - for slot in range(2): - ident = identity(f"{version}-{slot}", policy_version=str(version)) - token = registry.register(ident, "trial") - aid, message = gateway.generate(token, payload()) - registry.delivery(aid, message["id"], "sent") - registry.consume(aid, message["id"], message["content"], "session") - registry.close(aid) - turns, seal = registry.seal(aid) - persisted = (registry.attempts[aid]["root"] / "trace.json").read_bytes() - assert registry.attempts[aid]["turns"] == [] - assert registry.seal(aid) == (turns, seal) - turns[0]["input_ids"].append(999) - assert registry.seal(aid)[0][0]["input_ids"] == [1, 2] - assert (registry.attempts[aid]["root"] / "trace.json").read_bytes() == persisted - with pytest.raises(ValueError, match="registered"): - registry.register(ident, "other") - with pytest.raises(ValueError, match="unique"): - registry.register(identity(f"reuse-{version}-{slot}", policy_version=str(version)), - "other", credential=token) - registry.drain(.1) - next_version = str(version + 1) - registry.publish_version([{"healthy": True, "weight_version": next_version, - "endpoint": "engine"}], next_version) - assert len(registry.attempts) == 8 - assert all(not state["turns"] for state in registry.attempts.values()) - assert not registry.credentials - assert len(list((tmp_path / "traces").rglob("trace.json"))) == 8 - - -def test_seal_retains_payload_until_delivery_and_durable_write(tmp_path, monkeypatch): - from harborrl.gateway import trace - - registry, _, token, gateway = setup(tmp_path) - aid, message = gateway.generate(token, payload()) - registry.close(aid) - with pytest.raises(ValueError, match="drain"): - registry.seal(aid) - assert len(registry.attempts[aid]["turns"]) == 1 - registry.delivery(aid, message["id"], "sent") - original_publish = trace.publish - - def fail_trace(path, value): - if path.name == "trace.json": - raise OSError("disk unavailable") - return original_publish(path, value) - - with monkeypatch.context() as patch: - patch.setattr(trace, "publish", fail_trace) - with pytest.raises(OSError, match="disk unavailable"): - registry.seal(aid) - assert not registry.attempts[aid]["sealed"] - assert len(registry.attempts[aid]["turns"]) == 1 - turns, seal = registry.seal(aid) - assert len(turns) == 1 and seal["errors"] == ["unconfirmed_consumption"] - assert registry.attempts[aid]["turns"] == [] - - def test_session_audit_merges_split_assistant_blocks(tmp_path): class MultiBlockBackend(Backend): def generate(self, converted, ident, response_id): @@ -274,6 +203,56 @@ def test_pool_mismatch_keeps_admission_closed(tmp_path): assert not registry.accepting +def test_http_overload_rejects_before_any_attempt_state(tmp_path): + _registry, _ident, token, gateway = setup(tmp_path) + started, release = threading.Event(), threading.Event() + + class BlockingBackend(Backend): + def generate(self, converted, ident, response_id): + started.set() + release.wait(5) + return super().generate(converted, ident, response_id) + + gateway.backend = BlockingBackend() + try: + server = make_server(gateway, max_inflight=1) + except PermissionError: + pytest.skip("this CPU sandbox does not permit loopback socket binding") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + url = f"http://127.0.0.1:{server.server_port}/v1/messages" + opener = request.build_opener(request.ProxyHandler({})) + results, errors = [], [] + + def post(): + try: + with opener.open(request.Request(url, json.dumps(payload()).encode(), + {"x-api-key": token}), timeout=10) as response: + results.append(json.load(response)) + except Exception as exc: # noqa: BLE001 - test collector + errors.append(exc) + + try: + first = threading.Thread(target=post) + first.start() + assert started.wait(5), "first request never reached the backend" + with pytest.raises(error.HTTPError) as exc: + opener.open(request.Request(url, json.dumps(payload()).encode(), + {"x-api-key": token}), timeout=5) + assert exc.value.code == 503 + assert exc.value.read() # overloaded_error body is present + release.set() + first.join(10) + assert not errors and results and results[0]["role"] == "assistant" + # The rejected request never began a generation, so a retry is safe. + post() + assert len(results) == 2 + finally: + server.shutdown() + server.server_close() + thread.join(5) + + def test_http_sse_and_auth(tmp_path): _registry, _ident, token, gateway = setup(tmp_path) try: @@ -354,6 +333,69 @@ def test_qwen_finish_marker_maps_to_empty_end_turn(): assert output == [] +def test_model_adapter_preserves_unavailable_action_but_history_requires_native_error(): + from harborrl.gateway.messages import unavailable_tool_error + from harborrl.models.qwen import QwenModelAdapter + + tools = [{'name': 'Bash', 'input_schema': {'type': 'object'}}] + converted = convert(payload(tools=tools), model='policy') + text = '{"name":"Check","arguments":{"target":"solution.py"}}' + blocks = QwenModelAdapter().parse_output(text, response_id='r', tools=converted['tools']) + assert blocks == [{'type': 'tool_use', 'id': 'toolu_r_0', 'name': 'Check', + 'input': {'target': 'solution.py'}}] + history = [*payload()['messages'], {'role': 'assistant', 'content': blocks}, + {'role': 'user', 'content': [{'type': 'tool_result', 'tool_use_id': blocks[0]['id'], + 'content': unavailable_tool_error('Check'), 'is_error': True}]}] + result = convert(payload(tools=tools, messages=history), model='policy') + assert result['messages'][-1]['content'] == '[tool_error]\n' + unavailable_tool_error('Check') + history[-1]['content'][0]['is_error'] = False + with pytest.raises(ProtocolError, match='exact native rejection'): + convert(payload(tools=tools, messages=history), model='policy') + history[-1]['content'][0].update(is_error=True, content='command executed') + with pytest.raises(ProtocolError, match='exact native rejection'): + convert(payload(tools=tools, messages=history), model='policy') + + +@pytest.mark.parametrize('rejection', ['valid', 'absent', 'success', 'wrong_tool', 'duplicate']) +def test_unavailable_action_requires_audited_native_rejection(tmp_path, rejection): + from harborrl.gateway.messages import unavailable_tool_error + + class UnavailableBackend(Backend): + def generate(self, converted, ident, response_id): + turn = super().generate(converted, ident, response_id) + turn.update(finish_reason='tool_use', content=[{'type': 'tool_use', 'id': 'u1', + 'name': 'Check', 'input': {}}]) + return turn + + registry, _ident, token, gateway = setup(tmp_path) + gateway.backend = UnavailableBackend() + aid, message = gateway.generate(token, payload()) + registry.delivery(aid, message['id'], 'sent') + result = {'type': 'tool_result', 'tool_use_id': 'u1', 'content': unavailable_tool_error('Check'), + 'is_error': True} + if rejection == 'success': + result['is_error'] = False + if rejection == 'wrong_tool': + result['content'] = unavailable_tool_error('Other') + results = [] if rejection == 'absent' else [result] * (2 if rejection == 'duplicate' else 1) + path = tmp_path / 'session.jsonl' + path.write_text('\n'.join(map(json.dumps, [ + {'type': 'assistant', 'sessionId': 's', 'message': message}, + {'type': 'user', 'sessionId': 's', 'message': {'role': 'user', 'content': results}}]))) + if rejection == 'valid': + audit_session(path, registry, aid) + registry.close(aid) + turns, seal = registry.seal(aid) + assert seal['errors'] == [] + assert turns[0]['tool_rejection_ref'].startswith(str(path) + '#sha256=') + assert turns[0]['output_ids'] == [3] and turns[0]['logprobs'] == [-.2] + else: + with pytest.raises(ValueError, match='exact native rejection'): + audit_session(path, registry, aid) + registry.close(aid) + assert 'unconfirmed_tool_rejection' in registry.seal(aid)[1]['errors'] + + def test_sglang_generation_version_and_logprob_evidence(): from harborrl.gateway.sglang_backend import SGLangBackend class Tokenizer: @@ -421,8 +463,80 @@ def open(self, request, timeout): assert backend.weight_version() == "7" +def test_sglang_unavailable_tool_keeps_sampled_tokens_logprobs_and_text(): + from harborrl.gateway.sglang_backend import SGLangBackend + from harborrl.models.qwen import QwenModelAdapter + + text = '{"name":"Check","arguments":{}}' + tokenizer = SimpleNamespace(decode=lambda ids, **kwargs: text) + backend = SGLangBackend('http://engine', tokenizer, QwenModelAdapter(), + tokenizer_digest='tok', template_digest='tpl', audited_raw_logprobs=True) + converted = convert(payload(), model='policy') + output = {'text': text, 'output_ids': [13, 14], 'meta_info': { + 'weight_version': '1', 'output_token_logprobs': [[-.2, 13], [-.3, 14]], + 'finish_reason': {'type': 'stop'}}} + result = backend.decode(output, {'input_ids': [11, 12]}, converted, identity(), 'msg') + assert result['model_output_text'] == text + assert result['output_ids'] == [13, 14] and result['logprobs'] == [-.2, -.3] + assert result['finish_reason'] == 'tool_use' + assert result['unavailable_tool_calls'] == [{'id': 'toolu_msg_0', 'name': 'Check'}] + assert result['logprob_semantics'] == 'raw_model' + + +def test_sglang_request_seed_is_independent_of_dispatch_order_and_retry(): + from harborrl.gateway.sglang_backend import SGLangBackend + from harborrl.models.qwen import QwenModelAdapter + + sent = [] + + class Response: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def read(self): + return b'{}' + + class Opener: + def open(self, request, timeout): + sent.append(json.loads(request.data)) + return Response() + + backend = SGLangBackend('http://engine', SimpleNamespace(), QwenModelAdapter(), + tokenizer_digest='tok', template_digest='tpl', sampling_seed=0) + backend.opener = Opener() + backend.prepare = lambda converted: converted['input_ids'] + backend.decode = lambda output, serving, *args: serving + converted = {'input_ids': [11, 12], 'sampling': {'temperature': 1.0}} + first = backend.generate(converted, identity(), 'response-a') + backend.generate(converted, identity('1'), 'response-b') + retry = backend.generate(converted, identity(attempt_id='retry', run_id='other'), 'response-c') + assert first == retry + assert first['sampling_params']['sampling_seed'] != sent[1]['sampling_params']['sampling_seed'] + assert converted['sampling'] == {'temperature': 1.0} + for changed in (identity(batch_id='next'), identity(group_id='next'), identity(task_digest='next')): + assert backend.generate(converted, changed, 'response')['sampling_params'] != first['sampling_params'] + assert backend.generate({**converted, 'input_ids': [11, 13]}, identity(), 'response') != first + backend.sampling_seed = 1 + assert backend.generate(converted, identity(), 'response') != first + backend.sampling_seed = None + assert 'sampling_seed' not in backend.generate(converted, identity(), 'response')['sampling_params'] + + +@pytest.mark.parametrize('seed', [True, -1, 2**31, 1.5]) +def test_sglang_rejects_invalid_sampling_seed(seed): + from harborrl.gateway.sglang_backend import SGLangBackend + with pytest.raises(ValueError, match='sampling_seed'): + SGLangBackend('http://engine', None, None, tokenizer_digest='tok', + template_digest='tpl', sampling_seed=seed) + + def test_native_gateway_thread_does_not_own_asyncio_default_executor(): - from harborrl.rollout.native_generate import NativeRolloutRuntime + from types import SimpleNamespace + + from harborrl.rollout.runtime import GatewayShard, GatewayShardSpec started, stopped = threading.Event(), threading.Event() @@ -431,12 +545,13 @@ def serve_forever(self): started.set() stopped.wait(5) - runtime = object.__new__(NativeRolloutRuntime) - runtime.server = Server() - runtime._start_gateway() + shard = GatewayShard(GatewayShardSpec("gw0", 49184, "http://127.0.0.1:49184"), + registry=SimpleNamespace(), gateway=SimpleNamespace(), + server=Server(), trace_dir="unused") + shard.start() try: assert started.wait(1) - assert runtime.thread.daemon is True + assert shard.thread.daemon is True finally: stopped.set() - runtime.thread.join(1) + shard.thread.join(1) diff --git a/tests/harborrl/test_native_lifecycle.py b/tests/harborrl/test_native_lifecycle.py index 6813b1548..3193d6d74 100644 --- a/tests/harborrl/test_native_lifecycle.py +++ b/tests/harborrl/test_native_lifecycle.py @@ -12,7 +12,7 @@ from harborrl.rollout.harbor_job.coordinator import GroupSlots, run_attempt from harborrl.trajectories.native import RewardProfile -PROFILE = {"cli_version": "2.1.141", "model": "policy", "max_turns": 4, +PROFILE = {"cli_version": "2.1.0", "model": "policy", "max_turns": 4, "agent_timeout_sec": 60, "setup_timeout_sec": 60, "verifier_timeout_sec": 60} @@ -39,16 +39,11 @@ def test_profile_controls_reach_agent_and_host(): disabled = claude_config({**PROFILE, "disallowed_tools": "Agent,Task,WebSearch,WebFetch"}, "http://gateway:8000", "token") assert disabled["kwargs"]["disallowed_tools"] == "Agent,Task,WebSearch,WebFetch" - extended = claude_config({**PROFILE, "disallowed_tools": "WebSearch"}, - "http://gateway:8000", "token") - assert extended["kwargs"]["disallowed_tools"] == "Agent,Task,WebSearch" - for key in ("DISABLE_PROMPT_CACHING", "DISABLE_COMPACT", "DISABLE_AUTO_COMPACT"): - assert config["env"][key] == "1" - assert sanitized_runner_env({key: "0"})[key] == "1" + assert config["env"]["DISABLE_AUTO_COMPACT"] == "1" assert config["env"]["ANTHROPIC_MODEL"] == "policy" assert config["env"]["NO_PROXY"] == "gateway" assert config["env"]["no_proxy"] == "gateway" - assert config["kwargs"]["version"] == "2.1.141" + assert config["kwargs"]["version"] == "2.1.0" assert sanitized_runner_env({"PATH": "/bin", "ANTHROPIC_API_KEY": "external", "HTTPS_PROXY": "external"})["PATH"] == "/bin" assert "ANTHROPIC_API_KEY" not in sanitized_runner_env({"ANTHROPIC_API_KEY": "external"}) with pytest.raises(ValueError): @@ -57,6 +52,54 @@ def test_profile_controls_reach_agent_and_host(): claude_config(PROFILE, "http://gateway/v1", "token") +def test_explicit_harness_prompt_preserves_literal_content(): + prompt = "Use the declared tools.\nFinish with a final response; preserve 'quotes' and $literal." + config = claude_config({**PROFILE, "append_system_prompt": prompt}, "http://gateway", "token") + assert config["kwargs"]["append_system_prompt"] == prompt + assert "append_system_prompt" not in claude_config(PROFILE, "http://gateway", "token")["kwargs"] + + +@pytest.mark.parametrize('base', [None, 300, 600.0, 1200]) +def test_environment_timeout_is_separate_from_agent_setup(base): + from harborrl.rollout.harbor_job.bindings import environment_timeouts + + resources = {} if base is None else {'build_timeout_sec': base} + original = dict(resources) + assert environment_timeouts(PROFILE, resources) == {} + profile = {**PROFILE, 'environment_timeout_sec': 720} + config = claude_config(profile, 'http://gateway', 'token') + assert config['override_setup_timeout_sec'] == 60 + assert 'environment_timeout_sec' not in config['kwargs'] + multiplier = environment_timeouts(profile, resources)['environment_build_timeout_multiplier'] + assert multiplier * (600 if base is None else base) == pytest.approx(720) + assert resources == original + + +@pytest.mark.parametrize('budget', [None, False, 0, -1, 1.5, '720']) +def test_invalid_environment_timeout_is_rejected(budget): + from harborrl.rollout.harbor_job.bindings import environment_timeouts + + profile = {**PROFILE, 'environment_timeout_sec': budget} + with pytest.raises(ValueError, match='environment_timeout_sec'): + claude_config(profile, 'http://gateway', 'token') + with pytest.raises(ValueError, match='environment_timeout_sec'): + environment_timeouts(profile, {}) + + +@pytest.mark.parametrize('base', [None, False, 0, -1, float('nan'), float('inf'), '300']) +def test_invalid_task_environment_timeout_cannot_be_overridden(base): + from harborrl.rollout.harbor_job.bindings import environment_timeouts + + with pytest.raises(ValueError, match='task environment'): + environment_timeouts({**PROFILE, 'environment_timeout_sec': 720}, {'build_timeout_sec': base}) + + +@pytest.mark.parametrize("prompt", [None, False, 1, "", " \n", "text\x00text", ["text"]]) +def test_invalid_harness_prompt_is_rejected(prompt): + with pytest.raises(ValueError, match="append_system_prompt"): + claude_config({**PROFILE, "append_system_prompt": prompt}, "http://gateway", "token") + + def test_retry_budget_and_new_lineage(): group = GroupSlots([identity(), identity("1")], max_attempts=2) first = group.start("0") @@ -129,46 +172,55 @@ async def fake_spawn(*args, **kwargs): assert registry.attempts[identity().attempt_id]["closed"] assert (root / "failure.json").is_file() if mode != "wrong-event": - marker = json.loads((root / "cleanup-required.json").read_text()) - assert marker["identity"] == identity().to_dict() - assert marker["trial_id"] == "trial" - assert marker["runner_host"] is None + assert (root / "cleanup-required.json").is_file() assert json.loads((root / "runner-events.json").read_text())[0]["type"] == "prepared" -@pytest.mark.parametrize("runner_host", [None, "runner-host"]) -def test_cleanup_timeout_preserves_runner_marker_and_original_error(tmp_path, monkeypatch, runner_host): - worker = tmp_path / "worker.py" - worker.write_text('''import json, sys, time -from pathlib import Path -req = json.loads(sys.stdin.readline()) -print(json.dumps({"type": "prepared", "trial_id": "trial"}), flush=True) -sys.stdin.readline() -output = Path(req["output"]) -output.mkdir(parents=True) -(output / "cleanup-required.json").write_text('{"reason": "deletion failed"}') -print(json.dumps({"type": "event", "trial_id": "other"}), flush=True) -time.sleep(30) -''') +@pytest.mark.parametrize('mode,existing_marker', [('terminal', False), ('crash', False), ('crash', True)]) +def test_cleanup_marker_preserves_original_failure(tmp_path, monkeypatch, mode, existing_marker): + from harborrl.trajectories.native import digest, publish, read_json + + root = tmp_path / 'attempt' + marker = root / 'cleanup-required.json' + sandbox = {'create_timeout_s': 1, 'cleanup_timeout_s': 1} + worker = tmp_path / 'worker.py' + terminal = {'type': 'terminal', 'trial_id': 'trial'} + worker.write_text(chr(10).join(['import sys, json', 'json.loads(sys.stdin.readline())', + 'print(json.dumps({"type": "prepared", "trial_id": "trial"}), flush=True)', + 'sys.stdin.readline()', 'sys.exit(7)' if mode == 'crash' else + f'print({json.dumps(terminal)!r}, flush=True)', ''])) spawn = asyncio.create_subprocess_exec async def fake_spawn(*args, **kwargs): - return await spawn(sys.executable, "-S", str(worker), **kwargs) - - monkeypatch.setattr(asyncio, "create_subprocess_exec", fake_spawn) - root = tmp_path / "attempt" - with pytest.raises(ValueError, match="event trial identity mismatch"): + return await spawn(sys.executable, '-S', str(worker), **kwargs) + + monkeypatch.setattr(asyncio, 'create_subprocess_exec', fake_spawn) + registry = TraceRegistry(tmp_path / 'traces') + register = registry.register + prior = {'identity': identity().to_dict(), 'reason': 'Existing cleanup evidence'} + + def registered(*args, **kwargs): + publish(root / 'runner/cleanup-required.json', {'reason': 'Unconfirmed platform deletion'}) + publish(root / 'runner/sandbox-lifecycle.json', + {'state': 'delete_failed', 'runtime_digest': digest(sandbox)}) + if existing_marker: + publish(marker, prior) + return register(*args, **kwargs) + + monkeypatch.setattr(registry, 'register', registered) + reason = 'sandbox deletion unconfirmed' if mode == 'terminal' else 'Runner exited before terminal result' + with pytest.raises(RuntimeError, match=reason): asyncio.run(run_attempt(identity(), task_path=tmp_path, profile=PROFILE, - gateway_url="http://gateway", runner_python=sys.executable, root=root, - registry=TraceRegistry(tmp_path / "traces"), reward_profile=RewardProfile(), - timeout=2, cleanup_timeout=.1, runner_host=runner_host)) - marker = json.loads((root / "cleanup-required.json").read_text()) - assert marker["reason"] == "Inspect runner/cleanup-required.json" - assert marker["identity"] == identity().to_dict() - assert marker["trial_id"] == "trial" - assert marker["runner_host"] == runner_host - assert json.loads((root / "failure.json").read_text())["error_type"] == "ValueError" - assert (root / "runner-events.json").is_file() + gateway_url='http://gateway', runner_python=sys.executable, root=root, + registry=registry, reward_profile=RewardProfile(), sandbox=sandbox)) + expected = prior if existing_marker else {'identity': identity().to_dict(), 'reason': + 'Sandbox deletion unconfirmed' if mode == 'terminal' else 'Inspect runner/cleanup-required.json'} + assert read_json(marker) == expected + assert read_json(root / 'failure.json')['error_type'] == 'RuntimeError' + assert read_json(root / 'runner/cleanup-required.json') == {'reason': 'Unconfirmed platform deletion'} + assert registry.attempts[identity().attempt_id]['closed'] + assert (root / 'runner-events.json').exists() + assert not (root / 'trajectory.v2.json').exists() def test_offline_http_to_export_is_explicitly_synthetic(tmp_path): diff --git a/tests/harborrl/test_native_terminal.py b/tests/harborrl/test_native_terminal.py index ae7f3767a..52d0181d3 100644 --- a/tests/harborrl/test_native_terminal.py +++ b/tests/harborrl/test_native_terminal.py @@ -1,7 +1,10 @@ """Regression coverage for the full coordinator success path and empty finals.""" import asyncio import json +import os import sys +from copy import deepcopy +from pathlib import Path import pytest from test_native_contracts import identity, make_ir @@ -10,7 +13,7 @@ from harborrl.gateway.trace import TraceRegistry from harborrl.harnesses.claude_code import ClaudeCodeHarnessAdapter from harborrl.rollout.harbor_job.coordinator import run_attempt -from harborrl.trajectories.native import RewardProfile, digest +from harborrl.trajectories.native import RewardProfile, digest, require_ready def test_casefold_tool_alias_must_be_unique(): @@ -21,19 +24,39 @@ def test_casefold_tool_alias_must_be_unique(): {"function": {"name": "bash"}}]) -@pytest.mark.parametrize("empty", [False, True]) -def test_coordinator_publishes_complete_success(tmp_path, monkeypatch, empty): +@pytest.mark.parametrize("empty,budget,reward,continuations", [ + (False, False, 1, 0), (True, False, 1, 0), (False, True, 0, 0), (False, True, 1, 0), + (False, True, 0, 1), (False, True, 1, 1), (False, True, 0, 2), +]) +def test_coordinator_publishes_complete_success(tmp_path, monkeypatch, empty, budget, reward, continuations): ident = identity() trial_id = "trial-" + ident.attempt_id root = tmp_path / "attempt" trial = root / "runner/harbor/trial" - template = make_ir(tmp_path / "fixture", nturns=2) + template = make_ir(tmp_path / "fixture", nturns=2 + continuations) + profile = {**PROFILE, 'cli_version': '2.1.141', 'max_turns': 2} if budget else PROFILE class Registry(TraceRegistry): def register(self, ident, trial_id, **kwargs): token = super().register(ident, trial_id, **kwargs) turns = template["turns"] + messages = [{'role': 'user', 'content': 'Read the file.'}] for index, original in enumerate(turns): + if budget: + original['client_request']['messages'] = deepcopy(messages) + original['finish_reason'] = 'max_tokens' if index < continuations else 'tool_use' + original['client_request']['tools'] = [{'name': 'Read', 'input_schema': {'type': 'object'}}] + original['client_request_digest'] = digest(original['client_request']) + original['response'] = {'stop_reason': original['finish_reason'], + 'content': [{'type': 'text', 'text': 'Inspecting'}] if index < continuations else + [{'type': 'tool_use', 'id': f'tool-{index}', + 'name': 'Read', 'input': {'file_path': '/workspace/code.py'}}]} + original['response_digest'] = digest(original['response']) + messages.append({'role': 'assistant', 'content': original['response']['content']}) + messages.append({'role': 'user', 'content': [{'type': 'text', + 'text': ('Output token limit hit. Resume directly — no apology, no recap of what you were doing. ' + 'Pick up mid-thought if that is where the cut happened. ' + 'Break remaining work into smaller pieces.')}]}) aid, request, _, _ = self.begin(token) turn = {**original, "request_id": request, "delivery": "sent"} if empty and index == 1: @@ -47,17 +70,28 @@ def register(self, ident, trial_id, **kwargs): "type": "assistant", "sessionId": "session", "message": { "id": turn["response_id"], "role": "assistant", "content": turn["response"]["content"]}}) - for turn in turns[:1 if empty else 2])) - (trial / "agent/claude-code.txt").write_text(json.dumps({ + for turn in (turns[:1] if empty else turns))) + final = { "type": "result", "subtype": "success", "is_error": False, "result": "", "stop_reason": "end_turn", "session_id": "session", - "terminal_reason": "completed"})) + "terminal_reason": "completed"} + events = [] + if budget: + final.update(subtype='error_max_turns', is_error=True, stop_reason='tool_use', + terminal_reason='max_turns', num_turns=3, permission_denials=[], + errors=['Reached maximum number of turns (2)']) + events = [{'type': 'system', 'subtype': 'init', 'session_id': 'session', + 'claude_code_version': '2.1.141'}] + events += [{'type': 'assistant', 'session_id': 'session', + 'message': {'id': turn['response_id']}} for turn in turns] + (trial / "agent/claude-code.txt").write_text('\n'.join(map(json.dumps, [*events, final]))) (trial / "verifier").mkdir() - (trial / "verifier/reward.json").write_text('{"reward": 1}') + (trial / "verifier/reward.json").write_text(json.dumps({'reward': reward})) (trial / "result.json").write_text(json.dumps({ "id": trial_id, "finished_at": "2026-10-19T00:00:00Z", "verifier_environment_mode": "shared", - "verifier_result": {"rewards": {"reward": 1}}})) + "verifier_result": {"rewards": {"reward": reward}}, + **({'exception_info': {'exception_type': 'NonZeroAgentExitCodeError'}} if budget else {})})) return token worker = tmp_path / "worker.py" @@ -68,24 +102,80 @@ def register(self, ident, trial_id, **kwargs): "sys.stdin.readline()\n" f"print({json.dumps(terminal)!r}, flush=True)\n") spawn = asyncio.create_subprocess_exec + practice_path = str(tmp_path / 'deployment-practice') + monkeypatch.setenv('PYTHONPATH', practice_path) async def fake_spawn(*args, **kwargs): + expected = str(Path(__file__).resolve().parents[2]) + assert kwargs['env']['PYTHONPATH'].split(os.pathsep) == [expected, practice_path] return await spawn(sys.executable, "-S", str(worker), **kwargs) monkeypatch.setattr(asyncio, "create_subprocess_exec", fake_spawn) - registry = Registry(tmp_path / "trace") ir = asyncio.run(run_attempt( - ident, task_path=tmp_path, profile=PROFILE, gateway_url="http://gateway", - runner_python=sys.executable, root=root, registry=registry, + ident, task_path=tmp_path, profile=profile, gateway_url="http://gateway", + runner_python=sys.executable, root=root, registry=Registry(tmp_path / "trace"), reward_profile=RewardProfile())) assert ir["readiness"] == "CONTRACT_READY" assert json.loads((root / "trajectory.v2.json").read_text()) == ir assert (root / "evaluation.json").is_file() assert not (root / "cleanup-required.json").exists() - assert not (root / "trace-unsealed.json").exists() - assert registry.attempts[ident.attempt_id]["turns"] == [] - assert registry.seal(ident.attempt_id) == (ir["turns"], ir["trace_seal"]) assert all(turn["delivery"] == "consumed_confirmed" for turn in ir["turns"]) + assert ir['evaluation']['raw_reward'] == reward + assert ir['termination'] == ('budget_truncated' if budget else 'completed') + if budget: + assert ir['termination_evidence'] == ir['evaluation']['termination_evidence'] + assert ir['termination_evidence']['max_turns'] == 2 + assert ir['termination_evidence']['policy_response_count'] == 2 + continuations + assert len(ir['termination_evidence']['continuations']) == continuations + assert [turn['output_ids'] for turn in ir['turns']] == [[3, 4]] * (2 + continuations) + assert [turn['logprobs'] for turn in ir['turns']] == [[-.2, -.3]] * (2 + continuations) + for mutation in ('missing_receipt', 'result_digest', 'consumed', 'completed'): + changed = deepcopy(ir) + if mutation == 'missing_receipt': + changed['evaluation'].pop('termination_evidence') + elif mutation == 'result_digest': + changed['termination_evidence']['source_result_digest'] = 'different' + elif mutation == 'consumed': + changed['termination_evidence']['consumed_response_ids'] = [] + else: + changed['termination'] = 'completed' + with pytest.raises(ValueError, match='termination_evidence|truncation_evidence'): + require_ready(changed, allow_synthetic=True) + + +def test_coordinator_reports_harbor_failure_before_session_creation(tmp_path, monkeypatch): + ident = identity() + trial_id = 'trial-' + ident.attempt_id + root = tmp_path / 'attempt' + trial = root / 'runner/harbor/trial' + result_path = trial / 'result.json' + terminal = {'type': 'terminal', 'trial_id': trial_id, 'trial_dir': str(trial), + 'result_path': str(result_path)} + worker = tmp_path / 'worker.py' + worker.write_text('import sys\nsys.stdin.readline()\n' + f"print({json.dumps({'type': 'prepared', 'trial_id': trial_id})!r}, flush=True)\n" + 'sys.stdin.readline()\n' + f'print({json.dumps(terminal)!r}, flush=True)\n') + spawn = asyncio.create_subprocess_exec + + async def fake_spawn(*args, **kwargs): + return await spawn(sys.executable, '-S', str(worker), **kwargs) + + monkeypatch.setattr(asyncio, 'create_subprocess_exec', fake_spawn) + registry = TraceRegistry(tmp_path / 'trace') + original_register = registry.register + + def register(*args, **kwargs): + trial.mkdir(parents=True) + result_path.write_text(json.dumps({'exception_info': {'exception_type': 'EnvironmentStartTimeoutError'}})) + return original_register(*args, **kwargs) + + monkeypatch.setattr(registry, 'register', register) + with pytest.raises(RuntimeError, match='before session audit: EnvironmentStartTimeoutError'): + asyncio.run(run_attempt(ident, task_path=tmp_path, profile=PROFILE, gateway_url='http://gateway', + runner_python=sys.executable, root=root, registry=registry, reward_profile=RewardProfile())) + assert not (root / 'cleanup-required.json').exists() + assert not (root / 'trajectory.v2.json').exists() @pytest.mark.parametrize("change", [ @@ -113,3 +203,212 @@ def unconsumed_responses(self, aid): "terminal_reason": "completed", **change})) with pytest.raises(ValueError, match="matching session result"): ClaudeCodeHarnessAdapter().audit_empty_final(path, Registry(), "a", {"session_id": "session"}) + + +@pytest.mark.parametrize('mutation', [ + 'cli_version', 'init_version', 'init_session', 'missing_init', 'duplicate_result', + 'not_final', 'session', 'reason', 'subtype', 'errors', 'turn_count', 'consumed', + 'assistant_session', 'exception', 'permissions', 'is_error', +]) +def test_budget_termination_rejects_unverified_evidence(tmp_path, mutation): + from harborrl.rollout.harbor_job.audit import audit_termination + + profile = {**PROFILE, 'cli_version': '2.1.141', 'max_turns': 2} + audit = {'session_id': 'session', 'consumed_response_ids': ['r0', 'r1'], 'evidence_ref': 'session#sha'} + initial = {'type': 'system', 'subtype': 'init', 'session_id': 'session', 'claude_code_version': '2.1.141'} + responses = [{'type': 'assistant', 'session_id': 'session', 'message': {'id': rid}} + for rid in audit['consumed_response_ids']] + final = {'type': 'result', 'subtype': 'error_max_turns', 'is_error': True, + 'terminal_reason': 'max_turns', 'stop_reason': 'tool_use', 'num_turns': 3, + 'session_id': 'session', 'permission_denials': [], + 'errors': ['Reached maximum number of turns (2)']} + result = {'exception_info': {'exception_type': 'NonZeroAgentExitCodeError'}} + events = [initial, *responses, final] + if mutation == 'cli_version': + profile['cli_version'] = '2.1.0' + elif mutation == 'init_version': + initial['claude_code_version'] = '2.1.0' + elif mutation == 'init_session': + initial['session_id'] = 'other' + elif mutation == 'missing_init': + events.remove(initial) + elif mutation == 'duplicate_result': + events.append(final) + elif mutation == 'not_final': + events.append({'type': 'late_event'}) + elif mutation == 'session': + final['session_id'] = 'other' + elif mutation == 'reason': + final['terminal_reason'] = 'timeout' + elif mutation == 'subtype': + final['subtype'] = 'error_during_execution' + elif mutation == 'errors': + final['errors'].append('connection failure') + elif mutation == 'turn_count': + final['num_turns'] = 2 + elif mutation == 'consumed': + audit['consumed_response_ids'] = ['r0'] + elif mutation == 'assistant_session': + responses[-1]['session_id'] = 'other' + elif mutation == 'exception': + result['exception_info']['exception_type'] = 'AgentTimeoutError' + elif mutation == 'permissions': + final['permission_denials'] = ['Bash'] + else: + final['is_error'] = 1 + path, result_path = tmp_path / 'stream.jsonl', tmp_path / 'result.json' + path.write_text('\n'.join(map(json.dumps, events))) + result_path.write_text(json.dumps(result)) + with pytest.raises(ValueError): + audit_termination(path, result_path, profile, audit, turns=[ + {'response_id': 'r0', 'finish_reason': 'tool_use'}, + {'response_id': 'r1', 'finish_reason': 'tool_use'}]) + + +def test_collector_binds_budget_evidence_to_original_result(tmp_path): + from test_native_contracts import receipt + + from harborrl.rollout.harbor_job.collector import collect + + ident = identity() + receipt(tmp_path, ident, 0) + root = tmp_path / ident.attempt_id + result = json.loads((root / 'result.json').read_text()) + result['exception_info'] = {'exception_type': 'NonZeroAgentExitCodeError'} + (root / 'result.json').write_text(json.dumps(result)) + with pytest.raises(ValueError, match='verified termination'): + collect(ident, result['id'], root / 'result.json', root, RewardProfile(), + termination_evidence={'profile': 'claude_code_2.1.141_max_turns_v1', + 'source_result_digest': 'another result'}) + + +@pytest.mark.parametrize('mutation', [None, 'changed', 'missing', 'true', 'numeric', 'unknown', + 'required', 'schema_missing', 'tool', 'id', 'text']) +def test_native_consumption_only_accepts_declared_optional_defaults(tmp_path, mutation): + ident = identity() + registry = TraceRegistry(tmp_path / 'trace') + token = registry.register(ident, 'trial-' + ident.attempt_id) + aid, request, _, _ = registry.begin(token) + turn = make_ir(tmp_path / 'fixture')['turns'][0] + block = {'type': 'tool_use', 'id': 'tool-1', 'name': 'Edit', + 'input': {'file_path': '/workspace/code.py', 'old_string': 'a', 'new_string': 'b'}} + schema = {'type': 'object', 'properties': {'replace_all': {'type': 'boolean', 'default': False}}, + 'required': ['file_path', 'old_string', 'new_string']} + turn.update(request_id=request, delivery='sent', response={'content': [block]}, + client_request={'tools': [{'name': 'Edit', 'input_schema': schema}]}) + content = [deepcopy(block)] + content[0]['input']['replace_all'] = False + if mutation == 'changed': + content[0]['input']['new_string'] = 'c' + elif mutation == 'missing': + content[0]['input'].pop('old_string') + elif mutation == 'true': + content[0]['input']['replace_all'] = True + elif mutation == 'numeric': + content[0]['input']['replace_all'] = 0 + elif mutation == 'unknown': + content[0]['input']['extra'] = False + elif mutation == 'required': + schema['required'].append('replace_all') + elif mutation == 'schema_missing': + schema['properties']['replace_all'].pop('default') + elif mutation == 'tool': + content[0]['name'] = 'Write' + elif mutation == 'id': + content[0]['id'] = 'different' + elif mutation == 'text': + content.append({'type': 'text', 'text': 'injected'}) + registry.finish(aid, request, turn) + registry.delivery(aid, turn['response_id'], 'sent') + if mutation: + with pytest.raises(ValueError): + registry.consume(aid, turn['response_id'], content, 'session#sha') + return + registry.consume(aid, turn['response_id'], content, 'session#sha') + registry.close(aid) + turns, seal = registry.seal(aid) + assert not seal['errors'] + assert turns[0]['response'] == turn['response'] + assert 'replace_all' not in turns[0]['response']['content'][0]['input'] + proof = turns[0]['consumption_tool_defaults'] + assert proof['changes'][0]['defaults'] == {'replace_all': False} + assert proof['served_content_digest'] == digest(turn['response']['content']) + assert proof['consumed_content_digest'] == digest(content) + + +@pytest.mark.parametrize('tool', ['Edit', 'Write']) +@pytest.mark.parametrize('mutation', [None, 'wrong_version', 'missing_version', 'markdown', + 'field', 'indent', 'newline', 'path', 'name', 'id', 'schema', 'default', 'extra', 'old_string']) +def test_pinned_cli_whitespace_audit_preserves_exact_sample(tmp_path, tool, mutation): + from harborrl.rollout.harbor_job.audit import audit_session + + registry = TraceRegistry(tmp_path / 'trace') + ident = identity() + token = registry.register(ident, 'trial') + aid, request, _, _ = registry.begin(token) + turn = make_ir(tmp_path / 'fixture')['turns'][0] + field = 'new_string' if tool == 'Edit' else 'content' + path = '/workspace/code.MDX' if mutation == 'markdown' else '/workspace/code.py' + text = 'def f(): \r\n return 1\t\r \n# end\u00a0' + normalized = 'def f():\r\n return 1\r\n# end' + inputs = {'file_path': path, field: text} + properties = {'file_path': {'type': 'string'}, field: {'type': 'string'}} + if tool == 'Edit': + inputs['old_string'] = 'old \n' + properties.update(old_string={'type': 'string'}, replace_all={'type': 'boolean', 'default': False}) + block = {'type': 'tool_use', 'id': 'tool-1', 'name': tool, 'input': inputs} + turn.update(request_id=request, response={'content': [block]}, + client_request={'tools': [{'name': tool, 'input_schema': { + 'type': 'object', 'properties': properties, 'required': list(inputs)}}]}) + content = [deepcopy(block)] + consumed = content[0]['input'] + consumed[field] = normalized + if tool == 'Edit': + consumed['replace_all'] = False + version = '2.1.141' + if mutation == 'wrong_version': + version = '2.1.140' + elif mutation == 'missing_version': + version = None + elif mutation == 'field': + consumed[field] += 'changed' + elif mutation == 'indent': + consumed[field] = normalized.replace(' return', 'return') + elif mutation == 'newline': + consumed[field] = normalized.replace('\r\n', '\n') + elif mutation == 'path': + consumed['file_path'] = '/workspace/other.py' + elif mutation == 'name': + content[0]['name'] = 'Other' + elif mutation == 'id': + content[0]['id'] = 'other' + elif mutation == 'schema': + properties[field]['type'] = 'number' + elif mutation == 'default': + consumed['replace_all'] = True + elif mutation == 'extra': + consumed['extra'] = False + elif mutation == 'old_string': + consumed['old_string'] = 'old\n' + snapshot = deepcopy(turn) + registry.finish(aid, request, turn) + registry.delivery(aid, turn['response_id'], 'sent') + session = tmp_path / 'native-session.jsonl' + session.write_text(json.dumps({'type': 'assistant', 'sessionId': 'session', 'version': version, + 'message': {'role': 'assistant', 'id': turn['response_id'], 'content': content}}) + '\n') + if mutation: + with pytest.raises(ValueError): + audit_session(session, registry, aid, cli_version='2.1.141') + assert registry.attempts[aid]['turns'][0]['response'] == snapshot['response'] + return + audit_session(session, registry, aid, cli_version='2.1.141') + registry.close(aid) + turns, seal = registry.seal(aid) + assert not seal['errors'] + for key in ('response', 'output_ids', 'text', 'logprobs'): + assert turns[0].get(key) == snapshot.get(key) + proof = turns[0]['consumption_tool_normalization'] + assert proof['kind'] == 'claude_code_2.1.141_tool_inputs_v1' + assert proof['changes'][0]['field'] == field + assert proof['served_content_digest'] == digest(snapshot['response']['content']) + assert proof['consumed_content_digest'] == digest(content) diff --git a/tests/harborrl/test_native_training.py b/tests/harborrl/test_native_training.py index 14328c2d5..33100ede4 100644 --- a/tests/harborrl/test_native_training.py +++ b/tests/harborrl/test_native_training.py @@ -87,6 +87,22 @@ def test_schema2_launch_plan_and_strict_override(tmp_path, capsys): assert json.loads(plan["environment"]["HARBORRL_NATIVE_RUNNER_WORKERS"]) == ["worker-a", "worker-b"] assert plan["environment"]["N_SAMPLES"] == "2" assert plan["environment"]["TRAIN_PYTHON"] == sys.executable + assert 'HARBORRL_TRAINING_SEED' not in plan['environment'] + for seed in (0, 42, 2147483647): + seeded = launch_plan(load_config(path, [f'training.seed={seed}'])) + assert seeded['environment']['HARBORRL_TRAINING_SEED'] == str(seed) + for value in ('true', '-1', '2147483648', '1.5'): + with pytest.raises(ValueError, match='training.seed'): + load_config(path, [f'training.seed={value}']) + from harborrl.platform.evidence import run_environment + prompt_path = tmp_path / 'prompts.jsonl' + assert 'SAVE_HF' not in run_environment(plan, tmp_path, Path(catalog_path), prompt_path) + exported = launch_plan(load_config(path, ['training.export_hf=true'])) + assert run_environment(exported, tmp_path, Path(catalog_path), prompt_path)['SAVE_HF'] == str( + tmp_path / 'hf' / 'rollout-{rollout_id}') + for value in ('1', 'null', '"true"'): + with pytest.raises(ValueError, match='training.export_hf'): + load_config(path, [f'training.export_hf={value}']) assert json.loads(plan["environment"]["HARBORRL_NATIVE_TASK_WORKER_OVERRIDES"]) == {} assert plan["training_command"][0] == "bash" assert plan["launcher"].startswith("native Claude Code") @@ -146,13 +162,10 @@ def test_cli_command_mode_selects_experiment_dimensions(tmp_path, capsys, monkey ([], '0.0001'), (['training.learning_rate=0.000025'], '2.5e-05'), ]) -@pytest.mark.parametrize('layout', ['split', 'colocate']) -def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expected_lr, layout): +def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expected_lr): config_path, task = native_config(tmp_path) from harborrl.data.harbor.native_inspector import inspect_native - total_gpus = '8' if layout == 'split' else '4' - config = load_config(config_path, [*overrides, f'deployment.layout={layout}', - f'deployment.num_gpus={total_gpus}']) + config = load_config(config_path, overrides) catalog_path = Path(config['tasks']['catalog']) entries = json.loads(catalog_path.read_text()) entries[0]['task_digest'] = inspect_native(task)['task_digest'] @@ -174,9 +187,9 @@ def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expe "HARNESS_OPTION": "native-claude-code", "HARBORRL_NATIVE_ROLLOUT": "1", "HARBORRL_SKIP_GLOBAL_CLEANUP": "1", "ROLLOUT_PROMPT_DATA": str(prompt), "HF_CKPT": "/model", "REF_LOAD": "/reference", "MODEL_ARGS_FILE": "qwen3-8B", - "MODEL_TAG": "qwen3-8B", "NUM_GPUS": total_gpus, "ACTOR_GPUS": "4", + "MODEL_TAG": "qwen3-8B", "NUM_GPUS": "8", "ACTOR_GPUS": "4", "ROLLOUT_GPUS": "4", "TP_SIZE": "4", "ROLLOUT_NUM_GPUS_PER_ENGINE": "4", - "HARBORRL_GPU_LAYOUT": layout, "NUM_ROLLOUT": "1", "SAVE_INTERVAL": "1", + "HARBORRL_GPU_LAYOUT": "split", "NUM_ROLLOUT": "1", "SAVE_INTERVAL": "1", "RUN_DIR": str(root), "RUN_ID": "run", "RUNS_ROOT": str(tmp_path), "CKPT_ROOT": str(root / "checkpoints"), "HARBORRL_NATIVE_MAX_ATTEMPTS": "2", "HARBORRL_NATIVE_DRAIN_TIMEOUT": "1", @@ -185,6 +198,7 @@ def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expe "HARBORRL_NATIVE_SLIME_CONTRACT": "slime-v032", "SLIME_DIR": str(slime), "TRAIN_PYTHON": sys.executable, + "HARBORRL_TRAINING_SEED": "0", }) env['LR'] = plan['environment']['LR'] script = Path(__file__).resolve().parents[2] / "harborrl/platform/slime_train.sh" @@ -203,17 +217,17 @@ def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expe '--num-rollout': '1', '--save-interval': '1', '--rollout-batch-size': '1', '--n-samples-per-prompt': '2', '--rollout-max-response-len': '128', '--rollout-max-context-len': '1024', - '--num-gpus-per-node': total_gpus, '--actor-num-gpus-per-node': '4', + '--num-gpus-per-node': '8', '--actor-num-gpus-per-node': '4', '--rollout-num-gpus': '4', '--tensor-model-parallel-size': '4', '--rollout-num-gpus-per-engine': '4', '--hf-checkpoint': '/model', '--ref-load': '/reference', '--prompt-data': str(prompt), '--save': str(root / 'checkpoints' / 'run'), + '--seed': '0', '--rollout-seed': '0', } for flag, value in expected_flags.items(): assert arguments[arguments.index(flag) + 1] == value recorded = json.loads(next((root / 'dry-run-commands').glob('*.json')).read_text()) assert recorded['command'] == arguments - assert arguments.count('--colocate') == int(layout == 'colocate') assert recorded['mode'] == 'dry_run' for argument in ( "--rollout-function-path harborrl.backends.slime_v032.rollout.generate_rollout", @@ -225,8 +239,9 @@ def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expe assert argument in command assert "native_slime" not in command assert "--use-rollout-logprobs" in command + assert '--sglang-enable-deterministic-inference' in command assert "--megatron-to-hf-mode bridge" in command - assert "--no-save-optim" in command + assert "--no-save-optim" not in command assert "--load /model" in command assert command.count(" --load ") == 1 assert "--num-steps-per-rollout 2" in command @@ -241,7 +256,8 @@ def test_native_shell_dry_run_uses_official_v032_hooks(tmp_path, overrides, expe assert result.returncode == 0, result.stdout command = next(line for line in result.stdout.splitlines() if line.startswith("[dry-run] ")) assert "--load /checkpoint" in command - assert "--no-load-optim" in command + assert "--no-load-optim" not in command + assert "--use-checkpoint-opt-param-scheduler" in command def test_slime_runtime_config_maps_args_and_environment(tmp_path, monkeypatch): @@ -333,6 +349,8 @@ def test_native_slots_spread_workers_and_shift_retries(tmp_path, monkeypatch): "HARBORRL_NATIVE_MAX_ATTEMPTS": "2", "HARBORRL_NATIVE_GATEWAY_URL": "http://gateway:32123", "HARBORRL_NATIVE_GATEWAY_PORT": "32123", + # Fake runtimes carry no lease bookkeeping: pin static routing. + "HARBORRL_NATIVE_GATEWAY_ROUTING": "modulo", "HARBORRL_NATIVE_RUNNER_PYTHON": "/python", "HARBORRL_NATIVE_CATALOG": str(tmp_path / "catalog.json"), "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, @@ -371,8 +389,11 @@ async def fake_run_attempt(identity, *, runner_host, **kwargs): monkeypatch.setattr(rollout_runtime, "GroupSlots", fake_group_slots) monkeypatch.setattr(rollout_runtime, "run_attempt", fake_run_attempt) - runtime_ = SimpleNamespace(rollout_root=tmp_path, semaphore=asyncio.Semaphore(8), - registry=SimpleNamespace(), profile={"model": "policy"}) + runtime_ = SimpleNamespace( + rollout_root=tmp_path, semaphore=asyncio.Semaphore(8), + shards=(SimpleNamespace(name="gateway", url="http://gateway:32123", + registry=SimpleNamespace()),), + profile={"model": "policy"}) async def gather_slots(): return await asyncio.gather(*[ @@ -432,6 +453,11 @@ def test_native_dispatch_passes_materialized_prompt_data(tmp_path, monkeypatch): assert effective['environment']['N_SAMPLES'] == '2' assert effective['environment']['ROLLOUT_PROMPT_DATA'] == output.read_text() assert effective['components']['algorithm'] == 'grpo' + assert effective['update_semantics'] == { + 'advantage': 'reward_minus_complete_group_mean', + 'std_normalization': False, + 'reduction': 'trajectory_weighted_clipped_pg_v1', + } assert effective['deployment_environment']['SLIME_DIR'] == '/opt/slime-v0.3.2' assert 'ambient-provider-credential' not in effective_text assert 'user:credential' not in effective_text @@ -450,6 +476,69 @@ def test_all_schema_fields_have_declared_consumers(): } +@pytest.mark.parametrize('cancel_batch', [False, True]) +@pytest.mark.parametrize('uneven_cleanup', [False, True]) +@pytest.mark.parametrize('cancel_during_cleanup', [False, True]) +def test_native_batch_waits_for_all_group_cleanup( + monkeypatch, cancel_batch, uneven_cleanup, cancel_during_cleanup): + from harborrl.backends.slime_v032 import rollout + from harborrl.rollout import runtime as rollout_runtime + + monkeypatch.delenv('HARBORRL_ROLLOUT_SERVICE_URL', raising=False) + cleaned, closed = [], [] + failure = RuntimeError('one group failed') + + async def scenario(): + started = set() + all_started, blocked = asyncio.Event(), asyncio.Event() + cleanup_started = asyncio.Event() + + async def slot(config, runtime_, group, index, entry, reward_profile): + started.add((group, index)) + if len(started) == 4: + all_started.set() + try: + await all_started.wait() + if not cancel_batch and (group, index) == (0, 0): + raise failure + await blocked.wait() + finally: + # Cleanup must be awaited across both levels of group tasks. + if (group, index) == (1, 1): + cleanup_started.set() + await asyncio.sleep(0) + await asyncio.sleep(0.02 if uneven_cleanup and index else 0) + cleaned.append((group, index)) + + async def generate(args, group, params, rollout_id): + return await rollout_runtime.run_group( + SimpleNamespace(group_size=2), None, entry={}, group_index=group, + rollout_id=rollout_id, sampling_digest='sampling', policy_version='1') + + monkeypatch.setattr(rollout_runtime, '_group_slots', + lambda config, runtime_, group, *args: (group, None)) + monkeypatch.setattr(rollout_runtime, '_run_slot', slot) + monkeypatch.setattr(rollout, 'generate_group', generate) + monkeypatch.setattr(rollout, 'close_runtime', lambda: closed.append(set(cleaned))) + batch = asyncio.create_task(rollout.run_native_groups( + SimpleNamespace(rollout_max_response_len=128), [0, 1], '1', 0)) + await asyncio.wait_for(all_started.wait(), timeout=5) + if cancel_batch: + batch.cancel() + if cancel_during_cleanup: + await asyncio.wait_for(cleanup_started.wait(), timeout=5) + batch.cancel() + expected_error = asyncio.CancelledError if cancel_batch else RuntimeError + with pytest.raises(expected_error) as exc: + await asyncio.wait_for(batch, timeout=5) + if not cancel_batch: + assert exc.value is failure + assert closed == [{(0, 0), (0, 1), (1, 0), (1, 1)}] + assert len(cleaned) == 4 + + asyncio.run(scenario()) + + def test_schema_validation_does_not_read_harness_files(tmp_path): from harborrl.config.native import validate, validate_files config_path, _ = native_config(tmp_path) diff --git a/tests/harborrl/test_opd_training.py b/tests/harborrl/test_opd_training.py new file mode 100644 index 000000000..08d2fa996 --- /dev/null +++ b/tests/harborrl/test_opd_training.py @@ -0,0 +1,190 @@ +"""OPD contracts at the frozen-teacher and shared production-loss boundaries.""" +import json +import sys +from dataclasses import replace +from types import SimpleNamespace + +import pytest +import yaml +from test_math_rollout import samples + +from harborrl.backends.slime_v032.math_rollout import ( + build_batch, + convert_samples_to_train_data, +) +from harborrl.backends.slime_v032.postprocess import rollout_data_postprocess +from harborrl.config import ROOT, launch_plan, load_config +from harborrl.export.validate import validate_batch +from harborrl.math import teacher +from harborrl.platform.evidence import effective_config, run_environment +from harborrl.trajectories.native import digest + + +@pytest.fixture +def frozen(tmp_path): + checkpoint = tmp_path / 'teacher' + checkpoint.mkdir() + for name, data in {'config.json': '{}', 'model.safetensors': 'weight fixture', + 'tokenizer.json': '{}', 'tokenizer_config.json': '{}'}.items(): + (checkpoint / name).write_text(data, encoding='utf-8') + return teacher.freeze({'endpoint': 'http://127.0.0.1:30090', 'checkpoint': str(checkpoint), + 'weight_version': 'teacher-v1'}, checkpoint) + + +def test_teacher_fingerprint_and_tokenizer_compatibility(frozen, tmp_path): + root = tmp_path / 'teacher' + assert teacher.freeze({key: frozen[key] for key in ('endpoint', 'checkpoint', 'weight_version')}, root) == frozen + (root / 'model.safetensors').write_text('different weights', encoding='utf-8') + changed = teacher.freeze({key: frozen[key] for key in ('endpoint', 'checkpoint', 'weight_version')}, root) + assert changed['identity'] != frozen['identity'] + student = tmp_path / 'student' + student.mkdir() + (student / 'tokenizer.json').write_text('{"different":true}', encoding='utf-8') + (student / 'tokenizer_config.json').write_text('{}', encoding='utf-8') + with pytest.raises(ValueError, match='tokenizer files'): + teacher.freeze({key: frozen[key] for key in ('endpoint', 'checkpoint', 'weight_version')}, student) + + +@pytest.mark.parametrize('mutation', ['weights', 'tokenizer', 'version']) +def test_teacher_checks_live_identity(frozen, mutation): + info = {'model_path': frozen['checkpoint'], 'tokenizer_path': frozen['checkpoint'], + 'weight_version': frozen['weight_version'], 'is_generation': True} + assert teacher.check(frozen, transport=lambda url: info)['identity'] == frozen['identity'] + field = {'weights': 'model_path', 'tokenizer': 'tokenizer_path', 'version': 'weight_version'}[mutation] + info[field] += '-changed' + with pytest.raises(ValueError, match='changed'): + teacher.check(frozen, transport=lambda url: info) + + +@pytest.mark.parametrize('mutation', [None, 'version', 'token', 'length', 'nan', 'positive']) +def test_teacher_scores_exact_input_positions(frozen, mutation): + def transport(url, payload): + assert url.endswith('/generate') + assert payload['input_ids'] == [11, 12, 21, 22] + assert payload['sampling_params']['max_new_tokens'] == 0 + assert payload['logprob_start_len'] == 0 + rows = [[None, 11], [-0.1, 12], [-0.2, 21], [-0.3, 22]] + version = frozen['weight_version'] + if mutation == 'version': + version = 'changed' + elif mutation == 'token': + rows[2][1] = 99 + elif mutation == 'length': + rows.pop() + elif mutation == 'nan': + rows[2][0] = float('nan') + elif mutation == 'positive': + rows[2][0] = 0.1 + return {'meta_info': {'weight_version': version, 'input_token_logprobs': rows}} + + if mutation: + with pytest.raises(ValueError): + teacher.score_response(frozen, [11, 12], [21, 22], transport=transport) + else: + result = teacher.score_response(frozen, [11, 12], [21, 22], transport=transport) + assert result['logprobs'] == [-0.2, -0.3] + assert result['tokens_digest'] == digest([11, 12, 21, 22]) + + +def opd_samples(frozen, tmp_path, monkeypatch): + path = tmp_path / 'teacher.json' + path.write_text(json.dumps(frozen), encoding='utf-8') + monkeypatch.setenv('HARBORRL_MATH_ALGORITHM', 'opd') + monkeypatch.setenv('HARBORRL_OPD_TEACHER', str(path)) + rows = samples() + for index, sample in enumerate(rows): + sample.teacher_log_probs = [-0.2 - index] + sample.metadata['math_evidence'].update(algorithm='opd', teacher={ + 'identity': frozen['identity'], 'weight_version': frozen['weight_version'], + 'logprobs': list(sample.teacher_log_probs), 'tokens_digest': digest(sample.tokens), + }) + return rows + + +def test_opd_batch_preserves_token_scores_and_excludes_outcome_rewards(frozen, tmp_path, monkeypatch): + rows = opd_samples(frozen, tmp_path, monkeypatch) + args = SimpleNamespace(n_samples_per_prompt=2) + batch = build_batch(args, rows) + assert batch.algorithm == 'opd' + assert batch.teacher_identity == frozen['identity'] + assert [unit.advantage for unit in batch.groups[0].trajectories] == [0., 0.] + data = convert_samples_to_train_data(args, rows) + assert data['raw_reward'] == [1., -1.] + assert data['rewards'] == [0., 0.] + assert data['teacher_log_probs'] == [[-0.2], [-1.2]] + rollout_data_postprocess(args, 0, data) + del data['teacher_log_probs'] + with pytest.raises(ValueError, match='every sample'): + rollout_data_postprocess(args, 0, data) + with pytest.raises(ValueError, match='teacher_identity'): + validate_batch(replace(batch, teacher_identity=None)) + with pytest.raises(ValueError, match='reduction'): + validate_batch(replace(batch, algorithm='grpo')) + + +@pytest.mark.parametrize('mutation', ['identity', 'version', 'tokens', 'score', 'length', 'nan', 'positive']) +def test_opd_rejects_teacher_evidence_drift(frozen, tmp_path, monkeypatch, mutation): + rows = opd_samples(frozen, tmp_path, monkeypatch) + scores = rows[0].metadata['math_evidence']['teacher'] + if mutation in ('identity', 'version', 'tokens'): + key = {'identity': 'identity', 'version': 'weight_version', 'tokens': 'tokens_digest'}[mutation] + scores[key] = 'changed' + elif mutation == 'score': + rows[0].teacher_log_probs = [-3.] + else: + values = [] if mutation == 'length' else [float('nan') if mutation == 'nan' else 0.1] + rows[0].teacher_log_probs = values + scores['logprobs'] = values + with pytest.raises(ValueError): + build_batch(SimpleNamespace(n_samples_per_prompt=2), rows) + + +def test_opd_config_and_audit_are_explicit(tmp_path): + config = yaml.safe_load((ROOT / 'examples/math/train.yaml').read_text(encoding='utf-8')) + config['tasks']['path'] = str(ROOT / 'examples/math/tasks.jsonl') + config['algorithm'] = {'name': 'opd'} + path = tmp_path / 'train.yaml' + path.write_text(yaml.safe_dump(config), encoding='utf-8') + with pytest.raises(ValueError, match='teacher'): + load_config(path) + config['teacher'] = {'endpoint': 'http://127.0.0.1:30090', 'checkpoint': '/models/teacher', + 'weight_version': 'teacher-v1'} + path.write_text(yaml.safe_dump(config), encoding='utf-8') + plan = launch_plan(load_config(path)) + assert plan['math']['algorithm']['id'] == 'opd' + environment = run_environment(plan, tmp_path, path, path) + assert environment['HARBORRL_OPD_TEACHER'] == str(tmp_path / 'config' / 'teacher.json') + audit = effective_config(plan, environment, {}, {}, []) + assert audit['update_semantics']['reward_mixing'] is False + assert audit['update_semantics']['clipping'] is False + with pytest.raises(ValueError, match='teacher'): + load_config(path, ['algorithm.name=grpo']) + + +def test_production_opd_loss_matches_ragged_trajectory_mean_and_detaches(monkeypatch): + torch = pytest.importorskip('torch') + from harborrl.backends.slime_v032.opd import ( + compute_advantages_and_returns, + loss_function, + ) + + current = torch.tensor([-0.5, -1.0, -0.7, -1.5, -0.2], dtype=torch.float64, requires_grad=True) + old = torch.tensor([-0.6, -0.9, -0.8, -1.4, -0.3], dtype=torch.float64, requires_grad=True) + target = torch.tensor([-0.2, -1.1, -0.4, -1.0, -0.5], dtype=torch.float64, requires_grad=True) + batch = {'unconcat_tokens': [], 'total_lengths': [5, 4], 'response_lengths': [3, 2], + 'rollout_mask_sums': [2, 2], 'loss_masks': [torch.tensor([1, 0, 1]), torch.tensor([1, 1])], + 'rollout_log_probs': list(old.split([3, 2])), 'teacher_log_probs': list(target.split([3, 2])), + 'rewards': [1000., -1000.]} + monkeypatch.setitem(sys.modules, 'slime.backends.megatron_utils.loss', SimpleNamespace( + get_log_probs_and_entropy=lambda *args, **kwargs: (None, {'log_probs': list(current.split([3, 2]))}))) + compute_advantages_and_returns(None, batch) + assert not any(value.requires_grad for value in batch['advantages']) + loss, metrics = loss_function(SimpleNamespace(n_samples_per_prompt=2), batch, current, None) + # Two trajectories with two valid tokens each: each active token weighs 1/4. + expected = -(current[[0, 2, 3, 4]] * (target - old).detach()[[0, 2, 3, 4]]).mean() + assert loss.item() == pytest.approx(expected.item()) + loss.backward() + assert current.grad.tolist() == pytest.approx([-0.1, 0., -0.1, -0.1, 0.05]) + assert old.grad is None and target.grad is None + assert metrics['opd_weight_sum'].item() == pytest.approx(1.) + assert metrics['opd_token_count'].item() == 4 diff --git a/tests/harborrl/test_prm_training.py b/tests/harborrl/test_prm_training.py new file mode 100644 index 000000000..fdbc47e0e --- /dev/null +++ b/tests/harborrl/test_prm_training.py @@ -0,0 +1,279 @@ +"""PRM identity, boundaries, complete-group baselines and gradient contracts.""" +import json +import sys +from copy import deepcopy +from dataclasses import replace +from types import SimpleNamespace +from typing import ClassVar + +import pytest +import yaml +from test_math_rollout import samples + +from harborrl.backends.slime_v032 import math_rollout +from harborrl.config import ROOT, launch_plan, load_config +from harborrl.export.validate import validate_batch +from harborrl.math import prm +from harborrl.math.data import verify +from harborrl.platform.evidence import effective_config, run_environment +from harborrl.trajectories.native import digest + + +class Tokenizer: + pieces: ClassVar = {1: '2', 2: '\n\n', 3: '#### 2', 4: ' and next', 5: '\n\nNext', 6: '#### 999999', 99: ''} + + def decode(self, ids, skip_special_tokens=True): + return ''.join(self.pieces[token] for token in ids) + + +@pytest.fixture +def frozen(tmp_path): + model = tmp_path / 'prm' + model.mkdir() + config = {'architectures': ['Qwen2ForProcessRewardModel'], 'auto_map': { + 'AutoConfig': 'configuration_qwen2_rm.Qwen2RMConfig', + 'AutoModel': 'modeling_qwen2_rm.Qwen2ForProcessRewardModel'}} + (model / 'config.json').write_text(json.dumps(config), encoding='utf-8') + for name in ('tokenizer.json', 'tokenizer_config.json', 'configuration_qwen2_rm.py', + 'modeling_qwen2_rm.py', 'model.safetensors'): + (model / name).write_text('{}', encoding='utf-8') + return prm.freeze({'endpoint': 'http://127.0.0.1:30091', 'checkpoint': str(model), 'revision': 'a' * 40}) + + +def receipt(frozen, query, ids, probabilities): + def transport(url, payload): + assert url.endswith('/score') + return {'identity': frozen['identity'], 'request_digest': digest(payload), + 'scorer_input_digest': digest([7, 8, 9]), 'probabilities': probabilities} + return prm.score_response(frozen, query, ids, Tokenizer(), transport=transport) + + +def setup_samples(frozen, tmp_path, monkeypatch): + manifest = tmp_path / 'process_reward.json' + manifest.write_text(json.dumps(frozen), encoding='utf-8') + monkeypatch.setenv('HARBORRL_PROCESS_REWARD', str(manifest)) + monkeypatch.setenv('HARBORRL_MATH_ALGORITHM', 'grpo') + monkeypatch.setitem(sys.modules, 'slime.utils.processing_utils', SimpleNamespace( + load_tokenizer=lambda *args, **kwargs: Tokenizer())) + rows = samples() + for sample, scores in zip(rows, ([.9, .8], [.2, .4]), strict=True): + ids = [1, 2, 3 if sample.index == 0 else 6, 99] + sample.tokens = [42, 43] + ids + sample.response = Tokenizer().decode(ids) + sample.response_length = len(ids) + sample.loss_mask = [1] * len(ids) + sample.rollout_log_probs = [-.5] * len(ids) + sample.rollout_id = 0 + evidence = sample.metadata['math_evidence'] + evidence.update(tokens=list(sample.tokens), response=sample.response, response_length=len(ids), + verification=verify(sample.response, sample.metadata['math_task']['reference']), + logprobs=list(sample.rollout_log_probs), + process_reward=receipt(frozen, sample.metadata['math_task']['prompt'], ids, scores)) + return rows, SimpleNamespace(n_samples_per_prompt=2, hf_checkpoint='student') + + +def test_identity_pins_executable_source_and_weights(frozen): + config = {key: frozen[key] for key in prm.FIELDS} + assert prm.freeze({**config, 'endpoint': 'http://127.0.0.1:39999'})['identity'] == frozen['identity'] + from pathlib import Path + for name in ('modeling_qwen2_rm.py', 'model.safetensors', 'tokenizer.json'): + path = Path(frozen['checkpoint']) / name + original = path.read_bytes() + path.write_bytes(original + b'changed') + assert prm.freeze(config)['identity'] != frozen['identity'] + path.write_bytes(original) + + +@pytest.mark.parametrize('mutation', ['identity', 'revision', 'frozen', 'is_generation']) +def test_live_scorer_identity(frozen, mutation): + info = {key: frozen[key] for key in ('identity', 'revision', 'adapter', 'checkpoint')} + info.update(frozen=True, is_generation=False) + prm.check(frozen, transport=lambda url: info) + info[mutation] = not info[mutation] if type(info[mutation]) is bool else 'changed' + with pytest.raises(ValueError): + prm.check(frozen, transport=lambda url: info) + + +@pytest.mark.parametrize('ids,ends,steps', [ + ([1, 2, 3, 99], [2, 4], ['2\n\n', '#### 2']), + ([1, 5, 3, 99], [2, 4], ['2\n\nNext', '#### 2']), + ([2, 1, 2, 99], [4], ['\n\n2\n\n']), + ([1, 2, 99], [3], ['2\n\n']), +]) +def test_steps_keep_boundary_tokens_and_trailing_specials(ids, ends, steps): + assert prm.split_steps(ids, Tokenizer()) == (ends, steps) + assert ''.join(steps) == Tokenizer().decode(ids) + + +def test_partial_unicode_prefix_waits_for_complete_character(): + class ByteTokenizer: + def decode(self, ids, **kwargs): + return {1: 'x\n\n\ufffd', 2: 'x\n\n\u6570', 3: 'x\n\n\u6570y'}[len(ids)] + assert prm.split_steps([1, 2, 3], ByteTokenizer()) == ([2, 3], ['x\n\n\u6570', 'y']) + with pytest.raises(ValueError, match='no visible reasoning'): + prm.split_steps([99], Tokenizer()) + + +@pytest.mark.parametrize('ends,values', [([0, 2], [.5, .5]), ([2, 2], [.5, .5]), + ([1], [.5]), ([2], [float('nan')]), ([2], [1.1])]) +def test_invalid_token_partition_and_probability(ends, values): + with pytest.raises(ValueError): + prm.token_returns(ends, values, 2) + + +@pytest.mark.parametrize('mutation', ['identity', 'request_digest', 'response_tokens_digest', 'step_ends', + 'probabilities', 'token_returns', 'scorer_input_digest']) +def test_scoring_receipt_binds_request_tokens_and_returns(frozen, mutation): + scores = receipt(frozen, 'question', [1, 2, 3, 99], [.9, .2]) + scores[mutation] = {'step_ends': [1, 4], 'probabilities': [.9, .3], 'token_returns': [9.] * 4}.get(mutation, 'changed') + with pytest.raises(ValueError): + prm.validate_receipt(frozen, scores, 'question', [1, 2, 3, 99], Tokenizer()) + + +def test_prm_neutral_contract_and_scalar_outcome_separation(frozen, tmp_path, monkeypatch): + rows, args = setup_samples(frozen, tmp_path, monkeypatch) + batch = math_rollout.build_batch(args, rows) + assert batch.algorithm == 'grpo' + assert batch.process_reward_identity == frozen['identity'] + group = batch.groups[0] + assert [unit.reward for unit in group.trajectories] == pytest.approx([1.4, -.8]) + assert [unit.advantage for unit in group.trajectories] == pytest.approx([1.1, -1.1]) + assert group.trajectories[0].turns[0].process_rewards.token_returns == pytest.approx([1.4, 1.4, .6, .6]) + data = math_rollout.convert_samples_to_train_data(args, rows) + assert data['raw_reward'] == [1., -1.] + assert data['rewards'][0]['baseline'] == pytest.approx(.3) + assert 'teacher_log_probs' not in data + for changed in (replace(batch, process_reward_identity=None), replace(batch, process_reward_identity='changed')): + with pytest.raises(ValueError, match='PRM'): + validate_batch(changed) + corrupted = replace(group.trajectories[0].turns[0].process_rewards, token_returns=(0.,) * 4) + unit = replace(group.trajectories[0], turns=(replace(group.trajectories[0].turns[0], process_rewards=corrupted),)) + with pytest.raises(ValueError, match='reward-to-go'): + validate_batch(replace(batch, groups=(replace(group, trajectories=(unit, group.trajectories[1])),))) + + +def test_prm_config_uses_same_training_chain_and_explicit_semantics(tmp_path): + config = yaml.safe_load((ROOT / 'examples/math/train.yaml').read_text(encoding='utf-8')) + config['tasks']['path'] = str(ROOT / 'examples/math/tasks.jsonl') + config['process_reward'] = {'endpoint': 'http://127.0.0.1:30091', 'checkpoint': '/models/prm', 'revision': 'a' * 40} + path = tmp_path / 'train.yaml' + path.write_text(yaml.safe_dump(config), encoding='utf-8') + plan = launch_plan(load_config(path)) + environment = run_environment(plan, tmp_path, path, path) + assert environment['HARBORRL_PROCESS_REWARD'] == str(tmp_path / 'config' / 'process_reward.json') + semantics = effective_config(plan, environment, {}, {}, [])['update_semantics'] + assert semantics['outcome_reward_mixing'] is False + assert semantics['std_normalization'] is False + assert 'reward_to_go' in semantics['advantage'] + assert plan['training_command'][-1].endswith('slime_train.sh') + for override in ('algorithm.name=opd', 'algorithm.name=dapo', 'process_reward.revision=main'): + with pytest.raises(ValueError): + load_config(path, [override]) + + +def test_dp_partial_group_retains_full_group_baseline_and_validates_transport(frozen, tmp_path, monkeypatch): + torch = pytest.importorskip('torch') + from harborrl.backends.slime_v032.postprocess import rollout_data_postprocess + from harborrl.backends.slime_v032.prm import compute_advantages_and_returns + rows, args = setup_samples(frozen, tmp_path, monkeypatch) + original = math_rollout.convert_samples_to_train_data(args, rows) + + def shard(index): + data = {key: [deepcopy(values[index])] for key, values in original.items()} + for key in ('tokens', 'rollout_log_probs', 'loss_masks'): + data[key] = [torch.tensor(data[key][0])] + return data + + data = shard(0) + compute_advantages_and_returns(args, data) + assert data['advantages'][0].tolist() == pytest.approx([1.1, 1.1, .3, .3]) + assert data['returns'][0].tolist() == pytest.approx([1.4, 1.4, .6, .6]) + assert data['rewards'] == pytest.approx([1.1]) + rollout_data_postprocess(args, 0, data) + for key, value in (('identity', 'changed'), ('tokens_digest', 'changed'), ('baseline', 9.), + ('sample_index', 5), ('rollout_id', 5), ('group_totals', {'trajectory-0': 1.4})): + data = shard(0) + data['rewards'][0][key] = value + with pytest.raises(ValueError): + compute_advantages_and_returns(args, data) + + +def test_token_process_advantages_reach_ragged_clipped_policy_gradient(frozen, tmp_path, monkeypatch): + torch = pytest.importorskip('torch') + from harborrl.backends.slime_v032.loss import loss_function + from harborrl.backends.slime_v032.prm import compute_advantages_and_returns + rows, args = setup_samples(frozen, tmp_path, monkeypatch) + args.eps_clip = .2 + data = math_rollout.convert_samples_to_train_data(args, rows) + for key in ('tokens', 'rollout_log_probs', 'loss_masks'): + data[key] = [torch.tensor(values) for values in data[key]] + compute_advantages_and_returns(args, data) + data.update(unconcat_tokens=[], total_lengths=[6, 6]) + ratios = torch.tensor([.7, 1., 1.3, 1., .7, 1., 1.3, 1.], dtype=torch.float64) + current = (ratios.log() - .5).requires_grad_() + monkeypatch.setitem(sys.modules, 'slime.backends.megatron_utils.loss', SimpleNamespace( + get_log_probs_and_entropy=lambda *a, **k: (None, {'log_probs': list(current.split(4))}))) + loss, _ = loss_function(args, data, current, None) + advantages = torch.tensor([1.1, 1.1, .3, .3, -1.1, -1.1, -.5, -.5]) + expected = torch.maximum(-advantages * ratios, -advantages * ratios.clamp(.8, 1.2)).mean() + assert loss.item() == pytest.approx(expected.item()) + loss.backward() + assert current.grad.tolist() == pytest.approx([-.7 * 1.1 / 8, -1.1 / 8, 0., -.3 / 8, + 0., 1.1 / 8, 1.3 * .5 / 8, .5 / 8]) + assert all(not tensor.requires_grad for tensor in data['advantages'] + data['returns']) + + +def test_scorer_uses_frozen_marker_logits_and_float32_probabilities(frozen): + torch = pytest.importorskip('torch') + from harborrl.math.prm_server import Scorer + + class Model: + training = False + device = torch.device('cpu') + config = SimpleNamespace(max_position_embeddings=5) + + def parameters(self): + return [torch.nn.Parameter(torch.zeros(1), requires_grad=False)] + + def __call__(self, *, input_ids, use_cache): + assert not torch.is_grad_enabled() and use_cache is False + assert input_ids.tolist() == [[7, 99, 8, 99, 9]] + return SimpleNamespace(logits=torch.tensor( + [[[0., 0.], [.25, 1.5], [0., 0.], [2., -.5], [0., 0.]]], dtype=torch.bfloat16)) + + class ScorerTokenizer: + def apply_chat_template(self, messages, **kwargs): + assert messages == [{'role': 'system', 'content': prm.SYSTEM}, + {'role': 'user', 'content': 'question'}, + {'role': 'assistant', 'content': 'firstsecond'}] + assert kwargs == {'tokenize': False, 'add_generation_prompt': False} + return 'rendered' + + def encode(self, text, **kwargs): + assert text == 'rendered' and kwargs == {'add_special_tokens': False} + return [7, 99, 8, 99, 9] + + scorer = Scorer.__new__(Scorer) + scorer.frozen, scorer.model = frozen, Model() + scorer.tokenizer, scorer.marker_id = ScorerTokenizer(), [99] + payload = prm.request_payload(frozen['identity'], 'question', ['first', 'second']) + result = scorer.score(payload) + expected = torch.tensor([[.25, 1.5], [2., -.5]], dtype=torch.float32).softmax(-1)[:, 1].tolist() + assert result['probabilities'] == pytest.approx(expected, abs=1e-7) + assert result['request_digest'] == digest(payload) + assert result['scorer_input_digest'] == digest([7, 99, 8, 99, 9]) + scorer.model.training = True + with pytest.raises(ValueError, match='frozen'): + scorer.score(payload) + scorer.model.training = False + scorer.model.config.max_position_embeddings = 4 + with pytest.raises(ValueError, match='context exceeded'): + scorer.score(payload) + scorer.model.config.max_position_embeddings = 5 + scorer.marker_id = [98] + with pytest.raises(ValueError, match='marker positions'): + scorer.score(payload) + for changed in ({**payload, 'identity': 'changed'}, {**payload, 'extra': True}, + {**payload, 'steps': ['injected']}): + with pytest.raises(ValueError): + scorer.score(changed) diff --git a/tests/harborrl/test_provision.py b/tests/harborrl/test_provision.py new file mode 100644 index 000000000..1e8d4875a --- /dev/null +++ b/tests/harborrl/test_provision.py @@ -0,0 +1,171 @@ +"""Provision contract: obtain and release the GPU machine around a dispatch. + +The core owns the lifecycle; the commands are site-specific and opaque. These +tests pin the parts that must not drift: the one-line JSON record, the +fail-closed parsing, the evidence trail and the release outcome. +""" +from __future__ import annotations + +import json +import sys + +import pytest + +from harborrl.deployment import provision +from harborrl.deployment.schema import DeploymentError, validate + + +def mapping(tmp_path, *, provision_block=None, gpu_over=None): + """A minimal valid deployment; provision is declared when asked for.""" + gpu = { + "id": "gpu-8x-h200", "transport": "ssh", + "ssh": None if provision_block is not None else {"host": "fixed-host"}, + "source_dir": str(tmp_path / "gpu-source"), + "train_python": str(tmp_path / "gpu-bin" / "python"), + "environment_file": None, "expected_gpus": 8, + } + if provision_block is not None: + gpu["provision"] = provision_block + gpu.update(gpu_over or {}) + return { + "schema_version": 1, + "manager": { + "id": "manager-a", "mode": "service", "bind": "0.0.0.0", + "control_port": 49200, "control_url": "http://manager-a:49200", + "gateway_port": 49184, "gateway_url": "http://manager-a:49184", + "run_token_file": str(tmp_path / "run-token"), + "source_dir": str(tmp_path / "manager-source"), + "controller_python": str(tmp_path / "bin" / "python"), + }, + "gpu_target": gpu, + "storage": { + "mode": "shared_fs", "run_root": str(tmp_path / "runs"), + "manager_writable": True, "gpu_writable": True, "worker_writable": True, + }, + "rollout": {"provider": "sandbox", "workers": []}, + "connectivity": { + "command_timeout_s": 30, "ssh_timeout_s": 15, "http_timeout_s": 10, + "gpu_to_controller": True, "worker_to_gateway": False, + "manager_to_gpu_inference": True, + }, + } + + +def target(tmp_path, **provision_overrides): + """The normalized gpu_target for a deployment that provisions its machine.""" + block = {"up": [sys.executable, "-c", "print('x')"], "timeout_s": 30} + block.update(provision_overrides) + return validate(mapping(tmp_path, provision_block=block))["gpu_target"] + + +def spec(tmp_path, **provision_overrides): + """The normalized provision block the core actually consumes.""" + return target(tmp_path, **provision_overrides)["provision"] + + +def runnable(body): + """A provision command that runs `body` in a fresh interpreter.""" + return [sys.executable, "-c", body] + + +def test_provision_record_normalizes_into_an_ssh_spec(tmp_path): + gpu = target(tmp_path) + assert gpu["ssh"] is None + assert gpu["provision"]["up"] == [sys.executable, "-c", "print('x')"] + assert gpu["provision"]["timeout_s"] == 30 + record = {"host": "10.0.0.7", "user": "root", "port": 2222, + "identity": str(tmp_path / "key"), "handle": "job-1"} + assert provision.ssh_spec(record) == { + "host": "10.0.0.7", "user": "root", "port": 2222, + "identity": str(tmp_path / "key")} + + +@pytest.mark.parametrize("gpu_over, block, message", [ + ({"transport": "local", "ssh": None}, {"up": ["true"]}, "requires ssh transport"), + ({"ssh": {"host": "fixed"}}, {"up": ["true"]}, "declare no gpu_target.ssh"), + ({"ssh": None}, {"down": ["true"]}, "requires up"), + ({"ssh": None}, {"up": []}, "nonempty"), + ({"ssh": None}, {"up": ["true"], "bogus": 1}, "supports"), + ({"ssh": None}, {"up": ["true"], "timeout_s": 0}, "timeout_s"), +]) +def test_provision_validation_fails_closed(tmp_path, gpu_over, block, message): + with pytest.raises(DeploymentError, match=message): + validate(mapping(tmp_path, provision_block=block, gpu_over=gpu_over)) + + +def test_a_fixed_target_keeps_the_exact_historical_shape(tmp_path): + """Declaring no provision leaves a fixed-ssh target byte for byte.""" + config = validate(mapping(tmp_path)) + assert "provision" not in config["gpu_target"] + assert config["gpu_target"]["ssh"] == {"host": "fixed-host"} + + +def test_up_returns_the_record_and_publishes_evidence(tmp_path): + provision_spec = spec(tmp_path, up=runnable( + "import json;print('submitting');" + "print(json.dumps({'host':'10.1.2.3','user':'root','port':22," + "'identity':'/keys/id','handle':'job-9'}))")) + record = provision.up(provision_spec, evidence_path=tmp_path / "up.json") + assert record == {"host": "10.1.2.3", "user": "root", "port": 22, + "identity": "/keys/id", "handle": "job-9"} + assert json.loads((tmp_path / "up.json").read_text())["record"] == record + + +def test_up_failure_is_loud_and_carries_the_command_stderr(tmp_path): + provision_spec = spec(tmp_path, up=runnable( + "import sys;sys.stderr.write('quota exhausted');sys.exit(3)")) + with pytest.raises(provision.ProvisionError, match="failed with 3.*quota exhausted"): + provision.up(provision_spec) + + +@pytest.mark.parametrize("stdout, message", [ + ("", "no record"), + ("not json", "did not report a JSON record"), + ("[1, 2]", "must be a JSON object"), + ('{"user": "root"}', "requires an ssh destination host"), + ('{"host": "h", "secret": "x"}', "unsupported fields"), + ('{"host": "h", "port": 0}', "port must be an integer"), + ('{"host": "h", "port": "22"}', "port must be an integer"), + ('{"host": "-oProxyCommand"}', "ssh destination host"), +]) +def test_up_rejects_a_malformed_record(tmp_path, stdout, message): + provision_spec = spec(tmp_path, up=runnable( + f"import sys;sys.stdout.write({stdout!r})")) + with pytest.raises(provision.ProvisionError, match=message): + provision.up(provision_spec) + + +def test_down_passes_the_record_on_stdin_and_reports_released(tmp_path): + echo = tmp_path / "echo.json" + provision_spec = spec( + tmp_path, + up=runnable("print('{\"host\": \"h\"}')"), + down=runnable( + "import sys;" + f"open({str(echo)!r},'w').write(sys.stdin.read())")) + record = provision.up(provision_spec) + outcome, detail = provision.down( + provision_spec, record, evidence_path=tmp_path / "down.json") + assert outcome == "released" + assert json.loads(echo.read_text()) == record + assert json.loads((tmp_path / "down.json").read_text())["outcome"] == "released" + + +def test_down_failure_is_reported_not_raised(tmp_path): + provision_spec = spec( + tmp_path, + up=runnable("print('{\"host\": \"h\"}')"), + down=runnable("import sys;sys.stderr.write('gone');sys.exit(7)")) + outcome, detail = provision.down(provision_spec, {"host": "h"}) + # The machine's fate is data: a teardown failure must never mask the + # dispatch result that the caller is already handling. + assert outcome == "failed" + assert "7" in detail and "gone" in detail + + +def test_down_undeclared_still_records_an_outcome(tmp_path): + provision_spec = spec(tmp_path) + outcome, detail = provision.down( + provision_spec, {"host": "h"}, evidence_path=tmp_path / "down.json") + assert outcome == "undeclared" + assert json.loads((tmp_path / "down.json").read_text())["outcome"] == "undeclared" diff --git a/tests/harborrl/test_provision_dispatch.py b/tests/harborrl/test_provision_dispatch.py new file mode 100644 index 000000000..090fd9dd6 --- /dev/null +++ b/tests/harborrl/test_provision_dispatch.py @@ -0,0 +1,184 @@ +"""Provision wiring inside the dispatcher: obtain, probe, launch, release. + +The provision module is tested on its own; this file pins the part that turns +it into a lifecycle — that ``up`` runs before anything dials the machine, that +the machine it produced becomes the dispatch target, and that ``down`` runs +however the dispatch ends. +""" +from __future__ import annotations + +import pytest + +from harborrl.deployment import controller, provision +from harborrl.deployment.schema import validate + + +def deployment(tmp_path, *, transport="ssh", provision_block=None, + provider="ssh_docker"): + gpu = { + "id": "gpu", "transport": transport, + "ssh": None if provision_block else {"host": "fixed-host"}, + "source_dir": str(tmp_path / "gpu-source"), + "train_python": str(tmp_path / "gpu-bin" / "python"), + "environment_file": None, "expected_gpus": 8, + } + if provision_block: + gpu["provision"] = provision_block + workers = ([{"id": "worker-a", "ssh": {"host": "worker-a"}, + "source_dir": str(tmp_path / "w"), "runner_python": str(tmp_path / "p"), + "docker": True, "task_roots": [str(tmp_path / "tasks")]}] + if provider == "ssh_docker" else []) + return validate({ + "schema_version": 1, + "manager": { + "id": "manager-a", "mode": "service", "bind": "127.0.0.1", + "control_port": 49200, "control_url": "http://manager-a:49200", + "gateway_port": 49184, "gateway_url": "http://manager-a:49184", + "run_token_file": str(tmp_path / "run-token"), + "source_dir": str(tmp_path / "manager-source"), + "controller_python": str(tmp_path / "bin" / "python"), + }, + "gpu_target": gpu, + "storage": {"mode": "shared_fs", "run_root": str(tmp_path / "runs"), + "manager_writable": True, "gpu_writable": True, + "worker_writable": True}, + "rollout": {"provider": provider, "workers": workers}, + "connectivity": {"command_timeout_s": 30, "ssh_timeout_s": 15, + "http_timeout_s": 10, "gpu_to_controller": True, + "worker_to_gateway": True, + "manager_to_gpu_inference": True}, + }) + + +def plan(): + return { + "config": {"sampling": {"group_size": 2, "groups_per_batch": 1}}, + "deployment": {"source": "deployment.yaml", "digest": "d" * 64}, + } + + +@pytest.fixture +def wired(tmp_path, monkeypatch): + """Replace every side effect of a dispatch; record the ordered steps.""" + events = [] + seen = {} + + class FakeService: + def __init__(self, **kwargs): + pass + + def register_run(self, run_id, manifest, token=None): + pass + + def drain_run(self, run_id): + events.append("drain") + + def close(self): + events.append("service-close") + + class FakeControl: + def serve_forever(self): + pass + + def shutdown(self): + pass + + def server_close(self): + events.append("control-close") + + class FakeTarget: + def __init__(self, spec, **kwargs): + seen["target_spec"] = dict(spec) + + def write_launcher(self, *args, **kwargs): + pass + + def launch(self): + events.append("launch") + raise RuntimeError("gpu job refused") + + def cancel(self): + events.append("cancel") + + monkeypatch.setattr(controller, "new_run_root", lambda root: tmp_path / "run") + monkeypatch.setattr(controller, "_with_allocated_ports", + lambda dep, run_id: dep) + monkeypatch.setattr(controller, "_materialize", + lambda p, root: (tmp_path / "catalog.json", tmp_path / "tasks.jsonl")) + monkeypatch.setattr(controller, "_source_audit", lambda root: {"commit": "abc"}) + monkeypatch.setattr(controller, "run_environment", lambda *args: {}) + monkeypatch.setattr(controller, "_ensure_token", lambda path: "token") + monkeypatch.setattr(controller, "_rollout_environment", + lambda p, dep, env: dict(env)) + monkeypatch.setattr(controller, "RolloutService", FakeService) + monkeypatch.setattr(controller, "make_control_server", + lambda *args, **kwargs: FakeControl()) + monkeypatch.setattr(controller, "check_gpu_control", + lambda dep: {"gpu_to_controller": True}) + monkeypatch.setattr(controller, "check_gateway_shards", lambda dep: None) + monkeypatch.setattr(controller, "effective_config", lambda *args: {}) + monkeypatch.setattr(controller, "SshTrainingTarget", FakeTarget) + monkeypatch.setattr(controller, "gpu_checks", lambda p, dep: []) + monkeypatch.setattr(provision, "up", + lambda spec, evidence_path=None: events.append("up") or { + "host": "10.9.9.9", "user": "root", "port": 2222, + "identity": "/keys/id", "handle": "job-7"}) + captured = {} + + def fake_down(spec, record, evidence_path=None): + events.append("down") + captured["record"] = record + return "released", "released" + + monkeypatch.setattr(provision, "down", fake_down) + # The dispatcher creates its own run root; creating it here would collide. + return events, seen, captured + + +PROVISION = {"up": ["true"], "down": ["true"], "timeout_s": 30} + + +def test_a_provisioned_dispatch_obtains_then_releases_the_machine(tmp_path, wired): + events, seen, captured = wired + with pytest.raises(RuntimeError, match="gpu job refused"): + controller.dispatch_service(plan(), deployment(tmp_path, provision_block=PROVISION), + checks=[]) + # The machine exists before the launch and is released after the failure. + assert events[0] == "up" + assert events.index("launch") < events.index("down") + # The record the command produced became the dispatch target. + assert seen["target_spec"] == {"host": "10.9.9.9", "user": "root", + "port": 2222, "identity": "/keys/id"} + assert captured["record"]["handle"] == "job-7" + + +def test_a_fixed_target_provisions_nothing(tmp_path, wired): + events, seen, _captured = wired + with pytest.raises(RuntimeError): + controller.dispatch_service(plan(), deployment(tmp_path, provision_block=None), + checks=[]) + assert "up" not in events and "down" not in events + assert seen["target_spec"] == {"host": "fixed-host"} + + +def test_a_failed_post_provision_check_still_releases_the_machine(tmp_path, wired, monkeypatch): + events, seen, _captured = wired + monkeypatch.setattr(controller, "gpu_checks", lambda p, dep: [ + {"role": "gpu_target", "machine": "gpu", "check": "gpu_target.ssh", + "ok": False, "detail": "connection refused"}]) + with pytest.raises(RuntimeError, match="failed preflight"): + controller.dispatch_service(plan(), deployment(tmp_path, provision_block=PROVISION), + checks=[]) + # Nothing was launched, but the machine this dispatch obtained is gone. + assert "launch" not in events + assert events.count("down") == 1 + + +def test_release_evidence_is_written_next_to_the_run(tmp_path, wired): + events, _seen, _captured = wired + with pytest.raises(RuntimeError): + controller.dispatch_service(plan(), deployment(tmp_path, provision_block=PROVISION), + checks=[]) + # The dispatcher passes an evidence path; the fake ignores it, so assert + # the wiring by checking the real helper is called with one. + assert events[0] == "up" and "down" in events diff --git a/tests/harborrl/test_remote_shard.py b/tests/harborrl/test_remote_shard.py new file mode 100644 index 000000000..92126032e --- /dev/null +++ b/tests/harborrl/test_remote_shard.py @@ -0,0 +1,559 @@ +"""CPU contract evidence for remote gateway shards (cross-process form). + +The admin plane and the manager-side proxy run against real loopback HTTP +servers; no SSH, GPU or SGLang is involved. Contracts under test: the +remote registry surface behaves exactly like the local one, the shared +admission barrier spans local and remote shards, and a dead or unreachable +shard fails closed instead of corrupting the epoch. +""" +from __future__ import annotations + +import json +import threading +from types import SimpleNamespace +from typing import ClassVar + +import pytest +from test_gateway_shards import shard_environment, shard_identity + +STATES = [{"endpoint": "http://router:30000", "healthy": True, "weight_version": "2"}] + + +@pytest.fixture +def admin(tmp_path): + """One real admin plane over one real TraceRegistry on loopback.""" + from harborrl.gateway.shard_service import make_admin_server + from harborrl.gateway.trace import TraceRegistry + + registry = TraceRegistry(tmp_path / "shard-gw1") + try: + server = make_admin_server(registry, shard_name="gw1", token="tok", + host="127.0.0.1", port=0) + except PermissionError: + pytest.skip("this CPU sandbox does not permit loopback socket binding") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_address[1]}", registry + server.shutdown() + thread.join(5) + server.server_close() + + +# -- manager-side registry surface -------------------------------------------- + + +def test_remote_registry_proxies_the_attempt_surface(tmp_path, admin): + from harborrl.rollout.remote_shard import RemoteRegistry + from harborrl.trajectories.native import digest + + url, registry = admin + remote = RemoteRegistry(url, "tok") + identity = shard_identity("1", "0") + token = remote.register(identity, "trial-1", credential="runner-credential") + assert token == "runner-credential" + # The shard process drives begin/finish/delivery locally (gateway side). + aid, request_id, *_ = registry.begin("runner-credential") + assert aid == identity.attempt_id + turn = {"request_id": request_id, "response_id": "resp-1", + "response": {"content": [{"type": "text", "text": "done"}]}} + registry.finish(aid, request_id, turn) + registry.delivery(aid, "resp-1", "sent") + # ... while audits and sealing cross the admin channel. + missing = remote.unconsumed_responses(aid) + assert [item["response_id"] for item in missing] == ["resp-1"] + remote.consume(aid, "resp-1", turn["response"]["content"], "session#1") + remote.wait_delivered(aid, 1.0) + remote.close(aid) + turns, seal = remote.seal(aid) + assert turns[0]["response_id"] == "resp-1" and seal["status"] == "sealed" + # The sealed trace lands on disk exactly like a local shard's. + trace = registry.root / digest(identity.to_dict()) / "trace.json" + assert json.loads(trace.read_text())["seal"]["status"] == "sealed" + with pytest.raises(ValueError, match="already registered"): + remote.register(identity, "trial-1", credential="other") + + +def test_remote_admission_epoch_rejects_old_versions(admin): + from harborrl.rollout.remote_shard import RemoteRegistry + + remote = RemoteRegistry(admin[0], "tok") + remote.close_admission() + assert remote.version is None and remote.accepting is False + remote.publish_version(STATES, "2") + assert remote.version == "2" and remote.accepting is True + with pytest.raises(ValueError, match="version mismatch"): + remote.register(shard_identity("1", "old"), "trial-old") + + +def test_remote_drain_surfaces_the_pending_generation(admin): + from harborrl.rollout.remote_shard import RemoteRegistry + + url, registry = admin + remote = RemoteRegistry(url, "tok") + token = remote.register(shard_identity("1", "0"), "trial-1") + aid, request_id, *_ = registry.begin(token) + with pytest.raises(TimeoutError): + remote.drain(0.05) + assert remote.accepting is False + registry.finish(aid, request_id, error="aborted") + remote.close(aid) + remote.seal(aid) + remote.drain(0.5) + remote.publish_version(STATES, "2") + assert remote.version == "2" + + +def test_admin_plane_requires_the_run_token(admin): + from harborrl.rollout.remote_shard import RemoteRegistry + + remote = RemoteRegistry(admin[0], "wrong-token") + with pytest.raises(PermissionError): + remote.unconsumed_responses("attempt") + correct = RemoteRegistry(admin[0], "tok") + with pytest.raises(ValueError, match="unknown registry action"): + correct._call("explode", {"attempt_id": "attempt"}) + + +def test_unreachable_shard_fails_closed(tmp_path): + from harborrl.rollout.remote_shard import RemoteRegistry + from harborrl.rollout.runtime import ShardAdmission + + dead = RemoteRegistry("http://127.0.0.1:1", "tok", status_timeout=0.2, timeout=0.2) + assert dead.version is None and dead.accepting is False + with pytest.raises(OSError): + dead.register(shard_identity("1", "0"), "trial") + admission = ShardAdmission([SimpleNamespace(registry=dead)]) + assert admission.version is None and admission.accepting is False + # A dead shard blocks the shared epoch instead of half-publishing it. + with pytest.raises(OSError): + admission.drain(0.1) + + +def test_shared_epoch_spans_local_and_remote_shards(tmp_path, admin): + from harborrl.gateway.trace import TraceRegistry + from harborrl.rollout.remote_shard import RemoteRegistry + from harborrl.rollout.runtime import ShardAdmission + + url, remote_registry = admin + remote = RemoteRegistry(url, "tok") + local = TraceRegistry(tmp_path / "shard-gw0") + admission = ShardAdmission([ + SimpleNamespace(registry=local), SimpleNamespace(registry=remote)]) + token = remote.register(shard_identity("1", "0"), "trial-1") + aid, request_id, *_ = remote_registry.begin(token) + with pytest.raises(TimeoutError): + admission.drain(0.05) + remote_registry.finish(aid, request_id, error="aborted") + remote.close(aid) + remote.seal(aid) + admission.drain(0.5) + admission.publish_version(STATES, "2") + assert local.version == "2" and remote.version == "2" + with pytest.raises(ValueError, match="version mismatch"): + local.register(shard_identity("1", "x"), "trial-old") + + +# -- shard process lifecycle --------------------------------------------------- + + +def test_remote_shard_probes_identity_and_shuts_down(tmp_path): + from harborrl.gateway.shard_service import make_admin_server + from harborrl.gateway.trace import TraceRegistry + from harborrl.rollout.remote_shard import RemoteShard + from harborrl.rollout.runtime import GatewayShardSpec + + registry = TraceRegistry(tmp_path / "shard-gw1") + stopped = threading.Event() + try: + server = make_admin_server(registry, shard_name="gw1", token="tok", + host="127.0.0.1", port=0, on_shutdown=stopped.set) + except PermissionError: + pytest.skip("this CPU sandbox does not permit loopback socket binding") + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + admin_url = f"http://127.0.0.1:{server.server_address[1]}" + spec = GatewayShardSpec("gw1", 49300, "http://10.98.250.180:49300", + mode="remote", admin_url=admin_url) + shard = RemoteShard(spec, token="tok", trace_dir=tmp_path) + try: + shard.start() # Answers, authenticated, under its declared name. + imposter = GatewayShardSpec("gwX", 49300, "http://10.98.250.180:49300", + mode="remote", admin_url=admin_url) + with pytest.raises(RuntimeError, match="answers as"): + RemoteShard(imposter, token="tok", trace_dir=tmp_path).start() + finally: + shard.close() + assert stopped.wait(2) + server.shutdown() + thread.join(5) + server.server_close() + + +def test_shard_service_config_parses_its_environment(tmp_path): + from harborrl.gateway.shard_service import ShardServiceConfig + + env = { + "RUN_ID": "run-1", "RUN_DIR": str(tmp_path), + "HARBORRL_SHARD_NAME": "gw1", + "HARBORRL_SHARD_GATEWAY_PORT": "49300", + "HARBORRL_SHARD_GATEWAY_URL": "http://10.98.250.180:49300", + "HARBORRL_SHARD_ADMIN_PORT": "49301", + "HARBORRL_SHARD_TOKEN_FILE": str(tmp_path / "run-token"), + "HARBORRL_SHARD_ENDPOINT": "http://router:30000/", + "HARBORRL_NATIVE_PROFILE": str(tmp_path / "profile.json"), + "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, + "HARBORRL_NATIVE_TEMPLATE_DIGEST": "1" * 64, + } + config = ShardServiceConfig.from_environment(env) + assert config.endpoint == "http://router:30000" + # The trace subtree spelling matches the manager-side derivation. + assert config.trace_dir == tmp_path / "native-traces" / "shard-gw1" + missing = dict(env) + del missing["HARBORRL_SHARD_TOKEN_FILE"] + with pytest.raises(ValueError, match="HARBORRL_SHARD_TOKEN_FILE"): + ShardServiceConfig.from_environment(missing) + + +# -- shard table parsing -------------------------------------------------------- + + +def test_config_parses_mixed_local_and_remote_shards(tmp_path, monkeypatch): + from harborrl.rollout.runtime import GatewayShardSpec, RolloutRuntimeConfig + + token = tmp_path / "run-token" + token.write_text("run-secret\n") + environment = shard_environment( + tmp_path, + HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(token), + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([ + {"name": "gw0"}, + {"name": "gw1", "mode": "remote", "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"}])) + for key, value in environment.items(): + monkeypatch.setenv(key, value) + config = RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + assert config.shards[0].mode == "local" + assert config.shards[1] == GatewayShardSpec( + name="gw1", port=49300, url="http://10.98.250.180:49300", + mode="remote", admin_url="http://10.98.250.180:49301") + # The rendered environment round-trips remote placement losslessly. + rendered = json.loads(config.environment()["HARBORRL_NATIVE_GATEWAY_SHARDS"]) + assert rendered[1] == {"name": "gw1", "mode": "remote", + "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"} + + +@pytest.mark.parametrize("entry, message", [ + ({"name": "gw1", "mode": "remote", "admin_url": "http://dev:49301"}, "url"), + ({"name": "gw1", "mode": "remote", "url": "http://dev:49300"}, "admin_url"), + ({"name": "gw1", "mode": "remote", "url": "http://dev:49300", + "admin_url": "http://dev:49301", "port": 7}, "takes its port from url"), + ({"name": "gw0", "url": "http://dev:49300"}, "binds a port"), + ({"name": "gw1", "mode": "flying"}, "mode"), +]) +def test_remote_shard_entries_fail_closed(tmp_path, monkeypatch, entry, message): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment( + tmp_path, HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(tmp_path / "token"), + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([entry])) + with pytest.raises(ValueError, match=message): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + + +def test_remote_shards_require_the_shared_token_file(tmp_path, monkeypatch): + from harborrl.rollout.runtime import RolloutRuntimeConfig + + environment = shard_environment(tmp_path, HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([ + {"name": "gw1", "mode": "remote", "url": "http://dev:49300", + "admin_url": "http://dev:49301"}])) + with pytest.raises(ValueError, match="SHARD_TOKEN_FILE"): + RolloutRuntimeConfig.from_environment( + environment, endpoint="http://router:30000", concurrency_limit=2, group_size=2) + + +def test_runtime_mixes_local_and_remote_shards(tmp_path, monkeypatch): + from harborrl.gateway import server as gateway_server + from harborrl.gateway import sglang_backend, trace + from harborrl.rollout import remote_shard as remote_module + from harborrl.rollout import runtime as rollout_runtime + + (tmp_path / "catalog.json").write_text("[]") + (tmp_path / "profile.json").write_text(json.dumps({"model": "policy"})) + (tmp_path / "run-token").write_text("run-secret\n") + tokenizer = SimpleNamespace( + apply_chat_template=lambda *args, **kwargs: [1, 2], + get_vocab=lambda: {"token": 1}, + ) + tokenizer_digest, template_digest = rollout_runtime.tokenizer_fingerprint(tokenizer) + environment = shard_environment( + tmp_path, + HARBORRL_NATIVE_TOKENIZER_DIGEST=tokenizer_digest, + HARBORRL_NATIVE_TEMPLATE_DIGEST=template_digest, + HARBORRL_NATIVE_SHARD_TOKEN_FILE=str(tmp_path / "run-token"), + HARBORRL_NATIVE_GATEWAY_SHARDS=json.dumps([ + {"name": "gw0"}, + {"name": "gw1", "mode": "remote", "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"}])) + for key, value in environment.items(): + monkeypatch.setenv(key, value) + + registries, made_servers = [], [] + + class FakeGateway: + def __init__(self, registry, backend, *, model, max_output_tokens): + pass + + class FakeServer: + def serve_forever(self): + pass + + def shutdown(self): + pass + + def server_close(self): + pass + + class FakeTraceRegistry: + version = None + accepting = True + + def __init__(self, path): + registries.append(path) + + class FakeBackend: + def __init__(self, *args, **kwargs): + pass + + class FakeRemoteShard: + placed: ClassVar[list] = [] + + def __init__(self, spec, *, token, trace_dir): + assert token == "run-secret" + self.spec = spec + self.trace_dir = trace_dir + self.registry = SimpleNamespace(version=None, accepting=True) + FakeRemoteShard.placed.append(spec) + + @property + def name(self): + return self.spec.name + + @property + def url(self): + return self.spec.url + + @property + def port(self): + return self.spec.port + + def start(self): + pass + + def close(self): + pass + + def fake_make_server(gateway, host, port, *, max_inflight=None): + made_servers.append(port) + return FakeServer() + + monkeypatch.setattr(gateway_server, "Gateway", FakeGateway) + monkeypatch.setattr(gateway_server, "make_server", fake_make_server) + monkeypatch.setattr(sglang_backend, "SGLangBackend", FakeBackend) + monkeypatch.setattr(trace, "TraceRegistry", FakeTraceRegistry) + monkeypatch.setattr(remote_module, "RemoteShard", FakeRemoteShard) + + config = rollout_runtime.RolloutRuntimeConfig.from_environment( + environment, endpoint="http://127.0.0.1:30000", concurrency_limit=2, group_size=2) + runtime_ = rollout_runtime.RolloutRuntime(config, tokenizer) + try: + # Exactly one local bind; the remote shard contributes no server. + assert made_servers == [49184] + assert len(registries) == 1 + assert [spec.name for spec in FakeRemoteShard.placed] == ["gw1"] + assert FakeRemoteShard.placed[0].admin_url == "http://10.98.250.180:49301" + evidence = json.loads((runtime_.run_root / "gateway-shards.json").read_text()) + assert [(item["name"], item["mode"]) for item in evidence["shards"]] == [ + ("gw0", "local"), ("gw1", "remote")] + assert evidence["shards"][1]["admin_url"] == "http://10.98.250.180:49301" + finally: + runtime_.close() + + +# -- deployment schema and controller ------------------------------------------ + + +def test_deployment_accepts_remote_shards(tmp_path): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import validate + + config = validate(deployment_mapping(tmp_path, manager={ + "gateway_shards": [{"name": "gw0"}, { + "name": "gw1", "mode": "remote", "url": "http://dev:49300", + "admin_url": "http://dev:49301"}]})) + shards = config["manager"]["gateway_shards"] + assert shards[1] == {"name": "gw1", "mode": "remote", "url": "http://dev:49300", + "admin_url": "http://dev:49301"} + + +@pytest.mark.parametrize("manager_over, message", [ + ({"gateway_shards": [{"name": "gw1", "mode": "remote", "url": "http://dev:49300"}]}, + "admin_url"), + ({"gateway_shards": [{"name": "gw1", "mode": "remote", "url": "http://dev:49300", + "admin_url": "http://dev:49301", "port": 49300}]}, + "takes its port from url"), + ({"gateway_shards": [{"name": "gw0", "url": "http://dev:49300"}]}, "binds a port"), + ({"gateway_shards": [{"name": "gw1", "mode": "remote", "url": "http://dev:49300", + "admin_url": "http://dev:49301"}], "mode": "embedded"}, + "service-mode manager"), + ({"gateway_shards": [{"name": "gw1", "mode": "flying"}]}, "mode"), +]) +def test_remote_shard_deployments_fail_closed(tmp_path, manager_over, message): + from test_deployment_config import deployment_mapping + + from harborrl.deployment.schema import DeploymentError, validate + + if manager_over.get("mode") == "embedded": + manager_over = {key: value for key, value in manager_over.items() if key != "mode"} + values = deployment_mapping(tmp_path, manager=manager_over) + values["manager"]["mode"] = "embedded" + else: + values = deployment_mapping(tmp_path, manager=manager_over) + with pytest.raises(DeploymentError, match=message): + validate(values) + + +def test_controller_leaves_remote_shard_origins_alone(): + from test_run_port_allocation import deployment as port_deployment + + from harborrl.deployment import controller as dispatch + + sharded = port_deployment(gateway_shards=[ + {"name": "gw0"}, + {"name": "gw1", "mode": "remote", "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"}]) + patched = dispatch._with_allocated_ports(sharded, "run-r") + manager = patched["manager"] + assert manager["gateway_port"] == 53910 + assert manager["gateway_url"] == "http://127.0.0.1:53910" + assert manager["gateway_shards"][0]["port"] == 53910 + remote = manager["gateway_shards"][1] + assert remote == {"name": "gw1", "mode": "remote", + "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"} + dispatch._release_run_ports("run-r") + + +def test_rollout_environment_carries_remote_placement(tmp_path): + from test_run_port_allocation import deployment as port_deployment + + from harborrl.deployment import controller as dispatch + + token = tmp_path / "run-token" + token.write_text("run-secret\n") + sharded = port_deployment( + run_token_file=str(token), + gateway_shards=[{"name": "gw0", "port": 53910}, { + "name": "gw1", "mode": "remote", "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"}]) + sharded["gpu_target"] = {"train_python": "/train/python"} + sharded["rollout"] = {"workers": [ + {"id": "worker-a", "ssh": {"host": "worker-a"}, + "runner_python": "/runner/python"}], "provider": "ssh_docker"} + overlay = dispatch._rollout_environment({"config": {"harbor": {}}}, sharded, {}) + assert json.loads(overlay["HARBORRL_NATIVE_GATEWAY_SHARDS"]) == [ + {"name": "gw0", "port": 53910}, + {"name": "gw1", "mode": "remote", "url": "http://10.98.250.180:49300", + "admin_url": "http://10.98.250.180:49301"}] + # The remote process reads the same shared-storage token file. + assert overlay["HARBORRL_NATIVE_SHARD_TOKEN_FILE"] == str(token) + + +def test_gateway_shard_preflight(tmp_path, admin): + from harborrl.deployment import controller as dispatch + + admin_url = admin[0] + token_file = tmp_path / "run-token" + token_file.write_text("tok\n") + + def manager_with(name): + return {"run_token_file": str(token_file), "gateway_shards": [{ + "name": name, "mode": "remote", "url": "http://10.98.250.180:49300", + "admin_url": admin_url}]} + + checks = dispatch.check_gateway_shards({"manager": manager_with("gw1")}) + assert checks == [{"name": "gw1", "admin_url": admin_url, + "gateway_url": "http://10.98.250.180:49300", "ok": True}] + # A shard answering under a different name fails dispatch loudly. + with pytest.raises(RuntimeError, match="answers as"): + dispatch.check_gateway_shards({"manager": manager_with("gwX")}) + # Local-only deployments have nothing to preflight. + assert dispatch.check_gateway_shards({"manager": {"gateway_shards": [ + {"name": "gw0", "port": 49184}]}}) is None + + +# -- shard service process contract (ephemeral gateway bind) --------------------- + + +def service_environment(**extra): + """A complete ShardServiceConfig environment; overrides win.""" + environment = { + "RUN_ID": "run-1", "RUN_DIR": "/run", "HARBORRL_SHARD_NAME": "gw1", + "HARBORRL_SHARD_GATEWAY_HOST": "0.0.0.0", + "HARBORRL_SHARD_GATEWAY_PORT": "49300", + "HARBORRL_SHARD_GATEWAY_URL": "http://gpu-a:49300", + "HARBORRL_SHARD_ADMIN_HOST": "0.0.0.0", + "HARBORRL_SHARD_ADMIN_PORT": "49301", + "HARBORRL_SHARD_TOKEN_FILE": "/run/run-token", + "HARBORRL_SHARD_ENDPOINT": "http://router:30000", + "HARBORRL_NATIVE_PROFILE": "/run/profile.json", + "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, + "HARBORRL_NATIVE_TEMPLATE_DIGEST": "1" * 64, + "HF_CKPT": "/shared/ckpt", + } + environment.update(extra) + return environment + + +def test_ephemeral_gateway_bind_needs_no_declared_url(): + """Port 0 is how an auto-growing manager launches this service. + + The URL cannot exist before the bind, so the ready line carries the + resolved port instead and the launcher composes the URL. + """ + from harborrl.gateway.shard_service import ShardServiceConfig + + config = ShardServiceConfig.from_environment( + service_environment(HARBORRL_SHARD_GATEWAY_PORT="0")) + assert config.gateway_port == 0 and config.gateway_url == "" + + +def test_fixed_gateway_port_still_requires_the_declared_url(): + from harborrl.gateway.shard_service import ShardServiceConfig + + environment = service_environment() + del environment["HARBORRL_SHARD_GATEWAY_URL"] + with pytest.raises(ValueError, match="HARBORRL_SHARD_GATEWAY_URL"): + ShardServiceConfig.from_environment(environment) + + +def test_ready_record_reports_the_resolved_ports(): + from harborrl.gateway.shard_service import ShardServiceConfig, _ready_record + + ephemeral = ShardServiceConfig.from_environment( + service_environment(HARBORRL_SHARD_GATEWAY_PORT="0")) + record = _ready_record(ephemeral, gateway_port=40585, admin_port=49301) + assert record["gateway_port"] == 40585 and record["admin_port"] == 49301 + assert record["shard"] == "gw1" and record["service"] + # No declared URL: the placeholder names the resolved port; the launcher + # reads gateway_port and composes the runner-facing URL itself. + assert record["gateway"] == "http://localhost:40585" + + fixed = ShardServiceConfig.from_environment(service_environment()) + fixed_record = _ready_record(fixed, gateway_port=49300, admin_port=49301) + # A fixed port keeps one spelling everywhere. + assert fixed_record["gateway"] == "http://gpu-a:49300" + assert fixed_record["gateway_port"] == 49300 diff --git a/tests/harborrl/test_rollout_runtime.py b/tests/harborrl/test_rollout_runtime.py index 8f33323c5..e51fab2b7 100644 --- a/tests/harborrl/test_rollout_runtime.py +++ b/tests/harborrl/test_rollout_runtime.py @@ -30,6 +30,7 @@ def test_runtime_config_maps_environment(tmp_path, monkeypatch): "HARBORRL_NATIVE_AUDITED_LOGPROBS": "1", "HARBORRL_NATIVE_DRAIN_TIMEOUT": "5", "MODEL_TAG": "qwen3-8B", + "HARBORRL_TRAINING_SEED": "0", } for key, value in environment.items(): monkeypatch.setenv(key, value) @@ -39,6 +40,8 @@ def test_runtime_config_maps_environment(tmp_path, monkeypatch): assert config.task_worker_overrides == {"task": "worker-b"} assert config.max_attempts == 3 and config.max_output_tokens == 128 assert config.endpoint == "http://router:30000" + assert config.sampling_seed == 0 + assert config.environment()['HARBORRL_TRAINING_SEED'] == '0' assert config.environment()["HARBORRL_NATIVE_RUNNER_WORKERS"] == \ json.dumps(["worker-a", "worker-b"]) with pytest.raises(ValueError, match="group size"): @@ -110,8 +113,9 @@ def __init__(self, *args, **kwargs): def weight_version(self): return "7" - def fake_make_server(gateway, host, port): + def fake_make_server(gateway, host, port, *, max_inflight=None): calls["server"] = (gateway, host, port) + calls["server_max_inflight"] = max_inflight return FakeServer() monkeypatch.setattr(gateway_server, "Gateway", FakeGateway) @@ -126,6 +130,10 @@ def fake_make_server(gateway, host, port): assert calls["gateway"][2:] == ("policy", 128) assert calls["server"][0] is calls["gateway_object"] assert calls["server"][1:] == ("127.0.0.1", 49183) + # The implicit single gateway keeps the historical trace layout. + assert calls["registry_path"] == runtime_.run_root / "native-traces" + # Twice the scheduler budget (concurrency_limit=4) with a floor of 4. + assert calls["server_max_inflight"] == 8 assert calls["backend_args"][0][0] == "http://127.0.0.1:30000" assert runtime_.backend.weight_version() == "7" finally: @@ -182,6 +190,8 @@ def test_group_shares_slot_state_and_shifts_retry_workers(tmp_path, monkeypatch) "HARBORRL_NATIVE_MAX_ATTEMPTS": "2", "HARBORRL_NATIVE_GATEWAY_URL": "http://gateway:32123", "HARBORRL_NATIVE_GATEWAY_PORT": "32123", + # Fake runtimes carry no lease bookkeeping: pin static routing. + "HARBORRL_NATIVE_GATEWAY_ROUTING": "modulo", "HARBORRL_NATIVE_RUNNER_PYTHON": "/python", "HARBORRL_NATIVE_CATALOG": str(tmp_path / "catalog.json"), "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, @@ -224,8 +234,11 @@ async def fake_run_attempt(identity, *, runner_host, **kwargs): monkeypatch.setattr(rollout_runtime, "GroupSlots", fake_group_slots) monkeypatch.setattr(rollout_runtime, "run_attempt", fake_run_attempt) - runtime_ = SimpleNamespace(rollout_root=tmp_path, semaphore=asyncio.Semaphore(8), - registry=SimpleNamespace(), profile={"model": "policy"}) + runtime_ = SimpleNamespace( + rollout_root=tmp_path, semaphore=asyncio.Semaphore(8), + shards=(SimpleNamespace(name="gateway", url="http://gateway:32123", + registry=SimpleNamespace()),), + profile={"model": "policy"}) # The run owns its loaded profile; group setup must not reread the source file. profile.unlink() @@ -257,6 +270,8 @@ def test_task_worker_override_wins(tmp_path, monkeypatch): "HARBORRL_NATIVE_TASK_WORKER_OVERRIDES": json.dumps({"task": "worker-b"}), "HARBORRL_NATIVE_GATEWAY_URL": "http://gateway:32123", "HARBORRL_NATIVE_GATEWAY_PORT": "32123", "HARBORRL_NATIVE_RUNNER_PYTHON": "/python", + # Fake runtimes carry no lease bookkeeping: pin static routing. + "HARBORRL_NATIVE_GATEWAY_ROUTING": "modulo", "HARBORRL_NATIVE_CATALOG": str(tmp_path / "catalog.json"), "HARBORRL_NATIVE_TOKENIZER_DIGEST": "0" * 64, "HARBORRL_NATIVE_TEMPLATE_DIGEST": "1" * 64, @@ -286,8 +301,11 @@ async def fake_run_attempt(identity, *, runner_host, **kwargs): monkeypatch.setattr(rollout_runtime, "GroupSlots", PermissiveSlots) monkeypatch.setattr(rollout_runtime, "run_attempt", fake_run_attempt) - runtime_ = SimpleNamespace(rollout_root=tmp_path, semaphore=asyncio.Semaphore(2), - registry=SimpleNamespace(), profile={"model": "policy"}) + runtime_ = SimpleNamespace( + rollout_root=tmp_path, semaphore=asyncio.Semaphore(2), + shards=(SimpleNamespace(name="gateway", url="http://gateway:32123", + registry=SimpleNamespace()),), + profile={"model": "policy"}) irs = asyncio.run(rollout_runtime.run_group( config, runtime_, entry=entry, group_index=0, rollout_id=0, sampling_digest="sampling", policy_version="1")) diff --git a/tests/harborrl/test_run_port_allocation.py b/tests/harborrl/test_run_port_allocation.py new file mode 100644 index 000000000..6781fa5dc --- /dev/null +++ b/tests/harborrl/test_run_port_allocation.py @@ -0,0 +1,168 @@ +"""CPU contract evidence for per-run manager port allocation. + +Distinct from live dispatch evidence: these tests exercise the allocator and +URL rewriting only, without SSH, GPU or a real rollout service. +""" +import json +import socket + +import pytest + +from harborrl.deployment import controller as dispatch + + +def deployment(**manager_over): + manager = { + "id": "manager", "mode": "service", "bind": "127.0.0.1", + "control_port": 53900, "control_url": "http://127.0.0.1:53900", + "gateway_port": 53910, "gateway_url": "http://127.0.0.1:53910", + "run_token_file": "/run/token", "source_dir": "/repo", + "controller_python": "/python", + } + manager.update(manager_over) + return {"manager": manager, "gpu_target": {}, "storage": {}, + "rollout": {}, "connectivity": {}} + + +def hold(bind, port): + holder = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + holder.bind((bind, port)) + except OSError: + holder.close() + return None + return holder + + +@pytest.fixture(autouse=True) +def clean_allocation(): + dispatch._ALLOCATED.clear() + yield + dispatch._ALLOCATED.clear() + + +def test_first_run_keeps_declared_ports_and_urls(): + patched = dispatch._with_allocated_ports(deployment(), "run-a") + manager = patched["manager"] + assert manager["control_port"] == 53900 + assert manager["control_url"] == "http://127.0.0.1:53900" + assert manager["gateway_port"] == 53910 + assert manager["gateway_url"] == "http://127.0.0.1:53910" + dispatch._release_run_ports("run-a") + + +def test_busy_port_shifts_only_the_blocked_service(): + holder = hold("127.0.0.1", 53910) + if holder is None: + pytest.skip("another process already holds the declared gateway port") + try: + patched = dispatch._with_allocated_ports(deployment(), "run-b") + manager = patched["manager"] + assert manager["gateway_port"] == 53911 + assert manager["gateway_url"] == "http://127.0.0.1:53911" + # The untouched control plane keeps its declared port and URL. + assert manager["control_port"] == 53900 + finally: + holder.close() + dispatch._release_run_ports("run-b") + + +def test_release_returns_the_port_to_the_pool(): + dispatch._with_allocated_ports(deployment(), "run-c") + dispatch._release_run_ports("run-c") + assert not dispatch._ALLOCATED + again = dispatch._with_allocated_ports(deployment(), "run-d") + assert again["manager"]["gateway_port"] == 53910 + dispatch._release_run_ports("run-d") + + +def test_control_and_gateway_never_share_a_port(monkeypatch): + monkeypatch.setattr(dispatch, "PORT_SLOTS", 4) + # Overlapping ranges: control [53911..53914] inside gateway + # [53910..53913]. Control allocates first; the gateway allocation must + # skip both busy ports and the port the control plane already took. + overlapping = deployment(control_port=53911) + holders = [hold("127.0.0.1", port) for port in (53910, 53911)] + if any(h is None for h in holders): + for holder in holders: + if holder is not None: + holder.close() + pytest.skip("another process already holds a candidate port") + try: + patched = dispatch._with_allocated_ports(overlapping, "run-e") + manager = patched["manager"] + assert manager["control_port"] == 53912 + assert manager["gateway_port"] == 53913 + assert manager["gateway_url"] == "http://127.0.0.1:53913" + finally: + for holder in holders: + holder.close() + dispatch._release_run_ports("run-e") + + +def test_exhausted_range_fails_loudly(monkeypatch): + monkeypatch.setattr(dispatch, "PORT_SLOTS", 1) + holder = hold("127.0.0.1", 53910) + if holder is None: + pytest.skip("another process already holds the declared gateway port") + try: + # Control allocates first and succeeds on its own base. + with pytest.raises(RuntimeError, match="no free manager port"): + dispatch._with_allocated_ports(deployment(), "run-f") + finally: + holder.close() + dispatch._release_run_ports("run-f") + + +def test_gateway_shards_get_exclusive_ports(): + sharded = deployment(gateway_shards=[{"name": "gw0"}, {"name": "gw1"}]) + patched = dispatch._with_allocated_ports(sharded, "run-g") + manager = patched["manager"] + # Consecutive default slots after the declared gateway port. + assert [spec["port"] for spec in manager["gateway_shards"]] == [53910, 53911] + assert manager["gateway_port"] == 53910 + assert manager["gateway_url"] == "http://127.0.0.1:53910" + dispatch._release_run_ports("run-g") + + +def test_busy_shard_port_shifts_only_that_shard(): + holder = hold("127.0.0.1", 53911) + if holder is None: + pytest.skip("another process already holds the candidate shard port") + try: + sharded = deployment(gateway_shards=[{"name": "gw0"}, {"name": "gw1"}]) + patched = dispatch._with_allocated_ports(sharded, "run-h") + manager = patched["manager"] + assert [spec["port"] for spec in manager["gateway_shards"]] == [53910, 53912] + assert manager["gateway_port"] == 53910 + assert manager["gateway_url"] == "http://127.0.0.1:53910" + finally: + holder.close() + dispatch._release_run_ports("run-h") + + +def test_rollout_environment_carries_the_shard_table(): + workers = [{"id": "worker-a", "ssh": {"host": "worker-a", "user": "root"}, + "runner_python": "/runner/python"}, + {"id": "worker-b", "ssh": {"host": "worker-b"}, + "runner_python": "/runner/python"}] + sharded = deployment( + gateway_shards=[{"name": "gw0", "port": 53910}, {"name": "gw1", "port": 53911}], + gateway_routing="worker_affinity", + worker_shards={"worker-a": "gw0", "worker-b": "gw1"}) + sharded["gpu_target"] = {"train_python": "/train/python"} + sharded["rollout"] = {"workers": workers, "provider": "ssh_docker"} + overlay = dispatch._rollout_environment( + {"config": {"harbor": {}}}, sharded, {"BASE": "1"}) + assert json.loads(overlay["HARBORRL_NATIVE_GATEWAY_SHARDS"]) == [ + {"name": "gw0", "port": 53910}, {"name": "gw1", "port": 53911}] + assert overlay["HARBORRL_NATIVE_GATEWAY_ROUTING"] == "worker_affinity" + # Worker ids become the SSH destinations the runtime routes by. + assert json.loads(overlay["HARBORRL_NATIVE_GATEWAY_WORKER_SHARDS"]) == { + "root@worker-a": "gw0", "worker-b": "gw1"} + # A single-gateway deployment renders no shard table at all. + plain = deployment() + plain["gpu_target"] = {"train_python": "/train/python"} + plain["rollout"] = {"workers": workers, "provider": "ssh_docker"} + assert "HARBORRL_NATIVE_GATEWAY_SHARDS" not in dispatch._rollout_environment( + {"config": {"harbor": {}}}, plain, {"BASE": "1"}) diff --git a/tests/harborrl/test_runtime_cleanup_admission.py b/tests/harborrl/test_runtime_cleanup_admission.py new file mode 100644 index 000000000..1dd5b43f4 --- /dev/null +++ b/tests/harborrl/test_runtime_cleanup_admission.py @@ -0,0 +1,59 @@ +"""Cleanup markers block retries, queued slots and restarted runtimes.""" +import asyncio +from types import SimpleNamespace + +import pytest + +from harborrl.rollout import runtime + + +@pytest.mark.parametrize("marker_path", [ + "cleanup-required.json", "group/attempt/runner/cleanup-required.json", +]) +def test_runtime_restart_checks_cleanup_before_starting_gateway(tmp_path, monkeypatch, marker_path): + catalog, profile = tmp_path / "catalog.json", tmp_path / "profile.json" + catalog.write_text("[]") + profile.write_text('{}') + marker = tmp_path / "native-rollouts" / marker_path + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text('{}') + # Gateways start eagerly in __init__ here; the cleanup check runs before + # any gateway construction, so construction itself is the forbidden seam. + def forbidden(*_): + raise AssertionError("gateway must not start before cleanup is confirmed") + import harborrl.gateway.server as gateway_server + monkeypatch.setattr(gateway_server, "make_server", forbidden) + config = SimpleNamespace(run_root=tmp_path, catalog_path=catalog, profile_path=profile) + with pytest.raises(runtime.CleanupRequiredError, match="cleanup unresolved"): + runtime.RolloutRuntime(config, None) + marker.unlink() + runtime.CleanupAdmission(tmp_path / "native-rollouts").check() + + +@pytest.mark.parametrize("marker_path", ["cleanup-required.json", "runner/cleanup-required.json"]) +@pytest.mark.parametrize("raises", [False, True]) +def test_docker_cleanup_debt_blocks_retries_and_waiting_slots(tmp_path, monkeypatch, marker_path, raises): + calls = [] + async def attempt(identity, **kwargs): + calls.append(identity) + await asyncio.sleep(.01) + marker = kwargs["root"] / marker_path + marker.parent.mkdir(parents=True) + marker.write_text('{}') + if raises: + raise RuntimeError("Runner cleanup failed") + return {} + monkeypatch.setattr(runtime, "run_attempt", attempt) + config = SimpleNamespace(run_id="run", group_size=2, max_attempts=3, workers=(), + task_worker_overrides={}, runner_transport="local", sandbox=None, + gateway_routing="modulo", + gateway_url="http://gateway", runner_python="/python", max_output_tokens=128) + shard = SimpleNamespace(url="http://gateway", registry=None) + current = SimpleNamespace(profile={}, rollout_root=tmp_path, + semaphore=asyncio.Semaphore(1), registry=None, + shards=[shard]) + entry = {"id": "task", "path": "/task", "task_digest": "digest", "reward_profile": {}} + with pytest.raises(runtime.CleanupRequiredError, match="automatic retry blocked"): + asyncio.run(runtime.run_group(config, current, entry=entry, + group_index=0, rollout_id=0, sampling_digest="x", policy_version="1")) + assert len(calls) == 1 diff --git a/tests/harborrl/test_sandbox.py b/tests/harborrl/test_sandbox.py index b98d27197..a59d8e148 100644 --- a/tests/harborrl/test_sandbox.py +++ b/tests/harborrl/test_sandbox.py @@ -13,6 +13,7 @@ load_sandbox, task_binding, ) +from harborrl.tasks.content import tree_digest def settings(tmp_path): @@ -52,6 +53,47 @@ def test_sandbox_binding_and_secret_rejection(tmp_path): load_sandbox(path) +def test_task_digest_binding_preserves_original_files_and_rejects_changes(tmp_path): + task = tmp_path / 'task' + task.mkdir() + (task / 'task.toml').write_text('[environment]\n') + (task / 'instruction.md').write_text('Original task.') + original = tree_digest(task) + cfg = settings(tmp_path) + cfg['bindings'][0]['task_digests'] = [original] + path = tmp_path / 'sandbox.yaml' + path.write_text(yaml.safe_dump(cfg)) + config = load_sandbox(path) + assert task_binding(config, task)['image'] == cfg['bindings'][0]['image'] + assert tree_digest(task) == original + (task / 'instruction.md').write_text('Changed task.') + with pytest.raises(ValueError, match='exact binding'): + task_binding(config, task) + (task / 'task.toml').write_text('[environment]\ndocker_image = "' + cfg['bindings'][0]['image'] + '"\n') + with pytest.raises(ValueError, match='exact binding'): + task_binding(config, task) + + +@pytest.mark.parametrize('digests', [[], 'a' * 64, [1], ['invalid'], ['a' * 64, 'a' * 64]]) +def test_invalid_task_digest_bindings(tmp_path, digests): + cfg = settings(tmp_path) + cfg['bindings'][0]['task_digests'] = digests + path = tmp_path / 'sandbox.yaml' + path.write_text(yaml.safe_dump(cfg)) + with pytest.raises(ValueError, match='digest'): + load_sandbox(path) + + +def test_duplicate_task_digest_across_images(tmp_path): + cfg = settings(tmp_path) + cfg['bindings'][0]['task_digests'] = ['c' * 64] + cfg['bindings'].append({'image': 'other', 'max_live_instances': 1, 'task_digests': ['c' * 64]}) + path = tmp_path / 'sandbox.yaml' + path.write_text(yaml.safe_dump(cfg)) + with pytest.raises(ValueError, match='duplicate sandbox task digest'): + load_sandbox(path) + + def test_admission_across_groups_and_cleanup_debt(tmp_path): async def run(): cfg = {'max_live_instances': 2, 'bindings': [{'image': 'image', 'max_live_instances': 1}]} diff --git a/tests/harborrl/test_sandbox_allocation_transport.py b/tests/harborrl/test_sandbox_allocation_transport.py index db34cf098..5f71bbf5a 100644 --- a/tests/harborrl/test_sandbox_allocation_transport.py +++ b/tests/harborrl/test_sandbox_allocation_transport.py @@ -1,6 +1,8 @@ import asyncio +import json import httpx +import pytest from harborrl.rollout.harbor_job.sandbox_transport import SandboxStartupTransport @@ -22,3 +24,36 @@ def handler(request): assert calls == ['POST', 'GET'] and allocations == ['sbx-one'] asyncio.run(run()) + + +@pytest.mark.parametrize('kind', ['401', 'timeout', 'cancelled']) +def test_request_evidence_keeps_only_safe_lifecycle_metadata(kind): + async def run(): + records = [] + + async def handler(request): + if kind == 'timeout': + raise httpx.ReadTimeout('private signed URL and token') + if kind == 'cancelled': + raise asyncio.CancelledError('private SDK details') + return httpx.Response(401, text='private response body') + + transport = SandboxStartupTransport('https://sandbox.example/v1', lambda _: None, + transport=httpx.MockTransport(handler), on_request=records.append) + async with httpx.AsyncClient(transport=transport) as client: + url = 'https://sandbox.example/v1/sandboxes/sbx-one/endpoints/44772?token=private' + if kind == '401': + assert (await client.get(url, headers={'Authorization': 'private'})).status_code == 401 + else: + with pytest.raises(httpx.ReadTimeout if kind == 'timeout' else asyncio.CancelledError): + await client.get(url) + assert len(records) == 1 + record = records[0] + assert record.keys() == {'method', 'operation', 'http_status', 'error_type', 'elapsed_sec', 'time'} + assert record['method'] == 'GET' and record['operation'] == 'endpoint' + assert record['http_status'] == (401 if kind == '401' else None) + assert record['error_type'] == {'401': None, 'timeout': 'ReadTimeout', 'cancelled': 'CancelledError'}[kind] + assert record['elapsed_sec'] >= 0 + assert 'private' not in json.dumps(records) + + asyncio.run(run()) diff --git a/tests/harborrl/test_sandbox_cleanup.py b/tests/harborrl/test_sandbox_cleanup.py index 7364ac483..d1f8eaa86 100644 --- a/tests/harborrl/test_sandbox_cleanup.py +++ b/tests/harborrl/test_sandbox_cleanup.py @@ -34,3 +34,21 @@ async def handler(request): assert env.record['state'] == expected assert (tmp_path / 'cleanup-required.json').exists() == (expected == 'delete_failed') assert calls[:2] == [('GET', '/v1/sandboxes/sbx-owned'), ('DELETE', '/v1/sandboxes/sbx-owned')] + + +def test_startup_failure_preserves_cause_types_without_sdk_secrets(tmp_path): + import json + + env = object.__new__(ManagedOpenSandboxEnvironment) + env.cfg = {'create_timeout_s': 600, 'request_timeout_s': 180} + env.evidence = tmp_path + env.record = {'sandbox_id': 'sbx-owned'} + cause = httpx.ReadTimeout('private token and signed URL') + outer = RuntimeError('secret response body') + outer.__cause__ = cause + env.startup_failed('create', outer) + text = (tmp_path / 'sandbox-startup-failed.json').read_text() + assert 'private token' not in text and 'secret response' not in text + assert json.loads(text) == {'stage': 'create', 'sandbox_id': 'sbx-owned', + 'causes': [{'error_type': 'RuntimeError'}, {'error_type': 'ReadTimeout'}], + 'create_timeout_s': 600, 'request_timeout_s': 180} diff --git a/tests/harborrl/test_sandbox_runtime.py b/tests/harborrl/test_sandbox_runtime.py index f558af24b..027e43d54 100644 --- a/tests/harborrl/test_sandbox_runtime.py +++ b/tests/harborrl/test_sandbox_runtime.py @@ -37,97 +37,38 @@ def request(method, path, body=None): assert calls == ['POST', 'GET', 'GET'] -@pytest.mark.parametrize('provider', ['docker', 'sandbox']) -@pytest.mark.parametrize('marker_path', ['cleanup-required.json', 'runner/cleanup-required.json']) -def test_cleanup_debt_prevents_retry(tmp_path, monkeypatch, provider, marker_path): +def test_sandbox_cleanup_debt_prevents_retry(tmp_path, monkeypatch): from harborrl.rollout import runtime from harborrl.rollout.harbor_job import sandbox_config profile = tmp_path / 'profile.json' profile.write_text('{}') binding = {'image': 'image', 'max_live_instances': 1} - sandbox = {'max_live_instances': 1, 'bindings': [binding]} if provider == 'sandbox' else None + sandbox = {'max_live_instances': 1, 'bindings': [binding]} config = RolloutRuntimeConfig('run', str(tmp_path), '', str(profile), '127.0.0.1', 1, 'http://gateway', '/python', (), {}, 2, 128, 1024, 2, 2, 'tok', 'tpl', True, 1, 'qwen3-8B', 'http://router', sandbox=sandbox, runner_transport='local') - current = SimpleNamespace(rollout_root=tmp_path, semaphore=asyncio.Semaphore(2), - registry=None, profile={}) - if sandbox: - current.sandbox_admission = SandboxAdmission(sandbox) + current = SimpleNamespace( + rollout_root=tmp_path, semaphore=asyncio.Semaphore(2), registry=None, profile={}, + sandbox_admission=SandboxAdmission(sandbox), + shards=(SimpleNamespace(name="gateway", url="http://gateway", registry=None),)) calls = [] async def attempt(identity, **kwargs): calls.append(kwargs) assert kwargs['runner_host'] is None - assert kwargs.get('sandbox') == sandbox - marker = kwargs['root'] / marker_path - marker.parent.mkdir(parents=True) - marker.write_text('{}') + assert kwargs['sandbox'] == sandbox + kwargs['root'].mkdir(parents=True) + (kwargs['root'] / 'cleanup-required.json').write_text('{}') raise RuntimeError('deletion unconfirmed') monkeypatch.setattr(runtime, 'run_attempt', attempt) monkeypatch.setattr(sandbox_config, 'task_binding', lambda *args: binding) entry = {'id': 'task', 'path': '/task', 'task_digest': 'digest', 'reward_profile': {}} with pytest.raises(RuntimeError, match='cleanup unresolved'): asyncio.run(run_slot(config, current, 0, 0, 0, entry, 'sampling', '1')) - assert len(calls) == 1 - if sandbox: - assert current.sandbox_admission.blocked - with pytest.raises(runtime.CleanupRequiredError, match='cleanup-required.json'): - runtime.CleanupAdmission(tmp_path).check() - (calls[0]['root'] / marker_path).unlink() - current.cleanup.check() - runtime.CleanupAdmission(tmp_path).check() - if sandbox: - assert not current.sandbox_admission.blocked + assert len(calls) == 1 and current.sandbox_admission.blocked restored = replace(config, workers=('worker',), runner_transport='ssh') assert json.loads(restored.environment()['HARBORRL_NATIVE_SANDBOX']) == sandbox -def test_docker_cleanup_debt_blocks_waiting_slots(tmp_path, monkeypatch): - from harborrl.rollout import runtime - - calls = [] - - async def attempt(identity, **kwargs): - calls.append(identity) - await asyncio.sleep(.01) - kwargs['root'].mkdir(parents=True) - (kwargs['root'] / 'cleanup-required.json').write_text('{}') - raise RuntimeError('Runner cleanup failed') - - monkeypatch.setattr(runtime, 'run_attempt', attempt) - config = SimpleNamespace(run_id='run', group_size=2, max_attempts=3, workers=(), - task_worker_overrides={}, runner_transport='local', sandbox=None, - gateway_url='http://gateway', runner_python='/python', max_output_tokens=128) - current = SimpleNamespace(profile={}, rollout_root=tmp_path, - semaphore=asyncio.Semaphore(1), registry=None) - entry = {'id': 'task', 'path': '/task', 'task_digest': 'digest', 'reward_profile': {}} - with pytest.raises(runtime.CleanupRequiredError, match='automatic retry blocked') as error: - asyncio.run(runtime.run_group(config, current, entry=entry, - group_index=0, rollout_id=0, sampling_digest='x', policy_version='1')) - assert len(calls) == 1 - assert isinstance(error.value.__context__, RuntimeError) - - -def test_runtime_lifecycle_failure_is_not_retried(tmp_path, monkeypatch): - from harborrl.rollout import runtime - - calls = [] - - async def attempt(identity, **kwargs): - calls.append(identity) - raise runtime.RuntimeLifecycleError('wrong event loop') - - monkeypatch.setattr(runtime, 'run_attempt', attempt) - config = SimpleNamespace(run_id='run', group_size=2, max_attempts=3, workers=(), - task_worker_overrides={}, runner_transport='local', sandbox=None, - gateway_url='http://gateway', runner_python='/python', max_output_tokens=128) - current = SimpleNamespace(profile={}, rollout_root=tmp_path, - semaphore=asyncio.Semaphore(1), registry=None) - entry = {'id': 'task', 'path': '/task', 'task_digest': 'digest', 'reward_profile': {}} - with pytest.raises(runtime.RuntimeLifecycleError, match='wrong event loop'): - asyncio.run(run_slot(config, current, 0, 0, 0, entry, 'sampling', '1')) - assert len(calls) == 1 - - def test_group_failure_waits_for_sibling_cleanup(tmp_path, monkeypatch): from harborrl.rollout import runtime events = [] @@ -143,9 +84,11 @@ async def attempt(identity, **kwargs): monkeypatch.setattr(runtime, 'run_attempt', attempt) config = SimpleNamespace(run_id='run', group_size=2, max_attempts=1, workers=(), task_worker_overrides={}, runner_transport='local', sandbox=None, + gateway_routing='modulo', gateway_url='http://gateway', runner_python='/python', max_output_tokens=128) - current = SimpleNamespace(profile={}, rollout_root=tmp_path, - semaphore=asyncio.Semaphore(2), registry=None) + current = SimpleNamespace( + profile={}, rollout_root=tmp_path, semaphore=asyncio.Semaphore(2), registry=None, + shards=(SimpleNamespace(name="gateway", url="http://gateway", registry=None),)) entry = {'id': 'task', 'path': '/task', 'task_digest': 'digest', 'reward_profile': {}} with pytest.raises(RuntimeError, match='exhausted its retry budget'): asyncio.run(runtime.run_group(config, current, entry=entry, diff --git a/tests/harborrl/test_sandbox_startup.py b/tests/harborrl/test_sandbox_startup.py new file mode 100644 index 000000000..12ae5de93 --- /dev/null +++ b/tests/harborrl/test_sandbox_startup.py @@ -0,0 +1,195 @@ +"""Exercise Runner queue ownership with the pinned Harbor timer implementation.""" +import asyncio +import json +from datetime import datetime, timezone +from types import SimpleNamespace + +import pytest + +pytest.importorskip('harbor.environments.opensandbox') + +from harbor.models.trial.config import TrialConfig +from harbor.trial.trial import Trial +from test_native_contracts import identity + +from harborrl.rollout.harbor_job import runner +from harborrl.rollout.harbor_job.environments.opensandbox import ( + ManagedOpenSandboxEnvironment, + startup_slot, +) +from harborrl.trajectories.native import digest, read_json + + +@pytest.mark.parametrize('mode', ['ready', 'start_failed', 'cancel_waiting']) +def test_runner_queue_does_not_consume_harbor_environment_timer(tmp_path, monkeypatch, mode): + from harborrl.rollout.harbor_job import sandbox_config + + task = tmp_path / 'task' + task.mkdir() + profile = {'cli_version': '2.1.141', 'model': 'policy', 'max_turns': 2, + 'setup_timeout_sec': 1, 'agent_timeout_sec': 1, 'verifier_timeout_sec': 1} + sandbox = {'adapter': 'test:Environment', 'installation': None, 'bindings': []} + output = tmp_path / 'output' + monkeypatch.setattr(runner, 'probe', dict) + monkeypatch.setattr(runner.NATIVE_TASK_PROVIDER.adapter, 'inspect', lambda *a, **k: { + 'status': 'NEEDS_PROBE', 'path': str(task), 'resources': {}}) + monkeypatch.setattr(sandbox_config, 'environment_class', lambda _: ManagedOpenSandboxEnvironment) + monkeypatch.setattr(sandbox_config, 'task_binding', lambda *a: {}) + env = object.__new__(ManagedOpenSandboxEnvironment) + env.cfg = {'startup_concurrency': 1} + env.evidence = output + env.startup_dir = tmp_path / 'locks' + env._startup_reservation = None + env.record = {'sandbox_id': None, 'state': 'pending', 'events': []} + starts = [] + + async def start(force_build): + assert env._startup_reservation is not None + starts.append(force_build) + try: + await asyncio.sleep(.001) + if mode == 'start_failed': + raise RuntimeError('synthetic startup failure') + finally: + await env.release_startup() + + env.start = start + + class SyntheticTrial: + id = 'synthetic-trial' + _environment_build_timeout_sec = .05 + _now = staticmethod(lambda: datetime.now(timezone.utc)) + _start_agent_environment = Trial._start_agent_environment + + def __init__(self, config): + self.config = config + self.agent_environment = env + self.result = SimpleNamespace(exception_info=None) + self.paths = SimpleNamespace(result_path=output / 'harbor/result.json', trial_dir=output / 'harbor') + self.hooks = {} + + def add_hook(self, event, hook): + self.hooks[event] = hook + + async def _emit(self, event): + await self.hooks[event](SimpleNamespace( + event=event, trial_id=self.id, timestamp=self._now())) + + async def run(self): + try: + await Trial._setup_agent_environment(self) + except RuntimeError as exc: + self.result.exception_info = SimpleNamespace(exception_type=type(exc).__name__) + return self.result + + async def create(config): + assert isinstance(config, TrialConfig) + return SyntheticTrial(config) + + monkeypatch.setattr(Trial, 'create', create) + request = {'identity': identity(harness_digest=digest(profile)).to_dict(), + 'task_path': str(task), 'profile': profile, 'output': str(output), + 'sandbox': sandbox, 'runtime_digest': digest(sandbox), + 'gateway_url': 'http://gateway', 'credential': 'synthetic', 'max_output_tokens': 16} + + async def exercise(): + events = [] + run_sent = False + cancel = asyncio.Event() + + async def read(): + nonlocal run_sent + if not run_sent: + run_sent = True + return {'command': 'run'} + await cancel.wait() + return {'command': 'cancel'} + + async with startup_slot(env.startup_dir, 1): + running = asyncio.create_task(runner.execute(request, events.append, read)) + async with asyncio.timeout(2): + while env.record['state'] != 'waiting': + await asyncio.sleep(.001) + await asyncio.sleep(.08) + assert not running.done() and not starts + if mode == 'cancel_waiting': + cancel.set() + await asyncio.wait_for(running, 1) + assert events[-1]['type'] == 'cancelled' + if mode != 'cancel_waiting': + await asyncio.wait_for(running, 2) + assert len(starts) == 1 + assert events[-1]['exception'] == ('RuntimeError' if mode == 'start_failed' else None) + assert env._startup_reservation is None + async with asyncio.timeout(1), startup_slot(env.startup_dir, 1): + pass + lifecycle = read_json(output / 'sandbox-lifecycle.json') + assert lifecycle['sandbox_id'] is None + assert not (output / 'cleanup-required.json').exists() + assert [row['state'] for row in lifecycle['events']] == ( + ['waiting'] if mode == 'cancel_waiting' else ['waiting', 'reserved']) + + asyncio.run(exercise()) + + +@pytest.mark.parametrize('mode', ['ready', 'create_failed', 'readiness_failed', 'binding_failed']) +def test_sdk_startup_reuses_reservation_and_releases_on_failure(tmp_path, mode): + env = object.__new__(ManagedOpenSandboxEnvironment) + env.cfg = {'startup_concurrency': 1, 'create_timeout_s': 1, 'request_timeout_s': 1, + 'lifecycle_minutes': 1} + env.evidence, env.startup_dir = tmp_path, tmp_path / 'locks' + env._startup_reservation = None + env._sandbox = None + env._startup_transport = None + env._persistent_env = {} + env.binding = {'image': 'synthetic-image'} + env.record = {'sandbox_id': None, 'state': 'pending', 'events': [], 'runtime_digest': 'test'} + env._build_connection_config = lambda _: SimpleNamespace(model_copy=lambda **k: None) + env._build_resource = dict + env._build_network_policy = lambda _: None + + class Transport: + async def aclose(self): + pass + + env.startup_transport = lambda _: Transport() + + async def validate(): + if mode == 'binding_failed': + raise RuntimeError('synthetic binding failure') + + async def ready(_): + if mode == 'readiness_failed': + raise RuntimeError('synthetic readiness failure') + + async def stop(delete): + env.record_state('deleted') + + class Sandbox: + @staticmethod + async def create(*a, **k): + assert env._startup_reservation is not None + if mode == 'create_failed': + raise RuntimeError('synthetic create failure') + return SimpleNamespace(id='synthetic-sandbox') + + env.validate_binding, env._wait_until_ready, env.stop = validate, ready, stop + + async def exercise(): + await env.reserve_startup() + with pytest.raises(RuntimeError, match='already reserved'): + await env.reserve_startup() + if mode == 'ready': + assert (await env._create_sandbox({'Sandbox': Sandbox})).id == 'synthetic-sandbox' + else: + with pytest.raises(RuntimeError): + await env._create_sandbox({'Sandbox': Sandbox}) + assert env._startup_reservation is None + async with asyncio.timeout(1), startup_slot(env.startup_dir, 1): + pass + states = [row['state'] for row in env.record['events']] + assert states.count('reserved') == 1 + if mode == 'create_failed': + assert json.loads((tmp_path / 'cleanup-required.json').read_text())['reason'].startswith('Creation outcome unknown') + + asyncio.run(exercise()) diff --git a/tests/harborrl/test_service_recovery.py b/tests/harborrl/test_service_recovery.py new file mode 100644 index 000000000..293c70917 --- /dev/null +++ b/tests/harborrl/test_service_recovery.py @@ -0,0 +1,224 @@ +"""Restart restores durable group state, never launches an old Runner twice.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest +from test_rollout_service import ( + FakeRuntime, + activated_states, + group_request, + write_group_irs, +) + +from harborrl.rollout.runtime import CleanupAdmission, CleanupRequiredError +from harborrl.rollout.service import ( + GROUP_DIR, + RolloutService, + ServiceError, + request_digest, + validate_group_request, +) +from harborrl.trajectories.native import publish, read_json + +MANIFEST = {"run_id": "run-1", "group_size": 2} + + +def ready(service, root): + service.register_run("run-1", MANIFEST) + service.attach_runtime("run-1", SimpleNamespace(drain_timeout=1), FakeRuntime(root)) + service.activate_policy("run-1", activated_states(), "1") + + +def test_completed_group_survives_restart_without_regeneration(tmp_path): + calls = [] + async def executor(_config, _runtime, request): + calls.append(request["group_id"]) + return write_group_irs(tmp_path, request) + first = RolloutService(run_root=tmp_path, executor=executor, allow_synthetic=True) + try: + ready(first, tmp_path) + original = first.submit_group("run-1", group_request()) + finally: + first.close() + directory = tmp_path / GROUP_DIR / "run-1/group-0" + evidence = {p.name: p.read_bytes() for p in directory.glob("*.json")} + async def forbidden(*_): + raise AssertionError("recovered group must not run again") + second = RolloutService(run_root=tmp_path, executor=forbidden, allow_synthetic=True) + try: + reply = second.register_run("run-1", MANIFEST) + assert reply["existed"] and reply["recovered_groups"] == 1 + restored = second.submit_group("run-1", group_request()) + assert restored == original + assert calls == ["group-0"] + assert {p.name: p.read_bytes() for p in directory.glob("*.json")} == evidence + changed = group_request() + changed["sampling"]["max_new_tokens"] += 1 + with pytest.raises(ServiceError) as error: + second.submit_group("run-1", changed) + assert error.value.code == "conflict" + finally: + second.close() + + +@pytest.mark.parametrize("status,filename", [("failed", "failed.json"), ("cancelled", "cancelled.json")]) +def test_failed_and_cancelled_groups_do_not_silently_retry(tmp_path, status, filename): + request = group_request() + directory = tmp_path / GROUP_DIR / "run-1/group-0" + publish(directory / "request.json", request) + publish(directory / filename, {"group_id": "group-0", "request_digest": request_digest(request), + "status": status, "created_at": 1, "error": "original reason"}) + service = RolloutService(run_root=tmp_path) + try: + service.register_run("run-1", MANIFEST) + recovered = service.submit_group("run-1", request) + assert recovered["status"] == status and recovered["error"] == "original reason" + assert service.runs["run-1"].groups["group-0"].future is None + finally: + service.close() + + +@pytest.mark.parametrize("running_record", [False, True]) +def test_interrupted_group_restores_unknown_and_blocks_new_generation(tmp_path, running_record): + request = group_request() + directory = tmp_path / GROUP_DIR / "run-1/group-0" + publish(directory / "request.json", request) + if running_record: + publish(directory / "running.json", {"group_id": "group-0", "request_digest": request_digest(request), + "status": "running", "created_at": 1}) + for _ in range(2): + service = RolloutService(run_root=tmp_path) + try: + service.register_run("run-1", MANIFEST) + state = service.submit_group("run-1", request) + assert state["status"] == "unknown" and "cleanup" in state["error"] + assert service.runs["run-1"].groups["group-0"].future is None + with pytest.raises(ServiceError) as error: + service.submit_group("run-1", group_request(group_index=1), wait=False) + assert error.value.code == "admission_closed" + with pytest.raises(ServiceError) as error: + service.cancel_run("run-1") + assert error.value.code == "cleanup_required" + with pytest.raises(CleanupRequiredError): + CleanupAdmission(tmp_path / "native-rollouts").check() + finally: + service.close() + assert read_json(directory / "unknown.json")["status"] == "unknown" + assert read_json(directory / "request.json") == request + + +def test_run_directory_has_one_live_service_owner(tmp_path): + first = RolloutService(run_root=tmp_path) + try: + with pytest.raises(ServiceError) as error: + RolloutService(run_root=tmp_path) + assert error.value.code == "run_directory_busy" + finally: + first.close() + second = RolloutService(run_root=tmp_path) + second.close() + + +def test_restart_rejects_changed_registration_and_corrupt_group(tmp_path): + first = RolloutService(run_root=tmp_path) + first.register_run("run-1", MANIFEST) + first.close() + second = RolloutService(run_root=tmp_path) + try: + with pytest.raises(ServiceError) as error: + second.register_run("run-1", {**MANIFEST, "group_size": 3}) + assert error.value.code == "conflict" + request = group_request() + directory = tmp_path / GROUP_DIR / "run-1/group-0" + publish(directory / "request.json", request) + publish(directory / "failed.json", {"status": "failed", "group_id": "group-0", "request_digest": "wrong"}) + with pytest.raises(ServiceError) as error: + second.register_run("run-1", MANIFEST) + assert error.value.code == "invalid_persisted_group" + assert not second.runs + finally: + second.close() + + +def test_complete_is_reported_only_after_result_is_durable(tmp_path, monkeypatch): + from harborrl.rollout import service as module + publishing, release = threading.Event(), threading.Event() + original = module.publish + def delayed_result(path, value): + if path.name == "result.json": + publishing.set() + assert release.wait(5) + return original(path, value) + async def executor(_config, _runtime, request): + return write_group_irs(tmp_path, request) + service = RolloutService(run_root=tmp_path, executor=executor, allow_synthetic=True) + try: + ready(service, tmp_path) + monkeypatch.setattr(module, "publish", delayed_result) + service.submit_group("run-1", group_request(), wait=False) + assert publishing.wait(5) + state = service.group_status("run-1", "group-0") + assert state["status"] == "running" and "result" not in state + release.set() + service.runs["run-1"].groups["group-0"].future.result(timeout=5) + assert service.group_status("run-1", "group-0")["status"] == "complete" + finally: + release.set() + service.close() + + +def test_executor_setup_failure_becomes_terminal_not_stuck_running(tmp_path, monkeypatch): + from harborrl.rollout import service as module + original = module.publish + def failed_running(path, value): + if path.name == "running.json": + raise OSError("cannot persist running state") + return original(path, value) + async def forbidden(*_): + raise AssertionError("setup failed before allocation") + service = RolloutService(run_root=tmp_path, executor=forbidden) + try: + ready(service, tmp_path) + monkeypatch.setattr(module, "publish", failed_running) + result = service.submit_group("run-1", group_request()) + assert result["status"] == "failed" and "cannot persist" in result["error"] + finally: + service.close() + + +@pytest.mark.parametrize("bad_id", ["..", "../escape", "a/b", "a\\b"]) +def test_storage_identities_are_one_directory_component(tmp_path, bad_id): + request = group_request() + request["group_id"] = bad_id + with pytest.raises(ServiceError, match="directory component"): + validate_group_request(request) + service = RolloutService(run_root=tmp_path) + try: + with pytest.raises(ServiceError, match="directory component"): + service.register_run(bad_id, {"run_id": bad_id}) + finally: + service.close() + + +def test_running_group_is_cancelled_before_lock_is_released(tmp_path): + # The lock must also prevent a second owner while asynchronous cleanup runs. + started = asyncio.Event() + async def executor(*_): + started.set() + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(.01) + service = RolloutService(run_root=tmp_path, executor=executor) + try: + ready(service, tmp_path) + service.submit_group("run-1", group_request(), wait=False) + finally: + service.close() + recovered = RolloutService(run_root=tmp_path) + try: + recovered.register_run("run-1", MANIFEST) + assert recovered.group_status("run-1", "group-0")["status"] == "cancelled" + finally: + recovered.close() diff --git a/tests/harborrl/test_service_sampling_lock.py b/tests/harborrl/test_service_sampling_lock.py new file mode 100644 index 000000000..9f499c03f --- /dev/null +++ b/tests/harborrl/test_service_sampling_lock.py @@ -0,0 +1,79 @@ +"""Bind self-consistent group evidence to the sampling actually requested.""" +from copy import deepcopy + +import pytest +from test_rollout_service import FakeRuntime, group_request, service_ir + +from harborrl.backends.slime_v032.rollout_service_client import ( + ServiceRolloutError, + load_group_result, +) +from harborrl.rollout.service import build_group_result +from harborrl.trajectories.native import ( + Identity, + assemble, + canonical_sampling, + digest, + publish, + require_ready, +) + + +def result_for(root, request, *, actual_sampling=None): + sampling = canonical_sampling(actual_sampling or request["sampling"]) + runtime, trajectories = FakeRuntime(root), [] + for slot in request["slots"]: + source = deepcopy(service_ir(root / "receipts", slot["slot_id"])) + identity = Identity(**{**source["identity"], "sampling_digest": digest(sampling)}) + evaluation, seal, turns = source["evaluation"], source["trace_seal"], source["turns"] + evaluation["identity"] = seal["identity"] = identity.to_dict() + for turn in turns: + turn["identity"] = identity.to_dict() + turn["sampling"] = sampling + seal["turns_digest"] = digest(turns) + ir = assemble(identity, source["harbor_trial_id"], turns, seal, evaluation) + require_ready(ir, allow_synthetic=True) # internally valid, even for a different distribution + publish(runtime.manifest_path(ir["identity"]), ir) + trajectories.append(ir) + return build_group_result("run-1", request, trajectories, runtime=runtime) + + +@pytest.mark.parametrize("key,value", [ + ("temperature", .5), ("top_p", .8), ("top_k", 10), + ("max_new_tokens", 99), ("stop", ["END"]), +]) +def test_rejects_internally_consistent_but_unrequested_sampling(tmp_path, key, value): + request = group_request() + result = result_for(tmp_path, request, actual_sampling={**request["sampling"], key: value}) + with pytest.raises(ServiceRolloutError, match="sampling"): + load_group_result("run-1", request, result, allow_synthetic=True) + + +@pytest.mark.parametrize("probability", [1, 1.0]) +def test_accepts_semantically_equivalent_probability_numbers(tmp_path, probability): + request = group_request() + request["sampling"].update(temperature=probability, top_p=probability) + result = result_for(tmp_path, request) + assert len(load_group_result("run-1", request, result, allow_synthetic=True)) == 2 + assert type(request["sampling"]["temperature"]) is type(probability) + + +@pytest.mark.parametrize("missing", [False, True]) +def test_descriptor_sampling_cannot_be_changed_or_omitted(tmp_path, missing): + request = group_request() + result = result_for(tmp_path, request) + if missing: + del result["slots"][0]["sampling_digest"] + else: + result["slots"][0]["sampling_digest"] = "other-distribution" + with pytest.raises(ServiceRolloutError, match="sampling digest"): + load_group_result("run-1", request, result, allow_synthetic=True) + + +def test_requested_descriptor_does_not_mask_different_manifest_sampling(tmp_path): + request = group_request() + result = result_for(tmp_path, request, actual_sampling={**request["sampling"], "temperature": .5}) + for slot in result["slots"]: + slot["sampling_digest"] = digest(canonical_sampling(request["sampling"])) + with pytest.raises(ServiceRolloutError, match="manifest sampling"): + load_group_result("run-1", request, result, allow_synthetic=True) diff --git a/tests/harborrl/test_service_shutdown.py b/tests/harborrl/test_service_shutdown.py new file mode 100644 index 000000000..fa8f8fc5e --- /dev/null +++ b/tests/harborrl/test_service_shutdown.py @@ -0,0 +1,167 @@ +"""Service shutdown joins real async cleanup; no model, Docker or GPU allocation.""" +import asyncio +import threading +from types import SimpleNamespace + +import pytest +from test_rollout_service import FakeRegistry, group_request + +from harborrl.rollout.service import RolloutService, ServiceError + + +def service_for(root, executor, runtime_close=lambda: None, **kwargs): + service = RolloutService(run_root=root, executor=executor, poll_interval=.005, **kwargs) + service.register_run("run", {"run_id": "run"}) + runtime = SimpleNamespace(registry=FakeRegistry(), close=runtime_close) + service.attach_runtime("run", SimpleNamespace(drain_timeout=1), runtime) + service.activate_policy("run", [{"endpoint": "cpu", "healthy": True, "weight_version": "1"}], "1") + return service + + +def test_close_joins_runner_finally_before_runtime_and_loop(tmp_path): + started, cleaned = threading.Event(), threading.Event() + order = [] + children = [] + async def executor(*_): + import sys + child = await asyncio.create_subprocess_exec(sys.executable, "-c", "import time;time.sleep(60)") + children.append(child) + started.set() + try: + await asyncio.Event().wait() + finally: + await asyncio.sleep(.02) + child.terminate() + await child.wait() + order.append("runner cleaned") + cleaned.set() + service = service_for(tmp_path, executor, lambda: order.append("runtime closed"), cleanup_timeout=2) + try: + service.submit_group("run", group_request(), wait=False) + assert started.wait(5) + assert service.drain_run("run")["state"] == "drained" + assert not cleaned.is_set() # generation drain is not a group join + service.close() + assert order == ["runner cleaned", "runtime closed"] + assert cleaned.is_set() and children[0].returncode is not None + assert service.group_status("run", "group-0")["status"] == "cancelled" + assert service.loop.is_closed() and not service.thread.is_alive() + service.close() # idempotent; no scheduling on an already closed loop + finally: + service.close() + + +def test_cleanup_timeout_is_visible_and_retry_does_not_recancel_finally(tmp_path): + started, cleaning, release, cleaned = (threading.Event() for _ in range(4)) + closed = [] + async def executor(*_): + started.set() + try: + await asyncio.Event().wait() + finally: + cleaning.set() + while not release.is_set(): + await asyncio.sleep(.005) + cleaned.set() + service = service_for(tmp_path, executor, lambda: closed.append(True), cleanup_timeout=.03) + try: + service.submit_group("run", group_request(), wait=False) + assert started.wait(5) + with pytest.raises(ServiceError) as error: + service.cancel_run("run") + assert error.value.code == "cleanup_timeout" + assert cleaning.is_set() and not cleaned.is_set() + group = service.runs["run"].groups["group-0"] + assert not group.future.done() and not group.task.done() + with pytest.raises(ServiceError, match="still cleaning"): + service.close() + assert service.loop.is_running() and not closed + assert (tmp_path / "native-rollouts/cleanup-required.json").is_file() + with pytest.raises(ServiceError, match="closing"): + service.register_run("other", {"run_id": "other"}) + release.set() + service.cleanup_timeout = 2 + service.close() + assert cleaned.is_set() and closed == [True] + finally: + release.set() + service.cleanup_timeout = 2 + service.close() + + +def test_close_blocks_new_admission_while_cleanup_is_running(tmp_path): + started, cleaning, release = (threading.Event() for _ in range(3)) + failures = [] + async def executor(*_): + started.set() + try: + await asyncio.Event().wait() + finally: + cleaning.set() + while not release.is_set(): + await asyncio.sleep(.005) + service = service_for(tmp_path, executor, cleanup_timeout=2) + def close(): + try: + service.close() + except Exception as exc: # noqa: BLE001 - report worker failures to the test thread + failures.append(exc) + try: + service.submit_group("run", group_request(), wait=False) + assert started.wait(5) + closer = threading.Thread(target=close) + closer.start() + assert cleaning.wait(5) + with pytest.raises(ServiceError, match="closing"): + service.submit_group("run", group_request(group_index=1), wait=False) + with pytest.raises(ServiceError, match="closing"): + service.announce_serving("run", "http://cpu") + release.set() + closer.join(5) + assert not closer.is_alive() and not failures + finally: + release.set() + service.close() + + +def test_default_cleanup_budget_covers_runner_and_sandbox_deadlines(tmp_path): + service = RolloutService(run_root=tmp_path) + try: + assert service._cleanup_budget() > 120 + service.register_run("run", {"run_id": "run"}) + service.attach_runtime("run", SimpleNamespace(sandbox={"cleanup_timeout_s": 300}), None) + assert service._cleanup_budget() > 315 + finally: + service.close() + + +def test_close_does_not_interrupt_slower_sibling_cleanup(tmp_path, monkeypatch): + from harborrl.rollout import runtime as runtime_module + + started = threading.Event() + events, active = [], set() + async def slot(_config, _runtime, _slots, index, _entry, _reward): + active.add(index) + if len(active) == 2: + started.set() + try: + await asyncio.Event().wait() + finally: + if index == 1: + await asyncio.sleep(.05) + events.append(f"slot {index} cleaned") + monkeypatch.setattr(runtime_module, "_run_slot", slot) + async def executor(config, current, request): + return await runtime_module.run_group( + config, current, entry=request["task"], group_index=0, rollout_id=0, + sampling_digest="sampling", policy_version="1", group_size=2) + service = service_for(tmp_path, executor, lambda: events.append("runtime closed"), cleanup_timeout=2) + service.runs["run"].config = SimpleNamespace(run_id="run", max_attempts=1, drain_timeout=1) + service.runs["run"].runtime.profile = {} + try: + service.submit_group("run", group_request(), wait=False) + assert started.wait(5) + service.close() + assert events == ["slot 0 cleaned", "slot 1 cleaned", "runtime closed"] + finally: + service.close() diff --git a/tests/harborrl/test_sglang_text_evidence.py b/tests/harborrl/test_sglang_text_evidence.py new file mode 100644 index 000000000..11f006db2 --- /dev/null +++ b/tests/harborrl/test_sglang_text_evidence.py @@ -0,0 +1,54 @@ +"""No token prefix/suffix can masquerade as the text delivered to the Agent.""" +from types import SimpleNamespace + +import pytest + +from harborrl.gateway.sglang_backend import SGLangBackend + + +def decode(decoded, text, *, finish=None, stops=()): + tokenizer = SimpleNamespace(decode=lambda *_args, **_kwargs: decoded, eos_token_id=99) + model = SimpleNamespace(parse_output=lambda value, **_kwargs: [{"type": "text", "text": value}]) + backend = SGLangBackend("http://unused", tokenizer, model, + tokenizer_digest="fixture", template_digest="fixture") + return backend.decode( + {"text": text, "meta_info": {"weight_version": "1", "output_token_logprobs": [[-.2, 13]], + "finish_reason": finish or {"type": "stop", "matched": 99}}}, + {"input_ids": [1]}, {"tools": [], "sampling": {"stop": list(stops)}}, + SimpleNamespace(policy_version="1"), "response") + + +@pytest.mark.parametrize("decoded,text", [ + ("done hidden suffix", "done"), + ("done", "done extra visible text"), + ("done\n", "done"), + ("done", "done\n"), + ("hidden", ""), + ("", "unbacked text"), + (" ", ""), +]) +def test_rejects_prefixes_and_unproved_whitespace_changes(decoded, text): + with pytest.raises(ValueError, match="text disagrees"): + decode(decoded, text) + + +@pytest.mark.parametrize("text", ["done", "done\n", "\n done ", "你好", ""]) +def test_accepts_exact_visible_text(text): + result = decode(text, text) + assert result["content"] == [{"type": "text", "text": text}] + assert result["output_ids"] == [13] + + +def test_accepts_only_the_explicit_matched_terminal_stop_omission(): + finish = {"type": "stop", "matched": ""} + for decoded in ("done", "done"): + result = decode(decoded, "done", finish=finish, stops=[""]) + assert result["stop_sequence"] == "" + assert result["finish_reason"] == "stop_sequence" + for decoded in ("donehidden", "done", "done\n"): + with pytest.raises(ValueError, match="text disagrees"): + decode(decoded, "done", finish=finish, stops=[""]) + with pytest.raises(ValueError, match="text disagrees"): + decode("done", "done", finish=finish) # undeclared stop + with pytest.raises(ValueError, match="text disagrees"): + decode("done", "done", finish={"type": "length", "matched": ""}, stops=[""]) diff --git a/tests/harborrl/test_slime_v032_exporter.py b/tests/harborrl/test_slime_v032_exporter.py index 5fb0e406f..28e733bb3 100644 --- a/tests/harborrl/test_slime_v032_exporter.py +++ b/tests/harborrl/test_slime_v032_exporter.py @@ -24,7 +24,7 @@ from harborrl.backends.slime_v032.postprocess import rollout_data_postprocess from harborrl.backends.slime_v032.rollout import generate_rollout, to_slime_samples from harborrl.export import contract as export_contract -from harborrl.export.native import export_training_batch +from harborrl.export.native import export_training_batch, group_statistics from harborrl.export.validate import batch_weight_sum from harborrl.trajectories.native import digest @@ -104,7 +104,7 @@ def test_export_training_batch_goldel(tmp_path): assert batch.groups[1].reward_std == 0.0 assert [unit.advantage for unit in batch.groups[1].trajectories] == [0.0, 0.0] first_advantages = [unit.advantage for unit in batch.groups[0].trajectories] - assert first_advantages[0] < 0 < first_advantages[1] + assert first_advantages == [-0.5, 0.5] assert [unit.token_count for group in batch.groups for unit in group.trajectories] == [2, 6, 4, 2] for group in batch.groups: for unit in group.trajectories: @@ -119,6 +119,23 @@ def test_export_training_batch_goldel(tmp_path): assert batch_weight_sum(batch) == pytest.approx(2.0) +def test_centered_advantages_keep_reward_scale_and_reject_normalized_batch(tmp_path): + from dataclasses import replace + + from harborrl.export.validate import validate_batch + + assert group_statistics([0.0, 1.0])[2] == [-0.5, 0.5] + assert group_statistics([-2.0, 2.0])[2] == [-2.0, 2.0] + batch = export_training_batch(native_groups(tmp_path), rollout_ids=[10, 11]) + group = batch.groups[0] + normalized = replace(group, trajectories=tuple( + replace(unit, advantage=unit.advantage / group.reward_std) + for unit in group.trajectories + )) + with pytest.raises(ValueError, match="centered"): + validate_batch(replace(batch, groups=(normalized, *batch.groups[1:]))) + + def test_export_rejects_malformed_groups(tmp_path): groups = native_groups(tmp_path) with pytest.raises(ValueError, match="at least two"): @@ -355,13 +372,7 @@ def test_launcher_hooks_resolve_and_dry_run(): v032_launcher.hook_arguments("slime-legacy") -def test_rollout_hook_locks_policy_version_and_keeps_groups_atomic(tmp_path, monkeypatch, request): - import asyncio - - from harborrl.rollout.runtime import close_runtime - - close_runtime() - request.addfinalizer(close_runtime) +def test_rollout_hook_locks_policy_version_and_keeps_groups_atomic(tmp_path, monkeypatch): groups = native_groups(tmp_path) locked = [ [reidentify(json.loads(json.dumps(ir)), policy_version="7") for ir in group] @@ -374,18 +385,10 @@ def test_rollout_hook_locks_policy_version_and_keeps_groups_atomic(tmp_path, mon args = SimpleNamespace( rollout_batch_size=2, n_samples_per_prompt=2, rollout_max_response_len=128, ) - events, loops = [], [] - semaphore = asyncio.Semaphore(1) + events = [] async def runner(args_, groups_, policy_version): events.append(("generate", policy_version)) - loops.append(asyncio.get_running_loop()) - - async def queued_slot(): - async with semaphore: - await asyncio.sleep(.001) - - await asyncio.gather(queued_slot(), queued_slot()) return json.loads(json.dumps(locked)) class SlimeDataSource: @@ -418,13 +421,6 @@ def get_samples(self, count): ) assert all(sample.metadata["native_ir"]["identity"]["policy_version"] == "7" for group in output.samples for sample in group) - second = generate_rollout( - args, 4, SlimeDataSource(), runner=runner, output_factory=SimpleNamespace, - sample_factory=lambda **fields: SimpleNamespace(**fields)) - assert loops[0] is loops[1] - assert second.metrics["native_turn_count"] == 7.0 - assert {sample.rollout_id for group in second.samples for sample in group} == {8, 9} - async def stale_runner(args_, groups_, policy_version): return [[reidentify(json.loads(json.dumps(ir)), policy_version="8") for ir in group] for group in locked] @@ -443,6 +439,21 @@ async def stale_runner(args_, groups_, policy_version): ) +@pytest.mark.parametrize("dtype_name", ["int64", "bool", "float32"]) +def test_advantage_preserves_fractional_rewards_with_tensor_masks(dtype_name): + torch = pytest.importorskip("torch") + data = { + "rewards": [-0.25, 0.75], + "loss_masks": [torch.ones(2, dtype=getattr(torch, dtype_name)) for _ in range(2)], + "rollout_log_probs": [torch.zeros(2) for _ in range(2)], + } + compute_advantages_and_returns(SimpleNamespace(), data) + for key in ("advantages", "returns"): + assert data[key][0].tolist() == [-0.25, -0.25] + assert data[key][1].tolist() == [0.75, 0.75] + assert all(value.dtype == torch.float32 for value in data[key]) + + def test_to_slime_samples_requires_ir_metadata(tmp_path): groups = native_groups(tmp_path) batch = export_training_batch(groups, rollout_ids=[1, 2]) @@ -466,7 +477,7 @@ def test_export_package_has_no_backend_imports(): @pytest.mark.parametrize('group_index, expected_advantages, expected_counts', [ - (0, [-0.5 / 0.500001, 0.5 / 0.500001], [2, 6]), + (0, [-0.5, 0.5], [2, 6]), (1, [0.0, 0.0], [4, 2]), ]) def test_legacy_export_preserves_grpo_record_semantics( diff --git a/tests/test_consistency_torch.py b/tests/test_consistency_torch.py index 0f81e6660..e2b521fe1 100644 --- a/tests/test_consistency_torch.py +++ b/tests/test_consistency_torch.py @@ -1,12 +1,40 @@ -"""Distributed rejection tests; optional Torch dependency.""" +"""Real optimizer and distributed rejection tests; optional Torch dependency.""" from pathlib import Path import pytest -from harborrl.consistency.config import validate -from harborrl.consistency.report import assess, distributed_gate - torch = pytest.importorskip('torch') +pytest.importorskip('transformers') + +from harborrl.consistency.config import validate # noqa: E402 +from harborrl.consistency.report import assess, distributed_gate # noqa: E402 +from harborrl.distill.config import DEFAULTS # noqa: E402 +from harborrl.distill.models import load_student, parameter_digest, response_logprobs # noqa: E402 +from harborrl.distill.runner import optimize # noqa: E402 + + +def test_gate_rejects_before_optimizer_or_gradient(tmp_path): + from copy import deepcopy + config = deepcopy(DEFAULTS) + config['consistency'] = validate({'mode': 'enforce', 'max_abs_delta': .001}) + config['output']['directory'] = str(tmp_path) + config['method'] = 'reinforce' + config['training']['outcome_weight'] = 1. + model, tokenizer = load_student(config, torch.device('cpu')) + optimizer = torch.optim.AdamW(model.parameters(), lr=.001) + before = parameter_digest(model) + row = {'sample_id': '000000-000000', 'prompt_ids': [1, 4], 'response_ids': [5, 6], + 'behavior_logprobs': [-.01, -.01], 'loss_mask': [1, 1], 'verification': {'reward': 1.}} + with pytest.raises(ValueError, match='rejected'): + optimize(model, optimizer, [row], config, tokenizer, 1) + assert parameter_digest(model) == before + assert all(p.grad is None for p in model.parameters()) + assert not optimizer.state + row['sample_id'] = '000001-000000' + row['behavior_logprobs'] = response_logprobs(model, [[1, 4]], [[5, 6]], 0)[0].detach().tolist() + result = optimize(model, optimizer, [row], config, tokenizer, 1) + assert result['gradient_norm'] > 0 + assert parameter_digest(model) != before def rank_failure(rank, rendezvous, root):