From 3a9fb53ba490397dfc1208cf2684d5d0680d90b4 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Fri, 22 May 2026 14:51:02 -0700 Subject: [PATCH 1/4] [feat] single url --- examples/multiply_openai_agents/README.md | 104 +++++++++++++ examples/multiply_openai_agents/main.py | 140 ++++++++++++++++++ osmosis_ai/rollout/__init__.py | 2 - .../rollout/backend/harbor/agent_runner.py | 15 +- osmosis_ai/rollout/backend/harbor/backend.py | 44 +++--- .../rollout/backend/harbor/grader_runner.py | 50 ++++--- osmosis_ai/rollout/backend/local/backend.py | 13 +- osmosis_ai/rollout/context.py | 73 ++++----- osmosis_ai/rollout/driver.py | 13 +- .../integrations/agents/openai_agents.py | 102 ++++++------- .../rollout/integrations/agents/strands.py | 103 +++++-------- osmosis_ai/rollout/server/app.py | 33 +++-- osmosis_ai/rollout/types/__init__.py | 2 - osmosis_ai/rollout/types/protocol.py | 47 ++++-- osmosis_ai/rollout/types/sample.py | 17 ++- osmosis_ai/rollout/utils/rewards.py | 20 ++- 16 files changed, 517 insertions(+), 261 deletions(-) create mode 100644 examples/multiply_openai_agents/README.md create mode 100644 examples/multiply_openai_agents/main.py diff --git a/examples/multiply_openai_agents/README.md b/examples/multiply_openai_agents/README.md new file mode 100644 index 00000000..b284cbd5 --- /dev/null +++ b/examples/multiply_openai_agents/README.md @@ -0,0 +1,104 @@ +# multiply_openai_agents + +Minimal end-to-end Osmosis rollout server using the OpenAI Agents SDK. + +A single-file FastAPI server (`main.py`) that exposes the standard rollout +contract (`POST /rollout`, `GET /health`). On each rollout the workflow +spins up an `OsmosisAgent` with one `multiply` tool, runs it against the +trainer's session URL via LiteLLM, and the grader scores the final +`Answer: ` against the dataset label. + +Use this as the smoke test when you bring up multi-LoRA + remote rollout +on a fresh cluster. + +## Run standalone (sanity check) + +From this directory: + +```bash +uv venv .venv && source .venv/bin/activate +uv pip install -e '../..[server]' 'openai-agents>=0.14,<0.15' uvicorn + +python main.py & +curl http://localhost:8080/health # -> {"status":"ok", ...} +kill %1 +``` + +Standalone the server boots but won't drive a rollout end-to-end — there's +no chat-completions URL to call. The full flow is driven by the trainer +once you register an adapter that points at it. + +## Connect to the multi-LoRA + remote-rollout cluster + +Assumes you've already done `traingate multi-lora start ...` and the head +pod is reachable. Get the head pod name: + +```bash +POD=$(kubectl get pods -l 'skypilot-cluster=' -o name | head -n1) +``` + +Copy the example into the pod and start it on `localhost:8080`: + +```bash +tar czf /tmp/multiply.tgz -C .. multiply_openai_agents +kubectl cp /tmp/multiply.tgz "$POD":/tmp/multiply.tgz +kubectl exec "$POD" -- bash -lc ' + set -e + mkdir -p /tmp/multiply && tar xzf /tmp/multiply.tgz -C /tmp/multiply + cd /tmp/multiply/multiply_openai_agents + pip install -q -e /osmosis/workspace/sdk-python"[server]" \ + "openai-agents>=0.14,<0.15" + OPENAI_API_KEY='"$OPENAI_API_KEY"' ROLLOUT_PORT=8080 \ + nohup python main.py > /tmp/multiply.log 2>&1 & +' +kubectl exec "$POD" -- curl -fsS http://localhost:8080/health +``` + +Register an adapter pointing at it (run from `osmosis-traingate/`): + +```toml +# /tmp/multiply.toml +[config] +data = "/osmosis/data/datasets/multiply/train.parquet" +input_key = "messages" +label_key = "label" +rm_type = "math" +rank = 16 +alpha = 16 +num_row = 20 + +[config.metadata] +rollout_server_url = "http://localhost:8080" +``` + +```bash +uv run traingate multi-lora register multiply \ + --cluster --spec /tmp/multiply.toml +``` + +Tail both sides; you should see chat completions flowing within seconds: + +```bash +# rollout server +kubectl exec "$POD" -- tail -f /tmp/multiply.log + +# trainer +sky logs --follow | rg -i 'rollout|adapter|session-server' +``` + +## Dataset shape + +Each row in the parquet file needs the keys named in the adapter spec +(`input_key`, `label_key`). For this example: + +```python +{ + "messages": [ + {"role": "user", "content": "What is 23 times 41?"}, + ], + "label": "943", +} +``` + +The workflow treats `ctx.prompt` as the verbatim message list to seed the +agent's session. diff --git a/examples/multiply_openai_agents/main.py b/examples/multiply_openai_agents/main.py new file mode 100644 index 00000000..85c633a2 --- /dev/null +++ b/examples/multiply_openai_agents/main.py @@ -0,0 +1,140 @@ +"""Multiply rollout server: minimal end-to-end Osmosis rollout example. + +Self-contained FastAPI rollout server that drives a small OpenAI Agents SDK +workflow against the active rollout's chat-completions endpoint. The agent +has one tool (``multiply``) and a numeric grader that scores the final +answer against the dataset label. + +Bring up the multi-LoRA + remote-rollout cluster (see +``osmosis-traingate/specs/experiment/miles-multi-lora-service.toml``), then +point an adapter at this server via ``metadata.rollout_server_url`` and the +trainer's custom generate function will POST rollouts here. + +Standalone boot (sanity check): + + pip install -e '../..[server]' 'openai-agents>=0.14,<0.15' uvicorn + python main.py & + curl http://localhost:8080/health +""" + +from __future__ import annotations + +import logging +import os +import re +from typing import Any + +import uvicorn +from agents import ModelSettings, Runner, function_tool + +from osmosis_ai.rollout.agent_workflow import AgentWorkflow +from osmosis_ai.rollout.backend.local import LocalBackend +from osmosis_ai.rollout.context import AgentWorkflowContext, GraderContext +from osmosis_ai.rollout.grader import Grader +from osmosis_ai.rollout.integrations.agents.openai_agents import ( + OsmosisAgent, + OsmosisMemorySession, + OsmosisRolloutModel, +) +from osmosis_ai.rollout.server import create_rollout_server + +logger = logging.getLogger(__name__) + +MAX_TURNS = 8 +# Loose tolerance so float-y answers like "12.0" match integer labels. +ANSWER_TOLERANCE = 1e-2 + + +@function_tool +async def multiply(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b + + +class MultiplyWorkflow(AgentWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> None: + agent = OsmosisAgent( + name="multiply", + instructions=( + "Solve the user's multiplication problem. You may call the " + "`multiply` tool. End your final message with " + "`Answer: ` on its own line." + ), + model=OsmosisRolloutModel(), + model_settings=ModelSettings(temperature=1.0, max_tokens=512), + tools=[multiply], + ) + session = OsmosisMemorySession() + await Runner.run(agent, ctx.prompt, session=session, max_turns=MAX_TURNS) + + +_ANSWER_RE = re.compile(r"Answer:\s*(-?\d+(?:\.\d+)?)", re.IGNORECASE) + + +def _extract_text(content: Any) -> str: + """Flatten OpenAI Responses-API content into plain text. + + ``content`` can be a str, or a list of typed parts; we keep the text-ish + parts (``text`` / ``output_text``) and concatenate. + """ + if isinstance(content, str): + return content + if isinstance(content, list): + return "\n".join( + item.get("text", "") + for item in content + if isinstance(item, dict) and item.get("type") in ("text", "output_text") + ) + return "" + + +class MultiplyGrader(Grader): + """0/1 reward: match the trailing ``Answer: `` against ``ctx.label``.""" + + async def grade(self, ctx: GraderContext) -> None: + if ctx.sample is None or ctx.label is None: + ctx.set_reward(0.0) + return + + # Walk back to the most recent assistant turn with text content; + # skips tool-call / tool-result turns that have no human-readable body. + text = "" + for msg in reversed(ctx.sample.messages): + if msg.get("role") != "assistant": + continue + text = _extract_text(msg.get("content", "")) + if text: + break + + match = _ANSWER_RE.search(text) + if match is None: + ctx.set_reward(0.0) + return + + try: + predicted = float(match.group(1)) + target = float(ctx.label) + except ValueError: + ctx.set_reward(0.0) + return + + reward = 1.0 if abs(predicted - target) < ANSWER_TOLERANCE else 0.0 + ctx.set_reward(reward) + logger.info( + "multiply reward=%.1f predicted=%s target=%s", + reward, predicted, target, + ) + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, format="[multiply-rollout] %(message)s" + ) + backend = LocalBackend(workflow=MultiplyWorkflow, grader=MultiplyGrader) + app = create_rollout_server(backend=backend) + port = int(os.environ.get("ROLLOUT_PORT", "8080")) + uvicorn.run(app, host="0.0.0.0", port=port) + + +if __name__ == "__main__": + main() diff --git a/osmosis_ai/rollout/__init__.py b/osmosis_ai/rollout/__init__.py index 4805a45a..257da4fa 100644 --- a/osmosis_ai/rollout/__init__.py +++ b/osmosis_ai/rollout/__init__.py @@ -23,7 +23,6 @@ GraderCompleteRequest, GraderConfig, GraderStatus, - MultiTurnMode, RolloutCompleteRequest, RolloutErrorCategory, RolloutInitRequest, @@ -48,7 +47,6 @@ "GraderStatus", "HarborAgentWorkflowContext", "LocalBackend", - "MultiTurnMode", "OsmosisRolloutModel", "OsmosisStrandsAgent", "RolloutCompleteRequest", diff --git a/osmosis_ai/rollout/backend/harbor/agent_runner.py b/osmosis_ai/rollout/backend/harbor/agent_runner.py index 67a18f5d..07276750 100644 --- a/osmosis_ai/rollout/backend/harbor/agent_runner.py +++ b/osmosis_ai/rollout/backend/harbor/agent_runner.py @@ -66,16 +66,19 @@ async def run_workflow( with rollout_ctx: await workflow.run(ctx) - samples = await rollout_ctx.get_samples() - samples_data = {sid: s.model_dump() for sid, s in samples.items()} + sample = await rollout_ctx.get_sample() + sample_data = sample.model_dump() if sample is not None else None - (AGENT_LOGS_DIR / "samples.json").write_text( - json.dumps(samples_data, default=str) + (AGENT_LOGS_DIR / "sample.json").write_text( + json.dumps(sample_data, default=str) ) meta["status"] = "success" - meta["samples"] = samples_data - print(f"Agent runner complete: {len(samples)} samples collected") + meta["sample"] = sample_data + print( + "Agent runner complete: " + + ("1 sample collected" if sample is not None else "no sample collected") + ) except Exception as e: traceback.print_exc() diff --git a/osmosis_ai/rollout/backend/harbor/backend.py b/osmosis_ai/rollout/backend/harbor/backend.py index 8ad1b0ea..421ea0be 100644 --- a/osmosis_ai/rollout/backend/harbor/backend.py +++ b/osmosis_ai/rollout/backend/harbor/backend.py @@ -39,7 +39,7 @@ RolloutStatus, ) from osmosis_ai.rollout.utils.imports import to_import_path -from osmosis_ai.rollout.utils.rewards import validate_samples_have_rewards +from osmosis_ai.rollout.utils.rewards import validate_sample_has_reward logger: logging.Logger = logging.getLogger(__name__) @@ -53,7 +53,7 @@ set -e PYTHONPATH=/workspace python -m osmosis_ai.rollout.backend.harbor.grader_runner \ --config /logs/agent/rollout_config.json \ - --samples /logs/agent/samples.json + --sample /logs/agent/sample.json """ @@ -422,8 +422,8 @@ async def on_verification_start(self, event: TrialHookEvent) -> None: metadata = get_agent_metadata(event) if metadata and metadata.get("status") == "success": - samples = parse_samples(metadata.get("samples", {})) - result = ExecutionResult(status=RolloutStatus.SUCCESS, samples=samples) + sample = parse_sample(metadata.get("sample")) + result = ExecutionResult(status=RolloutStatus.SUCCESS, sample=sample) elif event.result and event.result.exception_info: err = event.result.exception_info result = ExecutionResult( @@ -467,15 +467,20 @@ async def on_trial_end(self, event: TrialHookEvent) -> None: if pending.on_grader_complete: metadata = get_agent_metadata(event) - samples = parse_samples(metadata.get("samples", {})) if metadata else {} + sample = parse_sample(metadata.get("sample") if metadata else None) if event.result and event.result.verifier_result: + # Harbor surfaces verifier output as ``rewards: dict[str, float]`` + # regardless of what shape the verifier wrote — for the single-sample + # contract we take any one value and apply it to our single sample. rewards = event.result.verifier_result.rewards or {} - for sid, reward in rewards.items(): - if sid in samples: - samples[sid].reward = float(reward) + reward_values = [ + float(v) for v in rewards.values() if v is not None + ] + if sample is not None and reward_values: + sample.reward = reward_values[0] try: - validate_samples_have_rewards(samples) + validate_sample_has_reward(sample) except ValueError as e: if not rewards: logger.warning( @@ -491,24 +496,24 @@ async def on_trial_end(self, event: TrialHookEvent) -> None: ) result = ExecutionResult( status=RolloutStatus.FAILURE, - samples=samples, + sample=sample, err_message=str(e), err_category=RolloutErrorCategory.VALIDATION_ERROR, ) else: result = ExecutionResult( - status=RolloutStatus.SUCCESS, samples=samples + status=RolloutStatus.SUCCESS, sample=sample ) elif event.result and event.result.exception_info: err = event.result.exception_info result = ExecutionResult( status=RolloutStatus.FAILURE, - samples=samples, + sample=sample, err_message=err.exception_message, err_category=RolloutErrorCategory.AGENT_ERROR, ) else: - result = ExecutionResult(status=RolloutStatus.FAILURE, samples=samples) + result = ExecutionResult(status=RolloutStatus.FAILURE, sample=sample) await pending.on_grader_complete(result) @@ -541,11 +546,14 @@ def get_agent_metadata(event: TrialHookEvent) -> dict[str, Any] | None: return None -def parse_samples(raw: dict[str, Any]) -> dict[str, RolloutSample]: - return { - sid: RolloutSample.model_validate(data) if isinstance(data, dict) else data - for sid, data in raw.items() - } +def parse_sample(raw: Any) -> RolloutSample | None: + if raw is None: + return None + if isinstance(raw, RolloutSample): + return raw + if isinstance(raw, dict): + return RolloutSample.model_validate(raw) + return None def rewrite_url_for_docker(url: str) -> str: diff --git a/osmosis_ai/rollout/backend/harbor/grader_runner.py b/osmosis_ai/rollout/backend/harbor/grader_runner.py index 9637e297..0416a160 100644 --- a/osmosis_ai/rollout/backend/harbor/grader_runner.py +++ b/osmosis_ai/rollout/backend/harbor/grader_runner.py @@ -1,7 +1,7 @@ -"""Grade agent samples inside a Harbor container. +"""Grade the rollout's single sample inside a Harbor container. Usage: - osmosis-grader-runner --config /workspace/rollout_config.json --samples /logs/agent/samples.json + osmosis-grader-runner --config /workspace/rollout_config.json --sample /logs/agent/sample.json """ from __future__ import annotations @@ -20,12 +20,12 @@ def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Grade agent samples") + parser = argparse.ArgumentParser(description="Grade the rollout sample") parser.add_argument( "--config", type=Path, required=True, help="Path to rollout_config.json" ) parser.add_argument( - "--samples", type=Path, required=True, help="Path to samples.json" + "--sample", type=Path, required=True, help="Path to sample.json" ) return parser.parse_args() @@ -34,14 +34,16 @@ def load_config(path: Path) -> dict[str, Any]: return json.loads(path.read_text()) -def load_samples(path: Path) -> dict[str, RolloutSample]: - raw: dict[str, Any] = json.loads(path.read_text()) - return {sid: RolloutSample.model_validate(data) for sid, data in raw.items()} +def load_sample(path: Path) -> RolloutSample | None: + raw: Any = json.loads(path.read_text()) + if raw is None: + return None + return RolloutSample.model_validate(raw) -def write_rewards(rewards: dict[str, float | None]) -> None: +def write_reward(reward: float | None) -> None: VERIFIER_LOGS_DIR.mkdir(parents=True, exist_ok=True) - (VERIFIER_LOGS_DIR / "reward.json").write_text(json.dumps(rewards)) + (VERIFIER_LOGS_DIR / "reward.json").write_text(json.dumps({"reward": reward})) def main() -> None: @@ -51,37 +53,39 @@ def main() -> None: if not label: print("No label in config, skipping grading") - write_rewards({}) + write_reward(None) return - if not args.samples.exists(): - print("No samples found, skipping grading") - write_rewards({}) + if not args.sample.exists(): + print("No sample found, skipping grading") + write_reward(None) return if "grader" not in config: print("No grader in config, skipping grading") - write_rewards({}) + write_reward(None) + return + + sample = load_sample(args.sample) + if sample is None: + print("Sample file is empty, skipping grading") + write_reward(None) return - samples = load_samples(args.samples) grader_cls = resolve_object(config["grader"]) grader_config = ( resolve_object(config["grader_config"]) if "grader_config" in config else None ) - ctx = GraderContext(label=label, samples=samples) + ctx = GraderContext(label=label, sample=sample) grader = grader_cls(grader_config) asyncio.run(grader.grade(ctx)) - graded = ctx.get_samples() - for sid, sample in graded.items(): - if sample.reward is None: - raise RuntimeError(f"Sample {sid} has no reward after grading") + if ctx.sample is None or ctx.sample.reward is None: + raise RuntimeError("Sample has no reward after grading") - rewards = {sid: sample.reward for sid, sample in graded.items()} - write_rewards(rewards) - print(f"Grading complete: {rewards}") + write_reward(ctx.sample.reward) + print(f"Grading complete: reward={ctx.sample.reward}") if __name__ == "__main__": diff --git a/osmosis_ai/rollout/backend/local/backend.py b/osmosis_ai/rollout/backend/local/backend.py index 81f9cdfa..d1b5badb 100644 --- a/osmosis_ai/rollout/backend/local/backend.py +++ b/osmosis_ai/rollout/backend/local/backend.py @@ -22,7 +22,7 @@ ) from osmosis_ai.rollout.utils.concurrency import ConcurrencyLimiter from osmosis_ai.rollout.utils.imports import resolve_object -from osmosis_ai.rollout.utils.rewards import validate_samples_have_rewards +from osmosis_ai.rollout.utils.rewards import validate_sample_has_reward logger: logging.Logger = logging.getLogger(__name__) @@ -114,7 +114,7 @@ async def run_workflow(self, request: ExecutionRequest) -> ExecutionResult: return ExecutionResult( status=RolloutStatus.SUCCESS, - samples=await rollout_ctx.get_samples(), + sample=await rollout_ctx.get_sample(), ) async def run_grader( @@ -123,21 +123,20 @@ async def run_grader( if not self.grader_cls: return result - grader_ctx = GraderContext(label=request.label, samples=result.samples) + grader_ctx = GraderContext(label=request.label, sample=result.sample) try: grader = self.grader_cls(self.grader_config) await grader.grade(grader_ctx) - samples = grader_ctx.get_samples() - validate_samples_have_rewards(samples) + validate_sample_has_reward(grader_ctx.sample) return ExecutionResult( status=RolloutStatus.SUCCESS, - samples=samples, + sample=grader_ctx.sample, ) except Exception as e: logger.error(traceback.format_exc()) return ExecutionResult( status=RolloutStatus.FAILURE, - samples=result.samples, + sample=result.sample, err_message=str(e), err_category=_categorize_exception(e), ) diff --git a/osmosis_ai/rollout/context.py b/osmosis_ai/rollout/context.py index ba10c294..fc2d6559 100644 --- a/osmosis_ai/rollout/context.py +++ b/osmosis_ai/rollout/context.py @@ -11,22 +11,18 @@ class SampleSource(ABC): - """Produces a ``RolloutSample`` for grading from a framework-specific object. + """Produces the rollout's ``RolloutSample`` for grading. - Each integration ships a small source class that wraps its underlying - object (an Agent, a Session, etc.) and produces a sample on demand. - The source is decoupled from the integration's framework classes so - they do not need to inherit from this ABC. - - RolloutContext indexes registered sources by name and calls - ``get_sample`` lazily at sample collection time, passing the - registration name so the source can stamp it onto the returned - ``RolloutSample.id``. + A rollout produces exactly one sample (one agent run, one reward). The + active ``RolloutContext`` holds at most one ``SampleSource``; the + integration registers a source that wraps its underlying object + (typically an Agent or a Session) and the backend pulls the sample at + grading time. """ @abstractmethod - async def get_sample(self, name: str) -> RolloutSample: - """Return the current rollout sample for the given registration name.""" + async def get_sample(self) -> RolloutSample: + """Return the rollout's current sample.""" rollout_contextvar: ContextVar["RolloutContext | None"] = ContextVar( @@ -50,7 +46,7 @@ class RolloutContext: chat_completions_url: str = "" api_key: str | None = None rollout_id: str = "" - sample_sources: dict[str, SampleSource] = field(default_factory=dict) + sample_source: SampleSource | None = None def __post_init__(self) -> None: if not self.chat_completions_url: @@ -67,24 +63,28 @@ def __enter__(self) -> "RolloutContext": def __exit__(self, *_: Any) -> None: rollout_contextvar.set(None) - def register_sample_source(self, name: str, source: SampleSource) -> None: - """Register a sample source under the given name. The context calls - ``source.get_sample(name)`` lazily when ``get_samples()`` is invoked, - so the source can keep reflecting state that mutates until the - rollout ends. + def set_sample_source(self, source: SampleSource) -> None: + """Register the single sample source for this rollout. + + A rollout has exactly one sample; integrations call this from their + agent/session constructors so the backend can pull the sample at + grading time. Attempting to register a second source raises, which + catches the common bug of putting two ``OsmosisMemorySession``s or + two ``OsmosisStrandsAgent``s in the same workflow. """ - if name in self.sample_sources: + if self.sample_source is not None: raise ValueError( - f"Session with {name} already exists, please give your session " - "or agent a unique name in the workflow" + "RolloutContext already has a sample source registered. " + "A rollout produces one sample; construct a single agent/" + "session per rollout." ) - self.sample_sources[name] = source + self.sample_source = source - async def get_samples(self) -> dict[str, RolloutSample]: - out: dict[str, RolloutSample] = {} - for name, source in self.sample_sources.items(): - out[name] = await source.get_sample(name) - return out + async def get_sample(self) -> RolloutSample | None: + """Return the rollout's sample, or ``None`` if no source was registered.""" + if self.sample_source is None: + return None + return await self.sample_source.get_sample() def get_rollout_context() -> RolloutContext | None: @@ -93,17 +93,20 @@ def get_rollout_context() -> RolloutContext | None: @dataclass class GraderContext: + """Context passed to ``Grader.grade``. + + Carries the single sample produced by the rollout. Graders set its + reward via :meth:`set_reward`. + """ + label: str | None = None - samples: dict[str, RolloutSample] = field(default_factory=dict) + sample: RolloutSample | None = None project_path: str | None = None - def get_samples(self) -> dict[str, RolloutSample]: - return self.samples - - def set_sample_reward(self, sample_id: str, reward: float) -> None: - if sample_id not in self.samples: - raise ValueError(f"Sample {sample_id} not found") - self.samples[sample_id].reward = reward + def set_reward(self, reward: float) -> None: + if self.sample is None: + raise ValueError("GraderContext has no sample to reward") + self.sample.reward = reward @dataclass diff --git a/osmosis_ai/rollout/driver.py b/osmosis_ai/rollout/driver.py index 84f24e82..8cfac4ea 100644 --- a/osmosis_ai/rollout/driver.py +++ b/osmosis_ai/rollout/driver.py @@ -17,22 +17,19 @@ class RolloutOutcome: """Result of a single rollout execution. - This is the eval-facing contract. Eval doesn't need to know - whether this came from in-process execution or an HTTP call. + Single-sample by design: one rollout = one agent run = one reward. + ``sample`` carries the conversation and reward (when grading succeeded), + ``rollout_id`` identifies the rollout in logs/cache rows, and the rest + are run-level metrics + diagnostics. """ status: RolloutStatus - samples: dict[str, RolloutSample] = field(default_factory=dict) + sample: RolloutSample | None = None error: str | None = None duration_ms: float = 0.0 tokens: int = 0 systemic_error: str | None = None rollout_id: str | None = None - controller_created_sample_ids: list[str] = field(default_factory=list) - completion_counts: dict[str, int] = field(default_factory=dict) - full_callback_sample_ids: list[str] = field(default_factory=list) - scored_sample_ids: list[str] = field(default_factory=list) - skipped_sample_ids: list[str] = field(default_factory=list) callback_diagnostics: dict[str, Any] = field(default_factory=dict) skipped: bool = False diff --git a/osmosis_ai/rollout/integrations/agents/openai_agents.py b/osmosis_ai/rollout/integrations/agents/openai_agents.py index 78f8b55b..5969f854 100644 --- a/osmosis_ai/rollout/integrations/agents/openai_agents.py +++ b/osmosis_ai/rollout/integrations/agents/openai_agents.py @@ -1,6 +1,5 @@ import uuid from collections.abc import AsyncIterator -from contextvars import ContextVar from typing import Any from agents import Agent @@ -17,36 +16,26 @@ from osmosis_ai.rollout.context import SampleSource, get_rollout_context from osmosis_ai.rollout.types import RolloutSample -current_sample_id: ContextVar[str | None] = ContextVar( - "osmosis_current_sample_id", default=None -) - class OsmosisMemorySession(SessionABC): - """In-memory session that doubles as the sample source for an Osmosis rollout. - - Behavior: - - Inside a ``RolloutContext``: registers itself for sample collection - (via ``SessionSampleSource``) and publishes its sample id - to a ContextVar each time the runner reads from or writes to it. - ``OsmosisRolloutModel`` reads that ContextVar to stamp per-rollout - headers. - - Outside a ``RolloutContext``: behaves as a plain in-memory - ``SessionABC`` implementation, so the same workflow code can run - locally without an Osmosis rollout. + """In-memory session that doubles as the rollout's sample source. + + Inside a ``RolloutContext`` this registers itself with the context so + the backend can pull the conversation at grading time. Outside a + ``RolloutContext`` it behaves as a plain in-memory ``SessionABC``, so + the same workflow code can run locally without an Osmosis rollout. Pass exactly one ``OsmosisMemorySession`` per ``Runner.run`` when using - ``OsmosisRolloutModel``. + ``OsmosisRolloutModel`` — a rollout produces one sample. """ - def __init__(self, name: str | None = None) -> None: - self.name: str = name or uuid.uuid4().hex - self.session_id: str = self.name + def __init__(self) -> None: + self.session_id: str = uuid.uuid4().hex self.items: list[TResponseInputItem] = [] self._registered_context_id: int | None = None ctx = get_rollout_context() if ctx is not None: - ctx.register_sample_source(self.name, SessionSampleSource(self)) + ctx.set_sample_source(SessionSampleSource(self)) self._registered_context_id = id(ctx) def _raise_if_unregistered_in_rollout_context(self) -> None: @@ -63,14 +52,12 @@ def _raise_if_unregistered_in_rollout_context(self) -> None: async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: self._raise_if_unregistered_in_rollout_context() - current_sample_id.set(self.name) if limit is None: return list(self.items) return list(self.items[-limit:]) async def add_items(self, items: list[TResponseInputItem]) -> None: self._raise_if_unregistered_in_rollout_context() - current_sample_id.set(self.name) self.items.extend(items) async def pop_item(self) -> TResponseInputItem | None: @@ -81,7 +68,7 @@ async def clear_session(self) -> None: class SessionSampleSource(SampleSource): - """Produces a ``RolloutSample`` from any OpenAI Agents SDK ``Session``. + """Produces the rollout's sample from any OpenAI Agents SDK ``Session``. Works against any ``SessionABC`` by returning the items the runner persisted via ``add_items`` (canonical Responses-API ``TResponseInputItem`` @@ -97,22 +84,23 @@ class SessionSampleSource(SampleSource): def __init__(self, session: SessionABC) -> None: self.session = session - async def get_sample(self, name: str) -> RolloutSample: + async def get_sample(self) -> RolloutSample: items = await self.session.get_items() - return RolloutSample(id=name, messages=items) + return RolloutSample(messages=items) class OsmosisRolloutModel(LitellmModel): - """Placeholder ``Model`` that carries litellm kwargs for use in workflow configs. + """Placeholder ``Model`` that carries litellm kwargs for workflow configs. - This is *not* a usable model on its own: ``OsmosisAgent`` replaces it - with a real :class:`OsmosisLitellmModel` (wired to the active - ``RolloutContext``) at agent construction time. Any direct call into - ``stream_response`` / ``get_response`` raises ``NotImplementedError`` - because the placeholder has no connection params. + Not a usable model on its own: ``OsmosisAgent`` replaces it with a real + :class:`OsmosisLitellmModel` (wired to the active ``RolloutContext``) + at agent construction time. Any direct call into ``stream_response`` / + ``get_response`` raises ``NotImplementedError`` because the placeholder + has no connection params. - Subclassing ``LitellmModel`` (without invoking its ``__init__``) is purely - a typing convenience so the placeholder satisfies ``Agent.model: Model``. + Subclassing ``LitellmModel`` (without invoking its ``__init__``) is + purely a typing convenience so the placeholder satisfies + ``Agent.model: Model``. """ def __init__(self, **litellm_kwargs: Any) -> None: @@ -135,19 +123,21 @@ async def get_response(self, *args: Any, **kwargs: Any) -> ModelResponse: class OsmosisLitellmModel(LitellmModel): - """Streaming-only LitellmModel pre-wired for the Osmosis completions server. - - Stateless across calls: the per-call session name is read from a - ContextVar that ``OsmosisMemorySession`` publishes when the runner interacts - with it. One instance can be safely shared across agents and runs (within - the rollout context that constructed it). - - Caveats this class handles for the user: - - The Osmosis completions server is streaming-only, so ``get_response`` - is implemented by aggregating ``stream_response`` events. This makes - ``Runner.run`` and ``Runner.run_sync`` work without forcing the user - to switch to ``Runner.run_streamed``. - - Per-rollout bookkeeping headers are injected on every model call. + """Streaming-only LitellmModel pre-wired to the active rollout's chat endpoint. + + The model server identifies the rollout entirely via the URL it was + handed (rollout id is baked into the ``chat_completions_url`` path), + so this class no longer stamps per-call headers. It does, however, + refuse to make a request when no ``OsmosisMemorySession`` has been + registered with the active ``RolloutContext`` — that catches the + common bug of calling ``Runner.run`` without + ``session=OsmosisMemorySession()`` (which would otherwise leave the + grader with no conversation to score). + + The Osmosis completions server is streaming-only, so ``get_response`` + is implemented by aggregating ``stream_response`` events. This makes + ``Runner.run`` and ``Runner.run_sync`` work without forcing the user + to switch to ``Runner.run_streamed``. """ def __init__(self, **litellm_kwargs: Any) -> None: @@ -158,7 +148,6 @@ def __init__(self, **litellm_kwargs: Any) -> None: "Construct it inside your workflow run, where the execution " "backend has set up the context." ) - self.rollout_id: str = ctx.rollout_id super().__init__( model="openai/osmosis-rollout", base_url=ctx.chat_completions_url, @@ -167,19 +156,18 @@ def __init__(self, **litellm_kwargs: Any) -> None: ) def _merge_headers(self, model_settings: ModelSettings) -> dict[str, str]: - sample_id = current_sample_id.get() - if sample_id is None: + # No per-call routing headers anymore — the URL carries rollout + # identity. We still gate the call on a registered sample source + # so users who forget ``session=OsmosisMemorySession()`` get a + # loud error instead of a silently-empty sample at grading time. + ctx = get_rollout_context() + if ctx is None or ctx.sample_source is None: raise RuntimeError( "OsmosisLitellmModel was called without an active OsmosisMemorySession.\n" "Pass `session=OsmosisMemorySession()` to Runner.run() so the " - "model can stamp per-rollout headers and the runner can " - "persist the conversation for grading." + "runner persists the conversation for grading." ) - return { - **super()._merge_headers(model_settings), - "x-sample-id": sample_id, - "x-rollout-id": self.rollout_id, - } + return super()._merge_headers(model_settings) async def get_response(self, *args: Any, **kwargs: Any) -> ModelResponse: output_items: list[Any] = [] diff --git a/osmosis_ai/rollout/integrations/agents/strands.py b/osmosis_ai/rollout/integrations/agents/strands.py index 56bcd8b9..2ae5c084 100644 --- a/osmosis_ai/rollout/integrations/agents/strands.py +++ b/osmosis_ai/rollout/integrations/agents/strands.py @@ -1,94 +1,57 @@ -import uuid -from collections.abc import AsyncGenerator -from typing import Any, TypeVar, cast +from typing import Any, cast from strands import Agent as StrandsAgent from strands.models.litellm import LiteLLMModel from strands.models.model import Model -from strands.types.content import Messages, SystemContentBlock -from strands.types.streaming import StreamEvent -from strands.types.tools import ToolChoice, ToolSpec +from strands.types.content import Messages from osmosis_ai.rollout.context import ( - RolloutContext, SampleSource, get_rollout_context, ) from osmosis_ai.rollout.types import RolloutSample from osmosis_ai.rollout.utils.messages import map_initial_messages_to_content_blocks -T = TypeVar("T") - class StrandsAgentSampleSource(SampleSource): - """Produces a ``RolloutSample`` from a Strands agent's ``messages`` field. + """Produces the rollout sample from a Strands agent's ``messages`` field. - Strands agents accumulate the conversation history on ``agent.messages`` - in chat-completion format, so the source just wraps that list. + Strands accumulates the conversation on ``agent.messages`` in + chat-completion format, so this just wraps that list. """ def __init__(self, agent: StrandsAgent) -> None: self.agent = agent - async def get_sample(self, name: str) -> RolloutSample: - return RolloutSample(id=name, messages=list(self.agent.messages)) + async def get_sample(self) -> RolloutSample: + return RolloutSample(messages=list(self.agent.messages)) class OsmosisRolloutModel(LiteLLMModel): - """Placeholder model that carries sampling params. - - At runtime, OsmosisStrandsAgent calls for_sample() which reads connection - info from the active RolloutContext to create a concrete LiteLLMModel. - """ - - def for_sample(self, sample_id: str, rollout_ctx: RolloutContext) -> LiteLLMModel: - headers = { - "x-sample-id": sample_id, - "x-rollout-id": rollout_ctx.rollout_id, - } - client_args = { - "api_base": rollout_ctx.chat_completions_url, - "api_key": rollout_ctx.api_key, - "extra_headers": headers, - } - return LiteLLMModel( - client_args=client_args, - model_id="openai/osmosis-rollout", - **self.config, # type: ignore - ) - - def update_config(self, **model_config: Any) -> None: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + """Placeholder ``Model`` that carries litellm kwargs for workflow configs. - def get_config(self) -> Any: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + Not a usable model on its own: ``OsmosisStrandsAgent`` replaces it with + a real ``LiteLLMModel`` (wired to the active ``RolloutContext``) at + agent construction time. It has no connection params, so any direct + call into the ``Model`` API will fail. - def structured_output( - self, - output_model: type[T], - prompt: Messages, - system_prompt: str | None = None, - **kwargs: Any, - ) -> AsyncGenerator[dict[str, T | Any], None]: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + Subclassing ``LiteLLMModel`` (without invoking its ``__init__``) is + purely a typing convenience so the placeholder satisfies + ``StrandsAgent.model: Model`` -- same pattern as + ``integrations.agents.openai_agents.OsmosisRolloutModel``. + """ - def stream( - self, - messages: Messages, - tool_specs: list[ToolSpec] | None = None, - system_prompt: str | None = None, - *, - tool_choice: ToolChoice | None = None, - system_prompt_content: list[SystemContentBlock] | None = None, - invocation_state: dict[str, Any] | None = None, - **kwargs: Any, - ) -> AsyncGenerator[StreamEvent, None]: - raise NotImplementedError("This should not be called for OsmosisRolloutModel") + def __init__(self, **litellm_kwargs: Any) -> None: + self.litellm_kwargs: dict[str, Any] = litellm_kwargs class OsmosisStrandsAgent(StrandsAgent): - """Drop-in replacement for StrandsAgent that handles rollout model swap and - sample registration transparently. The constructor signature matches StrandsAgent. + """Drop-in ``StrandsAgent`` that wires itself into the active rollout. + + If ``model`` is an ``OsmosisRolloutModel`` placeholder, this materializes + a real ``LiteLLMModel`` against the active ``RolloutContext`` and + registers the agent as the rollout's sample source. One + ``OsmosisStrandsAgent`` per rollout (matching the single-sample model). """ def __init__( @@ -106,11 +69,19 @@ def __init__( if rollout_ctx is None: raise RuntimeError( "OsmosisRolloutModel requires an active RolloutContext. " - "Ensure the execution backend sets up the context before running the workflow." + "Ensure the execution backend sets up the context before " + "running the workflow." ) - name = kwargs.get("name") or kwargs.get("agent_id") or uuid.uuid4().hex - model = model.for_sample(name, rollout_ctx) - rollout_ctx.register_sample_source(name, StrandsAgentSampleSource(self)) + litellm_model: Model = LiteLLMModel( + client_args={ + "api_base": rollout_ctx.chat_completions_url, + "api_key": rollout_ctx.api_key, + }, + model_id="openai/osmosis-rollout", + **model.litellm_kwargs, + ) + rollout_ctx.set_sample_source(StrandsAgentSampleSource(self)) + model = litellm_model super().__init__( *args, diff --git a/osmosis_ai/rollout/server/app.py b/osmosis_ai/rollout/server/app.py index 6717a09c..c72a320a 100644 --- a/osmosis_ai/rollout/server/app.py +++ b/osmosis_ai/rollout/server/app.py @@ -1,5 +1,6 @@ import logging import traceback +import uuid from typing import Any from fastapi import BackgroundTasks, FastAPI, HTTPException @@ -48,20 +49,24 @@ async def rollout( async def _handle_rollout( backend: ExecutionBackend, request: RolloutInitRequest ) -> None: + # Routing identity is in the URLs; ``rollout_id`` in the body is debug + # metadata. We prefer the caller's id (so logs/cache rows correlate + # across systems) and synthesize one only if the caller omits it. + rollout_id = request.rollout_id or uuid.uuid4().hex auth = ControllerAuth(api_key=request.controller_api_key) rollout_ctx = RolloutContext( chat_completions_url=request.chat_completions_url, api_key=request.controller_api_key, - rollout_id=request.rollout_id, + rollout_id=rollout_id, ) async def on_workflow_complete(result: ExecutionResult) -> None: await post_json_with_retry( url=request.completion_callback_url, payload=RolloutCompleteRequest( - rollout_id=request.rollout_id, status=result.status, + rollout_id=rollout_id, err_message=result.err_message, err_category=result.err_category, ).model_dump(), @@ -72,24 +77,24 @@ async def on_grader_complete(result: ExecutionResult) -> None: if not request.grader_callback_url: logger.info( "Skipping grader callback for %s: no grader_callback_url", - request.rollout_id, + rollout_id, ) return logger.info( - "Posting grader callback for %s to %s (status=%s, samples=%d)", - request.rollout_id, + "Posting grader callback for %s to %s (status=%s, has_sample=%s)", + rollout_id, request.grader_callback_url, result.status, - len(result.samples), + result.sample is not None, ) resp = await post_json_with_retry( url=request.grader_callback_url, payload=GraderCompleteRequest( - rollout_id=request.rollout_id, status=GraderStatus.SUCCESS if result.status == RolloutStatus.SUCCESS else GraderStatus.FAILURE, - samples=result.samples, + rollout_id=rollout_id, + sample=result.sample, err_message=result.err_message, err_category=result.err_category, ).model_dump(), @@ -97,7 +102,7 @@ async def on_grader_complete(result: ExecutionResult) -> None: ) logger.info( "Grader callback for %s completed: status=%d", - request.rollout_id, + rollout_id, resp.status_code, ) @@ -105,7 +110,7 @@ async def on_grader_complete(result: ExecutionResult) -> None: with rollout_ctx: await backend.execute( ExecutionRequest( - id=request.rollout_id, + id=rollout_id, prompt=request.initial_messages, label=request.label, agent_timeout_sec=request.agent_timeout_sec, @@ -116,17 +121,15 @@ async def on_grader_complete(result: ExecutionResult) -> None: if request.grader_callback_url else None, ) - logger.info("Rollout %s completed successfully", request.rollout_id) + logger.info("Rollout %s completed successfully", rollout_id) except Exception: - logger.error( - "Rollout %s failed: %s", request.rollout_id, traceback.format_exc() - ) + logger.error("Rollout %s failed: %s", rollout_id, traceback.format_exc()) try: await post_json_with_retry( url=request.completion_callback_url, payload=RolloutCompleteRequest( - rollout_id=request.rollout_id, status=RolloutStatus.FAILURE, + rollout_id=rollout_id, err_message="Internal server error", ).model_dump(), headers=auth.as_bearer_headers(), diff --git a/osmosis_ai/rollout/types/__init__.py b/osmosis_ai/rollout/types/__init__.py index 8e24b9d9..2b117342 100644 --- a/osmosis_ai/rollout/types/__init__.py +++ b/osmosis_ai/rollout/types/__init__.py @@ -17,7 +17,6 @@ ExecutionRequest, ExecutionResult, MessageDict, - MultiTurnMode, RolloutErrorCategory, RolloutSample, RolloutStatus, @@ -35,7 +34,6 @@ "GraderInitResponse", "GraderStatus", "MessageDict", - "MultiTurnMode", "RolloutCompleteRequest", "RolloutErrorCategory", "RolloutInitRequest", diff --git a/osmosis_ai/rollout/types/protocol.py b/osmosis_ai/rollout/types/protocol.py index fa0e0925..57aefb41 100644 --- a/osmosis_ai/rollout/types/protocol.py +++ b/osmosis_ai/rollout/types/protocol.py @@ -12,20 +12,27 @@ class RolloutInitRequest(BaseModel): - rollout_id: str + """Body of the POST to the SDK's /rollout endpoint. + + Routing identity lives in the URLs the caller hands us: the rollout + id is baked into ``chat_completions_url`` (session-scoped) and into + both callback URLs as a path segment. ``rollout_id`` is repeated in + the body for debug logging on the rollout-server side and for + correlation in user-side dashboards; the SDK does not rely on it for + routing, so it is optional. + """ - # These come from slime samples initial_messages: list[MessageDict] label: str | None = None metadata: dict[str, Any] | None = None - # These come from the rollout controller + rollout_id: str | None = None + chat_completions_url: str controller_api_key: str | None = None completion_callback_url: str grader_callback_url: str | None = None - # Per agent and grader timeout agent_timeout_sec: float | None = None grader_timeout_sec: float | None = None @@ -36,8 +43,15 @@ class RolloutInitResponse(BaseModel): ... class RolloutCompleteRequest(BaseModel): - rollout_id: str + """Body of the rollout-complete callback. + + ``rollout_id`` mirrors the id embedded in the callback URL purely for + debug/log correlation on the controller side. The receiver identifies + the rollout from the URL path, not this field. + """ + status: RolloutStatus + rollout_id: str | None = None extra_fields: dict[str, Any] | None = None @@ -46,8 +60,15 @@ class RolloutCompleteRequest(BaseModel): class GraderInitRequest(BaseModel): - rollout_id: str - samples: dict[str, RolloutSample] + """Body of a POST to a remote grader endpoint. + + A rollout produces a single sample (the agent's conversation), so this + carries that one sample directly rather than the legacy + ``dict[str, RolloutSample]``. + """ + + sample: RolloutSample + rollout_id: str | None = None completion_callback_url: str extra_fields: dict[str, Any] | None = None @@ -64,8 +85,16 @@ class GraderStatus(StrEnum): class GraderCompleteRequest(BaseModel): - rollout_id: str + """Body of the grader-complete callback. + + Carries the single graded sample (with its ``reward`` populated on + success) and nothing on failure. ``rollout_id`` mirrors the URL path + segment for debug/log correlation; the controller resolves the + rollout from the URL. + """ + status: GraderStatus - samples: dict[str, RolloutSample] + rollout_id: str | None = None + sample: RolloutSample | None = None err_message: str | None = None err_category: RolloutErrorCategory | None = None diff --git a/osmosis_ai/rollout/types/sample.py b/osmosis_ai/rollout/types/sample.py index 0996b263..25333d30 100644 --- a/osmosis_ai/rollout/types/sample.py +++ b/osmosis_ai/rollout/types/sample.py @@ -9,7 +9,15 @@ class RolloutSample(BaseModel): - id: str + """The conversation + grading artefacts produced by one rollout. + + A rollout produces exactly one sample (one agent run, one reward). There + used to be a per-sample id so the wire protocol could carry a + ``dict[str, RolloutSample]``; with the URL-routed single-sample wire + protocol the id is gone and callers identify rollouts via the URL paths + they hand the SDK. + """ + messages: Sequence[SampleMessage] = Field(default_factory=list) label: str | None = None reward: float | None = None @@ -33,11 +41,6 @@ class RolloutErrorCategory(StrEnum): AGENT_ERROR = "agent_error" -class MultiTurnMode(StrEnum): - MULTI_SAMPLE = "multi_sample" - SINGLE_SAMPLE = "single_sample" - - class ExecutionRequest(BaseModel): id: str prompt: list[MessageDict] @@ -48,6 +51,6 @@ class ExecutionRequest(BaseModel): class ExecutionResult(BaseModel): status: RolloutStatus - samples: dict[str, RolloutSample] = Field(default_factory=dict) + sample: RolloutSample | None = None err_message: str | None = None err_category: RolloutErrorCategory | None = None diff --git a/osmosis_ai/rollout/utils/rewards.py b/osmosis_ai/rollout/utils/rewards.py index d736d638..98dc4b19 100644 --- a/osmosis_ai/rollout/utils/rewards.py +++ b/osmosis_ai/rollout/utils/rewards.py @@ -1,10 +1,18 @@ from osmosis_ai.rollout.types import RolloutSample -def validate_samples_have_rewards(samples: dict[str, RolloutSample]) -> None: - if not samples: - raise ValueError("No samples after grading") +def validate_sample_has_reward(sample: RolloutSample | None) -> None: + """Raise if the rollout did not produce a graded sample. - for sid, sample in samples.items(): - if sample.reward is None: - raise ValueError(f"Sample {sid} has no reward after grading") + A rollout has at most one sample; an ungraded or missing sample is a + grader contract violation (or a workflow that forgot to register an + agent/session, in which case ``sample`` is ``None``). + """ + if sample is None: + raise ValueError( + "No sample to grade. The workflow must register a sample " + "source (e.g. via OsmosisStrandsAgent or OsmosisMemorySession) " + "before grading." + ) + if not sample.remove_sample and sample.reward is None: + raise ValueError("Sample has no reward after grading") From d725aa0ec334110447e3b5349b51ab6bd637e38b Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Fri, 22 May 2026 15:25:29 -0700 Subject: [PATCH 2/4] [fix] example --- examples/multiply_openai_agents/README.md | 104 -------------- examples/multiply_openai_agents/main.py | 131 ++---------------- .../multiply_rollout/__init__.py | 0 .../multiply_rollout/grader.py | 55 ++++++++ .../multiply_rollout/tools.py | 7 + .../multiply_rollout/utils.py | 9 ++ .../multiply_rollout/workflow.py | 50 +++++++ 7 files changed, 130 insertions(+), 226 deletions(-) delete mode 100644 examples/multiply_openai_agents/README.md create mode 100644 examples/multiply_openai_agents/multiply_rollout/__init__.py create mode 100644 examples/multiply_openai_agents/multiply_rollout/grader.py create mode 100644 examples/multiply_openai_agents/multiply_rollout/tools.py create mode 100644 examples/multiply_openai_agents/multiply_rollout/utils.py create mode 100644 examples/multiply_openai_agents/multiply_rollout/workflow.py diff --git a/examples/multiply_openai_agents/README.md b/examples/multiply_openai_agents/README.md deleted file mode 100644 index b284cbd5..00000000 --- a/examples/multiply_openai_agents/README.md +++ /dev/null @@ -1,104 +0,0 @@ -# multiply_openai_agents - -Minimal end-to-end Osmosis rollout server using the OpenAI Agents SDK. - -A single-file FastAPI server (`main.py`) that exposes the standard rollout -contract (`POST /rollout`, `GET /health`). On each rollout the workflow -spins up an `OsmosisAgent` with one `multiply` tool, runs it against the -trainer's session URL via LiteLLM, and the grader scores the final -`Answer: ` against the dataset label. - -Use this as the smoke test when you bring up multi-LoRA + remote rollout -on a fresh cluster. - -## Run standalone (sanity check) - -From this directory: - -```bash -uv venv .venv && source .venv/bin/activate -uv pip install -e '../..[server]' 'openai-agents>=0.14,<0.15' uvicorn - -python main.py & -curl http://localhost:8080/health # -> {"status":"ok", ...} -kill %1 -``` - -Standalone the server boots but won't drive a rollout end-to-end — there's -no chat-completions URL to call. The full flow is driven by the trainer -once you register an adapter that points at it. - -## Connect to the multi-LoRA + remote-rollout cluster - -Assumes you've already done `traingate multi-lora start ...` and the head -pod is reachable. Get the head pod name: - -```bash -POD=$(kubectl get pods -l 'skypilot-cluster=' -o name | head -n1) -``` - -Copy the example into the pod and start it on `localhost:8080`: - -```bash -tar czf /tmp/multiply.tgz -C .. multiply_openai_agents -kubectl cp /tmp/multiply.tgz "$POD":/tmp/multiply.tgz -kubectl exec "$POD" -- bash -lc ' - set -e - mkdir -p /tmp/multiply && tar xzf /tmp/multiply.tgz -C /tmp/multiply - cd /tmp/multiply/multiply_openai_agents - pip install -q -e /osmosis/workspace/sdk-python"[server]" \ - "openai-agents>=0.14,<0.15" - OPENAI_API_KEY='"$OPENAI_API_KEY"' ROLLOUT_PORT=8080 \ - nohup python main.py > /tmp/multiply.log 2>&1 & -' -kubectl exec "$POD" -- curl -fsS http://localhost:8080/health -``` - -Register an adapter pointing at it (run from `osmosis-traingate/`): - -```toml -# /tmp/multiply.toml -[config] -data = "/osmosis/data/datasets/multiply/train.parquet" -input_key = "messages" -label_key = "label" -rm_type = "math" -rank = 16 -alpha = 16 -num_row = 20 - -[config.metadata] -rollout_server_url = "http://localhost:8080" -``` - -```bash -uv run traingate multi-lora register multiply \ - --cluster --spec /tmp/multiply.toml -``` - -Tail both sides; you should see chat completions flowing within seconds: - -```bash -# rollout server -kubectl exec "$POD" -- tail -f /tmp/multiply.log - -# trainer -sky logs --follow | rg -i 'rollout|adapter|session-server' -``` - -## Dataset shape - -Each row in the parquet file needs the keys named in the adapter spec -(`input_key`, `label_key`). For this example: - -```python -{ - "messages": [ - {"role": "user", "content": "What is 23 times 41?"}, - ], - "label": "943", -} -``` - -The workflow treats `ctx.prompt` as the verbatim message list to seed the -agent's session. diff --git a/examples/multiply_openai_agents/main.py b/examples/multiply_openai_agents/main.py index 85c633a2..582728e6 100644 --- a/examples/multiply_openai_agents/main.py +++ b/examples/multiply_openai_agents/main.py @@ -1,139 +1,26 @@ -"""Multiply rollout server: minimal end-to-end Osmosis rollout example. - -Self-contained FastAPI rollout server that drives a small OpenAI Agents SDK -workflow against the active rollout's chat-completions endpoint. The agent -has one tool (``multiply``) and a numeric grader that scores the final -answer against the dataset label. - -Bring up the multi-LoRA + remote-rollout cluster (see -``osmosis-traingate/specs/experiment/miles-multi-lora-service.toml``), then -point an adapter at this server via ``metadata.rollout_server_url`` and the -trainer's custom generate function will POST rollouts here. - -Standalone boot (sanity check): - - pip install -e '../..[server]' 'openai-agents>=0.14,<0.15' uvicorn - python main.py & - curl http://localhost:8080/health -""" +"""Local OpenAI Agents multiply rollout server.""" from __future__ import annotations -import logging import os -import re -from typing import Any import uvicorn -from agents import ModelSettings, Runner, function_tool -from osmosis_ai.rollout.agent_workflow import AgentWorkflow +from multiply_rollout.grader import MultiplyGrader, multiply_grader_config +from multiply_rollout.workflow import MultiplyWorkflow, multiply_workflow_config from osmosis_ai.rollout.backend.local import LocalBackend -from osmosis_ai.rollout.context import AgentWorkflowContext, GraderContext -from osmosis_ai.rollout.grader import Grader -from osmosis_ai.rollout.integrations.agents.openai_agents import ( - OsmosisAgent, - OsmosisMemorySession, - OsmosisRolloutModel, -) from osmosis_ai.rollout.server import create_rollout_server -logger = logging.getLogger(__name__) - -MAX_TURNS = 8 -# Loose tolerance so float-y answers like "12.0" match integer labels. -ANSWER_TOLERANCE = 1e-2 - - -@function_tool -async def multiply(a: float, b: float) -> float: - """Multiply two numbers.""" - return a * b - - -class MultiplyWorkflow(AgentWorkflow): - async def run(self, ctx: AgentWorkflowContext) -> None: - agent = OsmosisAgent( - name="multiply", - instructions=( - "Solve the user's multiplication problem. You may call the " - "`multiply` tool. End your final message with " - "`Answer: ` on its own line." - ), - model=OsmosisRolloutModel(), - model_settings=ModelSettings(temperature=1.0, max_tokens=512), - tools=[multiply], - ) - session = OsmosisMemorySession() - await Runner.run(agent, ctx.prompt, session=session, max_turns=MAX_TURNS) - - -_ANSWER_RE = re.compile(r"Answer:\s*(-?\d+(?:\.\d+)?)", re.IGNORECASE) - - -def _extract_text(content: Any) -> str: - """Flatten OpenAI Responses-API content into plain text. - - ``content`` can be a str, or a list of typed parts; we keep the text-ish - parts (``text`` / ``output_text``) and concatenate. - """ - if isinstance(content, str): - return content - if isinstance(content, list): - return "\n".join( - item.get("text", "") - for item in content - if isinstance(item, dict) and item.get("type") in ("text", "output_text") - ) - return "" - - -class MultiplyGrader(Grader): - """0/1 reward: match the trailing ``Answer: `` against ``ctx.label``.""" - - async def grade(self, ctx: GraderContext) -> None: - if ctx.sample is None or ctx.label is None: - ctx.set_reward(0.0) - return - - # Walk back to the most recent assistant turn with text content; - # skips tool-call / tool-result turns that have no human-readable body. - text = "" - for msg in reversed(ctx.sample.messages): - if msg.get("role") != "assistant": - continue - text = _extract_text(msg.get("content", "")) - if text: - break - - match = _ANSWER_RE.search(text) - if match is None: - ctx.set_reward(0.0) - return - - try: - predicted = float(match.group(1)) - target = float(ctx.label) - except ValueError: - ctx.set_reward(0.0) - return - - reward = 1.0 if abs(predicted - target) < ANSWER_TOLERANCE else 0.0 - ctx.set_reward(reward) - logger.info( - "multiply reward=%.1f predicted=%s target=%s", - reward, predicted, target, - ) - def main() -> None: - logging.basicConfig( - level=logging.INFO, format="[multiply-rollout] %(message)s" + backend = LocalBackend( + workflow=MultiplyWorkflow, + workflow_config=multiply_workflow_config, + grader=MultiplyGrader, + grader_config=multiply_grader_config, ) - backend = LocalBackend(workflow=MultiplyWorkflow, grader=MultiplyGrader) app = create_rollout_server(backend=backend) - port = int(os.environ.get("ROLLOUT_PORT", "8080")) - uvicorn.run(app, host="0.0.0.0", port=port) + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("_OSMOSIS_ROLLOUT_PORT", "8000"))) if __name__ == "__main__": diff --git a/examples/multiply_openai_agents/multiply_rollout/__init__.py b/examples/multiply_openai_agents/multiply_rollout/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/multiply_openai_agents/multiply_rollout/grader.py b/examples/multiply_openai_agents/multiply_rollout/grader.py new file mode 100644 index 00000000..03000e37 --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/grader.py @@ -0,0 +1,55 @@ +import logging +from typing import Any + +from multiply_rollout.utils import extract_solution +from osmosis_ai.rollout.context import GraderContext +from osmosis_ai.rollout.grader import Grader +from osmosis_ai.rollout.types import GraderConfig + +logger = logging.getLogger(__name__) + + +class MultiplyGraderConfig(GraderConfig): + name: str = "MultiplyOpenAIAgentsGrader" + description: str = "Grades multiplication rollouts using OpenAI Agents" + + +multiply_grader_config = MultiplyGraderConfig() + + +class MultiplyGrader(Grader): + def compute_reward(self, solution_str: str, ground_truth: str) -> float: + extracted = extract_solution(solution_str) + try: + sol_val = float(extracted) + gt_val = float(ground_truth) + except (TypeError, ValueError): + return 0.0 + + if abs(gt_val - sol_val) < 1e-2: + return 1.0 + return 0.0 + + async def grade(self, ctx: GraderContext) -> None: + if ctx.sample is None: + ctx.set_reward(0.0) + return + + last_message = ctx.sample.messages[-1] + content = last_message.get("content", "") + reward = self.compute_reward(_extract_text(content), ctx.label or "") + ctx.set_reward(reward) + logger.info("[MultiplyGrader] reward = %.1f", reward) + + +def _extract_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + item.get("text", "") + for item in content + if isinstance(item, dict) and item.get("type") == "output_text" + ] + return "\n".join(parts) + return "" diff --git a/examples/multiply_openai_agents/multiply_rollout/tools.py b/examples/multiply_openai_agents/multiply_rollout/tools.py new file mode 100644 index 00000000..b6ce0c65 --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/tools.py @@ -0,0 +1,7 @@ +from agents import function_tool + + +@function_tool +async def multiply_tool(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b diff --git a/examples/multiply_openai_agents/multiply_rollout/utils.py b/examples/multiply_openai_agents/multiply_rollout/utils.py new file mode 100644 index 00000000..c008ed12 --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/utils.py @@ -0,0 +1,9 @@ +import re + + +def extract_solution(solution_str: str) -> str | None: + """Extract a final answer from the expected #### format.""" + solution = re.search(r"####\s*([-+]?\d*\.?\d+)", solution_str) + if not solution: + return None + return solution.group(1) diff --git a/examples/multiply_openai_agents/multiply_rollout/workflow.py b/examples/multiply_openai_agents/multiply_rollout/workflow.py new file mode 100644 index 00000000..d44ada0f --- /dev/null +++ b/examples/multiply_openai_agents/multiply_rollout/workflow.py @@ -0,0 +1,50 @@ +from typing import Any + +from agents import ModelSettings, Runner +from agents.models.interface import Model + +from multiply_rollout.tools import multiply_tool +from osmosis_ai.rollout.agent_workflow import AgentWorkflow +from osmosis_ai.rollout.context import AgentWorkflowContext +from osmosis_ai.rollout.integrations.agents.openai_agents import ( + OsmosisAgent, + OsmosisMemorySession, + OsmosisRolloutModel, +) +from osmosis_ai.rollout.types import AgentWorkflowConfig + +MAX_TURNS = 8 + + +class MultiplyAgentWorkflowConfig(AgentWorkflowConfig): + name: str = "MultiplyAgentWorkflow" + description: str = "Multiply two numbers using OpenAI Agents" + model: Model + model_settings: ModelSettings + tools: Any + + +multiply_workflow_config = MultiplyAgentWorkflowConfig( + model=OsmosisRolloutModel(), + model_settings=ModelSettings(temperature=1.0, top_p=1.0, max_tokens=4096), + tools=[multiply_tool], +) + + +class MultiplyWorkflow(AgentWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> None: + config = ctx.config + agent = OsmosisAgent( + name="multiply", + model=config.model, + model_settings=config.model_settings, + tools=config.tools, + ) + + session = OsmosisMemorySession() + await Runner.run( + agent, + ctx.prompt, + session=session, + max_turns=MAX_TURNS, + ) From 0b59fa1e358588260ce4230180144cd8b3a94345 Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sat, 23 May 2026 10:05:24 -0700 Subject: [PATCH 3/4] [misc] strands example --- examples/multiply_strands/main.py | 27 +++++++++ .../multiply_rollout/__init__.py | 0 .../multiply_rollout/grader.py | 55 +++++++++++++++++++ .../multiply_rollout/tools.py | 7 +++ .../multiply_rollout/utils.py | 9 +++ .../multiply_rollout/workflow.py | 53 ++++++++++++++++++ 6 files changed, 151 insertions(+) create mode 100644 examples/multiply_strands/main.py create mode 100644 examples/multiply_strands/multiply_rollout/__init__.py create mode 100644 examples/multiply_strands/multiply_rollout/grader.py create mode 100644 examples/multiply_strands/multiply_rollout/tools.py create mode 100644 examples/multiply_strands/multiply_rollout/utils.py create mode 100644 examples/multiply_strands/multiply_rollout/workflow.py diff --git a/examples/multiply_strands/main.py b/examples/multiply_strands/main.py new file mode 100644 index 00000000..5c0ff66d --- /dev/null +++ b/examples/multiply_strands/main.py @@ -0,0 +1,27 @@ +"""Local Strands multiply rollout server.""" + +from __future__ import annotations + +import os + +import uvicorn + +from multiply_rollout.grader import MultiplyGrader, multiply_grader_config +from multiply_rollout.workflow import MultiplyWorkflow, multiply_workflow_config +from osmosis_ai.rollout.backend.local import LocalBackend +from osmosis_ai.rollout.server import create_rollout_server + + +def main() -> None: + backend = LocalBackend( + workflow=MultiplyWorkflow, + workflow_config=multiply_workflow_config, + grader=MultiplyGrader, + grader_config=multiply_grader_config, + ) + app = create_rollout_server(backend=backend) + uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("_OSMOSIS_ROLLOUT_PORT", "8000"))) + + +if __name__ == "__main__": + main() diff --git a/examples/multiply_strands/multiply_rollout/__init__.py b/examples/multiply_strands/multiply_rollout/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/examples/multiply_strands/multiply_rollout/grader.py b/examples/multiply_strands/multiply_rollout/grader.py new file mode 100644 index 00000000..39eec48e --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/grader.py @@ -0,0 +1,55 @@ +import logging +from typing import Any + +from multiply_rollout.utils import extract_solution +from osmosis_ai.rollout.context import GraderContext +from osmosis_ai.rollout.grader import Grader +from osmosis_ai.rollout.types import GraderConfig + +logger = logging.getLogger(__name__) + + +class MultiplyGraderConfig(GraderConfig): + name: str = "MultiplyStrandsGrader" + description: str = "Grades multiplication rollouts using Strands" + + +multiply_grader_config = MultiplyGraderConfig() + + +class MultiplyGrader(Grader): + def compute_reward(self, solution_str: str, ground_truth: str) -> float: + extracted = extract_solution(solution_str) + try: + sol_val = float(extracted) + gt_val = float(ground_truth) + except (TypeError, ValueError): + return 0.0 + + if abs(gt_val - sol_val) < 1e-2: + return 1.0 + return 0.0 + + async def grade(self, ctx: GraderContext) -> None: + if ctx.sample is None: + ctx.set_reward(0.0) + return + + last_message = ctx.sample.messages[-1] + content = last_message.get("content", "") + reward = self.compute_reward(_extract_text(content), ctx.label or "") + ctx.set_reward(reward) + logger.info("[MultiplyGrader] reward = %.1f", reward) + + +def _extract_text(content: Any) -> str: + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [ + item.get("text", "") + for item in content + if isinstance(item, dict) and "text" in item + ] + return "\n".join(parts) + return "" diff --git a/examples/multiply_strands/multiply_rollout/tools.py b/examples/multiply_strands/multiply_rollout/tools.py new file mode 100644 index 00000000..a97079a5 --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/tools.py @@ -0,0 +1,7 @@ +from strands import tool + + +@tool +def multiply_tool(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b diff --git a/examples/multiply_strands/multiply_rollout/utils.py b/examples/multiply_strands/multiply_rollout/utils.py new file mode 100644 index 00000000..c008ed12 --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/utils.py @@ -0,0 +1,9 @@ +import re + + +def extract_solution(solution_str: str) -> str | None: + """Extract a final answer from the expected #### format.""" + solution = re.search(r"####\s*([-+]?\d*\.?\d+)", solution_str) + if not solution: + return None + return solution.group(1) diff --git a/examples/multiply_strands/multiply_rollout/workflow.py b/examples/multiply_strands/multiply_rollout/workflow.py new file mode 100644 index 00000000..c6f95b39 --- /dev/null +++ b/examples/multiply_strands/multiply_rollout/workflow.py @@ -0,0 +1,53 @@ +from typing import Any + +from strands.models.model import Model + +from multiply_rollout.tools import multiply_tool +from osmosis_ai.rollout.agent_workflow import AgentWorkflow +from osmosis_ai.rollout.context import AgentWorkflowContext +from osmosis_ai.rollout.integrations.agents.strands import ( + OsmosisRolloutModel, + OsmosisStrandsAgent, +) +from osmosis_ai.rollout.types import AgentWorkflowConfig + +MAX_ITERATIONS = 8 + + +class MultiplyAgentWorkflowConfig(AgentWorkflowConfig): + name: str = "MultiplyAgentWorkflow" + description: str = "Multiply two numbers using Strands" + model: Model + tools: Any + + +multiply_workflow_config = MultiplyAgentWorkflowConfig( + model=OsmosisRolloutModel(params={"temperature": 1.0, "top_p": 1.0, "max_tokens": 4096}), + tools=[multiply_tool], +) + + +class MultiplyWorkflow(AgentWorkflow): + async def run(self, ctx: AgentWorkflowContext) -> None: + config = ctx.config + system_prompt, user_prompt = _split_prompt(ctx.prompt) + agent = OsmosisStrandsAgent( + name="multiply", + model=config.model, + tools=config.tools, + system_prompt=system_prompt, + ) + await agent.invoke_async(user_prompt) + + +def _split_prompt(prompt: list[dict[str, Any]]) -> tuple[str | None, str]: + system: str | None = None + user_parts: list[str] = [] + for msg in prompt: + role = msg.get("role") + content = msg.get("content", "") + if role == "system" and isinstance(content, str): + system = content + elif role == "user" and isinstance(content, str): + user_parts.append(content) + return system, "\n".join(user_parts) From b934f186cbbda1159802a5034d6d288bcb52f50d Mon Sep 17 00:00:00 2001 From: Mathew Han Date: Sun, 24 May 2026 00:08:27 -0700 Subject: [PATCH 4/4] [fix] strands --- .../multiply_rollout/workflow.py | 33 ++++++++++--------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/examples/multiply_strands/multiply_rollout/workflow.py b/examples/multiply_strands/multiply_rollout/workflow.py index c6f95b39..dabc4822 100644 --- a/examples/multiply_strands/multiply_rollout/workflow.py +++ b/examples/multiply_strands/multiply_rollout/workflow.py @@ -1,8 +1,10 @@ from typing import Any +from strands.agent.agent_result import AgentResult from strands.models.model import Model from multiply_rollout.tools import multiply_tool +from multiply_rollout.utils import extract_solution from osmosis_ai.rollout.agent_workflow import AgentWorkflow from osmosis_ai.rollout.context import AgentWorkflowContext from osmosis_ai.rollout.integrations.agents.strands import ( @@ -28,26 +30,25 @@ class MultiplyAgentWorkflowConfig(AgentWorkflowConfig): class MultiplyWorkflow(AgentWorkflow): + async def check_done(self, result: AgentResult) -> bool: + content = result.message.get("content", "") + if not any("toolUse" in block for block in content): + return True + text_content = next((block for block in content if block.get("text")), None) + if text_content: + return extract_solution(text_content.get("text", "")) is not None + return False + async def run(self, ctx: AgentWorkflowContext) -> None: config = ctx.config - system_prompt, user_prompt = _split_prompt(ctx.prompt) agent = OsmosisStrandsAgent( name="multiply", model=config.model, tools=config.tools, - system_prompt=system_prompt, + messages=ctx.prompt, + callback_handler=None, ) - await agent.invoke_async(user_prompt) - - -def _split_prompt(prompt: list[dict[str, Any]]) -> tuple[str | None, str]: - system: str | None = None - user_parts: list[str] = [] - for msg in prompt: - role = msg.get("role") - content = msg.get("content", "") - if role == "system" and isinstance(content, str): - system = content - elif role == "user" and isinstance(content, str): - user_parts.append(content) - return system, "\n".join(user_parts) + for _ in range(MAX_ITERATIONS): + result = await agent.invoke_async() + if await self.check_done(result): + break