| |
|
|
| """Rollout collection helpers for synthetic dataset generation. |
| |
| Complements ``openenv.core.harness`` with: |
| |
| - ``EpisodeRecord`` — serializable view of one rollout + its verification. |
| - ``RolloutSerializer`` — append-only JSONL writer with metadata sidecar. |
| - ``CollectRunner`` — orchestrates repeated rollouts, supports resume, and |
| optional per-record filtering via ``should_keep``. |
| |
| The serialized schema is designed to be consumed directly by TRL's |
| ``SFTTrainer`` (``messages`` column) or by ``datasets.Dataset.from_json``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| import functools |
| import json |
| import threading |
| import warnings |
| from dataclasses import asdict, dataclass, field |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any, BinaryIO, Callable, Iterable, Iterator |
|
|
| from ..env_server.mcp_types import Tool |
| from ..llm_client import LLMClient |
| from .rollout import ( |
| _resolve_env_reward, |
| HarnessAdapter, |
| HarnessRolloutResult, |
| HarnessRunLimits, |
| ModelStep, |
| ModelStepResult, |
| ResourceSessionFactory, |
| VerifyResult, |
| ) |
|
|
| _SUPPORTED_SAMPLING_KEYS = frozenset({"temperature", "max_tokens", "top_p"}) |
|
|
| RESULTS_FILENAME = "results.jsonl" |
| METADATA_FILENAME = "metadata.json" |
|
|
|
|
| def _tool_trace_entry_to_plain(entry: Any) -> dict[str, Any]: |
| return { |
| "tool_name": entry.tool_name, |
| "arguments": dict(entry.arguments), |
| "result": { |
| "data": entry.result.data, |
| "done": entry.result.done, |
| "metadata": dict(entry.result.metadata), |
| "error": entry.result.error, |
| }, |
| } |
|
|
|
|
| def _tool_trace_to_plain(rollout: HarnessRolloutResult) -> list[dict[str, Any]]: |
| return [_tool_trace_entry_to_plain(entry) for entry in rollout.tool_trace] |
|
|
|
|
| @dataclass |
| class EpisodeRecord: |
| """Serializable view of one collected episode.""" |
|
|
| episode_id: str |
| messages: list[dict[str, Any]] |
| reward: float |
| done: bool |
| tool_trace: list[dict[str, Any]] |
| metrics: dict[str, Any] |
| verify_metrics: dict[str, Any] |
| artifacts: dict[str, Any] |
| task: Any = None |
| extra: dict[str, Any] = field(default_factory=dict) |
|
|
| @classmethod |
| def from_rollout( |
| cls, |
| episode_id: str, |
| rollout: HarnessRolloutResult, |
| verify: VerifyResult, |
| task: Any = None, |
| extra: dict[str, Any] | None = None, |
| ) -> "EpisodeRecord": |
| """Build a record from a harness rollout and its verification. |
| |
| Uses ``_resolve_env_reward`` so that any disagreement between the |
| reward emitted inside the environment and the one forwarded by |
| ``verify()`` raises — preserving the "rewards in env" invariant. |
| """ |
| reward = _resolve_env_reward(rollout, verify) |
| return cls( |
| episode_id=episode_id, |
| messages=[dict(m) for m in rollout.messages], |
| reward=reward, |
| done=bool(rollout.done or verify.done), |
| tool_trace=_tool_trace_to_plain(rollout), |
| metrics=dict(rollout.metrics), |
| verify_metrics=dict(verify.metrics), |
| artifacts=dict(verify.artifacts), |
| task=task, |
| extra=dict(extra or {}), |
| ) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return json.loads(json.dumps(asdict(self), default=str)) |
|
|
|
|
| class RolloutSerializer: |
| """Append-only JSONL writer for collected episodes.""" |
|
|
| def __init__(self, output_dir: Path | str): |
| self._output_dir = Path(output_dir) |
|
|
| @property |
| def output_dir(self) -> Path: |
| return self._output_dir |
|
|
| @property |
| def results_path(self) -> Path: |
| return self._output_dir / RESULTS_FILENAME |
|
|
| @property |
| def metadata_path(self) -> Path: |
| return self._output_dir / METADATA_FILENAME |
|
|
| def write_episode(self, record: EpisodeRecord) -> None: |
| self._output_dir.mkdir(parents=True, exist_ok=True) |
| line = (json.dumps(record.to_dict(), default=str) + "\n").encode("utf-8") |
| with self.results_path.open("a+b") as handle: |
| separator = self._append_separator(handle) |
| handle.write(separator + line) |
|
|
| def write_metadata(self, metadata: dict[str, Any]) -> None: |
| self._output_dir.mkdir(parents=True, exist_ok=True) |
| self.metadata_path.write_text( |
| json.dumps(metadata, indent=2, default=str), |
| encoding="utf-8", |
| ) |
|
|
| def collected_episode_ids(self) -> set[str]: |
| """Return episode ids already persisted on disk. Used for resume. |
| |
| Raises: |
| `ValueError`: If a non-empty line is not a valid episode record. |
| """ |
| if not self.results_path.exists(): |
| return set() |
|
|
| ids: set[str] = set() |
| with self.results_path.open("rb") as handle: |
| for line_number, raw in enumerate(handle, start=1): |
| raw = raw.strip() |
| if not raw: |
| continue |
| ids.add(self._episode_id_from_line(raw, line_number)) |
| return ids |
|
|
| def _append_separator(self, handle: BinaryIO) -> bytes: |
| """Return a newline when an existing complete final record lacks one.""" |
| end = handle.seek(0, 2) |
| if not end: |
| return b"" |
| handle.seek(-1, 2) |
| if handle.read(1) == b"\n": |
| return b"" |
|
|
| chunks: list[bytes] = [] |
| cursor = end |
| while cursor: |
| read_size = min(cursor, 8192) |
| cursor -= read_size |
| handle.seek(cursor) |
| chunk = handle.read(read_size) |
| newline = chunk.rfind(b"\n") |
| if newline >= 0: |
| chunks.append(chunk[newline + 1 :]) |
| break |
| chunks.append(chunk) |
| final_line = b"".join(reversed(chunks)) |
| if final_line.strip(): |
| self._episode_id_from_line(final_line, "final line") |
| return b"\n" |
|
|
| def _episode_id_from_line(self, raw: bytes, line_number: int | str) -> str: |
| """Parse and validate one non-empty JSONL episode record.""" |
| location = f"{self.results_path}:{line_number}" |
| try: |
| line = raw.decode("utf-8") |
| except UnicodeDecodeError as exc: |
| raise ValueError(f"{location}: invalid UTF-8: {exc}") from exc |
| try: |
| payload = json.loads(line) |
| except json.JSONDecodeError as exc: |
| raise ValueError(f"{location}: invalid JSON: {exc.msg}") from exc |
| if not isinstance(payload, dict): |
| raise ValueError(f"{location}: expected an object") |
| episode_id = payload.get("episode_id") |
| if not isinstance(episode_id, str): |
| raise ValueError(f"{location}: expected a string episode_id") |
| return episode_id |
|
|
|
|
| def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]: |
| """Build the structured payload passed to ``session.verify()``. |
| |
| Mirrors the shape used by ``build_harness_rollout_func`` so verifiers |
| can rely on the same contract regardless of caller (TRL vs. collector). |
| """ |
| return { |
| "done": rollout.done, |
| "metrics": dict(rollout.metrics), |
| "events": [ |
| {"type": event.type, "payload": dict(event.payload)} |
| for event in rollout.events |
| ], |
| "tool_trace": _tool_trace_to_plain(rollout), |
| } |
|
|
|
|
| @dataclass |
| class CollectResult: |
| """Summary returned by ``CollectRunner.run()``.""" |
|
|
| num_collected: int |
| num_skipped: int |
| num_dropped: int |
| num_failed: int |
| episode_ids: list[str] |
| avg_reward: float |
| success_rate: float |
|
|
|
|
| class CollectRunner: |
| """Drive repeated rollouts into an append-only JSONL dataset. |
| |
| The runner is stateless across ``run()`` calls — resume state lives on |
| disk via the ``RolloutSerializer``. A single runner instance can be |
| reused across runs targeting different output dirs. |
| """ |
|
|
| def __init__( |
| self, |
| *, |
| session_factory: ResourceSessionFactory, |
| harness_adapter: HarnessAdapter, |
| serializer: RolloutSerializer, |
| tasks: Iterable[Any] | None = None, |
| limits: HarnessRunLimits | None = None, |
| ): |
| self._session_factory = session_factory |
| self._harness_adapter = harness_adapter |
| self._serializer = serializer |
| self._tasks: Iterator[Any] | None = iter(tasks) if tasks is not None else None |
| self._limits = limits |
|
|
| def _next_task(self) -> Any: |
| if self._tasks is None: |
| return None |
| try: |
| return next(self._tasks) |
| except StopIteration: |
| return None |
|
|
| def run( |
| self, |
| *, |
| model_step: ModelStep, |
| num_episodes: int, |
| episode_id_prefix: str = "ep", |
| resume: bool = True, |
| should_keep: Callable[[EpisodeRecord], bool] | None = None, |
| ) -> CollectResult: |
| planned_ids = [f"{episode_id_prefix}-{i:06d}" for i in range(num_episodes)] |
|
|
| if resume: |
| already = self._serializer.collected_episode_ids() |
| else: |
| self._reset_results_file() |
| already = set() |
|
|
| num_skipped = 0 |
| num_collected = 0 |
| num_dropped = 0 |
| num_failed = 0 |
| rewards: list[float] = [] |
|
|
| from rich.progress import ( |
| BarColumn, |
| MofNCompleteColumn, |
| Progress, |
| TextColumn, |
| TimeElapsedColumn, |
| ) |
|
|
| progress = Progress( |
| TextColumn("[cyan]{task.description}"), |
| BarColumn(), |
| MofNCompleteColumn(), |
| TextColumn( |
| "collected={task.fields[collected]} reward={task.fields[avg_reward]:.2f}" |
| ), |
| TimeElapsedColumn(), |
| ) |
| task_id = progress.add_task( |
| "Collecting", total=len(planned_ids), collected=0, avg_reward=0.0 |
| ) |
|
|
| with progress: |
| for episode_id in planned_ids: |
| |
| task = self._next_task() |
| if episode_id in already: |
| num_skipped += 1 |
| progress.advance(task_id) |
| continue |
|
|
| session = None |
| try: |
| session = self._session_factory.create( |
| task=task, |
| episode_id=episode_id, |
| ) |
| rollout = self._harness_adapter.run_white_box( |
| model_step=model_step, |
| session=session, |
| limits=self._limits, |
| ) |
| verify = session.verify( |
| transcript=rollout.messages, |
| final_state=_rollout_final_state(rollout), |
| ) |
| record = EpisodeRecord.from_rollout( |
| episode_id=episode_id, |
| rollout=rollout, |
| verify=verify, |
| task=task, |
| ) |
| except Exception as exc: |
| num_failed += 1 |
| progress.console.print( |
| f"[red]failed[/red] {episode_id}: {type(exc).__name__}: {exc}" |
| ) |
| progress.advance(task_id) |
| continue |
| finally: |
| if session is not None: |
| session.close() |
|
|
| if should_keep is not None and not should_keep(record): |
| num_dropped += 1 |
| progress.advance(task_id) |
| continue |
|
|
| self._serializer.write_episode(record) |
| num_collected += 1 |
| rewards.append(record.reward) |
| avg = sum(rewards) / len(rewards) |
| progress.update( |
| task_id, advance=1, collected=num_collected, avg_reward=avg |
| ) |
|
|
| avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 |
| success_rate = ( |
| sum(1 for r in rewards if r > 0) / len(rewards) if rewards else 0.0 |
| ) |
|
|
| return CollectResult( |
| num_collected=num_collected, |
| num_skipped=num_skipped, |
| num_dropped=num_dropped, |
| num_failed=num_failed, |
| episode_ids=planned_ids, |
| avg_reward=avg_reward, |
| success_rate=success_rate, |
| ) |
|
|
| def _reset_results_file(self) -> None: |
| if self._serializer.results_path.exists(): |
| results_path = self._serializer.results_path |
| timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") |
| backup_path = results_path.with_name( |
| f"{results_path.stem}.{timestamp}.bak{results_path.suffix}" |
| ) |
| counter = 1 |
| while backup_path.exists(): |
| backup_path = results_path.with_name( |
| f"{results_path.stem}.{timestamp}.{counter}.bak" |
| f"{results_path.suffix}" |
| ) |
| counter += 1 |
| results_path.replace(backup_path) |
| warnings.warn( |
| f"resume=False moved existing results file to {backup_path}", |
| RuntimeWarning, |
| stacklevel=2, |
| ) |
|
|
|
|
| def _tool_to_mcp_dict(tool: Tool) -> dict[str, Any]: |
| return { |
| "name": tool.name, |
| "description": tool.description, |
| "inputSchema": tool.input_schema, |
| } |
|
|
|
|
| @functools.cache |
| def _model_step_loop() -> asyncio.AbstractEventLoop: |
| """Start the event loop shared by every model step, on its first use.""" |
| loop = asyncio.new_event_loop() |
| threading.Thread(target=loop.run_forever, daemon=True).start() |
| return loop |
|
|
|
|
| def build_model_step( |
| llm_client: LLMClient, |
| *, |
| system_prompt: str | None = None, |
| ) -> ModelStep: |
| """Adapt any OpenEnv ``LLMClient`` into a ``ModelStep`` for the harness. |
| |
| Works with every provider the ``LLMClient`` hierarchy already covers — |
| OpenAI and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, HF |
| Inference, Together, Groq, Fireworks, ...), Anthropic natively, and |
| any future subclass that implements ``complete_with_tools``. Provider |
| -specific schema conversion lives inside the client; this helper only |
| handles the sync/async adaptation and tool dict shape. |
| """ |
| |
| |
| |
| loop = _model_step_loop() |
|
|
| def model_step( |
| messages: list[dict[str, Any]], |
| tools: list[Tool], |
| sampling: dict[str, Any], |
| ) -> ModelStepResult: |
| effective_messages = list(messages) |
| if system_prompt and not any( |
| m.get("role") == "system" for m in effective_messages |
| ): |
| effective_messages.insert(0, {"role": "system", "content": system_prompt}) |
|
|
| tool_dicts = [_tool_to_mcp_dict(t) for t in tools] |
| dropped_keys = sorted(set(sampling) - _SUPPORTED_SAMPLING_KEYS) |
| if dropped_keys: |
| warnings.warn( |
| "Dropping unsupported sampling keys: " + ", ".join(dropped_keys), |
| RuntimeWarning, |
| stacklevel=2, |
| ) |
| filtered_sampling = { |
| k: v for k, v in sampling.items() if k in _SUPPORTED_SAMPLING_KEYS |
| } |
|
|
| response = asyncio.run_coroutine_threadsafe( |
| llm_client.complete_with_tools( |
| messages=effective_messages, |
| tools=tool_dicts, |
| **filtered_sampling, |
| ), |
| loop, |
| ).result() |
| return ModelStepResult(response=response) |
|
|
| return model_step |
|
|
|
|
| def _count_episodes(results_path: Path) -> int: |
| if not results_path.exists(): |
| return 0 |
| with results_path.open("r", encoding="utf-8") as handle: |
| return sum(1 for line in handle if line.strip()) |
|
|
|
|
| README_FILENAME = "README.md" |
|
|
| _DATASET_CARD_FRONTMATTER = """--- |
| configs: |
| - config_name: default |
| data_files: |
| - split: train |
| path: "results.jsonl" |
| --- |
| """ |
|
|
|
|
| def build_dataset_readme(output_dir: Path) -> str: |
| """Build the README.md written alongside ``results.jsonl`` on the Hub. |
| |
| The YAML front-matter pins ``results.jsonl`` as the only dataset file |
| so the HF Dataset Viewer does not try to cast ``metadata.json`` into |
| the same schema. |
| """ |
| output_dir = Path(output_dir) |
| results_path = output_dir / RESULTS_FILENAME |
| metadata_path = output_dir / METADATA_FILENAME |
|
|
| metadata: dict[str, Any] = {} |
| if metadata_path.exists(): |
| try: |
| metadata = json.loads(metadata_path.read_text(encoding="utf-8")) |
| except (json.JSONDecodeError, OSError): |
| metadata = {} |
|
|
| lines: list[str] = [ |
| _DATASET_CARD_FRONTMATTER.strip(), |
| "", |
| "# OpenEnv rollouts", |
| "", |
| "Collected with [OpenEnv](https://github.com/huggingface/OpenEnv) " |
| "(`openenv collect`).", |
| "", |
| f"- **Episodes:** {_count_episodes(results_path)}", |
| ] |
| if metadata: |
| lines.append("- **Run metadata:**") |
| lines.append("") |
| lines.append(" | key | value |") |
| lines.append(" | --- | --- |") |
| for key, value in metadata.items(): |
| lines.append(f" | `{key}` | `{value}` |") |
|
|
| lines.extend( |
| [ |
| "", |
| "## Schema", |
| "", |
| "Each line of `results.jsonl` is one episode:", |
| "", |
| "- `episode_id` (string)", |
| "- `messages` (chat transcript; TRL `SFTTrainer`-compatible)", |
| "- `reward` (float)", |
| "- `done` (bool)", |
| "- `tool_trace` (list of `{tool_name, arguments, result}`)", |
| "- `metrics`, `verify_metrics`, `artifacts`", |
| "- `task`, `extra` (optional call-site annotations)", |
| "", |
| "## Load", |
| "", |
| "```python", |
| "from datasets import load_dataset", |
| 'ds = load_dataset("<user>/<repo>", split="train")', |
| "```", |
| "", |
| ] |
| ) |
|
|
| return "\n".join(lines) |
|
|
|
|
| def push_to_hf_hub( |
| output_dir: Path | str, |
| repo_id: str, |
| *, |
| private: bool = False, |
| commit_message: str | None = None, |
| token: str | None = None, |
| ) -> str: |
| """Upload a collected rollout directory to the Hugging Face Hub as a dataset. |
| |
| Uploads ``results.jsonl``, ``metadata.json`` (if present), and an |
| auto-generated ``README.md`` whose YAML front-matter tells the HF |
| Dataset Viewer to treat ``results.jsonl`` as the sole dataset file. |
| Consumers load the result with ``load_dataset("<user>/<repo>", |
| split="train")``. |
| |
| Args: |
| output_dir: Directory previously populated by ``RolloutSerializer``. |
| repo_id: Destination repo in ``"user/name"`` form. |
| private: Create the dataset repo as private. |
| commit_message: Override the default commit message. |
| token: HF token. Defaults to the token resolved by ``huggingface_hub``. |
| |
| Returns: |
| The public URL of the dataset on the Hub. |
| |
| Raises: |
| FileNotFoundError: If ``results.jsonl`` is missing from ``output_dir``. |
| """ |
| output_dir = Path(output_dir) |
| results_path = output_dir / RESULTS_FILENAME |
| if not results_path.exists(): |
| raise FileNotFoundError( |
| f"No {RESULTS_FILENAME} found in {output_dir}. " |
| "Run a collect job before pushing to the Hub." |
| ) |
|
|
| readme_path = output_dir / README_FILENAME |
| readme_path.write_text(build_dataset_readme(output_dir), encoding="utf-8") |
|
|
| num_episodes = _count_episodes(results_path) |
| message = commit_message or f"Add {num_episodes} collected episode(s)" |
|
|
| from huggingface_hub import HfApi |
|
|
| api = HfApi(token=token) |
| api.create_repo( |
| repo_id=repo_id, |
| repo_type="dataset", |
| private=private, |
| exist_ok=True, |
| ) |
| api.upload_folder( |
| folder_path=str(output_dir), |
| repo_id=repo_id, |
| repo_type="dataset", |
| commit_message=message, |
| ) |
| return f"https://huggingface.co/datasets/{repo_id}" |
|
|
|
|
| __all__ = [ |
| "build_dataset_readme", |
| "build_model_step", |
| "CollectResult", |
| "CollectRunner", |
| "EpisodeRecord", |
| "push_to_hf_hub", |
| "RolloutSerializer", |
| "METADATA_FILENAME", |
| "README_FILENAME", |
| "RESULTS_FILENAME", |
| ] |
|
|
