diff --git a/docs/rollout-sdk.md b/docs/rollout-sdk.md index 0b90d82d..e8b33bc1 100644 --- a/docs/rollout-sdk.md +++ b/docs/rollout-sdk.md @@ -75,6 +75,28 @@ samples = await rollout_ctx.get_samples() # async -> {name: RolloutSa `RolloutSample` ([types/sample.py](../osmosis_ai/rollout/types/sample.py)) fields: `id`, `messages`, `label`, `reward`, `remove_sample`, `metrics`, `extra_fields`. +## Artifacts + +[../osmosis_ai/rollout/utils/file_artifacts.py](../osmosis_ai/rollout/utils/file_artifacts.py) + +Artifacts are files produced by a rollout — logs, traces, generated outputs, screenshots, or other large/binary data that does not belong in the sample. + +There is one rule: write files under `ctx.artifacts_dir` (available on both `AgentWorkflowContext` and `GraderContext`). It normally exists before your code runs, but is `Path | None` — it can be `None` when the host directory could not be created — so guard before use. In Harbor-backed rollouts it is the sandbox's `/logs/artifacts/`; `LocalBackend` provides an isolated per-rollout directory. If a file is produced elsewhere, copy it in (`shutil.copy2(path, ctx.artifacts_dir / "name")`). + +```python +import json + +async def grade(self, ctx: GraderContext) -> Any: + if ctx.artifacts_dir: + (ctx.artifacts_dir / "trace.json").write_text( + json.dumps({"score_reason": "matched rubric"}) + ) + for sample_id in ctx.get_samples(): + ctx.set_sample_reward(sample_id, 1.0) +``` + +After each rollout the artifacts land on the host under `~/.osmosis//artifacts/`, mirroring Harbor's collected-trial layout (user files sit at `.../artifacts/logs/artifacts/`). Artifact handling is best-effort and never affects rewards or rollout status. + ## Configs [../osmosis_ai/rollout/types/config.py](../osmosis_ai/rollout/types/config.py) diff --git a/osmosis_ai/rollout/backend/harbor/agent_runner.py b/osmosis_ai/rollout/backend/harbor/agent_runner.py index b0b5476e..9c0e085e 100644 --- a/osmosis_ai/rollout/backend/harbor/agent_runner.py +++ b/osmosis_ai/rollout/backend/harbor/agent_runner.py @@ -14,6 +14,7 @@ from typing import Any from osmosis_ai.rollout.context import AgentWorkflowContext, RolloutContext +from osmosis_ai.rollout.utils.file_artifacts import HARBOR_ARTIFACTS_DIR from osmosis_ai.rollout.utils.imports import resolve_object AGENT_LOGS_DIR = Path("/logs/agent") @@ -58,6 +59,7 @@ async def run_workflow( prompt=prompt, config=workflow_config, metadata=config.get("metadata"), + artifacts_dir=HARBOR_ARTIFACTS_DIR, ) meta: dict[str, Any] = {"id": config.get("id", "")} diff --git a/osmosis_ai/rollout/backend/harbor/backend.py b/osmosis_ai/rollout/backend/harbor/backend.py index 80f5efe3..eb3f9d4a 100644 --- a/osmosis_ai/rollout/backend/harbor/backend.py +++ b/osmosis_ai/rollout/backend/harbor/backend.py @@ -38,6 +38,7 @@ RolloutSample, RolloutStatus, ) +from osmosis_ai.rollout.utils.file_artifacts import default_artifact_root from osmosis_ai.rollout.utils.imports import to_import_path from osmosis_ai.rollout.utils.rewards import validate_samples_have_rewards @@ -142,6 +143,7 @@ def __init__( self.trials_dir: Path = ( trials_dir if trials_dir != Path("trials") else self.root_dir / "trials" ) + self.artifact_root: Path = default_artifact_root() self.pending: dict[str, PendingTrial] = {} self.prebuilt_image_tag: str | None = None @@ -444,6 +446,30 @@ async def on_verification_start(self, event: TrialHookEvent) -> None: pending.workflow_complete_called = True await pending.on_workflow_complete(result) + def _relocate_trial_artifacts(self, rollout_id: str, *, move: bool) -> bool: + """Persist collected artifacts to the artifact root; ``move`` before + cleanup, else copy. Best-effort: returns ``False`` to keep the source.""" + source_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" / "artifacts" + if not source_dir.is_dir(): + return True + dest_dir = self.artifact_root / rollout_id / "artifacts" + try: + if dest_dir.exists(): + shutil.rmtree(dest_dir) + dest_dir.parent.mkdir(parents=True, exist_ok=True) + if move: + shutil.move(source_dir, dest_dir) + else: + shutil.copytree(source_dir, dest_dir, symlinks=False) + except OSError: + logger.warning( + "Failed to relocate trial artifacts for rollout %s (best-effort)", + rollout_id, + exc_info=True, + ) + return False + return True + async def on_trial_end(self, event: TrialHookEvent) -> None: rollout_id = parse_rollout_id(event) pending = self.pending.pop(rollout_id, None) @@ -451,6 +477,14 @@ async def on_trial_end(self, event: TrialHookEvent) -> None: logger.error("No pending trial found for rollout %s", rollout_id) return + # Evacuate artifacts first; delete the source only if it succeeds. + delete_trial = bool( + self.cleanup_successful_trials + and event.result + and not event.result.exception_info + ) + relocated = self._relocate_trial_artifacts(rollout_id, move=delete_trial) + if not pending.workflow_complete_called: if event.result and event.result.exception_info: err = event.result.exception_info @@ -518,9 +552,9 @@ async def on_trial_end(self, event: TrialHookEvent) -> None: rollout_dir = self.rollouts_dir / rollout_id shutil.rmtree(rollout_dir, ignore_errors=True) - 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) + if delete_trial and relocated: + trial_dir = self.trials_dir / f"{TRIAL_NAME_PREFIX}{rollout_id}" + shutil.rmtree(trial_dir, ignore_errors=True) if not pending.done.done(): pending.done.set_result(None) diff --git a/osmosis_ai/rollout/backend/harbor/grader_runner.py b/osmosis_ai/rollout/backend/harbor/grader_runner.py index 06a3499d..a41c358e 100644 --- a/osmosis_ai/rollout/backend/harbor/grader_runner.py +++ b/osmosis_ai/rollout/backend/harbor/grader_runner.py @@ -14,6 +14,7 @@ from osmosis_ai.rollout.context import GraderContext from osmosis_ai.rollout.types import RolloutSample +from osmosis_ai.rollout.utils.file_artifacts import HARBOR_ARTIFACTS_DIR from osmosis_ai.rollout.utils.imports import resolve_object VERIFIER_LOGS_DIR = Path("/logs/verifier") @@ -71,7 +72,12 @@ def main() -> None: resolve_object(config["grader_config"]) if "grader_config" in config else None ) - ctx = GraderContext(label=label, samples=samples, metadata=metadata) + ctx = GraderContext( + label=label, + samples=samples, + metadata=metadata, + artifacts_dir=HARBOR_ARTIFACTS_DIR, + ) grader = grader_cls(grader_config) asyncio.run(grader.grade(ctx)) diff --git a/osmosis_ai/rollout/backend/local/backend.py b/osmosis_ai/rollout/backend/local/backend.py index 60b516b7..7c2d431f 100644 --- a/osmosis_ai/rollout/backend/local/backend.py +++ b/osmosis_ai/rollout/backend/local/backend.py @@ -1,6 +1,7 @@ import copy import logging import traceback +from pathlib import Path from typing import Any from osmosis_ai.rollout.agent_workflow import AgentWorkflow @@ -21,6 +22,10 @@ RolloutStatus, ) from osmosis_ai.rollout.utils.concurrency import ConcurrencyLimiter +from osmosis_ai.rollout.utils.file_artifacts import ( + HARBOR_COLLECTED_ARTIFACTS_SUBPATH, + default_artifact_root, +) from osmosis_ai.rollout.utils.imports import resolve_object from osmosis_ai.rollout.utils.rewards import validate_samples_have_rewards @@ -56,6 +61,8 @@ def __init__( max_concurrent=max_concurrent ) + self.artifact_root: Path = default_artifact_root() + @property def max_concurrency(self) -> int: return self.limiter.max_concurrent or 0 @@ -73,7 +80,9 @@ async def execute( on_grader_complete: ResultCallback | None = None, ) -> None: async with self.limiter.acquire(): - result = await self.run_workflow(request) + artifacts_dir = self._make_artifacts_dir(request.id) + + result = await self.run_workflow(request, artifacts_dir) await on_workflow_complete(result) if not on_grader_complete: @@ -84,17 +93,39 @@ async def execute( and (request.label is not None or request.metadata is not None) and result.status == RolloutStatus.SUCCESS ): - graded = await self.run_grader(request, result) + graded = await self.run_grader(request, result, artifacts_dir) await on_grader_complete(graded) else: await on_grader_complete(ExecutionResult(status=RolloutStatus.FAILURE)) - async def run_workflow(self, request: ExecutionRequest) -> ExecutionResult: + def _make_artifacts_dir(self, rollout_id: str) -> Path | None: + """Create the rollout's artifacts dir; ``None`` if it can't be created.""" + artifacts_dir = ( + self.artifact_root + / rollout_id + / "artifacts" + / HARBOR_COLLECTED_ARTIFACTS_SUBPATH + ) + try: + artifacts_dir.mkdir(parents=True, exist_ok=True) + except OSError: + logger.warning( + "Failed to create artifacts dir for rollout %s (best-effort)", + rollout_id, + exc_info=True, + ) + return None + return artifacts_dir + + async def run_workflow( + self, request: ExecutionRequest, artifacts_dir: Path | None = None + ) -> ExecutionResult: config = copy.deepcopy(self.workflow_config) ctx = AgentWorkflowContext( prompt=request.prompt, config=config, metadata=request.metadata, + artifacts_dir=artifacts_dir, ) rollout_ctx = get_rollout_context() @@ -119,7 +150,10 @@ async def run_workflow(self, request: ExecutionRequest) -> ExecutionResult: ) async def run_grader( - self, request: ExecutionRequest, result: ExecutionResult + self, + request: ExecutionRequest, + result: ExecutionResult, + artifacts_dir: Path | None = None, ) -> ExecutionResult: if not self.grader_cls: return result @@ -128,6 +162,7 @@ async def run_grader( label=request.label, samples=result.samples, metadata=request.metadata, + artifacts_dir=artifacts_dir, ) try: grader = self.grader_cls(self.grader_config) diff --git a/osmosis_ai/rollout/context.py b/osmosis_ai/rollout/context.py index 2a287684..6ba9c119 100644 --- a/osmosis_ai/rollout/context.py +++ b/osmosis_ai/rollout/context.py @@ -2,6 +2,7 @@ from abc import ABC, abstractmethod from contextvars import ContextVar from dataclasses import dataclass, field +from pathlib import Path from typing import Any from osmosis_ai.rollout.types import ( @@ -97,6 +98,7 @@ class GraderContext: samples: dict[str, RolloutSample] = field(default_factory=dict) project_path: str | None = None metadata: dict[str, Any] | None = None + artifacts_dir: Path | None = None def get_samples(self) -> dict[str, RolloutSample]: return self.samples @@ -112,16 +114,19 @@ class AgentWorkflowContext[TConfig: AgentWorkflowConfig]: prompt: list[dict[str, Any]] config: TConfig | None = None metadata: dict[str, Any] | None = None + artifacts_dir: Path | None = None def __init__( self, prompt: list[dict[str, Any]], config: TConfig | None = None, metadata: dict[str, Any] | None = None, + artifacts_dir: Path | None = None, ): self.prompt = prompt self.config = config self.metadata = metadata + self.artifacts_dir = artifacts_dir @dataclass @@ -143,6 +148,12 @@ def __init__( config: TConfig, environment: Any = None, metadata: dict[str, Any] | None = None, + artifacts_dir: Path | None = None, ): - super().__init__(prompt=prompt, config=config, metadata=metadata) + super().__init__( + prompt=prompt, + config=config, + metadata=metadata, + artifacts_dir=artifacts_dir, + ) self.environment = environment diff --git a/osmosis_ai/rollout/utils/file_artifacts.py b/osmosis_ai/rollout/utils/file_artifacts.py new file mode 100644 index 00000000..8ce4f22b --- /dev/null +++ b/osmosis_ai/rollout/utils/file_artifacts.py @@ -0,0 +1,21 @@ +"""Rollout file artifacts: user code writes under ``ctx.artifacts_dir``.""" + +from __future__ import annotations + +from pathlib import Path + +# Harbor's auto-collected convention directory inside the sandbox. +HARBOR_ARTIFACTS_DIR = Path("/logs/artifacts") + +# Host subpath where Harbor lands it; LocalBackend reuses it to match. +HARBOR_COLLECTED_ARTIFACTS_SUBPATH = HARBOR_ARTIFACTS_DIR.relative_to("/") + + +def default_artifact_root() -> Path: + """Host directory that holds one subdirectory per rollout id. + + Each rollout's outputs live at ``//`` — file + artifacts under ``artifacts/`` and the archived trajectory as + ``trajectory.json`` — which the platform persists to durable storage. + """ + return Path.home() / ".osmosis" diff --git a/osmosis_ai/templates/_scaffolds/rollout/main.py.tpl b/osmosis_ai/templates/_scaffolds/rollout/main.py.tpl index c47d772a..714e2c48 100644 --- a/osmosis_ai/templates/_scaffolds/rollout/main.py.tpl +++ b/osmosis_ai/templates/_scaffolds/rollout/main.py.tpl @@ -50,6 +50,11 @@ class MyGrader(Grader): # ctx.get_samples() returns {name: RolloutSample} for sources registered # by the workflow. Assign a scalar reward per sample with: # ctx.set_sample_reward(name, value) + # + # Optional: write file artifacts (artifacts_dir may be None, so guard): + # import json + # if ctx.artifacts_dir: + # (ctx.artifacts_dir / "debug.json").write_text(json.dumps({...})) for sample_id in ctx.get_samples(): ctx.set_sample_reward(sample_id, 0.0) diff --git a/tests/unit/rollout/test_context.py b/tests/unit/rollout/test_context.py index 7bee0796..22dee2fe 100644 --- a/tests/unit/rollout/test_context.py +++ b/tests/unit/rollout/test_context.py @@ -1,5 +1,6 @@ """Tests for osmosis_ai.rollout.context.""" +from pathlib import Path from unittest.mock import MagicMock import pytest @@ -210,3 +211,13 @@ def test_metadata_threaded_through_super_init(self): ) assert ctx.metadata == metadata assert ctx.environment is env + + def test_artifacts_dir_threaded_through_super_init(self, tmp_path: Path): + cfg = AgentWorkflowConfig(name="harbor-test") + ctx = HarborAgentWorkflowContext( + prompt=[{"role": "user", "content": "hi"}], + config=cfg, + environment=MagicMock(), + artifacts_dir=tmp_path, + ) + assert ctx.artifacts_dir == tmp_path diff --git a/tests/unit/rollout/test_harbor_backend.py b/tests/unit/rollout/test_harbor_backend.py index 7a908148..42222566 100644 --- a/tests/unit/rollout/test_harbor_backend.py +++ b/tests/unit/rollout/test_harbor_backend.py @@ -18,6 +18,7 @@ RolloutSample, RolloutStatus, ) +from osmosis_ai.rollout.utils.file_artifacts import HARBOR_ARTIFACTS_DIR # Module-level captures so importlib can resolve the workflow/grader by # ":" and round-trip them through the runners. @@ -38,6 +39,7 @@ async def run(self, ctx: AgentWorkflowContext) -> Any: from osmosis_ai.rollout.context import get_rollout_context _AGENT_CAPTURE["metadata"] = ctx.metadata + _AGENT_CAPTURE["artifacts_dir"] = ctx.artifacts_dir rollout_ctx = get_rollout_context() if rollout_ctx: rollout_ctx.register_sample_source( @@ -50,6 +52,7 @@ class MetadataCapturingGrader(Grader): async def grade(self, ctx: GraderContext) -> Any: _GRADER_CAPTURE["metadata"] = ctx.metadata _GRADER_CAPTURE["label"] = ctx.label + _GRADER_CAPTURE["artifacts_dir"] = ctx.artifacts_dir for sample_id in ctx.get_samples(): ctx.set_sample_reward(sample_id, 1.0) @@ -115,6 +118,100 @@ async def test_empty_verifier_rewards_logs_and_returns_validation_failure( assert result.err_category == RolloutErrorCategory.VALIDATION_ERROR assert "Harbor verifier returned empty rewards for rollout r1" in caplog.text + def _make_trial_end_backend(self, tmp_path, *, cleanup: bool) -> HarborBackend: + backend = HarborBackend.__new__(HarborBackend) + backend.pending = {} + backend.cleanup_successful_trials = cleanup + backend.trials_dir = tmp_path / "trials" + backend.rollouts_dir = tmp_path / "rollouts" + backend.artifact_root = tmp_path / "collected" + return backend + + def _seed_trial_artifacts(self, backend: HarborBackend) -> None: + trial_artifacts = backend.trials_dir / "trial-r1" / "artifacts" + convention = trial_artifacts / "logs" / "artifacts" + convention.mkdir(parents=True) + (convention / "output.txt").write_text("ok") + (trial_artifacts / "manifest.json").write_text("[]") + + @staticmethod + def _success_event() -> SimpleNamespace: + return SimpleNamespace( + config=SimpleNamespace(trial_name="trial-r1"), + result=SimpleNamespace( + agent_result=SimpleNamespace( + metadata={ + "status": "success", + "samples": { + "s1": RolloutSample(id="s1", messages=[]).model_dump() + }, + } + ), + verifier_result=SimpleNamespace(rewards={"s1": 1.0}), + exception_info=None, + ), + ) + + async def test_on_trial_end_relocates_artifacts_before_cleanup(self, tmp_path): + backend = self._make_trial_end_backend(tmp_path, cleanup=True) + self._seed_trial_artifacts(backend) + + on_grader = AsyncMock() + pending = PendingTrial(AsyncMock(), on_grader) + pending.workflow_complete_called = True + backend.pending["r1"] = pending + + await backend.on_trial_end(self._success_event()) + + result = on_grader.call_args.args[0] + assert result.status == RolloutStatus.SUCCESS + relocated = backend.artifact_root / "r1" / "artifacts" + assert (relocated / "logs" / "artifacts" / "output.txt").read_text() == "ok" + assert (relocated / "manifest.json").exists() + assert not (backend.trials_dir / "trial-r1").exists() + + async def test_on_trial_end_copies_artifacts_when_trial_kept(self, tmp_path): + backend = self._make_trial_end_backend(tmp_path, cleanup=False) + self._seed_trial_artifacts(backend) + + pending = PendingTrial(AsyncMock(), AsyncMock()) + pending.workflow_complete_called = True + backend.pending["r1"] = pending + + await backend.on_trial_end(self._success_event()) + + relocated = backend.artifact_root / "r1" / "artifacts" + assert (relocated / "logs" / "artifacts" / "output.txt").read_text() == "ok" + source = backend.trials_dir / "trial-r1" / "artifacts" + assert (source / "logs" / "artifacts" / "output.txt").exists() + + async def test_on_trial_end_keeps_trial_when_relocation_fails( + self, tmp_path, monkeypatch, caplog + ): + backend = self._make_trial_end_backend(tmp_path, cleanup=True) + self._seed_trial_artifacts(backend) + + def _boom(*_args, **_kwargs): + raise OSError("disk full") + + monkeypatch.setattr( + "osmosis_ai.rollout.backend.harbor.backend.shutil.move", _boom + ) + + pending = PendingTrial(AsyncMock(), AsyncMock()) + pending.workflow_complete_called = True + backend.pending["r1"] = pending + + with caplog.at_level( + logging.WARNING, logger="osmosis_ai.rollout.backend.harbor.backend" + ): + await backend.on_trial_end(self._success_event()) + + # Relocation failed: the source artifacts must survive cleanup. + source = backend.trials_dir / "trial-r1" / "artifacts" + assert (source / "logs" / "artifacts" / "output.txt").read_text() == "ok" + assert "Failed to relocate trial artifacts for rollout r1" in caplog.text + class TestBuildRolloutConfigMetadata: def test_metadata_written_when_present(self): @@ -162,6 +259,7 @@ async def test_metadata_surfaces_on_ctx(self, tmp_path, monkeypatch): assert meta["status"] == "success" assert _AGENT_CAPTURE["metadata"] == metadata + assert _AGENT_CAPTURE["artifacts_dir"] == HARBOR_ARTIFACTS_DIR class TestGraderRunnerRoundTrip: @@ -200,6 +298,7 @@ def test_grader_ctx_receives_metadata(self, tmp_path, monkeypatch): grader_runner.main() assert _GRADER_CAPTURE["metadata"] == metadata + assert _GRADER_CAPTURE["artifacts_dir"] == HARBOR_ARTIFACTS_DIR rewards = json.loads((verifier_dir / "reward.json").read_text()) assert rewards == {"sample-1": 1.0} diff --git a/tests/unit/rollout/test_local_backend.py b/tests/unit/rollout/test_local_backend.py index d1a5ab14..4c7a9df9 100644 --- a/tests/unit/rollout/test_local_backend.py +++ b/tests/unit/rollout/test_local_backend.py @@ -118,6 +118,73 @@ def test_health(self): assert h["status"] == "ok" assert "concurrency" in h + async def test_rollout_artifacts_land_under_rollout_id(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) # redirect ~/.osmosis + + class ArtifactWorkflow(StubWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> Any: + assert ctx.artifacts_dir is not None and ctx.artifacts_dir.is_dir() + (ctx.artifacts_dir / "trace.txt").write_text("trace") + await super().run(ctx) + + class WritingGrader(StubGrader): + async def grade(self, ctx: GraderContext) -> Any: + assert ctx.artifacts_dir is not None + (ctx.artifacts_dir / "grade_debug.json").write_text("{}") + await super().grade(ctx) + + backend = LocalBackend( + workflow=ArtifactWorkflow, + workflow_config=AgentWorkflowConfig(name="test"), + grader=WritingGrader, + ) + + request = ExecutionRequest( + id="r1", prompt=[{"role": "user", "content": "hi"}], label="x" + ) + await backend.execute(request, AsyncMock(), AsyncMock()) + + root = tmp_path / ".osmosis" + rollout_dir = root / "r1" / "artifacts" / "logs" / "artifacts" + assert (rollout_dir / "trace.txt").read_text() == "trace" + assert (rollout_dir / "grade_debug.json").read_text() == "{}" + + async def test_rollout_degrades_when_artifacts_dir_unavailable( + self, tmp_path, monkeypatch + ): + monkeypatch.setenv("HOME", str(tmp_path)) + + def _boom(self, *_args, **_kwargs): + raise OSError("read-only filesystem") + + monkeypatch.setattr("pathlib.Path.mkdir", _boom) + + class NoArtifactsWorkflow(StubWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> Any: + assert ctx.artifacts_dir is None + await super().run(ctx) + + class NoArtifactsGrader(StubGrader): + async def grade(self, ctx: GraderContext) -> Any: + assert ctx.artifacts_dir is None + await super().grade(ctx) + + backend = LocalBackend( + workflow=NoArtifactsWorkflow, + workflow_config=AgentWorkflowConfig(name="test"), + grader=NoArtifactsGrader, + ) + + request = ExecutionRequest( + id="r1", prompt=[{"role": "user", "content": "hi"}], label="x" + ) + on_complete = AsyncMock() + await backend.execute(request, on_complete, AsyncMock()) + + # mkdir failure degrades to "artifacts unavailable"; rollout still succeeds. + on_complete.assert_awaited_once() + assert on_complete.call_args[0][0].status == RolloutStatus.SUCCESS + async def test_execute_success_calls_callback(self): backend = self._make_backend() on_complete = AsyncMock()