diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 3c1e624c..1a6e6a34 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -52,11 +52,19 @@ jobs: # has no type stubs ("base class unknown"). # * convert_sample_to_trajectory: returns harbor's Trajectory which # has no type stubs ("return type is unknown"). + # * NativeHarborBackend.__init__ and resolve_task: their public + # parameters/return values are harbor config models, and harbor + # currently ships no type stubs. + # * NativeHarborBackend.prewarm/prewarm_lifespan: their public task + # lists intentionally use harbor TaskConfig, which has no type stubs. unexpected=$(echo "$output" | awk ' /^osmosis_ai\./ { sym = $0 } /error:/ { if (sym ~ /agent_adapter/) next if (sym ~ /convert_sample_to_trajectory/) next + if (sym ~ /native_harbor.*NativeHarborBackend.__init__/) next + if (sym ~ /native_harbor.*NativeHarborBackend.prewarm/) next + if (sym ~ /native_harbor.*resolve_task/) next print $0 } ') diff --git a/docs/README.md b/docs/README.md index 8fd28e3b..2a8112a7 100644 --- a/docs/README.md +++ b/docs/README.md @@ -27,6 +27,7 @@ The package (`osmosis_ai/`) is organized into top-level domains. See [architectu - [architecture.md](./architecture.md) — package layout, domain boundaries, import paths, lazy-loading rules, and the remote rollout protocol (controller <-> rollout server). Start here. - [rollout-sdk.md](./rollout-sdk.md) — the library API you implement against: `AgentWorkflow`, `Grader`, contexts, configs, `create_rollout_server`, execution backends, and framework integrations. +- [native-harbor-backend.md](./native-harbor-backend.md) — the `NativeHarborBackend`: run a self-contained Harbor task (instruction + environment + verifier) as a rollout, no `AgentWorkflow` / `Grader` to write. - [eval.md](./eval.md) — the `osmosis eval submit` config contract (SDK-vs-backend validation, submit flow), plus a brief note on the `evaluate_rubric` / `osmosis eval rubric` LLM-as-judge API. - [datasets.md](./datasets.md) — the dataset row contract enforced by the SDK validator. - [troubleshooting.md](./troubleshooting.md) — engineering issues (rollout timeouts, event-loop blocking, concurrency tuning). diff --git a/docs/architecture.md b/docs/architecture.md index ea92a81b..1e9c8b2d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -80,7 +80,7 @@ sequenceDiagram C-->>W: LLM response (tool_calls) Note over W: repeat until done S->>C: POST completion_callback_url (RolloutCompleteRequest) - S->>G: grade collected samples + S->>G: grade the collected sample S->>C: POST grader_callback_url (GraderCompleteRequest) ``` diff --git a/docs/native-harbor-backend.md b/docs/native-harbor-backend.md new file mode 100644 index 00000000..e0e95e3a --- /dev/null +++ b/docs/native-harbor-backend.md @@ -0,0 +1,354 @@ +# Native Harbor backend + +> The `NativeHarborBackend` execution backend. Anchored to [../osmosis_ai/rollout/backend/native_harbor/backend.py](../osmosis_ai/rollout/backend/native_harbor/backend.py). For the rollout protocol and the other backends see [architecture.md](./architecture.md) and [rollout-sdk.md](./rollout-sdk.md); for the dataset row contract see [datasets.md](./datasets.md). + +`NativeHarborBackend` turns each rollout into one native Harbor `Trial`: it resolves a Harbor task from the dataset row, runs the task's own agent against the controller-provided model endpoint, and maps the task's own verifier reward onto the rollout's single sample. You do **not** write an `AgentWorkflow`, a `Grader`, or a `SampleSource` — the Harbor task supplies the instruction, the environment, and the reward. + +It uses the SDK-pinned Harbor line (`harbor[daytona]>=0.20.0,<0.21`) and is **not** re-exported from `osmosis_ai.rollout`; import it from its subpackage: + +```python +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +``` + +## When to use it + +| You want | Use | +|----------|-----| +| Write the agent loop + grading in Python | `LocalBackend` ([rollout-sdk.md](./rollout-sdk.md)) | +| Run your Python `AgentWorkflow` inside a Harbor container | `HarborBackend` ([rollout-sdk.md](./rollout-sdk.md)) | +| Run an existing self-contained Harbor task (instruction + environment + tests) as the rollout | **`NativeHarborBackend`** | + +The clean fit is a task set like Terminal Bench, where every task is already a native Harbor task that bundles its own Docker environment and `tests/`. The rollout becomes "point Harbor at the task and read back its reward" — no glue code. + +## Shape: one Trial per rollout + +The agent is **fixed per backend** (chosen once at construction); only the **task** and (optionally) the **model** vary per rollout, carried on the dataset row's `metadata`. Each `execute()` builds a `TrialConfig`, runs it through a bounded [`TrialQueue`](../osmosis_ai/rollout/backend/native_harbor/backend.py), reads `result.verifier_result.rewards`, and fires the workflow + grader callbacks. A rollout produces exactly one sample and one reward. + +The dataset row's `system_prompt` / `user_prompt` (the wire `prompt` / `initial_messages`) are **ignored**: what enters training is the prompt the Harbor task's agent actually sends to the model endpoint, not the row text. Rows only need to carry the task reference (see [Dataset contract](#dataset-contract)). + +## Quickstart + +A native rollout server is the standard `create_rollout_server(backend=...)` wiring with a `NativeHarborBackend` instance. The resulting FastAPI app must be exposed as the module-level name `app`; `osmosis submit` imports that app to verify the actual backend binding without executing `main()`: + +```python +import os + +import uvicorn +from harbor.models.trial.config import AgentConfig + +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.server import create_rollout_server + + +backend = NativeHarborBackend( + agent=AgentConfig( + name="terminus-2", # the default; training-safe binding + model_name="openai/osmosis-rollout", + override_setup_timeout_sec=300, # setup/install, not the agent run + ), + max_concurrent=4, # running Trials + max_queue_depth=4, # accepted rollouts waiting for a slot +) +app = create_rollout_server(backend=backend) # required module-level ASGI app + + +def main() -> None: + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("_OSMOSIS_ROLLOUT_PORT", "8000"))) + + +if __name__ == "__main__": + main() +``` + +The model endpoint and key are **not** configured here — they arrive per rollout from the ambient `RolloutContext` (see [Model endpoint injection](#model-endpoint-injection)). To exercise the server locally, set the same env vars a training controller would: `OSMOSIS_CHAT_COMPLETIONS_URL`, `OSMOSIS_API_KEY`, `OSMOSIS_ROLLOUT_ID` ([context.py](../osmosis_ai/rollout/context.py)). + +The plain Harbor `Trial` path used here does not invoke Harbor's telemetry reporting call sites. Managed images should nevertheless set `HARBOR_TELEMETRY=0` as a belt-and-suspenders opt-out. + +Only agents the training path can consume are registered. The train and eval controllers each expose exactly one model-facing route per rollout, `POST /sessions/{id}/v1/chat/completions`, so an agent that speaks a different wire protocol (OpenAI Responses, Anthropic Messages) is not reachable and is not admitted. The SDK does not translate between protocols; see [Agents](#agents). + +### Supplying full Harbor configuration + +The Quickstart leaves `environment` and `verifier` at their defaults. The canonical constructor accepts Harbor's complete `AgentConfig`, `EnvironmentConfig`, and `VerifierConfig`, including skills, MCP servers, log filters, network settings, mounts, resource controls, and custom verifier fields: + +```python +from harbor.models.environment_type import EnvironmentType +from harbor.models.trial.config import AgentConfig, EnvironmentConfig, VerifierConfig + +backend = NativeHarborBackend( + agent=AgentConfig( + name="terminus-2", + model_name="openai/osmosis-rollout", + skills=["org/my-skill@sha256:..."], + include_logs=["*.json"], + extra_allowed_hosts=["models.example.com"], + ), + environment=EnvironmentConfig( + type=EnvironmentType.DAYTONA, + override_cpus=4, + ), + verifier=VerifierConfig( + max_timeout_sec=300, + include_logs=["reward.json"], + ), + max_concurrent=8, +) +``` + +Each Trial is one environment instance, so keep `max_concurrent` aligned with the host/remote capacity (see [Concurrency and trial directories](#concurrency-and-trial-directories)). For managed SkyPilot runs, an explicit `environment.kwargs["context_name"]` wins; otherwise the backend reads `HARBOR_SKYPILOT_CONTEXT`. On macOS local Docker, controller URLs using `localhost` or `127.0.0.1` are rewritten to `host.docker.internal` so the agent container can reach them. + +### Prewarm setup before the server becomes ready + +`NativeHarborBackend.prewarm()` runs one Harbor `TrialConfig(install_only=True)` per supplied `TaskConfig`. The convenience `prewarm_lifespan()` turns that into a startup gate for the same FastAPI server: + +```python +from pathlib import Path + +from harbor.models.trial.config import TaskConfig + + +prewarm_tasks = [ + TaskConfig(path=Path("./tasks/foo")), + TaskConfig(name="org/task", ref="sha256:9f2c..."), + TaskConfig( + git_url="https://example.com/org/tasks.git", + path=Path("tasks/bar"), + git_commit_id="abc123...", + ), +] + +app = create_rollout_server( + backend=backend, + lifespan=backend.prewarm_lifespan(prewarm_tasks), +) +``` + +The lifespan finishes prewarming before FastAPI accepts health checks or rollout requests. For an existing custom lifespan, call `await backend.prewarm(tasks)` before yielding. The backend clones the task list and its full agent/environment/verifier templates, assigns unique `native-prewarm-*` trial names, and uses the same `max_concurrent` queue with Harbor retries pinned to zero. It needs no `RolloutContext`, controller URL/key, or callbacks. Harbor resolves the task, starts and health-checks the environment, uploads configured skills, and runs agent setup/install; it skips the agent run and verification, so it makes no model call and produces no reward. + +Every configured task is attempted. Raised setup exceptions and Harbor `result.exception_info` failures are reported together, and any failure aborts server startup. The aggregate names each task and exception type, but keeps raw setup output out of startup logs because it may contain configured credentials. When Harbor got far enough to create a failed trial directory, the aggregate points to it for details; earlier resolution failures explicitly say that no directory was created. Successful prewarm trial directories follow `cleanup_successful_trials`; failed directories that exist are retained for inspection. Use immutable package digests and git commits in this startup list just as in dataset metadata. + +This list is a one-shot startup preparation plan, **not** a rollout-server work list: it does not sample, retry, grade, or schedule controller rollouts. Miles or eval still owns all real work. Harbor describes `install_only` as a fast setup compatibility check, and durable cache benefit depends on the environment provider and its lifecycle configuration. In particular, default local Docker uses `EnvironmentConfig.delete=True`; finalization removes the container, local image, and volumes, so an installed CLI is not guaranteed to survive into a later Trial. Prewarming can prime only the task/package/image caches that the chosen provider actually retains. It also does not run the real Miles/eval closed loop and makes no new training-safety or E2E claim. + +## Dataset contract + +Each row points at a Harbor task through a first-class `metadata` key. The dataset schema and validator are unchanged ([datasets.md](./datasets.md)); `system_prompt` / `user_prompt` stay required by the validator but are ignored at execution time. `resolve_task` ([backend.py](../osmosis_ai/rollout/backend/native_harbor/backend.py)) accepts three forms of `metadata["harbor_task"]`: + +```jsonc +// Local path — task directory shipped with the rollout server (recommended for v1). +// Triggered when harbor_task starts with "./", "/", or "~". +{ "system_prompt": "", "user_prompt": "", "metadata": { "harbor_task": "./tasks/foo" } } + +// Package — "org/name[@ref]" (must contain a "/"); resolved via Harbor's registry + cache. +{ "system_prompt": "", "user_prompt": "", "metadata": { "harbor_task": "org/name@sha256:9f2c..." } } + +// Git — set git_url; harbor_task can be any non-path marker (e.g. "git"). +{ "system_prompt": "", "user_prompt": "", + "metadata": { "harbor_task": "git", "git_url": "https://…", "task_path": "tasks/foo", "git_commit_id": "sha…" } } +``` + +`metadata["harbor_task"]` is **required** — a missing value raises `ValueError`. Keep its shape consistent across all rows (the dataset validator gates on a uniform `metadata` shape). Resolution, download, and the `~/.cache/harbor` content-hash cache are all handled by Harbor's `Trial.create()`; the backend never writes a loader. + +Pin network-resolved tasks so every rollout in a long run executes the same bytes. Package references without an `@ref`, or with explicit `@latest`, still resolve as before but log a warning; use the package's immutable `sha256:` digest. Git tasks without a non-blank `metadata["git_commit_id"]` likewise log an actionable warning; set it to the desired commit SHA. These are warnings, not validation errors, and warning logs never include the git URL. + +`metadata["harbor_model"]` (optional) overrides the backend's `model_name` for that single row. The selected binding still validates the provider prefix: Chat Completions bindings require `openai/...` so a row cannot silently switch the agent to a different wire protocol. + +> **v1 recommendation: local-path tasks.** Ship the task directories with the rollout server (baked into the image or mounted). It is offline, needs no registry auth, and matches the "everything lives on the rollout server" model. Package/git forms work but need network access (and, for packages, Harbor credentials) on the rollout host. + +Dataset construction has a few controller-owned constraints: + +- Platform datasets require `system_prompt` and `user_prompt` columns and at least four rows, even though Native ignores the prompt columns at execution. +- Eval samples 10% of rows with seed 42 when no limit is set. Set an explicit limit that covers the whole task set when full coverage matters. +- For training, one dataset row becomes one GRPO group (eight rollouts per row by default). +- As an acceptance check, run eval with the `oracle` binding; every valid task should receive reward `1.0`. + +## Constructor reference + +All arguments are keyword-only ([backend.py](../osmosis_ai/rollout/backend/native_harbor/backend.py)). + +| Argument | Default | Purpose | +|----------|---------|---------| +| `agent` | `AgentConfig(name="terminus-2")` | Complete Harbor agent configuration. Its name/import path selects the validated binding; all preserved fields are cloned per rollout. | +| `environment` | `EnvironmentConfig()` | Complete Harbor environment configuration (Docker, Daytona, SkyPilot, or a custom import). | +| `verifier` | `VerifierConfig()` | Complete Harbor verifier configuration. Native always enables it because it is the reward source. | +| `agent_setup_timeout_sec` | `None` | Compatibility overlay for `AgentConfig.override_setup_timeout_sec`; prefer setting the field on `agent`. | +| `binding` | agent name | Validated wire/identity binding. Must match the agent name for built-ins; import-path agents must select a custom binding explicitly (see [Bringing your own agent loop](#bringing-your-own-agent-loop)). | +| `model_name` | `agent.model_name`, else `"openai/osmosis-rollout"` | Compatibility/default overlay. Per-row `metadata["harbor_model"]` wins, subject to the binding's provider restriction. | +| `reward_key` | `"reward"` | Which named verifier channel becomes the scalar reward (see [Reward mapping](#reward-mapping)). | +| `trials_dir` | `Path("native_trials")` | Where Harbor writes trial directories. | +| `task_resolver` | `resolve_task` | Override `ExecutionRequest -> TaskConfig` resolution. | +| `max_concurrent` | `8` | In-flight Trial cap (`>= 1`). Each Trial is often a container, so this bounds host load. | +| `max_queue_depth` | `max_concurrent` | Maximum accepted rollouts waiting beyond the running cap (`>= 0`). Set `0` to reject whenever all Trial slots are occupied. | +| `cleanup_successful_trials` | `True` | Delete a successful trial only after its ATIF has been validated and Harbor-collected artifacts have been copied out. | +| `agent_name`, `agent_import_path`, `agent_kwargs`, `agent_env`, `environment_config` | `None` | Compatibility shims for the original reduced surface. They cannot be mixed with the corresponding canonical object. | + +### Configuration ownership and cloning + +The constructor deep-clones all three Harbor objects immediately and again for every rollout. Harbor may resolve agent skills in place, so no rollout can mutate the caller's objects or another rollout's nested dictionaries, lists, env, skills, or MCP definitions. Cloning uses Pydantic's `model_copy(deep=True)` rather than a serialization round trip; Harbor's serializers intentionally redact or templatize sensitive environment values. + +| Configuration field | Native policy | +|---|---| +| `AgentConfig.name` / `import_path` | Preserve; they select the binding. Setting both is rejected, and an `import_path` resolving to any Harbor built-in is rejected. | +| `model_name` | Overlay: dataset row `harbor_model` > explicit constructor `model_name` > `agent.model_name` > SDK default. | +| `kwargs` | Preserve, then overlay binding-owned `api_base`, `llm_kwargs.api_key`, and `extra_body.stream=False`. | +| `env` | Preserve, except that an environment-wired binding owns `OPENAI_BASE_URL` and `OPENAI_API_KEY`; setting either is rejected rather than silently overwritten. All other variables, including other providers' credentials, pass through untouched. | +| `skills`, `mcp_servers`, `include_logs`, `exclude_logs`, `extra_allowed_hosts` | Preserve. | +| `override_setup_timeout_sec` | Preserve unless `agent_setup_timeout_sec` explicitly overlays it. | +| `override_timeout_sec` | Preserve unless the rollout request supplies `agent_timeout_sec`. | +| `max_timeout_sec` | Preserve as the user's safety cap. With Native-owned timeout multipliers at `1.0`, it caps the request overlay. | +| `n_concurrent`, `concurrency_group`, `resume_trajectory=True` | Reject at construction; they conflict with the backend's queue/single-session ownership. | +| All `EnvironmentConfig` fields | Preserve. Managed SkyPilot fills `kwargs.context_name` only when the user left it unset. | +| `VerifierConfig.disable` | SDK-owned: `False` for real rollouts because the verifier produces the reward; Harbor sets it to `True` on setup-only prewarm Trials. | +| `VerifierConfig.override_timeout_sec` | Preserve unless the rollout request supplies `grader_timeout_sec`. | +| Other `VerifierConfig` fields, including `max_timeout_sec` | Preserve. | +| Trial-level task, name, directory, job/source/install flags, and timeout multipliers | SDK-owned; they are not constructor fields. | + +## Agents + +Agent support is binding-specific. A binding records the wire protocol, identity channel, training status, and an exact CLI version for installed agents. + +**Admission is training-parity.** A binding is registered only when the training path supports it. Eval deliberately does not get a wider agent set: widening it would mean carrying protocol support the trainer cannot use, for agents that could never graduate to training anyway. Anything not in the table below — including Harbor built-ins such as `codex`, `opencode`, and `claude-code` — fails at construction rather than inheriting generic `OPENAI_*` wiring. + +| Binding | Protocol / identity | Train | Status | +|---|---|---:|---| +| `terminus-2` | Chat Completions via `kwargs["api_base"]` and `kwargs["llm_kwargs"]["api_key"]` | ✓ | Summarization is off by default. | +| `oracle` | No model endpoint | ✗ | Drives no model at all, so the training-parity gate does not apply. Emits a construction warning; use it to validate datasets and verifiers. | +| `custom-chat-completions` | Chat Completions via `kwargs["api_base"]` / `kwargs["llm_kwargs"]["api_key"]`, for your own in-process `AgentConfig.import_path` agent | ✓ | Bring your own agent. | +| `custom-installed-chat-completions` | Chat Completions via `env["OPENAI_BASE_URL"]` / `env["OPENAI_API_KEY"]`, for your own container-side `AgentConfig.import_path` agent | ✓ | Bring your own agent loop — the shape a Harbor `BaseInstalledAgent` subclass takes. | + +### Bringing your own agent loop + +Harbor lets you run an agent you wrote instead of one of its built-ins, and Native supports that. Select it with `AgentConfig.import_path` plus the custom binding whose identity channel your agent actually accepts: + +- **`custom-installed-chat-completions`** — your agent subclasses Harbor's `BaseInstalledAgent` and runs inside the task container. It receives `OPENAI_BASE_URL` and `OPENAI_API_KEY` through the rollout-scoped `AgentConfig.env`, which Harbor's environment resolution prefers over host state. This is the usual shape for a custom loop. +- **`custom-chat-completions`** — your agent runs in the rollout server process and takes `api_base` / `llm_kwargs` constructor arguments, like `terminus-2`. + +Pick by how your agent takes configuration: an installed agent never sees constructor kwargs, and an in-process agent never sees the container environment, so the wrong binding silently delivers nothing. + +Both bindings own exactly two environment slots, `OPENAI_BASE_URL` and `OPENAI_API_KEY` (setting either yourself is rejected rather than silently overwritten). **Everything else in `agent.env` passes through untouched**, including other providers' credentials: a custom loop commonly routes only its policy model to the rollout endpoint and calls other providers directly for sub-agents, planners, or judges. + +Rollout identity is carried in the endpoint URL itself, so your agent does not need to stamp `x-rollout-id` / `x-sample-id` headers — pointing an OpenAI-compatible client at `OPENAI_BASE_URL` is enough. (Older Osmosis integrations sent those headers; they are no longer required.) + +Neither binding claims training safety on your behalf. Both warn at construction: the SDK wires identity into an agent it cannot inspect, so confirming the loop keeps one append-only trajectory is yours to do. + +Two reasons a Harbor **built-in** agent is absent: + +- **Unreachable protocol.** `codex` and `opencode` speak OpenAI Responses; `claude-code` speaks Anthropic Messages. The controller serves only `/v1/chat/completions`, so these cannot reach it without protocol translation, and translation would only ever buy eval. +- **Opaque context management.** Compaction, subagents, and session rewriting fork the token trajectory. That is a training blocker independent of protocol, so making these agents reachable would not make them trainable — see [Append-only trajectories](#append-only-trajectories-training-caveat). + +`import_path` cannot reintroduce them. The guard resolves the class against Harbor's own agent registry rather than the table above, so naming `harbor.agents.installed.codex:Codex` through a custom binding is rejected. Registered built-ins are likewise rejected by import path so their binding and CLI pin cannot be bypassed; select those with `agent_name`. + +Installed-agent versions are binding-owned and cannot be overridden through `agent.kwargs`. This prevents Harbor's default `@latest` installs from silently changing behavior during a long run. Chat Completions bindings also reject model prefixes for other providers. + +`AgentConfig.override_setup_timeout_sec` controls only Harbor's setup/install phase for each Trial. The controller-provided `ExecutionRequest.agent_timeout_sec` remains the separate agent **run** timeout and overlays `AgentConfig.override_timeout_sec` for that rollout. The compatibility `agent_setup_timeout_sec` argument can still overlay the setup value. + +### Model endpoint injection + +Endpoint and key come from the ambient `RolloutContext` ([context.py](../osmosis_ai/rollout/context.py)) — `chat_completions_url` and `api_key` — which the controller supplies per rollout (read from `OSMOSIS_CHAT_COMPLETIONS_URL` / `OSMOSIS_API_KEY` on a container host). + +Every registered binding reaches that endpoint directly; the SDK performs no protocol translation. The endpoint must be reachable by the Harbor environment, and local-Docker URL rewriting is applied to it. + +The backend overlays the selected binding's identity slots, so configuration cannot redirect model traffic. Kwargs-wired bindings have `agent.kwargs["api_base"]` overwritten with the rollout's endpoint even when the caller set it, and pin `extra_body.stream=False` until the controllers support that streaming path. Environment-wired bindings receive the endpoint and key in `agent.env` and reject a caller-set value for those two slots. The `oracle` binding is the exception: it invokes the task's reference solution and needs no model endpoint. + +## Append-only trajectories (training caveat) + +RL training needs a single, linear, **append-only** token trajectory. Anything that rewrites the running context mid-run — summarization, compaction, subagents — forks that trajectory and corrupts the training signal. The backend does **not** gate or police this; it only sets a safe default for the built-in default agent and otherwise stays out of the way ([backend.py](../osmosis_ai/rollout/backend/native_harbor/backend.py)): + +- **`terminus-2` (the default agent)** summarizes mid-run, so the backend defaults `enable_summarize=False` + `proactive_summarization_threshold=0` on it. These are *overridable defaults* — `agent.kwargs` wins — so set `AgentConfig(kwargs={"enable_summarize": True})` to get summarization back (e.g. for long-context eval). +- **`oracle` is not training-safe** and warns when constructed: it emits no model trajectory at all. +- **A custom agent loop's linearity is yours to verify.** The SDK wires identity into it but cannot inspect how it manages context, so both custom bindings warn rather than claim safety. +- **Harbor's installed coding CLIs are not registered.** Protocol reachability alone cannot prove that an opaque installed CLI keeps one append-only trajectory, so making them reachable would not make them trainable. This is why admission is training-parity rather than protocol-parity. + +| Agent | Run (eval — reward only) | Train (needs a linear token trajectory) | +|---|---|---| +| `terminus-2` (summarize off by default) | ✓ | ✓ | +| `oracle` | ✓ | ✗ (no model trajectory) | +| your own agent loop (either custom binding) | ✓ | ✓ once you have verified it is append-only | +| `codex` / `opencode` / `claude-code` | ✗ (not registered) | ✗ | + +## Reward mapping + +Harbor verifiers emit a **named-channel** dict (`dict[str, float]`, e.g. `{"reward": 1.0}`), not a scalar. `_pick_reward` ([backend.py](../osmosis_ai/rollout/backend/native_harbor/backend.py)) collapses it: it takes the `reward_key` channel if present, else the sole value when there is exactly one channel. If multiple channels exist and none matches `reward_key`, the reward is left unset and the sample fails grading with a logged warning — set `reward_key` to the channel you want. The reward is read from the in-memory trial-level `TrialResult`, so no `reward.json` parsing is needed. A defensive legacy step-result fallback remains internal; it does not make multi-step tasks supported. + +A Harbor `TrialResult.exception_info` is authoritative: both callbacks report failure and the sample reward remains unset, even if a verifier emitted a numeric reward before the trial failed in a later phase. Failed trials can never be revived into trainable or successful eval samples by a partial reward. + +The dataset row's `ground_truth` is **not** required for native tasks — the Harbor task's verifier is self-contained. + +## Structured diagnostics + +Native results carry an emit-only diagnostics object in callback `extra_fields` (`RolloutCompleteRequest` and, for failures discovered after agent completion, `GraderCompleteRequest`). Existing controllers ignore this unknown field, so it does not change callback handling, but it makes failures attributable without parsing Harbor log text: + +```json +{ + "backend": "native_harbor", + "phase": "verification", + "harbor_exception_type": "VerifierTimeoutError", + "category": "agent_error", + "timings_sec": { + "setup": 0.12, + "environment_setup": 8.4, + "agent_setup": 1.7, + "agent": 41.2, + "verification": 3.1, + "trial": 54.6 + } +} +``` + +The backend advances phase state from Harbor's `TrialQueue` lifecycle hooks. Possible phases are `setup`, `trial_setup`, `environment_setup`, `agent_setup`, `agent`, `verification`, `grading`, and `cancelled`. Hook durations cover the pre-trial intervals; when Harbor supplies its own `TimingInfo`, the exact environment, agent-setup, agent, verifier, and total-trial durations win. Values are non-negative seconds. Successful results use the same shape with `harbor_exception_type` and `category` set to `null`. + +On failure, the exact object sent to the callback is also written to the SDK log and archived as `~/.osmosis//diagnostics.json`. When a trajectory exists, the same object is additionally embedded at `trajectory.json.extra.osmosis.result_extra_fields`. The sidecar means setup or agent failures remain inspectable even when no valid ATIF document exists. Grader-only failures that happen after the workflow callback (for example a verifier timeout) cannot alter the already-sent completion callback; they still retain the structured payload in the log and final archive. + +## Native ATIF trajectories + +When the Harbor agent writes `agent/trajectory.json`, the backend treats that ATIF document as authoritative. It validates the document with Harbor's trajectory schema and passes its native steps, reasoning, tool calls, observations, subagent references, and metadata directly to trajectory persistence; it does not reconstruct them through `RolloutSample.messages`. Before deleting a successful trial, the backend durably writes a validated, redacted provisional document beside the collected artifacts. After `execute()` returns, the server overwrites it with the final document that normalizes the root `session_id` / `trajectory_id` to the rollout id, records the original ids under `extra.osmosis`, attaches rollout metadata and reward, and overlays exact per-call metrics reported by the controller. If provisional persistence fails, the source trial is retained. + +`agent.extra` is preserved. Before the document leaves the backend, credential-shaped leaves such as `api_key`, `authorization`, `password`, `secret`, and `token` are recursively replaced with `[REDACTED]`; other agent configuration remains intact. + +ATIF availability is still an agent capability: an agent that emits no `trajectory.json` can run and be graded, but there is no native document to persist. A malformed document causes the successful trial directory to be retained for inspection instead of being deleted. Multi-step tasks are unsupported; defensive loading code may preserve unexpected step-shaped output for inspection, but that is not a supported trajectory contract. + +## Artifacts + +The SDK does not scan the sandbox or decide which task files are artifacts. User or task code publishes selected files to Harbor's conventional `/logs/artifacts` directory (or declares additional artifacts in the Harbor task configuration), and Harbor downloads those files into the host trial directory. The backend only copies that already-collected tree to `~/.osmosis//artifacts`, alongside `trajectory.json`, before cleanup. Native has no supported multi-step artifact contract; defensive handling of unexpected step directories exists only to preserve evidence rather than discard it. + +## Concurrency and trial directories + +`max_concurrent` bounds running Trials through Harbor's `TrialQueue` semaphore; because each Trial is typically a container, `max_concurrent < 1` is rejected. `max_queue_depth` separately bounds requests already accepted by `POST /rollout` but waiting beyond those running slots. It defaults to `max_concurrent`, so the default backend accepts at most 16 rollouts: 8 running and 8 queued. Once that bound is full, `/rollout` returns HTTP 429 immediately instead of spending the controller's agent deadline in an unbounded SDK queue. Set `max_queue_depth=0` to admit no work beyond the current in-flight cap. The reservation is held through callbacks and trajectory persistence, then released on success, failure, or cancellation. Admission accounting is process-local; run one rollout-server worker per configured capacity, or budget each worker's limits independently. + +The server's `/health` response preserves the backend fields and adds a live capacity snapshot. Protocol fields describe this configured server instance; Chat Completions is the only reachable protocol. For example, an idle default server configured with eight running and eight queued slots reports: + +```json +{ + "status": "ok", + "backend": "native_harbor", + "agent": "terminus-2", + "binding": "terminus-2", + "binding_protocol": "OpenAI Chat Completions", + "protocol_capabilities": [ + "OpenAI Chat Completions" + ], + "training_supported": true, + "max_concurrency": 8, + "max_queue_depth": 8, + "capacity": { + "max_concurrent": 8, + "max_queue_depth": 8, + "in_flight": 0, + "queue_depth": 0, + "available": 16, + "accepting": true + } +} +``` + +Controllers currently ignore the additional health fields, so capacity and protocol surfacing are forward-compatible rather than full negotiation. The real Miles/eval capacity-mismatch measurement remains an E2E dependency. + +Native is explicitly single-step and single-agent; a task with scripted Harbor steps is unsupported rather than deferred. Harbor trial retries are hard-disabled: each rollout id owns exactly one Trial attempt and its linear model session. A single-step trial fires the workflow callback when verification starts and the grader callback when the trial finishes. Successful trials are removed only after artifact relocation and durable provisional ATIF persistence; failed trials and successful trials whose outputs could not be safely preserved are kept for inspection. Harbor reports in-trial failures via `result.exception_info` rather than by raising, and the backend always fires the grader callback even on failure so the trainer never hangs waiting on a missing reward. + +## Submit preflight + +`osmosis submit` normally requires a Python `AgentWorkflow` + `Grader` and rejects a rollout that has neither. Native rollouts have neither (reward comes from the Harbor verifier), so the contract check special-cases them: when the workflow fails to load, `discover_native_backend` ([eval/common/cli.py](../osmosis_ai/eval/common/cli.py)) imports the entrypoint, reads its module-level `app`, and verifies the backend marker recorded by `create_rollout_server`; only an app actually bound to a `NativeHarborBackend` (or subclass) skips the Grader requirement ([workspace_directory_contract.py](../osmosis_ai/platform/cli/workspace_directory_contract.py)). Merely importing or constructing the backend is insufficient, and constructing the app only inside `main()` is intentionally not part of the submit contract. The deeper checks (task resolves, agent exists, verifier present) remain runtime responsibilities inside `Trial.create().run()`. A self-deployed native server that never goes through `osmosis submit` is unaffected. + +## See also + +- [rollout-sdk.md](./rollout-sdk.md) — `create_rollout_server`, `ExecutionBackend`, `RolloutContext`, and the `LocalBackend` / `HarborBackend` alternatives. +- [architecture.md](./architecture.md) — the controller ↔ rollout-server protocol and execution model. +- [datasets.md](./datasets.md) — the dataset row contract the `metadata` task reference rides on. diff --git a/docs/rollout-sdk.md b/docs/rollout-sdk.md index 271a05ca..568e9188 100644 --- a/docs/rollout-sdk.md +++ b/docs/rollout-sdk.md @@ -2,7 +2,7 @@ > The library API you implement against. Anchored to [../osmosis_ai/rollout/__init__.py](../osmosis_ai/rollout/__init__.py). For how rollouts run end to end see [architecture.md](./architecture.md); for usage and the `osmosis rollout` CLI see [docs.osmosis.ai](https://docs.osmosis.ai/cli/rollout/overview). -A rollout has two halves you provide: an `AgentWorkflow` (the agent loop) and a `Grader` (turns the trajectory into rewards). The SDK runs them behind an execution backend and the FastAPI server. +A rollout has two halves you provide: an `AgentWorkflow` (the agent loop) and a `Grader` (turns the trajectory into one reward). The SDK runs them behind an execution backend and the FastAPI server. ## Public surface @@ -14,7 +14,7 @@ Everything below is re-exported from `osmosis_ai.rollout` unless noted. | `Grader` | [grader.py](../osmosis_ai/rollout/grader.py) | ABC you subclass; implement `async grade(ctx)` | | `AgentWorkflowContext`, `HarborAgentWorkflowContext`, `GraderContext`, `RolloutContext`, `get_rollout_context` | [context.py](../osmosis_ai/rollout/context.py) | Execution context passed to `run` / `grade` | | `AgentWorkflowConfig`, `GraderConfig`, `ConcurrencyConfig` | [types/config.py](../osmosis_ai/rollout/types/config.py) | Pydantic config models | -| `RolloutSample`, `RolloutStatus`, `RolloutErrorCategory`, `MultiTurnMode` | [types/sample.py](../osmosis_ai/rollout/types/sample.py) | Sample + status types | +| `RolloutSample`, `RolloutStatus`, `RolloutErrorCategory` | [types/sample.py](../osmosis_ai/rollout/types/sample.py) | Sample + status types | | `create_rollout_server`, `ControllerAuth` | [server/](../osmosis_ai/rollout/server/) | FastAPI factory + bearer auth | | `ExecutionBackend`, `LocalBackend` | [backend/](../osmosis_ai/rollout/backend/) | Execution backends | | `OsmosisStrandsAgent`, `OsmosisRolloutModel` | [integrations/agents/strands.py](../osmosis_ai/rollout/integrations/agents/strands.py) | Strands integration | @@ -32,7 +32,7 @@ class AgentWorkflow[TConfig: AgentWorkflowConfig](ABC): - `run` is **async** (enforced by [validator.py](../osmosis_ai/rollout/validator.py)). - `ctx.prompt` is the initial message list; `ctx.config` is your typed config. -- The return value is not the trajectory. Samples are collected from the active `RolloutContext` (see [Samples](#samples)); the integrations register sources for you. +- The return value is not the trajectory. The rollout's single sample is collected from the active `RolloutContext` (see [Sample](#sample)); the integrations register its source for you. ## Grader @@ -45,8 +45,8 @@ class Grader(ABC): [../osmosis_ai/rollout/grader.py](../osmosis_ai/rollout/grader.py) -- `ctx.get_samples()` returns the collected `dict[str, RolloutSample]` (**sync**). -- Attach rewards with `ctx.set_sample_reward(sample_id, reward)` — it raises `ValueError` for an unknown `sample_id` ([context.py](../osmosis_ai/rollout/context.py)). +- `ctx.sample` is the rollout's `RolloutSample`, or `None` if the workflow produced no sample. +- Attach the reward with `ctx.set_reward(reward)` — it raises `ValueError` when `ctx.sample` is `None` ([context.py](../osmosis_ai/rollout/context.py)). - `ctx.label` carries the dataset row's label (the ground-truth string). - `ctx.metadata` is the read-only input-side dataset row metadata. @@ -56,24 +56,24 @@ class Grader(ABC): - `AgentWorkflowContext` — `prompt: list[dict]`, `config`. - `HarborAgentWorkflowContext` — adds `environment` (Harbor `BaseEnvironment`) for `environment.exec()`, `environment.upload_file()`, etc. under `HarborBackend`. -- `GraderContext` — `label`, `samples`, `metadata` (input-side, read-only), plus `get_samples()` / `set_sample_reward()` (output-side). +- `GraderContext` — `label`, `sample`, `metadata`, `artifacts_dir`, plus `set_reward()` for the sample's output reward. - `RolloutContext` — ambient per-rollout context (chat completions URL, API key, rollout id). It is a context manager; the server enters it around execution. Local backends pass connection info directly; container runners read it from `OSMOSIS_CHAT_COMPLETIONS_URL` / `OSMOSIS_API_KEY` / `OSMOSIS_ROLLOUT_ID`. Fetch the current one with `get_rollout_context()`. -### Samples +### Sample -The workflow does not return samples; instead a `SampleSource` is registered on the **ambient** `RolloutContext` (fetched with `get_rollout_context()`, not the `ctx` passed to `run`) and called lazily at collection time: +The workflow does not return its sample; instead one `SampleSource` is registered on the **ambient** `RolloutContext` (fetched with `get_rollout_context()`, not the `ctx` passed to `run`) and called lazily at collection time: ```python from osmosis_ai.rollout import get_rollout_context -rollout_ctx = get_rollout_context() # the active RolloutContext -rollout_ctx.register_sample_source(name, source) # name must be unique per rollout -samples = await rollout_ctx.get_samples() # async -> {name: RolloutSample} +rollout_ctx = get_rollout_context() # the active RolloutContext +rollout_ctx.set_sample_source(source) # exactly one source per rollout +sample = await rollout_ctx.get_sample() # async -> RolloutSample | None ``` -`OsmosisStrandsAgent` registers a source automatically (keyed by the agent `name`/`agent_id`), so most workflows never call `register_sample_source` directly. +Registering a second source raises `ValueError`: one rollout is one agent execution, one sample, and one reward. `OsmosisStrandsAgent` registers its source automatically, so most workflows never call `set_sample_source` directly. -`RolloutSample` ([types/sample.py](../osmosis_ai/rollout/types/sample.py)) fields: `id`, `messages`, `label`, `reward`, `remove_sample`, `metrics`, `extra_fields`. +`RolloutSample` ([types/sample.py](../osmosis_ai/rollout/types/sample.py)) fields: `messages`, `label`, `reward`, `remove_sample`, `metrics`, `extra_fields`. It has no sample id; rollout identity comes from rollout-scoped URLs. The internal `trajectory_messages` field holds the normalized copy used for persistence. ## Artifacts @@ -91,8 +91,7 @@ async def grade(self, ctx: GraderContext) -> Any: (ctx.artifacts_dir / "trace.json").write_text( json.dumps({"score_reason": "matched rubric"}) ) - for sample_id in ctx.get_samples(): - ctx.set_sample_reward(sample_id, 1.0) + ctx.set_reward(1.0) ``` After each rollout the artifacts land on the host under `~/.osmosis//artifacts/`. `LocalBackend` writes your files at that root. Harbor mirrors its collected-trial layout, so the `/logs/artifacts/` convention dir lands at `.../artifacts/logs/artifacts/`, next to any paths you declare in the task's `artifacts` config. @@ -109,13 +108,11 @@ Layout per rollout, keyed by `rollout_id` (callers that need position semantics ``` ~/.osmosis// -├── trajectory.json # the rollout's ATIF document -│ # (trajectory-.json per sample while the -│ # transitional multi-sample protocol is still in use) +├── trajectory.json # the rollout's single ATIF document └── artifacts/... # file artifacts (see above) ``` -Each document carries a normalized, controller-compatible transcript as ATIF steps (tool calls fold into agent-step observations) and namespaces platform context under `extra.osmosis`: `rollout_id`, `sample_id`, `label`, `reward`, sample `metrics`/`extra_fields`, and the request's `metadata`/`extra_fields` (the natural channel for run identity such as an eval run id). +Each document carries a normalized, controller-compatible transcript as ATIF steps (tool calls fold into agent-step observations) and namespaces platform context under `extra.osmosis`: `rollout_id`, `label`, `reward`, sample `metrics`/`extra_fields`, and the request's `metadata`/`extra_fields` (the natural channel for run identity such as an eval run id). Built-in sample sources keep their framework-native `RolloutSample.messages` for graders and callbacks and prepare a separate `trajectory_messages` copy through the framework converter used for OpenAI-compatible `/chat/completions` traffic. `trajectory_messages` is SDK-internal: it crosses backend boundaries for persistence but is omitted from grader callbacks. This is not an exact wire replay because call-specific conversion arguments and separately supplied system instructions are not retained. Framework-native items omitted by the framework converter are outside the persisted transcript contract. Custom sample sources whose native history is already OpenAI chat-completions-shaped get the same behavior by default. A source with another native shape sets `RolloutSample.trajectory_messages` itself from `get_sample` (an explicit `None` marks conversion as unavailable and skips trajectory persistence for that sample). Like artifacts, conversion and saving are best-effort: failures are logged and never affect rewards, callbacks, or rollout status. @@ -126,7 +123,7 @@ ATIF has first-class slots for LLM operational data (`Step.metrics`, `Step.model 1. **Controller report (callback ack)** — the controller may attach a `trajectory` object to the JSON body of its completion/grader callback response ([report.py](../osmosis_ai/rollout/trajectory/report.py) defines the shape). Its LLM bridge serves every completion, so it is the party that has per-call usage. - **When to report**: snapshot the agent-phase calls into the **completion** ack, before resolving any internal future that triggers controller-side cleanup. Omit `trajectory` from the grader ack — an ack without a report keeps the earlier one, and grader-phase LLM calls (an LLM judge) would skew call counts and totals. A grader ack that does carry a report replaces the completion one entirely (no merge). - **Attribution**: `llm_call_metrics` map onto agent steps in dispatch order only when the counts match exactly; on a mismatch they are preserved under `extra.osmosis.unmatched_llm_call_metrics` instead of being mis-attributed, and totals still aggregate into `final_metrics`. The SDK always fills `final_metrics.total_steps` from the emitted ATIF steps. - - **Sample keys**: use the rollout's sample ids (the SDK integrations send them as the `x-sample-id` header on every completion). A controller that cannot know them may key its only entry arbitrarily — with exactly one sample and one entry they match regardless of key. Other unmatched entries are logged and, for single-sample rollouts, preserved under `extra.osmosis.unmatched_sample_reports`. + - **Report entry key**: `TrajectoryReport` retains a `samples` map for controller compatibility, but the SDK sample has no id. When the map contains exactly one entry, the SDK applies it to the rollout's sample regardless of the key. Multiple entries cannot be attributed and are preserved under `extra.osmosis.unmatched_sample_reports`. The SDK integrations do not stamp `x-rollout-id` or `x-sample-id` headers. ```jsonc // response body of POST or @@ -135,7 +132,7 @@ ATIF has first-class slots for LLM operational data (`Step.metrics`, `Step.model "trajectory": { "model_name": "openai/gpt-5-mini", "samples": { - "": { + "": { "llm_call_metrics": [ {"prompt_tokens": 120, "completion_tokens": 40, "cached_tokens": 0, "cost_usd": 0.0003, "logprobs": [-0.1], "model_name": "...", @@ -186,10 +183,12 @@ app = create_rollout_server(backend=backend) # FastAPI: POST /rollout, GET /he ``` - `create_rollout_server` ([server/app.py](../osmosis_ai/rollout/server/app.py)) wires the protocol: it runs the backend in a background task and posts the completion + grader callbacks. +- The controller supplies rollout-scoped `chat_completions_url`, `completion_callback_url`, and `grader_callback_url` values. Routing identity lives in those URLs; `rollout_id` in request and callback bodies is optional correlation metadata, and integrations do not add per-call rollout/sample routing headers. - `ControllerAuth` ([server/auth.py](../osmosis_ai/rollout/server/auth.py)) supplies the bearer headers for callbacks. - `ExecutionBackend` ([backend/base.py](../osmosis_ai/rollout/backend/base.py)) is the ABC; pick one: - `LocalBackend` ([backend/local/](../osmosis_ai/rollout/backend/local/)) — runs workflow + grader in-process. Re-exported from `osmosis_ai.rollout`. Used by the scaffold and eval. - `HarborBackend` ([backend/harbor/backend.py](../osmosis_ai/rollout/backend/harbor/backend.py)) — runs the agent inside a Harbor container; pairs with `HarborAgentWorkflowContext`. It is **not** re-exported (import `from osmosis_ai.rollout.backend.harbor.backend import HarborBackend`) and requires the external `harbor` dependency. + - `NativeHarborBackend` ([backend/native_harbor/backend.py](../osmosis_ai/rollout/backend/native_harbor/backend.py)) — turns one self-contained Harbor task into one rollout, one sample, and one verifier reward. Import it from `osmosis_ai.rollout.backend.native_harbor`; see [native-harbor-backend.md](./native-harbor-backend.md). ### Running a server @@ -227,9 +226,10 @@ class MyWorkflow(AgentWorkflow[MyConfig]): class MyGrader(Grader): async def grade(self, ctx: GraderContext) -> Any: - for sample_id, sample in ctx.get_samples().items(): - reward = 1.0 if str(ctx.label) in str(sample.messages[-1]) else 0.0 - ctx.set_sample_reward(sample_id, reward) + if ctx.sample is None: + raise ValueError("workflow produced no sample") + reward = 1.0 if str(ctx.label) in str(ctx.sample.messages[-1]) else 0.0 + ctx.set_reward(reward) ``` For complete, runnable rollouts (local Strands, local OpenAI Agents, Harbor) see the [Osmosis-AI/workspace-template](https://github.com/Osmosis-AI/workspace-template) `rollouts/` directory. diff --git a/osmosis_ai/eval/common/cli.py b/osmosis_ai/eval/common/cli.py index 2b7551d8..05645089 100644 --- a/osmosis_ai/eval/common/cli.py +++ b/osmosis_ai/eval/common/cli.py @@ -291,6 +291,31 @@ def load_workflow( return workflow_cls, workflow_config, entrypoint_module, None +def discover_native_backend( + rollout: str, + entrypoint: str, + *, + workspace_directory: Path | None = None, +) -> type | None: + """Return the native backend bound to the entrypoint's ASGI app, if any. + + Native submit preflight inspects the module-level app because native tasks have no Python workflow or grader. Load and inspection failures return ``None`` for the normal preflight path to report. + """ + try: + from osmosis_ai.rollout.backend.native_harbor.backend import ( + NativeHarborBackend, + ) + from osmosis_ai.rollout.server.app import _get_rollout_server_backend + + mod = _load_rollout_module( + rollout, entrypoint, workspace_directory=workspace_directory + ) + backend = _get_rollout_server_backend(vars(mod).get("app")) + return type(backend) if isinstance(backend, NativeHarborBackend) else None + except Exception: + return None + + def auto_discover_grader(module_name: str) -> tuple[type | None, Any]: """Discover a Grader subclass and its config from the entrypoint module. diff --git a/osmosis_ai/platform/cli/workspace_directory_contract.py b/osmosis_ai/platform/cli/workspace_directory_contract.py index 9d4a8345..7ed0bf6a 100644 --- a/osmosis_ai/platform/cli/workspace_directory_contract.py +++ b/osmosis_ai/platform/cli/workspace_directory_contract.py @@ -185,7 +185,11 @@ def validate_rollout_backend( Returns warnings for checks that could not run. Raises :class:`CLIError` only when the rollout is genuinely invalid. """ - from osmosis_ai.eval.common.cli import _resolve_grader, load_workflow + from osmosis_ai.eval.common.cli import ( + _resolve_grader, + discover_native_backend, + load_workflow, + ) from osmosis_ai.platform.cli.shared_config import validate_workspace_rollout_paths from osmosis_ai.rollout.validator import validate_backend @@ -226,6 +230,16 @@ def validate_rollout_backend( "The server validates it after installing the rollout's dependencies." ] if workflow_error or workflow_cls is None or entrypoint_module is None: + # Native Harbor tasks use verifier rewards instead of a Python Grader. + if ( + discover_native_backend( + rollout=rollout, + entrypoint=entrypoint, + workspace_directory=workspace_directory, + ) + is not None + ): + return [] raise CLIError( f"{command_label} preflight failed for `rollouts/{rollout}/{entrypoint}`.\n" f" {workflow_error or 'Failed to load workflow.'}" diff --git a/osmosis_ai/rollout/backend/base.py b/osmosis_ai/rollout/backend/base.py index ef8d8bc2..a92aeda2 100644 --- a/osmosis_ai/rollout/backend/base.py +++ b/osmosis_ai/rollout/backend/base.py @@ -22,5 +22,18 @@ def max_concurrency(self) -> int: """Max concurrent executions this backend supports. 0 = no limit.""" return 0 + @property + def max_queue_depth(self) -> int | None: + """Max queued executions beyond ``max_concurrency``. + + ``None`` is unbounded. Finite backends override this to reject excess work before controller deadlines expire. + """ + return None + + @property + def capture_final_result(self) -> bool: + """Whether the server should capture a final result without a grader URL.""" + return False + def health(self) -> dict[str, Any]: return {"status": "ok"} diff --git a/osmosis_ai/rollout/backend/native_harbor/__init__.py b/osmosis_ai/rollout/backend/native_harbor/__init__.py new file mode 100644 index 00000000..57c809a8 --- /dev/null +++ b/osmosis_ai/rollout/backend/native_harbor/__init__.py @@ -0,0 +1,9 @@ +from osmosis_ai.rollout.backend.native_harbor.backend import ( + NativeHarborBackend, + resolve_task, +) + +__all__ = [ + "NativeHarborBackend", + "resolve_task", +] diff --git a/osmosis_ai/rollout/backend/native_harbor/backend.py b/osmosis_ai/rollout/backend/native_harbor/backend.py new file mode 100644 index 00000000..5bc8b151 --- /dev/null +++ b/osmosis_ai/rollout/backend/native_harbor/backend.py @@ -0,0 +1,1403 @@ +"""Native Harbor execution backend. + +Run one Harbor Trial per rollout and map its verifier output to a single sample. The backend fixes the agent while the task and model may vary through metadata. +""" + +import asyncio +import copy +import importlib +import json +import logging +import math +import shutil +import traceback +import warnings +from collections.abc import AsyncIterator, Callable, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import dataclass, field +from enum import StrEnum +from pathlib import Path +from time import monotonic +from typing import Any +from uuid import uuid4 + +from harbor.agents.factory import AgentFactory +from harbor.models.agent.name import AgentName +from harbor.models.job.config import RetryConfig +from harbor.models.trajectories import Trajectory +from harbor.models.trial.config import ( + AgentConfig, + TaskConfig, + TrialConfig, + VerifierConfig, +) +from harbor.models.trial.config import ( + EnvironmentConfig as HarborEnvironmentConfig, +) +from harbor.models.trial.result import TrialResult +from harbor.trial.hooks import TrialHookEvent +from harbor.trial.queue import TrialQueue + +from osmosis_ai.rollout.backend.base import ExecutionBackend, ResultCallback +from osmosis_ai.rollout.backend.harbor.backend import ( + apply_managed_skypilot_placement, + log_trial_exception, + rewrite_url_for_docker, + uses_local_docker_runtime, +) +from osmosis_ai.rollout.context import ( + RolloutContext, + get_rollout_context, +) +from osmosis_ai.rollout.trajectory.save import _save_trajectories_with_status +from osmosis_ai.rollout.types import ( + ExecutionRequest, + ExecutionResult, + RolloutErrorCategory, + RolloutSample, + RolloutStatus, +) +from osmosis_ai.rollout.utils.file_artifacts import ( + copy_artifact_tree, + default_artifact_root, +) +from osmosis_ai.rollout.utils.rewards import validate_sample_has_reward + +logger: logging.Logger = logging.getLogger(__name__) + +HARBOR_TASK_KEY = "harbor_task" +HARBOR_MODEL_KEY = "harbor_model" +GIT_URL_KEY = "git_url" +GIT_TASK_PATH_KEY = "task_path" +GIT_COMMIT_KEY = "git_commit_id" + +DEFAULT_AGENT_NAME = "terminus-2" +DEFAULT_MODEL_NAME = "openai/osmosis-rollout" +DEFAULT_REWARD_KEY = "reward" +DEFAULT_MAX_CONCURRENT = 8 +TRIAL_NAME_PREFIX = "native-" +PREWARM_TRIAL_NAME_PREFIX = "native-prewarm-" +_BACKEND_DIAGNOSTIC_NAME = "native_harbor" + +# Disable summarization to preserve append-only trajectories. +_TERMINUS_2_DEFAULT_KWARGS: dict[str, Any] = { + "enable_summarize": False, + "proactive_summarization_threshold": 0, +} + + +class _AgentProtocol(StrEnum): + """Wire protocol used to reach the rollout controller. + + Only protocols exposed by the train and eval session routes can be admitted. + """ + + CHAT_COMPLETIONS = "OpenAI Chat Completions" + NONE = "none" + + +class _IdentityChannel(StrEnum): + """How an agent receives controller credentials. + + In-process agents use constructor kwargs; installed agents use scoped OpenAI environment variables. + """ + + KWARGS = "kwargs" + OPENAI_ENV = "openai_env" + NONE = "none" + + +@dataclass(frozen=True) +class _AgentBinding: + name: str + protocol: _AgentProtocol + identity_channel: _IdentityChannel + training_supported: bool + # ``oracle`` emits no model traffic or trajectory. + emits_model_traffic: bool = True + warning: str | None = None + allowed_model_providers: frozenset[str] | None = None + + +# Keep eval admission aligned with training; ``oracle`` is model-free. +_CUSTOM_CHAT_BINDING = "custom-chat-completions" +_CUSTOM_INSTALLED_CHAT_BINDING = "custom-installed-chat-completions" +_CUSTOM_BINDINGS = frozenset({_CUSTOM_CHAT_BINDING, _CUSTOM_INSTALLED_CHAT_BINDING}) +_AGENT_BINDINGS: dict[str, _AgentBinding] = { + DEFAULT_AGENT_NAME: _AgentBinding( + name=DEFAULT_AGENT_NAME, + protocol=_AgentProtocol.CHAT_COMPLETIONS, + identity_channel=_IdentityChannel.KWARGS, + training_supported=True, + allowed_model_providers=frozenset({"openai"}), + ), + "oracle": _AgentBinding( + name="oracle", + protocol=_AgentProtocol.NONE, + identity_channel=_IdentityChannel.NONE, + training_supported=False, + emits_model_traffic=False, + warning=( + "The native Harbor oracle binding runs the task's reference solution " + "and emits no model trajectory, so it is not training-safe. Use it to " + "validate datasets and verifiers." + ), + ), + # Caller-provided agents must preserve append-only trajectories themselves. + _CUSTOM_CHAT_BINDING: _AgentBinding( + name=_CUSTOM_CHAT_BINDING, + protocol=_AgentProtocol.CHAT_COMPLETIONS, + identity_channel=_IdentityChannel.KWARGS, + training_supported=True, + allowed_model_providers=frozenset({"openai"}), + warning=( + "The custom Chat Completions binding wires api_base/llm_kwargs identity " + "into an agent the SDK cannot inspect. Only use it with an in-process " + "agent that accepts that wiring, and confirm the agent keeps one " + "append-only trajectory before using it for training." + ), + ), + _CUSTOM_INSTALLED_CHAT_BINDING: _AgentBinding( + name=_CUSTOM_INSTALLED_CHAT_BINDING, + protocol=_AgentProtocol.CHAT_COMPLETIONS, + identity_channel=_IdentityChannel.OPENAI_ENV, + training_supported=True, + allowed_model_providers=frozenset({"openai"}), + warning=( + "The custom installed Chat Completions binding wires OPENAI_BASE_URL/" + "OPENAI_API_KEY into an agent the SDK cannot inspect. Only use it with " + "a container-side agent that reads that environment, and confirm the " + "agent keeps one append-only trajectory before using it for training." + ), + ), +} + +TaskResolver = Callable[[ExecutionRequest], TaskConfig] + +_REDACTED = "[REDACTED]" +_SENSITIVE_AGENT_EXTRA_KEYS = { + "api_key", + "apikey", + "authorization", + "credential", + "credentials", + "password", + "secret", + "token", +} + + +@dataclass +class _PendingNativeTrial: + request: ExecutionRequest + context: RolloutContext + on_workflow_complete: ResultCallback + workflow_complete_called: bool = False + workflow_result: ExecutionResult | None = None + preserve_trial: bool = False + phase: str = "setup" + phase_started_at: float | None = field(default_factory=lambda: monotonic()) + phase_timings_sec: dict[str, float] = field(default_factory=dict) + error_phase: str | None = None + error_payload: dict[str, Any] | None = None + + def transition_phase(self, phase: str) -> None: + now = monotonic() + self._finish_phase(now) + self.phase = phase + self.phase_started_at = now + + def finish_phase(self) -> None: + self._finish_phase(monotonic()) + + def timing_snapshot(self) -> dict[str, float]: + timings = dict(self.phase_timings_sec) + if self.phase_started_at is not None: + elapsed = max(0.0, monotonic() - self.phase_started_at) + timings[self.phase] = timings.get(self.phase, 0.0) + elapsed + return {key: round(value, 6) for key, value in timings.items()} + + def _finish_phase(self, now: float) -> None: + if self.phase_started_at is None: + return + elapsed = max(0.0, now - self.phase_started_at) + self.phase_timings_sec[self.phase] = ( + self.phase_timings_sec.get(self.phase, 0.0) + elapsed + ) + self.phase_started_at = None + + +def resolve_task(request: ExecutionRequest) -> TaskConfig: + """Resolve metadata["harbor_task"] to a TaskConfig: local path, package, or git.""" + md = request.metadata or {} + raw = md.get(HARBOR_TASK_KEY) + if not raw: + raise ValueError( + f"metadata[{HARBOR_TASK_KEY!r}] is required for the native harbor backend" + ) + + if isinstance(raw, str) and raw.startswith((".", "/", "~")): + return TaskConfig(path=Path(raw).expanduser()) + + if md.get(GIT_URL_KEY): + task_path = md.get(GIT_TASK_PATH_KEY) + git_commit_id = md.get(GIT_COMMIT_KEY) + if not git_commit_id or ( + isinstance(git_commit_id, str) and not git_commit_id.strip() + ): + logger.warning( + "Native Harbor rollout %s uses an unpinned git task; set " + "metadata[%r] to an immutable commit SHA", + request.id, + GIT_COMMIT_KEY, + ) + return TaskConfig( + git_url=md[GIT_URL_KEY], + path=Path(task_path) if task_path else None, + git_commit_id=git_commit_id, + ) + + name, _, ref = str(raw).partition("@") + if "/" not in name: + raise ValueError( + f"metadata[{HARBOR_TASK_KEY!r}]={raw!r} must be a local path " + "(./, /, ~), a git form (set git_url), or a package 'org/name[@ref]'" + ) + effective_ref = ref or "latest" + if effective_ref == "latest": + logger.warning( + "Native Harbor rollout %s uses package task %r with mutable ref " + "%r; pin metadata[%r] to an immutable sha256 digest", + request.id, + name, + effective_ref, + HARBOR_TASK_KEY, + ) + return TaskConfig(name=name, ref=effective_ref) + + +def _categorize_exception(exc: Exception) -> RolloutErrorCategory: + if isinstance(exc, TimeoutError): + return RolloutErrorCategory.TIMEOUT + if isinstance(exc, (ValueError, TypeError, AssertionError)): + return RolloutErrorCategory.VALIDATION_ERROR + return RolloutErrorCategory.AGENT_ERROR + + +def _is_sensitive_agent_extra_key(key: str) -> bool: + normalized = key.lower().replace("-", "_") + return normalized in _SENSITIVE_AGENT_EXTRA_KEYS or any( + normalized.endswith(f"_{suffix}") for suffix in _SENSITIVE_AGENT_EXTRA_KEYS + ) + + +def _redact_agent_extra(value: Any, *, api_key: str | None) -> Any: + """Preserve agent metadata while replacing credential-bearing leaves.""" + if isinstance(value, dict): + return { + key: _REDACTED + if _is_sensitive_agent_extra_key(str(key)) + else _redact_agent_extra(child, api_key=api_key) + for key, child in value.items() + } + if isinstance(value, list): + return [_redact_agent_extra(child, api_key=api_key) for child in value] + if isinstance(value, tuple): + return [_redact_agent_extra(child, api_key=api_key) for child in value] + if api_key and isinstance(value, str) and api_key in value: + return _REDACTED + return value + + +def _prewarm_task_label(task: TaskConfig) -> str: + """Return a safe task label for startup errors.""" + if task.name is not None: + return f"{task.name}@{task.ref}" if task.ref is not None else task.name + if task.git_url is not None: + commit = task.git_commit_id or "" + path = str(task.path) if task.path is not None else "" + # Avoid logging credential-bearing repository URLs. + return f"git:{path}@{commit}" + return str(task.path) + + +def _prewarm_failure_location(config: TrialConfig) -> str: + """Describe retained details without promising a directory was created.""" + trial_path = config.trials_dir / config.trial_name + if trial_path.exists(): + return f"inspect preserved trial {trial_path}" + return "no trial directory was created" + + +def _harbor_builtin_name_for_import_path(import_path: str) -> str | None: + """Resolve an import path against Harbor's built-in registry.""" + if ":" not in import_path: + return None + module_path, _, class_name = import_path.partition(":") + try: + imported = getattr(importlib.import_module(module_path), class_name) + except (ImportError, AttributeError): + return None + + for agent_name in AgentName: + try: + registered = AgentFactory.get_agent_class(agent_name) + except (KeyError, ValueError, ImportError, AttributeError): + continue + if imported is registered: + return agent_name.value + return None + + +def _identity_env_keys(binding: _AgentBinding) -> frozenset[str]: + """Return environment keys reserved by the binding.""" + if binding.identity_channel == _IdentityChannel.OPENAI_ENV: + return frozenset({"OPENAI_BASE_URL", "OPENAI_API_KEY"}) + return frozenset() + + +def _validate_agent_env(binding: _AgentBinding, agent_env: dict[str, str]) -> None: + conflicts = sorted(_identity_env_keys(binding).intersection(agent_env)) + if conflicts: + raise ValueError( + f"Native Harbor binding {binding.name!r} owns agent.env identity " + f"keys {conflicts!r}; remove them so model traffic cannot be " + "redirected. Other providers' credentials pass through untouched." + ) + + +def _validate_agent_config(agent: AgentConfig) -> None: + if agent.name is not None and agent.import_path is not None: + raise ValueError("agent config must set name or import_path, not both") + if agent.n_concurrent is not None: + raise ValueError( + "agent.n_concurrent is unsupported; NativeHarborBackend.max_concurrent " + "and TrialQueue own rollout concurrency" + ) + if agent.concurrency_group is not None: + raise ValueError( + "agent.concurrency_group is unsupported; NativeHarborBackend.max_concurrent " + "and TrialQueue own rollout concurrency" + ) + if agent.resume_trajectory: + raise ValueError( + "agent.resume_trajectory is unsupported for single-step native rollouts" + ) + + +def _validate_model_for_binding(binding: _AgentBinding, model_name: str) -> None: + allowed = binding.allowed_model_providers + if allowed is None: + return + provider, separator, _ = model_name.partition("/") + if not separator or provider not in allowed: + raise ValueError( + f"Native Harbor binding {binding.name!r} requires a model prefixed by " + f"one of {sorted(allowed)!r} so it uses {binding.protocol.value}; got " + f"{model_name!r}" + ) + + +def _resolve_binding( + *, + agent_name: str | None, + agent_import_path: str | None, + binding_name: str | None, +) -> _AgentBinding: + if agent_import_path is not None: + builtin_name = _harbor_builtin_name_for_import_path(agent_import_path) + if builtin_name is not None: + registered = builtin_name in _AGENT_BINDINGS + detail = ( + f"use agent_name={builtin_name!r} so its validated protocol " + "binding and CLI pin cannot be bypassed" + if registered + else f"{builtin_name!r} has no Native Harbor binding because the " + "controller cannot serve its wire protocol; import_path must not " + "reintroduce it" + ) + raise ValueError( + f"agent_import_path resolves to Harbor's built-in " + f"{builtin_name!r}; {detail}" + ) + if binding_name is None: + raise ValueError( + "agent_import_path requires an explicit custom binding so the " + "agent's wire protocol and identity channel are stated: " + f"{_CUSTOM_CHAT_BINDING!r} for an in-process agent taking " + f"api_base/llm_kwargs, {_CUSTOM_INSTALLED_CHAT_BINDING!r} for an " + "installed agent reading OPENAI_BASE_URL/OPENAI_API_KEY" + ) + if binding_name not in _CUSTOM_BINDINGS: + raise ValueError( + "agent_import_path only supports the custom bindings " + f"{', '.join(sorted(_CUSTOM_BINDINGS))}; got {binding_name!r}" + ) + else: + binding_name = binding_name or agent_name + if binding_name != agent_name: + raise ValueError( + f"binding {binding_name!r} does not match agent_name {agent_name!r}; " + "built-in agents must use their own validated binding" + ) + + selected = _AGENT_BINDINGS.get(binding_name or "") + if selected is None: + raise ValueError( + f"no validated Native Harbor binding for {binding_name!r}; supported " + f"bindings: {', '.join(sorted(_AGENT_BINDINGS))}" + ) + + if selected.emits_model_traffic and not selected.training_supported: + raise ValueError( + f"Native Harbor binding {selected.name!r} is not supported for " + "training, and eval does not admit a wider agent set; supported " + f"bindings: {', '.join(sorted(_AGENT_BINDINGS))}" + ) + + return selected + + +class NativeHarborBackend(ExecutionBackend): + """Execute one native Harbor Trial per rollout. + + The Harbor task owns the instruction, environment, and verifier; its reward maps to the rollout's single sample. + """ + + model_name: str + + def __init__( + self, + *, + agent: AgentConfig | None = None, + environment: HarborEnvironmentConfig | None = None, + verifier: VerifierConfig | None = None, + # Compatibility shims for the original reduced constructor surface. + agent_name: str | None = None, + agent_import_path: str | None = None, + agent_kwargs: dict[str, Any] | None = None, + agent_env: dict[str, str] | None = None, + agent_setup_timeout_sec: float | None = None, + binding: str | None = None, + model_name: str | None = None, + reward_key: str = DEFAULT_REWARD_KEY, + trials_dir: Path | str = Path("native_trials"), + task_resolver: TaskResolver | None = None, + environment_config: HarborEnvironmentConfig | None = None, + max_concurrent: int = DEFAULT_MAX_CONCURRENT, + max_queue_depth: int | None = None, + cleanup_successful_trials: bool = True, + ) -> None: + if max_concurrent < 1: + raise ValueError( + "max_concurrent must be >= 1; the native harbor backend spawns a " + "harbor Trial (often a container) per rollout, so unbounded " + "concurrency would exhaust the host." + ) + resolved_max_queue_depth = ( + max_concurrent if max_queue_depth is None else max_queue_depth + ) + if resolved_max_queue_depth < 0: + raise ValueError("max_queue_depth must be >= 0") + if agent_setup_timeout_sec is not None and ( + not math.isfinite(agent_setup_timeout_sec) or agent_setup_timeout_sec <= 0 + ): + raise ValueError("agent_setup_timeout_sec must be > 0 and finite") + + legacy_agent_fields = { + "agent_name": agent_name, + "agent_import_path": agent_import_path, + "agent_kwargs": agent_kwargs, + "agent_env": agent_env, + } + if agent is not None and any( + value is not None for value in legacy_agent_fields.values() + ): + supplied = sorted( + key for key, value in legacy_agent_fields.items() if value is not None + ) + raise ValueError( + f"agent cannot be combined with legacy constructor fields {supplied!r}" + ) + if environment is not None and environment_config is not None: + raise ValueError("environment and environment_config cannot both be set") + + if agent is None: + if agent_name is not None and agent_import_path is not None: + raise ValueError("set agent_name or agent_import_path, not both") + if agent_name is None and agent_import_path is None: + agent_name = DEFAULT_AGENT_NAME + agent_template = AgentConfig( + name=agent_name, + import_path=agent_import_path, + kwargs=copy.deepcopy(agent_kwargs or {}), + env=dict(agent_env or {}), + ) + else: + agent_template = agent.model_copy(deep=True) + + _validate_agent_config(agent_template) + agent_name = agent_template.name + agent_import_path = agent_template.import_path + resolved_binding = _resolve_binding( + agent_name=agent_name, + agent_import_path=agent_import_path, + binding_name=binding, + ) + _validate_agent_env(resolved_binding, agent_template.env) + + effective_model_name = ( + model_name + if model_name is not None + else agent_template.model_name + if agent_template.model_name is not None + else DEFAULT_MODEL_NAME + ) + _validate_model_for_binding(resolved_binding, effective_model_name) + agent_template.model_name = effective_model_name + + self._agent_name = agent_name + self._agent_import_path = agent_import_path + self._agent_config = agent_template + self._binding = resolved_binding + self.agent_setup_timeout_sec = agent_setup_timeout_sec + if resolved_binding.warning is not None: + warnings.warn(resolved_binding.warning, UserWarning, stacklevel=2) + self.model_name = effective_model_name + self.reward_key = reward_key + self.trials_dir: Path = Path(trials_dir) + self.task_resolver: TaskResolver = task_resolver or resolve_task + environment_template = ( + environment + if environment is not None + else environment_config + if environment_config is not None + else HarborEnvironmentConfig() + ).model_copy(deep=True) + self._environment_config: HarborEnvironmentConfig = ( + apply_managed_skypilot_placement(environment_template) + ) + verifier_template = ( + verifier.model_copy(deep=True) if verifier is not None else VerifierConfig() + ) + verifier_template.disable = False + self._verifier_config = verifier_template + self.cleanup_successful_trials = cleanup_successful_trials + self._max_concurrency = max_concurrent + self._max_queue_depth = resolved_max_queue_depth + self.artifact_root: Path = default_artifact_root() + self._pending: dict[str, _PendingNativeTrial] = {} + self._queue = TrialQueue( + n_concurrent=max_concurrent, + retry_config=RetryConfig(max_retries=0), + ) + self._queue.on_trial_started(self._on_trial_started) + self._queue.on_environment_started(self._on_environment_started) + self._queue.on_agent_started(self._on_agent_started) + self._queue.on_agent_ended(self._on_agent_ended) + self._queue.on_verification_started(self._on_verification_started) + self._queue.on_trial_ended(self._on_trial_ended) + self._queue.on_trial_cancelled(self._on_trial_cancelled) + + @property + def max_concurrency(self) -> int: + return self._max_concurrency + + @property + def max_queue_depth(self) -> int: + return self._max_queue_depth + + @property + def capture_final_result(self) -> bool: + # Harbor verification supplies the final reward. + return True + + @property + def agent_name(self) -> str | None: + return self._agent_name + + @property + def agent_import_path(self) -> str | None: + return self._agent_import_path + + @property + def binding(self) -> str: + return self._binding.name + + def health(self) -> dict[str, Any]: + return { + "status": "ok", + "backend": "native_harbor", + "agent": self.agent_name or self.agent_import_path, + "binding": self._binding.name, + "binding_protocol": self._binding.protocol.value, + "protocol_capabilities": [_AgentProtocol.CHAT_COMPLETIONS.value], + "training_supported": self._binding.training_supported, + "max_concurrency": self._max_concurrency, + "max_queue_depth": self._max_queue_depth, + } + + async def prewarm(self, tasks: Sequence[TaskConfig]) -> None: + """Run setup-only trials for every task before serving rollouts. + + Prewarm uses the backend's bounded zero-retry queue, attempts every task, and reports all setup failures together. + """ + if not tasks: + raise ValueError("prewarm requires at least one Harbor TaskConfig") + + configs = [self._build_prewarm_trial_config(task) for task in tasks] + logger.info("Prewarming %d native Harbor task(s)", len(configs)) + outcomes = await asyncio.gather( + *self._queue.submit_batch(configs), + return_exceptions=True, + ) + + failures: list[str] = [] + for config, outcome in zip(configs, outcomes, strict=True): + label = _prewarm_task_label(config.task) + if isinstance(outcome, BaseException): + # Preserve startup cancellation. + if not isinstance(outcome, Exception): + raise outcome + failures.append( + f"{label} [{type(outcome).__name__}]; " + f"{_prewarm_failure_location(config)}" + ) + continue + + exception_info = getattr(outcome, "exception_info", None) + if exception_info is not None: + exception_type = ( + getattr(exception_info, "exception_type", None) or "HarborError" + ) + failures.append( + f"{label} [{exception_type}]; {_prewarm_failure_location(config)}" + ) + continue + + if self.cleanup_successful_trials: + self._cleanup_trial(config.trial_name) + + if failures: + message = ( + f"Native Harbor prewarm failed for {len(failures)} of " + f"{len(configs)} task(s):\n" + + "\n".join(f" - {failure}" for failure in failures) + ) + logger.error(message) + raise RuntimeError(message) + + logger.info("Prewarmed %d native Harbor task(s)", len(configs)) + + def prewarm_lifespan( + self, tasks: Sequence[TaskConfig] + ) -> Callable[[object], AbstractAsyncContextManager[None]]: + """Return an ASGI lifespan that prewarms before startup. + + The task list is cloned when the lifespan is built so later caller mutations cannot change the startup plan. + """ + if not tasks: + raise ValueError("prewarm requires at least one Harbor TaskConfig") + frozen_tasks = tuple(task.model_copy(deep=True) for task in tasks) + + @asynccontextmanager + async def lifespan(_app: object) -> AsyncIterator[None]: + await self.prewarm(frozen_tasks) + yield + + return lifespan + + async def execute( + self, + request: ExecutionRequest, + on_workflow_complete: ResultCallback, + on_grader_complete: ResultCallback | None = None, + ) -> None: + ctx = get_rollout_context() or RolloutContext() + trial_name = f"{TRIAL_NAME_PREFIX}{request.id}" + pending = _PendingNativeTrial( + request=request, + context=ctx, + on_workflow_complete=on_workflow_complete, + ) + self._pending[trial_name] = pending + trial_result: TrialResult | None = None + setup_error: Exception | None = None + callback_error: Exception | None = None + try: + try: + task_cfg = self.task_resolver(request) + agent_cfg = self._build_agent_config(request, ctx) + trial_cfg = self._build_trial_config( + request, task_cfg, agent_cfg, trial_name + ) + trial_result = await self._queue.submit(trial_cfg) + except Exception as exc: + setup_error = exc + logger.error( + "Native trial %s failed to run: %s", + request.id, + traceback.format_exc(), + ) + + # Multi-step trials defer completion until the final result. + if not pending.workflow_complete_called: + workflow_result = ( + pending.workflow_result + or self._build_workflow_result( + pending, trial_result, setup_error=setup_error + ) + ) + pending.workflow_result = workflow_result + callback_error = await self._try_callback( + on_workflow_complete, workflow_result, request.id, "workflow" + ) + pending.workflow_complete_called = callback_error is None + + workflow_result = pending.workflow_result or ExecutionResult( + status=RolloutStatus.FAILURE, + err_message="Trial ended before the agent result was available", + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=self._diagnostic_payload( + pending, + trial_result, + category=RolloutErrorCategory.AGENT_ERROR, + ), + ) + final_workflow_result = ( + self._build_workflow_result( + pending, + trial_result, + setup_error=setup_error, + ) + if setup_error is not None + else workflow_result + ) + grader_result = self._build_grader_result( + pending, + final_workflow_result, + trial_result, + ) + result_to_persist = grader_result + if on_grader_complete is not None: + grader_callback_error = await self._try_callback( + on_grader_complete, grader_result, request.id, "grader" + ) + callback_error = callback_error or grader_callback_error + + successful = bool( + trial_result is not None + and getattr(trial_result, "exception_info", None) is None + ) + if result_to_persist.extra_fields is not None: + # Persist diagnostics even without an ATIF document. + persisted = await _save_trajectories_with_status( + rollout_id=request.id, + result=result_to_persist, + request_label=request.label, + request_metadata=request.metadata, + artifact_root=self.artifact_root, + ) + if successful and not persisted: + # Retain the source when archival fails. + pending.preserve_trial = True + + relocated = self._relocate_trial_artifacts(request.id) + if ( + successful + and relocated + and not pending.preserve_trial + and self.cleanup_successful_trials + ): + self._cleanup_trial(trial_name) + if callback_error is not None: + # Let the server perform final failure notification after archival. + raise callback_error + finally: + self._pending.pop(trial_name, None) + + @staticmethod + async def _try_callback( + callback: ResultCallback, result: ExecutionResult, rollout_id: str, label: str + ) -> Exception | None: + """Attempt an idempotent callback without aborting the trial. + + Delivery failures are retained until trial outputs are safe, then raised so the server can perform final notification fallback. + """ + try: + await callback(result) + except Exception as exc: + logger.error( + "Native %s callback for rollout %s failed: %s", + label, + rollout_id, + traceback.format_exc(), + ) + return exc + return None + + async def _on_trial_started(self, event: TrialHookEvent) -> None: + self._transition_pending_phase(event, "trial_setup") + + async def _on_environment_started(self, event: TrialHookEvent) -> None: + self._transition_pending_phase(event, "environment_setup") + + async def _on_agent_started(self, event: TrialHookEvent) -> None: + self._transition_pending_phase(event, "agent") + + async def _on_agent_ended(self, event: TrialHookEvent) -> None: + pending = self._pending.get(event.trial_name) + if pending is not None and pending.phase == "agent": + pending.finish_phase() + + async def _on_trial_ended(self, event: TrialHookEvent) -> None: + pending = self._pending.get(event.trial_name) + if pending is not None: + self._capture_error_phase(pending, event.result) + pending.finish_phase() + + async def _on_trial_cancelled(self, event: TrialHookEvent) -> None: + pending = self._pending.get(event.trial_name) + if pending is not None: + pending.transition_phase("cancelled") + pending.finish_phase() + + def _transition_pending_phase(self, event: TrialHookEvent, phase: str) -> None: + pending = self._pending.get(event.trial_name) + if pending is not None: + pending.transition_phase(phase) + + def _capture_error_phase( + self, pending: _PendingNativeTrial, trial_result: TrialResult + ) -> None: + if ( + pending.error_phase is None + and getattr(trial_result, "exception_info", None) is not None + ): + pending.error_phase = self._infer_phase(pending, trial_result) + + async def _on_verification_started(self, event: TrialHookEvent) -> None: + pending = self._pending.get(event.trial_name) + if pending is None: + return + # Preserve an agent failure before verification advances the phase. + self._capture_error_phase(pending, event.result) + pending.transition_phase("verification") + if pending.workflow_complete_called: + return + if event.result.step_results is not None: + # Multi-step verification is not workflow completion. + return + + result = self._build_workflow_result(pending, event.result) + pending.workflow_result = result + callback_error = await self._try_callback( + pending.on_workflow_complete, + result, + pending.request.id, + "workflow", + ) + pending.workflow_complete_called = callback_error is None + + def _diagnostic_payload( + self, + pending: _PendingNativeTrial, + trial_result: TrialResult | None, + *, + category: RolloutErrorCategory | None = None, + harbor_exception_type: str | None = None, + phase: str | None = None, + ) -> dict[str, Any]: + """Build stable callback and archive diagnostics. + + Failure payloads are cached so callbacks, logs, grader results, and archives keep identical content as later hooks update timing state. + """ + if category is not None and pending.error_payload is not None: + return pending.error_payload + + resolved_phase = ( + phase + or (pending.error_phase if category is not None else None) + or self._infer_phase(pending, trial_result) + ) + payload: dict[str, Any] = { + "backend": _BACKEND_DIAGNOSTIC_NAME, + "phase": resolved_phase, + "harbor_exception_type": harbor_exception_type, + "category": category.value if category is not None else None, + "timings_sec": self._phase_timings(pending, trial_result), + } + if category is not None: + pending.error_payload = payload + logger.error( + "Native Harbor structured error for rollout %s: %s", + pending.request.id, + json.dumps(payload, sort_keys=True), + ) + return payload + + @staticmethod + def _infer_phase( + pending: _PendingNativeTrial, trial_result: TrialResult | None + ) -> str: + phase = pending.phase + if trial_result is None: + return phase + + # Harbor exposes agent setup only through TimingInfo. + if ( + phase == "environment_setup" + and getattr(trial_result, "agent_setup", None) is not None + ): + return "agent_setup" + if phase != "setup": + return phase + + # Recover phase when older integrations omit hooks. + for attr, inferred in ( + ("verifier", "verification"), + ("agent_execution", "agent"), + ("agent_setup", "agent_setup"), + ("environment_setup", "environment_setup"), + ): + if getattr(trial_result, attr, None) is not None: + return inferred + return phase + + @classmethod + def _phase_timings( + cls, pending: _PendingNativeTrial, trial_result: TrialResult | None + ) -> dict[str, float]: + timings = pending.timing_snapshot() + if trial_result is None: + return timings + + exact: dict[str, float | None] = { + "trial": cls._duration_sec( + getattr(trial_result, "started_at", None), + getattr(trial_result, "finished_at", None), + ) + } + for attr, key in ( + ("environment_setup", "environment_setup"), + ("agent_setup", "agent_setup"), + ("agent_execution", "agent"), + ("verifier", "verification"), + ): + timing = getattr(trial_result, attr, None) + exact[key] = cls._duration_sec( + getattr(timing, "started_at", None), + getattr(timing, "finished_at", None), + ) + timings.update( + {key: round(value, 6) for key, value in exact.items() if value is not None} + ) + return timings + + @staticmethod + def _duration_sec(started_at: Any, finished_at: Any) -> float | None: + if started_at is None or finished_at is None: + return None + try: + duration = float((finished_at - started_at).total_seconds()) + except (AttributeError, TypeError, ValueError): + return None + if not math.isfinite(duration): + return None + return max(0.0, duration) + + def _build_workflow_result( + self, + pending: _PendingNativeTrial, + trial_result: TrialResult | None, + *, + setup_error: Exception | None = None, + ) -> ExecutionResult: + request = pending.request + if setup_error is not None: + category = _categorize_exception(setup_error) + return ExecutionResult( + status=RolloutStatus.FAILURE, + err_message=str(setup_error), + err_category=category, + extra_fields=self._diagnostic_payload( + pending, + trial_result, + category=category, + harbor_exception_type=type(setup_error).__name__, + ), + ) + + trajectory_document = self._load_native_trajectory( + request.id, pending.context.api_key + ) + if trajectory_document is None and self._native_trajectory_paths(request.id): + pending.preserve_trial = True + + sample = RolloutSample( + label=request.label, + trajectory_messages=None, + metrics=self._trial_metrics(trial_result), + ) + err = ( + getattr(trial_result, "exception_info", None) + if trial_result is not None + else None + ) + if err is not None: + if getattr(err, "exception_traceback", None) is not None: + log_trial_exception(request.id, err, phase="during native execution") + else: + logger.error( + "Native Harbor trial %s failed during execution [%s]: %s", + request.id, + getattr(err, "exception_type", "unknown"), + getattr(err, "exception_message", "unknown error"), + ) + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + trajectory_document=trajectory_document, + err_message=getattr(err, "exception_message", None) + or "Trial failed before completion", + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=self._diagnostic_payload( + pending, + trial_result, + category=RolloutErrorCategory.AGENT_ERROR, + harbor_exception_type=getattr(err, "exception_type", None), + ), + ) + return ExecutionResult( + status=RolloutStatus.SUCCESS, + sample=sample, + trajectory_document=trajectory_document, + extra_fields=self._diagnostic_payload(pending, trial_result), + ) + + @staticmethod + def _trial_metrics(trial_result: TrialResult | None) -> dict[str, Any]: + if trial_result is None: + return {} + try: + input_tokens, cache_tokens, output_tokens, cost_usd = ( + trial_result.compute_token_cost_totals() + ) + except AttributeError: + # Support older duck-typed results. + return {} + except Exception: + logger.warning("Failed to read native Harbor token totals", exc_info=True) + return {} + return { + key: value + for key, value in { + "input_tokens": input_tokens, + "cached_tokens": cache_tokens, + "output_tokens": output_tokens, + "cost_usd": cost_usd, + }.items() + if value is not None + } + + def _native_trajectory_paths(self, rollout_id: str) -> list[Path]: + trial_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + primary = trial_dir / "agent" / "trajectory.json" + paths: list[Path] = [] + if primary.is_file(): + paths.append(primary) + paths.extend(sorted((trial_dir / "steps").glob("*/agent/trajectory.json"))) + return paths + + def _load_native_trajectory( + self, + rollout_id: str, + api_key: str | None, + ) -> dict[str, Any] | None: + paths = self._native_trajectory_paths(rollout_id) + if not paths: + logger.info( + "Native Harbor agent emitted no ATIF trajectory for rollout %s", + rollout_id, + ) + return None + if len(paths) > 1: + logger.warning( + "Native Harbor rollout %s emitted %d independent trajectory " + "documents; preserving the trial directory instead of fabricating " + "a lossy merged trajectory", + rollout_id, + len(paths), + ) + return None + + try: + raw = json.loads(paths[0].read_text()) + trajectory = Trajectory.model_validate(raw) + except Exception: + logger.warning( + "Native Harbor emitted an invalid ATIF trajectory for rollout %s; " + "preserving the trial directory", + rollout_id, + exc_info=True, + ) + return None + + if trajectory.agent.extra is not None: + trajectory.agent.extra = _redact_agent_extra( + trajectory.agent.extra, api_key=api_key + ) + return trajectory.to_json_dict(exclude_none=True) + + def _relocate_trial_artifacts(self, rollout_id: str) -> bool: + """Copy only artifacts Harbor already collected from the environment.""" + trial_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + destinations: list[tuple[Path, Path]] = [] + source = trial_dir / "artifacts" + if source.is_dir(): + destinations.append((source, self.artifact_root / rollout_id / "artifacts")) + steps_dir = trial_dir / "steps" + if steps_dir.is_dir(): + for step_dir in sorted(steps_dir.iterdir(), key=lambda path: path.name): + step_artifacts = step_dir / "artifacts" + if step_artifacts.is_dir(): + destinations.append( + ( + step_artifacts, + self.artifact_root + / rollout_id + / "artifacts" + / "steps" + / step_dir.name, + ) + ) + + try: + for source_dir, destination_dir in destinations: + copy_artifact_tree( + source_dir, + destination_dir, + destination_root=self.artifact_root, + replace_destination=True, + ) + except Exception: + logger.warning( + "Failed to relocate native Harbor artifacts for rollout %s; " + "preserving the trial directory", + rollout_id, + exc_info=True, + ) + return False + return True + + def _cleanup_trial(self, trial_name: str) -> None: + shutil.rmtree(self.trials_dir / trial_name, ignore_errors=True) + + def _build_prewarm_trial_config(self, task_cfg: TaskConfig) -> TrialConfig: + """Clone constructor templates into one setup-only Harbor TrialConfig. + + Prewarm omits rollout credentials and disables verification while retaining agent setup defaults. + """ + agent_cfg = self._agent_config.model_copy(deep=True) + _validate_agent_config(agent_cfg) + _validate_agent_env(self._binding, agent_cfg.env) + + # Prewarm must not receive rollout credentials. + if agent_cfg.name == DEFAULT_AGENT_NAME: + agent_cfg.kwargs = { + **_TERMINUS_2_DEFAULT_KWARGS, + **agent_cfg.kwargs, + } + if self.agent_setup_timeout_sec is not None: + agent_cfg.override_setup_timeout_sec = self.agent_setup_timeout_sec + + verifier_cfg = self._verifier_config.model_copy(deep=True) + verifier_cfg.disable = True + return TrialConfig( + task=task_cfg.model_copy(deep=True), + trial_name=f"{PREWARM_TRIAL_NAME_PREFIX}{uuid4().hex}", + trials_dir=self.trials_dir, + install_only=True, + agent=agent_cfg, + verifier=verifier_cfg, + environment=self._environment_config.model_copy(deep=True), + ) + + def _build_trial_config( + self, + request: ExecutionRequest, + task_cfg: TaskConfig, + agent_cfg: AgentConfig, + trial_name: str, + ) -> TrialConfig: + verifier_cfg = self._verifier_config.model_copy(deep=True) + verifier_cfg.disable = False + if request.grader_timeout_sec is not None: + verifier_cfg.override_timeout_sec = request.grader_timeout_sec + return TrialConfig( + task=task_cfg, + trial_name=trial_name, + trials_dir=self.trials_dir, + agent=agent_cfg, + verifier=verifier_cfg, + environment=self._environment_config.model_copy(deep=True), + ) + + def _build_agent_config( + self, + request: ExecutionRequest, + ctx: RolloutContext, + ) -> AgentConfig: + md = request.metadata or {} + agent_cfg = self._agent_config.model_copy(deep=True) + _validate_agent_config(agent_cfg) + name = agent_cfg.name + model_name = md.get(HARBOR_MODEL_KEY, self.model_name) + if not isinstance(model_name, str) or not model_name.strip(): + raise ValueError( + f"rollout {request.id!r} has a harbor_model that is not a " + "non-empty string" + ) + _validate_model_for_binding(self._binding, model_name) + _validate_agent_env(self._binding, agent_cfg.env) + kwargs = agent_cfg.kwargs + env = agent_cfg.env + + endpoint: str | None = None + api_key = ctx.api_key + if self._binding.identity_channel != _IdentityChannel.NONE: + endpoint = ctx.chat_completions_url + if not endpoint: + raise ValueError(f"rollout {request.id!r} has no chat_completions_url") + if uses_local_docker_runtime(self._environment_config): + endpoint = rewrite_url_for_docker(endpoint) + + if self._binding.identity_channel == _IdentityChannel.OPENAI_ENV: + assert endpoint is not None + # AgentConfig.env overrides host credentials for this rollout. + env["OPENAI_BASE_URL"] = endpoint + if api_key: + env["OPENAI_API_KEY"] = api_key + elif self._binding.identity_channel == _IdentityChannel.KWARGS: + assert endpoint is not None + defaults = _TERMINUS_2_DEFAULT_KWARGS if name == DEFAULT_AGENT_NAME else {} + kwargs = {**defaults, **kwargs} + kwargs["api_base"] = endpoint + llm_kwargs: dict[str, Any] = dict(kwargs.get("llm_kwargs") or {}) + if api_key: + llm_kwargs["api_key"] = api_key + # Controllers require a non-streaming JSON response. + extra_body: dict[str, Any] = dict(llm_kwargs.get("extra_body") or {}) + extra_body["stream"] = False + llm_kwargs["extra_body"] = extra_body + kwargs["llm_kwargs"] = llm_kwargs + + agent_cfg.model_name = model_name + agent_cfg.kwargs = kwargs + agent_cfg.env = env + if self.agent_setup_timeout_sec is not None: + agent_cfg.override_setup_timeout_sec = self.agent_setup_timeout_sec + if request.agent_timeout_sec is not None: + agent_cfg.override_timeout_sec = request.agent_timeout_sec + return agent_cfg + + def _build_grader_result( + self, + pending: _PendingNativeTrial, + workflow_result: ExecutionResult, + trial_result: TrialResult | None, + ) -> ExecutionResult: + """Build the single-sample result from Harbor verification. + + Setup failures propagate the workflow error, and any Harbor exception invalidates a partial verifier reward. + """ + request = pending.request + sample = ( + workflow_result.sample.model_copy(deep=True) + if workflow_result.sample is not None + else RolloutSample(label=request.label, trajectory_messages=None) + ) + if trial_result is None: + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + trajectory_document=workflow_result.trajectory_document, + err_message=workflow_result.err_message, + err_category=workflow_result.err_category, + extra_fields=workflow_result.extra_fields, + ) + + err = getattr(trial_result, "exception_info", None) + if err is not None: + # A later Harbor failure invalidates any partial reward. + sample.reward = None + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + trajectory_document=workflow_result.trajectory_document, + err_message=getattr(err, "exception_message", None) + or "Trial failed before grading completed", + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=self._diagnostic_payload( + pending, + trial_result, + category=RolloutErrorCategory.AGENT_ERROR, + harbor_exception_type=getattr(err, "exception_type", None), + ), + ) + + reward_value = self._pick_reward(self._extract_rewards(trial_result)) + if reward_value is not None: + sample.reward = float(reward_value) + + try: + validate_sample_has_reward(sample) + except ValueError as e: + logger.warning("Native grading incomplete: %s", e) + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + trajectory_document=workflow_result.trajectory_document, + err_message=str(e), + err_category=RolloutErrorCategory.VALIDATION_ERROR, + extra_fields=self._diagnostic_payload( + pending, + trial_result, + category=RolloutErrorCategory.VALIDATION_ERROR, + phase="grading", + ), + ) + return ExecutionResult( + status=RolloutStatus.SUCCESS, + sample=sample, + trajectory_document=workflow_result.trajectory_document, + extra_fields=self._diagnostic_payload(pending, trial_result), + ) + + @staticmethod + def _extract_rewards(trial_result: TrialResult) -> dict[str, float | int] | None: + """Return trial-level rewards, falling back to the first step.""" + top = trial_result.verifier_result + if top is not None and top.rewards: + return top.rewards + for step in trial_result.step_results or []: + step_vr = step.verifier_result + if step_vr is not None and step_vr.rewards: + return step_vr.rewards + return None + + def _pick_reward( + self, rewards: dict[str, float | int] | None + ) -> float | int | None: + if not rewards: + return None + if self.reward_key in rewards: + return rewards[self.reward_key] + if len(rewards) == 1: + return next(iter(rewards.values())) + logger.warning( + "Native verifier returned rewards %s with no %r channel; reward unset", + sorted(rewards), + self.reward_key, + ) + return None diff --git a/osmosis_ai/rollout/server/app.py b/osmosis_ai/rollout/server/app.py index e5d21fdc..4832d919 100644 --- a/osmosis_ai/rollout/server/app.py +++ b/osmosis_ai/rollout/server/app.py @@ -1,4 +1,5 @@ import logging +import threading import traceback import uuid from typing import Any @@ -27,6 +28,60 @@ logger: logging.Logger = logging.getLogger(__name__) +_ROLLOUT_BACKEND_STATE_ATTR = "_osmosis_execution_backend" +_ROLLOUT_QUEUE_FULL_DETAIL = "Rollout queue is full" + + +class _RolloutAdmission: + """Process-local reservation accounting for accepted rollout requests.""" + + def __init__(self, backend: ExecutionBackend) -> None: + self._max_concurrent = backend.max_concurrency + self._max_queue_depth = backend.max_queue_depth + if self._max_queue_depth is not None and self._max_queue_depth < 0: + raise ValueError("backend.max_queue_depth must be >= 0 or None") + self._in_flight = 0 + self._lock = threading.Lock() + + @property + def _limit(self) -> int | None: + if self._max_concurrent <= 0 or self._max_queue_depth is None: + return None + return self._max_concurrent + self._max_queue_depth + + def reserve(self) -> bool: + """Atomically reserve one accepted request, or reject a full queue.""" + with self._lock: + limit = self._limit + if limit is not None and self._in_flight >= limit: + return False + self._in_flight += 1 + return True + + def release(self) -> None: + with self._lock: + if self._in_flight <= 0: + raise RuntimeError("rollout admission reservation underflow") + self._in_flight -= 1 + + def snapshot(self) -> dict[str, int | bool | None]: + with self._lock: + limit = self._limit + queue_depth = ( + max(0, self._in_flight - self._max_concurrent) + if self._max_concurrent > 0 + else 0 + ) + available = None if limit is None else max(0, limit - self._in_flight) + return { + "max_concurrent": self._max_concurrent, + "max_queue_depth": self._max_queue_depth, + "in_flight": self._in_flight, + "queue_depth": queue_depth, + "available": available, + "accepting": available is None or available > 0, + } + def _configure_default_logging() -> None: if logging.getLogger().handlers: @@ -52,25 +107,45 @@ def create_rollout_server( if configure_logging: _configure_default_logging() app = FastAPI(lifespan=lifespan) + setattr(app.state, _ROLLOUT_BACKEND_STATE_ATTR, backend) + admission = _RolloutAdmission(backend) @app.get("/health") async def health() -> dict[str, Any]: - return backend.health() + payload = dict(backend.health()) + payload["capacity"] = admission.snapshot() + return payload @app.post("/rollout") async def rollout( request: RolloutInitRequest, background_tasks: BackgroundTasks ) -> RolloutInitResponse: + if not admission.reserve(): + raise HTTPException(status_code=429, detail=_ROLLOUT_QUEUE_FULL_DETAIL) try: - background_tasks.add_task(_handle_rollout, backend, request) - return RolloutInitResponse() + response = RolloutInitResponse() + background_tasks.add_task( + _handle_admitted_rollout, + admission, + backend, + request, + ) except Exception as e: + admission.release() logger.error(traceback.format_exc()) raise HTTPException(status_code=500, detail=str(e)) from e + return response return app +def _get_rollout_server_backend(app: Any) -> ExecutionBackend | None: + """Return the backend recorded by ``create_rollout_server``, if any.""" + state = getattr(app, "state", None) + backend = getattr(state, _ROLLOUT_BACKEND_STATE_ATTR, None) + return backend if isinstance(backend, ExecutionBackend) else None + + async def _handle_rollout( backend: ExecutionBackend, request: RolloutInitRequest ) -> None: @@ -104,6 +179,7 @@ async def on_workflow_complete(result: ExecutionResult) -> None: payload=RolloutCompleteRequest( status=result.status, rollout_id=rollout_id, + extra_fields=result.extra_fields, err_message=result.err_message, err_category=result.err_category, ).model_dump(), @@ -135,6 +211,7 @@ async def on_grader_complete(result: ExecutionResult) -> None: if result.status == RolloutStatus.SUCCESS else GraderStatus.FAILURE, sample=result.sample, + extra_fields=result.extra_fields, err_message=result.err_message, err_category=result.err_category, ).model_dump(exclude={"sample": {"trajectory_messages"}}), @@ -160,7 +237,7 @@ async def on_grader_complete(result: ExecutionResult) -> None: ), on_workflow_complete=on_workflow_complete, on_grader_complete=on_grader_complete - if request.grader_callback_url + if request.grader_callback_url or backend.capture_final_result else None, ) logger.info("Rollout %s completed successfully", rollout_id) @@ -172,6 +249,9 @@ async def on_grader_complete(result: ExecutionResult) -> None: payload=RolloutCompleteRequest( status=RolloutStatus.FAILURE, rollout_id=rollout_id, + extra_fields=result_to_save.extra_fields + if result_to_save is not None + else None, err_message="Internal server error", ).model_dump(), headers=auth.as_bearer_headers(), @@ -198,3 +278,14 @@ async def on_grader_complete(result: ExecutionResult) -> None: request_extra_fields=request.extra_fields, report=report, ) + + +async def _handle_admitted_rollout( + admission: _RolloutAdmission, + backend: ExecutionBackend, + request: RolloutInitRequest, +) -> None: + try: + await _handle_rollout(backend, request) + finally: + admission.release() diff --git a/osmosis_ai/rollout/trajectory/converter.py b/osmosis_ai/rollout/trajectory/converter.py index 710fcb39..85a48624 100644 --- a/osmosis_ai/rollout/trajectory/converter.py +++ b/osmosis_ai/rollout/trajectory/converter.py @@ -45,6 +45,7 @@ def convert_sample_to_trajectory( request_label: str | None = None, request_metadata: dict[str, Any] | None = None, request_extra_fields: dict[str, Any] | None = None, + result_extra_fields: dict[str, Any] | None = None, report: SampleReport | None = None, default_model_name: str | None = None, unmatched_sample_reports: Mapping[str, SampleReport] | None = None, @@ -82,6 +83,7 @@ def convert_sample_to_trajectory( request_label=request_label, request_metadata=request_metadata, request_extra_fields=request_extra_fields, + result_extra_fields=result_extra_fields, unmatched_llm_call_metrics=unmatched_llm_call_metrics, unmatched_sample_reports=unmatched_reports, ), @@ -386,6 +388,7 @@ def _compose_extra( request_label: str | None, request_metadata: dict[str, Any] | None, request_extra_fields: dict[str, Any] | None, + result_extra_fields: dict[str, Any] | None, unmatched_llm_call_metrics: list[dict[str, Any]] | None, unmatched_sample_reports: dict[str, Any] | None, ) -> dict[str, Any]: @@ -398,6 +401,7 @@ def _compose_extra( "sample_extra_fields": sample.extra_fields or None, "request_metadata": request_metadata, "request_extra_fields": request_extra_fields, + "result_extra_fields": result_extra_fields, "unmatched_llm_call_metrics": unmatched_llm_call_metrics, "unmatched_sample_reports": unmatched_sample_reports, } diff --git a/osmosis_ai/rollout/trajectory/report.py b/osmosis_ai/rollout/trajectory/report.py index fcc574fd..4dfae855 100644 --- a/osmosis_ai/rollout/trajectory/report.py +++ b/osmosis_ai/rollout/trajectory/report.py @@ -8,7 +8,7 @@ Report in the completion ack: a grader ack without a report keeps the earlier one, one with a report replaces it wholesale. A worked example -plus timing and sample-key guidance live in docs/rollout-sdk.md. +plus timing and report-entry guidance live in docs/rollout-sdk.md. """ import logging diff --git a/osmosis_ai/rollout/trajectory/save.py b/osmosis_ai/rollout/trajectory/save.py index 95e560c9..d61adc6e 100644 --- a/osmosis_ai/rollout/trajectory/save.py +++ b/osmosis_ai/rollout/trajectory/save.py @@ -5,13 +5,20 @@ """ import asyncio +import json import logging +from collections.abc import Mapping from pathlib import Path from typing import Any +from harbor.models.trajectories import FinalMetrics, Trajectory from harbor.utils.trajectory_utils import format_trajectory_json -from osmosis_ai.rollout.trajectory.converter import convert_sample_to_trajectory +from osmosis_ai.rollout.trajectory.converter import ( + _apply_report, + _final_metrics_from_report, + convert_sample_to_trajectory, +) from osmosis_ai.rollout.trajectory.report import SampleReport, TrajectoryReport from osmosis_ai.rollout.types import ExecutionResult from osmosis_ai.rollout.utils.file_artifacts import default_artifact_root @@ -40,6 +47,82 @@ def _resolve_sample_report( return None, dict(report.samples) +def _prepare_native_trajectory( + document: dict[str, Any], + *, + rollout_id: str, + result: ExecutionResult, + request_label: str | None, + request_metadata: dict[str, Any] | None, + request_extra_fields: dict[str, Any] | None, + report: TrajectoryReport | None, +) -> Trajectory: + """Enrich a native ATIF trajectory without rebuilding its steps. + + Native agent, tool, and observation structure remains authoritative while controller metrics and Osmosis rollout metadata are overlaid. + """ + trajectory = Trajectory.model_validate(document) + native_session_id = trajectory.session_id + native_trajectory_id = trajectory.trajectory_id + + matched_report, unmatched_reports = _resolve_sample_report(report) + unmatched_llm_call_metrics = _apply_report(trajectory.steps, matched_report) + + controller_metrics = _final_metrics_from_report(matched_report) + if trajectory.final_metrics is None: + trajectory.final_metrics = FinalMetrics() + if controller_metrics is not None: + for field, value in controller_metrics.model_dump(exclude_none=True).items(): + if field != "total_steps": + setattr(trajectory.final_metrics, field, value) + trajectory.final_metrics.total_steps = len(trajectory.steps) + + reported_model = ( + matched_report.model_name if matched_report is not None else None + ) or (report.model_name if report is not None else None) + if reported_model: + trajectory.agent.model_name = reported_model + + # Preserve Harbor IDs before normalizing them for platform joins. + trajectory.session_id = rollout_id + trajectory.trajectory_id = rollout_id + + extra = dict(trajectory.extra or {}) + existing_osmosis = extra.get("osmosis") + osmosis = dict(existing_osmosis) if isinstance(existing_osmosis, Mapping) else {} + sample = result.sample + updates = { + "rollout_id": rollout_id, + "native_session_id": native_session_id + if native_session_id != rollout_id + else None, + "native_trajectory_id": native_trajectory_id + if native_trajectory_id not in (None, rollout_id) + else None, + "label": sample.label + if sample is not None and sample.label is not None + else request_label, + "reward": sample.reward if sample is not None else None, + "sample_metrics": sample.metrics if sample and sample.metrics else None, + "sample_extra_fields": sample.extra_fields + if sample and sample.extra_fields + else None, + "request_metadata": request_metadata, + "request_extra_fields": request_extra_fields, + "result_extra_fields": result.extra_fields, + "unmatched_llm_call_metrics": unmatched_llm_call_metrics, + "unmatched_sample_reports": { + key: value.model_dump(exclude_none=True) + for key, value in unmatched_reports.items() + } + or None, + } + osmosis.update({key: value for key, value in updates.items() if value is not None}) + extra["osmosis"] = osmosis + trajectory.extra = extra + return trajectory + + async def save_trajectories( *, rollout_id: str, @@ -51,6 +134,31 @@ async def save_trajectories( artifact_root: Path | None = None, ) -> None: """Save the rollout's sample as an ATIF document. Never raises.""" + await _save_trajectories_with_status( + rollout_id=rollout_id, + result=result, + request_label=request_label, + request_metadata=request_metadata, + request_extra_fields=request_extra_fields, + report=report, + artifact_root=artifact_root, + ) + + +async def _save_trajectories_with_status( + *, + rollout_id: str, + result: ExecutionResult, + request_label: str | None = None, + request_metadata: dict[str, Any] | None = None, + request_extra_fields: dict[str, Any] | None = None, + report: TrajectoryReport | None = None, + artifact_root: Path | None = None, +) -> bool: + """Save a trajectory without raising and report success. + + Native Harbor uses the result to decide whether the source trial is safe to delete; the public save API remains best effort. + """ try: await _save( rollout_id=rollout_id, @@ -61,12 +169,14 @@ async def save_trajectories( report=report, artifact_root=artifact_root or default_artifact_root(), ) + return True except Exception: logger.warning( "Failed to save trajectories for rollout %s (best-effort)", rollout_id, exc_info=True, ) + return False async def _save( @@ -79,6 +189,45 @@ async def _save( report: TrajectoryReport | None, artifact_root: Path, ) -> None: + if result.extra_fields is not None: + diagnostics_dest = artifact_root / rollout_id / "diagnostics.json" + diagnostics_data = json.dumps( + result.extra_fields, + ensure_ascii=False, + indent=2, + sort_keys=True, + ).encode() + await asyncio.to_thread( + _write_document, + diagnostics_dest, + diagnostics_data, + ) + logger.info( + "Saved rollout diagnostics for %s -> %s", + rollout_id, + diagnostics_dest, + ) + + if result.trajectory_document is not None: + trajectory = _prepare_native_trajectory( + result.trajectory_document, + rollout_id=rollout_id, + result=result, + request_label=request_label, + request_metadata=request_metadata, + request_extra_fields=request_extra_fields, + report=report, + ) + dest = artifact_root / rollout_id / "trajectory.json" + data = format_trajectory_json(trajectory.to_json_dict()).encode() + await asyncio.to_thread(_write_document, dest, data) + logger.info( + "Saved backend-native trajectory document for rollout %s -> %s", + rollout_id, + dest, + ) + return + sample = result.sample if sample is None: return @@ -107,6 +256,7 @@ async def _save( request_label=request_label, request_metadata=request_metadata, request_extra_fields=request_extra_fields, + result_extra_fields=result.extra_fields, report=matched_report, default_model_name=report.model_name if report else None, unmatched_sample_reports=unmatched_reports or None, diff --git a/osmosis_ai/rollout/types/protocol.py b/osmosis_ai/rollout/types/protocol.py index 1a50e234..f78b7981 100644 --- a/osmosis_ai/rollout/types/protocol.py +++ b/osmosis_ai/rollout/types/protocol.py @@ -105,5 +105,6 @@ class GraderCompleteRequest(BaseModel): status: GraderStatus rollout_id: str | None = None sample: RolloutSample | None = None + extra_fields: dict[str, Any] | None = None err_message: str | None = None err_category: RolloutErrorCategory | None = None diff --git a/osmosis_ai/rollout/types/sample.py b/osmosis_ai/rollout/types/sample.py index 52957ba6..30371732 100644 --- a/osmosis_ai/rollout/types/sample.py +++ b/osmosis_ai/rollout/types/sample.py @@ -80,5 +80,9 @@ def _validate_id(cls, value: str) -> str: class ExecutionResult(BaseModel): status: RolloutStatus sample: RolloutSample | None = None + # Backend diagnostics archived separately from request/sample extras. + extra_fields: dict[str, Any] | None = None + # Process-local native ATIF, excluded from callback serialization. + trajectory_document: dict[str, Any] | None = Field(default=None, exclude=True) err_message: str | None = None err_category: RolloutErrorCategory | None = None diff --git a/tests/unit/cli/test_native_preflight.py b/tests/unit/cli/test_native_preflight.py new file mode 100644 index 00000000..eb67e588 --- /dev/null +++ b/tests/unit/cli/test_native_preflight.py @@ -0,0 +1,195 @@ +"""Tests for native Harbor submit preflight. + +Native entrypoints expose a module-level ASGI app instead of a Python workflow and grader, so preflight validates the backend bound to that app. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from osmosis_ai.cli.errors import CLIError +from osmosis_ai.eval.common.cli import discover_native_backend +from osmosis_ai.platform.cli.workspace_directory_contract import ( + validate_rollout_backend, +) + +NATIVE_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.server import create_rollout_server + + +backend = NativeHarborBackend() +app = create_rollout_server(backend=backend) +""" + +NATIVE_SUBCLASS_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.server import create_rollout_server + + +class CustomNativeBackend(NativeHarborBackend): + pass + + +app = create_rollout_server(backend=CustomNativeBackend()) +""" + +MAIN_ONLY_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.server import create_rollout_server + + +def main(): + app = create_rollout_server(backend=NativeHarborBackend()) + return app +""" + +IMPORT_ONLY_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +""" + +UNREACHABLE_NATIVE_APP_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.server import create_rollout_server + + +if False: + app = create_rollout_server(backend=NativeHarborBackend()) +""" + +NON_NATIVE_APP_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.base import ExecutionBackend +from osmosis_ai.rollout.server import create_rollout_server + + +class OtherBackend(ExecutionBackend): + async def execute(self, request, on_workflow_complete, on_grader_complete=None): + return None + + +app = create_rollout_server(backend=OtherBackend()) +""" + +UNWIRED_NATIVE_ENTRYPOINT = """\ +from osmosis_ai.rollout.backend.base import ExecutionBackend +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.server import create_rollout_server + + +class OtherBackend(ExecutionBackend): + async def execute(self, request, on_workflow_complete, on_grader_complete=None): + return None + + +unused_native_backend = NativeHarborBackend() +app = create_rollout_server(backend=OtherBackend()) +""" + +EMPTY_ENTRYPOINT = "VALUE = 1\n" + + +def _make_rollout(workspace: Path, name: str, source: str) -> None: + rollout_dir = workspace / "rollouts" / name + rollout_dir.mkdir(parents=True, exist_ok=True) + (rollout_dir / "main.py").write_text(source, encoding="utf-8") + + +class TestDiscoverNativeBackend: + @pytest.mark.parametrize( + "source, expected_name", + [ + (NATIVE_ENTRYPOINT, "NativeHarborBackend"), + (NATIVE_SUBCLASS_ENTRYPOINT, "CustomNativeBackend"), + ], + ids=["native", "native-subclass"], + ) + def test_finds_backend_bound_to_module_app( + self, tmp_path: Path, source: str, expected_name: str + ) -> None: + _make_rollout(tmp_path, "native-rollout", source) + + cls = discover_native_backend( + rollout="native-rollout", + entrypoint="main.py", + workspace_directory=tmp_path, + ) + + assert cls is not None + assert cls.__name__ == expected_name + + @pytest.mark.parametrize( + "source", + [ + MAIN_ONLY_ENTRYPOINT, + IMPORT_ONLY_ENTRYPOINT, + UNREACHABLE_NATIVE_APP_ENTRYPOINT, + NON_NATIVE_APP_ENTRYPOINT, + UNWIRED_NATIVE_ENTRYPOINT, + EMPTY_ENTRYPOINT, + ], + ids=[ + "main-only", + "import-only", + "unreachable-native-app", + "non-native-app", + "unwired-native", + "empty", + ], + ) + def test_none_without_module_level_native_app( + self, tmp_path: Path, source: str + ) -> None: + _make_rollout(tmp_path, "native-rollout", source) + + assert ( + discover_native_backend( + rollout="native-rollout", + entrypoint="main.py", + workspace_directory=tmp_path, + ) + is None + ) + + def test_none_on_missing_entrypoint(self, tmp_path: Path) -> None: + _make_rollout(tmp_path, "native-rollout", NATIVE_ENTRYPOINT) + + assert ( + discover_native_backend( + rollout="native-rollout", + entrypoint="nope.py", + workspace_directory=tmp_path, + ) + is None + ) + + +class TestValidateRolloutBackendNative: + def test_native_app_passes_without_grader(self, tmp_path: Path) -> None: + _make_rollout(tmp_path, "native-rollout", NATIVE_ENTRYPOINT) + + validate_rollout_backend( + workspace_directory=tmp_path, + rollout="native-rollout", + entrypoint="main.py", + command_label="Test", + ) + + @pytest.mark.parametrize( + "source", + [IMPORT_ONLY_ENTRYPOINT, MAIN_ONLY_ENTRYPOINT, NON_NATIVE_APP_ENTRYPOINT], + ids=["import-only", "main-only", "non-native-app"], + ) + def test_non_native_contract_fails_preflight( + self, tmp_path: Path, source: str + ) -> None: + _make_rollout(tmp_path, "native-rollout", source) + + with pytest.raises(CLIError, match="preflight failed"): + validate_rollout_backend( + workspace_directory=tmp_path, + rollout="native-rollout", + entrypoint="main.py", + command_label="Test", + ) diff --git a/tests/unit/rollout/test_backend_concurrency.py b/tests/unit/rollout/test_backend_concurrency.py index a01dbd23..4db5d9c6 100644 --- a/tests/unit/rollout/test_backend_concurrency.py +++ b/tests/unit/rollout/test_backend_concurrency.py @@ -10,6 +10,7 @@ async def execute(self, request, on_workflow_complete, on_grader_complete=None): backend = StubBackend() assert backend.max_concurrency == 0 + assert backend.max_queue_depth is None def test_local_backend_max_concurrency(): diff --git a/tests/unit/rollout/test_native_harbor_backend.py b/tests/unit/rollout/test_native_harbor_backend.py new file mode 100644 index 00000000..f3ffc846 --- /dev/null +++ b/tests/unit/rollout/test_native_harbor_backend.py @@ -0,0 +1,1584 @@ +"""Unit tests for ``NativeHarborBackend``. + +Harbor queue submission is replaced with in-process trial results so configuration, callback, reward, and cleanup behavior can be tested without Docker. +""" + +import json +import logging +from datetime import UTC, datetime, timedelta +from pathlib import Path +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock + +import pytest +from harbor.models.environment_type import EnvironmentType +from harbor.models.task.config import MCPServerConfig +from harbor.models.trial.config import ( + AgentConfig, + EnvironmentConfig, + TaskConfig, + VerifierConfig, +) + +from osmosis_ai.rollout.backend.base import ExecutionBackend +from osmosis_ai.rollout.backend.native_harbor import backend as bmod +from osmosis_ai.rollout.backend.native_harbor.backend import ( + _AGENT_BINDINGS, + NativeHarborBackend, + _AgentProtocol, + resolve_task, +) +from osmosis_ai.rollout.context import RolloutContext +from osmosis_ai.rollout.types import ( + ExecutionRequest, + RolloutErrorCategory, + RolloutStatus, +) + + +def _trial_result( + rewards: dict[str, float | int] | None = None, + *, + steps: list[dict[str, float | int]] | None = None, + exc_message: str | None = None, + exc_type: str | None = None, +) -> Any: + """Build a duck-typed Harbor trial result.""" + top = SimpleNamespace(rewards=rewards) if rewards is not None else None + step_results = ( + [SimpleNamespace(verifier_result=SimpleNamespace(rewards=s)) for s in steps] + if steps is not None + else None + ) + exception_info = ( + SimpleNamespace(exception_message=exc_message, exception_type=exc_type) + if exc_message is not None + else None + ) + return SimpleNamespace( + verifier_result=top, + step_results=step_results, + exception_info=exception_info, + ) + + +def _patch_trial( + monkeypatch: pytest.MonkeyPatch, + *, + result: Any = None, + create_error: Exception | None = None, + capture: dict[str, Any] | None = None, +) -> None: + """Replace queue submission with an in-process fake. + + The fake optionally captures the generated TrialConfig and returns a duck-typed Harbor result. + """ + + async def _submit(self: Any, trial_config: Any) -> SimpleNamespace: + if capture is not None: + capture["config"] = trial_config + if create_error is not None: + raise create_error + return result if result is not None else _trial_result(rewards={"reward": 1.0}) + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + + +def _request(metadata: dict[str, Any] | None = None, **kw: Any) -> ExecutionRequest: + md = {"harbor_task": "/tmp/task"} if metadata is None else metadata + return ExecutionRequest( + id="ROLL", prompt=[{"role": "user", "content": "hi"}], metadata=md, **kw + ) + + +def _ctx() -> RolloutContext: + return RolloutContext( + chat_completions_url="http://ctrl:8080", api_key="sk-test", rollout_id="ROLL" + ) + + +def _native_trajectory(*, api_key: str = "sk-test") -> dict[str, Any]: + return { + "schema_version": "ATIF-v1.7", + "session_id": "native-session", + "trajectory_id": "native-trajectory", + "agent": { + "name": "terminus-2", + "version": "0.20.0", + "model_name": "model", + "extra": { + "llm_kwargs": { + "api_key": api_key, + "temperature": 0, + }, + "command": ["agent", "--token", api_key], + "safe": "kept", + }, + }, + "steps": [ + { + "step_id": 1, + "source": "agent", + "message": "done", + "llm_call_count": 1, + } + ], + } + + +class TestResolveTask: + @staticmethod + def _native_warnings(caplog: pytest.LogCaptureFixture) -> list[str]: + return [ + record.getMessage() + for record in caplog.records + if record.name == bmod.__name__ and record.levelno >= logging.WARNING + ] + + def test_local_path_does_not_warn(self, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger=bmod.__name__): + cfg = resolve_task(_request({"harbor_task": "/tmp/some/task"})) + assert cfg.path == Path("/tmp/some/task") + assert cfg.name is None + assert self._native_warnings(caplog) == [] + + @pytest.mark.parametrize( + "ref", + ["3", "sha256:0123456789abcdef"], + ) + def test_package_with_pinned_ref_does_not_warn( + self, ref: str, caplog: pytest.LogCaptureFixture + ): + with caplog.at_level(logging.WARNING, logger=bmod.__name__): + cfg = resolve_task(_request({"harbor_task": f"harbor/hello-world@{ref}"})) + assert cfg.name == "harbor/hello-world" + assert cfg.ref == ref + assert cfg.path is None + assert self._native_warnings(caplog) == [] + + @pytest.mark.parametrize( + "task_ref", + ["harbor/hello-world", "harbor/hello-world@latest"], + ) + def test_unpinned_package_warns_and_still_resolves_latest( + self, task_ref: str, caplog: pytest.LogCaptureFixture + ): + with caplog.at_level(logging.WARNING, logger=bmod.__name__): + cfg = resolve_task(_request({"harbor_task": task_ref})) + assert cfg.ref == "latest" + warnings = self._native_warnings(caplog) + assert len(warnings) == 1 + assert "mutable ref 'latest'" in warnings[0] + assert "sha256 digest" in warnings[0] + + def test_git_form_with_commit_does_not_warn(self, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger=bmod.__name__): + cfg = resolve_task( + _request( + { + "harbor_task": "git", + "git_url": "https://example.com/r.git", + "task_path": "tasks/foo", + "git_commit_id": "abc123", + } + ) + ) + assert cfg.git_url == "https://example.com/r.git" + assert cfg.path == Path("tasks/foo") + assert cfg.git_commit_id == "abc123" + assert self._native_warnings(caplog) == [] + + @pytest.mark.parametrize("commit", [None, "", " \t"]) + def test_unpinned_git_warns_without_logging_url( + self, commit: str | None, caplog: pytest.LogCaptureFixture + ): + git_url = "https://secret@example.com/private.git" + metadata = { + "harbor_task": "git", + "git_url": git_url, + "task_path": "tasks/foo", + } + if commit is not None: + metadata["git_commit_id"] = commit + with caplog.at_level(logging.WARNING, logger=bmod.__name__): + cfg = resolve_task(_request(metadata)) + + assert cfg.git_commit_id == commit + warnings = self._native_warnings(caplog) + assert len(warnings) == 1 + assert "unpinned git task" in warnings[0] + assert "git_commit_id" in warnings[0] + assert git_url not in warnings[0] + + def test_missing_raises(self): + with pytest.raises(ValueError, match="harbor_task"): + resolve_task(_request({})) + + def test_bare_package_name_without_org_raises(self): + with pytest.raises(ValueError, match="org/name"): + resolve_task(_request({"harbor_task": "helloworld"})) + + +class TestAgentConfig: + def test_in_process_url_and_endpoint_wiring(self): + backend = NativeHarborBackend() + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.name == "terminus-2" + assert ac.kwargs["api_base"] == "http://ctrl:8080" + assert "collect_rollout_details" not in ac.kwargs + assert ac.kwargs["llm_kwargs"]["api_key"] == "sk-test" + assert "extra_headers" not in ac.kwargs["llm_kwargs"] + assert ac.kwargs["llm_kwargs"]["extra_body"] == {"stream": False} + + def test_unknown_builtin_fails_without_validated_binding(self): + with pytest.raises(ValueError, match="no validated Native Harbor binding"): + NativeHarborBackend(agent_name="nop") + + def test_agent_kwargs_override_terminus_default(self): + backend = NativeHarborBackend( + agent_kwargs={"proactive_summarization_threshold": 4000} + ) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.kwargs["proactive_summarization_threshold"] == 4000 + assert ac.kwargs["enable_summarize"] is False + + def test_agent_kwargs_cannot_override_sdk_wiring(self): + backend = NativeHarborBackend(agent_kwargs={"api_base": "http://evil"}) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.kwargs["api_base"] == "http://ctrl:8080" + + def test_agent_kwargs_llm_kwargs_deep_merged(self): + backend = NativeHarborBackend( + agent_kwargs={"llm_kwargs": {"timeout": 30, "extra_body": {"foo": 1}}} + ) + ac = backend._build_agent_config(_request(), _ctx()) + llm = ac.kwargs["llm_kwargs"] + assert llm["timeout"] == 30 + assert llm["api_key"] == "sk-test" + assert llm["extra_body"] == {"foo": 1, "stream": False} + + def test_agent_kwargs_stream_is_binding_owned(self): + backend = NativeHarborBackend( + agent_kwargs={"llm_kwargs": {"extra_body": {"stream": True}}} + ) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.kwargs["llm_kwargs"]["extra_body"]["stream"] is False + + def test_agent_env_passthrough(self): + backend = NativeHarborBackend(agent_env={"FOO": "bar"}) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.env == {"FOO": "bar"} + + def test_agent_name_and_import_path_mutually_exclusive(self): + with pytest.raises(ValueError, match="not both"): + NativeHarborBackend( + agent_name="terminus-2", agent_import_path="my.pkg:MyAgent" + ) + + @pytest.mark.parametrize( + "agent_name", + ["codex", "opencode", "claude-code"], + ) + def test_agents_the_trainer_cannot_use_are_not_registered(self, agent_name: str): + with pytest.raises(ValueError, match="no validated Native Harbor binding"): + NativeHarborBackend(agent_name=agent_name) + + def test_registered_bindings_are_exactly_the_training_parity_set(self): + assert sorted(_AGENT_BINDINGS) == [ + "custom-chat-completions", + "custom-installed-chat-completions", + "oracle", + "terminus-2", + ] + + def test_every_registered_binding_speaks_a_reachable_protocol(self): + for binding in _AGENT_BINDINGS.values(): + assert binding.protocol in { + _AgentProtocol.CHAT_COMPLETIONS, + _AgentProtocol.NONE, + } + + def test_every_model_driving_binding_is_training_supported(self): + for binding in _AGENT_BINDINGS.values(): + if binding.emits_model_traffic: + assert binding.training_supported, binding.name + + def test_custom_agent_is_wired_but_not_injected(self): + with pytest.warns(UserWarning, match="custom Chat Completions"): + backend = NativeHarborBackend( + agent_import_path="my.custom.pkg:CustomAgent", + binding="custom-chat-completions", + ) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.import_path == "my.custom.pkg:CustomAgent" + assert ac.name is None + assert ac.kwargs["api_base"] == "http://ctrl:8080" + assert ac.kwargs["llm_kwargs"]["api_key"] == "sk-test" + assert "enable_summarize" not in ac.kwargs + + def test_custom_installed_agent_is_wired_through_env(self): + with pytest.warns(UserWarning, match="custom installed Chat Completions"): + backend = NativeHarborBackend( + agent_import_path="my.custom.pkg:CustomInstalledAgent", + binding="custom-installed-chat-completions", + ) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.import_path == "my.custom.pkg:CustomInstalledAgent" + assert ac.env["OPENAI_BASE_URL"] == "http://ctrl:8080" + assert ac.env["OPENAI_API_KEY"] == "sk-test" + assert ac.kwargs == {} + + def test_custom_installed_agent_passes_other_provider_credentials_through(self): + with pytest.warns(UserWarning, match="custom installed Chat Completions"): + backend = NativeHarborBackend( + agent_import_path="my.custom.pkg:CustomInstalledAgent", + binding="custom-installed-chat-completions", + agent_env={ + "ANTHROPIC_API_KEY": "user-anthropic", + "GEMINI_API_KEY": "user-gemini", + "EMAIL_SUBAGENT_MODEL": "anthropic:claude", + }, + ) + ac = backend._build_agent_config(_request(), _ctx()) + assert ac.env["ANTHROPIC_API_KEY"] == "user-anthropic" + assert ac.env["GEMINI_API_KEY"] == "user-gemini" + assert ac.env["EMAIL_SUBAGENT_MODEL"] == "anthropic:claude" + assert ac.env["OPENAI_BASE_URL"] == "http://ctrl:8080" + + @pytest.mark.parametrize("identity_key", ["OPENAI_BASE_URL", "OPENAI_API_KEY"]) + def test_custom_installed_agent_rejects_owned_identity_env(self, identity_key: str): + with pytest.raises(ValueError, match=r"owns agent\.env identity keys"): + NativeHarborBackend( + agent_import_path="my.custom.pkg:CustomInstalledAgent", + binding="custom-installed-chat-completions", + agent_env={identity_key: "user-owned"}, + ) + + def test_custom_agent_requires_an_explicit_custom_binding(self): + with pytest.raises(ValueError, match="requires an explicit custom binding"): + NativeHarborBackend(agent_import_path="my.custom.pkg:CustomAgent") + + def test_custom_agent_rejects_a_builtin_binding_name(self): + with pytest.raises(ValueError, match="only supports the custom bindings"): + NativeHarborBackend( + agent_import_path="my.custom.pkg:CustomAgent", + binding="terminus-2", + ) + + @pytest.mark.parametrize( + ("import_path", "agent_name"), + [ + ("harbor.agents.installed.opencode:OpenCode", "opencode"), + ("harbor.agents.installed.codex:Codex", "codex"), + ("harbor.agents.installed.claude_code:ClaudeCode", "claude-code"), + ], + ) + def test_import_path_cannot_reintroduce_an_unregistered_builtin( + self, import_path: str, agent_name: str + ): + with pytest.raises(ValueError, match=rf"built-in '{agent_name}'.*no Native"): + NativeHarborBackend( + agent_import_path=import_path, + binding="custom-chat-completions", + ) + + def test_import_path_cannot_bypass_a_registered_builtin_binding(self): + with pytest.raises(ValueError, match=r"built-in 'oracle'.*agent_name"): + NativeHarborBackend( + agent_import_path="harbor.agents.oracle:OracleAgent", + binding="custom-chat-completions", + ) + + def test_oracle_is_admitted_as_a_non_model_binding(self): + with pytest.warns(UserWarning, match="oracle.*not training-safe"): + backend = NativeHarborBackend(agent_name="oracle") + + ac = backend._build_agent_config( + _request(), RolloutContext(chat_completions_url="", api_key=None) + ) + + assert ac.name == "oracle" + assert ac.kwargs == {} + assert ac.env == {} + assert backend.health()["training_supported"] is False + + def test_resolved_agent_identity_is_read_only(self): + backend = NativeHarborBackend() + + with pytest.raises(AttributeError): + backend.agent_name = "claude-code" # type: ignore[misc] + with pytest.raises(AttributeError): + backend.binding = "claude-code" # type: ignore[misc] + + def test_chat_binding_rejects_non_openai_model_override(self): + backend = NativeHarborBackend() + request = _request( + {"harbor_task": "/tmp/task", "harbor_model": "anthropic/claude"} + ) + + with pytest.raises(ValueError, match=r"requires a model prefixed.*openai"): + backend._build_agent_config(request, _ctx()) + + def test_missing_endpoint_raises(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("OSMOSIS_CHAT_COMPLETIONS_URL", raising=False) + ctx = RolloutContext(chat_completions_url="", api_key="sk-test") + backend = NativeHarborBackend() + with pytest.raises(ValueError, match="no chat_completions_url"): + backend._build_agent_config(_request(), ctx) + + def test_metadata_overrides_model(self): + backend = NativeHarborBackend() + md = {"harbor_task": "/tmp/task", "harbor_model": "openai/custom"} + ac = backend._build_agent_config(_request(md), _ctx()) + assert ac.name == "terminus-2" + assert ac.model_name == "openai/custom" + + @pytest.mark.parametrize("invalid_model", [None, "", " ", 0, False]) + def test_explicit_invalid_metadata_model_is_rejected(self, invalid_model: Any): + backend = NativeHarborBackend() + request = _request({"harbor_task": "/tmp/task", "harbor_model": invalid_model}) + + with pytest.raises(ValueError, match=r"harbor_model.*non-empty string"): + backend._build_agent_config(request, _ctx()) + + def test_agent_timeout_forwarded(self): + backend = NativeHarborBackend() + ac = backend._build_agent_config(_request(agent_timeout_sec=42.0), _ctx()) + assert ac.override_timeout_sec == 42.0 + + def test_agent_setup_timeout_stored_and_forwarded(self): + backend = NativeHarborBackend(agent_setup_timeout_sec=180.5) + + ac = backend._build_agent_config(_request(), _ctx()) + + assert backend.agent_setup_timeout_sec == 180.5 + assert ac.override_setup_timeout_sec == 180.5 + + def test_agent_setup_timeout_is_separate_from_run_timeout(self): + backend = NativeHarborBackend(agent_setup_timeout_sec=180.0) + + ac = backend._build_agent_config( + _request(agent_timeout_sec=42.0), + _ctx(), + ) + + assert ac.override_setup_timeout_sec == 180.0 + assert ac.override_timeout_sec == 42.0 + + @pytest.mark.parametrize( + "invalid_timeout", + [0.0, -0.1, float("inf"), float("nan")], + ) + def test_agent_setup_timeout_must_be_positive_and_finite( + self, invalid_timeout: float + ): + with pytest.raises(ValueError, match="agent_setup_timeout_sec must be > 0"): + NativeHarborBackend(agent_setup_timeout_sec=invalid_timeout) + + +class TestFullConfigConstructor: + def test_full_configs_are_preserved_and_deep_cloned_per_rollout(self): + agent = AgentConfig( + name="terminus-2", + model_name="openai/config-model", + skills=["org/skill@1"], + mcp_servers=[ + MCPServerConfig( + name="shell-tools", + transport="stdio", + command="serve-tools", + args=["--safe"], + ) + ], + include_logs=["*.json"], + exclude_logs=["debug-*"], + extra_allowed_hosts=["models.example.com"], + override_setup_timeout_sec=11.0, + override_timeout_sec=12.0, + max_timeout_sec=13.0, + kwargs={"nested": {"items": ["original"]}}, + env={"CUSTOM_TOKEN": "agent-secret"}, + ) + environment = EnvironmentConfig( + type=EnvironmentType.DAYTONA, + force_build=True, + delete=False, + override_cpus=4, + env={"SANDBOX_TOKEN": "environment-secret"}, + kwargs={"nested": {"region": "us-west"}}, + extra_allowed_hosts=["packages.example.com"], + ) + verifier = VerifierConfig( + override_timeout_sec=21.0, + max_timeout_sec=22.0, + include_logs=["reward.json"], + exclude_logs=["verbose.log"], + env={"VERIFIER_TOKEN": "verifier-secret"}, + import_path="my.verifier:Verifier", + kwargs={"nested": {"threshold": 0.75}}, + disable=True, + ) + backend = NativeHarborBackend( + agent=agent, + environment=environment, + verifier=verifier, + ) + + agent_one = backend._build_agent_config(_request(), _ctx()) + agent_two = backend._build_agent_config(_request(), _ctx()) + trial_one = backend._build_trial_config( + _request(), TaskConfig(path=Path("/tmp/task")), agent_one, "trial-one" + ) + trial_two = backend._build_trial_config( + _request(), TaskConfig(path=Path("/tmp/task")), agent_two, "trial-two" + ) + + assert agent_one.model_name == "openai/config-model" + assert agent_one.skills == ["org/skill@1"] + assert agent_one.mcp_servers[0].command == "serve-tools" + assert agent_one.include_logs == ["*.json"] + assert agent_one.exclude_logs == ["debug-*"] + assert agent_one.extra_allowed_hosts == ["models.example.com"] + assert agent_one.override_setup_timeout_sec == 11.0 + assert agent_one.override_timeout_sec == 12.0 + assert agent_one.max_timeout_sec == 13.0 + assert agent_one.env["CUSTOM_TOKEN"] == "agent-secret" + assert trial_one.environment.type == EnvironmentType.DAYTONA + assert trial_one.environment.force_build is True + assert trial_one.environment.delete is False + assert trial_one.environment.override_cpus == 4 + assert trial_one.environment.env["SANDBOX_TOKEN"] == "environment-secret" + assert trial_one.verifier.disable is False + assert trial_one.verifier.override_timeout_sec == 21.0 + assert trial_one.verifier.max_timeout_sec == 22.0 + assert trial_one.verifier.include_logs == ["reward.json"] + assert trial_one.verifier.exclude_logs == ["verbose.log"] + assert trial_one.verifier.env["VERIFIER_TOKEN"] == "verifier-secret" + assert trial_one.verifier.import_path == "my.verifier:Verifier" + + agent_one.kwargs["nested"]["items"].append("mutated") + agent_one.skills.append("mutated-skill") + agent_one.env["CUSTOM_TOKEN"] = "mutated" + trial_one.environment.kwargs["nested"]["region"] = "mutated" + trial_one.verifier.kwargs["nested"]["threshold"] = 0.0 + + assert agent.kwargs == {"nested": {"items": ["original"]}} + assert agent.skills == ["org/skill@1"] + assert agent.env["CUSTOM_TOKEN"] == "agent-secret" + assert environment.kwargs == {"nested": {"region": "us-west"}} + assert environment.env["SANDBOX_TOKEN"] == "environment-secret" + assert verifier.kwargs == {"nested": {"threshold": 0.75}} + assert verifier.disable is True + assert agent_two.kwargs["nested"]["items"] == ["original"] + assert agent_two.skills == ["org/skill@1"] + assert agent_two.env["CUSTOM_TOKEN"] == "agent-secret" + assert trial_two.environment.kwargs == {"nested": {"region": "us-west"}} + assert trial_two.verifier.kwargs == {"nested": {"threshold": 0.75}} + assert agent_one is not agent_two + assert trial_one.environment is not trial_two.environment + assert trial_one.verifier is not trial_two.verifier + + def test_model_and_timeout_ownership_overlays_preserve_safety_caps(self): + agent = AgentConfig( + name="terminus-2", + model_name="openai/agent-default", + override_setup_timeout_sec=10.0, + override_timeout_sec=20.0, + max_timeout_sec=30.0, + ) + verifier = VerifierConfig( + override_timeout_sec=40.0, + max_timeout_sec=50.0, + ) + backend = NativeHarborBackend( + agent=agent, + verifier=verifier, + model_name="openai/constructor-default", + agent_setup_timeout_sec=15.0, + ) + request = _request( + { + "harbor_task": "/tmp/task", + "harbor_model": "openai/row-model", + }, + agent_timeout_sec=25.0, + grader_timeout_sec=45.0, + ) + + rollout_agent = backend._build_agent_config(request, _ctx()) + rollout = backend._build_trial_config( + request, + TaskConfig(path=Path("/tmp/task")), + rollout_agent, + "trial-overlays", + ) + + assert rollout.agent.model_name == "openai/row-model" + assert rollout.agent.override_setup_timeout_sec == 15.0 + assert rollout.agent.override_timeout_sec == 25.0 + assert rollout.agent.max_timeout_sec == 30.0 + assert rollout.verifier.override_timeout_sec == 45.0 + assert rollout.verifier.max_timeout_sec == 50.0 + assert agent.model_name == "openai/agent-default" + assert agent.override_setup_timeout_sec == 10.0 + assert agent.override_timeout_sec == 20.0 + assert verifier.override_timeout_sec == 40.0 + + default_request = _request() + default_agent = backend._build_agent_config(default_request, _ctx()) + default_trial = backend._build_trial_config( + default_request, + TaskConfig(path=Path("/tmp/task")), + default_agent, + "trial-defaults", + ) + assert default_agent.model_name == "openai/constructor-default" + assert default_agent.override_timeout_sec == 20.0 + assert default_trial.verifier.override_timeout_sec == 40.0 + + @pytest.mark.parametrize( + ("agent", "message"), + [ + ( + AgentConfig(name="terminus-2", n_concurrent=1), + "agent.n_concurrent is unsupported", + ), + ( + AgentConfig(name="terminus-2", concurrency_group="shared"), + "agent.concurrency_group is unsupported", + ), + ( + AgentConfig(name="terminus-2", resume_trajectory=True), + "agent.resume_trajectory is unsupported", + ), + ], + ) + def test_agent_owned_fields_are_rejected( + self, agent: AgentConfig, message: str + ) -> None: + with pytest.raises(ValueError, match=message): + NativeHarborBackend(agent=agent) + + def test_canonical_and_legacy_config_inputs_cannot_be_mixed(self): + with pytest.raises(ValueError, match="agent cannot be combined"): + NativeHarborBackend( + agent=AgentConfig(name="terminus-2"), + agent_kwargs={"temperature": 0}, + ) + with pytest.raises(ValueError, match="cannot both be set"): + NativeHarborBackend( + environment=EnvironmentConfig(type=EnvironmentType.DOCKER), + environment_config=EnvironmentConfig(type=EnvironmentType.DAYTONA), + ) + + def test_managed_skypilot_placement_does_not_mutate_caller( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("HARBOR_SKYPILOT_CONTEXT", "managed-cluster") + environment = EnvironmentConfig( + type=EnvironmentType.SKYPILOT, + kwargs={"nested": {"value": "preserved"}}, + ) + backend = NativeHarborBackend(environment=environment) + agent = backend._build_agent_config(_request(), _ctx()) + trial = backend._build_trial_config( + _request(), TaskConfig(path=Path("/tmp/task")), agent, "trial-sky" + ) + + assert "context_name" not in environment.kwargs + assert trial.environment.kwargs["context_name"] == "managed-cluster" + assert trial.environment.kwargs["nested"] == {"value": "preserved"} + + def test_explicit_empty_agent_config_preserves_oracle_default(self): + agent = AgentConfig() + with pytest.warns(UserWarning, match="oracle.*not training-safe"): + backend = NativeHarborBackend(agent=agent) + + built = backend._build_agent_config( + _request(), RolloutContext(chat_completions_url="", api_key=None) + ) + + assert agent.name == "oracle" + assert built.name == "oracle" + assert backend.agent_name == "oracle" + + +class TestPrewarm: + async def test_builds_install_only_configs_without_rollout_context( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + task = TaskConfig( + path=Path("/tmp/prewarm-task"), + overwrite=True, + source="prewarm-test", + ) + agent = AgentConfig( + name="terminus-2", + model_name="openai/prewarm-model", + skills=["org/skill@sha256:" + "a" * 64], + kwargs={"nested": {"items": ["original"]}}, + env={"CUSTOM_TOKEN": "agent-secret"}, + override_setup_timeout_sec=31.0, + ) + environment = EnvironmentConfig( + type=EnvironmentType.DAYTONA, + delete=False, + kwargs={"nested": {"region": "us-west"}}, + ) + verifier = VerifierConfig( + env={"VERIFIER_TOKEN": "verifier-secret"}, + kwargs={"nested": {"threshold": 0.75}}, + ) + captured: list[Any] = [] + queues: list[Any] = [] + + async def _submit(queue: Any, trial_config: Any) -> Any: + queues.append(queue) + captured.append(trial_config) + return _trial_result() + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + backend = NativeHarborBackend( + agent=agent, + environment=environment, + verifier=verifier, + trials_dir=tmp_path, + cleanup_successful_trials=False, + ) + + await backend.prewarm([task, task]) + + assert queues == [backend._queue, backend._queue] + assert backend._queue._retry_config.max_retries == 0 + assert len(captured) == 2 + assert len({config.trial_name for config in captured}) == 2 + for config in captured: + assert config.trial_name.startswith("native-prewarm-") + assert config.install_only is True + assert config.verifier.disable is True + assert config.task is not task + assert config.task.source == "prewarm-test" + assert config.agent is not agent + assert config.environment is not environment + assert config.verifier is not verifier + assert config.agent.model_name == "openai/prewarm-model" + assert config.agent.skills == ["org/skill@sha256:" + "a" * 64] + assert config.agent.env == {"CUSTOM_TOKEN": "agent-secret"} + assert config.agent.override_setup_timeout_sec == 31.0 + assert "api_base" not in config.agent.kwargs + assert "llm_kwargs" not in config.agent.kwargs + assert config.environment.type == EnvironmentType.DAYTONA + assert config.environment.delete is False + assert config.verifier.env == {"VERIFIER_TOKEN": "verifier-secret"} + + captured[0].task.path = Path("/tmp/mutated") + captured[0].agent.kwargs["nested"]["items"].append("mutated") + captured[0].environment.kwargs["nested"]["region"] = "mutated" + captured[0].verifier.kwargs["nested"]["threshold"] = 0.0 + + assert task.path == Path("/tmp/prewarm-task") + assert agent.kwargs == {"nested": {"items": ["original"]}} + assert environment.kwargs == {"nested": {"region": "us-west"}} + assert verifier.kwargs == {"nested": {"threshold": 0.75}} + assert verifier.disable is False + assert backend._verifier_config.disable is False + assert captured[1].task.path == Path("/tmp/prewarm-task") + assert captured[1].agent.kwargs["nested"] == {"items": ["original"]} + assert captured[1].environment.kwargs == {"nested": {"region": "us-west"}} + assert captured[1].verifier.kwargs == {"nested": {"threshold": 0.75}} + + async def test_prewarm_installs_agent_without_per_rollout_identity( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + captured: dict[str, Any] = {} + + async def _submit(_queue: Any, trial_config: Any) -> Any: + captured["config"] = trial_config + return _trial_result() + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + backend = NativeHarborBackend(cleanup_successful_trials=False) + + await backend.prewarm([TaskConfig(path=Path("/tmp/task"))]) + + config = captured["config"] + assert config.agent.kwargs["enable_summarize"] is False + assert "api_base" not in config.agent.kwargs + assert "llm_kwargs" not in config.agent.kwargs + assert config.agent.env == {} + + async def test_attempts_all_tasks_aggregates_failures_and_cleans_only_successes( + self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path + ) -> None: + tasks = [ + TaskConfig(path=Path("/tmp/success")), + TaskConfig(path=Path("/tmp/raised")), + TaskConfig(path=Path("/tmp/reported")), + ] + attempted: list[str] = [] + configs: dict[str, Any] = {} + + async def _submit(_queue: Any, trial_config: Any) -> Any: + task_name = trial_config.task.path.name + attempted.append(task_name) + configs[task_name] = trial_config + (tmp_path / trial_config.trial_name).mkdir() + if task_name == "raised": + raise OSError("image builder unavailable") + if task_name == "reported": + return _trial_result( + exc_message="agent install failed", + exc_type="AgentSetupError", + ) + return _trial_result() + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + backend = NativeHarborBackend(trials_dir=tmp_path) + + with pytest.raises(RuntimeError) as exc_info: + await backend.prewarm(tasks) + + assert sorted(attempted) == ["raised", "reported", "success"] + message = str(exc_info.value) + assert "failed for 2 of 3 task(s)" in message + assert "/tmp/raised [OSError]; inspect preserved trial" in message + assert "/tmp/reported [AgentSetupError]; inspect preserved trial" in message + assert "image builder unavailable" not in message + assert "agent install failed" not in message + assert not (tmp_path / configs["success"].trial_name).exists() + assert (tmp_path / configs["raised"].trial_name).is_dir() + assert (tmp_path / configs["reported"].trial_name).is_dir() + + async def test_rejects_empty_task_list(self) -> None: + backend = NativeHarborBackend() + + with pytest.raises(ValueError, match="at least one Harbor TaskConfig"): + await backend.prewarm([]) + with pytest.raises(ValueError, match="at least one Harbor TaskConfig"): + backend.prewarm_lifespan([]) + + async def test_failure_aggregate_omits_raw_setup_output_and_credentials( + self, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + ) -> None: + raised_url = "https://raised:secret@example.com/private.git" + reported_url = "https://reported:secret@example.com/private.git" + tasks = [ + TaskConfig( + git_url=raised_url, + path=Path("tasks/raised"), + git_commit_id="abc123", + ), + TaskConfig( + git_url=reported_url, + path=Path("tasks/reported"), + git_commit_id="def456", + ), + ] + + agent_env_secret = "agent-env-secret" + agent_kwarg_secret = "agent-kwarg-secret" + + async def _submit(_queue: Any, trial_config: Any) -> Any: + if trial_config.task.path.name == "raised": + raise OSError(f"clone of {raised_url} failed with {agent_env_secret}") + return _trial_result( + exc_message=( + f"download of {reported_url} failed with {agent_kwarg_secret}" + ), + exc_type="TaskDownloadError", + ) + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + backend = NativeHarborBackend( + agent_env={"CUSTOM_TOKEN": agent_env_secret}, + agent_kwargs={"llm_kwargs": {"api_key": agent_kwarg_secret}}, + ) + + with caplog.at_level(logging.ERROR): + with pytest.raises(RuntimeError) as exc_info: + await backend.prewarm(tasks) + + message = str(exc_info.value) + assert "git:tasks/raised@abc123 [OSError]" in message + assert "git:tasks/reported@def456 [TaskDownloadError]" in message + assert message.count("no trial directory was created") == 2 + for secret in ( + "raised:secret", + "reported:secret", + agent_env_secret, + agent_kwarg_secret, + ): + assert secret not in message + assert secret not in caplog.text + + async def test_lifespan_clones_tasks_and_awaits_prewarm_before_serving( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from osmosis_ai.rollout.server import create_rollout_server + + backend = NativeHarborBackend() + task = TaskConfig(path=Path("/tmp/original")) + events: list[str] = [] + received: list[TaskConfig] = [] + + async def _prewarm(tasks: Any) -> None: + events.append("prewarm") + received.extend(tasks) + + monkeypatch.setattr(backend, "prewarm", _prewarm) + app = create_rollout_server( + backend=backend, + lifespan=backend.prewarm_lifespan([task]), + ) + task.path = Path("/tmp/mutated-after-app-creation") + + assert events == [] + async with app.router.lifespan_context(app): + assert events == ["prewarm"] + assert received[0].path == Path("/tmp/original") + events.append("serving") + assert events == ["prewarm", "serving"] + + async def test_lifespan_failure_aborts_startup( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from osmosis_ai.rollout.server import create_rollout_server + + backend = NativeHarborBackend() + + async def _prewarm(_tasks: Any) -> None: + raise RuntimeError("prewarm failed") + + monkeypatch.setattr(backend, "prewarm", _prewarm) + app = create_rollout_server( + backend=backend, + lifespan=backend.prewarm_lifespan([TaskConfig(path=Path("/tmp/task"))]), + ) + entered = False + + with pytest.raises(RuntimeError, match="prewarm failed"): + async with app.router.lifespan_context(app): + entered = True + assert entered is False + + +class TestRewardPicking: + def test_named_channel(self): + assert NativeHarborBackend()._pick_reward({"reward": 1}) == 1 + + def test_sole_value_fallback(self): + assert NativeHarborBackend()._pick_reward({"accuracy": 0.8}) == 0.8 + + def test_ambiguous_returns_none(self): + backend = NativeHarborBackend() + assert backend._pick_reward({"a": 1, "b": 2}) is None + + def test_custom_reward_key(self): + backend = NativeHarborBackend(reward_key="score") + assert backend._pick_reward({"score": 0.3, "reward": 0.9}) == 0.3 + + def test_extract_top_level(self): + assert NativeHarborBackend()._extract_rewards( + _trial_result(rewards={"reward": 1.0}) + ) == {"reward": 1.0} + + def test_extract_multi_step_fallback(self): + assert NativeHarborBackend()._extract_rewards( + _trial_result(steps=[{"reward": 0.5}]) + ) == {"reward": 0.5} + + +class TestExecute: + def test_is_execution_backend(self): + assert isinstance(NativeHarborBackend(), ExecutionBackend) + + async def test_success_single_sample_rewarded(self, monkeypatch): + capture: dict[str, Any] = {} + _patch_trial( + monkeypatch, result=_trial_result(rewards={"reward": 1.0}), capture=capture + ) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + + on_wf.assert_awaited_once() + on_gr.assert_awaited_once() + wf_result = on_wf.call_args.args[0] + gr_result = on_gr.call_args.args[0] + assert wf_result.status == RolloutStatus.SUCCESS + assert gr_result.status == RolloutStatus.SUCCESS + + assert gr_result.sample.reward == 1.0 + assert "extra_headers" not in capture["config"].agent.kwargs["llm_kwargs"] + assert capture["config"].verifier.disable is False + + async def test_int_reward_coerced_to_float(self, monkeypatch): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1})) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + sample = on_gr.call_args.args[0].sample + assert sample.reward == 1.0 + assert isinstance(sample.reward, float) + + async def test_agent_failure_fires_both_callbacks(self, monkeypatch): + _patch_trial(monkeypatch, result=_trial_result(exc_message="boom")) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + wf_result = on_wf.call_args.args[0] + gr_result = on_gr.call_args.args[0] + assert wf_result.status == RolloutStatus.FAILURE + assert wf_result.err_message == "boom" + assert wf_result.err_category == RolloutErrorCategory.AGENT_ERROR + assert gr_result.status == RolloutStatus.FAILURE + assert gr_result.sample.reward is None + + async def test_trial_create_raises(self, monkeypatch): + _patch_trial(monkeypatch, create_error=RuntimeError("docker down")) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + assert on_wf.call_args.args[0].status == RolloutStatus.FAILURE + assert "docker down" in on_wf.call_args.args[0].err_message + assert on_wf.call_args.args[0].extra_fields["phase"] == "setup" + assert ( + on_wf.call_args.args[0].extra_fields["harbor_exception_type"] + == "RuntimeError" + ) + assert on_gr.call_args.args[0].status == RolloutStatus.FAILURE + + @pytest.mark.parametrize("queue_fails", [False, True]) + async def test_controller_identity_is_wired_per_rollout( + self, + monkeypatch: pytest.MonkeyPatch, + queue_fails: bool, + ) -> None: + backend = NativeHarborBackend() + captured: dict[str, Any] = {} + + async def _submit(queue: Any, trial_config: Any) -> Any: + captured["agent"] = trial_config.agent + if queue_fails: + raise RuntimeError("trial failed") + return _trial_result(rewards={"reward": 1.0}) + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + + agent = captured["agent"] + assert agent.kwargs["api_base"] == "http://ctrl:8080" + assert agent.kwargs["llm_kwargs"]["api_key"] == "sk-test" + + async def test_missing_task_is_validation_error(self, monkeypatch): + _patch_trial(monkeypatch) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request({}), on_wf, on_gr) + assert ( + on_wf.call_args.args[0].err_category + == RolloutErrorCategory.VALIDATION_ERROR + ) + + async def test_no_grader_callback_only_workflow(self, monkeypatch): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend() + on_wf = AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, None) + on_wf.assert_awaited_once() + + async def test_empty_rewards_is_validation_failure(self, monkeypatch): + _patch_trial(monkeypatch, result=_trial_result(rewards={})) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + assert on_wf.call_args.args[0].status == RolloutStatus.SUCCESS + gr_result = on_gr.call_args.args[0] + assert gr_result.status == RolloutStatus.FAILURE + assert gr_result.err_category == RolloutErrorCategory.VALIDATION_ERROR + + async def test_reward_does_not_revive_failed_trial(self, monkeypatch): + _patch_trial( + monkeypatch, + result=_trial_result( + rewards={"reward": 0.7}, exc_message="post-verify upload failed" + ), + ) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + workflow_result = on_wf.call_args.args[0] + gr_result = on_gr.call_args.args[0] + assert workflow_result.status == RolloutStatus.FAILURE + assert workflow_result.sample.reward is None + assert gr_result.status == RolloutStatus.FAILURE + assert gr_result.sample.reward is None + + async def test_swallowed_exception_is_agent_error(self, monkeypatch): + # Harbor reports in-trial failures through exception_info. + _patch_trial( + monkeypatch, + result=_trial_result( + exc_message="verifier timed out", exc_type="VerifierTimeoutError" + ), + ) + backend = NativeHarborBackend() + on_wf, on_gr = AsyncMock(), AsyncMock() + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + assert on_wf.call_args.args[0].err_category == RolloutErrorCategory.AGENT_ERROR + + async def test_hook_phase_timings_and_error_payload_are_reused( + self, monkeypatch, caplog + ): + ticks = iter(float(value) for value in range(20)) + monkeypatch.setattr(bmod, "monotonic", lambda: next(ticks)) + started_at = datetime(2026, 1, 1, tzinfo=UTC) + result = _trial_result( + rewards={"reward": 0.7}, + exc_message="agent timed out", + exc_type="AgentTimeoutError", + ) + result.started_at = started_at + result.finished_at = started_at + timedelta(seconds=10) + result.environment_setup = SimpleNamespace( + started_at=started_at, + finished_at=started_at + timedelta(seconds=2), + ) + result.agent_setup = SimpleNamespace( + started_at=started_at + timedelta(seconds=2), + finished_at=started_at + timedelta(seconds=3), + ) + result.agent_execution = SimpleNamespace( + started_at=started_at + timedelta(seconds=3), + finished_at=started_at + timedelta(seconds=6), + ) + result.verifier = SimpleNamespace( + started_at=started_at + timedelta(seconds=6), + finished_at=started_at + timedelta(seconds=10), + ) + + async def _submit(queue: Any, trial_config: Any) -> Any: + for event_name in ( + "start", + "environment-start", + "agent-start", + "agent-end", + "verification-start", + "end", + ): + event = next( + event for event in queue._hooks if event.value == event_name + ) + hook_event = SimpleNamespace( + trial_name=trial_config.trial_name, + result=result, + ) + for callback in queue._hooks[event]: + await callback(hook_event) + return result + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + caplog.set_level( + logging.ERROR, + logger="osmosis_ai.rollout.backend.native_harbor.backend", + ) + on_wf, on_gr = AsyncMock(), AsyncMock() + + with _ctx(): + await NativeHarborBackend().execute(_request(), on_wf, on_gr) + + workflow_payload = on_wf.call_args.args[0].extra_fields + grader_payload = on_gr.call_args.args[0].extra_fields + assert grader_payload == workflow_payload + assert workflow_payload == { + "backend": "native_harbor", + "phase": "agent", + "harbor_exception_type": "AgentTimeoutError", + "category": "agent_error", + "timings_sec": { + "setup": 1.0, + "trial_setup": 1.0, + "environment_setup": 2.0, + "agent_setup": 1.0, + "agent": 3.0, + "verification": 4.0, + "trial": 10.0, + }, + } + assert json.dumps(workflow_payload, sort_keys=True) in caplog.text + + async def test_late_verifier_failure_keeps_callback_timing_and_phase( + self, monkeypatch + ): + result = _trial_result(rewards={"reward": 0.7}) + + async def _submit(queue: Any, trial_config: Any) -> Any: + hook_event = SimpleNamespace( + trial_name=trial_config.trial_name, + result=result, + ) + verification_event = next( + event for event in queue._hooks if event.value == "verification-start" + ) + for callback in queue._hooks[verification_event]: + await callback(hook_event) + + result.exception_info = SimpleNamespace( + exception_message="verifier timed out", + exception_type="VerifierTimeoutError", + ) + end_event = next(event for event in queue._hooks if event.value == "end") + for callback in queue._hooks[end_event]: + await callback(hook_event) + return result + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + on_wf, on_gr = AsyncMock(), AsyncMock() + + with _ctx(): + await NativeHarborBackend().execute(_request(), on_wf, on_gr) + + workflow_result = on_wf.call_args.args[0] + grader_result = on_gr.call_args.args[0] + assert workflow_result.status == RolloutStatus.SUCCESS + assert workflow_result.extra_fields["harbor_exception_type"] is None + assert grader_result.status == RolloutStatus.FAILURE + assert grader_result.extra_fields["phase"] == "verification" + assert ( + grader_result.extra_fields["harbor_exception_type"] + == "VerifierTimeoutError" + ) + + async def test_late_failure_without_grader_callback_is_archived( + self, monkeypatch, tmp_path + ): + result = _trial_result(rewards={"reward": 0.7}) + + async def _submit(queue: Any, trial_config: Any) -> Any: + hook_event = SimpleNamespace( + trial_name=trial_config.trial_name, + result=result, + ) + verification_event = next( + event for event in queue._hooks if event.value == "verification-start" + ) + for callback in queue._hooks[verification_event]: + await callback(hook_event) + result.exception_info = SimpleNamespace( + exception_message="verifier timed out", + exception_type="VerifierTimeoutError", + ) + end_event = next(event for event in queue._hooks if event.value == "end") + for callback in queue._hooks[end_event]: + await callback(hook_event) + return result + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + backend = NativeHarborBackend() + backend.artifact_root = tmp_path + on_wf = AsyncMock() + + with _ctx(): + await backend.execute(_request(), on_wf, None) + + assert on_wf.call_args.args[0].status == RolloutStatus.SUCCESS + diagnostics = json.loads((tmp_path / "ROLL" / "diagnostics.json").read_text()) + assert diagnostics["phase"] == "verification" + assert diagnostics["harbor_exception_type"] == "VerifierTimeoutError" + assert diagnostics["category"] == "agent_error" + + async def test_post_callback_queue_exception_is_archived_and_graded( + self, monkeypatch, tmp_path + ): + result = _trial_result(rewards={"reward": 0.7}) + + async def _submit(queue: Any, trial_config: Any) -> Any: + hook_event = SimpleNamespace( + trial_name=trial_config.trial_name, + result=result, + ) + verification_event = next( + event for event in queue._hooks if event.value == "verification-start" + ) + for callback in queue._hooks[verification_event]: + await callback(hook_event) + raise RuntimeError("failed to finalize trial result") + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + backend = NativeHarborBackend() + backend.artifact_root = tmp_path + on_wf, on_gr = AsyncMock(), AsyncMock() + + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + + assert on_wf.call_args.args[0].status == RolloutStatus.SUCCESS + grader_result = on_gr.call_args.args[0] + assert grader_result.status == RolloutStatus.FAILURE + assert grader_result.extra_fields["phase"] == "verification" + assert grader_result.extra_fields["harbor_exception_type"] == "RuntimeError" + diagnostics = json.loads((tmp_path / "ROLL" / "diagnostics.json").read_text()) + assert diagnostics == grader_result.extra_fields + + async def test_grader_callback_failure_propagates_after_trial(self, monkeypatch): + # The server owns the final notification fallback. + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend() + on_wf = AsyncMock() + on_gr = AsyncMock(side_effect=RuntimeError("controller down")) + with _ctx(): + with pytest.raises(RuntimeError, match="controller down"): + await backend.execute(_request(), on_wf, on_gr) + on_wf.assert_awaited_once() + on_gr.assert_awaited_once() + + async def test_failed_verification_hook_callback_retries_after_trial( + self, monkeypatch + ): + result = _trial_result(rewards={"reward": 1.0}) + + async def _submit(queue: Any, trial_config: Any) -> Any: + hook = next( + callback + for event, callbacks in queue._hooks.items() + if event.value == "verification-start" + for callback in callbacks + ) + await hook( + SimpleNamespace(trial_name=trial_config.trial_name, result=result) + ) + return result + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + on_wf = AsyncMock(side_effect=[RuntimeError("temporary outage"), None]) + + with _ctx(): + await NativeHarborBackend().execute(_request(), on_wf) + + assert on_wf.await_count == 2 + + async def test_single_step_workflow_callback_precedes_verifier_completion( + self, monkeypatch + ): + events: list[str] = [] + result = _trial_result(rewards={"reward": 1.0}) + + async def _submit(queue: Any, trial_config: Any) -> Any: + hook = next( + callback + for event, callbacks in queue._hooks.items() + if event.value == "verification-start" + for callback in callbacks + ) + await hook( + SimpleNamespace(trial_name=trial_config.trial_name, result=result) + ) + events.append("verifier-finished") + return result + + async def _workflow_callback(result: Any) -> None: + events.append("workflow") + + async def _grader_callback(result: Any) -> None: + events.append("grader") + + monkeypatch.setattr(bmod.TrialQueue, "submit", _submit) + with _ctx(): + await NativeHarborBackend().execute( + _request(), _workflow_callback, _grader_callback + ) + + assert events == ["workflow", "verifier-finished", "grader"] + + +class TestConcurrencyAndLifecycle: + def test_retry_config_is_not_a_constructor_argument(self): + with pytest.raises(TypeError, match="retry_config"): + NativeHarborBackend(retry_config=bmod.RetryConfig(max_retries=1)) # type: ignore[call-arg] + + def test_trial_queue_retries_are_hard_disabled(self): + backend = NativeHarborBackend() + assert backend._queue._retry_config.max_retries == 0 + + def test_unbounded_concurrency_rejected(self): + with pytest.raises(ValueError, match="max_concurrent must be >= 1"): + NativeHarborBackend(max_concurrent=0) + + def test_negative_queue_depth_rejected(self): + with pytest.raises(ValueError, match="max_queue_depth must be >= 0"): + NativeHarborBackend(max_queue_depth=-1) + + def test_health_reports_capacity_and_binding_capabilities(self): + backend = NativeHarborBackend(max_concurrent=3) + assert backend.max_concurrency == 3 + assert backend.max_queue_depth == 3 + assert backend.health() == { + "status": "ok", + "backend": "native_harbor", + "agent": "terminus-2", + "binding": "terminus-2", + "binding_protocol": "OpenAI Chat Completions", + "protocol_capabilities": ["OpenAI Chat Completions"], + "training_supported": True, + "max_concurrency": 3, + "max_queue_depth": 3, + } + + async def test_successful_trial_dir_cleaned_up(self, monkeypatch, tmp_path): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend(trials_dir=tmp_path) + trial_dir = tmp_path / "native-ROLL" + trial_dir.mkdir() + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + assert not trial_dir.exists() + + async def test_failed_trial_dir_kept(self, monkeypatch, tmp_path): + _patch_trial(monkeypatch, result=_trial_result(exc_message="boom")) + backend = NativeHarborBackend(trials_dir=tmp_path) + trial_dir = tmp_path / "native-ROLL" + trial_dir.mkdir() + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + assert trial_dir.exists() + + async def test_native_atif_preserved_with_agent_extra_redacted( + self, monkeypatch, tmp_path + ): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend(trials_dir=tmp_path) + backend.artifact_root = tmp_path / "saved" + trajectory_path = tmp_path / "native-ROLL" / "agent" / "trajectory.json" + trajectory_path.parent.mkdir(parents=True) + trajectory_path.write_text(json.dumps(_native_trajectory())) + on_wf, on_gr = AsyncMock(), AsyncMock() + + with _ctx(): + await backend.execute(_request(), on_wf, on_gr) + + workflow_document = on_wf.call_args.args[0].trajectory_document + grader_document = on_gr.call_args.args[0].trajectory_document + assert grader_document == workflow_document + assert workflow_document["steps"] == _native_trajectory()["steps"] + extra = workflow_document["agent"]["extra"] + assert extra["llm_kwargs"] == { + "api_key": "[REDACTED]", + "temperature": 0, + } + assert extra["command"] == ["agent", "--token", "[REDACTED]"] + assert extra["safe"] == "kept" + saved_path = backend.artifact_root / "ROLL" / "trajectory.json" + assert saved_path.is_file() + saved = json.loads(saved_path.read_text()) + diagnostics = saved["extra"]["osmosis"]["result_extra_fields"] + assert diagnostics["backend"] == "native_harbor" + assert diagnostics["phase"] == "setup" + assert diagnostics["harbor_exception_type"] is None + assert diagnostics["category"] is None + assert not trajectory_path.exists() + + async def test_native_atif_persistence_failure_preserves_successful_trial( + self, monkeypatch, tmp_path + ): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend(trials_dir=tmp_path) + trajectory_path = tmp_path / "native-ROLL" / "agent" / "trajectory.json" + trajectory_path.parent.mkdir(parents=True) + trajectory_path.write_text(json.dumps(_native_trajectory())) + persist = AsyncMock(return_value=False) + monkeypatch.setattr(bmod, "_save_trajectories_with_status", persist) + + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + + assert trajectory_path.exists() + persist.assert_awaited_once() + + async def test_invalid_native_atif_preserves_successful_trial( + self, monkeypatch, tmp_path + ): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend(trials_dir=tmp_path) + trajectory_path = tmp_path / "native-ROLL" / "agent" / "trajectory.json" + trajectory_path.parent.mkdir(parents=True) + trajectory_path.write_text('{"not": "atif"}') + + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + + assert trajectory_path.exists() + + async def test_harbor_collected_artifacts_relocated_before_cleanup( + self, monkeypatch, tmp_path + ): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + trials_dir = tmp_path / "trials" + artifact_root = tmp_path / "saved" + backend = NativeHarborBackend(trials_dir=trials_dir) + backend.artifact_root = artifact_root + artifact = ( + trials_dir + / "native-ROLL" + / "artifacts" + / "logs" + / "artifacts" + / "result.txt" + ) + artifact.parent.mkdir(parents=True) + artifact.write_text("user-selected output") + + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + + relocated = ( + artifact_root / "ROLL" / "artifacts" / "logs" / "artifacts" / "result.txt" + ) + assert relocated.read_text() == "user-selected output" + assert not (trials_dir / "native-ROLL").exists() + + async def test_artifact_copy_failure_preserves_successful_trial( + self, monkeypatch, tmp_path + ): + _patch_trial(monkeypatch, result=_trial_result(rewards={"reward": 1.0})) + backend = NativeHarborBackend(trials_dir=tmp_path) + source = tmp_path / "native-ROLL" / "artifacts" / "out.txt" + source.parent.mkdir(parents=True) + source.write_text("keep me") + + def _fail_copy(*args: Any, **kwargs: Any) -> int: + raise OSError("destination unavailable") + + monkeypatch.setattr(bmod, "copy_artifact_tree", _fail_copy) + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + + assert source.exists() + + async def test_environment_config_threaded_into_trial(self, monkeypatch): + from harbor.models.environment_type import EnvironmentType + from harbor.models.trial.config import EnvironmentConfig + + capture: dict[str, Any] = {} + _patch_trial( + monkeypatch, result=_trial_result(rewards={"reward": 1.0}), capture=capture + ) + backend = NativeHarborBackend( + environment_config=EnvironmentConfig(type=EnvironmentType.DAYTONA) + ) + with _ctx(): + await backend.execute(_request(), AsyncMock(), AsyncMock()) + assert capture["config"].environment.type == EnvironmentType.DAYTONA diff --git a/tests/unit/rollout/test_server_admission.py b/tests/unit/rollout/test_server_admission.py new file mode 100644 index 00000000..51dd94f9 --- /dev/null +++ b/tests/unit/rollout/test_server_admission.py @@ -0,0 +1,252 @@ +"""Admission-control coverage for ``create_rollout_server``.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import httpx +import pytest + +from osmosis_ai.rollout.backend.base import ExecutionBackend, ResultCallback +from osmosis_ai.rollout.server import app as app_module +from osmosis_ai.rollout.server.app import create_rollout_server +from osmosis_ai.rollout.types import ExecutionRequest, RolloutInitRequest + + +class _CapacityBackend(ExecutionBackend): + def __init__( + self, + *, + max_concurrent: int = 1, + max_queue_depth: int | None = None, + ) -> None: + self._max_concurrent = max_concurrent + self._max_queue_depth = max_queue_depth + + @property + def max_concurrency(self) -> int: + return self._max_concurrent + + @property + def max_queue_depth(self) -> int | None: + return self._max_queue_depth + + def health(self) -> dict[str, Any]: + return {"status": "ok", "backend": "capacity-test"} + + async def execute( + self, + request: ExecutionRequest, + on_workflow_complete: ResultCallback, + on_grader_complete: ResultCallback | None = None, + ) -> None: # pragma: no cover - _handle_rollout is replaced in these tests + raise AssertionError("execute should not be called") + + +class _BlockingHandler: + def __init__(self, expected_entries: int) -> None: + self.expected_entries = expected_entries + self.entries = 0 + self.entered = asyncio.Event() + self.release = asyncio.Event() + + async def __call__( + self, + backend: ExecutionBackend, + request: RolloutInitRequest, + ) -> None: + del backend, request + self.entries += 1 + if self.entries >= self.expected_entries: + self.entered.set() + await self.release.wait() + + +def _payload(rollout_id: str) -> dict[str, Any]: + return { + "rollout_id": rollout_id, + "initial_messages": [{"role": "user", "content": "hi"}], + "chat_completions_url": f"http://controller/sessions/{rollout_id}/v1", + "completion_callback_url": ( + f"http://controller/v1/rollout/{rollout_id}/completed" + ), + } + + +async def _wait_for(event: asyncio.Event) -> None: + await asyncio.wait_for(event.wait(), timeout=5) + + +async def _health(client: httpx.AsyncClient) -> dict[str, Any]: + response = await client.get("/health") + assert response.status_code == 200 + return response.json() + + +class TestBoundedAdmission: + async def test_full_queue_returns_429_and_health_tracks_release( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + handler = _BlockingHandler(expected_entries=2) + monkeypatch.setattr(app_module, "_handle_rollout", handler) + app = create_rollout_server( + backend=_CapacityBackend(max_concurrent=1, max_queue_depth=1), + configure_logging=False, + ) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://rollout.test" + ) as client: + first = asyncio.create_task(client.post("/rollout", json=_payload("r1"))) + second = asyncio.create_task(client.post("/rollout", json=_payload("r2"))) + try: + await _wait_for(handler.entered) + health = await _health(client) + assert health["backend"] == "capacity-test" + assert health["capacity"] == { + "max_concurrent": 1, + "max_queue_depth": 1, + "in_flight": 2, + "queue_depth": 1, + "available": 0, + "accepting": False, + } + + rejected = await asyncio.wait_for( + client.post("/rollout", json=_payload("r3")), timeout=5 + ) + assert rejected.status_code == 429 + assert rejected.json() == {"detail": "Rollout queue is full"} + finally: + handler.release.set() + responses = await asyncio.wait_for( + asyncio.gather(first, second), timeout=5 + ) + + assert [response.status_code for response in responses] == [200, 200] + assert (await _health(client))["capacity"] == { + "max_concurrent": 1, + "max_queue_depth": 1, + "in_flight": 0, + "queue_depth": 0, + "available": 2, + "accepting": True, + } + + async def test_zero_queue_depth_rejects_while_execution_is_active( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + handler = _BlockingHandler(expected_entries=1) + monkeypatch.setattr(app_module, "_handle_rollout", handler) + app = create_rollout_server( + backend=_CapacityBackend(max_concurrent=1, max_queue_depth=0), + configure_logging=False, + ) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://rollout.test" + ) as client: + active = asyncio.create_task(client.post("/rollout", json=_payload("r1"))) + try: + await _wait_for(handler.entered) + rejected = await asyncio.wait_for( + client.post("/rollout", json=_payload("r2")), timeout=5 + ) + assert rejected.status_code == 429 + assert (await _health(client))["capacity"] == { + "max_concurrent": 1, + "max_queue_depth": 0, + "in_flight": 1, + "queue_depth": 0, + "available": 0, + "accepting": False, + } + finally: + handler.release.set() + response = await asyncio.wait_for(active, timeout=5) + assert response.status_code == 200 + + async def test_background_error_releases_reservation( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + async def fail( + backend: ExecutionBackend, + request: RolloutInitRequest, + ) -> None: + del backend, request + raise RuntimeError("background failed") + + monkeypatch.setattr(app_module, "_handle_rollout", fail) + app = create_rollout_server( + backend=_CapacityBackend(max_concurrent=1, max_queue_depth=0), + configure_logging=False, + ) + transport = httpx.ASGITransport(app=app, raise_app_exceptions=False) + async with httpx.AsyncClient( + transport=transport, base_url="http://rollout.test" + ) as client: + response = await client.post("/rollout", json=_payload("r1")) + assert response.status_code == 200 + assert (await _health(client))["capacity"]["in_flight"] == 0 + assert (await _health(client))["capacity"]["accepting"] is True + + async def test_background_cancellation_releases_reservation( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + handler = _BlockingHandler(expected_entries=1) + monkeypatch.setattr(app_module, "_handle_rollout", handler) + app = create_rollout_server( + backend=_CapacityBackend(max_concurrent=1, max_queue_depth=0), + configure_logging=False, + ) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://rollout.test" + ) as client: + active = asyncio.create_task(client.post("/rollout", json=_payload("r1"))) + await _wait_for(handler.entered) + active.cancel() + with pytest.raises(asyncio.CancelledError): + await active + + capacity = (await _health(client))["capacity"] + assert capacity["in_flight"] == 0 + assert capacity["accepting"] is True + + +class TestUnboundedAdmission: + async def test_backend_without_queue_bound_remains_unbounded( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + handler = _BlockingHandler(expected_entries=3) + monkeypatch.setattr(app_module, "_handle_rollout", handler) + app = create_rollout_server( + backend=_CapacityBackend(max_concurrent=1, max_queue_depth=None), + configure_logging=False, + ) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient( + transport=transport, base_url="http://rollout.test" + ) as client: + requests = [ + asyncio.create_task(client.post("/rollout", json=_payload(f"r{i}"))) + for i in range(3) + ] + try: + await _wait_for(handler.entered) + assert (await _health(client))["capacity"] == { + "max_concurrent": 1, + "max_queue_depth": None, + "in_flight": 3, + "queue_depth": 2, + "available": None, + "accepting": True, + } + finally: + handler.release.set() + responses = await asyncio.wait_for(asyncio.gather(*requests), timeout=5) + assert [response.status_code for response in responses] == [200, 200, 200] diff --git a/tests/unit/rollout/test_server_app_logging.py b/tests/unit/rollout/test_server_app_logging.py index f1ce8531..bc6dfeea 100644 --- a/tests/unit/rollout/test_server_app_logging.py +++ b/tests/unit/rollout/test_server_app_logging.py @@ -5,7 +5,10 @@ import logging from osmosis_ai.rollout.backend.base import ExecutionBackend, ResultCallback -from osmosis_ai.rollout.server.app import create_rollout_server +from osmosis_ai.rollout.server.app import ( + _get_rollout_server_backend, + create_rollout_server, +) from osmosis_ai.rollout.types import ExecutionRequest @@ -54,3 +57,11 @@ def test_opting_out_installs_nothing(self, monkeypatch): assert root.handlers == [] assert root.level == logging.WARNING + + +def test_create_rollout_server_records_backend_marker() -> None: + backend = _Backend() + + app = create_rollout_server(backend=backend, configure_logging=False) + + assert _get_rollout_server_backend(app) is backend diff --git a/tests/unit/rollout/test_server_app_trajectory.py b/tests/unit/rollout/test_server_app_trajectory.py index 6f70d390..2cf11d92 100644 --- a/tests/unit/rollout/test_server_app_trajectory.py +++ b/tests/unit/rollout/test_server_app_trajectory.py @@ -6,6 +6,8 @@ from typing import Any from osmosis_ai.rollout.backend.base import ExecutionBackend, ResultCallback +from osmosis_ai.rollout.backend.native_harbor import NativeHarborBackend +from osmosis_ai.rollout.backend.native_harbor import backend as native_backend_module from osmosis_ai.rollout.server.app import _handle_rollout from osmosis_ai.rollout.types import ( ExecutionRequest, @@ -119,6 +121,115 @@ async def test_records_graded_result(tmp_path: Path, monkeypatch) -> None: assert "trajectory_messages" not in grader_payload["sample"] +async def test_result_extra_fields_are_posted_and_archived( + tmp_path: Path, monkeypatch +) -> None: + posted = patch_callbacks(monkeypatch) + patch_artifact_root(monkeypatch, tmp_path) + diagnostics = { + "backend": "native_harbor", + "phase": "agent", + "harbor_exception_type": "AgentTimeoutError", + "category": "timeout", + "timings_sec": {"agent": 42.0}, + } + backend = StubBackend( + workflow_result=ExecutionResult( + status=RolloutStatus.FAILURE, + sample=make_sample(), + extra_fields=diagnostics, + ), + grader_result=ExecutionResult( + status=RolloutStatus.FAILURE, + sample=make_sample(), + extra_fields=diagnostics, + ), + ) + + await _handle_rollout(backend, make_request()) + + completion_payload = posted[0][1] + assert completion_payload["extra_fields"] == diagnostics + grader_payload = posted[1][1] + assert grader_payload["extra_fields"] == diagnostics + doc = json.loads((tmp_path / "r1" / "trajectory.json").read_text()) + assert doc["extra"]["osmosis"]["result_extra_fields"] == diagnostics + assert json.loads((tmp_path / "r1" / "diagnostics.json").read_text()) == ( + diagnostics + ) + + +async def test_result_extra_fields_are_archived_without_a_sample( + tmp_path: Path, monkeypatch +) -> None: + posted = patch_callbacks(monkeypatch) + patch_artifact_root(monkeypatch, tmp_path) + diagnostics = { + "backend": "native_harbor", + "phase": "setup", + "harbor_exception_type": "RuntimeError", + "category": "agent_error", + } + backend = StubBackend( + workflow_result=ExecutionResult( + status=RolloutStatus.FAILURE, + extra_fields=diagnostics, + ) + ) + + await _handle_rollout(backend, make_request(grader_callback_url=None)) + + assert posted[0][1]["extra_fields"] == diagnostics + assert json.loads((tmp_path / "r1" / "diagnostics.json").read_text()) == ( + diagnostics + ) + assert not (tmp_path / "r1" / "trajectory.json").exists() + + +async def test_native_late_failure_without_grader_url_keeps_final_diagnostics( + tmp_path: Path, monkeypatch +) -> None: + posted = patch_callbacks(monkeypatch) + patch_artifact_root(monkeypatch, tmp_path) + backend = NativeHarborBackend(trials_dir=tmp_path / "trials") + backend.artifact_root = tmp_path + result = SimpleNamespace( + verifier_result=SimpleNamespace(rewards={"reward": 0.7}), + step_results=None, + exception_info=None, + ) + + async def fake_submit(_queue: Any, trial_config: Any) -> Any: + hook_event = SimpleNamespace( + trial_name=trial_config.trial_name, + result=result, + ) + await backend._on_verification_started(hook_event) + result.exception_info = SimpleNamespace( + exception_message="verifier timed out", + exception_type="VerifierTimeoutError", + ) + await backend._on_trial_ended(hook_event) + return result + + monkeypatch.setattr(native_backend_module.TrialQueue, "submit", fake_submit) + + await _handle_rollout( + backend, + make_request( + grader_callback_url=None, + metadata={"harbor_task": "/tmp/task"}, + ), + ) + + assert len(posted) == 1 + assert posted[0][1]["status"] == "success" + diagnostics = json.loads((tmp_path / "r1" / "diagnostics.json").read_text()) + assert diagnostics["phase"] == "verification" + assert diagnostics["harbor_exception_type"] == "VerifierTimeoutError" + assert diagnostics["category"] == "agent_error" + + async def test_records_workflow_result_without_grader_callback( tmp_path: Path, monkeypatch ) -> None: diff --git a/tests/unit/rollout/test_trajectory_converter.py b/tests/unit/rollout/test_trajectory_converter.py index 80d6c21b..dc714158 100644 --- a/tests/unit/rollout/test_trajectory_converter.py +++ b/tests/unit/rollout/test_trajectory_converter.py @@ -303,6 +303,7 @@ def test_extra_carries_platform_context() -> None: request_label="request label", request_metadata={"dataset_row": {"q": "x"}}, request_extra_fields={"eval_run_id": "er-1", "row_index": 3}, + result_extra_fields={"backend": "native_harbor", "phase": "agent"}, ) osmosis = trajectory.extra["osmosis"] @@ -313,6 +314,10 @@ def test_extra_carries_platform_context() -> None: assert osmosis["sample_extra_fields"] == {"custom": True} assert osmosis["request_metadata"] == {"dataset_row": {"q": "x"}} assert osmosis["request_extra_fields"] == {"eval_run_id": "er-1", "row_index": 3} + assert osmosis["result_extra_fields"] == { + "backend": "native_harbor", + "phase": "agent", + } def test_request_label_used_when_sample_has_none() -> None: diff --git a/tests/unit/rollout/test_trajectory_save.py b/tests/unit/rollout/test_trajectory_save.py index 3caa4968..1af488bd 100644 --- a/tests/unit/rollout/test_trajectory_save.py +++ b/tests/unit/rollout/test_trajectory_save.py @@ -56,6 +56,29 @@ async def test_save_without_sample_writes_nothing(tmp_path: Path) -> None: assert not (tmp_path / "r1").exists() +async def test_save_diagnostics_without_sample_writes_sidecar(tmp_path: Path) -> None: + diagnostics = { + "backend": "native_harbor", + "phase": "setup", + "harbor_exception_type": "RuntimeError", + "category": "agent_error", + } + + await save_trajectories( + rollout_id="r1", + result=ExecutionResult( + status=RolloutStatus.FAILURE, + extra_fields=diagnostics, + ), + artifact_root=tmp_path, + ) + + assert json.loads((tmp_path / "r1" / "diagnostics.json").read_text()) == ( + diagnostics + ) + assert not (tmp_path / "r1" / "trajectory.json").exists() + + async def test_save_skips_sample_without_trajectory_messages( tmp_path: Path, caplog ) -> None: @@ -163,3 +186,95 @@ async def test_multi_entry_report_is_preserved_not_guessed( "judge": {"llm_call_metrics": [{"prompt_tokens": 99}]}, } assert any("preserving them under" in r.getMessage() for r in caplog.records) + + +async def test_native_atif_is_preserved_and_enriched_without_message_roundtrip( + tmp_path: Path, +) -> None: + native_document = { + "schema_version": "ATIF-v1.7", + "session_id": "harbor-session", + "trajectory_id": "harbor-trajectory", + "agent": { + "name": "terminus-2", + "version": "0.20.0", + "model_name": "agent-reported-model", + "extra": {"runtime": "native"}, + }, + "steps": [ + { + "step_id": 1, + "source": "agent", + "message": "calling tool", + "reasoning_content": "native reasoning", + "tool_calls": [ + { + "tool_call_id": "call-1", + "function_name": "shell", + "arguments": {"command": "true"}, + } + ], + "observation": { + "results": [ + {"source_call_id": "call-1", "content": "native output"} + ] + }, + "metrics": {"prompt_tokens": 999}, + "llm_call_count": 1, + "extra": {"native-step": True}, + } + ], + "final_metrics": {"total_prompt_tokens": 999, "total_steps": 1}, + "extra": {"harbor": {"preserved": True}}, + } + result = ExecutionResult( + status=RolloutStatus.SUCCESS, + sample=RolloutSample( + label="expected", + reward=0.75, + trajectory_messages=None, + ), + trajectory_document=native_document, + extra_fields={"backend": "native_harbor", "phase": "verification"}, + ) + report = TrajectoryReport( + model_name="controller-model", + samples={ + "single": SampleReport( + llm_call_metrics=[LlmCallMetrics(prompt_tokens=11, completion_tokens=4)] + ) + }, + ) + + await save_trajectories( + rollout_id="r1", + result=result, + request_metadata={"harbor_task": "org/task"}, + report=report, + artifact_root=tmp_path, + ) + + doc = json.loads((tmp_path / "r1" / "trajectory.json").read_text()) + assert doc["session_id"] == "r1" + assert doc["trajectory_id"] == "r1" + assert doc["agent"]["model_name"] == "controller-model" + assert doc["agent"]["extra"] == {"runtime": "native"} + assert doc["steps"][0]["reasoning_content"] == "native reasoning" + assert doc["steps"][0]["tool_calls"][0]["arguments"] == {"command": "true"} + assert doc["steps"][0]["observation"]["results"][0]["content"] == ("native output") + assert doc["steps"][0]["extra"] == {"native-step": True} + assert doc["steps"][0]["metrics"]["prompt_tokens"] == 11 + assert doc["final_metrics"] == { + "total_prompt_tokens": 11, + "total_completion_tokens": 4, + "total_steps": 1, + } + assert doc["extra"]["harbor"] == {"preserved": True} + assert doc["extra"]["osmosis"]["native_session_id"] == "harbor-session" + assert doc["extra"]["osmosis"]["native_trajectory_id"] == "harbor-trajectory" + assert doc["extra"]["osmosis"]["reward"] == 0.75 + assert doc["extra"]["osmosis"]["request_metadata"] == {"harbor_task": "org/task"} + assert doc["extra"]["osmosis"]["result_extra_fields"] == { + "backend": "native_harbor", + "phase": "verification", + }