From cb65ba72355ff590706844b70f361818e7247854 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Fri, 29 May 2026 17:57:46 +0000 Subject: [PATCH 01/16] examples: add terminus async grpo training --- examples/end_to_end/tbench2_pi_trl/README.md | 9 + .../skills/terminus-terminal-task/SKILL.md | 26 ++ .../tbench2_pi_trl/agents/terminus/AGENTS.md | 15 + .../end_to_end/tbench2_pi_trl/pi_harness.py | 83 ++++++ .../tbench2_pi_trl/train_terminus_grpo.py | 259 ++++++++++++++++++ 5 files changed, 392 insertions(+) create mode 100644 examples/end_to_end/tbench2_pi_trl/README.md create mode 100644 examples/end_to_end/tbench2_pi_trl/agents/terminus/.pi/skills/terminus-terminal-task/SKILL.md create mode 100644 examples/end_to_end/tbench2_pi_trl/agents/terminus/AGENTS.md create mode 100644 examples/end_to_end/tbench2_pi_trl/pi_harness.py create mode 100644 examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py diff --git a/examples/end_to_end/tbench2_pi_trl/README.md b/examples/end_to_end/tbench2_pi_trl/README.md new file mode 100644 index 000000000..db426e21b --- /dev/null +++ b/examples/end_to_end/tbench2_pi_trl/README.md @@ -0,0 +1,9 @@ +# Terminus + TRL Async GRPO + +Start a Terminus server and a compatible vLLM server, then run: + +```bash +TERMINUS_ENV_URL=http://localhost:8000 \ +TERMINUS_VLLM_SERVER_URL=http://localhost:8001 \ +uv run train_terminus_grpo.py +``` diff --git a/examples/end_to_end/tbench2_pi_trl/agents/terminus/.pi/skills/terminus-terminal-task/SKILL.md b/examples/end_to_end/tbench2_pi_trl/agents/terminus/.pi/skills/terminus-terminal-task/SKILL.md new file mode 100644 index 000000000..021809664 --- /dev/null +++ b/examples/end_to_end/tbench2_pi_trl/agents/terminus/.pi/skills/terminus-terminal-task/SKILL.md @@ -0,0 +1,26 @@ +--- +name: terminus-terminal-task +description: Use inside a Terminus environment session when solving one sandboxed terminal task with the terminal tool. +--- + +# Terminus Terminal Task + +Use this skill only inside a Terminus task session. + +## Workflow + +1. Read the task. +2. Use the `terminal` tool for each terminal action. +3. Pass `command` to inspect and modify the sandbox. +4. Check command output before choosing the next command. +5. When the task is complete, pass `final_answer` exactly once. + +## Guardrails + +- Do not change hidden checks or task configuration. +- Do not claim completion until the visible task requirements are satisfied. +- Stay focused on the current task and terminal outputs. +- Do not include both `command` and `final_answer` in the same tool call. +- For simple file writes, prefer commands like `printf %s 'text' > path`. +- If a command fails, inspect the error and continue with a smaller diagnostic + command. diff --git a/examples/end_to_end/tbench2_pi_trl/agents/terminus/AGENTS.md b/examples/end_to_end/tbench2_pi_trl/agents/terminus/AGENTS.md new file mode 100644 index 000000000..8defe822c --- /dev/null +++ b/examples/end_to_end/tbench2_pi_trl/agents/terminus/AGENTS.md @@ -0,0 +1,15 @@ +# Terminus Task Instructions + +You are solving one task inside a Terminus terminal environment. + +Use the available `terminal` tool to inspect and modify the sandbox. Prefer +short, direct shell commands. Read command output before deciding the next +action. While working, call `terminal` with a `command`. When done, call +`terminal` with a `final_answer`. Do not include both arguments in the same +tool call. + +When the requested task is complete, submit exactly one final answer. The final +answer should be concise and should not include implementation notes. + +Do not change hidden task checks or environment configuration. Stay focused on +the current task and the available terminal tool. diff --git a/examples/end_to_end/tbench2_pi_trl/pi_harness.py b/examples/end_to_end/tbench2_pi_trl/pi_harness.py new file mode 100644 index 000000000..3e97afbf0 --- /dev/null +++ b/examples/end_to_end/tbench2_pi_trl/pi_harness.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +"""Pi-style context loading for the Terminus async GRPO example.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +Message = dict[str, Any] + +DEFAULT_PI_SYSTEM_PROMPT = """You are running inside a Pi-compatible environment session. + +Follow the project AGENTS.md and available skill instructions. Use the available +tools exactly as described by their schemas. When the task is complete, submit +the environment's final tool call. +""" + + +@dataclass +class PiContext: + """Load model-facing project instructions and skills.""" + + project_root: Path + agents_md_path: str = "AGENTS.md" + skills_dir: str = ".pi/skills" + max_context_chars: int = 6000 + extra_system_prompts: list[str] = field(default_factory=list) + + def messages(self) -> list[Message]: + parts = [DEFAULT_PI_SYSTEM_PROMPT.strip(), *self.extra_system_prompts] + for label, path, text in self.context_files(): + parts.append(f"{label} ({path}):\n{text}") + return [{"role": "system", "content": "\n\n".join(parts)}] + + def loaded_files(self) -> list[str]: + return [path for _, path, _ in self.context_files()] + + def context_files(self) -> list[tuple[str, str, str]]: + files = [ + ("Project instructions", self._resolve(self.agents_md_path)), + *( + (f"Skill: {path.parent.name}", path) + for path in _iter_skill_files(self._resolve(self.skills_dir)) + ), + ] + loaded = [] + for label, path in files: + if text := _read_bounded(path, self.max_context_chars): + loaded.append((label, self._display(path), text)) + return loaded + + def _resolve(self, path: str) -> Path: + resolved = Path(path) + return resolved if resolved.is_absolute() else self.project_root / resolved + + def _display(self, path: Path) -> str: + try: + return path.relative_to(self.project_root).as_posix() + except ValueError: + return path.as_posix() + + +def _iter_skill_files(skills_path: Path) -> list[Path]: + if not skills_path.exists(): + return [] + if skills_path.is_file(): + return [skills_path] + skill_files = set(skills_path.rglob("SKILL.md")) + skill_files.update(path for path in skills_path.glob("*.md") if path.is_file()) + return sorted(skill_files, key=lambda path: path.as_posix()) + + +def _read_bounded(path: Path, max_chars: int) -> str: + if not path.is_file(): + return "" + text = path.read_text(encoding="utf-8").strip() + if len(text) <= max_chars: + return text + return text[:max_chars].rstrip() + "\n\n[truncated by PiContext]" + + +__all__ = ["PiContext"] diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py new file mode 100644 index 000000000..04d03f1e7 --- /dev/null +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -0,0 +1,259 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "aiohttp>=3.9.0", +# "accelerate>=1.10.0", +# "datasets>=3.0.0", +# "huggingface-hub>=0.35.0", +# "openenv-core", +# "openenv-terminus-env", +# "peft>=0.17.0", +# "trackio<0.25.0", +# "transformers @ git+https://github.com/huggingface/transformers.git@e1a37d29cd4822d74f4f3323289fb69e1eec61a0", +# "trl @ git+https://github.com/huggingface/trl.git@a7ba987d05b1e9dbbdbd2e9091264623746e3528", +# "vllm==0.19.1", +# ] +# [tool.uv.sources] +# openenv-core = { git = "https://github.com/burtenshaw/OpenEnv.git", branch = "codex/terminus-pi-trl-space" } +# openenv-terminus-env = { git = "https://github.com/burtenshaw/OpenEnv.git", branch = "codex/terminus-env-harness", subdirectory = "envs/terminus_env" } +# /// +"""Run the Terminus async GRPO environment-factory training example.""" + +from __future__ import annotations + +import json +import os +import re +from pathlib import Path +from typing import Any + +from datasets import load_dataset +from terminus_env.client import TerminusEnv +from terminus_env.harness import TerminusSessionFactory +from transformers import AutoTokenizer +from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer + +TASK_DATASET_ID = "burtenshaw/terminus-pi-trl-tasks" +MODEL = "Qwen/Qwen3.5-4B" +ENV_URL = os.environ.get("TERMINUS_ENV_URL", "http://localhost:8000") +OUTPUT_DIR = Path(os.environ.get("TERMINUS_OUTPUT_DIR", "/tmp/terminus-pi-trl-output")) +HUB_MODEL_ID = "burtenshaw/terminus-pi-trl-async-grpo-qwen35-4b" +TRACKIO_PROJECT = "terminus-pi-trl" +REPORT_TO = "trackio" +RUN_NAME = f"{os.environ.get('JOB_ID', 'local')}-terminus" +VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") +REWARD_RE = re.compile(r"reward=([+-]?(?:\d+(?:\.\d*)?|\.\d+))") + +os.environ["TRACKIO_PROJECT"] = TRACKIO_PROJECT + + +class TerminusHarnessEnvironment: + """Small TRL environment wrapper backed by an OpenEnv resource session.""" + + def __init__(self, session_factory: TerminusSessionFactory): + self._session_factory = session_factory + self._session = None + + def terminal(self, command: str = "", final_answer: str = "") -> str: + """Run a shell command or submit final_answer inside Terminus. + + Args: + command: Shell command to run in the sandbox. + final_answer: Final answer to submit after the task is complete. + + Returns: + A JSON string with the tool output, reward, done flag, and error. + """ + if self._session is None: + return json.dumps( + { + "done": True, + "error": "environment was not reset", + "output": "", + "reward": 0.0, + }, + sort_keys=True, + ) + + command = str(command or "") + final_answer = str(final_answer or "") + if command.strip() and final_answer.strip(): + first = self._session.call_tool("terminal", {"command": command}) + result = first + if not first.done: + result = self._session.call_tool( + "terminal", + {"final_answer": final_answer}, + ) + first_data = first.data if isinstance(first.data, dict) else {} + result_data = result.data if isinstance(result.data, dict) else {} + result.data = { + **result_data, + "output": "\n".join( + str(part) + for part in ( + first_data.get("output"), + result_data.get("output"), + ) + if part + ), + } + elif command.strip(): + result = self._session.call_tool("terminal", {"command": command}) + elif final_answer.strip(): + result = self._session.call_tool( + "terminal", + {"final_answer": final_answer}, + ) + else: + result = self._session.call_tool("terminal", {"command": ""}) + + reward = result.metadata.get("reward") + data = result.data if isinstance(result.data, dict) else {"output": result.data} + return json.dumps( + { + "done": result.done, + "error": result.error, + "output": data.get("output"), + "reward": 0.0 if reward is None else float(reward), + }, + sort_keys=True, + default=str, + ) + + def reset(self, prompt: Any = None, **_: Any) -> None: + self._close_session() + self._session = self._session_factory.create(task=prompt) + + def _close_session(self) -> None: + if self._session is not None: + self._session.close() + self._session = None + + +def terminus_reward(completions=None, **kwargs) -> list[float]: + del kwargs + rewards = [] + for completion in completions or []: + reward = 0.0 + for message in completion if isinstance(completion, list) else []: + if not isinstance(message, dict) or message.get("role") != "tool": + continue + parsed = _parse_tool_reward(str(message.get("content", ""))) + if parsed is not None: + reward = parsed + rewards.append(reward) + return rewards + + +def _parse_tool_reward(content: str) -> float | None: + try: + payload = json.loads(content) + except json.JSONDecodeError: + payload = None + if isinstance(payload, dict) and payload.get("reward") is not None: + try: + return float(payload["reward"]) + except (TypeError, ValueError): + return None + match = REWARD_RE.search(content) + if match is None: + return None + return float(match.group(1)) + + +class TerminusAsyncGRPOTrainer(AsyncGRPOTrainer): + """Map local Qwen3.5 weight names to the vLLM server wrapper.""" + + def _streaming_iter(self): + for name, tensor in super()._streaming_iter(): + yield _vllm_weight_name(name), tensor + + +def _install_vllm_weight_name_adapter() -> None: + import trl.experimental.async_grpo.async_grpo_trainer as trainer_module + + original_worker = trainer_module.AsyncRolloutWorker + if getattr(original_worker, "_terminus_weight_name_adapter", False): + return + + class TerminusRolloutWorker(original_worker): + _terminus_weight_name_adapter = True + + def __init__(self, *args: Any, **kwargs: Any) -> None: + names = kwargs.get("weight_names") + if names is not None: + kwargs = {**kwargs, "weight_names": [_vllm_weight_name(n) for n in names]} + super().__init__(*args, **kwargs) + + trainer_module.AsyncRolloutWorker = TerminusRolloutWorker + + +def _vllm_weight_name(name: str) -> str: + name = name.removeprefix("module.") + if name.startswith(("model.", "lm_head.")): + return f"language_model.{name}" + return name + + +def main() -> None: + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + + task_dataset = load_dataset(TASK_DATASET_ID, split="train") + task = task_dataset[0] + train_dataset = task_dataset.select_columns(["prompt"]) + session_factory = TerminusSessionFactory( + client_factory=lambda: TerminusEnv( + base_url=ENV_URL, + connect_timeout_s=30.0, + message_timeout_s=600.0, + ).sync(), + default_verify=list(task["verify"]), + ) + + tokenizer = AutoTokenizer.from_pretrained(MODEL) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + _install_vllm_weight_name_adapter() + + trainer = TerminusAsyncGRPOTrainer( + model=MODEL, + args=AsyncGRPOConfig( + output_dir=str(OUTPUT_DIR), + max_steps=task["max_steps"], + per_device_train_batch_size=task["batch_size"], + gradient_accumulation_steps=1, + num_generations=task["num_generations"], + max_completion_length=task["max_completion_length"], + max_tool_calling_iterations=task["max_turns"], + max_inflight_tasks=task["num_generations"], + learning_rate=1e-6, + logging_steps=1, + logging_strategy="steps", + log_completions=True, + report_to=REPORT_TO, + run_name=RUN_NAME, + project=TRACKIO_PROJECT, + save_strategy="steps", + save_steps=1, + save_total_limit=2, + push_to_hub=True, + hub_always_push=True, + hub_model_id=HUB_MODEL_ID, + chat_template_kwargs={"enable_thinking": False}, + vllm_server_base_url=VLLM_SERVER_URL, + request_timeout=600, + ), + processing_class=tokenizer, + train_dataset=train_dataset, + reward_funcs=terminus_reward, + environment_factory=lambda: TerminusHarnessEnvironment(session_factory), + ) + trainer.train() + trainer.save_model() + trainer.push_to_hub(commit_message=f"Async GRPO Terminus run {RUN_NAME}") + + +if __name__ == "__main__": + main() From 56fdccbfc3ade310d048969872a615f4c59040ab Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Fri, 29 May 2026 19:07:51 +0000 Subject: [PATCH 02/16] Use Qwen3 4B in Terminus async GRPO example --- .../tbench2_pi_trl/train_terminus_grpo.py | 43 ++----------------- 1 file changed, 4 insertions(+), 39 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index 04d03f1e7..beb740d70 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -35,13 +35,13 @@ from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer TASK_DATASET_ID = "burtenshaw/terminus-pi-trl-tasks" -MODEL = "Qwen/Qwen3.5-4B" +MODEL = "Qwen/Qwen3-4B" ENV_URL = os.environ.get("TERMINUS_ENV_URL", "http://localhost:8000") OUTPUT_DIR = Path(os.environ.get("TERMINUS_OUTPUT_DIR", "/tmp/terminus-pi-trl-output")) -HUB_MODEL_ID = "burtenshaw/terminus-pi-trl-async-grpo-qwen35-4b" +HUB_MODEL_ID = "burtenshaw/terminus-pi-trl-async-grpo-qwen3-4b" TRACKIO_PROJECT = "terminus-pi-trl" REPORT_TO = "trackio" -RUN_NAME = f"{os.environ.get('JOB_ID', 'local')}-terminus" +RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") REWARD_RE = re.compile(r"reward=([+-]?(?:\d+(?:\.\d*)?|\.\d+))") @@ -163,40 +163,6 @@ def _parse_tool_reward(content: str) -> float | None: return float(match.group(1)) -class TerminusAsyncGRPOTrainer(AsyncGRPOTrainer): - """Map local Qwen3.5 weight names to the vLLM server wrapper.""" - - def _streaming_iter(self): - for name, tensor in super()._streaming_iter(): - yield _vllm_weight_name(name), tensor - - -def _install_vllm_weight_name_adapter() -> None: - import trl.experimental.async_grpo.async_grpo_trainer as trainer_module - - original_worker = trainer_module.AsyncRolloutWorker - if getattr(original_worker, "_terminus_weight_name_adapter", False): - return - - class TerminusRolloutWorker(original_worker): - _terminus_weight_name_adapter = True - - def __init__(self, *args: Any, **kwargs: Any) -> None: - names = kwargs.get("weight_names") - if names is not None: - kwargs = {**kwargs, "weight_names": [_vllm_weight_name(n) for n in names]} - super().__init__(*args, **kwargs) - - trainer_module.AsyncRolloutWorker = TerminusRolloutWorker - - -def _vllm_weight_name(name: str) -> str: - name = name.removeprefix("module.") - if name.startswith(("model.", "lm_head.")): - return f"language_model.{name}" - return name - - def main() -> None: OUTPUT_DIR.mkdir(parents=True, exist_ok=True) @@ -215,9 +181,8 @@ def main() -> None: tokenizer = AutoTokenizer.from_pretrained(MODEL) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - _install_vllm_weight_name_adapter() - trainer = TerminusAsyncGRPOTrainer( + trainer = AsyncGRPOTrainer( model=MODEL, args=AsyncGRPOConfig( output_dir=str(OUTPUT_DIR), From 9a433be4821158db7ee7bca78d45043f242a2af8 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Fri, 29 May 2026 19:50:54 +0000 Subject: [PATCH 03/16] Avoid periodic checkpoint saves in Terminus example --- examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index beb740d70..8cd4ceb4a 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -200,11 +200,8 @@ def main() -> None: report_to=REPORT_TO, run_name=RUN_NAME, project=TRACKIO_PROJECT, - save_strategy="steps", - save_steps=1, - save_total_limit=2, + save_strategy="no", push_to_hub=True, - hub_always_push=True, hub_model_id=HUB_MODEL_ID, chat_template_kwargs={"enable_thinking": False}, vllm_server_base_url=VLLM_SERVER_URL, From 68368034e978d9503ea661bfbd630450e80ecb36 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 09:00:10 +0000 Subject: [PATCH 04/16] Move Pi context into OpenEnv core --- .../tbench2_pi_trl/train_terminus_grpo.py | 51 +++++-------------- src/openenv/core/__init__.py | 3 ++ .../pi_harness.py => src/openenv/core/pi.py | 11 ++-- tests/core/test_pi_context.py | 37 ++++++++++++++ 4 files changed, 62 insertions(+), 40 deletions(-) rename examples/end_to_end/tbench2_pi_trl/pi_harness.py => src/openenv/core/pi.py (89%) create mode 100644 tests/core/test_pi_context.py diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index 8cd4ceb4a..bf3a07af4 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -24,13 +24,12 @@ import json import os -import re from pathlib import Path from typing import Any from datasets import load_dataset from terminus_env.client import TerminusEnv -from terminus_env.harness import TerminusSessionFactory +from terminus_env.harness import TerminusSessionFactory, terminus_reward from transformers import AutoTokenizer from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer @@ -43,7 +42,6 @@ REPORT_TO = "trackio" RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") -REWARD_RE = re.compile(r"reward=([+-]?(?:\d+(?:\.\d*)?|\.\d+))") os.environ["TRACKIO_PROJECT"] = TRACKIO_PROJECT @@ -65,9 +63,19 @@ def terminal(self, command: str = "", final_answer: str = "") -> str: Returns: A JSON string with the tool output, reward, done flag, and error. """ + command = str(command or "") + final_answer = str(final_answer or "") + arguments = { + key: value + for key, value in (("command", command), ("final_answer", final_answer)) + if value.strip() + } + if self._session is None: return json.dumps( { + "tool_name": "terminal", + "arguments": arguments, "done": True, "error": "environment was not reset", "output": "", @@ -76,8 +84,6 @@ def terminal(self, command: str = "", final_answer: str = "") -> str: sort_keys=True, ) - command = str(command or "") - final_answer = str(final_answer or "") if command.strip() and final_answer.strip(): first = self._session.call_tool("terminal", {"command": command}) result = first @@ -108,11 +114,14 @@ def terminal(self, command: str = "", final_answer: str = "") -> str: ) else: result = self._session.call_tool("terminal", {"command": ""}) + arguments = {"command": ""} reward = result.metadata.get("reward") data = result.data if isinstance(result.data, dict) else {"output": result.data} return json.dumps( { + "tool_name": "terminal", + "arguments": arguments, "done": result.done, "error": result.error, "output": data.get("output"), @@ -131,38 +140,6 @@ def _close_session(self) -> None: self._session.close() self._session = None - -def terminus_reward(completions=None, **kwargs) -> list[float]: - del kwargs - rewards = [] - for completion in completions or []: - reward = 0.0 - for message in completion if isinstance(completion, list) else []: - if not isinstance(message, dict) or message.get("role") != "tool": - continue - parsed = _parse_tool_reward(str(message.get("content", ""))) - if parsed is not None: - reward = parsed - rewards.append(reward) - return rewards - - -def _parse_tool_reward(content: str) -> float | None: - try: - payload = json.loads(content) - except json.JSONDecodeError: - payload = None - if isinstance(payload, dict) and payload.get("reward") is not None: - try: - return float(payload["reward"]) - except (TypeError, ValueError): - return None - match = REWARD_RE.search(content) - if match is None: - return None - return float(match.group(1)) - - def main() -> None: OUTPUT_DIR.mkdir(parents=True, exist_ok=True) diff --git a/src/openenv/core/__init__.py b/src/openenv/core/__init__.py index 96065d6a8..2821aeed1 100644 --- a/src/openenv/core/__init__.py +++ b/src/openenv/core/__init__.py @@ -26,6 +26,7 @@ ToolCall, ) from .mcp_client import MCPClientBase, MCPToolClient + from .pi import PiContext from .sync_client import SyncEnvClient __all__ = [ @@ -40,6 +41,7 @@ "LLMResponse", "OpenAIClient", "ToolCall", + "PiContext", "create_llm_client", ] + env_server.__all__ # type: ignore @@ -56,6 +58,7 @@ "LLMResponse": (".llm_client", "LLMResponse"), "OpenAIClient": (".llm_client", "OpenAIClient"), "ToolCall": (".llm_client", "ToolCall"), + "PiContext": (".pi", "PiContext"), "create_llm_client": (".llm_client", "create_llm_client"), } diff --git a/examples/end_to_end/tbench2_pi_trl/pi_harness.py b/src/openenv/core/pi.py similarity index 89% rename from examples/end_to_end/tbench2_pi_trl/pi_harness.py rename to src/openenv/core/pi.py index 3e97afbf0..570dd3456 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_harness.py +++ b/src/openenv/core/pi.py @@ -1,5 +1,10 @@ -#!/usr/bin/env python3 -"""Pi-style context loading for the Terminus async GRPO example.""" +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Pi-compatible model context helpers.""" from __future__ import annotations @@ -80,4 +85,4 @@ def _read_bounded(path: Path, max_chars: int) -> str: return text[:max_chars].rstrip() + "\n\n[truncated by PiContext]" -__all__ = ["PiContext"] +__all__ = ["DEFAULT_PI_SYSTEM_PROMPT", "PiContext"] diff --git a/tests/core/test_pi_context.py b/tests/core/test_pi_context.py new file mode 100644 index 000000000..df3085b31 --- /dev/null +++ b/tests/core/test_pi_context.py @@ -0,0 +1,37 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from openenv.core import PiContext + + +def test_pi_context_loads_project_instructions_and_skills(tmp_path): + (tmp_path / "AGENTS.md").write_text("Project rules", encoding="utf-8") + skill_dir = tmp_path / ".pi" / "skills" / "shell" + skill_dir.mkdir(parents=True) + (skill_dir / "SKILL.md").write_text("Use the shell skill", encoding="utf-8") + + context = PiContext( + project_root=tmp_path, + extra_system_prompts=["Extra context"], + ) + + assert context.loaded_files() == ["AGENTS.md", ".pi/skills/shell/SKILL.md"] + message = context.messages()[0] + assert message["role"] == "system" + assert "Pi-compatible environment session" in message["content"] + assert "Extra context" in message["content"] + assert "Project rules" in message["content"] + assert "Use the shell skill" in message["content"] + + +def test_pi_context_truncates_long_files(tmp_path): + (tmp_path / "AGENTS.md").write_text("abcdef", encoding="utf-8") + + message = PiContext(project_root=tmp_path, max_context_chars=3).messages()[0] + + assert "abc\n\n[truncated by PiContext]" in message["content"] From 804a919b06f2c05b7d8bf4ba84b7ffbbc75f48d2 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 09:04:00 +0000 Subject: [PATCH 05/16] Expose Pi context from core harness --- src/openenv/core/__init__.py | 3 --- src/openenv/core/harness/__init__.py | 3 +++ tests/core/test_pi_context.py | 9 ++++++++- 3 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/openenv/core/__init__.py b/src/openenv/core/__init__.py index 2821aeed1..96065d6a8 100644 --- a/src/openenv/core/__init__.py +++ b/src/openenv/core/__init__.py @@ -26,7 +26,6 @@ ToolCall, ) from .mcp_client import MCPClientBase, MCPToolClient - from .pi import PiContext from .sync_client import SyncEnvClient __all__ = [ @@ -41,7 +40,6 @@ "LLMResponse", "OpenAIClient", "ToolCall", - "PiContext", "create_llm_client", ] + env_server.__all__ # type: ignore @@ -58,7 +56,6 @@ "LLMResponse": (".llm_client", "LLMResponse"), "OpenAIClient": (".llm_client", "OpenAIClient"), "ToolCall": (".llm_client", "ToolCall"), - "PiContext": (".pi", "PiContext"), "create_llm_client": (".llm_client", "create_llm_client"), } diff --git a/src/openenv/core/harness/__init__.py b/src/openenv/core/harness/__init__.py index 8c1412ca4..2f397bf36 100644 --- a/src/openenv/core/harness/__init__.py +++ b/src/openenv/core/harness/__init__.py @@ -22,6 +22,7 @@ from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool from ..env_server.types import State from ..llm_client import LLMResponse +from ..pi import DEFAULT_PI_SYSTEM_PROMPT, PiContext Message = dict[str, Any] RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"}) @@ -703,6 +704,7 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: __all__ = [ "CLIHarnessAdapter", + "DEFAULT_PI_SYSTEM_PROMPT", "HarnessAdapter", "HarnessRolloutResult", "HarnessRunLimits", @@ -710,6 +712,7 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: "Message", "ModelStep", "ModelStepResult", + "PiContext", "RESERVED_TOOL_NAMES", "ResourceSession", "ResourceSessionFactory", diff --git a/tests/core/test_pi_context.py b/tests/core/test_pi_context.py index df3085b31..f70777224 100644 --- a/tests/core/test_pi_context.py +++ b/tests/core/test_pi_context.py @@ -6,7 +6,9 @@ from __future__ import annotations -from openenv.core import PiContext +import openenv.core +from openenv.core.harness import PiContext as HarnessPiContext +from openenv.core.pi import PiContext def test_pi_context_loads_project_instructions_and_skills(tmp_path): @@ -29,6 +31,11 @@ def test_pi_context_loads_project_instructions_and_skills(tmp_path): assert "Use the shell skill" in message["content"] +def test_pi_context_is_exported_from_pi_and_harness_modules(): + assert HarnessPiContext is PiContext + assert "PiContext" not in openenv.core.__all__ + + def test_pi_context_truncates_long_files(tmp_path): (tmp_path / "AGENTS.md").write_text("abcdef", encoding="utf-8") From 90763f503cea27dbbb1c990a0c511c79cb147aba Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 09:40:50 +0000 Subject: [PATCH 06/16] Add Pi CLI harness adapter --- examples/end_to_end/tbench2_pi_trl/README.md | 2 +- .../tbench2_pi_trl/train_terminus_grpo.py | 102 ++++--- src/openenv/core/harness/__init__.py | 3 + src/openenv/core/harness/pi_cli.py | 289 ++++++++++++++++++ tests/core/test_harness_runtime.py | 79 +++++ 5 files changed, 427 insertions(+), 48 deletions(-) create mode 100644 src/openenv/core/harness/pi_cli.py diff --git a/examples/end_to_end/tbench2_pi_trl/README.md b/examples/end_to_end/tbench2_pi_trl/README.md index db426e21b..84ddf1734 100644 --- a/examples/end_to_end/tbench2_pi_trl/README.md +++ b/examples/end_to_end/tbench2_pi_trl/README.md @@ -1,6 +1,6 @@ # Terminus + TRL Async GRPO -Start a Terminus server and a compatible vLLM server, then run: +Start a Terminus server, a compatible vLLM server, and configure the `pi` CLI, then run: ```bash TERMINUS_ENV_URL=http://localhost:8000 \ diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index bf3a07af4..99b22f6b5 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -28,6 +28,7 @@ from typing import Any from datasets import load_dataset +from openenv.core.harness import HarnessRunLimits, PiCLIHarnessAdapter from terminus_env.client import TerminusEnv from terminus_env.harness import TerminusSessionFactory, terminus_reward from transformers import AutoTokenizer @@ -42,33 +43,39 @@ REPORT_TO = "trackio" RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") +PI_PROVIDER = os.environ.get("TERMINUS_PI_PROVIDER") +PI_MODEL = os.environ.get("TERMINUS_PI_MODEL") os.environ["TRACKIO_PROJECT"] = TRACKIO_PROJECT class TerminusHarnessEnvironment: - """Small TRL environment wrapper backed by an OpenEnv resource session.""" - - def __init__(self, session_factory: TerminusSessionFactory): + """Small TRL environment wrapper backed by the Pi CLI harness.""" + + def __init__( + self, + session_factory: TerminusSessionFactory, + pi_harness: PiCLIHarnessAdapter, + limits: HarnessRunLimits, + ): self._session_factory = session_factory + self._pi_harness = pi_harness + self._limits = limits self._session = None def terminal(self, command: str = "", final_answer: str = "") -> str: - """Run a shell command or submit final_answer inside Terminus. + """Run the Pi CLI harness against the current Terminus task. Args: - command: Shell command to run in the sandbox. - final_answer: Final answer to submit after the task is complete. + command: Ignored compatibility field. + final_answer: Ignored compatibility field. Returns: A JSON string with the tool output, reward, done flag, and error. """ - command = str(command or "") - final_answer = str(final_answer or "") + del command, final_answer arguments = { - key: value - for key, value in (("command", command), ("final_answer", final_answer)) - if value.strip() + "harness": "pi_cli", } if self._session is None: @@ -84,47 +91,37 @@ def terminal(self, command: str = "", final_answer: str = "") -> str: sort_keys=True, ) - if command.strip() and final_answer.strip(): - first = self._session.call_tool("terminal", {"command": command}) - result = first - if not first.done: - result = self._session.call_tool( - "terminal", - {"final_answer": final_answer}, - ) - first_data = first.data if isinstance(first.data, dict) else {} - result_data = result.data if isinstance(result.data, dict) else {} - result.data = { - **result_data, - "output": "\n".join( - str(part) - for part in ( - first_data.get("output"), - result_data.get("output"), - ) - if part - ), - } - elif command.strip(): - result = self._session.call_tool("terminal", {"command": command}) - elif final_answer.strip(): - result = self._session.call_tool( - "terminal", - {"final_answer": final_answer}, + try: + rollout = self._pi_harness.run_black_box( + session=self._session, + limits=self._limits, + ) + verify = self._session.verify( + transcript=rollout.messages, + final_state={ + "done": rollout.done, + "metrics": dict(rollout.metrics), + }, ) - else: - result = self._session.call_tool("terminal", {"command": ""}) - arguments = {"command": ""} + result = rollout.tool_trace[-1].result if rollout.tool_trace else None + data = result.data if result is not None else {} + output = data.get("output") if isinstance(data, dict) else data + reward = verify.env_reward + done = verify.done or rollout.done + error = None if result is None else result.error + except Exception as exc: + output = "" + reward = 0.0 + done = True + error = str(exc) - reward = result.metadata.get("reward") - data = result.data if isinstance(result.data, dict) else {"output": result.data} return json.dumps( { "tool_name": "terminal", "arguments": arguments, - "done": result.done, - "error": result.error, - "output": data.get("output"), + "done": done, + "error": error, + "output": output, "reward": 0.0 if reward is None else float(reward), }, sort_keys=True, @@ -140,12 +137,14 @@ def _close_session(self) -> None: self._session.close() self._session = None + def main() -> None: OUTPUT_DIR.mkdir(parents=True, exist_ok=True) task_dataset = load_dataset(TASK_DATASET_ID, split="train") task = task_dataset[0] train_dataset = task_dataset.select_columns(["prompt"]) + limits = HarnessRunLimits(max_turns=task["max_turns"]) session_factory = TerminusSessionFactory( client_factory=lambda: TerminusEnv( base_url=ENV_URL, @@ -154,6 +153,11 @@ def main() -> None: ).sync(), default_verify=list(task["verify"]), ) + pi_harness = PiCLIHarnessAdapter( + provider=PI_PROVIDER, + model=PI_MODEL, + timeout_s=600.0, + ) tokenizer = AutoTokenizer.from_pretrained(MODEL) if tokenizer.pad_token is None: @@ -187,7 +191,11 @@ def main() -> None: processing_class=tokenizer, train_dataset=train_dataset, reward_funcs=terminus_reward, - environment_factory=lambda: TerminusHarnessEnvironment(session_factory), + environment_factory=lambda: TerminusHarnessEnvironment( + session_factory, + pi_harness, + limits, + ), ) trainer.train() trainer.save_model() diff --git a/src/openenv/core/harness/__init__.py b/src/openenv/core/harness/__init__.py index 2f397bf36..8911d4f9d 100644 --- a/src/openenv/core/harness/__init__.py +++ b/src/openenv/core/harness/__init__.py @@ -702,6 +702,8 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: return rollout_func +from .pi_cli import PiCLIHarnessAdapter + __all__ = [ "CLIHarnessAdapter", "DEFAULT_PI_SYSTEM_PROMPT", @@ -712,6 +714,7 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: "Message", "ModelStep", "ModelStepResult", + "PiCLIHarnessAdapter", "PiContext", "RESERVED_TOOL_NAMES", "ResourceSession", diff --git a/src/openenv/core/harness/pi_cli.py b/src/openenv/core/harness/pi_cli.py new file mode 100644 index 000000000..bdca50c2a --- /dev/null +++ b/src/openenv/core/harness/pi_cli.py @@ -0,0 +1,289 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Pi CLI harness adapter.""" + +from __future__ import annotations + +import json +import os +import subprocess +import tempfile +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Callable + +from ..env_server.mcp_types import Tool +from . import ( + _serialize_for_message, + CLIHarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + Message, + ResourceSession, + RolloutEvent, + SessionMCPBridge, + ToolResult, + ToolTraceEntry, +) + + +def _messages_to_prompt(messages: list[Message]) -> str: + if len(messages) == 1 and isinstance(messages[0].get("content"), str): + return str(messages[0]["content"]) + parts = [] + for message in messages: + role = str(message.get("role", "message")) + content = message.get("content", "") + parts.append(f"{role}:\n{_serialize_for_message(content)}") + return "\n\n".join(parts) + + +def _extension_source(tools: list[Tool]) -> str: + tool_payload = [ + { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + for tool in tools + ] + return ( + 'import { Type } from "typebox";\n\n' + f"const tools = {json.dumps(tool_payload, sort_keys=True)};\n" + "const bridgeUrl = process.env.OPENENV_PI_BRIDGE_URL;\n\n" + "export default function(pi) {\n" + " for (const tool of tools) {\n" + " pi.registerTool({\n" + " name: tool.name,\n" + " label: tool.name,\n" + " description: tool.description || tool.name,\n" + " parameters: Type.Unsafe(tool.inputSchema || { type: 'object', properties: {} }),\n" + " async execute(toolCallId, params) {\n" + " const response = await fetch(bridgeUrl, {\n" + " method: 'POST',\n" + " headers: { 'content-type': 'application/json' },\n" + " body: JSON.stringify({\n" + " jsonrpc: '2.0',\n" + " id: toolCallId,\n" + " method: 'tools/call',\n" + " params: { name: tool.name, arguments: params || {} },\n" + " }),\n" + " });\n" + " const payload = await response.json();\n" + " if (payload.error) {\n" + " throw new Error(payload.error.message || JSON.stringify(payload.error));\n" + " }\n" + " const result = payload.result || {};\n" + " return {\n" + " content: [{ type: 'text', text: JSON.stringify(result.data ?? result) }],\n" + " details: result,\n" + " };\n" + " },\n" + " });\n" + " }\n" + "}\n" + ) + + +def _json_events(stdout: str) -> list[dict[str, Any]]: + events = [] + for line in stdout.splitlines(): + line = line.strip() + if not line: + continue + try: + event = json.loads(line) + except json.JSONDecodeError: + event = {"type": "raw", "text": line} + if isinstance(event, dict): + events.append(event) + return events + + +def _messages_from_events(events: list[dict[str, Any]], stdout: str) -> list[Message]: + for event in reversed(events): + messages = event.get("messages") + if event.get("type") == "agent_end" and isinstance(messages, list): + return [message for message in messages if isinstance(message, dict)] + + messages = [ + event["message"] + for event in events + if event.get("type") == "message_end" and isinstance(event.get("message"), dict) + ] + if messages: + return messages + return [{"role": "assistant", "content": stdout}] + + +def _response_body(handler: BaseHTTPRequestHandler, status: int, payload: Any) -> None: + body = json.dumps(payload, default=str).encode("utf-8") + handler.send_response(status) + handler.send_header("content-type", "application/json") + handler.send_header("content-length", str(len(body))) + handler.end_headers() + handler.wfile.write(body) + + +def _bridge_handler( + bridge: SessionMCPBridge, + tool_trace: list[ToolTraceEntry], +) -> type[BaseHTTPRequestHandler]: + class BridgeHandler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: Any) -> None: + return None + + def do_POST(self) -> None: + try: + content_length = int(self.headers.get("content-length", "0")) + request = json.loads(self.rfile.read(content_length).decode("utf-8")) + except Exception as exc: + _response_body( + self, + 400, + { + "jsonrpc": "2.0", + "id": None, + "error": {"code": -32700, "message": str(exc)}, + }, + ) + return + + response = bridge.handle_request(request) + if request.get("method") == "tools/call": + params = request.get("params", {}) or {} + result = response.get("result") + if isinstance(result, dict): + tool_trace.append( + ToolTraceEntry( + tool_name=str(params.get("name", "")), + arguments=dict(params.get("arguments", {}) or {}), + result=ToolResult( + data=result.get("data"), + done=bool(result.get("done")), + metadata=dict(result.get("metadata", {}) or {}), + error=result.get("error"), + ), + ) + ) + _response_body(self, 200, response) + + return BridgeHandler + + +class PiCLIHarnessAdapter(CLIHarnessAdapter): + """Black-box harness adapter that drives an OpenEnv session with the Pi CLI.""" + + def __init__( + self, + *, + pi_command: str | list[str] | tuple[str, ...] = "pi", + provider: str | None = None, + model: str | None = None, + cwd: str | None = None, + timeout_s: float = 600.0, + extra_args: list[str] | None = None, + command_runner: Callable[..., subprocess.CompletedProcess[str]] | None = None, + ): + self._pi_command = ( + [pi_command] if isinstance(pi_command, str) else list(pi_command) + ) + self._provider = provider + self._model = model + self._cwd = cwd + self._timeout_s = timeout_s + self._extra_args = list(extra_args or []) + self._command_runner = command_runner or subprocess.run + super().__init__(runner=self._run_pi) + + def _run_pi( + self, + bridge: SessionMCPBridge, + session: ResourceSession, + limits: HarnessRunLimits, + ) -> HarnessRolloutResult: + del limits + tools = session.list_tools() + tool_trace: list[ToolTraceEntry] = [] + + with tempfile.TemporaryDirectory(prefix="openenv-pi-") as temp_dir: + extension_path = os.path.join(temp_dir, "openenv-tools.mjs") + with open(extension_path, "w", encoding="utf-8") as extension: + extension.write(_extension_source(tools)) + + server = ThreadingHTTPServer( + ("127.0.0.1", 0), + _bridge_handler(bridge, tool_trace), + ) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + bridge_url = f"http://127.0.0.1:{server.server_port}" + + env = dict(os.environ) + env["OPENENV_PI_BRIDGE_URL"] = bridge_url + + command = [ + *self._pi_command, + "--mode", + "json", + "--print", + "--no-session", + "--no-builtin-tools", + "--no-extensions", + "--no-skills", + "--no-prompt-templates", + "--no-context-files", + "--extension", + extension_path, + "--tools", + ",".join(tool.name for tool in tools), + ] + if self._provider: + command.extend(["--provider", self._provider]) + if self._model: + command.extend(["--model", self._model]) + command.extend(self._extra_args) + command.append(_messages_to_prompt(session.initial_messages())) + + try: + completed = self._command_runner( + command, + cwd=self._cwd, + env=env, + text=True, + capture_output=True, + timeout=self._timeout_s, + ) + finally: + server.shutdown() + server.server_close() + thread.join(timeout=1.0) + + stdout = completed.stdout or "" + stderr = completed.stderr or "" + if completed.returncode != 0: + raise RuntimeError( + "pi CLI failed with exit code " + f"{completed.returncode}: {(stderr or stdout).strip()}" + ) + + events = _json_events(stdout) + return HarnessRolloutResult( + messages=_messages_from_events(events, stdout), + tool_trace=tool_trace, + events=[RolloutEvent(type="pi_event", payload=event) for event in events], + done=bool(tool_trace and tool_trace[-1].result.done), + metrics={ + "harness": "pi_cli", + "pi_events": len(events), + "tool_calls": len(tool_trace), + "stderr": stderr, + }, + ) + + +__all__ = ["PiCLIHarnessAdapter"] diff --git a/tests/core/test_harness_runtime.py b/tests/core/test_harness_runtime.py index 388f70845..582835ebd 100644 --- a/tests/core/test_harness_runtime.py +++ b/tests/core/test_harness_runtime.py @@ -8,6 +8,9 @@ from __future__ import annotations +import json +import subprocess +import urllib.request from dataclasses import dataclass, field from typing import Any @@ -22,6 +25,7 @@ HarnessRunLimits, MCPHarnessAdapter, ModelStepResult, + PiCLIHarnessAdapter, RESERVED_TOOL_NAMES, RolloutEvent, SessionMCPBridge, @@ -700,3 +704,78 @@ def runner(bridge, current_session, limits): assert result.done is True assert result.metrics["mode"] == "black_box" assert result.events[0].payload["result"]["done"] is True + + +class TestPiCLIHarnessAdapter: + """Tests for the black-box Pi CLI harness adapter.""" + + def test_pi_cli_adapter_forwards_pi_tool_calls_to_session_bridge(self): + env = FakeSyncEnv() + session = StepEnvSessionAdapter( + client=env, + task="pi-task", + tool_specs=[ + Tool( + name="finish", + description="Finish the task", + input_schema={"type": "object", "properties": {}}, + ) + ], + action_builder=lambda name, arguments: name, + initial_messages_builder=lambda result, task: [ + {"role": "user", "content": f"Solve {task}"} + ], + ) + seen_commands: list[list[str]] = [] + + def fake_pi(command, cwd, env, text, capture_output, timeout): + del cwd, text, capture_output, timeout + seen_commands.append(list(command)) + bridge_url = env["OPENENV_PI_BRIDGE_URL"] + request = urllib.request.Request( + bridge_url, + data=json.dumps( + { + "jsonrpc": "2.0", + "id": "call-1", + "method": "tools/call", + "params": {"name": "finish", "arguments": {}}, + } + ).encode("utf-8"), + headers={"content-type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(request, timeout=5) as response: + payload = json.loads(response.read().decode("utf-8")) + + assert payload["result"]["done"] is True + stdout = "\n".join( + [ + json.dumps({"type": "session", "id": "pi-session"}), + json.dumps( + { + "type": "agent_end", + "messages": [{"role": "assistant", "content": "done"}], + } + ), + ] + ) + return subprocess.CompletedProcess(command, 0, stdout=stdout, stderr="") + + adapter = PiCLIHarnessAdapter( + pi_command="pi", + model="test-model", + command_runner=fake_pi, + ) + + result = adapter.run_black_box(session=session, limits=HarnessRunLimits()) + + assert result.done is True + assert result.metrics["harness"] == "pi_cli" + assert result.tool_trace[0].tool_name == "finish" + assert result.tool_trace[0].result.metadata["reward"] == 1.0 + assert result.messages == [{"role": "assistant", "content": "done"}] + assert seen_commands[0][:3] == ["pi", "--mode", "json"] + assert "--extension" in seen_commands[0] + assert "--no-context-files" in seen_commands[0] + assert seen_commands[0][-1] == "Solve pi-task" From 1a8dfb834436432022ea065bf3f4309c9517854f Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 09:52:00 +0000 Subject: [PATCH 07/16] Simplify Pi CLI harness integration --- .../tbench2_pi_trl/train_terminus_grpo.py | 92 ++++++------------- pyproject.toml | 1 + src/openenv/core/harness/__init__.py | 3 - src/openenv/core/harness/pi_bridge.mjs | 47 ++++++++++ src/openenv/core/harness/pi_cli.py | 59 +----------- src/openenv/core/pi.py | 88 ------------------ tests/core/test_harness_runtime.py | 44 ++++++--- tests/core/test_pi_context.py | 44 --------- 8 files changed, 107 insertions(+), 271 deletions(-) create mode 100644 src/openenv/core/harness/pi_bridge.mjs delete mode 100644 src/openenv/core/pi.py delete mode 100644 tests/core/test_pi_context.py diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index 99b22f6b5..bb6d0688d 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -43,8 +43,6 @@ REPORT_TO = "trackio" RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") -PI_PROVIDER = os.environ.get("TERMINUS_PI_PROVIDER") -PI_MODEL = os.environ.get("TERMINUS_PI_MODEL") os.environ["TRACKIO_PROJECT"] = TRACKIO_PROJECT @@ -64,78 +62,44 @@ def __init__( self._session = None def terminal(self, command: str = "", final_answer: str = "") -> str: - """Run the Pi CLI harness against the current Terminus task. - - Args: - command: Ignored compatibility field. - final_answer: Ignored compatibility field. - - Returns: - A JSON string with the tool output, reward, done flag, and error. - """ + """Compatibility tool that runs Pi against the current Terminus task.""" del command, final_answer - arguments = { - "harness": "pi_cli", - } - if self._session is None: - return json.dumps( - { - "tool_name": "terminal", - "arguments": arguments, - "done": True, - "error": "environment was not reset", - "output": "", - "reward": 0.0, - }, - sort_keys=True, - ) - - try: - rollout = self._pi_harness.run_black_box( - session=self._session, - limits=self._limits, - ) - verify = self._session.verify( - transcript=rollout.messages, - final_state={ - "done": rollout.done, - "metrics": dict(rollout.metrics), - }, - ) - result = rollout.tool_trace[-1].result if rollout.tool_trace else None - data = result.data if result is not None else {} - output = data.get("output") if isinstance(data, dict) else data - reward = verify.env_reward - done = verify.done or rollout.done - error = None if result is None else result.error - except Exception as exc: - output = "" - reward = 0.0 - done = True - error = str(exc) + raise RuntimeError("environment was not reset") + + rollout = self._pi_harness.run_black_box( + session=self._session, + limits=self._limits, + ) + verify = self._session.verify( + transcript=rollout.messages, + final_state={ + "done": rollout.done, + "metrics": dict(rollout.metrics), + }, + ) + result = rollout.tool_trace[-1].result if rollout.tool_trace else None + data = result.data if result is not None else {} return json.dumps( { "tool_name": "terminal", - "arguments": arguments, - "done": done, - "error": error, - "output": output, - "reward": 0.0 if reward is None else float(reward), + "arguments": {"harness": "pi_cli"}, + "done": verify.done or rollout.done, + "error": None if result is None else result.error, + "output": data.get("output") if isinstance(data, dict) else data, + "reward": 0.0 + if verify.env_reward is None + else float(verify.env_reward), }, sort_keys=True, default=str, ) def reset(self, prompt: Any = None, **_: Any) -> None: - self._close_session() - self._session = self._session_factory.create(task=prompt) - - def _close_session(self) -> None: if self._session is not None: self._session.close() - self._session = None + self._session = self._session_factory.create(task=prompt) def main() -> None: @@ -153,11 +117,7 @@ def main() -> None: ).sync(), default_verify=list(task["verify"]), ) - pi_harness = PiCLIHarnessAdapter( - provider=PI_PROVIDER, - model=PI_MODEL, - timeout_s=600.0, - ) + pi_harness = PiCLIHarnessAdapter(timeout_s=600.0) tokenizer = AutoTokenizer.from_pretrained(MODEL) if tokenizer.pad_token is None: @@ -172,7 +132,7 @@ def main() -> None: gradient_accumulation_steps=1, num_generations=task["num_generations"], max_completion_length=task["max_completion_length"], - max_tool_calling_iterations=task["max_turns"], + max_tool_calling_iterations=1, max_inflight_tasks=task["num_generations"], learning_rate=1e-6, logging_steps=1, diff --git a/pyproject.toml b/pyproject.toml index 8da12c016..798c4b187 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,6 +85,7 @@ include-package-data = true [tool.setuptools.package-data] "openenv.cli" = ["templates/**/*"] +"openenv.core.harness" = ["*.mjs"] [tool.setuptools.exclude-package-data] "*" = ["*.pyc", "*.pyo", "__pycache__/*"] diff --git a/src/openenv/core/harness/__init__.py b/src/openenv/core/harness/__init__.py index 8911d4f9d..36a92120b 100644 --- a/src/openenv/core/harness/__init__.py +++ b/src/openenv/core/harness/__init__.py @@ -22,7 +22,6 @@ from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool from ..env_server.types import State from ..llm_client import LLMResponse -from ..pi import DEFAULT_PI_SYSTEM_PROMPT, PiContext Message = dict[str, Any] RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"}) @@ -706,7 +705,6 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: __all__ = [ "CLIHarnessAdapter", - "DEFAULT_PI_SYSTEM_PROMPT", "HarnessAdapter", "HarnessRolloutResult", "HarnessRunLimits", @@ -715,7 +713,6 @@ def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: "ModelStep", "ModelStepResult", "PiCLIHarnessAdapter", - "PiContext", "RESERVED_TOOL_NAMES", "ResourceSession", "ResourceSessionFactory", diff --git a/src/openenv/core/harness/pi_bridge.mjs b/src/openenv/core/harness/pi_bridge.mjs new file mode 100644 index 000000000..fea5ebf2c --- /dev/null +++ b/src/openenv/core/harness/pi_bridge.mjs @@ -0,0 +1,47 @@ +import { Type } from "typebox"; + +const bridgeUrl = process.env.OPENENV_PI_BRIDGE_URL; + +async function callBridge(method, params = {}, id = method) { + if (!bridgeUrl) { + throw new Error("OPENENV_PI_BRIDGE_URL is not set"); + } + const response = await fetch(bridgeUrl, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ + jsonrpc: "2.0", + id, + method, + params, + }), + }); + const payload = await response.json(); + if (!response.ok || payload.error) { + throw new Error(payload.error?.message || response.statusText); + } + return payload.result || {}; +} + +export default async function(pi) { + const { tools = [] } = await callBridge("tools/list"); + for (const tool of tools) { + pi.registerTool({ + name: tool.name, + label: tool.name, + description: tool.description || tool.name, + parameters: Type.Unsafe(tool.inputSchema || { type: "object", properties: {} }), + async execute(toolCallId, params) { + const result = await callBridge( + "tools/call", + { name: tool.name, arguments: params || {} }, + toolCallId, + ); + return { + content: [{ type: "text", text: JSON.stringify(result.data ?? result) }], + details: result, + }; + }, + }); + } +} diff --git a/src/openenv/core/harness/pi_cli.py b/src/openenv/core/harness/pi_cli.py index bdca50c2a..725bc92d5 100644 --- a/src/openenv/core/harness/pi_cli.py +++ b/src/openenv/core/harness/pi_cli.py @@ -11,12 +11,11 @@ import json import os import subprocess -import tempfile import threading from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from importlib.resources import as_file, files from typing import Any, Callable -from ..env_server.mcp_types import Tool from . import ( _serialize_for_message, CLIHarnessAdapter, @@ -42,53 +41,6 @@ def _messages_to_prompt(messages: list[Message]) -> str: return "\n\n".join(parts) -def _extension_source(tools: list[Tool]) -> str: - tool_payload = [ - { - "name": tool.name, - "description": tool.description, - "inputSchema": tool.input_schema, - } - for tool in tools - ] - return ( - 'import { Type } from "typebox";\n\n' - f"const tools = {json.dumps(tool_payload, sort_keys=True)};\n" - "const bridgeUrl = process.env.OPENENV_PI_BRIDGE_URL;\n\n" - "export default function(pi) {\n" - " for (const tool of tools) {\n" - " pi.registerTool({\n" - " name: tool.name,\n" - " label: tool.name,\n" - " description: tool.description || tool.name,\n" - " parameters: Type.Unsafe(tool.inputSchema || { type: 'object', properties: {} }),\n" - " async execute(toolCallId, params) {\n" - " const response = await fetch(bridgeUrl, {\n" - " method: 'POST',\n" - " headers: { 'content-type': 'application/json' },\n" - " body: JSON.stringify({\n" - " jsonrpc: '2.0',\n" - " id: toolCallId,\n" - " method: 'tools/call',\n" - " params: { name: tool.name, arguments: params || {} },\n" - " }),\n" - " });\n" - " const payload = await response.json();\n" - " if (payload.error) {\n" - " throw new Error(payload.error.message || JSON.stringify(payload.error));\n" - " }\n" - " const result = payload.result || {};\n" - " return {\n" - " content: [{ type: 'text', text: JSON.stringify(result.data ?? result) }],\n" - " details: result,\n" - " };\n" - " },\n" - " });\n" - " }\n" - "}\n" - ) - - def _json_events(stdout: str) -> list[dict[str, Any]]: events = [] for line in stdout.splitlines(): @@ -209,12 +161,9 @@ def _run_pi( del limits tools = session.list_tools() tool_trace: list[ToolTraceEntry] = [] + bridge_resource = files("openenv.core.harness").joinpath("pi_bridge.mjs") - with tempfile.TemporaryDirectory(prefix="openenv-pi-") as temp_dir: - extension_path = os.path.join(temp_dir, "openenv-tools.mjs") - with open(extension_path, "w", encoding="utf-8") as extension: - extension.write(_extension_source(tools)) - + with as_file(bridge_resource) as extension_path: server = ThreadingHTTPServer( ("127.0.0.1", 0), _bridge_handler(bridge, tool_trace), @@ -238,7 +187,7 @@ def _run_pi( "--no-prompt-templates", "--no-context-files", "--extension", - extension_path, + str(extension_path), "--tools", ",".join(tool.name for tool in tools), ] diff --git a/src/openenv/core/pi.py b/src/openenv/core/pi.py deleted file mode 100644 index 570dd3456..000000000 --- a/src/openenv/core/pi.py +++ /dev/null @@ -1,88 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -"""Pi-compatible model context helpers.""" - -from __future__ import annotations - -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -Message = dict[str, Any] - -DEFAULT_PI_SYSTEM_PROMPT = """You are running inside a Pi-compatible environment session. - -Follow the project AGENTS.md and available skill instructions. Use the available -tools exactly as described by their schemas. When the task is complete, submit -the environment's final tool call. -""" - - -@dataclass -class PiContext: - """Load model-facing project instructions and skills.""" - - project_root: Path - agents_md_path: str = "AGENTS.md" - skills_dir: str = ".pi/skills" - max_context_chars: int = 6000 - extra_system_prompts: list[str] = field(default_factory=list) - - def messages(self) -> list[Message]: - parts = [DEFAULT_PI_SYSTEM_PROMPT.strip(), *self.extra_system_prompts] - for label, path, text in self.context_files(): - parts.append(f"{label} ({path}):\n{text}") - return [{"role": "system", "content": "\n\n".join(parts)}] - - def loaded_files(self) -> list[str]: - return [path for _, path, _ in self.context_files()] - - def context_files(self) -> list[tuple[str, str, str]]: - files = [ - ("Project instructions", self._resolve(self.agents_md_path)), - *( - (f"Skill: {path.parent.name}", path) - for path in _iter_skill_files(self._resolve(self.skills_dir)) - ), - ] - loaded = [] - for label, path in files: - if text := _read_bounded(path, self.max_context_chars): - loaded.append((label, self._display(path), text)) - return loaded - - def _resolve(self, path: str) -> Path: - resolved = Path(path) - return resolved if resolved.is_absolute() else self.project_root / resolved - - def _display(self, path: Path) -> str: - try: - return path.relative_to(self.project_root).as_posix() - except ValueError: - return path.as_posix() - - -def _iter_skill_files(skills_path: Path) -> list[Path]: - if not skills_path.exists(): - return [] - if skills_path.is_file(): - return [skills_path] - skill_files = set(skills_path.rglob("SKILL.md")) - skill_files.update(path for path in skills_path.glob("*.md") if path.is_file()) - return sorted(skill_files, key=lambda path: path.as_posix()) - - -def _read_bounded(path: Path, max_chars: int) -> str: - if not path.is_file(): - return "" - text = path.read_text(encoding="utf-8").strip() - if len(text) <= max_chars: - return text - return text[:max_chars].rstrip() + "\n\n[truncated by PiContext]" - - -__all__ = ["DEFAULT_PI_SYSTEM_PROMPT", "PiContext"] diff --git a/tests/core/test_harness_runtime.py b/tests/core/test_harness_runtime.py index 582835ebd..df3af1001 100644 --- a/tests/core/test_harness_runtime.py +++ b/tests/core/test_harness_runtime.py @@ -732,22 +732,35 @@ def fake_pi(command, cwd, env, text, capture_output, timeout): del cwd, text, capture_output, timeout seen_commands.append(list(command)) bridge_url = env["OPENENV_PI_BRIDGE_URL"] - request = urllib.request.Request( - bridge_url, - data=json.dumps( - { - "jsonrpc": "2.0", - "id": "call-1", - "method": "tools/call", - "params": {"name": "finish", "arguments": {}}, - } - ).encode("utf-8"), - headers={"content-type": "application/json"}, - method="POST", + + def post(payload): + request = urllib.request.Request( + bridge_url, + data=json.dumps(payload).encode("utf-8"), + headers={"content-type": "application/json"}, + method="POST", + ) + with urllib.request.urlopen(request, timeout=5) as response: + return json.loads(response.read().decode("utf-8")) + + tools = post( + { + "jsonrpc": "2.0", + "id": "tools", + "method": "tools/list", + "params": {}, + } + ) + payload = post( + { + "jsonrpc": "2.0", + "id": "call-1", + "method": "tools/call", + "params": {"name": "finish", "arguments": {}}, + } ) - with urllib.request.urlopen(request, timeout=5) as response: - payload = json.loads(response.read().decode("utf-8")) + assert tools["result"]["tools"][0]["name"] == "finish" assert payload["result"]["done"] is True stdout = "\n".join( [ @@ -776,6 +789,7 @@ def fake_pi(command, cwd, env, text, capture_output, timeout): assert result.tool_trace[0].result.metadata["reward"] == 1.0 assert result.messages == [{"role": "assistant", "content": "done"}] assert seen_commands[0][:3] == ["pi", "--mode", "json"] - assert "--extension" in seen_commands[0] + extension_path = seen_commands[0][seen_commands[0].index("--extension") + 1] + assert extension_path.endswith("pi_bridge.mjs") assert "--no-context-files" in seen_commands[0] assert seen_commands[0][-1] == "Solve pi-task" diff --git a/tests/core/test_pi_context.py b/tests/core/test_pi_context.py deleted file mode 100644 index f70777224..000000000 --- a/tests/core/test_pi_context.py +++ /dev/null @@ -1,44 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import openenv.core -from openenv.core.harness import PiContext as HarnessPiContext -from openenv.core.pi import PiContext - - -def test_pi_context_loads_project_instructions_and_skills(tmp_path): - (tmp_path / "AGENTS.md").write_text("Project rules", encoding="utf-8") - skill_dir = tmp_path / ".pi" / "skills" / "shell" - skill_dir.mkdir(parents=True) - (skill_dir / "SKILL.md").write_text("Use the shell skill", encoding="utf-8") - - context = PiContext( - project_root=tmp_path, - extra_system_prompts=["Extra context"], - ) - - assert context.loaded_files() == ["AGENTS.md", ".pi/skills/shell/SKILL.md"] - message = context.messages()[0] - assert message["role"] == "system" - assert "Pi-compatible environment session" in message["content"] - assert "Extra context" in message["content"] - assert "Project rules" in message["content"] - assert "Use the shell skill" in message["content"] - - -def test_pi_context_is_exported_from_pi_and_harness_modules(): - assert HarnessPiContext is PiContext - assert "PiContext" not in openenv.core.__all__ - - -def test_pi_context_truncates_long_files(tmp_path): - (tmp_path / "AGENTS.md").write_text("abcdef", encoding="utf-8") - - message = PiContext(project_root=tmp_path, max_context_chars=3).messages()[0] - - assert "abc\n\n[truncated by PiContext]" in message["content"] From 6db1d80ad0b7d3d923de7c5c14700601ff312c7a Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 09:58:38 +0000 Subject: [PATCH 08/16] Point Pi harness at training vLLM model --- examples/end_to_end/tbench2_pi_trl/README.md | 2 +- .../tbench2_pi_trl/train_terminus_grpo.py | 10 +++++- src/openenv/core/harness/pi_bridge.mjs | 31 +++++++++++++++++++ src/openenv/core/harness/pi_cli.py | 20 ++++++++++-- tests/core/test_harness_runtime.py | 6 ++++ 5 files changed, 65 insertions(+), 4 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/README.md b/examples/end_to_end/tbench2_pi_trl/README.md index 84ddf1734..db426e21b 100644 --- a/examples/end_to_end/tbench2_pi_trl/README.md +++ b/examples/end_to_end/tbench2_pi_trl/README.md @@ -1,6 +1,6 @@ # Terminus + TRL Async GRPO -Start a Terminus server, a compatible vLLM server, and configure the `pi` CLI, then run: +Start a Terminus server and a compatible vLLM server, then run: ```bash TERMINUS_ENV_URL=http://localhost:8000 \ diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index bb6d0688d..641f8b69b 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -43,6 +43,10 @@ REPORT_TO = "trackio" RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") +PI_VLLM_BASE_URL = os.environ.get( + "TERMINUS_PI_VLLM_BASE_URL", + VLLM_SERVER_URL.rstrip("/") + "/v1", +) os.environ["TRACKIO_PROJECT"] = TRACKIO_PROJECT @@ -117,7 +121,11 @@ def main() -> None: ).sync(), default_verify=list(task["verify"]), ) - pi_harness = PiCLIHarnessAdapter(timeout_s=600.0) + pi_harness = PiCLIHarnessAdapter( + model=MODEL, + model_base_url=PI_VLLM_BASE_URL, + timeout_s=600.0, + ) tokenizer = AutoTokenizer.from_pretrained(MODEL) if tokenizer.pad_token is None: diff --git a/src/openenv/core/harness/pi_bridge.mjs b/src/openenv/core/harness/pi_bridge.mjs index fea5ebf2c..5d89a7bc9 100644 --- a/src/openenv/core/harness/pi_bridge.mjs +++ b/src/openenv/core/harness/pi_bridge.mjs @@ -1,6 +1,35 @@ import { Type } from "typebox"; const bridgeUrl = process.env.OPENENV_PI_BRIDGE_URL; +const modelBaseUrl = process.env.OPENENV_PI_MODEL_BASE_URL; +const modelId = process.env.OPENENV_PI_MODEL_ID; +const modelProvider = process.env.OPENENV_PI_MODEL_PROVIDER || "openenv-vllm"; +const modelApiKey = process.env.OPENENV_PI_MODEL_API_KEY || "openenv"; + +function registerModelProvider(pi) { + if (!modelBaseUrl || !modelId) { + return; + } + pi.registerProvider(modelProvider, { + baseUrl: modelBaseUrl, + apiKey: modelApiKey, + api: "openai-completions", + compat: { + supportsDeveloperRole: false, + supportsReasoningEffort: false, + thinkingFormat: "qwen-chat-template", + }, + models: [{ + id: modelId, + name: modelId, + reasoning: false, + input: ["text"], + contextWindow: 32768, + maxTokens: 4096, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + }], + }); +} async function callBridge(method, params = {}, id = method) { if (!bridgeUrl) { @@ -24,6 +53,8 @@ async function callBridge(method, params = {}, id = method) { } export default async function(pi) { + registerModelProvider(pi); + const { tools = [] } = await callBridge("tools/list"); for (const tool of tools) { pi.registerTool({ diff --git a/src/openenv/core/harness/pi_cli.py b/src/openenv/core/harness/pi_cli.py index 725bc92d5..fe530d59f 100644 --- a/src/openenv/core/harness/pi_cli.py +++ b/src/openenv/core/harness/pi_cli.py @@ -136,6 +136,9 @@ def __init__( pi_command: str | list[str] | tuple[str, ...] = "pi", provider: str | None = None, model: str | None = None, + model_base_url: str | None = None, + model_api_key: str = "openenv", + model_provider: str = "openenv-vllm", cwd: str | None = None, timeout_s: float = 600.0, extra_args: list[str] | None = None, @@ -146,6 +149,9 @@ def __init__( ) self._provider = provider self._model = model + self._model_base_url = model_base_url + self._model_api_key = model_api_key + self._model_provider = model_provider self._cwd = cwd self._timeout_s = timeout_s self._extra_args = list(extra_args or []) @@ -174,6 +180,13 @@ def _run_pi( env = dict(os.environ) env["OPENENV_PI_BRIDGE_URL"] = bridge_url + if self._model_base_url: + if not self._model: + raise ValueError("model is required when model_base_url is set") + env["OPENENV_PI_MODEL_BASE_URL"] = self._model_base_url + env["OPENENV_PI_MODEL_ID"] = self._model + env["OPENENV_PI_MODEL_PROVIDER"] = self._provider or self._model_provider + env["OPENENV_PI_MODEL_API_KEY"] = self._model_api_key command = [ *self._pi_command, @@ -191,8 +204,11 @@ def _run_pi( "--tools", ",".join(tool.name for tool in tools), ] - if self._provider: - command.extend(["--provider", self._provider]) + provider = self._provider or ( + self._model_provider if self._model_base_url else None + ) + if provider: + command.extend(["--provider", provider]) if self._model: command.extend(["--model", self._model]) command.extend(self._extra_args) diff --git a/tests/core/test_harness_runtime.py b/tests/core/test_harness_runtime.py index df3af1001..9d6149625 100644 --- a/tests/core/test_harness_runtime.py +++ b/tests/core/test_harness_runtime.py @@ -732,6 +732,9 @@ def fake_pi(command, cwd, env, text, capture_output, timeout): del cwd, text, capture_output, timeout seen_commands.append(list(command)) bridge_url = env["OPENENV_PI_BRIDGE_URL"] + assert env["OPENENV_PI_MODEL_BASE_URL"] == "http://trainer.example/v1" + assert env["OPENENV_PI_MODEL_ID"] == "test-model" + assert env["OPENENV_PI_MODEL_PROVIDER"] == "openenv-vllm" def post(payload): request = urllib.request.Request( @@ -778,6 +781,7 @@ def post(payload): adapter = PiCLIHarnessAdapter( pi_command="pi", model="test-model", + model_base_url="http://trainer.example/v1", command_runner=fake_pi, ) @@ -788,6 +792,8 @@ def post(payload): assert result.tool_trace[0].tool_name == "finish" assert result.tool_trace[0].result.metadata["reward"] == 1.0 assert result.messages == [{"role": "assistant", "content": "done"}] + assert seen_commands[0][seen_commands[0].index("--provider") + 1] == "openenv-vllm" + assert seen_commands[0][seen_commands[0].index("--model") + 1] == "test-model" assert seen_commands[0][:3] == ["pi", "--mode", "json"] extension_path = seen_commands[0][seen_commands[0].index("--extension") + 1] assert extension_path.endswith("pi_bridge.mjs") From b4abfc7c4e983726ad635e7c9b0dda2623a13706 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 10:21:04 +0000 Subject: [PATCH 09/16] Add PI interception rollout worker --- examples/end_to_end/tbench2_pi_trl/README.md | 2 +- .../tbench2_pi_trl/pi_rollout_worker.py | 834 ++++++++++++++++++ .../tbench2_pi_trl/train_terminus_grpo.py | 113 +-- 3 files changed, 871 insertions(+), 78 deletions(-) create mode 100644 examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py diff --git a/examples/end_to_end/tbench2_pi_trl/README.md b/examples/end_to_end/tbench2_pi_trl/README.md index db426e21b..07096a696 100644 --- a/examples/end_to_end/tbench2_pi_trl/README.md +++ b/examples/end_to_end/tbench2_pi_trl/README.md @@ -1,6 +1,6 @@ # Terminus + TRL Async GRPO -Start a Terminus server and a compatible vLLM server, then run: +Start a Terminus server and a vLLM server with weight transfer enabled, then run: ```bash TERMINUS_ENV_URL=http://localhost:8000 \ diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py new file mode 100644 index 000000000..03c3a89e0 --- /dev/null +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -0,0 +1,834 @@ +"""Pi rollout worker for the Terminus async GRPO example.""" + +from __future__ import annotations + +import json +import logging +import queue +import re +import threading +import time +import uuid +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any, Callable, Iterator, Sequence, cast + +import requests +from openenv.core.harness import HarnessRunLimits +from openenv.core.harness.pi_cli import PiCLIHarnessAdapter + +try: + from trl.chat_template_utils import ( + add_response_schema, + get_training_chat_template, + is_chat_template_prefix_preserving, + parse_response, + ) +except Exception: # pragma: no cover - optional across TRL revisions + add_response_schema = None + get_training_chat_template = None + is_chat_template_prefix_preserving = None + parse_response = None + +try: + from vllm.distributed.weight_transfer.nccl_engine import ( + NCCLTrainerSendWeightsArgs, + NCCLWeightTransferEngine, + ) + from vllm.utils.network_utils import get_ip, get_open_port +except Exception: # pragma: no cover - optional outside training runtime + NCCLTrainerSendWeightsArgs = None + NCCLWeightTransferEngine = None + get_ip = None + get_open_port = None + +logger = logging.getLogger(__name__) + + +def _vllm_version() -> tuple[int, ...]: + try: + import vllm + + return tuple(int(part) for part in vllm.__version__.split(".")[:3]) + except Exception: + return (0, 0, 0) + + +_VLLM_NEEDS_WEIGHT_UPDATE_LIFECYCLE = _vllm_version() >= (0, 21, 0) + + +@dataclass +class RolloutSample: + input_ids: list[int] + completion_mask: list[int] + old_log_probs: list[float] + advantage: float + model_version: int + metrics: dict[str, Any] = field(default_factory=dict) + + +@dataclass(frozen=True) +class WorkerConfig: + max_inflight: int = 2 + queue_maxsize: int = 64 + max_turns: int = 8 + max_completion_tokens: int = 512 + temperature: float = 1.0 + request_timeout_s: float = 600.0 + server_timeout_s: float = 600.0 + idle_sleep_s: float = 0.25 + + +class InterceptionServer: + """Minimal OpenAI-compatible gate for PI chat completion requests.""" + + def __init__(self, *, host: str = "127.0.0.1", port: int = 0, secret: str = "openenv"): + self.host = host + self.port = port + self.secret = secret + self._server: ThreadingHTTPServer | None = None + self._thread: threading.Thread | None = None + self._lock = threading.RLock() + self._rollouts: dict[str, queue.Queue[str]] = {} + self._intercepts: dict[str, dict[str, Any]] = {} + + @property + def base_url(self) -> str: + if self._server is None: + raise RuntimeError("interception server is not running") + return f"http://127.0.0.1:{self.port}" + + def start(self) -> None: + if self._server is not None: + return + server = ThreadingHTTPServer((self.host, self.port), self._handler()) + self._server = server + self.port = int(server.server_port) + self._thread = threading.Thread( + target=server.serve_forever, + daemon=True, + name="terminus-pi-interception", + ) + self._thread.start() + + def stop(self) -> None: + server = self._server + if server is None: + return + server.shutdown() + server.server_close() + if self._thread is not None: + self._thread.join(timeout=2.0) + self._server = None + self._thread = None + + def register_rollout(self, rollout_id: str) -> queue.Queue[str]: + request_queue: queue.Queue[str] = queue.Queue() + with self._lock: + self._rollouts[rollout_id] = request_queue + return request_queue + + def unregister_rollout(self, rollout_id: str) -> None: + with self._lock: + self._rollouts.pop(rollout_id, None) + intercepts = [ + key + for key, intercept in self._intercepts.items() + if intercept.get("rollout_id") == rollout_id + ] + for key in intercepts: + intercept = self._intercepts.pop(key) + intercept["response"] = _error_response("rollout cancelled") + intercept["event"].set() + + def get_intercept(self, request_id: str) -> dict[str, Any] | None: + with self._lock: + return self._intercepts.get(request_id) + + def deliver(self, intercept: dict[str, Any], response: dict[str, Any]) -> None: + intercept["response"] = response + intercept["event"].set() + + def _authorized(self, headers: Any) -> bool: + auth = headers.get("Authorization", "") + api_key = headers.get("x-api-key", "") + return auth == f"Bearer {self.secret}" or api_key == self.secret + + def _handler(self) -> type[BaseHTTPRequestHandler]: + outer = self + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format: str, *args: Any) -> None: + return None + + def do_GET(self) -> None: + if self.path == "/health": + self._json({"status": "ok"}) + return + self._json({"error": "not found"}, status=404) + + def do_POST(self) -> None: + if not outer._authorized(self.headers): + self._json({"error": "unauthorized"}, status=401) + return + match = re.fullmatch( + r"/rollout/([^/]+)/v1/chat/completions", + self.path.split("?", 1)[0], + ) + if match is None: + self._json({"error": "not found"}, status=404) + return + rollout_id = match.group(1) + try: + length = int(self.headers.get("content-length", "0")) + body = json.loads(self.rfile.read(length).decode("utf-8")) + except Exception as exc: + self._json({"error": f"invalid JSON: {exc}"}, status=400) + return + + with outer._lock: + request_queue = outer._rollouts.get(rollout_id) + if request_queue is None: + self._json({"error": "rollout not found"}, status=404) + return + + request_id = f"req_{uuid.uuid4().hex[:8]}" + intercept = { + "request_id": request_id, + "rollout_id": rollout_id, + "messages": body.get("messages"), + "tools": body.get("tools"), + "body": body, + "stream": bool(body.get("stream")), + "event": threading.Event(), + "response": None, + } + with outer._lock: + outer._intercepts[request_id] = intercept + request_queue.put(request_id) + + if not intercept["event"].wait(timeout=900): + self._json({"error": "interception timeout"}, status=504) + return + + with outer._lock: + outer._intercepts.pop(request_id, None) + response = intercept["response"] or _error_response("empty response") + if intercept["stream"]: + self._sse(response) + else: + self._json(response) + + def _json(self, payload: dict[str, Any], *, status: int = 200) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("content-type", "application/json") + self.send_header("content-length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def _sse(self, payload: dict[str, Any]) -> None: + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("cache-control", "no-cache") + self.end_headers() + for chunk in _stream_chunks(payload): + self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode("utf-8")) + self.wfile.write(b"data: [DONE]\n\n") + self.wfile.flush() + + return Handler + + +class TerminusPiRolloutWorker: + """TRL rollout worker that lets PI drive tools while trainer owns generation.""" + + def __init__( + self, + *, + session_factory: Any, + tasks: Sequence[Any], + tokenizer: Any, + vllm_base_url: str, + vllm_model: str, + vllm_api_key: str = "openenv", + config: WorkerConfig | None = None, + chat_template_kwargs: dict[str, Any] | None = None, + pi_command: str = "pi", + command_runner: Callable[..., Any] | None = None, + ): + self._session_factory = session_factory + self._tasks = list(tasks) + if not self._tasks: + raise ValueError("tasks must not be empty") + self._tokenizer = tokenizer + if add_response_schema is not None: + try: + self._tokenizer = add_response_schema(tokenizer) + except Exception: + logger.debug("could not add response schema to tokenizer", exc_info=True) + self._chat_template = None + if ( + get_training_chat_template is not None + and is_chat_template_prefix_preserving is not None + ): + try: + if not is_chat_template_prefix_preserving(self._tokenizer): + self._chat_template = get_training_chat_template(self._tokenizer) + except Exception: + logger.debug("could not inspect chat template", exc_info=True) + self._chat_template_kwargs = dict(chat_template_kwargs or {}) + self._vllm_base_url = vllm_base_url.rstrip("/") + self._vllm_model = vllm_model + self._vllm_api_key = vllm_api_key + self._config = config or WorkerConfig() + self._pi_command = pi_command + self._command_runner = command_runner + + self.rollout_buffer: queue.Queue[RolloutSample] = queue.Queue( + maxsize=self._config.queue_maxsize, + ) + self._interception = InterceptionServer(secret=vllm_api_key) + self._threads: list[threading.Thread] = [] + self._stop = threading.Event() + self._pause = threading.Event() + self._lock = threading.Lock() + self._weight_sync_lock = threading.Lock() + self._task_index = 0 + self._model_version = 0 + self._failed = threading.Event() + self._exception: BaseException | None = None + self._last_heartbeat_s = time.monotonic() + self._model_update_group: Any | None = None + + self._wait_for_server_ready() + self._init_weight_transfer() + + def start(self) -> None: + self._interception.start() + self._stop.clear() + self._last_heartbeat_s = time.monotonic() + with self._lock: + if self._threads: + return + for index in range(max(1, self._config.max_inflight)): + thread = threading.Thread( + target=self._loop, + args=(index,), + daemon=True, + name=f"terminus-pi-rollout-{index}", + ) + thread.start() + self._threads.append(thread) + + def stop(self) -> None: + self._stop.set() + for thread in self._threads: + thread.join(timeout=5.0) + with self._lock: + self._threads = [] + self._interception.stop() + self._destroy_model_update_group() + + def pause(self) -> None: + self._pause.set() + if self._model_update_group is not None: + self._post_json("/pause", params={"mode": "keep"}, timeout=60) + + def resume(self) -> None: + if self._model_update_group is not None: + self._post_json("/resume", timeout=60) + self._pause.clear() + + def send_weights(self, iterator: Iterator[tuple[str, Any]]) -> None: + items = list(iterator) + if not items: + return + if self._model_update_group is None: + raise RuntimeError("vLLM weight-transfer group is not initialized") + + update_info = { + "names": [name for name, _ in items], + "dtype_names": [ + str(getattr(tensor, "dtype", "float32")).split(".")[-1] + for _, tensor in items + ], + "shapes": [list(getattr(tensor, "shape", [])) for _, tensor in items], + "packed": True, + "is_checkpoint_format": True, + } + + with self._weight_sync_lock: + if _VLLM_NEEDS_WEIGHT_UPDATE_LIFECYCLE: + self._post_json( + "/start_weight_update", + json_body={"is_checkpoint_format": True}, + timeout=60, + ) + + post_error: list[Exception] = [] + + def post_update() -> None: + try: + self._post_json( + "/update_weights", + json_body={"update_info": update_info}, + timeout=1800, + ) + except Exception as exc: # noqa: BLE001 + post_error.append(exc) + + update_thread = threading.Thread(target=post_update, daemon=True) + update_thread.start() + + assert NCCLTrainerSendWeightsArgs is not None + assert NCCLWeightTransferEngine is not None + NCCLWeightTransferEngine.trainer_send_weights( + iterator=iter(items), + trainer_args=NCCLTrainerSendWeightsArgs( + group=self._model_update_group, + packed=True, + ), + ) + + update_thread.join(timeout=1800) + if update_thread.is_alive(): + raise TimeoutError("timed out waiting for vLLM /update_weights") + if post_error: + raise RuntimeError("vLLM /update_weights failed") from post_error[0] + if _VLLM_NEEDS_WEIGHT_UPDATE_LIFECYCLE: + self._post_json("/finish_weight_update", timeout=120) + + def update_model_version(self, version: int) -> None: + with self._lock: + self._model_version = version + + def check_health(self, stale_after_s: float) -> None: + if self._failed.is_set(): + raise RuntimeError("Terminus PI rollout worker failed") from self._exception + if not self._threads or not any(thread.is_alive() for thread in self._threads): + raise RuntimeError("Terminus PI rollout worker is not running") + age = time.monotonic() - self._last_heartbeat_s + if age > stale_after_s: + raise RuntimeError( + f"Terminus PI rollout worker heartbeat stale: {age:.0f}s" + ) + + def _wait_for_server_ready(self) -> None: + start = time.time() + while True: + try: + response = requests.get(f"{self._vllm_base_url}/health", timeout=5) + if response.status_code == 200: + return + except requests.RequestException: + pass + if time.time() - start >= self._config.server_timeout_s: + raise TimeoutError( + f"timed out waiting for vLLM server at {self._vllm_base_url}" + ) + time.sleep(2.0) + + def _post_json( + self, + path: str, + *, + timeout: float, + json_body: dict[str, Any] | None = None, + params: dict[str, Any] | None = None, + ) -> requests.Response: + response = requests.post( + f"{self._vllm_base_url}{path}", + headers={"Authorization": f"Bearer {self._vllm_api_key}"}, + json=json_body, + params=params, + timeout=timeout, + ) + if response.status_code != 200: + raise RuntimeError( + f"{path} returned {response.status_code}: {response.text[:400]}" + ) + return response + + def _init_weight_transfer(self) -> None: + if ( + NCCLTrainerSendWeightsArgs is None + or NCCLWeightTransferEngine is None + or get_ip is None + or get_open_port is None + ): + raise RuntimeError("vLLM NCCL weight-transfer modules are unavailable") + + response = requests.get( + f"{self._vllm_base_url}/get_world_size", + headers={"Authorization": f"Bearer {self._vllm_api_key}"}, + timeout=10, + ) + if response.status_code != 200: + raise RuntimeError( + "vLLM weight sync requires /get_world_size. Start vLLM with " + 'VLLM_SERVER_DEV_MODE=1 and --weight-transfer-config \'{"backend":"nccl"}\'.' + ) + + inference_world_size = int(response.json()["world_size"]) + init_info = { + "master_address": get_ip(), + "master_port": get_open_port(), + "rank_offset": 1, + "world_size": inference_world_size + 1, + } + post_error: list[Exception] = [] + + def post_init() -> None: + try: + self._post_json( + "/init_weight_transfer_engine", + json_body={"init_info": init_info}, + timeout=120, + ) + except Exception as exc: # noqa: BLE001 + post_error.append(exc) + + init_thread = threading.Thread(target=post_init, daemon=True) + init_thread.start() + self._model_update_group = NCCLWeightTransferEngine.trainer_init( + { + "master_address": init_info["master_address"], + "master_port": init_info["master_port"], + "world_size": init_info["world_size"], + } + ) + init_thread.join(timeout=120) + if init_thread.is_alive(): + raise TimeoutError("timed out waiting for vLLM weight-transfer init") + if post_error: + raise RuntimeError("vLLM weight-transfer init failed") from post_error[0] + + def _destroy_model_update_group(self) -> None: + group = self._model_update_group + if group is None: + return + try: + group.group.store = None + group.group.socket = None + except Exception: + logger.debug("could not destroy vLLM weight-transfer group", exc_info=True) + self._model_update_group = None + + def _loop(self, worker_index: int) -> None: + while not self._stop.is_set(): + while self._pause.is_set() and not self._stop.is_set(): + time.sleep(0.05) + if self._stop.is_set(): + return + task = self._next_task() + try: + sample = self._rollout( + task, + f"terminus-{worker_index}-{uuid.uuid4().hex[:8]}", + ) + self.rollout_buffer.put(sample, timeout=2.0) + self._last_heartbeat_s = time.monotonic() + except Exception: + logger.exception("terminus PI rollout failed") + time.sleep(self._config.idle_sleep_s) + + def _next_task(self) -> Any: + with self._lock: + task = self._tasks[self._task_index % len(self._tasks)] + self._task_index += 1 + return task + + def _current_model_version(self) -> int: + with self._lock: + return self._model_version + + def _rollout(self, task: Any, rollout_id: str) -> RolloutSample: + session = self._session_factory.create(task=task, episode_id=rollout_id) + request_queue = self._interception.register_rollout(rollout_id) + result_box: dict[str, Any] = {} + error_box: list[BaseException] = [] + + adapter = PiCLIHarnessAdapter( + pi_command=self._pi_command, + model=self._vllm_model, + model_base_url=f"{self._interception.base_url}/rollout/{rollout_id}/v1", + model_api_key=self._vllm_api_key, + timeout_s=self._config.request_timeout_s, + command_runner=self._command_runner, + ) + + def run_pi() -> None: + try: + result_box["rollout"] = adapter.run_black_box( + session=session, + limits=HarnessRunLimits(max_turns=self._config.max_turns), + ) + except BaseException as exc: # noqa: BLE001 + error_box.append(exc) + + pi_thread = threading.Thread(target=run_pi, daemon=True, name=f"pi-{rollout_id}") + pi_thread.start() + + all_ids: list[int] = [] + all_mask: list[int] = [] + all_logprobs: list[float] = [] + previous_prompt_and_turn: list[int] | None = None + turns = 0 + + try: + while turns < self._config.max_turns: + self._last_heartbeat_s = time.monotonic() + if error_box: + raise RuntimeError("pi subprocess failed") from error_box[0] + try: + request_id = request_queue.get(timeout=0.5) + except queue.Empty: + if not pi_thread.is_alive(): + break + continue + + intercept = self._interception.get_intercept(request_id) + if intercept is None: + continue + prompt_ids = self._render_prompt_ids(intercept) + if previous_prompt_and_turn is None: + all_ids.extend(prompt_ids) + all_mask.extend([0] * len(prompt_ids)) + all_logprobs.extend([0.0] * len(prompt_ids)) + else: + suffix = prompt_ids[len(previous_prompt_and_turn) :] + all_ids.extend(suffix) + all_mask.extend([0] * len(suffix)) + all_logprobs.extend([0.0] * len(suffix)) + + turn_ids, turn_logprobs, text, finish_reason = self._generate(prompt_ids) + all_ids.extend(turn_ids) + all_mask.extend([1] * len(turn_ids)) + all_logprobs.extend(turn_logprobs) + previous_prompt_and_turn = prompt_ids + turn_ids + turns += 1 + + assistant_message = _parse_assistant_message( + tokenizer=self._tokenizer, + completion_ids=turn_ids, + fallback_text=text, + ) + self._interception.deliver( + intercept, + _chat_response( + assistant_message, + model=self._vllm_model, + finish_reason=finish_reason, + ), + ) + + pi_thread.join(timeout=2.0) + if error_box: + raise RuntimeError("pi subprocess failed") from error_box[0] + rollout = result_box.get("rollout") + verify = session.verify( + transcript=[] if rollout is None else rollout.messages, + final_state=None + if rollout is None + else {"done": rollout.done, "metrics": dict(rollout.metrics)}, + ) + reward = float(verify.env_reward or 0.0) + if not all_ids: + pad_id = getattr(self._tokenizer, "pad_token_id", None) or 0 + all_ids = [pad_id] + all_mask = [1] + all_logprobs = [0.0] + return RolloutSample( + input_ids=all_ids, + completion_mask=all_mask, + old_log_probs=all_logprobs, + advantage=reward, + model_version=self._current_model_version(), + metrics={"reward": reward, "turns": float(turns)}, + ) + finally: + self._interception.unregister_rollout(rollout_id) + pi_thread.join(timeout=1.0) + session.close() + + def _render_prompt_ids(self, intercept: dict[str, Any]) -> list[int]: + body = intercept.get("body") or {} + messages = body.get("messages") or intercept.get("messages") + if not isinstance(messages, list): + raise RuntimeError("intercepted request did not include messages") + kwargs: dict[str, Any] = { + "add_generation_prompt": True, + "return_dict": False, + **self._chat_template_kwargs, + } + if self._chat_template is not None: + kwargs["chat_template"] = self._chat_template + tools = body.get("tools") or intercept.get("tools") + if tools: + kwargs["tools"] = tools + try: + return cast(list[int], self._tokenizer.apply_chat_template(messages, **kwargs)) + except TypeError: + kwargs.pop("tools", None) + return cast(list[int], self._tokenizer.apply_chat_template(messages, **kwargs)) + + def _generate(self, prompt_ids: list[int]) -> tuple[list[int], list[float], str, str | None]: + response = requests.post( + f"{self._vllm_base_url}/v1/completions", + headers={ + "Authorization": f"Bearer {self._vllm_api_key}", + "Content-Type": "application/json", + }, + json={ + "model": self._vllm_model, + "prompt": prompt_ids, + "max_tokens": self._config.max_completion_tokens, + "temperature": self._config.temperature, + "n": 1, + "return_token_ids": True, + "logprobs": 0, + }, + timeout=self._config.request_timeout_s, + ) + if response.status_code != 200: + raise RuntimeError(f"vLLM {response.status_code}: {response.text[:400]}") + choice = response.json()["choices"][0] + token_ids = list(choice["token_ids"]) + logprobs = choice.get("logprobs", {}).get("token_logprobs", []) + if len(logprobs) != len(token_ids): + logprobs = [0.0] * len(token_ids) + return ( + token_ids, + [0.0 if value is None else float(value) for value in logprobs], + str(choice.get("text", "")), + choice.get("finish_reason"), + ) + + +def _parse_assistant_message( + *, + tokenizer: Any, + completion_ids: list[int], + fallback_text: str, +) -> dict[str, Any]: + if parse_response is None: + return {"role": "assistant", "content": fallback_text} + try: + parsed = parse_response(tokenizer, completion_ids) + except Exception: + return {"role": "assistant", "content": fallback_text} + if not isinstance(parsed, dict): + return {"role": "assistant", "content": fallback_text} + message = {"role": "assistant", "content": str(parsed.get("content") or "")} + tool_calls = _normalize_tool_calls(parsed.get("tool_calls")) + if tool_calls: + message["tool_calls"] = tool_calls + return message + + +def _normalize_tool_calls(raw_tool_calls: Any) -> list[dict[str, Any]]: + if not isinstance(raw_tool_calls, list): + return [] + tool_calls = [] + for raw_call in raw_tool_calls: + if not isinstance(raw_call, dict): + continue + function = raw_call.get("function") + if not isinstance(function, dict) or not function.get("name"): + continue + arguments = function.get("arguments") or {} + tool_calls.append( + { + "id": str(raw_call.get("id") or f"call_{uuid.uuid4().hex[:8]}"), + "type": "function", + "function": { + "name": str(function["name"]), + "arguments": arguments + if isinstance(arguments, str) + else json.dumps(arguments), + }, + } + ) + return tool_calls + + +def _chat_response( + assistant_message: dict[str, Any], + *, + model: str, + finish_reason: str | None, +) -> dict[str, Any]: + return { + "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", + "object": "chat.completion", + "created": int(time.time()), + "model": model, + "choices": [ + { + "index": 0, + "message": assistant_message, + "finish_reason": finish_reason or "stop", + } + ], + } + + +def _stream_chunks(response: dict[str, Any]) -> list[dict[str, Any]]: + chunks = [] + for choice in response.get("choices", []): + message = choice.get("message") or {} + chunks.append( + { + "id": response.get("id", ""), + "object": "chat.completion.chunk", + "created": response.get("created", int(time.time())), + "model": response.get("model", ""), + "choices": [ + { + "index": choice.get("index", 0), + "delta": { + "role": "assistant", + "content": message.get("content"), + "tool_calls": message.get("tool_calls"), + }, + "finish_reason": None, + } + ], + } + ) + chunks.append( + { + "id": response.get("id", ""), + "object": "chat.completion.chunk", + "created": response.get("created", int(time.time())), + "model": response.get("model", ""), + "choices": [ + { + "index": choice.get("index", 0), + "delta": {}, + "finish_reason": choice.get("finish_reason"), + } + ], + } + ) + return chunks + + +def _error_response(message: str) -> dict[str, Any]: + return { + "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", + "object": "chat.completion", + "created": int(time.time()), + "model": "openenv-error", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": message}, + "finish_reason": "stop", + } + ], + } + + +__all__ = ["TerminusPiRolloutWorker", "WorkerConfig"] diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index 641f8b69b..f8af35ee0 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -18,22 +18,21 @@ # openenv-core = { git = "https://github.com/burtenshaw/OpenEnv.git", branch = "codex/terminus-pi-trl-space" } # openenv-terminus-env = { git = "https://github.com/burtenshaw/OpenEnv.git", branch = "codex/terminus-env-harness", subdirectory = "envs/terminus_env" } # /// -"""Run the Terminus async GRPO environment-factory training example.""" +"""Run Terminus async GRPO with PI rollouts owned by TRL.""" from __future__ import annotations -import json import os from pathlib import Path -from typing import Any from datasets import load_dataset -from openenv.core.harness import HarnessRunLimits, PiCLIHarnessAdapter from terminus_env.client import TerminusEnv -from terminus_env.harness import TerminusSessionFactory, terminus_reward +from terminus_env.harness import TerminusSessionFactory from transformers import AutoTokenizer from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer +from pi_rollout_worker import TerminusPiRolloutWorker, WorkerConfig + TASK_DATASET_ID = "burtenshaw/terminus-pi-trl-tasks" MODEL = "Qwen/Qwen3-4B" ENV_URL = os.environ.get("TERMINUS_ENV_URL", "http://localhost:8000") @@ -43,76 +42,17 @@ REPORT_TO = "trackio" RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") -PI_VLLM_BASE_URL = os.environ.get( - "TERMINUS_PI_VLLM_BASE_URL", - VLLM_SERVER_URL.rstrip("/") + "/v1", -) +VLLM_API_KEY = os.environ.get("VLLM_API_KEY", "openenv") os.environ["TRACKIO_PROJECT"] = TRACKIO_PROJECT -class TerminusHarnessEnvironment: - """Small TRL environment wrapper backed by the Pi CLI harness.""" - - def __init__( - self, - session_factory: TerminusSessionFactory, - pi_harness: PiCLIHarnessAdapter, - limits: HarnessRunLimits, - ): - self._session_factory = session_factory - self._pi_harness = pi_harness - self._limits = limits - self._session = None - - def terminal(self, command: str = "", final_answer: str = "") -> str: - """Compatibility tool that runs Pi against the current Terminus task.""" - del command, final_answer - if self._session is None: - raise RuntimeError("environment was not reset") - - rollout = self._pi_harness.run_black_box( - session=self._session, - limits=self._limits, - ) - verify = self._session.verify( - transcript=rollout.messages, - final_state={ - "done": rollout.done, - "metrics": dict(rollout.metrics), - }, - ) - result = rollout.tool_trace[-1].result if rollout.tool_trace else None - data = result.data if result is not None else {} - - return json.dumps( - { - "tool_name": "terminal", - "arguments": {"harness": "pi_cli"}, - "done": verify.done or rollout.done, - "error": None if result is None else result.error, - "output": data.get("output") if isinstance(data, dict) else data, - "reward": 0.0 - if verify.env_reward is None - else float(verify.env_reward), - }, - sort_keys=True, - default=str, - ) - - def reset(self, prompt: Any = None, **_: Any) -> None: - if self._session is not None: - self._session.close() - self._session = self._session_factory.create(task=prompt) - - def main() -> None: OUTPUT_DIR.mkdir(parents=True, exist_ok=True) task_dataset = load_dataset(TASK_DATASET_ID, split="train") task = task_dataset[0] train_dataset = task_dataset.select_columns(["prompt"]) - limits = HarnessRunLimits(max_turns=task["max_turns"]) session_factory = TerminusSessionFactory( client_factory=lambda: TerminusEnv( base_url=ENV_URL, @@ -121,16 +61,37 @@ def main() -> None: ).sync(), default_verify=list(task["verify"]), ) - pi_harness = PiCLIHarnessAdapter( - model=MODEL, - model_base_url=PI_VLLM_BASE_URL, - timeout_s=600.0, - ) tokenizer = AutoTokenizer.from_pretrained(MODEL) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token + rank = int( + os.environ.get("RANK") + or os.environ.get("ACCELERATE_PROCESS_INDEX") + or os.environ.get("SLURM_PROCID") + or "0" + ) + worker = None + if rank == 0: + worker = TerminusPiRolloutWorker( + session_factory=session_factory, + tasks=list(task_dataset), + tokenizer=tokenizer, + vllm_base_url=VLLM_SERVER_URL, + vllm_model=MODEL, + vllm_api_key=VLLM_API_KEY, + chat_template_kwargs={"enable_thinking": False}, + config=WorkerConfig( + max_inflight=task["num_generations"], + max_turns=task["max_turns"], + max_completion_tokens=task["max_completion_length"], + ), + ) + + def unused_reward(**kwargs: object) -> list[float]: + return [0.0] * len(kwargs.get("prompts", [])) + trainer = AsyncGRPOTrainer( model=MODEL, args=AsyncGRPOConfig( @@ -140,9 +101,10 @@ def main() -> None: gradient_accumulation_steps=1, num_generations=task["num_generations"], max_completion_length=task["max_completion_length"], - max_tool_calling_iterations=1, max_inflight_tasks=task["num_generations"], learning_rate=1e-6, + temperature=1.0, + weight_sync_steps=1, logging_steps=1, logging_strategy="steps", log_completions=True, @@ -155,15 +117,12 @@ def main() -> None: chat_template_kwargs={"enable_thinking": False}, vllm_server_base_url=VLLM_SERVER_URL, request_timeout=600, + vllm_server_timeout=600, ), processing_class=tokenizer, train_dataset=train_dataset, - reward_funcs=terminus_reward, - environment_factory=lambda: TerminusHarnessEnvironment( - session_factory, - pi_harness, - limits, - ), + reward_funcs=unused_reward, + rollout_worker=worker, ) trainer.train() trainer.save_model() From 140f272054abe83a014894c872744a3d1e2b763c Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 18:49:01 +0000 Subject: [PATCH 10/16] Allow run-specific Terminus tracking ids --- .../tbench2_pi_trl/train_terminus_grpo.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index f8af35ee0..1b52311c6 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -37,10 +37,19 @@ MODEL = "Qwen/Qwen3-4B" ENV_URL = os.environ.get("TERMINUS_ENV_URL", "http://localhost:8000") OUTPUT_DIR = Path(os.environ.get("TERMINUS_OUTPUT_DIR", "/tmp/terminus-pi-trl-output")) -HUB_MODEL_ID = "burtenshaw/terminus-pi-trl-async-grpo-qwen3-4b" +HUB_MODEL_ID = os.environ.get( + "TERMINUS_HUB_MODEL_ID", + "burtenshaw/terminus-pi-trl-async-grpo-qwen3-4b", +) TRACKIO_PROJECT = "terminus-pi-trl" +TRACKIO_SPACE_ID = os.environ.get( + "TRACKIO_SPACE_ID", + "burtenshaw/terminus-pi-trl-trackio", +) REPORT_TO = "trackio" -RUN_NAME = os.environ.get("JOB_ID", "local") + "-terminus" +RUN_NAME = os.environ.get("TERMINUS_RUN_NAME") or ( + os.environ.get("JOB_ID", "local") + "-terminus" +) VLLM_SERVER_URL = os.environ.get("TERMINUS_VLLM_SERVER_URL", "http://localhost:8001") VLLM_API_KEY = os.environ.get("VLLM_API_KEY", "openenv") @@ -111,6 +120,7 @@ def unused_reward(**kwargs: object) -> list[float]: report_to=REPORT_TO, run_name=RUN_NAME, project=TRACKIO_PROJECT, + trackio_space_id=TRACKIO_SPACE_ID, save_strategy="no", push_to_hub=True, hub_model_id=HUB_MODEL_ID, From 0e322847530f6a5369fd5dea3da8590fe523f2fe Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 19:08:11 +0000 Subject: [PATCH 11/16] Normalize Terminus PI text tool calls --- .../tbench2_pi_trl/pi_rollout_worker.py | 67 +++++++++++++++++-- 1 file changed, 62 insertions(+), 5 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py index 03c3a89e0..ed636471b 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -17,6 +17,11 @@ from openenv.core.harness import HarnessRunLimits from openenv.core.harness.pi_cli import PiCLIHarnessAdapter +try: + from terminus_env.harness import build_terminal_tool_call +except Exception: # pragma: no cover - optional outside the Terminus example + build_terminal_tool_call = None + try: from trl.chat_template_utils import ( add_response_schema, @@ -639,13 +644,27 @@ def run_pi() -> None: all_ids = [pad_id] all_mask = [1] all_logprobs = [0.0] + metrics = {"reward": reward, "turns": float(turns)} + if rollout is not None: + metrics.update( + { + "pi/tool_calls": float(len(rollout.tool_trace)), + "pi/events": float(rollout.metrics.get("pi_events", 0.0)), + "pi/done": float(bool(rollout.done)), + } + ) + for name, value in (verify.metrics or {}).items(): + if isinstance(value, bool): + metrics[f"verify/{name}"] = float(value) + elif isinstance(value, (int, float)): + metrics[f"verify/{name}"] = float(value) return RolloutSample( input_ids=all_ids, completion_mask=all_mask, old_log_probs=all_logprobs, advantage=reward, model_version=self._current_model_version(), - metrics={"reward": reward, "turns": float(turns)}, + metrics=metrics, ) finally: self._interception.unregister_rollout(rollout_id) @@ -713,20 +732,58 @@ def _parse_assistant_message( fallback_text: str, ) -> dict[str, Any]: if parse_response is None: - return {"role": "assistant", "content": fallback_text} + return _assistant_message_from_text(fallback_text) try: parsed = parse_response(tokenizer, completion_ids) except Exception: - return {"role": "assistant", "content": fallback_text} + return _assistant_message_from_text(fallback_text) if not isinstance(parsed, dict): - return {"role": "assistant", "content": fallback_text} - message = {"role": "assistant", "content": str(parsed.get("content") or "")} + return _assistant_message_from_text(fallback_text) + content = str(parsed.get("content") or "") + message = {"role": "assistant", "content": content} tool_calls = _normalize_tool_calls(parsed.get("tool_calls")) + if not tool_calls: + tool_calls = _terminal_tool_call_from_text(content or fallback_text) if tool_calls: + message["content"] = "" message["tool_calls"] = tool_calls return message +def _assistant_message_from_text(text: str) -> dict[str, Any]: + tool_calls = _terminal_tool_call_from_text(text) + if tool_calls: + return {"role": "assistant", "content": "", "tool_calls": tool_calls} + return {"role": "assistant", "content": text} + + +def _terminal_tool_call_from_text(text: str) -> list[dict[str, Any]]: + if build_terminal_tool_call is None or not text.strip(): + return [] + try: + tool_call = build_terminal_tool_call( + text, + call_id=f"call_{uuid.uuid4().hex[:8]}", + ) + except Exception: + logger.debug("could not parse Terminus terminal text", exc_info=True) + return [] + arguments = getattr(tool_call, "args", None) or {} + name = getattr(tool_call, "name", "") + if name != "terminal" or not arguments: + return [] + return [ + { + "id": str(getattr(tool_call, "id", "") or f"call_{uuid.uuid4().hex[:8]}"), + "type": "function", + "function": { + "name": name, + "arguments": json.dumps(arguments), + }, + } + ] + + def _normalize_tool_calls(raw_tool_calls: Any) -> list[dict[str, Any]]: if not isinstance(raw_tool_calls, list): return [] From 7e662a83d6bf1db0168d3384090cc65199563594 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 19:15:25 +0000 Subject: [PATCH 12/16] Pass Terminus tasks as plain instructions --- .../tbench2_pi_trl/pi_rollout_worker.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py index ed636471b..65c535a5a 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -549,7 +549,10 @@ def _current_model_version(self) -> int: return self._model_version def _rollout(self, task: Any, rollout_id: str) -> RolloutSample: - session = self._session_factory.create(task=task, episode_id=rollout_id) + session = self._session_factory.create( + task=_session_task(task), + episode_id=rollout_id, + ) request_queue = self._interception.register_rollout(rollout_id) result_box: dict[str, Any] = {} error_box: list[BaseException] = [] @@ -750,6 +753,19 @@ def _parse_assistant_message( return message +def _session_task(task: Any) -> Any: + if not isinstance(task, dict) or not isinstance(task.get("prompt"), list): + return task + instruction = "\n\n".join( + str(message.get("content", "")) + for message in task["prompt"] + if isinstance(message, dict) and message.get("content") + ) + if not instruction: + return task + return {**task, "instruction": instruction} + + def _assistant_message_from_text(text: str) -> dict[str, Any]: tool_calls = _terminal_tool_call_from_text(text) if tool_calls: From 70a6bb27b13212da9b05845070207db8fed43477 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Sat, 30 May 2026 19:27:49 +0000 Subject: [PATCH 13/16] Normalize PI chat content for Qwen --- .../tbench2_pi_trl/pi_rollout_worker.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py index 65c535a5a..29e1b4c57 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -679,6 +679,7 @@ def _render_prompt_ids(self, intercept: dict[str, Any]) -> list[int]: messages = body.get("messages") or intercept.get("messages") if not isinstance(messages, list): raise RuntimeError("intercepted request did not include messages") + messages = _normalize_chat_messages(messages) kwargs: dict[str, Any] = { "add_generation_prompt": True, "return_dict": False, @@ -753,6 +754,23 @@ def _parse_assistant_message( return message +def _normalize_chat_messages(messages: list[Any]) -> list[dict[str, Any]]: + normalized = [] + for message in messages: + if not isinstance(message, dict): + continue + item = dict(message) + content = item.get("content") + if isinstance(content, list): + item["content"] = "\n".join( + str(part.get("text", "")) + for part in content + if isinstance(part, dict) and part.get("text") is not None + ) + normalized.append(item) + return normalized + + def _session_task(task: Any) -> Any: if not isinstance(task, dict) or not isinstance(task.get("prompt"), list): return task From 4ae641b02334b05e8813f336ced657999a212f9a Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Mon, 1 Jun 2026 09:48:34 +0000 Subject: [PATCH 14/16] Simplify Terminus PI rollout worker --- .../tbench2_pi_trl/pi_rollout_worker.py | 92 +++---------------- 1 file changed, 11 insertions(+), 81 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py index 29e1b4c57..9ba0d9128 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -204,7 +204,6 @@ def do_POST(self) -> None: "messages": body.get("messages"), "tools": body.get("tools"), "body": body, - "stream": bool(body.get("stream")), "event": threading.Event(), "response": None, } @@ -219,10 +218,7 @@ def do_POST(self) -> None: with outer._lock: outer._intercepts.pop(request_id, None) response = intercept["response"] or _error_response("empty response") - if intercept["stream"]: - self._sse(response) - else: - self._json(response) + self._json(response) def _json(self, payload: dict[str, Any], *, status: int = 200) -> None: body = json.dumps(payload).encode("utf-8") @@ -232,16 +228,6 @@ def _json(self, payload: dict[str, Any], *, status: int = 200) -> None: self.end_headers() self.wfile.write(body) - def _sse(self, payload: dict[str, Any]) -> None: - self.send_response(200) - self.send_header("content-type", "text/event-stream") - self.send_header("cache-control", "no-cache") - self.end_headers() - for chunk in _stream_chunks(payload): - self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode("utf-8")) - self.wfile.write(b"data: [DONE]\n\n") - self.wfile.flush() - return Handler @@ -301,8 +287,6 @@ def __init__( self._weight_sync_lock = threading.Lock() self._task_index = 0 self._model_version = 0 - self._failed = threading.Event() - self._exception: BaseException | None = None self._last_heartbeat_s = time.monotonic() self._model_update_group: Any | None = None @@ -409,8 +393,6 @@ def update_model_version(self, version: int) -> None: self._model_version = version def check_health(self, stale_after_s: float) -> None: - if self._failed.is_set(): - raise RuntimeError("Terminus PI rollout worker failed") from self._exception if not self._threads or not any(thread.is_alive() for thread in self._threads): raise RuntimeError("Terminus PI rollout worker is not running") age = time.monotonic() - self._last_heartbeat_s @@ -544,10 +526,6 @@ def _next_task(self) -> Any: self._task_index += 1 return task - def _current_model_version(self) -> int: - with self._lock: - return self._model_version - def _rollout(self, task: Any, rollout_id: str) -> RolloutSample: session = self._session_factory.create( task=_session_task(task), @@ -661,12 +639,14 @@ def run_pi() -> None: metrics[f"verify/{name}"] = float(value) elif isinstance(value, (int, float)): metrics[f"verify/{name}"] = float(value) + with self._lock: + model_version = self._model_version return RolloutSample( input_ids=all_ids, completion_mask=all_mask, old_log_probs=all_logprobs, advantage=reward, - model_version=self._current_model_version(), + model_version=model_version, metrics=metrics, ) finally: @@ -735,23 +715,21 @@ def _parse_assistant_message( completion_ids: list[int], fallback_text: str, ) -> dict[str, Any]: - if parse_response is None: - return _assistant_message_from_text(fallback_text) + parsed: dict[str, Any] = {} try: - parsed = parse_response(tokenizer, completion_ids) + if parse_response is not None: + parsed = parse_response(tokenizer, completion_ids) except Exception: - return _assistant_message_from_text(fallback_text) + logger.debug("could not parse TRL response schema", exc_info=True) if not isinstance(parsed, dict): - return _assistant_message_from_text(fallback_text) + parsed = {} content = str(parsed.get("content") or "") - message = {"role": "assistant", "content": content} tool_calls = _normalize_tool_calls(parsed.get("tool_calls")) if not tool_calls: tool_calls = _terminal_tool_call_from_text(content or fallback_text) if tool_calls: - message["content"] = "" - message["tool_calls"] = tool_calls - return message + return {"role": "assistant", "content": "", "tool_calls": tool_calls} + return {"role": "assistant", "content": content or fallback_text} def _normalize_chat_messages(messages: list[Any]) -> list[dict[str, Any]]: @@ -784,13 +762,6 @@ def _session_task(task: Any) -> Any: return {**task, "instruction": instruction} -def _assistant_message_from_text(text: str) -> dict[str, Any]: - tool_calls = _terminal_tool_call_from_text(text) - if tool_calls: - return {"role": "assistant", "content": "", "tool_calls": tool_calls} - return {"role": "assistant", "content": text} - - def _terminal_tool_call_from_text(text: str) -> list[dict[str, Any]]: if build_terminal_tool_call is None or not text.strip(): return [] @@ -865,47 +836,6 @@ def _chat_response( } -def _stream_chunks(response: dict[str, Any]) -> list[dict[str, Any]]: - chunks = [] - for choice in response.get("choices", []): - message = choice.get("message") or {} - chunks.append( - { - "id": response.get("id", ""), - "object": "chat.completion.chunk", - "created": response.get("created", int(time.time())), - "model": response.get("model", ""), - "choices": [ - { - "index": choice.get("index", 0), - "delta": { - "role": "assistant", - "content": message.get("content"), - "tool_calls": message.get("tool_calls"), - }, - "finish_reason": None, - } - ], - } - ) - chunks.append( - { - "id": response.get("id", ""), - "object": "chat.completion.chunk", - "created": response.get("created", int(time.time())), - "model": response.get("model", ""), - "choices": [ - { - "index": choice.get("index", 0), - "delta": {}, - "finish_reason": choice.get("finish_reason"), - } - ], - } - ) - return chunks - - def _error_response(message: str) -> dict[str, Any]: return { "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", From d1d2c5660c753d4fd144c0ac972acdf03af7aadd Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Mon, 1 Jun 2026 18:19:33 +0000 Subject: [PATCH 15/16] Run Terminus async GRPO with Qwen3.6 --- .../tbench2_pi_trl/pi_rollout_worker.py | 7 ++++ .../tbench2_pi_trl/train_terminus_grpo.py | 40 +++++++++++++++---- 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py index 9ba0d9128..0eb219d0a 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -82,6 +82,7 @@ class WorkerConfig: request_timeout_s: float = 600.0 server_timeout_s: float = 600.0 idle_sleep_s: float = 0.25 + vllm_weight_name_prefix: str = "" class InterceptionServer: @@ -333,6 +334,12 @@ def send_weights(self, iterator: Iterator[tuple[str, Any]]) -> None: items = list(iterator) if not items: return + if self._config.vllm_weight_name_prefix: + prefix = self._config.vllm_weight_name_prefix + items = [ + (name if name.startswith(prefix) else f"{prefix}{name}", tensor) + for name, tensor in items + ] if self._model_update_group is None: raise RuntimeError("vLLM weight-transfer group is not initialized") diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index 1b52311c6..afd923337 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -26,22 +26,25 @@ from pathlib import Path from datasets import load_dataset +from huggingface_hub import HfApi from terminus_env.client import TerminusEnv from terminus_env.harness import TerminusSessionFactory -from transformers import AutoTokenizer +from transformers import AutoTokenizer, TrainerCallback from trl.experimental.async_grpo import AsyncGRPOConfig, AsyncGRPOTrainer from pi_rollout_worker import TerminusPiRolloutWorker, WorkerConfig TASK_DATASET_ID = "burtenshaw/terminus-pi-trl-tasks" -MODEL = "Qwen/Qwen3-4B" +MODEL = "Qwen/Qwen3.6-27B" +TRAINER_MODEL = os.environ.get("TERMINUS_TRAINER_MODEL", MODEL) ENV_URL = os.environ.get("TERMINUS_ENV_URL", "http://localhost:8000") OUTPUT_DIR = Path(os.environ.get("TERMINUS_OUTPUT_DIR", "/tmp/terminus-pi-trl-output")) HUB_MODEL_ID = os.environ.get( "TERMINUS_HUB_MODEL_ID", - "burtenshaw/terminus-pi-trl-async-grpo-qwen3-4b", + "burtenshaw/terminus-pi-trl-async-grpo-qwen3-6-27b", ) TRACKIO_PROJECT = "terminus-pi-trl" +MAX_STEPS = 4 TRACKIO_SPACE_ID = os.environ.get( "TRACKIO_SPACE_ID", "burtenshaw/terminus-pi-trl-trackio", @@ -95,6 +98,7 @@ def main() -> None: max_inflight=task["num_generations"], max_turns=task["max_turns"], max_completion_tokens=task["max_completion_length"], + vllm_weight_name_prefix="language_model.", ), ) @@ -102,10 +106,10 @@ def unused_reward(**kwargs: object) -> list[float]: return [0.0] * len(kwargs.get("prompts", [])) trainer = AsyncGRPOTrainer( - model=MODEL, + model=TRAINER_MODEL, args=AsyncGRPOConfig( output_dir=str(OUTPUT_DIR), - max_steps=task["max_steps"], + max_steps=MAX_STEPS, per_device_train_batch_size=task["batch_size"], gradient_accumulation_steps=1, num_generations=task["num_generations"], @@ -134,9 +138,31 @@ def unused_reward(**kwargs: object) -> list[float]: reward_funcs=unused_reward, rollout_worker=worker, ) + + class SaveAndUploadCallback(TrainerCallback): + def __init__(self) -> None: + self.saved = False + + def on_step_end(self, args, state, control, **kwargs): + if self.saved or state.global_step < MAX_STEPS: + return control + self.saved = True + state_dict = trainer.accelerator.get_state_dict(trainer.model) + if rank == 0: + trainer._save(str(OUTPUT_DIR), state_dict=state_dict) + del state_dict + api = HfApi() + api.create_repo(repo_id=HUB_MODEL_ID, repo_type="model", exist_ok=True) + api.upload_large_folder( + repo_id=HUB_MODEL_ID, + repo_type="model", + folder_path=OUTPUT_DIR, + num_workers=8, + ) + return control + + trainer.add_callback(SaveAndUploadCallback()) trainer.train() - trainer.save_model() - trainer.push_to_hub(commit_message=f"Async GRPO Terminus run {RUN_NAME}") if __name__ == "__main__": From e881580f2e9eb3bd22bc40a449b7f78d2c4eb033 Mon Sep 17 00:00:00 2001 From: Ben Burtenshaw Date: Tue, 2 Jun 2026 10:21:21 +0000 Subject: [PATCH 16/16] examples: use harder Terminus PI GRPO tasks --- .../tbench2_pi_trl/pi_rollout_worker.py | 87 ++++++++++++++++++- .../tbench2_pi_trl/train_terminus_grpo.py | 12 +-- 2 files changed, 91 insertions(+), 8 deletions(-) diff --git a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py index 0eb219d0a..9dc2434a0 100644 --- a/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py +++ b/examples/end_to_end/tbench2_pi_trl/pi_rollout_worker.py @@ -219,6 +219,9 @@ def do_POST(self) -> None: with outer._lock: outer._intercepts.pop(request_id, None) response = intercept["response"] or _error_response("empty response") + if body.get("stream"): + self._sse(response) + return self._json(response) def _json(self, payload: dict[str, Any], *, status: int = 200) -> None: @@ -229,6 +232,56 @@ def _json(self, payload: dict[str, Any], *, status: int = 200) -> None: self.end_headers() self.wfile.write(body) + def _sse(self, payload: dict[str, Any]) -> None: + self.send_response(200) + self.send_header("content-type", "text/event-stream") + self.send_header("cache-control", "no-cache") + self.end_headers() + for choice in payload.get("choices") or []: + message = choice.get("message") or {} + tool_calls = [ + {"index": index, **tool_call} + for index, tool_call in enumerate(message.get("tool_calls") or []) + ] + self._sse_data( + { + "id": payload.get("id", ""), + "object": "chat.completion.chunk", + "created": payload.get("created", int(time.time())), + "model": payload.get("model", ""), + "choices": [ + { + "index": choice.get("index", 0), + "delta": { + "role": "assistant", + "content": message.get("content"), + "tool_calls": tool_calls or None, + }, + "finish_reason": None, + } + ], + } + ) + self._sse_data( + { + "id": payload.get("id", ""), + "object": "chat.completion.chunk", + "created": payload.get("created", int(time.time())), + "model": payload.get("model", ""), + "choices": [ + { + "index": choice.get("index", 0), + "delta": {}, + "finish_reason": choice.get("finish_reason") or "stop", + } + ], + } + ) + self.wfile.write(b"data: [DONE]\n\n") + + def _sse_data(self, payload: dict[str, Any]) -> None: + self.wfile.write(f"data: {json.dumps(payload)}\n\n".encode("utf-8")) + return Handler @@ -746,16 +799,43 @@ def _normalize_chat_messages(messages: list[Any]) -> list[dict[str, Any]]: continue item = dict(message) content = item.get("content") - if isinstance(content, list): + if content is None: + item["content"] = "" + elif isinstance(content, list): item["content"] = "\n".join( str(part.get("text", "")) for part in content if isinstance(part, dict) and part.get("text") is not None ) + tool_calls = item.get("tool_calls") + if isinstance(tool_calls, list): + item["tool_calls"] = _normalize_message_tool_calls(tool_calls) normalized.append(item) return normalized +def _normalize_message_tool_calls(raw_tool_calls: list[Any]) -> list[dict[str, Any]]: + tool_calls = [] + for raw_call in raw_tool_calls: + if not isinstance(raw_call, dict): + continue + call = dict(raw_call) + function = call.get("function") + if not isinstance(function, dict): + tool_calls.append(call) + continue + function = dict(function) + arguments = function.get("arguments") + if isinstance(arguments, str): + try: + function["arguments"] = json.loads(arguments) + except json.JSONDecodeError: + function["arguments"] = {"command": arguments} + call["function"] = function + tool_calls.append(call) + return tool_calls + + def _session_task(task: Any) -> Any: if not isinstance(task, dict) or not isinstance(task.get("prompt"), list): return task @@ -828,6 +908,9 @@ def _chat_response( model: str, finish_reason: str | None, ) -> dict[str, Any]: + choice_finish_reason = ( + "tool_calls" if assistant_message.get("tool_calls") else finish_reason or "stop" + ) return { "id": f"chatcmpl-{uuid.uuid4().hex[:8]}", "object": "chat.completion", @@ -837,7 +920,7 @@ def _chat_response( { "index": 0, "message": assistant_message, - "finish_reason": finish_reason or "stop", + "finish_reason": choice_finish_reason, } ], } diff --git a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py index afd923337..31cb1b637 100644 --- a/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py +++ b/examples/end_to_end/tbench2_pi_trl/train_terminus_grpo.py @@ -34,17 +34,16 @@ from pi_rollout_worker import TerminusPiRolloutWorker, WorkerConfig -TASK_DATASET_ID = "burtenshaw/terminus-pi-trl-tasks" -MODEL = "Qwen/Qwen3.6-27B" +TASK_DATASET_ID = "burtenshaw/terminus-pi-trl-hard-tasks" +MODEL = "Qwen/Qwen3.5-4B" TRAINER_MODEL = os.environ.get("TERMINUS_TRAINER_MODEL", MODEL) ENV_URL = os.environ.get("TERMINUS_ENV_URL", "http://localhost:8000") OUTPUT_DIR = Path(os.environ.get("TERMINUS_OUTPUT_DIR", "/tmp/terminus-pi-trl-output")) HUB_MODEL_ID = os.environ.get( "TERMINUS_HUB_MODEL_ID", - "burtenshaw/terminus-pi-trl-async-grpo-qwen3-6-27b", + "burtenshaw/terminus-pi-trl-async-grpo-qwen3-5-4b-hard", ) TRACKIO_PROJECT = "terminus-pi-trl" -MAX_STEPS = 4 TRACKIO_SPACE_ID = os.environ.get( "TRACKIO_SPACE_ID", "burtenshaw/terminus-pi-trl-trackio", @@ -64,6 +63,7 @@ def main() -> None: task_dataset = load_dataset(TASK_DATASET_ID, split="train") task = task_dataset[0] + max_steps = int(task["max_steps"]) train_dataset = task_dataset.select_columns(["prompt"]) session_factory = TerminusSessionFactory( client_factory=lambda: TerminusEnv( @@ -109,7 +109,7 @@ def unused_reward(**kwargs: object) -> list[float]: model=TRAINER_MODEL, args=AsyncGRPOConfig( output_dir=str(OUTPUT_DIR), - max_steps=MAX_STEPS, + max_steps=max_steps, per_device_train_batch_size=task["batch_size"], gradient_accumulation_steps=1, num_generations=task["num_generations"], @@ -144,7 +144,7 @@ def __init__(self) -> None: self.saved = False def on_step_end(self, args, state, control, **kwargs): - if self.saved or state.global_step < MAX_STEPS: + if self.saved or state.global_step < max_steps: return control self.saved = True state_dict = trainer.accelerator.get_state_dict(trainer.model)