diff --git a/benchmarks/container_lifecycle/bench_harness/bench_harness/__init__.py b/benchmarks/container_lifecycle/bench_harness/bench_harness/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/benchmarks/container_lifecycle/bench_harness/bench_harness/grade.py b/benchmarks/container_lifecycle/bench_harness/bench_harness/grade.py new file mode 100644 index 00000000..c2962aab --- /dev/null +++ b/benchmarks/container_lifecycle/bench_harness/bench_harness/grade.py @@ -0,0 +1,8 @@ +"""Bench grader: constant reward; exists to exercise the in-env grader path.""" + +from osmosis_ai.rollout import Grader, GraderContext + + +class BenchGrader(Grader): + async def grade(self, ctx: GraderContext) -> None: + ctx.set_reward(1.0) diff --git a/benchmarks/container_lifecycle/bench_harness/bench_harness/solver.py b/benchmarks/container_lifecycle/bench_harness/bench_harness/solver.py new file mode 100644 index 00000000..a9abd15c --- /dev/null +++ b/benchmarks/container_lifecycle/bench_harness/bench_harness/solver.py @@ -0,0 +1,25 @@ +"""Bench workflow: two chat calls against the rollout context's model URL.""" + +import httpx + +from osmosis_ai.rollout import AgentWorkflow, AgentWorkflowContext, get_rollout_context + + +class BenchWorkflow(AgentWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> list[dict]: + rollout_ctx = get_rollout_context() + url = f"{rollout_ctx.chat_completions_url.rstrip('/')}/chat/completions" + headers = {"Authorization": f"Bearer {rollout_ctx.api_key}"} + messages = list(ctx.prompt) + + async with httpx.AsyncClient(timeout=60) as client: + for _ in range(2): + response = await client.post( + url, + json={"model": "bench", "messages": messages}, + headers=headers, + ) + response.raise_for_status() + messages.append(response.json()["choices"][0]["message"]) + + return messages diff --git a/benchmarks/container_lifecycle/bench_harness/pyproject.toml b/benchmarks/container_lifecycle/bench_harness/pyproject.toml new file mode 100644 index 00000000..198b1c0d --- /dev/null +++ b/benchmarks/container_lifecycle/bench_harness/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "bench-harness" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "osmosis-ai @ https://github.com/Osmosis-AI/osmosis-sdk-python/archive/refs/heads/feat/harbor-backend-v2.tar.gz", +] + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +include = ["bench_harness*"] diff --git a/benchmarks/container_lifecycle/container_lifecycle_bench.py b/benchmarks/container_lifecycle/container_lifecycle_bench.py new file mode 100644 index 00000000..eebc954f --- /dev/null +++ b/benchmarks/container_lifecycle/container_lifecycle_bench.py @@ -0,0 +1,319 @@ +"""Benchmark the Harbor rollout backend end to end on local Docker. + +Starts a stub OpenAI endpoint, a controller that receives the protocol +callbacks, and a real rollout server, then drives concurrent rollouts through +the full container pipeline for several consecutive runs. + +Usage: + uv run benchmarks/container_lifecycle/container_lifecycle_bench.py \\ + --runs 5 --concurrency 20 + +Point it at any Harbor task folder instead of the generated one: + + uv run benchmarks/container_lifecycle/container_lifecycle_bench.py \\ + --tasks-dir path/to/task + +The bench harness declares its SDK source in bench_harness/pyproject.toml, +so any image with python3 and pip works; the bundle install pulls the SDK. +""" + +from __future__ import annotations + +import argparse +import asyncio +import socket +import statistics +import subprocess +import sys +import tempfile +import time +import uuid +from pathlib import Path + +import httpx +import uvicorn +from fastapi import FastAPI, Request +from harbor.trial.queue import TrialQueue + +from osmosis_ai.rollout.backend.harbor import HarborBackendV2 +from osmosis_ai.rollout.server import create_rollout_server + +HARNESS_DIR = Path(__file__).resolve().parent / "bench_harness" +sys.path.insert(0, str(HARNESS_DIR)) # a real server runs inside its own project +WORKFLOW = "bench_harness.solver:BenchWorkflow" +GRADER = "bench_harness.grade:BenchGrader" + +DOCKERFILE = """\ +FROM python:3.12-slim +""" + + +def free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def stub_llm(latency: float) -> FastAPI: + app = FastAPI() + + @app.post("/v1/chat/completions") + async def completions(request: Request) -> dict: + await asyncio.sleep(latency) + return { + "id": "bench", + "object": "chat.completion", + "model": "bench", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "bench response"}, + "finish_reason": "stop", + } + ], + } + + return app + + +class Controller: + """The protocol's controller half: submits rollouts, receives callbacks.""" + + def __init__(self, port: int, rollout_url: str, stub_url: str): + self.port = port + self.rollout_url = rollout_url + self.stub_url = stub_url + self.outcomes: dict[str, asyncio.Future[str]] = {} + self.app = FastAPI() + + @self.app.post("/rollout/{rollout_id}/completed") + async def completed(rollout_id: str, request: Request) -> dict: + return {"status": "ok"} + + @self.app.post("/grader/{rollout_id}/completed") + async def graded(rollout_id: str, request: Request) -> dict: + body = await request.json() + future = self.outcomes.get(rollout_id) + if future and not future.done(): + future.set_result(body.get("status", "unknown")) + return {"status": "ok"} + + async def submit( + self, client: httpx.AsyncClient, rollout_id: str, timeout: float + ) -> tuple[str, float]: + """Run one rollout; return its outcome status and wall time.""" + self.outcomes[rollout_id] = asyncio.get_running_loop().create_future() + base = f"http://127.0.0.1:{self.port}" + start = time.monotonic() + response = await client.post( + f"{self.rollout_url}/rollout", + json={ + "rollout_id": rollout_id, + "initial_messages": [{"role": "user", "content": "bench"}], + "label": "bench", + "chat_completions_url": f"{self.stub_url}/v1", + "completion_callback_url": f"{base}/rollout/{rollout_id}/completed", + "grader_callback_url": f"{base}/grader/{rollout_id}/completed", + "controller_api_key": "bench", + "agent_timeout_sec": timeout, + "grader_timeout_sec": timeout, + }, + ) + response.raise_for_status() + try: + status = await asyncio.wait_for(self.outcomes[rollout_id], timeout) + except TimeoutError: + status = "timeout" + return status, time.monotonic() - start + + +async def serve(app: FastAPI, port: int) -> tuple[uvicorn.Server, asyncio.Task]: + server = uvicorn.Server( + uvicorn.Config(app, host="127.0.0.1", port=port, log_level="warning") + ) + task = asyncio.create_task(server.serve()) + while not server.started: + await asyncio.sleep(0.05) + return server, task + + +def prepare_task(work_dir: Path) -> Path: + task_dir = work_dir / "bench-task" + env_dir = task_dir / "environment" + env_dir.mkdir(parents=True) + (env_dir / "Dockerfile").write_text(DOCKERFILE) + (task_dir / "task.toml").write_text('[task]\nname = "osmosis/bench"\n') + return task_dir + + +def make_backend( + task_dir: Path, + concurrency: int, + keep_trials: bool, + patch_dockerfile_with_sdk: bool = True, +): + return HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=concurrency), + tasks_dir=task_dir, + agent=WORKFLOW, + grader=GRADER, + cleanup_successful_trials=not keep_trials, + patch_dockerfile_with_sdk=patch_dockerfile_with_sdk, + ) + + +async def run_series( + task_dir: Path, + stub_url: str, + runs: int, + concurrency: int, + timeout: float, + keep_trials: bool, + patch_dockerfile_with_sdk: bool = True, +) -> dict: + setup_start = time.monotonic() + backend = make_backend( + task_dir, concurrency, keep_trials, patch_dockerfile_with_sdk + ) + setup = time.monotonic() - setup_start + + controller_port, rollout_port = free_port(), free_port() + controller = Controller( + controller_port, + rollout_url=f"http://127.0.0.1:{rollout_port}", + stub_url=stub_url, + ) + servers = [ + await serve(controller.app, controller_port), + await serve(create_rollout_server(backend=backend), rollout_port), + ] + + walls: list[float] = [] + warm_latencies: list[float] = [] + failures = 0 + # Unique per invocation: rollout ids become docker compose project names + # and trial dirs, so two overlapping bench runs must never share them. + nonce = uuid.uuid4().hex[:6] + try: + async with httpx.AsyncClient(timeout=30) as client: + for run in range(1, runs + 1): + start = time.monotonic() + outcomes = await asyncio.gather( + *( + controller.submit( + client, f"bench-{nonce}-run{run}-{i}", timeout + ) + for i in range(concurrency) + ) + ) + wall = time.monotonic() - start + walls.append(wall) + if run > 1: + warm_latencies.extend(latency for status, latency in outcomes) + succeeded = sum( + 1 for status, latency in outcomes if status == "success" + ) + failures += concurrency - succeeded + print( + f"run {run}: {wall:.1f}s, {succeeded}/{concurrency} ok, " + f"{concurrency / wall:.2f} rollouts/s" + ) + finally: + for server, _task in servers: + server.should_exit = True + await asyncio.gather(*(task for server, task in servers)) + + return { + "setup": setup, + "walls": walls, + "warm_latencies": warm_latencies, + "failures": failures, + } + + +def percentile(values: list[float], fraction: float) -> float: + ordered = sorted(values) + return ordered[round(fraction * (len(ordered) - 1))] + + +def report(result: dict, concurrency: int) -> None: + header = ( + f"{'setup':>7} {'cold':>7} " + f"{'warm mean':>10} {'warm max':>9} {'warm rps':>9} " + f"{'lat p50':>8} {'lat p95':>8} {'lat max':>8}" + ) + print(f"\nwarm = runs 2+; lat = per-rollout submit->graded seconds\n{header}") + warm = result["walls"][1:] + latencies = result["warm_latencies"] + if not warm: + print(f"{result['setup']:>6.1f}s {result['walls'][0]:>6.1f}s (single run)") + else: + mean = statistics.mean(warm) + print( + f"{result['setup']:>6.1f}s {result['walls'][0]:>6.1f}s " + f"{mean:>9.1f}s {max(warm):>8.1f}s {concurrency / mean:>9.2f} " + f"{percentile(latencies, 0.5):>7.1f}s " + f"{percentile(latencies, 0.95):>7.1f}s " + f"{max(latencies):>7.1f}s" + ) + if result["failures"]: + raise SystemExit(f"failures: {result['failures']}") + + +async def bench(args: argparse.Namespace) -> None: + work_dir = Path(tempfile.mkdtemp(prefix="harbor-bench-")) + print(f"work dir: {work_dir}") + + stub_port = free_port() + stub_server, stub_task = await serve(stub_llm(args.latency), stub_port) + stub_url = f"http://127.0.0.1:{stub_port}" + + try: + result = await run_series( + args.tasks_dir or prepare_task(work_dir), + stub_url, + args.runs, + args.concurrency, + args.timeout, + args.keep_trials, + args.patch_dockerfile_with_sdk, + ) + finally: + stub_server.should_exit = True + await stub_task + + report(result, args.concurrency) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--tasks-dir", + type=Path, + default=None, + help="benchmark an arbitrary Harbor task folder instead of the " + "generated default (see module docstring for image requirements)", + ) + parser.add_argument( + "--no-patch-dockerfile-with-sdk", + dest="patch_dockerfile_with_sdk", + action="store_false", + help="disable the default Dockerfile patch that pre-installs the " + "harness's dependencies into the task image; dependencies then " + "download inside every trial container", + ) + parser.add_argument("--runs", type=int, default=5) + parser.add_argument("--concurrency", type=int, default=20) + parser.add_argument("--latency", type=float, default=0.2) + parser.add_argument("--timeout", type=float, default=600) + parser.add_argument("--keep-trials", action="store_true") + args = parser.parse_args() + if subprocess.run(["docker", "info"], capture_output=True).returncode != 0: + raise SystemExit("docker daemon is required") + asyncio.run(bench(args)) + + +if __name__ == "__main__": + main() diff --git a/osmosis_ai/packaging.py b/osmosis_ai/packaging.py new file mode 100644 index 00000000..b21aa545 --- /dev/null +++ b/osmosis_ai/packaging.py @@ -0,0 +1,260 @@ +"""Build an installable wheel from an agent/harness project dir. + +Stages the user's project (their pyproject.toml, dependencies, and build +backend included), injects a generated shim with literal imports, and exposes +``-agent`` / ``-grade`` console scripts. The wheel installs +anywhere with one ``pip install`` — a rollout container or a user's own box. +""" + +from __future__ import annotations + +import hashlib +import shutil +import subprocess +import tempfile +import tomllib +import zipfile +from dataclasses import dataclass +from importlib.metadata import PathDistribution +from pathlib import Path + +import platformdirs +import toml +from packaging.requirements import Requirement +from packaging.utils import canonicalize_name + +BUNDLES_DIR = platformdirs.user_cache_path("osmosis") / "bundles" + +AGENT_MAIN_TEMPLATE = """\ + + +def agent_main(): + runner.agent_main({workflow_class}, {workflow_config}) +""" + +GRADER_MAIN_TEMPLATE = """\ + + +def grader_main(): + runner.grader_main({grader_class}, {grader_config}) +""" + +EXCLUDE_DIRS = { + "__pycache__", + ".git", + ".venv", + "venv", + "dist", + "build", + "node_modules", + ".pytest_cache", + ".ruff_cache", +} + + +@dataclass(frozen=True) +class BundleInfo: + wheel: Path + agent_script: str | None + grader_script: str | None + requirements: list[str] + + +def agent_script_name(package: str) -> str: + return f"{package.replace('_', '-')}-agent" + + +def grader_script_name(package: str) -> str: + return f"{package.replace('_', '-')}-grade" + + +def content_hash(path: Path, *, extra: str = "") -> str: + digest = hashlib.sha256(extra.encode()) + for file in sorted( + p for p in path.rglob("*") if p.is_file() and not (set(p.parts) & EXCLUDE_DIRS) + ): + digest.update(str(file.relative_to(path)).encode()) + digest.update(file.read_bytes()) + return digest.hexdigest()[:32] + + +def find_package(code_dir: Path) -> str: + packages = sorted( + d.name + for d in code_dir.iterdir() + if d.is_dir() and (d / "__init__.py").exists() and d.name not in EXCLUDE_DIRS + ) + if len(packages) != 1: + raise ValueError( + f"expected exactly one python package in {code_dir}, " + f"found {packages or 'none'}; pass package= explicitly" + ) + return packages[0] + + +def requirement_name(spec: str) -> str: + return canonicalize_name(Requirement(spec).name) + + +def split_ref(ref: str, label: str) -> tuple[str, str]: + if ":" not in ref: + raise ValueError(f"{label} must be 'module:attr', got {ref!r}") + module, attr = ref.rsplit(":", 1) + return module, attr + + +def project_dir_for(obj: object) -> Path: + """Locate the project dir (containing pyproject.toml) of *obj*'s package.""" + import inspect + + package_dir = Path( + inspect.getfile(type(obj) if not isinstance(obj, type) else obj) + ).parent + while (package_dir.parent / "__init__.py").exists(): + package_dir = package_dir.parent + project_dir = package_dir.parent + if not (project_dir / "pyproject.toml").is_file(): + raise ValueError( + f"no pyproject.toml next to package {package_dir.name!r} " + f"(looked in {project_dir}); pass code_dir= explicitly" + ) + return project_dir + + +def extra_gated(spec: str) -> bool: + marker = Requirement(spec).marker + return marker is not None and "extra ==" in str(marker) + + +def inspect_bundle(wheel: Path) -> BundleInfo: + """Read the bundle's console scripts and dependencies from wheel metadata. + + Extra-gated dependencies are dropped; environment markers (python_version, + sys_platform, …) ship verbatim since pip evaluates them inside the + container, where they may resolve differently than on this host. + """ + dist_info = next( + entry + for entry in zipfile.Path(wheel).iterdir() + if entry.name.endswith(".dist-info") + ) + dist = PathDistribution(dist_info) + scripts = { + ep.name: ep.value for ep in dist.entry_points if ep.group == "console_scripts" + } + agent = next( + (n for n, t in scripts.items() if t.endswith(".bundle_main:agent_main")), None + ) + grader = next( + (n for n, t in scripts.items() if t.endswith(".bundle_main:grader_main")), None + ) + if agent is None and grader is None: + raise ValueError(f"{wheel.name} is not an osmosis bundle (no runner scripts)") + return BundleInfo( + wheel=wheel, + agent_script=agent, + grader_script=grader, + requirements=[spec for spec in dist.requires or [] if not extra_gated(spec)], + ) + + +def build_bundle( + code_dir: Path, + *, + workflow: str | None = None, + grader: str | None = None, + workflow_config: str | None = None, + grader_config: str | None = None, + deps: list[str] | None = None, + package: str | None = None, + bundles_dir: Path = BUNDLES_DIR, +) -> Path: + """Package the project at *code_dir* into a wheel; rebuild only on change. + + workflow/grader are 'module:Class' refs; the config refs point at + module-level instances and are passed to the runner as-is. Grader-only + bundles (workflow=None) carry just the grading script. Entries in *deps* + replace same-named dependencies from the project's pyproject.toml. + """ + code_dir = code_dir.resolve() + pyproject_path = code_dir / "pyproject.toml" + if not pyproject_path.is_file(): + raise ValueError(f"project dir must contain pyproject.toml: {code_dir}") + if workflow is None and grader is None: + raise ValueError("pass workflow and/or grader") + package = package or find_package(code_dir) + + imports = [] + references = {} + for label, ref in ( + ("workflow", workflow), + ("grader", grader), + ("workflow_config", workflow_config), + ("grader_config", grader_config), + ): + if ref is None: + references[label] = "None" + continue + module, attr = split_ref(ref, label) + imports.append(f"from {module} import {attr}") + references[label] = attr + + shim = "\n".join(imports) + "\nfrom osmosis_ai.rollout.container import runner\n" + scripts = {} + if workflow: + shim += AGENT_MAIN_TEMPLATE.format( + workflow_class=references["workflow"], + workflow_config=references["workflow_config"], + ) + scripts[agent_script_name(package)] = f"{package}.bundle_main:agent_main" + if grader: + shim += GRADER_MAIN_TEMPLATE.format( + grader_class=references["grader"], + grader_config=references["grader_config"], + ) + scripts[grader_script_name(package)] = f"{package}.bundle_main:grader_main" + + project = tomllib.loads(pyproject_path.read_text()).get("project", {}) + name = project.get("name") or f"osmosis-harness-{package}" + bundle_key = content_hash( + code_dir, extra=pyproject_path.read_text() + shim + repr(deps or []) + ) + bundles_dir.mkdir(parents=True, exist_ok=True) + wheel_glob = f"{name.replace('-', '_')}-*.whl" + marker = bundles_dir / f"{name}-{bundle_key}.key" + cached = sorted(bundles_dir.glob(wheel_glob)) + if cached and marker.exists(): + return cached[0] + + for old in bundles_dir.glob(f"{name}-*.key"): + old.unlink() + with tempfile.TemporaryDirectory() as staging: + stage = Path(staging) + shutil.copytree( + code_dir, + stage, + dirs_exist_ok=True, + ignore=shutil.ignore_patterns(*EXCLUDE_DIRS, "*.pyc"), + ) + (stage / package / "bundle_main.py").write_text(shim) + staged = tomllib.loads((stage / "pyproject.toml").read_text()) + staged_project = staged.setdefault("project", {}) + overridden = {requirement_name(d) for d in deps or []} + staged_project["dependencies"] = [ + d + for d in staged_project.get("dependencies", []) + if requirement_name(d) not in overridden + ] + list(deps or []) + staged_project["scripts"] = {**staged_project.get("scripts", {}), **scripts} + (stage / "pyproject.toml").write_text(toml.dumps(staged)) + subprocess.run( + ["uv", "build", "--wheel", "--out-dir", str(bundles_dir), str(stage)], + check=True, + capture_output=True, + ) + + marker.touch() + wheels = sorted(bundles_dir.glob(wheel_glob)) + if not wheels: + raise RuntimeError(f"uv build produced no wheel in {bundles_dir}") + return wheels[0] diff --git a/osmosis_ai/rollout/backend/harbor/__init__.py b/osmosis_ai/rollout/backend/harbor/__init__.py index 2b692b4a..e27e43a7 100644 --- a/osmosis_ai/rollout/backend/harbor/__init__.py +++ b/osmosis_ai/rollout/backend/harbor/__init__.py @@ -1,7 +1,11 @@ from osmosis_ai.rollout.backend.harbor.agent_adapter import OsmosisInstalledAgent from osmosis_ai.rollout.backend.harbor.backend import HarborBackend +from osmosis_ai.rollout.backend.harbor.backend_v2 import HarborBackendV2 +from osmosis_ai.rollout.backend.harbor.tasks import TaskMode __all__ = [ "HarborBackend", + "HarborBackendV2", "OsmosisInstalledAgent", + "TaskMode", ] diff --git a/osmosis_ai/rollout/backend/harbor/artifacts.py b/osmosis_ai/rollout/backend/harbor/artifacts.py new file mode 100644 index 00000000..405c7fe3 --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/artifacts.py @@ -0,0 +1,60 @@ +"""Host-side artifact movement around a finished trial.""" + +from __future__ import annotations + +import logging +import shutil +from pathlib import Path + +from osmosis_ai.rollout.backend.harbor.backend import TRIAL_NAME_PREFIX +from osmosis_ai.rollout.utils.file_artifacts import ( + GRADER_ARTIFACTS_SNAPSHOT_DIRNAME, + HARBOR_ARTIFACTS_DIR, + copy_artifact_tree, +) + +logger = logging.getLogger(__name__) + + +def relocate_trial_artifacts( + trials_dir: Path, artifact_root: Path, rollout_id: str, *, move: bool +) -> bool: + source_dir = trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" / "artifacts" + if not source_dir.is_dir(): + return True + try: + copy_artifact_tree( + source_dir, + artifact_root / rollout_id / "artifacts", + destination_root=artifact_root, + replace_destination=True, + ) + if move: + shutil.rmtree(source_dir) + except Exception: + logger.warning( + "Failed to relocate trial artifacts for rollout %s (best-effort)", + rollout_id, + exc_info=True, + ) + return False + return True + + +def merge_grader_artifacts(trials_dir: Path, rollout_id: str) -> None: + trial_dir = trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + source_dir = trial_dir / "verifier" / GRADER_ARTIFACTS_SNAPSHOT_DIRNAME + if not source_dir.is_dir(): + return + try: + copy_artifact_tree( + source_dir, + trial_dir / "artifacts" / HARBOR_ARTIFACTS_DIR.relative_to("/"), + destination_root=trials_dir, + ) + except Exception: + logger.warning( + "Failed to merge grader artifacts for rollout %s (best-effort)", + rollout_id, + exc_info=True, + ) diff --git a/osmosis_ai/rollout/backend/harbor/backend.py b/osmosis_ai/rollout/backend/harbor/backend.py index a8d4f261..4f94d5ea 100644 --- a/osmosis_ai/rollout/backend/harbor/backend.py +++ b/osmosis_ai/rollout/backend/harbor/backend.py @@ -12,6 +12,7 @@ from typing import Any import toml +from harbor.environments.definition import environment_content_hash from harbor.models.environment_type import EnvironmentType from harbor.models.trial.config import ( AgentConfig as HarborAgentConfig, @@ -118,6 +119,8 @@ def __init__( self.on_workflow_complete = on_workflow_complete self.on_grader_complete = on_grader_complete self.workflow_complete_called = False + self.preserve_trial = False + self.api_key: str | None = None self.done: asyncio.Future[None] = asyncio.get_event_loop().create_future() @@ -252,8 +255,16 @@ def prepare_shared_env(self) -> None: ) def build_image(self) -> str: - """Build the Docker image once from the shared env dir and return the tag.""" - image_tag = f"osmosis-harbor-{self.task_dir.name}:latest" + """Content-addressed build: skip when the shared env dir is unchanged.""" + image_tag = f"osmosis-harbor-{environment_content_hash(self.shared_env_dir)}" + inspect = subprocess.run( + ["docker", "image", "inspect", image_tag], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + if inspect.returncode == 0: + logger.info("Reusing prebuilt image %s", image_tag) + return image_tag logger.info( "Building prebuilt image %s from %s", image_tag, self.shared_env_dir ) diff --git a/osmosis_ai/rollout/backend/harbor/backend_v2.py b/osmosis_ai/rollout/backend/harbor/backend_v2.py new file mode 100644 index 00000000..1455ac3c --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/backend_v2.py @@ -0,0 +1,643 @@ +"""Harbor execution backend v2. Host-side only. + +``agent=`` selects the track: a registered native agent name ("terminus-2", +"mini-swe-agent", "oracle") runs Harbor's own agent with the rollout endpoint +injected; an AgentWorkflow class (or "module:Class" path) is packaged into a +wheel and installed in the task container at trial start. ``grader=None`` +makes the task's own tests the reward source; a Grader class is delivered as +the verifier instead. + +Tasks come from ``tasks_dir`` (template or dataset mode), or per rollout via +``metadata["harbor_task"]``: a local path, a registry package "org/name[@ref]", +or a git checkout (with ``metadata["git_url"]`` and, ideally, a pinned +``metadata["git_commit_id"]``). ``metadata["harbor_model"]`` overrides the +model per rollout. + +Task images stay pure task environments either way; with the pinned harbor +0.20.0 each trial builds its own compose-named tag (fast via docker's layer +cache, removed at trial teardown), and newer harbor releases share one +content-addressed hb__ image across trials. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import shutil +import traceback +import uuid +from collections.abc import AsyncIterator, Callable, Sequence +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from pathlib import Path +from typing import Any + +from harbor.models.trajectories import Trajectory +from harbor.models.trial.config import ( + AgentConfig as HarborAgentConfig, +) +from harbor.models.trial.config import ( + EnvironmentConfig as HarborEnvironmentConfig, +) +from harbor.models.trial.config import ( + TaskConfig, + TrialConfig, + VerifierConfig, +) +from harbor.tasks.client import TaskClient +from harbor.trial.hooks import TrialEvent, TrialHookEvent +from harbor.trial.queue import TrialQueue + +from osmosis_ai.rollout.backend.base import ExecutionBackend, ResultCallback +from osmosis_ai.rollout.backend.harbor.artifacts import ( + merge_grader_artifacts, + relocate_trial_artifacts, +) +from osmosis_ai.rollout.backend.harbor.backend import ( + TRIAL_NAME_PREFIX, + PendingTrial, + apply_managed_skypilot_placement, + ensure_import_path, + get_agent_metadata, + log_trial_exception, + parse_rollout_id, + rewrite_url_for_docker, + uses_local_docker_runtime, +) +from osmosis_ai.rollout.backend.harbor.bundling import resolve_backend_bundle +from osmosis_ai.rollout.backend.harbor.diagnostics import ( + diagnostic_payload, + failure_phase, + redact_secrets, + trial_metrics, + trial_timings, +) +from osmosis_ai.rollout.backend.harbor.native_agents import ( + native_agent_config, + native_binding, +) +from osmosis_ai.rollout.backend.harbor.tasks import ( + HarborTask, + TaskMode, + parse_task_ref, +) +from osmosis_ai.rollout.container.files import ContainerInput, ContainerResult +from osmosis_ai.rollout.container.trajectories import messages_from_trajectory +from osmosis_ai.rollout.context import get_rollout_context +from osmosis_ai.rollout.types import ( + ExecutionRequest, + ExecutionResult, + RolloutErrorCategory, + RolloutSample, + RolloutStatus, +) +from osmosis_ai.rollout.utils.file_artifacts import default_artifact_root +from osmosis_ai.rollout.utils.rewards import validate_sample_has_reward + +logger: logging.Logger = logging.getLogger(__name__) + +HARNESS_AGENT_IMPORT_PATH = ( + "osmosis_ai.rollout.backend.harbor.harness_agent:OsmosisHarnessInstalledAgent" +) +PREWARM_PREFIX = "prewarm-" + + +class HarborBackendV2(ExecutionBackend): + def __init__( + self, + *, + orchestrator: TrialQueue, + tasks_dir: Path, + agent: str | type | None = None, + task_mode: TaskMode | str = TaskMode.TEMPLATE, + model_name: str = "openai/osmosis-rollout", + grader: type | str | None = None, + workflow_config: Any = None, + grader_config: Any = None, + code_dir: Path | None = None, + bundle: Path | None = None, + environment_config: HarborEnvironmentConfig | None = None, + trials_dir: Path | None = None, + cleanup_successful_trials: bool = True, + patch_dockerfile_with_sdk: bool | None = None, + agent_setup_timeout_sec: float | None = None, + ) -> None: + self.orchestrator = orchestrator + self.tasks_dir = Path(tasks_dir) + self.task_mode = TaskMode(task_mode) + self.model_name = model_name + self.agent = agent + self.agent_setup_timeout_sec = agent_setup_timeout_sec + self.native = native_binding(agent) + if self.native and not self.native.trainable: + logger.warning( + "native agent %r emits no model trajectory; use it to validate " + "datasets and verifiers, never for training", + agent, + ) + self.bundle = resolve_backend_bundle( + agent=agent, + grader=grader, + workflow_config=workflow_config, + grader_config=grader_config, + code_dir=code_dir, + bundle=bundle, + native=self.native is not None, + ) + self.environment_config = apply_managed_skypilot_placement( + environment_config or HarborEnvironmentConfig() + ) + if patch_dockerfile_with_sdk and self.bundle is None: + raise ValueError("patch_dockerfile_with_sdk requires a bundle") + if patch_dockerfile_with_sdk is None: + patch_dockerfile_with_sdk = self.bundle is not None + # The bundle's declared dependencies (stable) are pre-installed into the + # task image; the bundle itself (volatile user code) installs per trial. + self.sdk_requirements = ( + self.bundle.requirements + if patch_dockerfile_with_sdk and self.bundle + else None + ) + + root = Path(f"/tmp/osmosis-harbor-{self.tasks_dir.name}") + self.rollouts_dir = root / "rollouts" + self.rollouts_dir.mkdir(parents=True, exist_ok=True) + self.trials_dir = trials_dir or root / "trials" + self.artifact_root = default_artifact_root() + self.cleanup_successful_trials = cleanup_successful_trials + self.pending: dict[str, PendingTrial] = {} + self.fetch_locks: dict[str, asyncio.Lock] = {} + self.running = 0 + + orchestrator.add_hook(TrialEvent.START, self.on_trial_started) + orchestrator.add_hook(TrialEvent.VERIFICATION_START, self.on_verification_start) + orchestrator.add_hook(TrialEvent.END, self.on_trial_end) + + def health(self) -> dict[str, Any]: + return { + "status": "ok", + "backend": "harbor-v2", + "agent": self.agent + if isinstance(self.agent, str) + else ensure_import_path(self.agent), + "in_flight": len(self.pending), + "running": self.running, + "queued": max(0, len(self.pending) - self.running), + } + + async def resolve_task(self, request: ExecutionRequest) -> HarborTask: + metadata = request.metadata or {} + if ref := metadata.get("harbor_task"): + return await self.fetch_task(str(ref), metadata) + return self.select_task(request) + + def select_task(self, request: ExecutionRequest) -> HarborTask: + if self.task_mode == TaskMode.TEMPLATE: + return HarborTask(self.tasks_dir) + task_id = (request.metadata or {}).get("harbor_task_id") + if not task_id: + raise ValueError( + "dataset mode requires metadata['harbor_task_id'] or " + "metadata['harbor_task']" + ) + return HarborTask.from_dataset(self.tasks_dir, task_id) + + async def fetch_task(self, ref: str, metadata: dict[str, Any]) -> HarborTask: + """Download (or reuse from cache) a git/package/local task ref.""" + lock = self.fetch_locks.setdefault(ref, asyncio.Lock()) + async with lock: + batch = await TaskClient().download_tasks([parse_task_ref(ref, metadata)]) + return HarborTask(batch.paths[0]) + + def build_input(self, request: ExecutionRequest) -> ContainerInput: + container_input = ContainerInput( + rollout_id=request.id, + # Dataset tasks carry their own instruction.md; row prompts only + # drive template mode. + prompt=(request.prompt or []) + if self.task_mode == TaskMode.TEMPLATE + else [], + label=request.label, + metadata=request.metadata, + ) + ctx = get_rollout_context() + if ctx: + url = ctx.chat_completions_url or "" + if url and uses_local_docker_runtime(self.environment_config): + url = rewrite_url_for_docker(url) + container_input.chat_completions_url = url + container_input.api_key = ctx.api_key + return container_input + + def build_agent_config( + self, task_dir: Path, container_input: ContainerInput + ) -> HarborAgentConfig: + if self.native is not None and isinstance(self.agent, str): + model = (container_input.metadata or {}).get( + "harbor_model" + ) or self.model_name + return native_agent_config( + self.agent, + self.native, + model, + container_input.chat_completions_url, + container_input.api_key or "dummy", + ) + if self.bundle is None: + raise ValueError("workflow agents require a bundle") + return HarborAgentConfig( + import_path=HARNESS_AGENT_IMPORT_PATH, + kwargs={ + "bundle_path": str(self.bundle.wheel), + "agent_script": self.bundle.agent_script, + "input_path": str(task_dir / "container_input.json"), + }, + ) + + def build_trial_config( + self, task_dir: Path, request: ExecutionRequest, container_input: ContainerInput + ) -> TrialConfig: + agent_config = self.build_agent_config(task_dir, container_input) + if request.agent_timeout_sec is not None: + agent_config.override_timeout_sec = request.agent_timeout_sec + if self.agent_setup_timeout_sec is not None: + agent_config.override_setup_timeout_sec = self.agent_setup_timeout_sec + + verifier_config = VerifierConfig(disable=not (task_dir / "tests").is_dir()) + if request.grader_timeout_sec is not None: + verifier_config.override_timeout_sec = request.grader_timeout_sec + + return TrialConfig( + task=TaskConfig(path=task_dir), + trial_name=f"{TRIAL_NAME_PREFIX}{request.id}", + trials_dir=self.trials_dir, + agent=agent_config, + environment=self.environment_config.model_copy(deep=True), + verifier=verifier_config, + ) + + def materialize_task( + self, task: HarborTask, rollout_id: str, container_input: ContainerInput + ) -> Path: + return task.materialize( + self.rollouts_dir / rollout_id, + container_input, + grader_script=self.bundle.grader_script if self.bundle else None, + grader_wheel=self.bundle.wheel if self.bundle and self.native else None, + sdk_requirements=self.sdk_requirements, + ) + + async def execute( + self, + request: ExecutionRequest, + on_workflow_complete: ResultCallback, + on_grader_complete: ResultCallback | None = None, + ) -> None: + pending = PendingTrial(on_workflow_complete, on_grader_complete) + self.pending[request.id] = pending + try: + container_input = self.build_input(request) + pending.api_key = container_input.api_key + task = await self.resolve_task(request) + task_dir = self.materialize_task(task, request.id, container_input) + await self.orchestrator.submit( + self.build_trial_config(task_dir, request, container_input) + ) + await pending.done + except Exception as e: + self.pending.pop(request.id, None) + logger.error("Failed trial %s: %s", request.id, e) + await on_workflow_complete( + ExecutionResult( + status=RolloutStatus.FAILURE, + err_message=str(e), + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=diagnostic_payload( + phase="setup", + category=RolloutErrorCategory.AGENT_ERROR, + exception_type=type(e).__name__, + timings={}, + ), + ) + ) + + async def prewarm(self, task_ids: Sequence[str] | None = None) -> None: + """Build every task image and run agent setup before serving rollouts. + + Prewarm trials are install-only, carry no rollout credentials, and + report all failing tasks together. + """ + if self.task_mode == TaskMode.DATASET: + if not task_ids: + raise ValueError("dataset mode prewarm requires task ids") + tasks = [HarborTask.from_dataset(self.tasks_dir, t) for t in task_ids] + else: + tasks = [HarborTask(self.tasks_dir)] + + configs = [self.prewarm_trial_config(task) for task in tasks] + logger.info("Prewarming %d harbor task(s)", len(configs)) + outcomes = await asyncio.gather( + *(self.orchestrator.submit(config) for config in configs), + return_exceptions=True, + ) + + failures = [] + for config, outcome in zip(configs, outcomes, strict=True): + label = config.task.path.name if config.task.path else config.trial_name + if isinstance(outcome, BaseException): + failures.append(f"{label}: {type(outcome).__name__}: {outcome}") + elif (err := getattr(outcome, "exception_info", None)) is not None: + failures.append(f"{label}: {err.exception_type}") + if failures: + raise RuntimeError( + f"prewarm failed for {len(failures)} of {len(configs)} task(s):\n" + + "\n".join(f" - {failure}" for failure in failures) + ) + logger.info("Prewarmed %d harbor task(s)", len(configs)) + + def prewarm_trial_config(self, task: HarborTask) -> TrialConfig: + rollout_id = f"{PREWARM_PREFIX}{uuid.uuid4().hex[:8]}" + container_input = ContainerInput( + rollout_id=rollout_id, + prompt=[{"role": "user", "content": "prewarm"}] + if self.task_mode == TaskMode.TEMPLATE + else [], + ) + task_dir = self.materialize_task(task, rollout_id, container_input) + config = self.build_trial_config( + task_dir, ExecutionRequest(id=rollout_id, prompt=[]), container_input + ) + config.install_only = True + config.verifier.disable = True + return config + + def prewarm_lifespan( + self, task_ids: Sequence[str] | None = None + ) -> Callable[[object], AbstractAsyncContextManager[None]]: + """An ASGI lifespan that prewarms before the server accepts traffic.""" + + @asynccontextmanager + async def lifespan(app: object) -> AsyncIterator[None]: + await self.prewarm(task_ids) + yield + + return lifespan + + async def try_callback( + self, + callback: ResultCallback, + result: ExecutionResult, + rollout_id: str, + label: str, + ) -> bool: + """Callback delivery failures must never abort trial archival.""" + try: + await callback(result) + return True + except Exception: + logger.error( + "%s callback for rollout %s failed: %s", + label, + rollout_id, + traceback.format_exc(), + ) + return False + + async def on_trial_started(self, event: TrialHookEvent) -> None: + self.running += 1 + + def event_diagnostics( + self, event: TrialHookEvent, category: RolloutErrorCategory | None = None + ) -> dict[str, Any]: + err = event.result.exception_info if event.result else None + return diagnostic_payload( + phase=failure_phase(event.result), + category=category, + exception_type=err.exception_type if err else None, + timings=trial_timings(event.result), + ) + + def container_result(self, event: TrialHookEvent) -> ContainerResult | None: + metadata = get_agent_metadata(event) + if not metadata: + return None + try: + return ContainerResult.model_validate(metadata) + except ValueError: + return None + + def native_sample( + self, event: TrialHookEvent, rollout_id: str, pending: PendingTrial + ) -> RolloutSample | None: + trial_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + paths = [p for p in (trial_dir / "agent" / "trajectory.json",) if p.is_file()] + paths += sorted(trial_dir.glob("steps/*/agent/trajectory.json")) + if not paths: + return None + if len(paths) > 1: + logger.warning( + "rollout %s emitted %d trajectory documents; preserving the " + "trial instead of fabricating a merged trajectory", + rollout_id, + len(paths), + ) + pending.preserve_trial = True + return None + try: + trajectory = Trajectory.model_validate(json.loads(paths[0].read_text())) + except Exception: + logger.warning( + "rollout %s emitted an invalid ATIF trajectory; preserving the trial", + rollout_id, + exc_info=True, + ) + pending.preserve_trial = True + return None + if trajectory.agent.extra: + trajectory.agent.extra = redact_secrets( + trajectory.agent.extra, pending.api_key + ) + messages = messages_from_trajectory(trajectory.to_json_dict(exclude_none=True)) + if not messages: + return None + return RolloutSample(messages=messages, metrics=trial_metrics(event.result)) + + def primary_sample( + self, event: TrialHookEvent, rollout_id: str, pending: PendingTrial + ) -> RolloutSample | None: + if self.native is not None: + return self.native_sample(event, rollout_id, pending) + result = self.container_result(event) + output = result.output if result else None + if output is None: + return None + messages = output.primary_messages() + if messages is None: + return None + return RolloutSample( + messages=messages, + metrics={**trial_metrics(event.result), **output.metrics}, + ) + + def agent_succeeded(self, event: TrialHookEvent) -> tuple[bool, str | None]: + if event.result and event.result.exception_info: + return False, event.result.exception_info.exception_message + if self.native is not None: + return True, None + result = self.container_result(event) + if result is not None and result.status == RolloutStatus.SUCCESS: + return True, None + return False, result.err_message if result else "Unknown error" + + async def on_verification_start(self, event: TrialHookEvent) -> None: + rollout_id = parse_rollout_id(event) + pending = self.pending.get(rollout_id) + if not pending: + logger.error("No pending trial found for rollout %s", rollout_id) + return + + succeeded, err_message = self.agent_succeeded(event) + if succeeded: + outcome = ExecutionResult( + status=RolloutStatus.SUCCESS, + sample=self.primary_sample(event, rollout_id, pending), + extra_fields=self.event_diagnostics(event), + ) + else: + outcome = ExecutionResult( + status=RolloutStatus.FAILURE, + err_message=err_message, + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=self.event_diagnostics( + event, RolloutErrorCategory.AGENT_ERROR + ), + ) + + pending.workflow_complete_called = await self.try_callback( + pending.on_workflow_complete, outcome, rollout_id, "workflow" + ) + + def grader_outcome( + self, event: TrialHookEvent, rollout_id: str, pending: PendingTrial + ) -> ExecutionResult: + sample = self.primary_sample(event, rollout_id, pending) + + if event.result and event.result.verifier_result: + rewards = event.result.verifier_result.rewards or {} + reward = rewards.get("reward") + if sample is not None and reward is not None: + sample.reward = float(reward) + try: + validate_sample_has_reward(sample) + except ValueError as e: + logger.warning( + "Verifier rewards for rollout %s missing 'reward': %s", + rollout_id, + e, + ) + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + err_message=str(e), + err_category=RolloutErrorCategory.VALIDATION_ERROR, + extra_fields=self.event_diagnostics( + event, RolloutErrorCategory.VALIDATION_ERROR + ), + ) + return ExecutionResult( + status=RolloutStatus.SUCCESS, + sample=sample, + extra_fields=self.event_diagnostics(event), + ) + + if event.result and event.result.exception_info: + err = event.result.exception_info + log_trial_exception(rollout_id, err, phase="during grading") + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + err_message=err.exception_message, + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=self.event_diagnostics( + event, RolloutErrorCategory.AGENT_ERROR + ), + ) + return ExecutionResult( + status=RolloutStatus.FAILURE, + sample=sample, + extra_fields=self.event_diagnostics( + event, RolloutErrorCategory.AGENT_ERROR + ), + ) + + async def on_trial_end(self, event: TrialHookEvent) -> None: + rollout_id = parse_rollout_id(event) + self.running = max(0, self.running - 1) + if rollout_id.startswith(PREWARM_PREFIX): + if event.result and not event.result.exception_info: + trial_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + shutil.rmtree(trial_dir, ignore_errors=True) + shutil.rmtree(self.rollouts_dir / rollout_id, ignore_errors=True) + return + + pending = self.pending.pop(rollout_id, None) + if not pending: + logger.error("No pending trial found for rollout %s", rollout_id) + return + + try: + merge_grader_artifacts(self.trials_dir, rollout_id) + grader_result = ( + self.grader_outcome(event, rollout_id, pending) + if pending.on_grader_complete + else None + ) + delete_trial = bool( + self.cleanup_successful_trials + and event.result + and not event.result.exception_info + and not pending.preserve_trial + ) + relocated = relocate_trial_artifacts( + self.trials_dir, self.artifact_root, rollout_id, move=delete_trial + ) + + if not pending.workflow_complete_called: + if event.result and event.result.exception_info: + err = event.result.exception_info + log_trial_exception( + rollout_id, err, phase="before the agent completed" + ) + message = err.exception_message + else: + message = "Trial ended before agent completed" + pending.workflow_complete_called = await self.try_callback( + pending.on_workflow_complete, + ExecutionResult( + status=RolloutStatus.FAILURE, + err_message=message, + err_category=RolloutErrorCategory.AGENT_ERROR, + extra_fields=self.event_diagnostics( + event, RolloutErrorCategory.AGENT_ERROR + ), + ), + rollout_id, + "workflow", + ) + + if pending.on_grader_complete and grader_result is not None: + await self.try_callback( + pending.on_grader_complete, grader_result, rollout_id, "grader" + ) + + if self.cleanup_successful_trials and not pending.preserve_trial: + shutil.rmtree(self.rollouts_dir / rollout_id, ignore_errors=True) + if delete_trial and relocated and not pending.preserve_trial: + trial_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + shutil.rmtree(trial_dir, ignore_errors=True) + finally: + timings = trial_timings(event.result) + if timings: + logger.info("rollout %s phase timings: %s", rollout_id, timings) + if not pending.done.done(): + pending.done.set_result(None) diff --git a/osmosis_ai/rollout/backend/harbor/bundling.py b/osmosis_ai/rollout/backend/harbor/bundling.py new file mode 100644 index 00000000..e0d8a0a7 --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/bundling.py @@ -0,0 +1,49 @@ +"""Resolve the code bundle a backend needs for its agent and grader.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from osmosis_ai.packaging import ( + BundleInfo, + build_bundle, + inspect_bundle, + project_dir_for, +) +from osmosis_ai.rollout.backend.harbor.backend import ensure_import_path +from osmosis_ai.rollout.utils.imports import resolve_object + + +def resolve_backend_bundle( + *, + agent: str | type | None, + grader: type | str | None, + workflow_config: Any = None, + grader_config: Any = None, + code_dir: Path | None = None, + bundle: Path | None = None, + native: bool = False, +) -> BundleInfo | None: + """Native agents need a bundle only when a grader is delivered; + workflow agents always need one.""" + if bundle is not None: + return inspect_bundle(Path(bundle)) + if native: + if grader is None: + return None + anchor = resolve_object(grader) + else: + if agent is None: + raise ValueError("pass agent (a native name or an AgentWorkflow)") + anchor = resolve_object(agent) + wheel = build_bundle( + code_dir or project_dir_for(anchor), + workflow=None if native else ensure_import_path(agent), + grader=ensure_import_path(grader) if grader else None, + workflow_config=ensure_import_path(workflow_config) + if workflow_config + else None, + grader_config=ensure_import_path(grader_config) if grader_config else None, + ) + return inspect_bundle(wheel) diff --git a/osmosis_ai/rollout/backend/harbor/diagnostics.py b/osmosis_ai/rollout/backend/harbor/diagnostics.py new file mode 100644 index 00000000..4530f169 --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/diagnostics.py @@ -0,0 +1,123 @@ +"""Trial observability: phase timings, token totals, failure diagnostics, +and secret redaction. Everything reads Harbor's TrialResult duck-typed, so +tests can pass any object with the same fields. +""" + +from __future__ import annotations + +import logging +from typing import Any + +from osmosis_ai.rollout.types import RolloutErrorCategory + +logger: logging.Logger = logging.getLogger(__name__) + +REDACTED = "[REDACTED]" +SENSITIVE_KEYS = frozenset( + { + "api_key", + "apikey", + "authorization", + "credential", + "credentials", + "password", + "secret", + "token", + } +) + + +def sensitive_key(key: str) -> bool: + normalized = key.lower().replace("-", "_") + return normalized in SENSITIVE_KEYS or normalized.endswith( + tuple(f"_{name}" for name in SENSITIVE_KEYS) + ) + + +def redact_secrets(value: Any, api_key: str | None = None) -> Any: + """Replace credential-bearing dict leaves and api-key substrings.""" + if isinstance(value, dict): + return { + key: REDACTED if sensitive_key(str(key)) else redact_secrets(child, api_key) + for key, child in value.items() + } + if isinstance(value, (list, tuple)): + return [redact_secrets(child, api_key) for child in value] + if api_key and isinstance(value, str) and api_key in value: + return REDACTED + return value + + +def span_seconds(info: Any) -> float | None: + if info is None or info.started_at is None or info.finished_at is None: + return None + return max(0.0, (info.finished_at - info.started_at).total_seconds()) + + +def trial_timings(result: Any) -> dict[str, float]: + """Per-phase durations from Harbor's own TimingInfo records.""" + if result is None: + return {} + spans = { + "environment_setup": result.environment_setup, + "agent_setup": result.agent_setup, + "agent": result.agent_execution, + "verifier": result.verifier, + "total": result, + } + return { + name: round(seconds, 2) + for name, info in spans.items() + if (seconds := span_seconds(info)) is not None + } + + +def failure_phase(result: Any) -> str: + """The furthest phase the trial reached; a failure happened there.""" + if result is None: + return "setup" + for name, info in ( + ("verifier", result.verifier), + ("agent", result.agent_execution), + ("agent_setup", result.agent_setup), + ("environment_setup", result.environment_setup), + ): + if info is not None: + return name + return "setup" + + +def trial_metrics(result: Any) -> dict[str, Any]: + """Token and cost totals Harbor accumulated across the trial.""" + if result is None: + return {} + try: + input_tokens, cached_tokens, output_tokens, cost_usd = ( + result.compute_token_cost_totals() + ) + except Exception: + logger.warning("Failed to read token totals from trial result", exc_info=True) + return {} + values = { + "input_tokens": input_tokens, + "cached_tokens": cached_tokens, + "output_tokens": output_tokens, + "cost_usd": cost_usd, + } + return {key: value for key, value in values.items() if value is not None} + + +def diagnostic_payload( + *, + phase: str, + category: RolloutErrorCategory | None, + exception_type: str | None, + timings: dict[str, float], +) -> dict[str, Any]: + return { + "backend": "harbor-v2", + "phase": phase, + "harbor_exception_type": exception_type, + "category": category.value if category else None, + "timings_sec": timings, + } diff --git a/osmosis_ai/rollout/backend/harbor/harness_agent.py b/osmosis_ai/rollout/backend/harbor/harness_agent.py new file mode 100644 index 00000000..fe833a79 --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/harness_agent.py @@ -0,0 +1,81 @@ +"""Harbor installed agent that runs a bundled workflow. + +install() bootstraps uv if the image lacks it, then installs the bundle wheel, +so any task image with python works unmodified. run() ships the ContainerInput +into the container and executes the bundle's agent console script; results come +back through the ContainerResult file. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from harbor.agents.installed.base import BaseInstalledAgent +from harbor.models.agent.context import AgentContext +from harbor.models.trial.paths import EnvironmentPaths + +from osmosis_ai.rollout.backend.harbor.tasks import ( + SDK_UV, + SDK_VENV, + venv_or_fallback_script, +) +from osmosis_ai.rollout.container.files import ( + INPUT_FILENAME, + RESULT_FILENAME, + ContainerInput, +) + + +class OsmosisHarnessInstalledAgent(BaseInstalledAgent): + def __init__( + self, + logs_dir: Path, + *args: Any, + bundle_path: str, + agent_script: str, + input_path: str, + **kwargs: Any, + ): + super().__init__(logs_dir, *args, **kwargs) + self.bundle_path = Path(bundle_path) + self.agent_script = agent_script + self.input_path = Path(input_path) + + @staticmethod + def name() -> str: + return "osmosis-harness-agent" + + async def install(self, environment: Any) -> None: + wheel = f"/tmp/{self.bundle_path.name}" + await environment.upload_file(self.bundle_path, wheel) + await self.exec_as_agent( + environment, + f"if [ -x {SDK_VENV}/bin/python ]; then " + f"{SDK_UV} pip install --python {SDK_VENV}/bin/python --no-deps {wheel}; " + f"else command -v uv >/dev/null || python3 -m pip install --quiet uv; " + f"uv pip install --system {wheel}; fi", + ) + + async def run(self, instruction: Any, environment: Any, context: Any) -> None: + container_input = ContainerInput.read(self.input_path) + if not container_input.prompt: + container_input.prompt = [{"role": "user", "content": instruction}] + + host_input = self.logs_dir / INPUT_FILENAME + container_input.write(host_input) + if not environment.capabilities.mounted: + agent_dir = EnvironmentPaths.for_os(environment.os).agent_dir + await environment.upload_file( + host_input, (agent_dir / INPUT_FILENAME).as_posix() + ) + + await self.exec_as_agent( + environment, venv_or_fallback_script(self.agent_script) + ) + + def populate_context_post_run(self, context: AgentContext) -> None: + result_path = self.logs_dir / RESULT_FILENAME + if result_path.exists(): + context.metadata = json.loads(result_path.read_text()) diff --git a/osmosis_ai/rollout/backend/harbor/native_agents.py b/osmosis_ai/rollout/backend/harbor/native_agents.py new file mode 100644 index 00000000..19a3a7cb --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/native_agents.py @@ -0,0 +1,68 @@ +"""Registered native Harbor agents and how each receives the model endpoint.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal + +from harbor.models.trial.config import AgentConfig as HarborAgentConfig + + +@dataclass(frozen=True) +class NativeAgentBinding: + wiring: Literal["env", "kwargs", "none"] + trainable: bool = True + env: dict[str, str] = field(default_factory=dict) + kwargs: dict[str, Any] = field(default_factory=dict) + + +NATIVE_AGENTS: dict[str, NativeAgentBinding] = { + "terminus-2": NativeAgentBinding( + wiring="kwargs", + # Summarization rewrites the running context, which forks the token + # trajectory RL training needs to stay append-only. + kwargs={"enable_summarize": False}, + ), + "mini-swe-agent": NativeAgentBinding( + wiring="env", + env={"MSWEA_COST_TRACKING": "ignore_errors"}, + ), + # Runs the task's reference solution with no model traffic; validates + # datasets and verifiers, never produces training data. + "oracle": NativeAgentBinding(wiring="none", trainable=False), +} + + +def native_binding(agent: Any) -> NativeAgentBinding | None: + """The binding for a registered native agent name; None for workflow agents.""" + if not isinstance(agent, str) or ":" in agent: + return None + binding = NATIVE_AGENTS.get(agent) + if binding is None: + raise ValueError( + f"unknown native agent {agent!r}; registered: {sorted(NATIVE_AGENTS)}" + ) + return binding + + +def native_agent_config( + name: str, + binding: NativeAgentBinding, + model_name: str, + url: str, + api_key: str, +) -> HarborAgentConfig: + if binding.wiring == "none": + return HarborAgentConfig( + name=name, + model_name=model_name, + env=dict(binding.env), + kwargs=dict(binding.kwargs), + ) + if binding.wiring == "env": + env = {**binding.env, "OPENAI_API_BASE": url, "OPENAI_API_KEY": api_key} + return HarborAgentConfig( + name=name, model_name=model_name, env=env, kwargs=dict(binding.kwargs) + ) + kwargs = {**binding.kwargs, "api_base": url, "api_key": api_key} + return HarborAgentConfig(name=name, model_name=model_name, kwargs=kwargs) diff --git a/osmosis_ai/rollout/backend/harbor/tasks.py b/osmosis_ai/rollout/backend/harbor/tasks.py new file mode 100644 index 00000000..301cd308 --- /dev/null +++ b/osmosis_ai/rollout/backend/harbor/tasks.py @@ -0,0 +1,190 @@ +"""Harbor task directories. Host-side only — nothing here runs in a container. + +A HarborTask is one task directory on disk (instruction.md, environment/, +tests/). materialize() copies it into a per-rollout run directory and stages +that rollout's files; Harbor then builds the container from the copy. Harbor's +content-addressed image cache dedupes builds across identical copies. +""" + +from __future__ import annotations + +import json +import logging +import re +import shutil +from enum import StrEnum +from pathlib import Path +from typing import Any + +from harbor.models.task.id import GitTaskId, LocalTaskId, PackageTaskId + +from osmosis_ai.rollout.container.files import INPUT_FILENAME, ContainerInput + +logger: logging.Logger = logging.getLogger(__name__) + +# Harness dependencies install into this venv so the task's python is never touched. +SDK_UV = "/opt/osmosis/uv" +SDK_VENV = "/opt/osmosis/venv" +SDK_REQUIREMENTS_FILENAME = "osmosis-requirements.txt" + + +class TaskMode(StrEnum): + TEMPLATE = "template" + DATASET = "dataset" + + +def parse_task_ref( + ref: str, metadata: dict[str, Any] +) -> GitTaskId | LocalTaskId | PackageTaskId: + """metadata["harbor_task"]: a local path, git checkout, or package ref.""" + if git_url := metadata.get("git_url"): + if not metadata.get("git_commit_id"): + logger.warning( + "git task %r is unpinned; set metadata['git_commit_id'] to an " + "immutable commit sha", + ref, + ) + return GitTaskId( + git_url=git_url, + git_commit_id=metadata.get("git_commit_id"), + path=Path(ref), + ) + if ref.startswith((".", "/", "~")): + return LocalTaskId(path=Path(ref)) + name, _, version = ref.partition("@") + org, slash, task = name.partition("/") + if not slash: + raise ValueError( + f"harbor_task {ref!r} must be a local path (./, /, ~), a package " + "'org/name[@ref]', or a git checkout (set metadata['git_url'])" + ) + if (version or "latest") == "latest": + logger.warning( + "package task %r uses the mutable ref 'latest'; pin a sha256 digest", ref + ) + return PackageTaskId(org=org, name=task, ref=version or "latest") + + +def venv_or_fallback_install(wheel: str) -> str: + """Shell command installing *wheel* into the SDK venv when the image has + one, else into the system python.""" + return ( + f"if [ -x {SDK_VENV}/bin/python ]; then " + f"{SDK_UV} pip install --python {SDK_VENV}/bin/python --no-deps {wheel}; " + f"else uv pip install --system {wheel} || python3 -m pip install {wheel}; fi" + ) + + +def venv_or_fallback_script(script: str) -> str: + """Shell command running *script* from the SDK venv when present.""" + return ( + f"if [ -x {SDK_VENV}/bin/{script} ]; then {SDK_VENV}/bin/{script}; " + f"else {script}; fi" + ) + + +def patch_dockerfile_with_sdk(env_dir: Path, requirements: list[str]) -> None: + """Pre-install *requirements* into an isolated venv in the task's image. + + Appends to the final stage: a static uv binary, the requirements file, and + a venv at /opt/osmosis with its own managed python — the task's own + packages and runtime user are left untouched (USER root is scoped to the + install and the stage's original USER is restored). + """ + dockerfile = env_dir / "Dockerfile" + if not dockerfile.is_file(): + raise ValueError(f"cannot patch Dockerfile: none found in {env_dir}") + lines = dockerfile.read_text().splitlines() + stage_start = max( + (i for i, line in enumerate(lines) if re.match(r"\s*FROM\s", line, re.I)), + default=0, + ) + original_user = next( + ( + line.strip() + for line in reversed(lines[stage_start:]) + if re.match(r"\s*USER\s", line, re.I) + ), + None, + ) + + (env_dir / SDK_REQUIREMENTS_FILENAME).write_text("\n".join(requirements) + "\n") + ignore = env_dir / ".dockerignore" + if ignore.is_file(): + ignore.write_text( + ignore.read_text().rstrip() + f"\n!{SDK_REQUIREMENTS_FILENAME}\n" + ) + + block = [ + "", + "USER root", + f"COPY --from=ghcr.io/astral-sh/uv:latest /uv {SDK_UV}", + f"COPY {SDK_REQUIREMENTS_FILENAME} /opt/osmosis/requirements.txt", + f"RUN {SDK_UV} venv {SDK_VENV} --python 3.12 && " + f"{SDK_UV} pip install --python {SDK_VENV}/bin/python " + "-r /opt/osmosis/requirements.txt", + ] + if original_user: + block.append(original_user) + dockerfile.write_text("\n".join([*lines, *block]) + "\n") + + +class HarborTask: + def __init__(self, path: Path): + self.path = path.resolve() + if not self.path.is_dir(): + raise ValueError(f"harbor task directory not found: {self.path}") + + @classmethod + def from_dataset(cls, root: Path, task_id: str) -> HarborTask: + """Look up a task under *root* by id, rejecting path escapes.""" + root = root.resolve() + path = (root / task_id).resolve() + if not path.is_relative_to(root) or not path.is_dir(): + raise ValueError(f"unknown harbor task id: {task_id!r}") + return cls(path) + + def materialize( + self, + out_dir: Path, + container_input: ContainerInput, + grader_script: str | None = None, + grader_wheel: Path | None = None, + sdk_requirements: list[str] | None = None, + ) -> Path: + """Copy this task into *out_dir* and stage one rollout's files. + + With grader_wheel the generated test.sh installs the wheel first, so + grading works even when the agent phase installed nothing (native + Harbor agents); the container input ships in tests/ for the same reason. + With sdk_requirements the copied Dockerfile pre-installs them into an + isolated venv, so per-trial installs stop downloading dependencies. + """ + task_dir = out_dir / self.path.name + shutil.rmtree(task_dir, ignore_errors=True) + shutil.copytree(self.path, task_dir) + if sdk_requirements: + patch_dockerfile_with_sdk(task_dir / "environment", sdk_requirements) + + if container_input.prompt: + (task_dir / "instruction.md").write_text( + json.dumps(container_input.prompt, default=str) + ) + elif not (task_dir / "instruction.md").exists(): + raise ValueError(f"task {self.path.name} has no instruction and no prompt") + + container_input.write(task_dir / INPUT_FILENAME) + + test_sh = task_dir / "tests" / "test.sh" + if grader_script and not test_sh.exists(): + test_sh.parent.mkdir(parents=True, exist_ok=True) + lines = ["#!/bin/bash", "set -e"] + if grader_wheel is not None: + shutil.copy2(grader_wheel, test_sh.parent / grader_wheel.name) + container_input.write(test_sh.parent / INPUT_FILENAME) + lines.append(venv_or_fallback_install(f"/tests/{grader_wheel.name}")) + lines.append(venv_or_fallback_script(grader_script)) + test_sh.write_text("\n".join(lines) + "\n") + test_sh.chmod(0o755) + + return task_dir diff --git a/osmosis_ai/rollout/container/__init__.py b/osmosis_ai/rollout/container/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/osmosis_ai/rollout/container/files.py b/osmosis_ai/rollout/container/files.py new file mode 100644 index 00000000..51cec60f --- /dev/null +++ b/osmosis_ai/rollout/container/files.py @@ -0,0 +1,62 @@ +"""The two files exchanged with the rollout container. + +ContainerInput is staged into the container before the agent phase; +ContainerResult comes back when it ends. Both ends are SDK code — user +harnesses and native agents never touch these. The verifier reward file is +Harbor's own contract and has one writer here. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from pydantic import BaseModel + +from osmosis_ai.rollout.types import RolloutStatus +from osmosis_ai.rollout.types.output import AgentWorkflowOutput, Messages + +AGENT_LOGS_DIR = Path("/logs/agent") +VERIFIER_LOGS_DIR = Path("/logs/verifier") +INPUT_FILENAME = "container_input.json" +RESULT_FILENAME = "container_result.json" + + +class ContainerInput(BaseModel): + """Everything one agent run needs inside the container.""" + + version: int = 1 + rollout_id: str + prompt: Messages = [] + label: str | None = None + metadata: dict[str, Any] | None = None + chat_completions_url: str = "" + api_key: str | None = None + + @classmethod + def read(cls, path: Path) -> ContainerInput: + return cls.model_validate_json(path.read_text()) + + def write(self, path: Path) -> None: + path.write_text(self.model_dump_json()) + + +class ContainerResult(BaseModel): + """How the agent phase went: status, error, and the workflow's output.""" + + status: RolloutStatus + output: AgentWorkflowOutput | None = None + err_message: str | None = None + + @classmethod + def read(cls, path: Path) -> ContainerResult: + return cls.model_validate_json(path.read_text()) + + def write(self, path: Path) -> None: + path.write_text(self.model_dump_json()) + + +def write_reward(reward: float) -> None: + VERIFIER_LOGS_DIR.mkdir(parents=True, exist_ok=True) + (VERIFIER_LOGS_DIR / "reward.json").write_text(json.dumps({"reward": reward})) diff --git a/osmosis_ai/rollout/container/runner.py b/osmosis_ai/rollout/container/runner.py new file mode 100644 index 00000000..519e0fa1 --- /dev/null +++ b/osmosis_ai/rollout/container/runner.py @@ -0,0 +1,162 @@ +"""In-container entrypoints for bundled workflows and graders. + +The bundle's generated shim calls ``agent_main(WorkflowClass, config)`` and +``grader_main(GraderClass, config)`` directly — the class is bound at package +time, so nothing is resolved at runtime. Both read the ContainerInput staged +by the backend; the grader additionally reads the agent phase's +ContainerResult. +""" + +from __future__ import annotations + +import asyncio +import json +import shutil +import sys +import traceback +from pathlib import Path +from typing import Any + +from osmosis_ai.rollout.container.files import ( + AGENT_LOGS_DIR, + INPUT_FILENAME, + RESULT_FILENAME, + VERIFIER_LOGS_DIR, + ContainerInput, + ContainerResult, + write_reward, +) +from osmosis_ai.rollout.container.trajectories import messages_from_trajectory +from osmosis_ai.rollout.context import ( + AgentWorkflowContext, + GraderContext, + RolloutContext, +) +from osmosis_ai.rollout.types import RolloutSample, RolloutStatus +from osmosis_ai.rollout.types.output import AgentWorkflowOutput, coerce_output +from osmosis_ai.rollout.utils.file_artifacts import ( + GRADER_ARTIFACTS_SNAPSHOT_DIRNAME, + HARBOR_ARTIFACTS_DIR, + ArtifactFileState, + artifact_tree_state, + copy_artifact_tree, +) + + +async def run_agent(workflow_cls: Any, workflow_config: Any) -> ContainerResult: + container_input = ContainerInput.read(AGENT_LOGS_DIR / INPUT_FILENAME) + rollout_ctx = RolloutContext( + chat_completions_url=container_input.chat_completions_url, + api_key=container_input.api_key, + rollout_id=container_input.rollout_id, + ) + ctx = AgentWorkflowContext( + prompt=container_input.prompt, + config=workflow_config, + metadata=container_input.metadata, + artifacts_dir=HARBOR_ARTIFACTS_DIR, + ) + workflow = workflow_cls(workflow_config) + with rollout_ctx: + returned = await workflow.run(ctx) + output = coerce_output(returned) + if output is None: + sample = await rollout_ctx.get_sample() + if sample is not None: + output = AgentWorkflowOutput( + samples={"default": [dict(m) for m in sample.messages]}, + metrics=sample.metrics or {}, + ) + else: + output = AgentWorkflowOutput() + return ContainerResult(status=RolloutStatus.SUCCESS, output=output) + + +def agent_main(workflow_cls: Any, workflow_config: Any = None) -> None: + try: + result = asyncio.run(run_agent(workflow_cls, workflow_config)) + except Exception as e: + traceback.print_exc() + result = ContainerResult(status=RolloutStatus.FAILURE, err_message=str(e)) + result.write(AGENT_LOGS_DIR / RESULT_FILENAME) + print(f"Agent runner complete: status={result.status}") + + +def snapshot_grader_artifacts(baseline: ArtifactFileState) -> None: + """Stage grader-authored artifact changes where Harbor returns verifier files.""" + snapshot_dir = VERIFIER_LOGS_DIR / GRADER_ARTIFACTS_SNAPSHOT_DIRNAME + try: + shutil.rmtree(snapshot_dir, ignore_errors=True) + if not HARBOR_ARTIFACTS_DIR.is_dir(): + return + copied = copy_artifact_tree( + HARBOR_ARTIFACTS_DIR, + snapshot_dir, + destination_root=VERIFIER_LOGS_DIR, + baseline=baseline, + ) + if not copied: + shutil.rmtree(snapshot_dir, ignore_errors=True) + except Exception as e: + print(f"Failed to stage grader artifacts (best-effort): {e}", file=sys.stderr) + + +TESTS_DIR = Path("/tests") + + +def read_container_input() -> ContainerInput: + for directory in (AGENT_LOGS_DIR, TESTS_DIR): + path = directory / INPUT_FILENAME + if path.exists(): + return ContainerInput.read(path) + raise FileNotFoundError( + f"{INPUT_FILENAME} not found in {AGENT_LOGS_DIR} or {TESTS_DIR}" + ) + + +def load_messages() -> tuple[list[dict[str, Any]] | None, dict[str, float]]: + """Messages from the workflow result, else from the agent's trajectory.""" + result_path = AGENT_LOGS_DIR / RESULT_FILENAME + if result_path.exists(): + output = ContainerResult.read(result_path).output + if output is not None: + return output.primary_messages(), output.metrics + for name in sorted(AGENT_LOGS_DIR.glob("*trajectory*.json")): + try: + messages = messages_from_trajectory(json.loads(name.read_text())) + except (ValueError, OSError): + continue + if messages: + return messages, {} + return None, {} + + +def grader_main(grader_cls: Any, grader_config: Any = None) -> None: + container_input = read_container_input() + messages, metrics = load_messages() + if messages is None: + print("No messages from the agent phase, skipping grading") + write_reward(0.0) + return + + sample = RolloutSample(messages=messages, metrics=metrics) + try: + baseline = artifact_tree_state(HARBOR_ARTIFACTS_DIR) + except Exception: + baseline = {} + try: + ctx = GraderContext( + label=container_input.label, + sample=sample, + metadata=container_input.metadata, + artifacts_dir=HARBOR_ARTIFACTS_DIR, + ) + grader = grader_cls(grader_config) + asyncio.run(grader.grade(ctx)) + finally: + snapshot_grader_artifacts(baseline) + + if ctx.sample is None or ctx.sample.reward is None: + raise RuntimeError("Sample has no reward after grading") + write_reward(ctx.sample.reward) + print(f"Grading complete: reward={ctx.sample.reward}") diff --git a/osmosis_ai/rollout/container/trajectories.py b/osmosis_ai/rollout/container/trajectories.py new file mode 100644 index 00000000..1f4e24d8 --- /dev/null +++ b/osmosis_ai/rollout/container/trajectories.py @@ -0,0 +1,24 @@ +"""Extract chat messages from the trajectory documents agents leave behind. + +Shared by the in-container grader runner and the host-side backend, so it +must not import Harbor. +""" + +from __future__ import annotations + +from typing import Any + +ATIF_ROLE_BY_SOURCE = {"user": "user", "agent": "assistant", "system": "system"} + + +def messages_from_trajectory(document: dict[str, Any]) -> list[dict[str, Any]]: + """Accept ATIF (steps) or raw harness formats (a messages list).""" + if isinstance(document.get("messages"), list): + return document["messages"] + messages = [] + for step in document.get("steps") or []: + content = step.get("message") + if content is not None: + role = ATIF_ROLE_BY_SOURCE.get(step.get("source"), "tool") + messages.append({"role": role, "content": content}) + return messages diff --git a/osmosis_ai/rollout/types/output.py b/osmosis_ai/rollout/types/output.py new file mode 100644 index 00000000..010fda77 --- /dev/null +++ b/osmosis_ai/rollout/types/output.py @@ -0,0 +1,44 @@ +"""Return type for AgentWorkflow.run.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + +Messages = list[dict[str, Any]] + + +class AgentWorkflowOutput(BaseModel): + """What a workflow hands back: message histories plus optional measurements. + + ``samples`` maps a name to one agent's message history (multi-agent + workflows return several). ``info`` carries workflow-produced context for + the grader; the rollout request's ``metadata`` is a separate, input-side + field. + """ + + samples: dict[str, Messages] = Field(default_factory=dict) + metrics: dict[str, float] = Field(default_factory=dict) + info: dict[str, Any] = Field(default_factory=dict) + + def primary_messages(self) -> Messages | None: + if not self.samples: + return None + if "default" in self.samples: + return self.samples["default"] + return next(iter(self.samples.values())) + + +def coerce_output(value: Any) -> AgentWorkflowOutput | None: + """Normalize a run() return value; None means "use the fallback source".""" + if value is None: + return None + if isinstance(value, AgentWorkflowOutput): + return value + if isinstance(value, list): + return AgentWorkflowOutput(samples={"default": value}) + raise TypeError( + "run() must return AgentWorkflowOutput, a message list, or None; " + f"got {type(value).__name__}" + ) diff --git a/osmosis_ai/rollout/types/sample.py b/osmosis_ai/rollout/types/sample.py index 52957ba6..b63c7663 100644 --- a/osmosis_ai/rollout/types/sample.py +++ b/osmosis_ai/rollout/types/sample.py @@ -82,3 +82,5 @@ class ExecutionResult(BaseModel): sample: RolloutSample | None = None err_message: str | None = None err_category: RolloutErrorCategory | None = None + # Backend diagnostics (failure phase, timings); not part of the wire protocol. + extra_fields: dict[str, Any] | None = None diff --git a/pyproject.toml b/pyproject.toml index 1a3ed51c..922d0924 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -46,6 +46,7 @@ dependencies = [ # Requirement parsing for the submit preflight. Transitive via setuptools and # litellm, but declared because the CLI imports it directly. "packaging>=24.0", + "platformdirs>=4.0", # CLI framework. <0.27: cli/_click_compat.py imports typer._click internals. "typer>=0.26,<0.27", # Do not directly constrain typer-slim: versions <0.22 shared typer/* with diff --git a/tests/unit/rollout/test_harbor_backend_v2.py b/tests/unit/rollout/test_harbor_backend_v2.py new file mode 100644 index 00000000..733fba0c --- /dev/null +++ b/tests/unit/rollout/test_harbor_backend_v2.py @@ -0,0 +1,573 @@ +"""Bundle backend: contract round-trip, task materialization, trial config.""" + +import json +from pathlib import Path + +import pytest +from harbor.trial.queue import TrialQueue + +from osmosis_ai.packaging import build_bundle, inspect_bundle +from osmosis_ai.rollout.backend.harbor.backend_v2 import HarborBackendV2 +from osmosis_ai.rollout.backend.harbor.tasks import ( + SDK_REQUIREMENTS_FILENAME, + HarborTask, + TaskMode, + patch_dockerfile_with_sdk, + venv_or_fallback_script, +) +from osmosis_ai.rollout.container.files import ContainerInput, ContainerResult +from osmosis_ai.rollout.types import ExecutionRequest, RolloutStatus +from osmosis_ai.rollout.types.output import AgentWorkflowOutput, coerce_output + +PYPROJECT = """\ +[project] +name = "bench" +version = "0.1.0" +dependencies = [] + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +include = ["bench*"] +""" + + +@pytest.fixture(scope="module") +def bundle(tmp_path_factory): + code_dir = tmp_path_factory.mktemp("harness") / "project" + package = code_dir / "bench" + package.mkdir(parents=True) + (package / "__init__.py").touch() + (package / "solver.py").write_text("class W: pass\nclass G: pass\n") + (code_dir / "pyproject.toml").write_text(PYPROJECT) + return build_bundle( + code_dir, + workflow="bench.solver:W", + grader="bench.solver:G", + bundles_dir=tmp_path_factory.mktemp("bundles"), + ) + + +@pytest.fixture +def template_task(tmp_path): + task = tmp_path / "template-task" + (task / "environment").mkdir(parents=True) + (task / "environment" / "Dockerfile").write_text("FROM python:3.12-slim\n") + (task / "task.toml").write_text('[task]\nname = "template-task"\n') + return task + + +def request_for(prompt=None, metadata=None) -> ExecutionRequest: + return ExecutionRequest(id="r1", prompt=prompt or [], metadata=metadata) + + +class TestContract: + def test_spec_round_trip(self, tmp_path): + container_input = ContainerInput( + rollout_id="r1", + prompt=[{"role": "user", "content": "hi"}], + label="4", + metadata={"k": 1}, + chat_completions_url="http://x/v1", + api_key="secret", + ) + container_input.write(tmp_path / "input.json") + assert ContainerInput.read(tmp_path / "input.json") == container_input + + def test_result_round_trip(self, tmp_path): + output = AgentWorkflowOutput( + samples={"default": [{"role": "assistant", "content": "y"}]}, + metrics={"turns": 1.0}, + ) + result = ContainerResult(status=RolloutStatus.SUCCESS, output=output) + result.write(tmp_path / "result.json") + loaded = ContainerResult.read(tmp_path / "result.json") + assert loaded.status == RolloutStatus.SUCCESS + assert loaded.output.primary_messages() == output.primary_messages() + assert loaded.output.metrics == {"turns": 1.0} + + +class TestAgentWorkflowOutput: + def test_coerce_none_passes_through(self): + assert coerce_output(None) is None + + def test_coerce_messages_wraps_as_default(self): + messages = [{"role": "assistant", "content": "hi"}] + output = coerce_output(messages) + assert output.samples == {"default": messages} + assert output.primary_messages() == messages + + def test_coerce_output_object_passes_through(self): + output = AgentWorkflowOutput(samples={"solver": [], "critic": []}) + assert coerce_output(output) is output + + def test_coerce_rejects_other_types(self): + import pytest as pytest_module + + with pytest_module.raises(TypeError, match="run\\(\\) must return"): + coerce_output("a string") + + def test_primary_prefers_default_key(self): + default = [{"role": "assistant", "content": "d"}] + output = AgentWorkflowOutput(samples={"z": [], "default": default}) + assert output.primary_messages() == default + + def test_empty_output_has_no_primary(self): + assert AgentWorkflowOutput().primary_messages() is None + + +class TestHarborTask: + def test_template_writes_prompt_and_input(self, template_task, tmp_path): + prompt = [{"role": "user", "content": "solve"}] + container_input = ContainerInput(rollout_id="r1", prompt=prompt) + + task_dir = HarborTask(template_task).materialize( + tmp_path / "r1", container_input + ) + + assert json.loads((task_dir / "instruction.md").read_text()) == prompt + assert ContainerInput.read(task_dir / "container_input.json") == container_input + assert (task_dir / "environment" / "Dockerfile").exists() + + def test_without_prompt_or_instruction_rejected(self, template_task, tmp_path): + with pytest.raises(ValueError, match="no instruction"): + HarborTask(template_task).materialize( + tmp_path / "r1", ContainerInput(rollout_id="r1") + ) + + def test_grader_script_generates_test_sh(self, template_task, tmp_path): + prompt = [{"role": "user", "content": "x"}] + task_dir = HarborTask(template_task).materialize( + tmp_path / "r1", + ContainerInput(rollout_id="r1", prompt=prompt), + grader_script="bench-grade", + ) + assert "bench-grade" in (task_dir / "tests" / "test.sh").read_text() + + def test_task_native_tests_win(self, template_task, tmp_path): + (template_task / "tests").mkdir() + (template_task / "tests" / "test.sh").write_text("#!/bin/bash\nnative\n") + prompt = [{"role": "user", "content": "x"}] + task_dir = HarborTask(template_task).materialize( + tmp_path / "r1", + ContainerInput(rollout_id="r1", prompt=prompt), + grader_script="bench-grade", + ) + assert "native" in (task_dir / "tests" / "test.sh").read_text() + + def test_from_dataset_routes_by_task_id(self, tmp_path): + root = tmp_path / "dataset" + task = root / "task-a" + (task / "environment").mkdir(parents=True) + (task / "instruction.md").write_text("fix the bug") + + task_dir = HarborTask.from_dataset(root, "task-a").materialize( + tmp_path / "r1", ContainerInput(rollout_id="r1") + ) + assert (task_dir / "instruction.md").read_text() == "fix the bug" + + def test_from_dataset_rejects_traversal_and_unknown_ids(self, tmp_path): + root = tmp_path / "dataset" + root.mkdir() + with pytest.raises(ValueError, match="unknown harbor task id"): + HarborTask.from_dataset(root, "../../etc") + with pytest.raises(ValueError, match="unknown harbor task id"): + HarborTask.from_dataset(root, "nope") + + +class TestPatchDockerfileWithSdk: + def test_patch_appends_isolated_venv(self, tmp_path): + env = tmp_path / "environment" + env.mkdir() + (env / "Dockerfile").write_text( + "FROM builder AS build\nRUN make\n" + 'FROM python:3.12-slim\nUSER agent\nCMD ["bash"]\n' + ) + patch_dockerfile_with_sdk(env, ["pydantic>=2", "httpx"]) + + dockerfile = (env / "Dockerfile").read_text() + reqs = (env / SDK_REQUIREMENTS_FILENAME).read_text() + assert reqs == "pydantic>=2\nhttpx\n" + assert "uv venv /opt/osmosis/venv" in dockerfile + # USER root is scoped to the install; the stage's user is restored last + assert dockerfile.rstrip().endswith("USER agent") + assert dockerfile.index("USER root") < dockerfile.index("uv venv") + + def test_patch_without_final_user_adds_no_restore(self, tmp_path): + env = tmp_path / "environment" + env.mkdir() + (env / "Dockerfile").write_text("FROM python:3.12-slim\n") + patch_dockerfile_with_sdk(env, ["httpx"]) + assert (env / "Dockerfile").read_text().count("USER") == 1 + + def test_patch_negates_dockerignore(self, tmp_path): + env = tmp_path / "environment" + env.mkdir() + (env / "Dockerfile").write_text("FROM python:3.12-slim\n") + (env / ".dockerignore").write_text("*\n") + patch_dockerfile_with_sdk(env, ["httpx"]) + assert f"!{SDK_REQUIREMENTS_FILENAME}" in (env / ".dockerignore").read_text() + + def test_patch_requires_dockerfile(self, tmp_path): + with pytest.raises(ValueError, match="cannot patch Dockerfile"): + patch_dockerfile_with_sdk(tmp_path, ["httpx"]) + + def test_materialize_patches_dockerfile(self, template_task, tmp_path): + task_dir = HarborTask(template_task).materialize( + tmp_path / "r1", + ContainerInput(rollout_id="r1", prompt=[{"role": "user", "content": "x"}]), + sdk_requirements=["httpx"], + ) + assert "uv venv" in (task_dir / "environment" / "Dockerfile").read_text() + # the source task stays pristine + assert ( + "uv venv" not in (template_task / "environment" / "Dockerfile").read_text() + ) + + def test_bundle_requirements_skips_extras(self, bundle): + assert inspect_bundle(bundle).requirements == [] + + def test_backend_flag_requires_bundle(self, template_task): + with pytest.raises(ValueError, match="requires a bundle"): + HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + agent="mini-swe-agent", + patch_dockerfile_with_sdk=True, + ) + + def test_patch_defaults_on_with_bundle(self, bundle, template_task): + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + bundle=bundle, + ) + assert backend.sdk_requirements is not None + + def test_patch_defaults_off_without_bundle(self, template_task): + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + agent="mini-swe-agent", + ) + assert backend.sdk_requirements is None + + def test_patch_opt_out(self, bundle, template_task): + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + bundle=bundle, + patch_dockerfile_with_sdk=False, + ) + assert backend.sdk_requirements is None + + +class TestBundleBackend: + @pytest.fixture + def backend(self, bundle, template_task): + return HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + bundle=bundle, + ) + + def prepare(self, backend, request): + container_input = backend.build_input(request) + task_dir = backend.select_task(request).materialize( + backend.rollouts_dir / request.id, + container_input, + grader_script=backend.bundle.grader_script if backend.bundle else None, + ) + return task_dir, container_input + + def test_trial_config_wires_harness_agent(self, backend): + request = request_for([{"role": "user", "content": "x"}]) + task_dir, container_input = self.prepare(backend, request) + config = backend.build_trial_config(task_dir, request, container_input) + + assert config.agent.import_path.endswith(":OsmosisHarnessInstalledAgent") + assert config.agent.kwargs["agent_script"] == "bench-agent" + assert config.agent.kwargs["bundle_path"] == str(backend.bundle.wheel) + assert config.verifier.disable is False # generated test.sh enables it + + def test_verifier_disabled_without_tests(self, bundle, template_task): + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + bundle=bundle, + ) + request = request_for([{"role": "user", "content": "x"}]) + container_input = backend.build_input(request) + task_dir = backend.select_task(request).materialize( + backend.rollouts_dir / request.id, container_input + ) + config = backend.build_trial_config(task_dir, request, container_input) + assert config.verifier.disable is True + + def test_build_input_carries_request_fields(self, backend): + request = ExecutionRequest( + id="r9", + prompt=[{"role": "user", "content": "q"}], + label="42", + metadata={"harbor_task_id": "t"}, + ) + container_input = backend.build_input(request) + assert container_input.rollout_id == "r9" + assert container_input.label == "42" + assert container_input.metadata == {"harbor_task_id": "t"} + + def test_environment_config_cloned_per_trial(self, backend): + request = request_for([{"role": "user", "content": "x"}]) + task_dir, container_input = self.prepare(backend, request) + first = backend.build_trial_config(task_dir, request, container_input) + second = backend.build_trial_config(task_dir, request, container_input) + assert first.environment is not second.environment + assert first.environment is not backend.environment_config + + +class TestNativeAgents: + def backend_for(self, agent, template_task, **kwargs): + return HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + agent=agent, + **kwargs, + ) + + def test_unknown_native_agent_rejected(self, template_task): + with pytest.raises(ValueError, match="unknown native agent"): + self.backend_for("claude-code", template_task) + + def test_env_wired_agent_receives_endpoint(self, template_task): + backend = self.backend_for("mini-swe-agent", template_task) + container_input = ContainerInput( + rollout_id="r1", + chat_completions_url="http://trainer:30000/sessions/abc/v1", + api_key="k", + ) + config = backend.build_agent_config(template_task, container_input) + + assert config.name == "mini-swe-agent" + assert config.env["OPENAI_API_BASE"] == "http://trainer:30000/sessions/abc/v1" + assert config.env["OPENAI_API_KEY"] == "k" + assert config.env["MSWEA_COST_TRACKING"] == "ignore_errors" + + def test_kwargs_wired_agent_receives_endpoint(self, template_task): + backend = self.backend_for("terminus-2", template_task, model_name="openai/m") + container_input = ContainerInput( + rollout_id="r1", chat_completions_url="http://t/v1" + ) + config = backend.build_agent_config(template_task, container_input) + + assert config.name == "terminus-2" + assert config.model_name == "openai/m" + assert config.kwargs["api_base"] == "http://t/v1" + assert config.kwargs["enable_summarize"] is False + + def test_native_without_grader_needs_no_bundle(self, template_task): + backend = self.backend_for("mini-swe-agent", template_task) + assert backend.bundle is None + + def test_dataset_mode_keeps_task_instruction(self, tmp_path, bundle): + root = tmp_path / "dataset" + task = root / "task-a" + (task / "environment").mkdir(parents=True) + (task / "instruction.md").write_text("real instruction") + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=root, + task_mode=TaskMode.DATASET, + bundle=bundle, + ) + request = request_for( + [{"role": "user", "content": "row prompt"}], + metadata={"harbor_task_id": "task-a"}, + ) + container_input = backend.build_input(request) + assert container_input.prompt == [] + + task_dir = backend.select_task(request).materialize( + backend.rollouts_dir / request.id, container_input + ) + assert (task_dir / "instruction.md").read_text() == "real instruction" + + def test_grader_wheel_ships_in_tests_dir(self, template_task, tmp_path, bundle): + from osmosis_ai.packaging import inspect_bundle + + info = inspect_bundle(bundle) + prompt = [{"role": "user", "content": "x"}] + task_dir = HarborTask(template_task).materialize( + tmp_path / "r1", + ContainerInput(rollout_id="r1", prompt=prompt), + grader_script=info.grader_script, + grader_wheel=info.wheel, + ) + test_sh = (task_dir / "tests" / "test.sh").read_text() + assert f"pip install /tests/{info.wheel.name}" in test_sh + # the grader script runs from the SDK venv when the image has one + assert test_sh.rstrip().endswith(venv_or_fallback_script(info.grader_script)) + assert (task_dir / "tests" / info.wheel.name).exists() + assert (task_dir / "tests" / "container_input.json").exists() + + def test_oracle_binding_gets_no_endpoint(self, template_task): + backend = self.backend_for("oracle", template_task) + config = backend.build_agent_config( + template_task, + ContainerInput(rollout_id="r1", chat_completions_url="http://t/v1"), + ) + assert config.name == "oracle" + assert "OPENAI_API_BASE" not in config.env + assert "api_base" not in config.kwargs + + def test_harbor_model_metadata_overrides_model(self, template_task): + backend = self.backend_for( + "terminus-2", template_task, model_name="openai/default" + ) + config = backend.build_agent_config( + template_task, + ContainerInput( + rollout_id="r1", + chat_completions_url="http://t/v1", + metadata={"harbor_model": "openai/override"}, + ), + ) + assert config.model_name == "openai/override" + + +class TestDiagnostics: + def result_with_timings(self): + from datetime import UTC, datetime, timedelta + from types import SimpleNamespace + + from harbor.models.trial.result import TimingInfo + + start = datetime(2026, 1, 1, tzinfo=UTC) + + def span(offset, seconds): + return TimingInfo( + started_at=start + timedelta(seconds=offset), + finished_at=start + timedelta(seconds=offset + seconds), + ) + + return SimpleNamespace( + started_at=start, + finished_at=start + timedelta(seconds=100), + environment_setup=span(0, 80), + agent_setup=span(80, 10), + agent_execution=span(90, 6), + verifier=span(96, 4), + ) + + def test_trial_timings_reads_harbor_spans(self): + from osmosis_ai.rollout.backend.harbor.diagnostics import trial_timings + + assert trial_timings(self.result_with_timings()) == { + "environment_setup": 80.0, + "agent_setup": 10.0, + "agent": 6.0, + "verifier": 4.0, + "total": 100.0, + } + + def test_failure_phase_is_furthest_span_reached(self): + from osmosis_ai.rollout.backend.harbor.diagnostics import failure_phase + + result = self.result_with_timings() + assert failure_phase(result) == "verifier" + result.verifier = None + result.agent_execution = None + assert failure_phase(result) == "agent_setup" + assert failure_phase(None) == "setup" + + def test_redact_secrets_scrubs_keys_and_api_key(self): + from osmosis_ai.rollout.backend.harbor.diagnostics import redact_secrets + + redacted = redact_secrets( + { + "llm_kwargs": {"api_key": "sk-1", "extra": ["Bearer sk-1", "safe"]}, + "model": "gpt", + "session_token": "t0", + }, + api_key="sk-1", + ) + assert redacted == { + "llm_kwargs": {"api_key": "[REDACTED]", "extra": ["[REDACTED]", "safe"]}, + "model": "gpt", + "session_token": "[REDACTED]", + } + + +class TestTaskRefs: + def test_local_path_ref(self): + from harbor.models.task.id import LocalTaskId + + from osmosis_ai.rollout.backend.harbor.tasks import parse_task_ref + + assert parse_task_ref("./tasks/t1", {}) == LocalTaskId(path=Path("./tasks/t1")) + + def test_package_ref_with_version(self): + from harbor.models.task.id import PackageTaskId + + from osmosis_ai.rollout.backend.harbor.tasks import parse_task_ref + + assert parse_task_ref("laude/swe-bench@sha256:abc", {}) == PackageTaskId( + org="laude", name="swe-bench", ref="sha256:abc" + ) + + def test_git_ref_uses_metadata(self): + from harbor.models.task.id import GitTaskId + + from osmosis_ai.rollout.backend.harbor.tasks import parse_task_ref + + task_id = parse_task_ref( + "tasks/t1", + {"git_url": "https://github.com/org/tasks.git", "git_commit_id": "abc123"}, + ) + assert task_id == GitTaskId( + git_url="https://github.com/org/tasks.git", + git_commit_id="abc123", + path=Path("tasks/t1"), + ) + + def test_bare_name_rejected(self): + from osmosis_ai.rollout.backend.harbor.tasks import parse_task_ref + + with pytest.raises(ValueError, match="must be a local path"): + parse_task_ref("not-a-ref", {}) + + +class TestPrewarm: + def test_prewarm_config_is_install_only_without_credentials( + self, bundle, template_task + ): + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=template_task, + bundle=bundle, + agent_setup_timeout_sec=120, + ) + config = backend.prewarm_trial_config(HarborTask(template_task)) + + assert config.install_only is True + assert config.verifier.disable is True + assert config.trial_name.startswith("trial-prewarm-") + assert config.agent.override_setup_timeout_sec == 120 + container_input = ContainerInput.read( + Path(config.task.path) / "container_input.json" + ) + assert container_input.api_key is None + assert container_input.chat_completions_url in ("", None) + + async def test_dataset_prewarm_requires_task_ids(self, bundle, tmp_path): + root = tmp_path / "dataset" + root.mkdir() + backend = HarborBackendV2( + orchestrator=TrialQueue(n_concurrent=1), + tasks_dir=root, + task_mode=TaskMode.DATASET, + bundle=bundle, + ) + with pytest.raises(ValueError, match="requires task ids"): + await backend.prewarm() diff --git a/tests/unit/test_packaging.py b/tests/unit/test_packaging.py new file mode 100644 index 00000000..18ee4e8a --- /dev/null +++ b/tests/unit/test_packaging.py @@ -0,0 +1,157 @@ +"""Packaging: wheel build, caching, and bundle inspection.""" + +import zipfile + +import pytest + +from osmosis_ai.packaging import build_bundle, inspect_bundle + +PYPROJECT = """\ +[project] +name = "my-harness" +version = "0.1.0" +dependencies = ["httpx>=0.27"] + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +include = ["my_harness*"] +""" + + +@pytest.fixture +def project(tmp_path): + code_dir = tmp_path / "harness" + package = code_dir / "my_harness" + package.mkdir(parents=True) + (package / "__init__.py").touch() + (package / "solver.py").write_text("class MyWorkflow: pass\n") + (package / "grade.py").write_text("class MyGrader: pass\n") + (code_dir / "pyproject.toml").write_text(PYPROJECT) + return code_dir + + +def test_build_produces_scripts_and_keeps_deps(project, tmp_path): + wheel = build_bundle( + project, + workflow="my_harness.solver:MyWorkflow", + grader="my_harness.grade:MyGrader", + bundles_dir=tmp_path / "bundles", + ) + with zipfile.ZipFile(wheel) as archive: + names = archive.namelist() + metadata = archive.read( + next(n for n in names if n.endswith("METADATA")) + ).decode() + shim = archive.read("my_harness/bundle_main.py").decode() + + assert "my_harness/solver.py" in names + assert "Requires-Dist: httpx>=0.27" in metadata + assert "from my_harness.solver import MyWorkflow" in shim + assert "runner.agent_main(MyWorkflow, None)" in shim + + info = inspect_bundle(wheel) + assert info.agent_script == "my-harness-agent" + assert info.grader_script == "my-harness-grade" + + +def test_deps_override_same_named_pyproject_entries(project, tmp_path): + wheel = build_bundle( + project, + workflow="my_harness.solver:MyWorkflow", + deps=["httpx @ https://example.com/httpx.tar.gz"], + bundles_dir=tmp_path / "bundles", + ) + with zipfile.ZipFile(wheel) as archive: + metadata = archive.read( + next(n for n in archive.namelist() if n.endswith("METADATA")) + ).decode() + assert "Requires-Dist: httpx>=0.27" not in metadata + assert "Requires-Dist: httpx @ https://example.com/httpx.tar.gz" in metadata + + +def test_requirements_keep_markers_drop_extras(tmp_path): + code_dir = tmp_path / "harness" + package = code_dir / "my_harness" + package.mkdir(parents=True) + (package / "__init__.py").touch() + (package / "solver.py").write_text("class W: pass\n") + (code_dir / "pyproject.toml").write_text( + """\ +[project] +name = "my-harness" +version = "0.1.0" +dependencies = ["httpx>=0.27", "tomli; python_version < '3.11'"] + +[project.optional-dependencies] +dev = ["pytest"] + +[build-system] +requires = ["setuptools>=61.0"] +build-backend = "setuptools.build_meta" + +[tool.setuptools.packages.find] +include = ["my_harness*"] +""" + ) + wheel = build_bundle( + code_dir, workflow="my_harness.solver:W", bundles_dir=tmp_path / "bundles" + ) + requirements = inspect_bundle(wheel).requirements + assert "httpx>=0.27" in requirements + assert any(r.startswith("tomli") for r in requirements) + assert not any("pytest" in r for r in requirements) + + +def test_grader_optional(project, tmp_path): + wheel = build_bundle( + project, + workflow="my_harness.solver:MyWorkflow", + bundles_dir=tmp_path / "bundles", + ) + info = inspect_bundle(wheel) + assert info.grader_script is None + + +def test_cache_hits_until_source_changes(project, tmp_path): + bundles_dir = tmp_path / "bundles" + kwargs = dict(workflow="my_harness.solver:MyWorkflow", bundles_dir=bundles_dir) + + first = build_bundle(project, **kwargs) + first_mtime = first.stat().st_mtime_ns + assert build_bundle(project, **kwargs).stat().st_mtime_ns == first_mtime + + (project / "my_harness" / "solver.py").write_text( + "class MyWorkflow:\n changed = True\n" + ) + assert build_bundle(project, **kwargs).stat().st_mtime_ns != first_mtime + + +def test_rejects_bad_refs_and_missing_pyproject(project, tmp_path): + with pytest.raises(ValueError, match="module:attr"): + build_bundle(project, workflow="not-a-ref", bundles_dir=tmp_path / "b") + with pytest.raises(ValueError, match=r"pyproject\.toml"): + build_bundle(tmp_path / "nowhere", workflow="a:B", bundles_dir=tmp_path / "b") + + +def test_project_dir_for_locates_bench_harness(): + import sys + from pathlib import Path + + from osmosis_ai.packaging import project_dir_for + + harness_dir = ( + Path(__file__).parents[2] + / "benchmarks" + / "container_lifecycle" + / "bench_harness" + ) + sys.path.insert(0, str(harness_dir)) + try: + from bench_harness.solver import BenchWorkflow + + assert project_dir_for(BenchWorkflow) == harness_dir + finally: + sys.path.remove(str(harness_dir)) diff --git a/uv.lock b/uv.lock index dd3d5588..dfaf9141 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.12" resolution-markers = [ "python_full_version >= '3.15' and sys_platform == 'win32'", @@ -1944,6 +1944,7 @@ dependencies = [ { name = "openai-agents", extra = ["litellm"] }, { name = "orjson" }, { name = "packaging" }, + { name = "platformdirs" }, { name = "pydantic" }, { name = "python-dotenv" }, { name = "questionary" }, @@ -1998,6 +1999,7 @@ requires-dist = [ { name = "osmosis-ai", extras = ["server"], marker = "extra == 'dev'" }, { name = "osmosis-ai", extras = ["server"], marker = "extra == 'full'" }, { name = "packaging", specifier = ">=24.0" }, + { name = "platformdirs", specifier = ">=4.0" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=4.6.0,<5.0.0" }, { name = "pyarrow", marker = "extra == 'platform'", specifier = ">=23.0.1" }, { name = "pydantic", specifier = ">=2.0.0,<3.0.0" },