diff --git a/.gitignore b/.gitignore index 9526bd1..1c89e32 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,223 @@ swe-rl/scripts/logs swe.pem __pycache__/ *.pyc +logs +.DS_Store + + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py.cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +# Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +# poetry.lock +# poetry.toml + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +# pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python. +# https://pdm-project.org/en/latest/usage/project/#working-with-version-control +# pdm.lock +# pdm.toml +.pdm-python +.pdm-build/ + +# pixi +# Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control. +# pixi.lock +# Pixi creates a virtual environment in the .pixi directory, just like venv module creates one +# in the .venv directory. It is recommended not to include this directory in version control. +.pixi + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# Redis +*.rdb +*.aof +*.pid + +# RabbitMQ +mnesia/ +rabbitmq/ +rabbitmq-data/ + +# ActiveMQ +activemq-data/ + +# SageMath parsed files +*.sage.py + +# Environments +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +# .idea/ + +# Abstra +# Abstra is an AI-powered process automation framework. +# Ignore directories containing user credentials, local state, and settings. +# Learn more at https://abstra.io/docs +.abstra/ + +# Visual Studio Code +# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore +# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore +# and can be added to the global gitignore or merged into this file. However, if you prefer, +# you could uncomment the following to ignore the entire vscode folder +# .vscode/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc + +# Marimo +marimo/_static/ +marimo/_lsp/ +__marimo__/ + +# Streamlit +.streamlit/secrets.toml \ No newline at end of file diff --git a/terminal-rl/README.md b/terminal-rl/README.md new file mode 100644 index 0000000..15d01bf --- /dev/null +++ b/terminal-rl/README.md @@ -0,0 +1,105 @@ +# Terminal RL + +RL training for terminal agents. The agent interacts with Docker-hosted environments and is trained with GRPO (optional PRM). + +This workflow has two independent components: + +- **Training machine** runs task router + Ray + training scripts, and connects to workers via `WORKER_URLS` +- **Remote workers:** run the pool server and execute tasks (Docker required): [remote/README.md](remote/README.md) + +--- + +## Prerequisites + +- **Training machine:** GPU node/cluster with the required training dependencies. +- **Remote workers:** Docker-capable hosts reachable from the training machine (default pool server port **18081**). Setup: [remote/README.md](remote/README.md). + +--- + +## Instructions + +### 0. Start remote workers (pool server) + +Follow [remote/README.md](remote/README.md) on each worker to start `pool_server` (it should be reachable at e.g. `http://:18081`). + +### 1. Clone the repo + +From a directory of your choice: + +```bash +git clone https://github.com/Gen-Verse/OpenClaw-RL.git +cd OpenClaw-RL +``` + +### 2. Prepare dataset (download + convert) + +Download a supported dataset under `terminal-rl/dataset/`: + +```bash +export DATASET_DIR="terminal-rl/dataset" +python terminal-rl/data_utils/download.py seta_env +``` + +Convert tasks into training JSONL: + +```bash +python terminal-rl/data_utils/convert_task_to_dataset.py \ + --tasks_dir terminal-rl/dataset/seta_env +``` + +The `seta_env` dataset corresponds to the task dataset published in: [camel-ai/seta-env](https://github.com/camel-ai/seta-env/tree/main/Dataset). + +### 3. Run training + +On the training machine, set the required environment variables: + +```bash +# Hugging Face cache / model paths +export HF_HOME="/path/to/huggingface" +export MODEL_CKPT="/path/to/model" +export REF_LOAD="/path/to/reference_model_dir" +export SAVE_CKPT="/path/to/save/checkpoints" + +# Dataset + workers +export ROLLOUT_PROMPT_DATA="/path/to/train.jsonl" +export WORKER_URLS="http://worker1:18081,http://worker2:18081" + +# Logging +export WANDB_KEY="your-wandb-key" +``` + +Then run (from repo root): + +```bash +bash terminal-rl/terminal_qwen3_8b_rl.sh +``` + +--- + +### PRM training (optional) + +To enable PRM scoring with the 2-node script, add: + +```bash +export PRM_ENABLE=1 +export PRM_MODEL_PATH="/path/to/prm-model" +export PRM_M=3 +export PRM_STEP_COEF=1.0 +export PRM_TEMPERATURE=0.0 +export PRM_MAX_NEW_TOKENS=4096 +# Optional: use an external PRM endpoint instead of framework-hosted engines +export PRM_SGLANG_URL="http://:" +``` + +Then run: + +```bash +bash terminal-rl/terminal_qwen3_8b_prm_rl_2nodes.sh +``` + +--- + +## Notes + +- `WORKER_URLS` must point to already-running pool servers. +- As an example, one rollout agent implementation in this repo is based on **CAMEL** (see `terminal-rl/agent/camel_agent.py` and [CAMEL](https://github.com/camel-ai/camel)). diff --git a/terminal-rl/agent/camel_agent.py b/terminal-rl/agent/camel_agent.py new file mode 100644 index 0000000..53ef155 --- /dev/null +++ b/terminal-rl/agent/camel_agent.py @@ -0,0 +1,440 @@ +from __future__ import annotations + +import asyncio +import datetime +import json +import logging +import re +import threading +import time +import uuid +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type, Union + +from openai import AsyncStream, Stream +from openai.lib.streaming.chat import ( + AsyncChatCompletionStreamManager, + ChatCompletionStreamManager, +) +from openai.types.chat import ChatCompletion +from pydantic import BaseModel + +from camel.agents import ChatAgent +from camel.messages import BaseMessage, FunctionCallingMessage, OpenAIMessage +from camel.models import BaseModelBackend +from camel.responses import ChatAgentResponse +from camel.types import ChatCompletionChunk, ModelType, OpenAIBackendRole +from camel.types.agents import ToolCallingRecord +from camel.utils import OpenAITokenCounter +from camel.utils.token_counting import BaseTokenCounter + +from inference_client import SGLangTurnClient + +from .prompts import get_developer_agent_prompt + +logger = logging.getLogger(__name__) + + +if TYPE_CHECKING: + from transformers.tokenization_utils_fast import PreTrainedTokenizerFast + + +class HFTokenCounter(BaseTokenCounter): + """Token counter backed by the same HF tokenizer used for generation.""" + + def __init__( + self, tokenizer: "PreTrainedTokenizerFast", tokens_per_message: int = 3 + ) -> None: + self.tokenizer = tokenizer + self.tokens_per_message = tokens_per_message + + def count_tokens_from_messages(self, messages: List[OpenAIMessage]) -> int: + num_tokens = 0 + for message in messages: + num_tokens += self.tokens_per_message + for _, value in message.items(): + if not isinstance(value, list): + num_tokens += len(self.tokenizer.encode(str(value))) + continue + for item in value: + if isinstance(item, dict) and item.get("type") == "text": + num_tokens += len( + self.tokenizer.encode(str(item.get("text", ""))) + ) + else: + num_tokens += len(self.tokenizer.encode(str(item))) + num_tokens += 3 + return num_tokens + + def encode(self, text: str) -> List[int]: + return self.tokenizer.encode(text) + + def decode(self, token_ids: List[int]) -> str: + return self.tokenizer.decode(token_ids) + + +class CamelAgentBackend(BaseModelBackend): + """Backend adapter that can reuse the external SGLang turn client.""" + + def __init__( + self, + model_type: ModelType | str, + *, + sglang_client: SGLangTurnClient | None = None, + model_config_dict: Dict[str, Any] | None = None, + token_counter: BaseTokenCounter | None = None, + ) -> None: + super().__init__( + model_type=model_type, + model_config_dict=model_config_dict or {}, + api_key=None, + url=None, + token_counter=token_counter, + timeout=30.0, + max_retries=0, + ) + self._sglang_client = sglang_client + self._turn_counter = 0 + self.cache: dict[str, Any] = {} + + @property + def token_counter(self) -> BaseTokenCounter: + if not self._token_counter: + hf_tokenizer = None + if self._sglang_client is not None: + hf_tokenizer = getattr(self._sglang_client, "tokenizer", None) + if hf_tokenizer is not None: + self._token_counter = HFTokenCounter(hf_tokenizer) + else: + self._token_counter = OpenAITokenCounter(ModelType.GPT_4O_MINI) + return self._token_counter + + @property + def stream(self) -> bool: + return False + + def _run( + self, + messages: list[OpenAIMessage], + response_format: type[BaseModel] | None = None, + tools: list[dict[str, Any]] | None = None, + ) -> ( + ChatCompletion + | Stream[ChatCompletionChunk] + | ChatCompletionStreamManager[BaseModel] + ): + _ = (messages, response_format, tools) + raise RuntimeError("CamelAgentBackend._run is not used by AgentRunner.") + + async def _arun( + self, + messages: list[OpenAIMessage], + response_format: type[BaseModel] | None = None, + tools: list[dict[str, Any]] | None = None, + ) -> ( + ChatCompletion + | AsyncStream[ChatCompletionChunk] + | AsyncChatCompletionStreamManager[BaseModel] + ): + _ = response_format + if self._sglang_client is None: + raise RuntimeError("CamelAgentBackend has no SGLang client configured.") + + chat_completion, interaction = await self._sglang_client.generate_turn( + messages=messages, + tools=tools, + turn_idx=self._turn_counter, + ) + self.cache[chat_completion.id] = interaction + self._turn_counter += 1 + return chat_completion + + +# Adapted from https://github.com/camel-ai/seta/blob/main/training/tbench_areal_workflow/chat_agent_trace.py +class CamelAgent(ChatAgent): + """ChatAgent extension used by AgentRunner's rollout loop.""" + + def __init__( + self, + *, + model_type: str, + sglang_client: SGLangTurnClient, + non_think_mode: bool, + max_total_tokens: int, + max_parse_errors: int | None = None, + system: str = "Linux (in Docker)", + machine: str = "x86_64", + is_workforce: bool = False, + current_date: str | None = None, + ) -> None: + prompt_date = current_date or str(datetime.date.today()) + system_prompt = get_developer_agent_prompt( + current_date=prompt_date, + system=system, + machine=machine, + is_workforce=is_workforce, + non_think_mode=non_think_mode, + ) + backend = CamelAgentBackend(model_type=model_type, sglang_client=sglang_client) + + super().__init__( + system_message=BaseMessage.make_assistant_message( + role_name="Developer Agent", + content=system_prompt, + ), + model=backend, + tools=[], + token_limit=max_total_tokens, + ) + super().reset() + + self.max_parse_errors = max(1, int(max_parse_errors or 3)) + self.parse_error_count = 0 + self._tool_call_records: List[Any] = [] + self._accumulated_context_tokens = 0 + self._step_token_usage = self._create_token_usage_tracker() + self._original_response_format: Optional[Type[BaseModel]] = None + self._used_prompt_formatting: bool = False + + def set_max_parse_errors(self, max_parse_errors: int) -> None: + self.max_parse_errors = max(1, int(max_parse_errors)) + + def start_turn_loop( + self, + input_message: Union[BaseMessage, str], + response_format: Optional[Type[BaseModel]] = None, + ) -> None: + self.parse_error_count = 0 + self._tool_call_records = [] + self._accumulated_context_tokens = 0 + self._step_token_usage = self._create_token_usage_tracker() + + self._original_response_format = response_format + input_message, _response_format, used_prompt_formatting = ( + self._handle_response_format_with_non_strict_tools( + input_message, + response_format, + ) + ) + self._used_prompt_formatting = used_prompt_formatting + + if isinstance(input_message, str): + input_message = BaseMessage.make_user_message( + role_name="User", content=input_message + ) + self.update_memory(input_message, OpenAIBackendRole.USER) + + async def _wait_if_paused(self) -> None: + if self.pause_event is None or self.pause_event.is_set(): + return + if isinstance(self.pause_event, asyncio.Event): + await self.pause_event.wait() + return + if isinstance(self.pause_event, threading.Event): + loop = asyncio.get_running_loop() + await loop.run_in_executor(None, self.pause_event.wait) + + async def get_turn_context( + self, + ) -> tuple[Optional[List[dict[str, Any]]], Optional[ChatAgentResponse]]: + await self._wait_if_paused() + try: + context_messages, num_tokens = self.memory.get_context() + self._accumulated_context_tokens += num_tokens + return context_messages, None + except RuntimeError as exc: + terminated_response = self._step_terminate( + exc.args[1], + self._tool_call_records, + "max_tokens_exceeded", + ) + return None, terminated_response + + async def consume_completion( + self, chat_completion: Any + ) -> tuple[Optional[Any], List[Any], bool, Optional[ChatAgentResponse]]: + model_response = self._handle_batch_response(chat_completion) + self._update_token_usage_tracker( + self._step_token_usage, model_response.usage_dict + ) + + if self.stop_event and self.stop_event.is_set(): + logger.info("Termination triggered by stop_event") + terminated_response = self._step_terminate( + self._accumulated_context_tokens, + self._tool_call_records, + "termination_triggered", + ) + return model_response, [], False, terminated_response + + if model_response.tool_call_requests: + return model_response, list(model_response.tool_call_requests), False, None + + parse_error_record = await self.adetect_tool_calls_parse_error(model_response) + if parse_error_record: + logger.warning( + f"Detected tool call parse error, prompting model to correct." + ) + self._tool_call_records.append(parse_error_record) + return model_response, [], True, None + + return model_response, [], False, None + + async def adetect_tool_calls_parse_error(self, response): + r""" + Asynchronously detect tool calls in the response content using Qwen25Detector. + if the model is Qwen 2.5 or Qwen 3. + if there's tool call tokens detected, but got json parse failure, format the information into a tool call record, + so that the agent can handle the error next step. + add a self.count_parse_error, so that we can limit the number of parse errors we handle in one step. if max reached, just + break the loop. + + Args: + response: The model response to check for parse errors + + Returns: + Optional[ToolCallingRecord]: A tool calling record with error information if parse error detected, None otherwise + """ + bot_token = "\n" + eot_token = "\n" + + # Check if we've reached max parse errors + if self.parse_error_count >= self.max_parse_errors: + logger.warning( + f"Max parse errors ({self.max_parse_errors}) reached, stopping error handling" + ) + return None + + # Extract content from response + if not response.output_messages: + return None + + content = response.output_messages[0].content + if not content or bot_token not in content: + return None + + # Find all potential tool call blocks + pattern = rf"{re.escape(bot_token)}(.*?){re.escape(eot_token)}" + matches = re.findall(pattern, content, re.DOTALL) + + if not matches: + return None + + # Check each match for JSON parse errors + for match_text in matches: + try: + # Try to parse the JSON + json.loads(match_text.strip()) + # If successful, no error for this match + continue + except json.JSONDecodeError as e: + # Found a parse error + self.parse_error_count += 1 + logger.warning( + f"Detected JSON parse error (count: {self.parse_error_count}/{self.max_parse_errors}): {str(e)}" + ) + logger.warning(f"Problematic content: {match_text[:200]}...") + + # Create an error tool calling record + error_message = ( + f"JSON Parse Error: {str(e)}\n" + f"The tool call format is incorrect. Please ensure:\n" + f"1. The JSON is valid and properly formatted\n" + f"2. All quotes are properly escaped\n" + f"3. The structure matches: {{'name': 'function_name', 'arguments': {{}}}}\n" + f"Problematic content (first 200 chars): {match_text[:200]}..." + ) + + # Generate a unique error tool call ID + error_tool_call_id = f"error_{uuid.uuid4().hex[:8]}" + + # Create the error record + error_record = ToolCallingRecord( + tool_name="json_parse_error", + args={"raw_content": match_text, "error": str(e)}, + result=error_message, + tool_call_id=error_tool_call_id, + ) + + # Record this in memory so the model can see the error + assist_msg = FunctionCallingMessage( + role_name=self.role_name, + role_type=self.role_type, + meta_dict=None, + content="", + func_name="json_parse_error", + args={"raw_content": match_text[:200], "error": str(e)}, + tool_call_id=error_tool_call_id, + ) + + func_msg = FunctionCallingMessage( + role_name=self.role_name, + role_type=self.role_type, + meta_dict=None, + content="", + func_name="json_parse_error", + result=error_message, + tool_call_id=error_tool_call_id, + ) + + # Use precise timestamps + current_time_ns = time.time_ns() + base_timestamp = current_time_ns / 1_000_000_000 + + self.update_memory( + assist_msg, OpenAIBackendRole.ASSISTANT, timestamp=base_timestamp + ) + self.update_memory( + func_msg, + OpenAIBackendRole.FUNCTION, + timestamp=base_timestamp + 1e-6, + ) + + return error_record + + return None + + def record_tool_result(self, tool_call_request: Any, raw_result: Any) -> None: + func_name = tool_call_request.tool_name + args = tool_call_request.args + tool_call_id = tool_call_request.tool_call_id + + if self.mask_tool_output: + with self._secure_result_store_lock: + self._secure_result_store[tool_call_id] = raw_result + result = ( + "[The tool has been executed successfully, but the " + "output from the tool is masked. You can move forward]" + ) + else: + result = raw_result + + tool_record = self._record_tool_calling( + func_name, + args, + result, + tool_call_id, + mask_output=self.mask_tool_output, + extra_content=tool_call_request.extra_content, + ) + self._tool_call_records.append(tool_record) + + def finalize_response(self, model_response: Any) -> ChatAgentResponse: + if self._used_prompt_formatting and self._original_response_format: + self._apply_prompt_based_parsing( + model_response, self._original_response_format + ) + + self._record_final_output(model_response.output_messages) + + if self.prune_tool_calls_from_memory and self._tool_call_records: + self.memory.clean_tool_calls() + + return self._convert_to_chatagent_response( + model_response, + self._tool_call_records, + self._accumulated_context_tokens, + None, + self._step_token_usage["prompt_tokens"], + self._step_token_usage["completion_tokens"], + self._step_token_usage["total_tokens"], + ) diff --git a/terminal-rl/agent/prm_agent.py b/terminal-rl/agent/prm_agent.py new file mode 100644 index 0000000..3126485 --- /dev/null +++ b/terminal-rl/agent/prm_agent.py @@ -0,0 +1,217 @@ +import json +import random +from typing import Any, Dict, List, Optional + + +PRM_SYSTEM = """You are an evaluator for a terminal agent. +You are provided with: +1) the agent's task instruction, +2) the interaction history, and +3) the agent's most recent step to evaluate. +""" + +USER_INSTRUCTION = """ +Evaluate ONLY the single most recent step using the information above. + + +Assign a score of +1 if ALL of the following are true: +- The current assistant message is a correct/helpful step that advances the task; +- The tool-call format is valid; +- Tool usage is appropriate for the step; +- Tool results (if any) are consistent with making progress. + +Otherwise assign a score of -1, for example if: +- The step is incorrect, misleading, or does not advance the task; +- Tool-call format is broken (invalid JSON / parse error); +- Tool usage is clearly wrong or irrelevant; +- Tool results show failure or clearly no progress. + +Think carefully, then provide your reasoning and put the final score in \\boxed{}. +""" + + +import re + +_PRM_BOXED_PATTERN = re.compile(r"\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}", re.DOTALL) +_PRM_STRICT_NUMBER_PATTERN = re.compile(r"^\s*([-+]?\d+(?:\.\d+)?)\s*$") + + +def _extract_prm_sign_from_text(text: str) -> int: + if not text: + return 0 + match = _PRM_BOXED_PATTERN.search(text) + if not match: + return 0 + boxed_content = match.group(1).strip() + strict_number_match = _PRM_STRICT_NUMBER_PATTERN.fullmatch(boxed_content) + if not strict_number_match: + return 0 + try: + value = float(strict_number_match.group(1)) + except ValueError: + return 0 + if abs(value - 1.0) < 1e-9: + return 1 + if abs(value + 1.0) < 1e-9: + return -1 + return 0 + + +def _truncate(text: str, limit: int = 2000) -> str: + if not text: + return "" + return text if len(text) <= limit else (text[:limit] + "...") + + +class TerminalPRMAgent: + """ + PRM agent. + + history_mode: + - "last": last k turns + - "random": random sample k turns + - "head_tail": first k turns + last k turns + """ + + def __init__( + self, + *, + sglang_client, + task_instruction: str, + history_k: int = 3, + history_mode: str = "head_tail", + head_k: int = 2, + tail_k: int = 2, + history_include_assistant: bool = False, + current_truncate: int = 2000, + history_truncate: int = 5000, # characters not tokens + ): + self._sglang_client = sglang_client + self.task_instruction = task_instruction + + self.history_k = history_k + self.history_mode = history_mode + self.head_k = head_k + self.tail_k = tail_k + self.history_include_assistant = history_include_assistant + + self.current_truncate = current_truncate + self.history_truncate = history_truncate + + self._history: Dict[int, Dict[str, Any]] = {} + + def record_model_turn( + self, + turn_idx: int, + *, + assistant_text: str, + tool_calls: Optional[List[Dict[str, Any]]] = None, + parse_error_recorded: bool = False, + finish_reason: Optional[str] = None, + ) -> None: + rec = self._history.setdefault(turn_idx, {}) + rec["assistant_text"] = assistant_text + rec["tool_calls"] = tool_calls + rec["parse_error_recorded"] = bool(parse_error_recorded) + rec["finish_reason"] = finish_reason + rec.setdefault("tool_results", []) + + def record_tool_result( + self, turn_idx: int, tool_call_request, raw_result: Any + ) -> None: + rec = self._history.setdefault(turn_idx, {}) + lst = rec.setdefault("tool_results", []) + lst.append( + { + "name": tool_call_request.tool_name, + "args": tool_call_request.args, + "result": raw_result, + } + ) + + def get_history(self, current_turn_idx: int) -> List[Dict[str, Any]]: + prev = sorted(t for t in self._history.keys() if t < current_turn_idx) + if not prev: + return [] + + if self.history_mode == "last": + if self.history_k <= 0: + return [] + chosen = prev[-self.history_k :] + + elif self.history_mode == "random": + if self.history_k <= 0: + return [] + k = min(self.history_k, len(prev)) + chosen = random.sample(prev, k=k) + chosen = sorted(chosen) + + elif self.history_mode == "head_tail": + head = prev[: max(0, self.head_k)] + tail = prev[-max(0, self.tail_k) :] if self.tail_k > 0 else [] + # de-dup while preserving order + seen = set() + chosen = [] + for t in head + tail: + if t not in seen: + seen.add(t) + chosen.append(t) + + else: + raise ValueError(f"Invalid history mode: {self.history_mode}") + + hist: List[Dict[str, Any]] = [] + for t in chosen: + r = self._history.get(t, {}) + item = { + "turn_idx": t, + "assistant_text": ( + _truncate(r["assistant_text"], self.history_truncate) + if self.history_include_assistant + else "[OMITTED]" + ), + "tool_calls": r["tool_calls"], + "tool_results": r["tool_results"], + "parse_error_recorded": r["parse_error_recorded"], + } + hist.append(item) + return hist + + def _build_messages(self, turn_idx: int) -> List[Dict[str, str]]: + cur = self._history.get(turn_idx, {}) + history = self.get_history(turn_idx) + + payload = { + "task_instruction": self.task_instruction, + "history": history, + "current": { + "turn_idx": turn_idx, + "assistant_text": _truncate( + cur["assistant_text"], self.current_truncate + ), + "tool_calls": cur["tool_calls"], + "tool_results": cur["tool_results"], + "parse_error_recorded": cur["parse_error_recorded"], + }, + } + + return [ + {"role": "system", "content": PRM_SYSTEM}, + { + "role": "user", + "content": json.dumps(payload, ensure_ascii=False) + + "\n\n" + + USER_INSTRUCTION, + }, + ] + + async def judge_turn(self, turn_idx: int) -> int: + messages = self._build_messages(turn_idx) + _cc, interaction = await self._sglang_client.generate_turn( + messages=messages, + tools=None, + turn_idx=turn_idx, + ) + text = interaction.output_text + score = _extract_prm_sign_from_text(text[-50:]) + return text, score diff --git a/terminal-rl/agent/prompts.py b/terminal-rl/agent/prompts.py new file mode 100644 index 0000000..e6e046b --- /dev/null +++ b/terminal-rl/agent/prompts.py @@ -0,0 +1,149 @@ +def get_developer_agent_prompt(current_date:str, system:str, machine:str, is_workforce:bool, non_think_mode:bool=True): + """ + Generate the prompt for the Lead Software Engineer agent. + Args: + current_date (str): The current date. + system (str): The operating system. (e.g., "Linux", "Darwin", "Windows", "Linux (in Docker)"...) + machine (str): The machine type. (e.g., "x86_64", "arm64") + is_workforce (bool): Whether the agent is part of a workforce with other agents or standalone. + Returns: + str: The prompt for the Lead Software Engineer agent. + """ + LEAD_SDE_ROLE_PROMPT = f""" + +You are a Lead Software Engineer, a master-level coding assistant with a +powerful and unrestricted terminal. Your primary role is to solve any +technical task by analyzing the problem, making plans, +writing and executing code, installing necessary libraries, +interacting with the operating system, and deploying applications. You are the +team's go-to expert for all technical implementation. + +""" + TEAM_STRUCTURE_PROMPT = f"" + + OPERATING_ENVIRONMENT_PROMPT = f""" + +- **System**: {system} ({machine}). +""" \ ++ \ +(""" +Note that the terminal commands and file system operations you perform will be +executed inside a Docker container. But note taking tools will operate on the host system. +""") if "Docker" in system else ""\ ++ \ +f""" +- **Current Date**: {current_date}. + +""" + + MANDATORY_INSTRUCTIONS_PROMPT = f""" + +- You MUST use analyze, plan and review requirements and your work. +- When you complete your task, your final response must be a comprehensive +summary of your work and the outcome, presented in a clear, detailed, and +easy-to-read format. Avoid using markdown tables for presenting data; use +plain text formatting instead. +- You MUST use tools and follow tool schemas precisely for every response, +- You MUST be concise about your reasoning and planning, and limit within 600 tokens. +- You MUST try diverse tools available in toolkits. + +""" + CAPABILITIES_PROMPT = """ + +Your capabilities are extensive and powerful: +- **Unrestricted Code Execution**: You can write and execute code in any +language to solve a task. +- For multi-line code, You MUST use tool (shell_write_content_to_file) to first save your code +to somewhere on the system (e.g.,`script.py`) and then run it from the terminal (e.g., +`python script.py`). Beware of the code that includes quotes\"\'; ensure proper +escaping when writing arguments for toolkit. Make sure it can be parsed by JSON. +- **Full Terminal Control**: You have root-level access to the terminal. You +can run any command-line tool, manage files, and interact with the OS. If +a tool is missing, you MUST install it with the appropriate package manager +(e.g., `pip3`, `uv`, or `apt-get`). Your capabilities include: + - **Text & Data Processing**: `awk`, `sed`, `grep`, `jq`. + - **File System & Execution**: `find`, `xargs`, `tar`, `zip`, `unzip`, + `chmod`. + - **Networking & Web**: `curl`, `wget` for web requests; `ssh` for + remote access. +- **IMPORTANT**: Always complete the full automation workflow—do not just +prepare or suggest actions. Execute them to completion. +- **Solution Verification**: You can immediately test and verify your +solutions by executing them in the terminal. +""" + \ +""" + +""" + + PHILOSOPHY_PROMPT = """ + +- **Bias for Action**: Your purpose is to take action. Don't just suggest +solutions—implement them. Write code, run commands, and build things. +- **Complete the Full Task**: When automating GUI applications, always finish +what you start. If the task involves sending something, send it. If it +involves submitting data, submit it. Never stop at just preparing or +drafting—execute the complete workflow to achieve the desired outcome. +- **Embrace Challenges**: Never say "I can't." If you +encounter a limitation, find a way to overcome it. +- **Resourcefulness**: If a tool is missing, install it. If information is +lacking, find it. You have the full power of a terminal to acquire any +resource you need. +- **Think Like an Engineer**: Approach problems methodically. Analyze +requirements, execute it, and verify the results. Your +strength lies in your ability to engineer solutions. +- ** Use Absolute Paths**: You can access files from any place in the file +system. For all file system operations, you MUST use absolute paths to ensure +precision and avoid ambiguity. +- ** Check current directory**: Always check your current directory with `pwd` and list +files with `ls -la` before performing file operations. This helps you +understand your context and avoid mistakes. +- ** Search for Files**: If you need a file but cannot find it in the current directory, +use commands like `find / -name "filename"` or search in directories common for the System +to locate it anywhere in the file system. This ensures you can always access the resources you need. +- ** Adhere to the initial task instruction**: Always keep the original task instruction in mind, make sure to understand +all requirements and useful information. Make sure finish every subtask mentioned in the instruction. + +""" + + SHELL_TIPS_PROMPT = f""" + +The terminal tools are session-based, identified by a unique `id`. Master +these tips to maximize your effectiveness: + +- **Command-Line Best Practices**: +- **Be Creative**: The terminal is your most powerful tool. Use it boldly. +- **Automate Confirmation**: Use `-y` or `-f` flags to avoid interactive +prompts. +- **Manage Output**: Redirect long outputs to a file (e.g., `> output.txt`). +- **Chain Commands**: Use `&&` to link several commands for sequential execution. +But also avoid chaining too many commands in one line +to avoid json parse errors due to complex escaping issues. +- **Piping**: Use `|` to pass output from one command to another. +- **Permissions**: Use `ls -F` to check file permissions. +- **Installation**: Use `pip3 install` or `apt-get install` for new +packages. +- **Time Management**: `shell_exec` commands come with block or non-block mode. The block mode + has a time limit, and only suitable for very quick commands. If you expect a command to take a long time, or + you have experienced a timeout for a command, you MUST use non-block mode by setting `block=False`. + The non-block mode allows commands to run in the background. You can check the status using `shell_view`, + send in further input using `shell_write_to_process`, and kill it using `shell_kill_process` if needed. + + +""" + COLLABORATION_AND_ASSISTANCE_PROMPT = f""" + """ + + FINAL_INSTRUCTIONS_PROMPT = f""" + {LEAD_SDE_ROLE_PROMPT} + {TEAM_STRUCTURE_PROMPT} + {OPERATING_ENVIRONMENT_PROMPT} + {MANDATORY_INSTRUCTIONS_PROMPT} + {CAPABILITIES_PROMPT} + {PHILOSOPHY_PROMPT} + {SHELL_TIPS_PROMPT} + {COLLABORATION_AND_ASSISTANCE_PROMPT} + """ + if non_think_mode: + FINAL_INSTRUCTIONS_PROMPT = rf"{FINAL_INSTRUCTIONS_PROMPT} /no_think" + + return FINAL_INSTRUCTIONS_PROMPT diff --git a/terminal-rl/agent_runner.py b/terminal-rl/agent_runner.py new file mode 100644 index 0000000..f9d44b0 --- /dev/null +++ b/terminal-rl/agent_runner.py @@ -0,0 +1,137 @@ +from __future__ import annotations + +from typing import Any, Dict, List, Optional, Protocol + +from custom_types import TurnContext, TurnResult +from inference_client import SGLangTurnClient + + +class RolloutAgent(Protocol): + @property + def parse_error_count(self) -> int: ... + + def set_max_parse_errors(self, max_parse_errors: int) -> None: ... + + def start_turn_loop(self, input_message: Any) -> None: ... + + async def get_turn_context( + self, + ) -> tuple[Optional[List[dict[str, Any]]], Optional[Any]]: ... + + async def consume_completion( + self, chat_completion: Any + ) -> tuple[Optional[Any], List[Any], bool, Optional[Any]]: ... + + def record_tool_result(self, tool_call_request: Any, raw_result: Any) -> None: ... + + def finalize_response(self, model_response: Any) -> Any: ... + + +class AgentRunner: + def __init__( + self, + *, + rollout_agent: RolloutAgent, + sglang_client: SGLangTurnClient, + tool_schemas: List[Dict[str, Any]], + ) -> None: + self._rollout_agent = rollout_agent + self._sglang_client = sglang_client + self._tool_schemas = tool_schemas + self._model_turn_count = 0 + self._max_iterations = 10 + self._max_parse_errors = 3 + + @property + def model_turn_count(self) -> int: + return self._model_turn_count + + @property + def parse_error_count(self) -> int: + return self._rollout_agent.parse_error_count + + @property + def max_iterations(self) -> int: + return self._max_iterations + + @property + def max_parse_errors(self) -> int: + return self._max_parse_errors + + def reset(self, input_message: Any) -> None: + self._model_turn_count = 0 + self._rollout_agent.start_turn_loop(input_message) + + def set_max_parse_errors(self, max_parse_errors: int) -> None: + self._max_parse_errors = max(1, int(max_parse_errors)) + self._rollout_agent.set_max_parse_errors(self._max_parse_errors) + + def set_max_iterations(self, max_iterations: int) -> None: + self._max_iterations = max(1, int(max_iterations)) + + def reached_iteration_limit(self) -> bool: + return self._model_turn_count >= self._max_iterations + + def reached_parse_error_limit(self) -> bool: + return self.parse_error_count >= self._max_parse_errors + + async def get_turn_context(self) -> TurnContext: + messages, terminated = await self._rollout_agent.get_turn_context() + return TurnContext(context_messages=messages, terminated_response=terminated) + + async def run_model_turn( + self, context_messages: List[dict[str, Any]] + ) -> TurnResult: + chat_completion, interaction = await self._sglang_client.generate_turn( + messages=context_messages, + tools=self._tool_schemas, + turn_idx=self._model_turn_count, + ) + self._model_turn_count += 1 + + model_response, tool_call_requests, parse_error_recorded, terminated = ( + await self._rollout_agent.consume_completion(chat_completion) + ) + return TurnResult( + interaction=interaction, + model_response=model_response, + tool_call_requests=tool_call_requests, + parse_error_recorded=parse_error_recorded, + terminated_response=terminated, + ) + + def record_tool_result(self, tool_call_request: Any, raw_result: Any) -> None: + self._rollout_agent.record_tool_result(tool_call_request, raw_result) + + def finalize_response(self, model_response: Any) -> Any: + return self._rollout_agent.finalize_response(model_response) + + +def create_agent_runner( + *, + agent_type: str, + sglang_client: SGLangTurnClient, + model_type: str, + tool_schemas: List[Dict[str, Any]], + non_think_mode: bool, + max_total_tokens: int, +) -> AgentRunner: + if agent_type == "camel_agent": + from agent.camel_agent import CamelAgent + + rollout_agent = CamelAgent( + model_type=model_type, + sglang_client=sglang_client, + non_think_mode=non_think_mode, + max_total_tokens=max_total_tokens, + ) + else: + raise ValueError( + f"Unsupported agent type: {agent_type!r}. Expected 'camel_agent'." + ) + + return AgentRunner( + rollout_agent=rollout_agent, + sglang_client=sglang_client, + tool_schemas=tool_schemas, + ) diff --git a/terminal-rl/custom_types.py b/terminal-rl/custom_types.py new file mode 100644 index 0000000..b4cb5b7 --- /dev/null +++ b/terminal-rl/custom_types.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Optional, List + + +@dataclass(frozen=True) +class TaskSpec: + task_name: str + task_path: str + instruction: str + + +@dataclass(frozen=True) +class RunContext: + uid: str + group_index: int + sample_index: int + log_dir: Path + + def to_payload(self) -> dict[str, Any]: + return { + "uid": self.uid, + "group_index": self.group_index, + "sample_index": self.sample_index, + "log_dir": str(self.log_dir), + } + + +@dataclass +class TaskTimeouts: + ensure_image: float = 300.0 + reset_session: float = 300.0 + close_session: float = 60.0 + eval: float = 600.0 + + def to_payload(self) -> dict[str, float]: + return { + "ensure_image": float(self.ensure_image), + "reset_session": float(self.reset_session), + "close_session": float(self.close_session), + "eval": float(self.eval), + } + + +from openai.types.chat.chat_completion import ChatCompletion +from openai.types.chat.chat_completion_message_param import ( + ChatCompletionMessageParam as OpenAIMessage, +) + + +@dataclass +class Interaction: + turn_idx: int = 0 + completion: ChatCompletion | None = None + input_ids: list[int] = field(default_factory=list) + output_token_ids: list[int] = field(default_factory=list) + output_token_logprobs: list[float] = field(default_factory=list) + output_text: str = "" + finish_reason: str = "" + messages: list[OpenAIMessage] = field(default_factory=list) + latency_ms: float = 0.0 + + +@dataclass +class TurnContext: + context_messages: Optional[List[dict[str, Any]]] + terminated_response: Optional[Any] = None + + +@dataclass +class TurnResult: + interaction: Interaction + model_response: Optional[Any] + tool_call_requests: List[Any] + parse_error_recorded: bool + terminated_response: Optional[Any] = None diff --git a/terminal-rl/data_utils/convert_task_to_dataset.py b/terminal-rl/data_utils/convert_task_to_dataset.py new file mode 100644 index 0000000..f4390af --- /dev/null +++ b/terminal-rl/data_utils/convert_task_to_dataset.py @@ -0,0 +1,172 @@ +"""Convert Terminal Bench tasks to RLLM/VERL format.""" + +import os +import json +import pandas as pd +from pathlib import Path +from typing import List, Optional +from tqdm import tqdm + +# Add project to path +import sys +from load_tasks import TBenchTrainingTask, load_terminal_bench_tasks + +DATASET_DIR = Path(os.getenv("DATASET_DIR", "./terminal-rl/dataset")) + + +def convert_tasks( + tasks_dir: List[Path], + train_split: Optional[float] = None, + system_prompt: Optional[str] = None, + task_names: Optional[List[str]] = None, + test_tasks_dir: Optional[Path] = None, + format: Optional[str] = "jsonl", + output_dir: Optional[Path] = None, +) -> None: + """Convert terminal bench tasks to parquet format for VERL training. + + Args: + tasks_dir: Directory containing terminal bench tasks (or train tasks if test_tasks_dir is provided) + train_split: Fraction of data for training (ignored if test_tasks_dir is provided) + system_prompt: System prompt to use + task_names: Specific task names to convert + test_tasks_dir: Directory containing test tasks for validation set + """ + + # Load tasks + print(f"Loading tasks from {tasks_dir}") + tasks = [] + for dir_path in tasks_dir: + tasks.extend(load_terminal_bench_tasks(dir_path, task_names)) + print(f"Loaded {len(tasks)} tasks") + + if output_dir is None: + output_dir = DATASET_DIR + + # Create output directory + output_dir.mkdir(parents=True, exist_ok=True) + + # Prepare data for parquet + data_records = [] + + print(f"len tasks {len(tasks)}") + + for task in tqdm(tasks, desc="Converting tasks"): + # find path relative to outdir + task_path = task.task_path.relative_to(DATASET_DIR) + print(f"Processing task: {task.task_name} at {task_path}") + record = { + "task_name": task.task_name, + "task_path": str(task_path), + "instruction": task.instruction, + "data_source": "terminal_bench", # For reward_fn_key + } + + data_records.append(record) + + # Create DataFrame + df = pd.DataFrame(data_records) + + # Split into train and validation + # Use train_split parameter + if train_split is None: + train_split = 1.0 + n_train = int(len(df) * train_split) + train_df = df[:n_train] + val_df = df[n_train:] + + # Save to file(s) + train_path = output_dir / f"train.{format}" + val_path = output_dir / f"val.{format}" + + if format == "jsonl": + train_wrapped = pd.DataFrame({"task": train_df.to_dict(orient="records")}) + val_wrapped = pd.DataFrame({"task": val_df.to_dict(orient="records")}) + train_wrapped.to_json( + train_path, orient="records", lines=True, force_ascii=False + ) + val_wrapped.to_json(val_path, orient="records", lines=True, force_ascii=False) + elif format == "parquet": + train_df.to_parquet(train_path, index=False) + val_df.to_parquet(val_path, index=False) + + print(f"Saved {len(train_df)} training examples to {train_path}") + print(f"Saved {len(val_df)} validation examples to {val_path}") + + +def main(tasks_dir, test_tasks_dir, train_split, depth, output_dir): + + # Check if directories exist + if not tasks_dir.exists(): + print(f"Error: {tasks_dir} directory not found") + return + + if depth == 0: + tasks_dirs = [tasks_dir] + elif depth == 1: + tasks_dirs = [p for p in tasks_dir.iterdir() if p.is_dir()] + print(f"Found {tasks_dirs} task directories at depth {depth}") + + if output_dir is None: + output_dir = DATASET_DIR / f"{tasks_dir.name}_convert" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"Writing converted dataset to {output_dir}") + + # Convert to parquet only (with extra_info) + convert_tasks( + tasks_dir=tasks_dirs, + train_split=train_split, + system_prompt=None, + task_names=None, + test_tasks_dir=test_tasks_dir, + format="jsonl", + output_dir=output_dir, + ) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Convert Terminal Bench tasks to RLLM/VERL dataset format." + ) + parser.add_argument( + "--tasks_dir", + type=str, + help="Directory containing terminal bench tasks (or train tasks if --test_tasks_dir is provided)", + ) + parser.add_argument( + "--test_tasks_dir", + type=str, + default=None, + help="Directory containing test tasks for validation set (if provided)", + ) + parser.add_argument( + "--train_split", + type=float, + default=1.0, + help="Fraction of data for training (ignored if --test_tasks_dir is provided)", + ) + parser.add_argument( + "-d", + "--depth", + type=int, + default=0, + help="Depth of directory traversal ", + ) + parser.add_argument( + "--output_dir", + type=str, + default=None, + help="Output directory for converted dataset (default: /_convert)", + ) + args = parser.parse_args() + main( + tasks_dir=Path(args.tasks_dir), + test_tasks_dir=Path(args.test_tasks_dir) if args.test_tasks_dir else None, + train_split=args.train_split, + depth=args.depth, + output_dir=args.output_dir, + ) diff --git a/terminal-rl/data_utils/download.py b/terminal-rl/data_utils/download.py new file mode 100644 index 0000000..7e7fdbd --- /dev/null +++ b/terminal-rl/data_utils/download.py @@ -0,0 +1,139 @@ +# Script to download and prepare currently supported datasets for terminal agent training + +import os +import shutil +import subprocess +from pathlib import Path + +DATASET_DIR = Path(os.getenv("DATASET_DIR", "./terminal-rl/dataset")) + + +def _download_github_folder( + repo_url, sparse_path, target_dir, branch="main", temp_suffix="temp" +): + """ + General function to download a specific folder from a GitHub repository. + + Args: + repo_url: GitHub repository URL (.git) + sparse_path: Path within the repo to download + target_dir: Local destination directory + branch: Git branch to checkout (default: "main") + temp_suffix: Suffix for temporary directory name + """ + if target_dir.exists(): + print(f"Dataset already exists at {target_dir}. Skipping download.") + return + + DATASET_DIR.mkdir(parents=True, exist_ok=True) + temp_dir = DATASET_DIR / f"temp_{temp_suffix}" + + try: + # Clone with sparse checkout + subprocess.run( + [ + "git", + "clone", + "--depth", + "1", + "--filter=blob:none", + "--sparse", + repo_url, + str(temp_dir), + "-b", + branch, + ], + check=True, + ) + subprocess.run( + ["git", "-C", str(temp_dir), "sparse-checkout", "set", sparse_path], + check=True, + ) + + # Move downloaded folder to target location + shutil.move(str(temp_dir / sparse_path), str(target_dir)) + print(f"Successfully downloaded to {target_dir}") + finally: + if temp_dir.exists(): + shutil.rmtree(temp_dir) + + +def download_seta_env(): + url = "https://github.com/camel-ai/seta-env.git" + target_dir = DATASET_DIR / "seta_env" + _download_github_folder( + url, "Dataset", target_dir, branch="main", temp_suffix="seta_env" + ) + + +def download_tbench_core(): + url = "https://github.com/laude-institute/terminal-bench.git" + target_dir = DATASET_DIR / "tbench_core" + _download_github_folder( + url, "tasks", target_dir, branch="main", temp_suffix="tbench_core" + ) + + +def download_tbench_test(): + url = "https://github.com/laude-institute/terminal-bench.git" + target_dir = DATASET_DIR / "tbench_test" + _download_github_folder( + url, + "tasks", + target_dir, + branch="dataset/terminal-bench-core/v0.1.x", + temp_suffix="tbench_test", + ) + + +def download_tbench_adapted(): + url = "https://github.com/laude-institute/terminal-bench-datasets.git" + raw_dir = DATASET_DIR / "tbench_adapted_raw" + target_dir = DATASET_DIR / "tbench_adapted" + + if target_dir.exists(): + print(f"Dataset already exists at {target_dir}. Skipping download.") + return + + # Download the raw datasets + _download_github_folder( + url, "datasets", raw_dir, branch="main", temp_suffix="tbench_adapted" + ) + + # Create target directory + target_dir.mkdir(parents=True, exist_ok=True) + + # Create symbolic links with prefixed names + for subfolder in raw_dir.iterdir(): + if subfolder.is_dir(): + subfolder_name = subfolder.name + for task_folder in subfolder.iterdir(): + if task_folder.is_dir(): + task_name = task_folder.name + prefixed_name = f"{subfolder_name}_{task_name}" + symlink_path = target_dir / prefixed_name + symlink_path.symlink_to(task_folder, target_is_directory=True) + + print(f"Successfully created symlinks in {target_dir}") + + +def download_data(ds_name): + DATASET_DOWNLOADERS = { + "seta_env": download_seta_env, + "tbench_core": download_tbench_core, + "tbench_test": download_tbench_test, + "tbench_adapted": download_tbench_adapted, + } + if ds_name not in DATASET_DOWNLOADERS: + raise ValueError(f"Dataset {ds_name} is not supported.") + DATASET_DOWNLOADERS[ds_name]() + + +if __name__ == "__main__": + import sys + + if len(sys.argv) > 1: + download_data(sys.argv[1]) + else: + print("Available datasets: seta_env, tbench_core, tbench_test, tbench_adapted") + print("Usage: python download_data.py ") diff --git a/terminal-rl/data_utils/load_tasks.py b/terminal-rl/data_utils/load_tasks.py new file mode 100644 index 0000000..d0cbc09 --- /dev/null +++ b/terminal-rl/data_utils/load_tasks.py @@ -0,0 +1,119 @@ +import json +from pathlib import Path +from typing import List, Optional +from pydantic import BaseModel +import yaml + + +class TBenchTrainingTask(BaseModel): + """Data model for a task which follows the same format as terminal-bench.""" + + task_name: str + task_path: Path + instruction: str + # test_weights: dict + # dockerfile_contents: str + # py_test_file_contents: str + # max_test_timeout_sec: int = 300 # Default timeout + # additional_files: Optional[dict] = None # Maps file paths to contents + + +def load_terminal_bench_tasks( + tasks_dir: Path, + task_names: Optional[List[str]] = None, +) -> List[TBenchTrainingTask]: + if task_names is None: + print(f"tasks_dir {tasks_dir}") + import os + + print(f"len {len(os.listdir(str(tasks_dir)))}") + task_names = [p.name for p in tasks_dir.iterdir() if p.is_dir()] + + print(f"task names {task_names}") + tasks = [] + from tqdm import tqdm + + for task_name in tqdm(task_names): + try: + task_path = tasks_dir / task_name + task_yaml = task_path / "task.yaml" + + if not task_yaml.exists(): + raise FileNotFoundError(f"Task YAML file not found: {task_yaml}") + + with open(task_yaml, "r", encoding="utf-8") as f: + task_data = yaml.safe_load(f) + + instruction = task_data.get("instruction") + if not instruction: + raise ValueError(f"Instruction not found in task YAML: {task_yaml}") + + # # Get max test timeout if specified + # max_test_timeout_sec = task_data.get("max_test_timeout_sec", 300) + + # # Load test weights + # test_weights_path = task_path / "test_weights.json" + # if test_weights_path.exists(): + # with open(test_weights_path, "r", encoding="utf-8") as f: + # test_weights = json.load(f) + # else: + # test_weights = {"default": 1.0} + + # # Load Dockerfile + # dockerfile_path = task_path / "Dockerfile" + # if not dockerfile_path.exists(): + # continue + # with open(dockerfile_path, "r", encoding="utf-8") as f: + # dockerfile_contents = f.read() + + # # Load Python test file if it exists + # py_test_file_path = task_path / "tests" / "test_outputs.py" + # if not py_test_file_path.exists(): + # # raise FileNotFoundError(f"Python test file not found: {py_test_file_path}") + # continue + # with open(py_test_file_path, "r", encoding="utf-8") as f: + # py_test_file_contents = f.read() + # if not py_test_file_contents: + # # raise ValueError(f"Python test file is empty for task: {task_name}") + # continue + + # # Load additional files if they exist + # additional_files = {} + # # List all files in the task directory (excluding standard files) + # standard_files = {'Dockerfile', 'task.yaml', 'test_weights.json'} + # standard_dirs = {'tests', '__pycache__'} + + # for item in task_path.iterdir(): + # if item.is_file() and item.name not in standard_files: + # # Read the file and store with relative path + # rel_path = item.relative_to(task_path) + # with open(item, "r", encoding="utf-8") as f: + # additional_files[str(rel_path)] = f.read() + # elif item.is_dir() and item.name not in standard_dirs: + # # Recursively read files from subdirectories + # for subfile in item.rglob("*"): + # if subfile.is_file(): + # rel_path = subfile.relative_to(task_path) + # try: + # with open(subfile, "r", encoding="utf-8") as f: + # additional_files[str(rel_path)] = f.read() + # except UnicodeDecodeError: + # # Skip binary files for now + # pass + + tasks.append( + TBenchTrainingTask( + task_name=task_name, + task_path=task_path, + instruction=instruction, + # test_weights=None, + # dockerfile_contents=None, + # py_test_file_contents=None, + # max_test_timeout_sec=None, + # additional_files=None, + ) + ) + except Exception as e: + print(f"Error loading task {task_name}: {e}") + continue + return tasks diff --git a/terminal-rl/env_client.py b/terminal-rl/env_client.py new file mode 100644 index 0000000..3a0cf50 --- /dev/null +++ b/terminal-rl/env_client.py @@ -0,0 +1,116 @@ +from __future__ import annotations + +import logging +import os +from typing import Any + +from slime.utils.http_utils import post + +logger = logging.getLogger(__name__) + + +class TerminalEnvClient: + + def __init__(self, base_url: str): + self.base_url = base_url.rstrip("/") + self.default_max_retries = int(os.getenv("ENV_HTTP_MAX_RETRIES", "10")) + self.allocate_max_retries = int(os.getenv("ENV_ALLOCATE_MAX_RETRIES", "100")) + self.evaluate_max_retries = int(os.getenv("ENV_EVALUATE_MAX_RETRIES", "1")) + self.close_max_retries = int(os.getenv("ENV_CLOSE_MAX_RETRIES", "3")) + self.exec_tool_max_retries = int(os.getenv("ENV_EXEC_TOOL_MAX_RETRIES", "3")) + + async def allocate( + self, + task_key: str, + request_id: str | None = None, + ) -> dict[str, Any]: + out = await post( + f"{self.base_url}/allocate", + {"task_key": task_key, "request_id": request_id}, + max_retries=self.allocate_max_retries, + ) + if not out.get("ok", False): + raise RuntimeError(f"allocate failed: {out}") + return out + + async def heartbeat(self, lease_id: str) -> None: + out = await post( + f"{self.base_url}/heartbeat", + {"lease_id": lease_id}, + max_retries=self.default_max_retries, + ) + if not out.get("ok", False): + raise RuntimeError(f"heartbeat failed: {out}") + + async def reset( + self, + lease_id: str, + task_meta: dict[str, Any], + run_ctx: dict[str, Any], + task_timeouts: dict[str, Any] | None = None, + ) -> dict[str, Any]: + out = await post( + f"{self.base_url}/reset", + { + "lease_id": lease_id, + "task_meta": task_meta, + "run_ctx": run_ctx, + "task_timeouts": task_timeouts, + }, + max_retries=self.default_max_retries, + ) + if not out.get("ok", False): + raise RuntimeError(f"reset failed: {out}") + return out + + async def exec_tool( + self, lease_id: str, tool_name: str, arguments: dict[str, Any] + ) -> str: + out = await post( + f"{self.base_url}/exec_tool", + { + "lease_id": lease_id, + "tool_call": {"name": tool_name, "arguments": arguments}, + }, + max_retries=self.exec_tool_max_retries, + ) + if not out.get("ok", False): + raise RuntimeError(f"exec_tool failed: {out}") + return str(out.get("observation", "")) + + async def evaluate(self, lease_id: str) -> float: + out = await post( + f"{self.base_url}/evaluate", + {"lease_id": lease_id}, + max_retries=self.evaluate_max_retries, + ) + if not out.get("ok", False): + raise RuntimeError(f"evaluate failed: {out}") + return float(out.get("score", 0.0)) + + async def close(self, lease_id: str) -> None: + try: + out = await post( + f"{self.base_url}/close", + {"lease_id": lease_id}, + max_retries=self.close_max_retries, + ) + except Exception as exc: + error_str = str(exc) + resp_text = "" + if hasattr(exc, "response"): + try: + resp_text = exc.response.text + except Exception: + pass + combined = f"{error_str} {resp_text}" + if "Unknown run_lease_id" in combined or "Unknown lease" in combined: + logger.debug("close(%s): lease already gone, nothing to do.", lease_id) + return + raise + if not out.get("ok", False): + error_msg = str(out.get("error", "")) + if "Unknown" in error_msg and "lease" in error_msg.lower(): + logger.debug("close(%s): lease already gone, nothing to do.", lease_id) + return + raise RuntimeError(f"close failed: {out}") diff --git a/terminal-rl/generate.py b/terminal-rl/generate.py new file mode 100644 index 0000000..55db352 --- /dev/null +++ b/terminal-rl/generate.py @@ -0,0 +1,658 @@ +from __future__ import annotations + +import logging +import os +import uuid +from copy import deepcopy +from pathlib import Path +from typing import Any, Dict, List, Optional +import asyncio + +from slime.rollout.sglang_rollout import GenerateState +from slime.utils.types import Sample + +from agent.prm_agent import TerminalPRMAgent +from custom_types import ( + Interaction, + RunContext, + TaskSpec, + TaskTimeouts, + TurnContext, + TurnResult, +) +from inference_client import SGLangTurnClient +from agent_runner import create_agent_runner +from env_client import TerminalEnvClient + +logger = logging.getLogger(__name__) + + +def _extract_task_meta(sample: Sample) -> Dict[str, Any]: + if isinstance(sample.prompt, dict): + return sample.prompt + + metadata = sample.metadata or {} + task_meta = metadata.get("task_meta") if isinstance(metadata, dict) else None + if isinstance(task_meta, dict): + return task_meta + + if isinstance(metadata, dict): + return metadata + + return {} + + +def _make_task_spec(meta: Dict[str, Any]) -> TaskSpec: + return TaskSpec( + task_name=meta.get("task_name", "unknown"), + task_path=meta.get("task_path", ""), + instruction=meta.get("instruction", ""), + ) + + +def _build_samples( + interactions: List[Interaction], + base_sample: Sample, + outcome: float, + status: Sample.Status, + prm_turn_scores: dict[int, float] | None = None, + prm_coef: float = 1.0, + discount: float = 1.0, + encourage: bool = False, +) -> List[Sample]: + """Create one Sample per interaction with discounted reward.""" + num_turns = len(interactions) + samples: List[Sample] = [] + + accuracy = float(outcome) + raw_score = accuracy + (accuracy == 1.0) * int(encourage) + base_outcome = 2.0 * accuracy - 1.0 + + for interaction in interactions: + turn_idx = interaction.turn_idx + s = deepcopy(base_sample) + s.tokens = interaction.input_ids + interaction.output_token_ids + s.response_length = len(interaction.output_token_ids) + s.loss_mask = [1] * s.response_length + s.rollout_log_probs = list(interaction.output_token_logprobs) + s.response = interaction.output_text + s.status = status + + s.metadata.update( + { + "turn_idx": turn_idx, + "num_turns": num_turns, + "finish_reason": interaction.finish_reason, + "latency_ms": interaction.latency_ms, + } + ) + + steps_from_end = num_turns - 1 - turn_idx + discounted_base = base_outcome * (discount**steps_from_end) + + if prm_turn_scores is not None: + prm = prm_turn_scores.get(turn_idx, 0.0) + final = discounted_base + prm_coef * prm + s.metadata["step_wise"] = { + "step_scores": [prm], + "step_scores_with_outcome": [final], + "step_indices": [turn_idx], + "step_token_spans": [[0, s.response_length]], + } + else: + final = discounted_base + + s.reward = { + "accuracy": accuracy, + "raw_score": raw_score, + "base_score": discounted_base, + "score": final, + } + + if prm_turn_scores is not None: + s.reward["prm_turn_score"] = prm + samples.append(s) + + return samples + + +def _mark_non_trainable_samples(samples: List[Sample]) -> None: + for sample in samples: + if sample.status in {Sample.Status.ABORTED, Sample.Status.FAILED}: + if sample.reward is None: + sample.reward = {"score": 0.0} + sample.remove_sample = True + + +def _infer_completion_budget(sampling_params: Dict[str, Any]) -> int: + for key in ("max_new_tokens", "max_tokens", "max_completion_tokens"): + raw_value = sampling_params.get(key) + if raw_value is None: + continue + try: + parsed = int(raw_value) + except (TypeError, ValueError): + continue + if parsed > 0: + return parsed + return 0 + + +def _normalize_tool_schemas(raw_tools: List[Any]) -> List[Dict[str, Any]]: + schemas: List[Dict[str, Any]] = [] + for tool in raw_tools: + if hasattr(tool, "get_openai_tool_schema") and callable( + tool.get_openai_tool_schema + ): + schemas.append(tool.get_openai_tool_schema()) + elif isinstance(tool, dict): + schemas.append(tool) + else: + raise TypeError(f"Unsupported tool schema object type: {type(tool)!r}") + return schemas + + +async def _create_env_client( + task_spec: TaskSpec, + run_ctx: RunContext, +) -> tuple[TerminalEnvClient, str]: + env_server_url = os.getenv("ENV_SERVER_URL", "") + if not env_server_url: + raise RuntimeError("ENV_SERVER_URL is empty.") + + env_client = TerminalEnvClient(env_server_url) + task_key = f"{task_spec.task_name}:{task_spec.task_path}" + request_id = ( + f"{task_key}:{run_ctx.uid}:{run_ctx.group_index}:{run_ctx.sample_index}" + ) + lease = await env_client.allocate(task_key=task_key, request_id=request_id) + lease_id = str(lease["lease_id"]) + logger.info( + "Using remote terminal env backend lease=%s server=%s", lease_id, env_server_url + ) + return env_client, lease_id + + +def _create_sglang_client( + args: Any, + tokenizer: Any, + sampling_params: Dict[str, Any], + max_total_tokens: int, + enable_sglang_non_think: bool, + *, + sglang_url: str | None = None, + max_retries: int = 30, +) -> SGLangTurnClient: + if not sglang_url: + sglang_url = ( + f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate" + ) + client_template_kwargs = { + "chat_template_type": getattr(args, "chat_template_type", "hf"), + "chat_template_kwargs": getattr(args, "chat_template_kwargs", None), + "messages_delimiter_start": getattr( + args, "messages_delimiter_start", "<|im_start|>" + ), + "messages_delimiter_end": getattr(args, "messages_delimiter_end", "<|im_end|>"), + "tool_call_parser": getattr(args, "tool_call_parser", "qwen25"), + } + if enable_sglang_non_think: + raw_chat_template_kwargs = client_template_kwargs.get("chat_template_kwargs") + if isinstance(raw_chat_template_kwargs, dict): + merged_chat_template_kwargs = dict(raw_chat_template_kwargs) + else: + merged_chat_template_kwargs = {} + merged_chat_template_kwargs["enable_thinking"] = False + client_template_kwargs["chat_template_kwargs"] = merged_chat_template_kwargs + + completion_budget = _infer_completion_budget(sampling_params) + effective_context_limit = max_total_tokens + for maybe_cap in ( + getattr(args, "rollout_max_context_len", None), + getattr(args, "sglang_max_context_len", None), + ): + try: + parsed_cap = int(maybe_cap) + except (TypeError, ValueError): + continue + if parsed_cap > 0: + effective_context_limit = min(effective_context_limit, parsed_cap) + max_input_tokens = max(1, effective_context_limit - completion_budget) + logger.info( + "SGLang client: url=%s context_limit=%d, completion_budget=%d, max_input_tokens=%d", + sglang_url, + effective_context_limit, + completion_budget, + max_input_tokens, + ) + raw_request_timeout = getattr(args, "sglang_request_timeout", None) + if raw_request_timeout in (None, "", 0, 0.0): + raw_request_timeout = os.getenv("SGLANG_REQUEST_TIMEOUT") + try: + request_timeout = ( + float(raw_request_timeout) if raw_request_timeout is not None else None + ) + except (TypeError, ValueError): + request_timeout = None + if request_timeout is not None and request_timeout <= 0: + request_timeout = None + + return SGLangTurnClient( + model_type=None, + tokenizer=tokenizer, + sampling_params=sampling_params, + url=sglang_url, + session_id=None, + max_input_tokens=max_input_tokens, + request_timeout=request_timeout, + max_retries=max_retries, + **client_template_kwargs, + ) + + +async def generate( + args, + sample: Sample, + sampling_params: Dict[str, Any], + evaluation: bool = False, +) -> List[Sample]: + _ = evaluation + state = GenerateState(args) + + task_meta = _extract_task_meta(sample) + uid = (sample.metadata or {}).get("uid") or uuid.uuid4().hex[:8] + group_index = int(sample.group_index) if sample.group_index is not None else -1 + sample_index = int(sample.index) if sample.index is not None else -1 + task_spec = _make_task_spec(task_meta) + run_ctx = RunContext( + uid=uid, + group_index=group_index, + sample_index=sample_index, + log_dir=Path(getattr(args, "tbench_output_root", "build_outputs")) + / "AgentRunner_Output", + ) + run_ctx_payload = run_ctx.to_payload() + timeouts = TaskTimeouts( + ensure_image=getattr(args, "ensure_image_timeout", 300.0), + reset_session=getattr(args, "reset_session_timeout", 300.0), + close_session=getattr(args, "close_session_timeout", 60.0), + eval=getattr(args, "eval_timeout", 600.0), + ) + timeouts_payload = timeouts.to_payload() + + env_client: Optional[TerminalEnvClient] = None + lease_id: Optional[str] = None + + prm_enable = bool(getattr(args, "prm_enable", False)) and (not evaluation) + prm_coef = float(getattr(args, "prm_turn_coef", 1.0)) + prm_agent: TerminalPRMAgent | None = None + prm_pending: list[tuple[int, asyncio.Task]] = [] + prm_turn_scores: dict[int, float] = {} + prm_turn_details: list[dict[str, Any]] = [] + + _log_tag = f"[task={task_spec.task_name} uid={run_ctx.uid} group_idx={run_ctx.group_index} sample_idx={run_ctx.sample_index}]" + + try: + env_client, lease_id = await _create_env_client(task_spec, run_ctx) + reset_payload = await env_client.reset( + lease_id=lease_id, + task_meta=task_meta, + run_ctx=run_ctx_payload, + task_timeouts=timeouts_payload, + ) + user_msg = str(reset_payload.get("user_msg", "")) + raw_tools = list(reset_payload.get("tool_schemas", [])) + logger.info("%s Start terminal rollout", _log_tag) + + tool_schemas = _normalize_tool_schemas(raw_tools) + agent_type = str(getattr(args, "terminal_agent_type", "camel_agent")) + model_type = str(getattr(args, "model_type", "slime-sglang")) + non_think_mode = bool(getattr(args, "non_think_mode", True)) + non_think_mode_source = str( + getattr(args, "non_think_mode_source", "prompt") + ).lower() + if non_think_mode_source not in {"prompt", "sglang", "both"}: + non_think_mode_source = "prompt" + enable_prompt_non_think = non_think_mode and non_think_mode_source in { + "prompt", + "both", + } + enable_sglang_non_think = non_think_mode and non_think_mode_source in { + "sglang", + "both", + } + + terminal_max_iterations = max(1, int(getattr(args, "max_iteration", 10))) + terminal_max_parse_errors = max(1, int(getattr(args, "max_parse_errors", 3))) + max_total_tokens = int(getattr(args, "max_total_tokens", 32768)) + sglang_client = _create_sglang_client( + args=args, + tokenizer=state.tokenizer, + sampling_params=sampling_params, + max_total_tokens=max_total_tokens, + enable_sglang_non_think=enable_sglang_non_think, + ) + + if prm_enable: + prm_router_ip = getattr(args, "prm_router_ip", None) + prm_router_port = getattr(args, "prm_router_port", None) + if prm_router_ip and prm_router_port: + prm_sglang_url = f"http://{prm_router_ip}:{prm_router_port}/generate" + else: + prm_sglang_url = getattr(args, "prm_sglang_url", None) or os.getenv( + "PRM_SGLANG_URL", "" + ) + if not prm_sglang_url: + raise RuntimeError( + "prm_enable=True but no PRM endpoint: set prm_router_ip/port, " + "prm_sglang_url, or PRM_SGLANG_URL env var." + ) + prm_sampling_params = { + "temperature": float(getattr(args, "prm_temperature", 0.0)), + "max_new_tokens": int(getattr(args, "prm_max_new_tokens", 4096)), + } + prm_max_total_tokens = int(getattr(args, "prm_max_total_tokens", 16384)) + prm_sglang_client = _create_sglang_client( + args=args, + tokenizer=state.tokenizer, + sampling_params=prm_sampling_params, + max_total_tokens=prm_max_total_tokens, + enable_sglang_non_think=True, + sglang_url=prm_sglang_url, + max_retries=10, + ) + prm_agent = TerminalPRMAgent( + sglang_client=prm_sglang_client, + task_instruction=task_spec.instruction, + history_mode=str(getattr(args, "prm_history_mode", "head_tail")), + ) + logger.info( + "%s PRM enabled: url=%s coef=%.3f", _log_tag, prm_sglang_url, prm_coef + ) + + agent_runner = create_agent_runner( + agent_type=agent_type, + sglang_client=sglang_client, + model_type=model_type, + tool_schemas=tool_schemas, + non_think_mode=enable_prompt_non_think, + max_total_tokens=max_total_tokens, + ) + agent_runner.reset(user_msg) + agent_runner.set_max_parse_errors(terminal_max_parse_errors) + agent_runner.set_max_iterations(terminal_max_iterations) + + # Loop + interactions: List[Interaction] = [] + final_model_response = None + final_response = None + reached_iteration_limit = False + reached_parse_error_limit = False + + while True: + context_result: TurnContext = await agent_runner.get_turn_context() + if context_result.terminated_response is not None: + logger.warning("%s Rollout pre-terminated before model turn.", _log_tag) + final_response = context_result.terminated_response + break + if context_result.context_messages is None: + logger.warning("%s Rollout context is empty; aborting loop.", _log_tag) + break + + turn_state: TurnResult = await agent_runner.run_model_turn( + context_result.context_messages + ) + interaction = turn_state.interaction + turn_idx = int(interaction.turn_idx) + interactions.append(interaction) + + if prm_agent is not None: + tool_calls_for_prm = [ + {"tool_name": tc.tool_name, "args": tc.args} + for tc in (turn_state.tool_call_requests or []) + ] + prm_agent.record_model_turn( + turn_idx, + assistant_text=interaction.output_text or "", + tool_calls=tool_calls_for_prm or None, + parse_error_recorded=turn_state.parse_error_recorded, + finish_reason=interaction.finish_reason, + ) + + if turn_state.terminated_response is not None: + logger.warning( + "%s Rollout terminated during model turn %d.", _log_tag, turn_idx + ) + final_response = turn_state.terminated_response + break + if turn_state.model_response is None: + logger.warning( + "%s Model turn %d returned empty model_response.", + _log_tag, + turn_idx, + ) + break + + should_continue_loop = False + if tool_call_requests := turn_state.tool_call_requests: + logger.info( + "%s Turn %d: executing %d tool call(s).", + _log_tag, + turn_idx, + len(tool_call_requests), + ) + for tool_call_request in tool_call_requests: + assert env_client is not None and lease_id is not None + await env_client.heartbeat(lease_id) + raw_result = await env_client.exec_tool( + lease_id, + tool_call_request.tool_name, + tool_call_request.args, + ) + agent_runner.record_tool_result(tool_call_request, raw_result) + if prm_agent is not None: + prm_agent.record_tool_result( + turn_idx, tool_call_request, raw_result + ) + should_continue_loop = True + + if turn_state.parse_error_recorded: + logger.warning( + "%s Turn %d: tool-call parse error.", + _log_tag, + turn_idx, + ) + should_continue_loop = True + + if prm_agent is not None: + task = asyncio.create_task(prm_agent.judge_turn(turn_idx)) + prm_pending.append((turn_idx, task)) + + if should_continue_loop: + if ( + turn_state.parse_error_recorded + and agent_runner.reached_parse_error_limit() + ): + logger.error( + "%s Max parse errors (%d) reached at turn %d.", + _log_tag, + agent_runner.max_parse_errors, + turn_idx, + ) + reached_parse_error_limit = True + final_model_response = turn_state.model_response + break + if agent_runner.reached_iteration_limit(): + logger.warning( + "%s Max iterations (%d) reached.", + _log_tag, + agent_runner.max_iterations, + ) + reached_iteration_limit = True + final_model_response = turn_state.model_response + break + continue + + final_model_response = turn_state.model_response + break + + if final_response is None and final_model_response is not None: + final_response = agent_runner.finalize_response(final_model_response) + + if final_response is None: + logger.error( + "%s No final response produced; mark sample aborted.", _log_tag + ) + sample.status = Sample.Status.ABORTED + sample.remove_sample = True + sample.reward = {"score": 0.0} + return [sample] + + finish_reasons = final_response.info.get("termination_reasons", []) + is_aborted = not final_response.msg + + if final_response.terminated and "max_tokens_exceeded" in finish_reasons: + status = Sample.Status.TRUNCATED + elif reached_iteration_limit: + status = Sample.Status.TRUNCATED + elif reached_parse_error_limit: + status = Sample.Status.FAILED + elif is_aborted: + status = Sample.Status.ABORTED + else: + status = Sample.Status.COMPLETED + logger.info( + "%s Rollout finished: status=%s turns=%d parse_errors=%d", + _log_tag, + status, + agent_runner.model_turn_count, + agent_runner.parse_error_count, + ) + + # Evaluation + reward = 0.0 + eval_error: str | None = None + should_evaluate = (not is_aborted) and status != Sample.Status.FAILED + if should_evaluate: + try: + assert env_client is not None and lease_id is not None + await env_client.heartbeat(lease_id) + raw_score = await env_client.evaluate(lease_id) + reward = float(raw_score) + logger.info("%s Evaluation reward=%.4f", _log_tag, reward) + except Exception as exc: + eval_error = f"{type(exc).__name__}: {exc}" + status = Sample.Status.FAILED + reward = 0.0 + logger.error( + "%s Evaluation failed, marking FAILED: %s", + _log_tag, + eval_error, + ) + + if not interactions: + logger.warning("%s No interactions recorded; remove sample.", _log_tag) + sample.status = status + sample.remove_sample = True + sample.reward = {"score": 0.0} + return [sample] + + if prm_agent is not None and prm_pending: + for turn_idx, prm_task in prm_pending: + try: + output_text, score = await prm_task + prm_turn_scores[turn_idx] = float(score) + prm_turn_details.append( + { + "turn_idx": turn_idx, + "score": float(score), + "output_text": output_text, + } + ) + logger.info( + "%s PRM judge turn %d score=%.4f, output_text=%s", + _log_tag, + turn_idx, + float(score), + output_text.replace("\n", ""), + ) + except Exception as exc: + logger.warning( + "%s PRM judge failed for turn %d (ignored): %s", + _log_tag, + turn_idx, + exc, + ) + prm_turn_scores[turn_idx] = 0.0 + prm_turn_details.append( + {"turn_idx": turn_idx, "score": 0.0, "error": str(exc)} + ) + + if prm_agent is not None: + sample.metadata["prm"] = { + "enabled": True, + "coef": prm_coef, + "turn_scores": prm_turn_scores, + "turn_details": prm_turn_details, + } + + # Build training samples + samples = _build_samples( + interactions=interactions, + base_sample=sample, + outcome=reward, + status=status, + prm_turn_scores=(prm_turn_scores if prm_agent is not None else None), + prm_coef=prm_coef, + discount=1.0, + encourage=False, + ) + for s in samples: + s.metadata["model_turn_count"] = agent_runner.model_turn_count + s.metadata["parse_error_count"] = agent_runner.parse_error_count + if eval_error is not None: + s.metadata["evaluation_failed"] = True + s.metadata["evaluation_error"] = eval_error + _mark_non_trainable_samples(samples) + return samples + + except Exception as exc: + logger.error( + "%s Generate failed (%s): %s", + _log_tag, + type(exc).__name__, + exc, + exc_info=True, + ) + sample.status = Sample.Status.FAILED + sample.remove_sample = True + sample.reward = {"score": 0.0} + + eos = state.tokenizer.eos_token_id + if eos is None: + sample.tokens = [] + sample.response_length = 0 + sample.rollout_log_probs = [] + sample.loss_mask = [] + else: + sample.tokens = [eos, eos] + sample.response_length = 1 + sample.rollout_log_probs = [0.0] + sample.loss_mask = [0] + return [sample] + + finally: + for _turn_idx, t in prm_pending: + if not t.done(): + t.cancel() + + if env_client is not None and lease_id is not None: + try: + await env_client.close(lease_id) + except Exception as exc: + logger.debug( + "%s Best-effort remote close failed lease=%s: %s", + _log_tag, + lease_id, + exc, + ) diff --git a/terminal-rl/inference_client.py b/terminal-rl/inference_client.py new file mode 100644 index 0000000..b00cb6b --- /dev/null +++ b/terminal-rl/inference_client.py @@ -0,0 +1,333 @@ +from __future__ import annotations + +import asyncio +import datetime +import inspect +import logging +import time +import traceback +import uuid +from copy import deepcopy +from typing import Any, Dict, List + +from openai.types.chat import ChatCompletion, ChatCompletionMessage +from openai.types.chat.chat_completion import Choice +from openai.types.chat.chat_completion_message_function_tool_call import ( + ChatCompletionMessageFunctionToolCall, + Function, +) +from openai.types.completion_usage import CompletionUsage +from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall + +from slime.utils.http_utils import post as async_post + +from custom_types import Interaction + +logger = logging.getLogger(__name__) + + +def process_tool_calls( + text: str, + tools: list[Any], + tool_call_parser: str | None, + finish_reason: str, + use_responses: bool = False, +) -> tuple[ + list[ChatCompletionMessageFunctionToolCall | ResponseFunctionToolCall] | None, + str, + str, +]: + from sglang.srt.entrypoints.openai.protocol import Function as SglFunction + from sglang.srt.entrypoints.openai.protocol import Tool as SglTool + from sglang.srt.function_call.function_call_parser import FunctionCallParser + + if use_responses: + tools = [ + SglTool( + type=tool["type"], + function=SglFunction( + name=tool.get("name"), + description=tool.get("description"), + parameters=tool.get("parameters"), + ), + ) + for tool in tools + ] + else: + tools = [ + SglTool(type=tool["type"], function=SglFunction(**tool["function"])) + for tool in tools + ] + + parser = FunctionCallParser(tools, tool_call_parser) + if parser.has_tool_call(text): + if finish_reason == "stop": + finish_reason = "tool_calls" + try: + text, call_info_list = parser.parse_non_stream(text) + + if use_responses: + tool_calls = [ + ResponseFunctionToolCall( + type="function_call", + id=f"fc-{uuid.uuid4().hex[:24]}", + call_id=f"call_{uuid.uuid4().hex[:24]}", + name=call_info.name, + arguments=call_info.parameters, + status="completed", + ) + for call_info in call_info_list + ] + else: + tool_calls = [ + ChatCompletionMessageFunctionToolCall( + type="function", + id=f"call_{uuid.uuid4().hex[:24]}", + function=Function( + name=call_info.name, arguments=call_info.parameters + ), + ) + for call_info in call_info_list + ] + return tool_calls, text, finish_reason + except Exception as exc: + logger.error("Tool call parsing error: %s", exc) + traceback.print_exc() + return None, text, finish_reason + + return None, text, finish_reason + + +def _ensure_stop_token_ids( + tokenizer, sampling_params: Dict[str, Any] +) -> Dict[str, Any]: + normalized = dict(sampling_params) + if "stop_token_ids" in normalized: + return normalized + + stop_ids: set[int] = set() + if tokenizer.eos_token_id is not None: + stop_ids.add(tokenizer.eos_token_id) + if tokenizer.pad_token_id is not None: + stop_ids.add(tokenizer.pad_token_id) + if stop_ids: + normalized["stop_token_ids"] = list(stop_ids) + return normalized + + +def _to_positive_int(value: Any) -> int | None: + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed > 0 else None + + +class SGLangTurnClient: + def __init__( + self, + *, + model_type: str | None = None, + tokenizer, + sampling_params: Dict[str, Any], + url: str, + chat_template_type: str = "hf", + chat_template_kwargs: Dict[str, Any] | None = None, + messages_delimiter_start: str = "<|im_start|>", + messages_delimiter_end: str = "<|im_end|>", + session_id: str | None = None, + tool_call_parser: str | None = None, + max_input_tokens: int | None = None, + request_timeout: float | None = None, + max_retries: int = 30, + ) -> None: + self.model_type = model_type + self.tokenizer = tokenizer + self.sampling_params = _ensure_stop_token_ids(tokenizer, sampling_params) + self.url = url + self.chat_template_type = chat_template_type + self.chat_template_kwargs = chat_template_kwargs or {} + self.messages_delimiter_start = messages_delimiter_start + self.messages_delimiter_end = messages_delimiter_end + self.session_id = session_id + self.tool_call_parser = tool_call_parser + self.max_input_tokens = _to_positive_int(max_input_tokens) + self.request_timeout = ( + request_timeout if (request_timeout and request_timeout > 0) else None + ) + self.max_retries = max(1, max_retries) + + self.marker = "\n[OMITTED MIDDLE]\n" + self.sep_ids = self.tokenizer.encode(self.marker, add_special_tokens=False) + + def _truncate_input_ids(self, input_ids: List[int]) -> List[int]: + max_toks = self.max_input_tokens + if max_toks is None or len(input_ids) <= max_toks: + return input_ids + + dropped = len(input_ids) - max_toks + logger.warning( + "Prompt is too long for configured budget: input=%d, budget=%d. Truncating %d token(s) from the left and right.", + len(input_ids), + max_toks, + dropped, + ) + + keep_head_ratio = getattr(self, "keep_head_ratio", 0.3) + head = max(1, int(max_toks * keep_head_ratio)) + tail = max_toks - head - len(self.sep_ids) + if tail <= 0: + logger.warning( + f"tail is not positive: tail={tail}, head={head}, max_toks={max_toks}, len(self.sep_ids)={len(self.sep_ids)}" + ) + tail = 1 + + return input_ids[:head] + self.sep_ids + input_ids[-tail:] + + async def generate_turn( + self, + *, + messages: List[dict[str, Any]], + tools: List[dict[str, Any]] | None, + turn_idx: int, + ) -> tuple[ChatCompletion, Interaction]: + input_ids = self._apply_chat_template(messages, tools) + input_ids = self._truncate_input_ids(input_ids) + payload: Dict[str, Any] = { + "input_ids": input_ids, + "sampling_params": self.sampling_params, + "return_logprob": True, + } + headers: Dict[str, str] | None = None + if self.session_id: + headers = {"X-SMG-Routing-Key": self.session_id} + + t0 = time.monotonic() + supports_headers = "headers" in inspect.signature(async_post).parameters + + async def _do_post(): + if headers and supports_headers: + return await async_post( + self.url, payload, max_retries=self.max_retries, headers=headers + ) + else: + if headers and not supports_headers: + logger.warning( + "async_post() does not accept headers; routing key will be ignored for this request." + ) + return await async_post(self.url, payload, max_retries=self.max_retries) + + if self.request_timeout: + try: + output = await asyncio.wait_for( + _do_post(), timeout=self.request_timeout + ) + except asyncio.TimeoutError: + elapsed = (time.monotonic() - t0) * 1000.0 + raise TimeoutError( + f"SGLang generate request timed out after {self.request_timeout}s " + f"(elapsed={elapsed:.0f}ms, turn_idx={turn_idx})" + ) + else: + output = await _do_post() + latency_ms = (time.monotonic() - t0) * 1000.0 + + output_text: str = output["text"] + raw_output_text = output_text + meta_info = output["meta_info"] + finish_reason: str = meta_info["finish_reason"]["type"] + + if "output_token_logprobs" in meta_info: + raw_logprobs = meta_info["output_token_logprobs"] + if raw_logprobs and logger.isEnabledFor(logging.DEBUG): + logger.debug( + "output_token_logprobs sample element: %s", raw_logprobs[0] + ) + output_token_ids: list[int] = [x[1] for x in raw_logprobs] + output_token_logprobs: list[float] = [x[0] for x in raw_logprobs] + else: + output_token_ids = [] + output_token_logprobs = [] + + tool_calls = None + if tools: + tool_calls, output_text, finish_reason = process_tool_calls( + output_text, + tools, + self.tool_call_parser, + finish_reason, + ) + + completion_id = f"chatcmpl-{uuid.uuid4().hex[:29]}" + current_time = int(datetime.datetime.now().timestamp()) + chat_completion = ChatCompletion( + id=completion_id, + choices=[ + Choice( + finish_reason=finish_reason, + index=0, + logprobs=None, + message=ChatCompletionMessage( + content=output_text, + role="assistant", + tool_calls=tool_calls, + ), + ) + ], + created=current_time, + model=self.model_type or "unknown", + object="chat.completion", + service_tier=None, + system_fingerprint=None, + usage=CompletionUsage( + prompt_tokens=len(input_ids), + completion_tokens=len(output_token_ids), + total_tokens=len(input_ids) + len(output_token_ids), + ), + ) + + interaction = Interaction( + turn_idx=turn_idx, + completion=deepcopy(chat_completion), + input_ids=list(input_ids), + output_token_ids=output_token_ids, + output_token_logprobs=output_token_logprobs, + output_text=raw_output_text, + finish_reason=finish_reason, + messages=deepcopy(messages), + latency_ms=latency_ms, + ) + return chat_completion, interaction + + def _apply_chat_template( + self, + messages: List[dict[str, Any]], + tools: List[dict[str, Any]] | None, + ) -> List[int]: + if self.chat_template_type == "hf": + try: + return self.tokenizer.apply_chat_template( + messages, + tools=tools or None, + add_generation_prompt=True, + tokenize=True, + **self.chat_template_kwargs, + ) + except Exception: + return self.tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=True, + **self.chat_template_kwargs, + ) + + if self.chat_template_type == "concat": + start = self.messages_delimiter_start + end = self.messages_delimiter_end + message_strs: List[str] = [] + for msg in messages: + message_strs.append(f"{start}{msg['role']}\n{msg['content']}{end}\n") + message_strs.append(f"{start}assistant\n") + return self.tokenizer.encode("".join(message_strs)) + + raise ValueError(f"Unsupported chat_template_type: {self.chat_template_type!r}") diff --git a/terminal-rl/remote/README.md b/terminal-rl/remote/README.md new file mode 100644 index 0000000..9d14446 --- /dev/null +++ b/terminal-rl/remote/README.md @@ -0,0 +1,127 @@ +# Remote worker (pool server) + +This directory runs on the **remote worker**: a pool server that manages Docker containers and executes terminal tasks. + +## Prerequisites + +Set up a machine that will act as a worker node: a cloud VM (e.g. AWS EC2, GCP, or any provider), a bare-metal server, or any host where: + +- You can install **Docker** (and Docker Compose). +- You have network connectivity so that the **training cluster** (where the router runs) can reach this host on the port you use for the pool server (default **18081**). + +### GPU (optional) + +A GPU is not required to run the pool server, but may be required by some tasks. + +--- + +## Instructions + +### 1. Clone the repo + +From a directory of your choice: + +```bash +git clone https://github.com/Gen-Verse/OpenClaw-RL.git +cd OpenClaw-RL +``` + +### 2. Install dependencies + +Install Docker, a Python 3.12 environment, and the Python packages required by the pool server. You can use the provided script (run from **repo root**): + +```bash +bash terminal-rl/remote/setup.sh +``` + +This will: + +- Install Docker and Docker Compose if missing. +- Install [uv](https://github.com/astral-sh/uv) and create a virtualenv at repo root (`.venv`). +- Install other required packages. + +### 3. Download dataset + +To download a dataset: + +```bash +source .venv/bin/activate +export DATASET_DIR="terminal-rl/dataset" +python terminal-rl/data_utils/download.py seta_env +``` + +The `seta_env` dataset corresponds to the task dataset published in: [camel-ai/seta-env](https://github.com/camel-ai/seta-env/tree/main/Dataset). + +### 4. Run the pool server + +From the **repo root**: + +```bash +bash terminal-rl/remote/run_pool_server.sh +``` + +This script: + +- Activates the venv if `.venv` exists. +- Sets `DATASET_DIR` and `TBENCH_OUTPUT_ROOT` under `terminal-rl/` by default. +- Starts the pool server with `python -m terminal-rl.remote.pool_server` on `0.0.0.0:18081` (overridable via `ENV_SERVER_PORT`, `WORKER_MAX_TASKS`, `WORKER_MAX_RUNS_PER_TASK`). + +Run in background / under a process manager as needed. Example (nohup`): + +```bash +nohup bash terminal-rl/remote/run_pool_server.sh > pool_server.log 2>&1 & +``` + +### 5. Tell the training side the worker URL + +On the training machine (router host), set `WORKER_URLS` to include this worker: + +```bash +export WORKER_URLS="http://:18081" +``` + +For multiple workers, use a comma-separated list: + +```bash +export WORKER_URLS="http://worker1:18081,http://worker2:18081" +``` + +Then start the router and training as described in the main Terminal RL docs; the router forwards requests to these pool servers. + +--- + +## Optional environment variables + +When running the pool server (via `run_pool_server.sh` or `python -m terminal-rl.remote.pool_server`), the following variables are supported: + + +| Variable | Default | Description | +| ---------------------------- | --------------------------- | ----------------------------------------------------------------------------------------------------------- | +| `DATASET_DIR` | `terminal-rl/dataset` | Path to the task dataset directory. | +| `TBENCH_OUTPUT_ROOT` | `terminal-rl/build_outputs` | Root directory for build/output artifacts. | +| `ENV_SERVER_PORT` | `18081` | Port the pool server listens on. | +| `WORKER_MAX_TASKS` | `16` | Max tasks allocate to per worker. | +| `WORKER_MAX_RUNS_PER_TASK` | `8` | Max concurrent runs per task. | +| `TBENCH_DOCKER_IMAGE_SOURCE` | `build` | `build` or `pull` — build images locally or pull from a registry. | +| `TBENCH_DOCKER_PULL_PREFIX` | — | Image name prefix used in `pull` mode; the task name is appended (e.g., `task-1374` → `task-1374`). | +| `COMPOSE_OVERRIDE_PATH` | — | Optional Docker Compose override file. | + + +Example with custom port and limits: + +```bash +export ENV_SERVER_PORT=18082 +export WORKER_MAX_TASKS=10 +export WORKER_MAX_RUNS_PER_TASK=8 +bash terminal-rl/remote/run_pool_server.sh +``` + +Example using pre-built images from a registry (pull mode). Set the image source and prefix; you can build and push your own: + +```bash +export TBENCH_DOCKER_IMAGE_SOURCE=pull +export TBENCH_DOCKER_PULL_PREFIX="ghcr.io//:task-" +export COMPOSE_OVERRIDE_PATH="terminal-rl/remote/compose_override.yaml" +bash terminal-rl/remote/run_pool_server.sh +``` + diff --git a/terminal-rl/remote/__init__.py b/terminal-rl/remote/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/terminal-rl/remote/compose_override.yaml b/terminal-rl/remote/compose_override.yaml new file mode 100644 index 0000000..46951be --- /dev/null +++ b/terminal-rl/remote/compose_override.yaml @@ -0,0 +1,9 @@ +services: + client: + environment: + - http_proxy=${http_proxy} + - https_proxy=${https_proxy} + - no_proxy=${no_proxy} + - HTTP_PROXY=${http_proxy} + - HTTPS_PROXY=${https_proxy} + - NO_PROXY=${no_proxy} diff --git a/terminal-rl/remote/docker_compose_utils.py b/terminal-rl/remote/docker_compose_utils.py new file mode 100644 index 0000000..29a8778 --- /dev/null +++ b/terminal-rl/remote/docker_compose_utils.py @@ -0,0 +1,374 @@ +from __future__ import annotations + +import logging +import os +import shutil +import subprocess +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +from terminal_bench.handlers.trial_handler import TrialHandler +from terminal_bench.terminal.docker_compose_manager import DockerComposeManager +from terminal_bench.terminal.terminal import Terminal + + +@dataclass(frozen=True) +class ImagePreparationResult: + mode: Literal["build", "pull"] + client_image_name: str | None = None + + +def _shorten_output(text: str | None, max_chars: int = 4000) -> str: + if not text: + return "" + stripped = text.strip() + if len(stripped) <= max_chars: + return stripped + return f"{stripped[:max_chars]}...(truncated, total={len(stripped)} chars)" + + +def _build_docker_pull_error_message( + *, + image: str, + cmd: list[str], + return_code: int | None = None, + timeout: float | None = None, + stdout: str | None = None, + stderr: str | None = None, +) -> str: + lines = [ + f"Docker image pull failed for image '{image}'.", + f"Command: {' '.join(cmd)}", + ] + if return_code is not None: + lines.append(f"Exit code: {return_code}") + if timeout is not None: + lines.append(f"Timeout: {timeout:.1f}s") + + out = _shorten_output(stdout) + err = _shorten_output(stderr) + if out: + lines.append(f"STDOUT:\n{out}") + if err: + lines.append(f"STDERR:\n{err}") + + lines.append( + "Hints: verify task_name/task image tag, run docker login for the registry, and ensure image exists." + ) + return "\n".join(lines) + + +def _build_compose_up_error_message( + *, + cmd: list[str], + return_code: int | None = None, + timeout: float | None = None, + stdout: str | None = None, + stderr: str | None = None, + note: str | None = None, +) -> str: + lines = [ + "Docker compose up --no-build failed.", + f"Command: {' '.join(cmd)}", + ] + if return_code is not None: + lines.append(f"Exit code: {return_code}") + if timeout is not None: + lines.append(f"Timeout: {timeout:.1f}s") + + out = _shorten_output(stdout) + err = _shorten_output(stderr) + if out: + lines.append(f"STDOUT:\n{out}") + if err: + lines.append(f"STDERR:\n{err}") + if note: + lines.append(f"Note: {note}") + + lines.append( + "Hints: verify `docker compose version`; if unavailable, install Compose plugin or ensure `docker-compose` is on PATH." + ) + return "\n".join(lines) + + +def _compose_plugin_maybe_missing(stderr: str | None) -> bool: + if not stderr: + return False + lowered = stderr.lower() + return ( + "unknown shorthand flag: 'p' in -p" in lowered + or "docker: 'compose' is not a docker command" in lowered + or 'unknown command "compose"' in lowered + ) + + +def build_docker_image(task: dict[str, Any], timeout: float = 1200.0) -> None: + dataset_dir = str(os.getenv("DATASET_DIR", "")).strip() + if not dataset_dir: + raise ValueError("DATASET_DIR is required") + + task_path = Path(dataset_dir) / str(task.get("task_path", "")) + trial_handler = TrialHandler( + trial_name="build_run", + input_path=task_path, + output_path=Path("build_outputs"), + ) + + compose_manager = DockerComposeManager( + client_container_name=trial_handler.client_container_name, + client_image_name=trial_handler.client_image_name, + docker_image_name_prefix=trial_handler.docker_image_name_prefix, + docker_compose_path=trial_handler.task_paths.docker_compose_path, + no_rebuild=True, + cleanup=False, + sessions_logs_path=trial_handler.trial_paths.sessions_path, + agent_logs_path=trial_handler.trial_paths.agent_logging_dir, + ) + compose_manager.build(timeout=timeout) + + +def _resolve_pull_image(task: dict[str, Any]) -> str: + prefix = str(os.getenv("TBENCH_DOCKER_PULL_PREFIX", "")).strip() + if not prefix: + raise ValueError("TBENCH_DOCKER_PULL_PREFIX is required in pull mode") + + task_name = str(task.get("task_name", "")).strip() + if not task_name: + raise ValueError("task_name is required to resolve pull image") + if "<" in task_name and ">" in task_name: + raise ValueError( + "task_name appears to still be a placeholder " + f"('{task_name}'). Please provide a concrete task_name." + ) + + return f"{prefix}{task_name}" + + +def _docker_image_exists_locally(image: str, timeout: float = 30.0) -> bool: + cmd = ["docker", "image", "inspect", image] + result = subprocess.run( + cmd, + check=False, + capture_output=True, + text=True, + timeout=timeout, + ) + return result.returncode == 0 + + +def pull_docker_image(image: str, timeout: float = 1200.0) -> None: + if _docker_image_exists_locally(image): + return + + cmd = ["docker", "pull", image] + try: + subprocess.run( + cmd, + check=True, + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as exc: + raise RuntimeError( + _build_docker_pull_error_message( + image=image, + cmd=cmd, + timeout=timeout, + stdout=exc.stdout, + stderr=exc.stderr, + ) + ) from exc + except subprocess.CalledProcessError as exc: + raise RuntimeError( + _build_docker_pull_error_message( + image=image, + cmd=cmd, + return_code=exc.returncode, + stdout=exc.stdout, + stderr=exc.stderr, + ) + ) from exc + + +def prepare_task_docker_image( + task: dict[str, Any], + timeout: float = 1200.0, +) -> ImagePreparationResult: + raw_mode = str(os.getenv("TBENCH_DOCKER_IMAGE_SOURCE", "")).strip().lower() + if not raw_mode: + raise ValueError("TBENCH_DOCKER_IMAGE_SOURCE is required") + + if raw_mode in {"build", "docker_build"}: + build_docker_image(task=task, timeout=timeout) + return ImagePreparationResult(mode="build", client_image_name=None) + + if raw_mode in {"pull", "docker_pull"}: + image = _resolve_pull_image(task=task) + pull_docker_image(image=image, timeout=timeout) + return ImagePreparationResult(mode="pull", client_image_name=image) + + raise ValueError( + f"Unsupported docker image source '{raw_mode}'. Expected one of: build, pull" + ) + + +_DEFAULT_CONTAINER_MEMORY_LIMIT = os.getenv("CONTAINER_MEMORY_LIMIT", "16g") + + +def _apply_container_memory_limit( + container_name: str, + memory_limit: str, + logger: logging.Logger | None = None, +) -> None: + """Best-effort ``docker update --memory`` on a running container.""" + if not memory_limit: + return + cmd = [ + "docker", + "update", + f"--memory={memory_limit}", + f"--memory-swap={memory_limit}", + container_name, + ] + try: + subprocess.run(cmd, check=True, capture_output=True, text=True, timeout=30.0) + if logger is not None: + logger.info( + "Applied memory limit %s to container %s", memory_limit, container_name + ) + except Exception as exc: + if logger is not None: + logger.warning( + "Failed to apply memory limit to container %s: %s", container_name, exc + ) + + +def compose_up_no_build( + terminal: Terminal, + *, + timeout: float, + container_name: str, + logger: logging.Logger | None = None, +) -> None: + compose_manager = getattr(terminal, "_compose_manager") + compose_override_path = str(os.getenv("COMPOSE_OVERRIDE_PATH", "")).strip() + compose_command = ["up", "-d", "--no-build"] + if compose_override_path: + compose_command = ["-f", compose_override_path, *compose_command] + + command = compose_manager.get_docker_compose_command(compose_command) + if logger is not None: + logger.info("Running docker compose command: %s", " ".join(command)) + if compose_override_path: + logger.info("Using compose override file: %s", compose_override_path) + + compose_manager.env["http_proxy"] = os.getenv("HTTP_PROXY", "") + compose_manager.env["https_proxy"] = os.getenv("HTTPS_PROXY", "") + + try: + subprocess.run( + command, + env=compose_manager.env, + check=True, + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as exc: + if logger is not None: + logger.error( + "Docker compose up --no-build timed out after %.1f sec", timeout + ) + if exc.stdout: + logger.error("STDOUT: %s", exc.stdout) + if exc.stderr: + logger.error("STDERR: %s", exc.stderr) + raise RuntimeError( + _build_compose_up_error_message( + cmd=command, + timeout=timeout, + stdout=exc.stdout, + stderr=exc.stderr, + ) + ) from exc + except subprocess.CalledProcessError as exc: + if logger is not None: + logger.error( + "Docker compose up --no-build failed with code %s", exc.returncode + ) + if exc.stdout: + logger.error("STDOUT: %s", exc.stdout) + if exc.stderr: + logger.error("STDERR: %s", exc.stderr) + if _compose_plugin_maybe_missing(exc.stderr): + fallback_command = ["docker-compose", *command[2:]] + if shutil.which("docker-compose"): + if logger is not None: + logger.warning( + "docker compose plugin may be unavailable; falling back to: %s", + " ".join(fallback_command), + ) + try: + subprocess.run( + fallback_command, + env=compose_manager.env, + check=True, + capture_output=True, + text=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as fallback_exc: + raise RuntimeError( + _build_compose_up_error_message( + cmd=fallback_command, + timeout=timeout, + stdout=fallback_exc.stdout, + stderr=fallback_exc.stderr, + note=f"Initial command {' '.join(command)} failed with exit code {exc.returncode}.", + ) + ) from fallback_exc + except subprocess.CalledProcessError as fallback_exc: + raise RuntimeError( + _build_compose_up_error_message( + cmd=fallback_command, + return_code=fallback_exc.returncode, + stdout=fallback_exc.stdout, + stderr=fallback_exc.stderr, + note=f"Initial command {' '.join(command)} failed with exit code {exc.returncode}.", + ) + ) from fallback_exc + container = compose_manager._client.containers.get(container_name) + terminal.container = container + compose_manager._client_container = container + _apply_container_memory_limit( + container_name, _DEFAULT_CONTAINER_MEMORY_LIMIT, logger=logger + ) + return + + raise RuntimeError( + _build_compose_up_error_message( + cmd=command, + return_code=exc.returncode, + stdout=exc.stdout, + stderr=exc.stderr, + note="Compose plugin appears unavailable and `docker-compose` binary was not found.", + ) + ) from exc + + raise RuntimeError( + _build_compose_up_error_message( + cmd=command, + return_code=exc.returncode, + stdout=exc.stdout, + stderr=exc.stderr, + ) + ) from exc + + container = compose_manager._client.containers.get(container_name) + terminal.container = container + compose_manager._client_container = container + _apply_container_memory_limit( + container_name, _DEFAULT_CONTAINER_MEMORY_LIMIT, logger=logger + ) diff --git a/terminal-rl/remote/pool_server.py b/terminal-rl/remote/pool_server.py new file mode 100644 index 0000000..b7801a7 --- /dev/null +++ b/terminal-rl/remote/pool_server.py @@ -0,0 +1,679 @@ +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import time +import uuid +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from ..custom_types import RunContext, TaskSpec, TaskTimeouts +from ..request_utils import json_payload +from .terminal_env import TerminalEnv + +logger = logging.getLogger("terminal.env.worker") +app = FastAPI() + + +def _parse_timeout_overrides( + base: TaskTimeouts, payload: dict[str, Any] | None +) -> TaskTimeouts: + if not isinstance(payload, dict): + return base + + def _pick(key: str, default: float) -> float: + raw = payload.get(key, default) + try: + value = float(raw) + except (TypeError, ValueError): + return default + return value if value > 0 else default + + return TaskTimeouts( + ensure_image=_pick("ensure_image", base.ensure_image), + reset_session=_pick("reset_session", base.reset_session), + close_session=_pick("close_session", base.close_session), + eval=_pick("eval", base.eval), + ) + + +def _build_task_spec(task_meta: dict[str, Any]) -> TaskSpec: + return TaskSpec( + task_name=str(task_meta.get("task_name", "unknown")), + task_path=str(task_meta.get("task_path", "")), + instruction=str(task_meta.get("instruction", "")), + ) + + +def _build_run_ctx( + run_ctx_payload: dict[str, Any] | None, default_log_dir: Path +) -> RunContext: + payload = run_ctx_payload if isinstance(run_ctx_payload, dict) else {} + uid = str(payload.get("uid") or uuid.uuid4().hex[:8]) + try: + group_index = int(payload.get("group_index") or 0) + except (TypeError, ValueError): + group_index = 0 + try: + sample_index = int(payload.get("sample_index") or 0) + except (TypeError, ValueError): + sample_index = 0 + + log_dir_raw = payload.get("log_dir") + if isinstance(log_dir_raw, str) and log_dir_raw: + log_dir = Path(log_dir_raw).resolve() + else: + log_dir = default_log_dir.resolve() + + return RunContext( + uid=uid, + group_index=group_index, + sample_index=sample_index, + log_dir=log_dir, + ) + + +class CapacityError(Exception): + def __init__(self, code: str, message: str): + self.code = code + self.message = message + super().__init__(message) + + +@dataclass +class RunSlot: + run_lease_id: str + task_key: str + env: TerminalEnv + last_used_ts: float = field(default_factory=time.time) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + + +@dataclass +class TaskSlot: + task_key: str + runs: dict[str, RunSlot] = field(default_factory=dict) + created_ts: float = field(default_factory=time.time) + last_used_ts: float = field(default_factory=time.time) + + +class WorkerPool: + def __init__( + self, + *, + max_tasks: int, + max_runs_per_task: int, + run_idle_ttl: int, + output_root: str, + default_timeouts: TaskTimeouts, + idempotency_ttl: int = 300, + max_concurrent_closes: int = 8, + ) -> None: + self.max_tasks = max_tasks + self.max_runs_per_task = max_runs_per_task + self.run_idle_ttl = run_idle_ttl + self.output_root = Path(output_root).resolve() + self.output_root.mkdir(parents=True, exist_ok=True) + self.default_timeouts = default_timeouts + self.idempotency_ttl = idempotency_ttl + + self._tasks: dict[str, TaskSlot] = {} + self._run_to_task: dict[str, str] = {} + self._idempotency: dict[tuple[str, str], tuple[str, float]] = {} + self._lock = asyncio.Lock() + + self._close_sem = asyncio.Semaphore(max_concurrent_closes) + self._closing_tasks: set[asyncio.Task] = set() + + def _new_env(self) -> TerminalEnv: + return TerminalEnv() + + async def _close_run_slot( + self, task_key: str, run_lease_id: str, run_slot: RunSlot, *, reason: str + ) -> None: + logger.warning("%s %s (task=%s)", reason, run_lease_id, task_key) + async with self._close_sem: + async with run_slot.lock: + try: + await run_slot.env.close() + except Exception: + logger.exception("Failed to close run session %s", run_lease_id) + + def _schedule_close( + self, task_key: str, run_lease_id: str, run_slot: RunSlot, *, reason: str + ) -> None: + task = asyncio.create_task( + self._close_run_slot(task_key, run_lease_id, run_slot, reason=reason) + ) + self._closing_tasks.add(task) + task.add_done_callback(self._closing_tasks.discard) + + def _reap_idle_locked(self) -> list[tuple[str, str, RunSlot]]: + now = time.time() + expired_slots: list[tuple[str, str, RunSlot]] = [] + + expired_idem = [ + k + for k, (_, ts) in self._idempotency.items() + if now - ts > self.idempotency_ttl + ] + for k in expired_idem: + self._idempotency.pop(k, None) + + for task_key, task_slot in list(self._tasks.items()): + expired_runs: list[str] = [] + for rid, rslot in task_slot.runs.items(): + if now - rslot.last_used_ts > self.run_idle_ttl: + expired_runs.append(rid) + + for rid in expired_runs: + rslot = task_slot.runs.pop(rid, None) + self._run_to_task.pop(rid, None) + if rslot is not None: + expired_slots.append((task_key, rid, rslot)) + + if task_slot.runs: + task_slot.last_used_ts = max( + r.last_used_ts for r in task_slot.runs.values() + ) + else: + logger.info("Reaping empty task slot: %s", task_key) + self._tasks.pop(task_key, None) + + return expired_slots + + def _get_run_slot(self, run_lease_id: str) -> RunSlot: + task_key = self._run_to_task.get(run_lease_id) + if task_key is None: + raise KeyError(f"Unknown run_lease_id: {run_lease_id}") + task_slot = self._tasks.get(task_key) + if task_slot is None: + raise KeyError(f"Run {run_lease_id} points to missing task slot") + run_slot = task_slot.runs.get(run_lease_id) + if run_slot is None: + raise KeyError(f"Run {run_lease_id} not found in task slot") + return run_slot + + async def allocate( + self, task_key: str, request_id: str | None = None + ) -> dict[str, Any]: + async with self._lock: + expired_slots = self._reap_idle_locked() + + if request_id: + idem_key = (task_key, request_id) + cached = self._idempotency.get(idem_key) + if cached is not None: + run_lease_id, _ = cached + if run_lease_id in self._run_to_task: + return {"lease_id": run_lease_id, "reused": True} + + task_slot = self._tasks.get(task_key) + if task_slot is None: + if len(self._tasks) >= self.max_tasks: + raise CapacityError( + "TASK_SLOTS_EXHAUSTED", + f"Worker at task capacity: {len(self._tasks)}/{self.max_tasks}", + ) + task_slot = TaskSlot(task_key=task_key) + self._tasks[task_key] = task_slot + + if len(task_slot.runs) >= self.max_runs_per_task: + raise CapacityError( + "RUN_SLOTS_EXHAUSTED", + f"Task {task_key} at run capacity: {len(task_slot.runs)}/{self.max_runs_per_task}", + ) + + env = self._new_env() + run_lease_id = f"run-{uuid.uuid4().hex[:16]}" + run_slot = RunSlot(run_lease_id=run_lease_id, task_key=task_key, env=env) + task_slot.runs[run_lease_id] = run_slot + task_slot.last_used_ts = time.time() + self._run_to_task[run_lease_id] = task_key + + if request_id: + self._idempotency[(task_key, request_id)] = (run_lease_id, time.time()) + + for tk, rid, rslot in expired_slots: + self._schedule_close(tk, rid, rslot, reason="Reaping idle run slot") + + return {"lease_id": run_lease_id, "reused": False} + + async def heartbeat(self, run_lease_id: str) -> None: + async with self._lock: + run_slot = self._get_run_slot(run_lease_id) + async with run_slot.lock: + run_slot.last_used_ts = time.time() + + async def reset( + self, + run_lease_id: str, + task_meta: dict[str, Any], + run_ctx_payload: dict[str, Any] | None = None, + task_timeouts: dict[str, Any] | None = None, + ) -> dict[str, Any]: + if not isinstance(task_meta, dict): + raise ValueError("task_meta must be a dict") + + async with self._lock: + run_slot = self._get_run_slot(run_lease_id) + + run_ctx = _build_run_ctx( + run_ctx_payload, default_log_dir=self.output_root / "AgentRunner_Output" + ) + timeouts = _parse_timeout_overrides(self.default_timeouts, task_timeouts) + task_spec = _build_task_spec(task_meta) + + async with run_slot.lock: + user_msg, tool_schemas = await run_slot.env.reset( + task_meta=task_meta, + task_spec=task_spec, + run_ctx=run_ctx, + timeouts=timeouts, + ) + run_slot.last_used_ts = time.time() + return {"user_msg": user_msg, "tool_schemas": tool_schemas} + + async def exec_tool( + self, run_lease_id: str, tool_name: str, arguments: dict[str, Any] | None = None + ) -> str: + async with self._lock: + run_slot = self._get_run_slot(run_lease_id) + async with run_slot.lock: + observation = await run_slot.env.exec_tool(tool_name, arguments or {}) + run_slot.last_used_ts = time.time() + return str(observation) + + async def evaluate(self, run_lease_id: str) -> float: + async with self._lock: + run_slot = self._get_run_slot(run_lease_id) + async with run_slot.lock: + score = await run_slot.env.evaluate() + run_slot.last_used_ts = time.time() + return float(score) + + async def close_run(self, run_lease_id: str) -> bool: + async with self._lock: + task_key = self._run_to_task.pop(run_lease_id, None) + if task_key is None: + logger.debug( + "close_run: lease %s already gone, nothing to do.", run_lease_id + ) + return False + task_slot = self._tasks.get(task_key) + run_slot = task_slot.runs.pop(run_lease_id, None) if task_slot else None + if task_slot is not None and not task_slot.runs: + self._tasks.pop(task_key, None) + logger.info("Removed empty task slot: %s", task_key) + + if run_slot is not None: + self._schedule_close( + task_key, run_lease_id, run_slot, reason="Closing run slot" + ) + return True + + async def status(self) -> dict[str, Any]: + async with self._lock: + tasks_info: dict[str, Any] = {} + total_runs = 0 + for tk, ts in self._tasks.items(): + tasks_info[tk] = {"active_runs": len(ts.runs)} + total_runs += len(ts.runs) + + return { + "max_tasks": self.max_tasks, + "active_tasks": len(self._tasks), + "max_runs_per_task": self.max_runs_per_task, + "total_active_runs": total_runs, + "pending_closes": len(self._closing_tasks), + "tasks": tasks_info, + } + + async def periodic_reap(self, interval: float = 60.0) -> None: + while True: + await asyncio.sleep(interval) + try: + async with self._lock: + expired_slots = self._reap_idle_locked() + for tk, rid, rslot in expired_slots: + self._schedule_close( + tk, rid, rslot, reason="Periodic reaper: idle run slot" + ) + if expired_slots: + logger.info( + "Periodic reaper cleaned up %d idle run slots", + len(expired_slots), + ) + except Exception: + logger.exception("Periodic reaper error") + + async def shutdown(self) -> None: + async with self._lock: + slots_to_close: list[tuple[str, str, RunSlot]] = [] + for task_key, task_slot in self._tasks.items(): + for run_lease_id, run_slot in task_slot.runs.items(): + slots_to_close.append((task_key, run_lease_id, run_slot)) + self._tasks.clear() + self._run_to_task.clear() + self._idempotency.clear() + + for task_key, run_lease_id, run_slot in slots_to_close: + self._schedule_close( + task_key, + run_lease_id, + run_slot, + reason="Closing run slot during shutdown", + ) + + if self._closing_tasks: + logger.info( + "Shutdown: waiting for %d pending close tasks...", + len(self._closing_tasks), + ) + await asyncio.gather(*self._closing_tasks, return_exceptions=True) + + +POOL: WorkerPool | None = None + + +@app.get("/healthz") +async def healthz() -> dict[str, Any]: + return {"ok": True} + + +@app.get("/status") +async def status() -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + return JSONResponse({"ok": True, "pool": await POOL.status()}) + + +@app.post("/allocate") +async def allocate(request: Request) -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + task_key = data.get("task_key", "") + request_id = data.get("request_id") + + if not task_key: + return JSONResponse( + {"ok": False, "error": "task_key is required"}, status_code=400 + ) + + try: + result = await POOL.allocate(task_key=str(task_key), request_id=request_id) + return JSONResponse({"ok": True, **result}) + except CapacityError as exc: + return JSONResponse( + {"ok": False, "error": exc.message, "code": exc.code}, status_code=429 + ) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +@app.post("/heartbeat") +async def heartbeat(request: Request) -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + lease_id = data.get("lease_id") + if not lease_id: + return JSONResponse( + {"ok": False, "error": "lease_id is required"}, status_code=400 + ) + + try: + await POOL.heartbeat(str(lease_id)) + return JSONResponse({"ok": True}) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +@app.post("/reset") +async def reset(request: Request) -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + lease_id = data.get("lease_id") + task_meta = data.get("task_meta") + run_ctx_payload = data.get("run_ctx") + task_timeouts = data.get("task_timeouts") + + if not lease_id: + return JSONResponse( + {"ok": False, "error": "lease_id is required"}, status_code=400 + ) + if not isinstance(task_meta, dict): + return JSONResponse( + {"ok": False, "error": "task_meta dict is required"}, status_code=400 + ) + + try: + out = await POOL.reset( + run_lease_id=str(lease_id), + task_meta=task_meta, + run_ctx_payload=run_ctx_payload, + task_timeouts=task_timeouts, + ) + return JSONResponse({"ok": True, **out}) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +@app.post("/exec_tool") +async def exec_tool(request: Request) -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + lease_id = data.get("lease_id") + tool_call = data.get("tool_call") + + if not lease_id: + return JSONResponse( + {"ok": False, "error": "lease_id is required"}, status_code=400 + ) + if not isinstance(tool_call, dict): + return JSONResponse( + {"ok": False, "error": "tool_call dict is required"}, status_code=400 + ) + + tool_name = tool_call.get("name") + arguments = tool_call.get("arguments") + + if not isinstance(tool_name, str) or not tool_name: + return JSONResponse( + {"ok": False, "error": "tool_call.name is required"}, status_code=400 + ) + if arguments is not None and not isinstance(arguments, dict): + return JSONResponse( + {"ok": False, "error": "tool_call.arguments must be a dict"}, + status_code=400, + ) + + try: + observation = await POOL.exec_tool( + str(lease_id), tool_name, arguments=arguments + ) + return JSONResponse({"ok": True, "observation": observation}) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +@app.post("/evaluate") +async def evaluate(request: Request) -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + lease_id = data.get("lease_id") + + if not lease_id: + return JSONResponse( + {"ok": False, "error": "lease_id is required"}, status_code=400 + ) + + try: + score = await POOL.evaluate(str(lease_id)) + return JSONResponse({"ok": True, "score": score}) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +@app.post("/close") +async def close(request: Request) -> JSONResponse: + if POOL is None: + return JSONResponse( + {"ok": False, "error": "Pool is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + lease_id = data.get("lease_id") + if not lease_id: + return JSONResponse( + {"ok": False, "error": "lease_id is required"}, status_code=400 + ) + + try: + found = await POOL.close_run(str(lease_id)) + return JSONResponse({"ok": True, "found": found}) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +_REAPER_TASK: asyncio.Task | None = None + + +@app.on_event("startup") +async def _on_startup() -> None: + global _REAPER_TASK + if POOL is not None: + _REAPER_TASK = asyncio.create_task(POOL.periodic_reap(interval=60.0)) + + +@app.on_event("shutdown") +async def _on_shutdown() -> None: + global POOL, _REAPER_TASK + if _REAPER_TASK is not None: + _REAPER_TASK.cancel() + _REAPER_TASK = None + if POOL is not None: + await POOL.shutdown() + POOL = None + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="C-layer: terminal env worker server") + + parser.add_argument("--host", type=str, default="0.0.0.0") + parser.add_argument( + "--port", type=int, default=int(os.getenv("ENV_SERVER_PORT", "18081")) + ) + + parser.add_argument( + "--max-tasks", type=int, default=int(os.getenv("WORKER_MAX_TASKS", "16")) + ) + parser.add_argument( + "--max-runs-per-task", + type=int, + default=int(os.getenv("WORKER_MAX_RUNS_PER_TASK", "8")), + ) + parser.add_argument( + "--run-idle-ttl", + type=int, + default=int(os.getenv("WORKER_RUN_IDLE_TTL", "600")), + help="Seconds before an idle RunSlot is reaped", + ) + + parser.add_argument( + "--output-root", + type=str, + default=os.getenv("TBENCH_OUTPUT_ROOT", "build_outputs"), + ) + + parser.add_argument( + "--ensure-image-timeout", + type=float, + default=float(os.getenv("ENSURE_IMAGE_TIMEOUT", "300.0")), + ) + parser.add_argument( + "--reset-session-timeout", + type=float, + default=float(os.getenv("RESET_SESSION_TIMEOUT", "300.0")), + ) + parser.add_argument( + "--close-session-timeout", + type=float, + default=float(os.getenv("CLOSE_SESSION_TIMEOUT", "60.0")), + ) + parser.add_argument( + "--eval-timeout", type=float, default=float(os.getenv("EVAL_TIMEOUT", "600.0")) + ) + parser.add_argument( + "--max-concurrent-closes", + type=int, + default=int(os.getenv("WORKER_MAX_CONCURRENT_CLOSES", "10")), + help="Max concurrent Docker stop operations", + ) + + return parser.parse_args() + + +def main() -> None: + global POOL + args = parse_args() + + logging.basicConfig( + level=logging.INFO, format="[%(asctime)s %(levelname)s %(name)s] %(message)s" + ) + + POOL = WorkerPool( + max_tasks=args.max_tasks, + max_runs_per_task=args.max_runs_per_task, + run_idle_ttl=args.run_idle_ttl, + output_root=args.output_root, + default_timeouts=TaskTimeouts( + ensure_image=float(args.ensure_image_timeout), + reset_session=float(args.reset_session_timeout), + close_session=float(args.close_session_timeout), + eval=float(args.eval_timeout), + ), + max_concurrent_closes=args.max_concurrent_closes, + ) + + logger.info( + "Starting worker server on %s:%s max_tasks=%s max_runs_per_task=%s", + args.host, + args.port, + args.max_tasks, + args.max_runs_per_task, + ) + + uvicorn.run(app, host=args.host, port=args.port, log_level="info") + + +if __name__ == "__main__": + main() diff --git a/terminal-rl/remote/run_pool_server.sh b/terminal-rl/remote/run_pool_server.sh new file mode 100644 index 0000000..567c63e --- /dev/null +++ b/terminal-rl/remote/run_pool_server.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash + +set -euo pipefail + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +TERMINAL_RL="$(cd -- "${SCRIPT_DIR}/.." &>/dev/null && pwd)" +REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." &>/dev/null && pwd)" + +cd "${REPO_ROOT}" + +export DATASET_DIR="${DATASET_DIR:-${TERMINAL_RL}/dataset}" +export TBENCH_OUTPUT_ROOT="${TBENCH_OUTPUT_ROOT:-${TERMINAL_RL}/build_outputs}" + +export TBENCH_DOCKER_IMAGE_SOURCE="${TBENCH_DOCKER_IMAGE_SOURCE:-build}" +export TBENCH_DOCKER_PULL_PREFIX="${TBENCH_DOCKER_PULL_PREFIX:-}" +export COMPOSE_OVERRIDE_PATH="${COMPOSE_OVERRIDE_PATH:-}" + +if [ -d "${REPO_ROOT}/.venv" ]; then + source .venv/bin/activate +fi + +# Start the pool server +exec python -m terminal-rl.remote.pool_server \ + --host 0.0.0.0 \ + --port "${ENV_SERVER_PORT:-18081}" \ + --max-tasks "${WORKER_MAX_TASKS:-16}" \ + --max-runs-per-task "${WORKER_MAX_RUNS_PER_TASK:-8}" \ + --output-root "${TBENCH_OUTPUT_ROOT}" diff --git a/terminal-rl/remote/setup.sh b/terminal-rl/remote/setup.sh new file mode 100644 index 0000000..4d57659 --- /dev/null +++ b/terminal-rl/remote/setup.sh @@ -0,0 +1,55 @@ +# System dependencies +if ! command -v docker &> /dev/null; then + echo "Docker not found, installing..." + sudo apt-get update + sudo apt-get install -y docker.io + sudo apt-get install -y docker-compose-v2 + sudo apt-get install -y docker-compose-plugin + sudo apt-get install -y docker-compose +else + echo "Docker found, skipping installation." +fi + +# Modify docker to increase network address pool +DOCKER_DAEMON_CONFIG='/etc/docker/daemon.json' + +# Backup existing daemon.json if it exists +if [ -f "$DOCKER_DAEMON_CONFIG" ]; then + echo "Backing up existing Docker daemon configuration..." + sudo cp "$DOCKER_DAEMON_CONFIG" "${DOCKER_DAEMON_CONFIG}.backup.$(date +%Y%m%d_%H%M%S)" +fi + +# Create or update daemon.json with network pool settings +echo "Configuring Docker daemon..." +sudo tee "$DOCKER_DAEMON_CONFIG" > /dev/null </dev/null; then + echo "uv not found, installing..." + curl -LsSf https://astral.sh/uv/install.sh | sh + export PATH="${HOME}/.local/bin:${PATH}" +else + echo "uv found, skipping installation." +fi + + + +# Setup Python env +uv venv .venv --python 3.12 +source .venv/bin/activate +uv pip install terminal-bench +uv pip install camel-ai +uv pip install fastapi \ No newline at end of file diff --git a/terminal-rl/remote/terminal_env.py b/terminal-rl/remote/terminal_env.py new file mode 100644 index 0000000..ccaa6eb --- /dev/null +++ b/terminal-rl/remote/terminal_env.py @@ -0,0 +1,335 @@ +from __future__ import annotations + +import asyncio +import json +import logging +import os +from functools import partial +from pathlib import Path +from typing import Any + +from camel.toolkits import FunctionTool, TerminalToolkit + +from terminal_bench.handlers.trial_handler import TrialHandler +from terminal_bench.parsers.base_parser import UnitTestStatus +from terminal_bench.parsers.parser_factory import ParserFactory +from terminal_bench.terminal.docker_compose_manager import DockerComposeManager +from terminal_bench.terminal.terminal import Terminal + +from ..custom_types import RunContext, TaskSpec, TaskTimeouts + +from .docker_compose_utils import compose_up_no_build, prepare_task_docker_image + +logger = logging.getLogger(__name__) + + +def _stop_terminal_compat(terminal: Terminal, timeout: float) -> None: + try: + terminal.stop(timeout=timeout) + except TypeError as exc: + if "unexpected keyword argument 'timeout'" not in str(exc): + raise + logger.warning( + "Terminal.stop(timeout=...) is unsupported; retrying with Terminal.stop()." + ) + terminal.stop() + + +def _drain_toolkit_sessions(toolkit: Any) -> None: + sessions = getattr(toolkit, "shell_sessions", None) + if not isinstance(sessions, dict): + return + lock = getattr(toolkit, "_session_lock", None) + try: + if lock is not None: + lock.acquire() + for session in sessions.values(): + proc = session.get("process") + if proc is not None: + try: + if hasattr(proc, "terminate"): + proc.terminate() + elif hasattr(proc, "close"): + proc.close() + except Exception: + pass + q = session.get("output_stream") + if q is not None: + try: + while not q.empty(): + q.get_nowait() + except Exception: + pass + sessions.clear() + finally: + if lock is not None: + try: + lock.release() + except RuntimeError: + pass + + +class TerminalEnv: + def __init__(self) -> None: + self._closed = False + self._task_spec: TaskSpec | None = None + self._run_ctx: RunContext | None = None + self._timeouts: TaskTimeouts | None = None + + self._trial_handler: TrialHandler | None = None + self._terminal: Terminal | None = None + self._parser = None + self._terminal_toolkit: TerminalToolkit | None = None + self._tools: dict[str, Any] = {} + + async def reset( + self, + *, + task_meta: dict[str, Any], + task_spec: TaskSpec, + run_ctx: RunContext, + timeouts: TaskTimeouts, + ) -> tuple[str, list[dict[str, Any]]]: + await self.close() + + self._closed = False + self._task_spec = task_spec + self._run_ctx = run_ctx + self._timeouts = timeouts + + image_prep = await asyncio.to_thread( + prepare_task_docker_image, + task=task_meta, + timeout=self._timeouts.ensure_image, + ) + + dataset_dir = str(os.getenv("DATASET_DIR", "")).strip() + if not dataset_dir: + raise ValueError("DATASET_DIR is required") + task_path = Path(dataset_dir) / self._task_spec.task_path + output_path = Path(self._run_ctx.log_dir).resolve() + output_path.mkdir(parents=True, exist_ok=True) + + def _sync_reset() -> tuple[str, list[dict[str, Any]]]: + self._trial_handler = TrialHandler( + trial_name=f"{self._task_spec.task_name}.{self._run_ctx.uid}.slime-run", + input_path=task_path, + output_path=output_path, + ) + task_config = self._trial_handler.task + self._parser = ParserFactory.get_parser(task_config.parser_name) + client_image_name = ( + image_prep.client_image_name or self._trial_handler.client_image_name + ) + + self._terminal = Terminal( + client_container_name=self._trial_handler.client_container_name, + client_image_name=client_image_name, + docker_compose_path=self._trial_handler.task_paths.docker_compose_path, + docker_image_name_prefix=self._trial_handler.docker_image_name_prefix, + sessions_logs_path=self._trial_handler.trial_paths.sessions_path, + agent_logs_path=self._trial_handler.trial_paths.agent_logging_dir, + no_rebuild=True, + cleanup=False, + ) + if image_prep.mode == "pull": + compose_up_no_build( + self._terminal, + timeout=self._timeouts.reset_session, + container_name=self._trial_handler.client_container_name, + logger=logger, + ) + else: + self._terminal.start(timeout=self._timeouts.reset_session) + try: + from .docker_compose_utils import ( + _DEFAULT_CONTAINER_MEMORY_LIMIT, + _apply_container_memory_limit, + ) + + _apply_container_memory_limit( + self._trial_handler.client_container_name, + _DEFAULT_CONTAINER_MEMORY_LIMIT, + logger=logger, + ) + except Exception: + pass + + session_logs_dir = ( + self._trial_handler.trial_paths.sessions_path + / "terminal_toolkit_session_logs" + ) + self._terminal_toolkit = TerminalToolkit( + timeout=20.0, + working_directory=None, + use_docker_backend=True, + docker_container_name=self._trial_handler.client_container_name, + session_logs_dir=session_logs_dir, + safe_mode=False, + ) + self._tools = { + "shell_exec": self._terminal_toolkit.shell_exec, + "shell_view": self._terminal_toolkit.shell_view, + "shell_write_to_process": self._terminal_toolkit.shell_write_to_process, + "shell_write_content_to_file": self._terminal_toolkit.shell_write_content_to_file, + } + + user_msg = f"Task name:{self._task_spec.task_name}\nTask instruction: {self._task_spec.instruction}" + function_tools = [FunctionTool(fn) for fn in self._tools.values()] + tool_schemas = [ + func_tool.get_openai_tool_schema() for func_tool in function_tools + ] + return user_msg, tool_schemas + + return await asyncio.to_thread(_sync_reset) + + async def exec_tool(self, name: str, arguments: dict[str, Any]) -> str: + if not self._tools: + raise RuntimeError("env is not initialized; call reset first") + + if name not in self._tools: + return f"[TOOL_ERROR] unknown tool: {name}" + + fn = self._tools[name] + + try: + if asyncio.iscoroutinefunction(fn): + result = await fn(**arguments) + elif hasattr(fn, "async_call") and callable(fn.async_call): + result = await fn.async_call(**arguments) + else: + result = await asyncio.to_thread(partial(fn, **arguments)) + except Exception as exc: + return f"[TOOL_ERROR] {name}: {type(exc).__name__}: {exc}" + + if isinstance(result, str): + return result + return json.dumps(result, ensure_ascii=False) + + async def evaluate(self) -> float: + if ( + self._trial_handler is None + or self._terminal is None + or self._parser is None + or self._timeouts is None + ): + raise RuntimeError("env is not initialized; call reset first") + + def _sync_eval() -> float: + task_name = ( + self._task_spec.task_name if self._task_spec is not None else "unknown" + ) + paths: list[Path] = [self._trial_handler.task_paths.run_tests_path] + if self._trial_handler.task_paths.test_dir.exists(): + paths.append(self._trial_handler.task_paths.test_dir) + + self._terminal.copy_to_container( + paths=paths, + container_dir=str(DockerComposeManager.CONTAINER_TEST_DIR), + ) + + test_session = self._terminal.create_session( + "tests", + is_active_stream=False, + as_configured_user=False, + ) + test_script_path = str( + DockerComposeManager.CONTAINER_TEST_DIR / "run-tests.sh" + ) + test_timeout_sec = min( + self._timeouts.eval, + 4 * self._trial_handler.task.max_test_timeout_sec, + ) + try: + test_session.send_keys( + [f"bash {test_script_path}", "Enter"], + block=True, + max_timeout_sec=test_timeout_sec, + ) + except TimeoutError as exc: + logger.warning( + "Evaluation tests timed out for task=%s after %.1fs.", + task_name, + test_timeout_sec, + ) + raise RuntimeError( + f"Evaluation tests timed out for task={task_name} after {test_timeout_sec:.1f}s" + ) from exc + + test_output = test_session.capture_pane(capture_entire=True) + try: + parser_results = self._parser.parse(test_output) + except Exception as exc: + tail = test_output[-2000:] if test_output else "" + logger.warning( + "Failed to parse test output for task=%s with parser=%s: %s. Output tail:\n%s", + task_name, + type(self._parser).__name__, + exc, + tail, + ) + raise RuntimeError( + f"Failed to parse test output for task={task_name} with parser={type(self._parser).__name__}: {exc}" + ) from exc + + if not parser_results: + return 0.0 + passed = sum( + 1 + for status in parser_results.values() + if status == UnitTestStatus.PASSED + ) + reward = ( + float(passed / len(parser_results)) if len(parser_results) > 0 else 0.0 + ) + return reward + + return await asyncio.wait_for( + asyncio.to_thread(_sync_eval), + timeout=self._timeouts.eval + 30.0, + ) + + async def close(self) -> None: + trial_name = ( + self._trial_handler.trial_name + if self._trial_handler is not None + else "unknown" + ) + if self._closed: + logger.warning("TerminalEnv %s already closed", trial_name) + return + self._closed = True + + terminal = self._terminal + timeouts = self._timeouts + toolkit = self._terminal_toolkit + + self._tools = {} + self._terminal = None + self._trial_handler = None + self._parser = None + self._terminal_toolkit = None + self._task_spec = None + self._run_ctx = None + self._timeouts = None + + if toolkit is not None: + try: + await asyncio.to_thread(toolkit.cleanup) + except Exception: + logger.exception( + "Failed to cleanup terminal toolkit for %s", trial_name + ) + try: + await asyncio.to_thread(_drain_toolkit_sessions, toolkit) + except Exception: + logger.exception("Failed to drain toolkit sessions for %s", trial_name) + + if terminal is not None and timeouts is not None: + try: + await asyncio.to_thread( + _stop_terminal_compat, terminal, timeouts.close_session + ) + logger.info("TerminalEnv %s closed", trial_name) + except Exception: + logger.exception("Failed to stop terminal session during close") diff --git a/terminal-rl/request_utils.py b/terminal-rl/request_utils.py new file mode 100644 index 0000000..4bc8fac --- /dev/null +++ b/terminal-rl/request_utils.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from typing import Any, Dict + +from fastapi import Request + + +async def json_payload(request: Request) -> Dict[str, Any]: + """Safely parse JSON body from a FastAPI request. + + Returns an empty dict if parsing fails or the payload is not a JSON object. + """ + try: + data = await request.json() + except Exception: + return {} + return data if isinstance(data, dict) else {} diff --git a/terminal-rl/rollout_log.py b/terminal-rl/rollout_log.py new file mode 100644 index 0000000..884498f --- /dev/null +++ b/terminal-rl/rollout_log.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import logging +from typing import Any, Dict, List + +import wandb +from slime.utils import logging_utils +from slime.utils.types import Sample +from slime.ray.rollout import compute_rollout_step + +logger = logging.getLogger(__name__) + + +def _ensure_terminal_step_metric(args) -> None: + if not getattr(args, "use_wandb", False): + return + try: + wandb.define_metric("terminal/*", step_metric="rollout/step") + except Exception as e: + logger.warning("Failed to define wandb step metric for terminal/*: %s", e) + + +def rollout_log(rollout_id, args, samples, rollout_extra_metrics, rollout_time): + + trainable = [s for s in samples if not getattr(s, "remove_sample", False)] + non_trainable = [s for s in samples if getattr(s, "remove_sample", False)] + + log_dict: Dict[str, Any] = {} + + total = len(samples) + n_failed = sum(1 for s in samples if s.status == Sample.Status.FAILED) + n_aborted = sum(1 for s in samples if s.status == Sample.Status.ABORTED) + n_truncated = sum(1 for s in samples if s.status == Sample.Status.TRUNCATED) + n_completed = sum(1 for s in samples if s.status == Sample.Status.COMPLETED) + + log_dict["terminal/total_samples"] = total + log_dict["terminal/completed"] = n_completed + log_dict["terminal/truncated"] = n_truncated + log_dict["terminal/failed"] = n_failed + log_dict["terminal/aborted"] = n_aborted + log_dict["terminal/failed_ratio"] = n_failed / total if total else 0.0 + log_dict["terminal/non_trainable_ratio"] = ( + len(non_trainable) / total if total else 0.0 + ) + + if trainable: + trainable_rewards = [s.reward["score"] for s in trainable] + log_dict["terminal/reward_mean"] = sum(trainable_rewards) / len( + trainable_rewards + ) + log_dict["terminal/reward_min"] = min(trainable_rewards) + log_dict["terminal/reward_max"] = max(trainable_rewards) + + trainable_accs = [] + for s in trainable: + if isinstance(s.reward, dict) and "accuracy" in s.reward: + trainable_accs.append(float(s.reward["accuracy"])) + if trainable_accs: + log_dict["terminal/accuracy"] = sum(trainable_accs) / len(trainable_accs) + + trainable_prm = [] + for s in trainable: + if isinstance(s.reward, dict) and "prm_turn_score" in s.reward: + trainable_prm.append(float(s.reward["prm_turn_score"])) + if trainable_prm: + log_dict["terminal/prm_turn_score"] = sum(trainable_prm) / len( + trainable_prm + ) + + log_dict["terminal/rollout_time"] = rollout_time + + step = compute_rollout_step(args, rollout_id) + log_dict["rollout/step"] = step + _ensure_terminal_step_metric(args) + logging_utils.log(args, log_dict, step_key="rollout/step") + + return False diff --git a/terminal-rl/router_server.py b/terminal-rl/router_server.py new file mode 100644 index 0000000..5c6b66b --- /dev/null +++ b/terminal-rl/router_server.py @@ -0,0 +1,443 @@ +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +from hashlib import sha1 +from typing import Any + +import aiohttp +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse + +from .request_utils import json_payload + +logger = logging.getLogger("terminal.env.router") +app = FastAPI() + + +def _format_error(exc: BaseException) -> str: + detail = str(exc).strip() + if detail: + return f"{type(exc).__name__}: {detail}" + return type(exc).__name__ + + +def _status_from_payload(payload: dict[str, Any], default: int) -> int: + raw = payload.get("status_code") + if isinstance(raw, int): + return raw + return default + + +class Router: + def __init__( + self, + worker_urls: list[str], + forward_timeout: float = 600.0, + forward_retries: int = 1, + forward_retry_backoff: float = 0.2, + ): + if not worker_urls: + raise ValueError("At least one worker URL is required") + self.workers = [u.rstrip("/") for u in worker_urls] + self.forward_timeout = float(forward_timeout) + self.forward_retries = max(0, int(forward_retries)) + self.forward_retry_backoff = max(0.0, float(forward_retry_backoff)) + self._session: aiohttp.ClientSession | None = None + + @property + def num_workers(self) -> int: + return len(self.workers) + + async def startup(self) -> None: + if self._session is None or self._session.closed: + timeout = aiohttp.ClientTimeout(total=self.forward_timeout) + connector = aiohttp.TCPConnector(limit=0, ttl_dns_cache=300) + self._session = aiohttp.ClientSession(timeout=timeout, connector=connector) + + async def shutdown(self) -> None: + if self._session is not None and not self._session.closed: + await self._session.close() + self._session = None + + def select_worker(self, task_key: str) -> tuple[int, str]: + digest = sha1(task_key.encode("utf-8")).digest() + idx = ( + int.from_bytes(digest[:8], byteorder="big", signed=False) % self.num_workers + ) + return idx, self.workers[idx] + + @staticmethod + def encode_lease(worker_idx: int, worker_lease: str) -> str: + return f"{worker_idx}:{worker_lease}" + + @staticmethod + def decode_lease(global_lease: str) -> tuple[int, str]: + sep = global_lease.index(":") + return int(global_lease[:sep]), global_lease[sep + 1 :] + + def worker_url(self, worker_idx: int) -> str: + return self.workers[worker_idx] + + def iter_worker_candidates(self, start_idx: int) -> list[tuple[int, str]]: + return [ + ( + (start_idx + offset) % self.num_workers, + self.workers[(start_idx + offset) % self.num_workers], + ) + for offset in range(self.num_workers) + ] + + async def _request( + self, + method: str, + worker_url: str, + path: str, + payload: dict[str, Any] | None, + timeout: float | None, + ) -> tuple[dict[str, Any], int]: + if self._session is None: + raise RuntimeError("Router HTTP session is not initialized") + + kwargs: dict[str, Any] = {} + if payload is not None: + kwargs["json"] = payload + if timeout is not None: + kwargs["timeout"] = aiohttp.ClientTimeout(total=float(timeout)) + + max_attempts = self.forward_retries + 1 + for attempt in range(1, max_attempts + 1): + try: + async with self._session.request( + method, f"{worker_url}{path}", **kwargs + ) as resp: + status = resp.status + try: + body = await resp.json(content_type=None) + except Exception: + raw_text = await resp.text() + body = { + "ok": False, + "error": "Worker returned non-JSON response", + "raw_text": raw_text, + "status_code": status, + } + return body, status + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + if attempt >= max_attempts: + raise + logger.warning( + "Upstream request failed (%s %s) worker=%s attempt=%d/%d err=%s", + method, + path, + worker_url, + attempt, + max_attempts, + _format_error(exc), + ) + backoff = self.forward_retry_backoff * attempt + if backoff > 0: + await asyncio.sleep(backoff) + + async def forward( + self, + worker_url: str, + path: str, + payload: dict[str, Any], + timeout: float | None = None, + ) -> tuple[dict[str, Any], int]: + return await self._request("POST", worker_url, path, payload, timeout) + + async def forward_by_lease( + self, + global_lease: str, + path: str, + payload: dict[str, Any], + timeout: float | None = None, + ) -> tuple[dict[str, Any], int]: + worker_idx, worker_lease = self.decode_lease(global_lease) + url = self.worker_url(worker_idx) + forwarded_payload = dict(payload) + forwarded_payload["lease_id"] = worker_lease + return await self.forward(url, path, forwarded_payload, timeout) + + async def worker_status( + self, worker_url: str, timeout: float = 10.0 + ) -> tuple[dict[str, Any], int]: + return await self._request("GET", worker_url, "/status", None, timeout) + + +ROUTER: Router | None = None + + +def _worker_unreachable( + *, + worker_idx: int, + worker_url: str, + path: str, + exc: BaseException, + lease_id: str | None = None, + task_key: str | None = None, +) -> JSONResponse: + payload: dict[str, Any] = { + "ok": False, + "error": f"Worker unreachable: {_format_error(exc)}", + "worker_idx": worker_idx, + "worker_url": worker_url, + "path": path, + } + if lease_id: + payload["lease_id"] = lease_id + if task_key: + payload["task_key"] = task_key + return JSONResponse(payload, status_code=502) + + +@app.get("/healthz") +async def healthz() -> dict[str, Any]: + return {"ok": True} + + +@app.get("/status") +async def status() -> JSONResponse: + if ROUTER is None: + return JSONResponse( + {"ok": False, "error": "Router is not initialized"}, status_code=500 + ) + + async def _fetch(idx: int, url: str) -> dict[str, Any]: + try: + data, _ = await ROUTER.worker_status(url, timeout=10) + return {"worker_idx": idx, "url": url, **data} + except Exception as exc: + return { + "worker_idx": idx, + "url": url, + "ok": False, + "error": _format_error(exc), + } + + workers = await asyncio.gather( + *[_fetch(idx, url) for idx, url in enumerate(ROUTER.workers)] + ) + return JSONResponse( + {"ok": True, "num_workers": ROUTER.num_workers, "workers": workers} + ) + + +@app.post("/allocate") +async def allocate(request: Request) -> JSONResponse: + if ROUTER is None: + return JSONResponse( + {"ok": False, "error": "Router is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + task_key = data.get("task_key", "") + request_id = data.get("request_id") + + if not task_key: + return JSONResponse( + {"ok": False, "error": "task_key is required"}, status_code=400 + ) + + try: + payload = {"task_key": task_key, "request_id": request_id} + primary_idx, _ = ROUTER.select_worker(str(task_key)) + upstream_errors: list[dict[str, Any]] = [] + for worker_idx, worker_url in ROUTER.iter_worker_candidates(primary_idx): + try: + result, code = await ROUTER.forward(worker_url, "/allocate", payload) + if worker_idx != primary_idx: + logger.warning( + "Primary worker unreachable for /allocate task_key=%s; fallback worker_idx=%d url=%s", + task_key, + worker_idx, + worker_url, + ) + if result.get("ok") and "lease_id" in result: + result["lease_id"] = Router.encode_lease( + worker_idx, str(result["lease_id"]) + ) + result["worker_idx"] = worker_idx + + return JSONResponse( + result, status_code=_status_from_payload(result, code) + ) + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + logger.warning( + "Worker unreachable for /allocate task_key=%s worker_idx=%d url=%s err=%s", + task_key, + worker_idx, + worker_url, + _format_error(exc), + ) + upstream_errors.append( + { + "worker_idx": worker_idx, + "worker_url": worker_url, + "detail": _format_error(exc), + } + ) + + return JSONResponse( + { + "ok": False, + "error": "Worker unreachable: all candidates failed for /allocate", + "task_key": task_key, + "primary_worker_idx": primary_idx, + "upstream_errors": upstream_errors, + }, + status_code=502, + ) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +async def _lease_proxy(path: str, request: Request) -> JSONResponse: + if ROUTER is None: + return JSONResponse( + {"ok": False, "error": "Router is not initialized"}, status_code=500 + ) + + data = await json_payload(request) + global_lease = data.get("lease_id", "") + if not global_lease: + return JSONResponse( + {"ok": False, "error": "lease_id is required"}, status_code=400 + ) + + try: + worker_idx, worker_lease = ROUTER.decode_lease(str(global_lease)) + worker_url = ROUTER.worker_url(worker_idx) + except (ValueError, IndexError) as exc: + return JSONResponse( + {"ok": False, "error": f"Invalid lease_id format: {exc}"}, status_code=400 + ) + + payload = dict(data) + payload["lease_id"] = worker_lease + + try: + result, code = await ROUTER.forward(worker_url, path, payload) + return JSONResponse(result, status_code=_status_from_payload(result, code)) + except (aiohttp.ClientError, asyncio.TimeoutError) as exc: + return _worker_unreachable( + worker_idx=worker_idx, + worker_url=worker_url, + path=path, + exc=exc, + lease_id=str(global_lease), + ) + except Exception as exc: + return JSONResponse({"ok": False, "error": str(exc)}, status_code=500) + + +@app.post("/heartbeat") +async def heartbeat(request: Request) -> JSONResponse: + return await _lease_proxy("/heartbeat", request) + + +@app.post("/reset") +async def reset(request: Request) -> JSONResponse: + return await _lease_proxy("/reset", request) + + +@app.post("/exec_tool") +async def exec_tool(request: Request) -> JSONResponse: + return await _lease_proxy("/exec_tool", request) + + +@app.post("/evaluate") +async def evaluate(request: Request) -> JSONResponse: + return await _lease_proxy("/evaluate", request) + + +@app.post("/close") +async def close(request: Request) -> JSONResponse: + return await _lease_proxy("/close", request) + + +@app.on_event("startup") +async def _on_startup() -> None: + if ROUTER is not None: + await ROUTER.startup() + + +@app.on_event("shutdown") +async def _on_shutdown() -> None: + if ROUTER is not None: + await ROUTER.shutdown() + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="B-layer: terminal env router server") + parser.add_argument("--host", type=str, default="0.0.0.0") + parser.add_argument( + "--port", type=int, default=int(os.getenv("ROUTER_PORT", "18080")) + ) + parser.add_argument( + "--workers", + type=str, + default=os.getenv("WORKER_URLS", ""), + help="Comma-separated worker URLs, e.g. http://w0:18081,http://w1:18081", + ) + parser.add_argument( + "--forward-timeout", + type=float, + default=float(os.getenv("ROUTER_FORWARD_TIMEOUT", "600.0")), + help="HTTP timeout (seconds) when forwarding to a worker", + ) + parser.add_argument( + "--forward-retries", + type=int, + default=int(os.getenv("ROUTER_FORWARD_RETRIES", "1")), + help="Retries for transient worker connection errors", + ) + parser.add_argument( + "--forward-retry-backoff", + type=float, + default=float(os.getenv("ROUTER_FORWARD_RETRY_BACKOFF", "0.2")), + help="Linear backoff (seconds) between worker retries", + ) + return parser.parse_args() + + +def main() -> None: + global ROUTER + args = parse_args() + + logging.basicConfig( + level=logging.INFO, format="[%(asctime)s %(levelname)s %(name)s] %(message)s" + ) + + worker_urls = [u.strip() for u in args.workers.split(",") if u.strip()] + if not worker_urls: + raise SystemExit( + "ERROR: --workers (or WORKER_URLS env) must list at least one worker URL" + ) + + ROUTER = Router( + worker_urls=worker_urls, + forward_timeout=args.forward_timeout, + forward_retries=args.forward_retries, + forward_retry_backoff=args.forward_retry_backoff, + ) + logger.info( + "Starting router on %s:%s workers=%s forward_timeout=%s forward_retries=%s forward_retry_backoff=%s", + args.host, + args.port, + worker_urls, + args.forward_timeout, + args.forward_retries, + args.forward_retry_backoff, + ) + + uvicorn.run(app, host=args.host, port=args.port, log_level="info") + + +if __name__ == "__main__": + main() diff --git a/terminal-rl/terminal_qwen3_8b_prm_rl_2nodes.sh b/terminal-rl/terminal_qwen3_8b_prm_rl_2nodes.sh new file mode 100755 index 0000000..e4d931c --- /dev/null +++ b/terminal-rl/terminal_qwen3_8b_prm_rl_2nodes.sh @@ -0,0 +1,427 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +log() { echo "[$(date +'%F %T')] $*"; } + +require_cmd() { command -v "$1" >/dev/null 2>&1 || { echo "[ERROR] missing cmd: $1"; exit 1; }; } + +export PYTHONUNBUFFERED=1 +export PYTHONFAULTHANDLER=1 + +export RAY_health_check_failure_threshold=${RAY_health_check_failure_threshold:-20} +export RAY_health_check_period_ms=${RAY_health_check_period_ms:-5000} +export RAY_health_check_timeout_ms=${RAY_health_check_timeout_ms:-30000} +export RAY_num_heartbeats_timeout=${RAY_num_heartbeats_timeout:-60} + +require_env() { + local name="$1" + if [[ -z "${!name:-}" ]]; then + echo "[ERROR] missing env: ${name}" + exit 1 + fi +} + +NUM_NODES=${NUM_NODES:-2} +NUM_GPUS_PER_NODE=${NUM_GPUS_PER_NODE:-8} +ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-1} +ACTOR_GPUS_PER_NODE=${ACTOR_GPUS_PER_NODE:-4} +ROLLOUT_GPUS_TOTAL=${ROLLOUT_GPUS_TOTAL:-4} +ROLLOUT_NUM_GPUS_PER_ENGINE=${ROLLOUT_NUM_GPUS_PER_ENGINE:-2} +PRM_GPUS_TOTAL=${PRM_GPUS_TOTAL:-8} +PRM_GPUS_PER_ENGINE=${PRM_GPUS_PER_ENGINE:-1} + +PRM_ENABLE="${PRM_ENABLE:-1}" +PRM_MODEL_PATH="${PRM_MODEL_PATH:-}" +PRM_TEMPERATURE="${PRM_TEMPERATURE:-0.0}" +PRM_MAX_NEW_TOKENS="${PRM_MAX_NEW_TOKENS:-4096}" +PRM_M="${PRM_M:-1}" +PRM_STEP_COEF="${PRM_STEP_COEF:-1.0}" +PRM_SGLANG_URL="${PRM_SGLANG_URL:-}" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}" +export REPO_ROOT + +SLIME_DIR="${SLIME_DIR:-${REPO_ROOT}/slime}" +MEGATRON_LM_PATH="${MEGATRON_LM_PATH:-${REPO_ROOT}/Megatron-LM}" + +export SLIME_PKG_DIR="${REPO_ROOT}/slime" +export MEGATRON_DIR="${REPO_ROOT}/Megatron-LM" + +source "${SLIME_DIR}/scripts/models/qwen3-8B.sh" + +# Paths: set/export before running (no built-in defaults). +HF_HOME="${HF_HOME:-}" +HF_CKPT="${HF_CKPT:-}" +REF_LOAD="${REF_LOAD:-}" +SAVE_CKPT="${SAVE_CKPT:-}" +RESUME_LOAD="${RESUME_LOAD:-${SAVE_CKPT}}" +ROLLOUT_PROMPT_DATA="${ROLLOUT_PROMPT_DATA:-}" + +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-max_split_size_mb:2048,expandable_segments:True}" + + +MLP_ROLE_INDEX=${MLP_ROLE_INDEX:-0} +HEAD_ADDR="${MLP_WORKER_0_HOST:-${MASTER_ADDR:-$(hostname -I | awk '{print $1}')}}" +_WORKER_IP_VAR="MLP_WORKER_${MLP_ROLE_INDEX}_HOST" +NODE_IP="${!_WORKER_IP_VAR:-${WORKER_IP:-$(hostname -I | awk '{print $1}')}}" + +unset MASTER_ADDR +export no_proxy="127.0.0.1,${HEAD_ADDR}" +log "MLP_ROLE_INDEX=${MLP_ROLE_INDEX}, HEAD_ADDR=${HEAD_ADDR}, NODE_IP=${NODE_IP}" + +export USE_REMOTE_ENV="${USE_REMOTE_ENV:-1}" +export PROVIDER_NAME="${PROVIDER_NAME:-pull}" +export ENV_SERVER_BIND_HOST="${ENV_SERVER_BIND_HOST:-0.0.0.0}" +export ENV_SERVER_PORT="${ENV_SERVER_PORT:-18080}" +export ENV_SERVER_HOST="${ENV_SERVER_HOST:-${HEAD_ADDR}}" +export ENV_SERVER_URL="${ENV_SERVER_URL:-}" +export START_ENV_POOL_SERVER="${START_ENV_POOL_SERVER:-0}" + +export RAY_TMPDIR=/tmp/ray_${MLP_ROLE_INDEX} + +export WORKER_URLS="${WORKER_URLS:-}" + +ROUTER_SESSION_NAME="${ROUTER_SESSION_NAME:-terminal_router}" +ROUTER_CONDA_ENV_PATH="${ROUTER_CONDA_ENV_PATH:-}" +ROUTER_PROJECT_DIR="${ROUTER_PROJECT_DIR:-${REPO_ROOT}}" +ROUTER_HOST="${ROUTER_HOST:-0.0.0.0}" +ROUTER_PORT="${ROUTER_PORT:-${ENV_SERVER_PORT}}" + +CHECK_HOST="${CHECK_HOST:-127.0.0.1}" +CHECK_WAIT_SECS="${CHECK_WAIT_SECS:-60}" +ROUTER_RESTART="${ROUTER_RESTART:-1}" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CKPT}" + --ref-load "${REF_LOAD}" + --load "${RESUME_LOAD}" + --save "${SAVE_CKPT}" + --save-interval 3 + --rotary-base 1000000 +) + +ROLLOUT_ARGS=( + --prompt-data "${ROLLOUT_PROMPT_DATA}" + --input-key task + --rollout-shuffle + --reward-key score + --num-rollout 2000 + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-max-context-len 16384 + --rollout-temperature 1 + + --num-steps-per-rollout 2 + --balance-data +) + +EVAL_ARGS=( + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 + --log-probs-chunk-size 1024 +) + +GRPO_ARGS=( + --advantage-estimator step_wise + --dynamic_history + --use-kl-loss + --kl-loss-coef 0.01 + --kl-loss-type k3 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_KEY_VALUE=${WANDB_KEY:-${WANDB_API_KEY:-}} +if [ -n "${WANDB_KEY_VALUE}" ]; then + WANDB_ARGS=( + --use-wandb + --wandb-project slime + --wandb-group qwen3-8B-prm-2nodes-rl_terminal + --wandb-key ${WANDB_KEY_VALUE} + ) +else + WANDB_ARGS=() +fi + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine ${ROLLOUT_NUM_GPUS_PER_ENGINE} + --sglang-mem-fraction-static 0.6 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +CUSTOM_ARGS=( + --custom-generate-function-path generate.generate + --custom-rollout-log-function-path rollout_log.rollout_log +) + +PRM_ARGS=( + --prm-m "${PRM_M}" + --prm-temperature "${PRM_TEMPERATURE}" + --prm-max-new-tokens "${PRM_MAX_NEW_TOKENS}" + --prm-num-gpus "${PRM_GPUS_TOTAL}" + --prm-num-gpus-per-engine "${PRM_GPUS_PER_ENGINE}" + --prm-step-coef "${PRM_STEP_COEF}" +) + +if [[ "${PRM_ENABLE:-0}" == "1" ]]; then + PRM_ARGS+=(--prm-enable) + PRM_ARGS+=(--prm-model-path "${PRM_MODEL_PATH}") +fi +if [[ -n "${PRM_SGLANG_URL}" ]]; then + PRM_ARGS+=(--prm-sglang-url "${PRM_SGLANG_URL}") +fi + +check_gpus() { + local total_available=$((NUM_NODES * NUM_GPUS_PER_NODE)) + local total_requested=$((ACTOR_NUM_NODES * ACTOR_GPUS_PER_NODE + ROLLOUT_GPUS_TOTAL + PRM_GPUS_TOTAL)) + if (( total_requested > total_available )); then + echo "Requested GPUs exceed cluster capacity." + echo "requested=${total_requested}, available=${total_available}" + echo "actor=$((ACTOR_NUM_NODES * ACTOR_GPUS_PER_NODE)), rollout=${ROLLOUT_GPUS_TOTAL}, prm=${PRM_GPUS_TOTAL}" + exit 1 + fi + if [[ "${PRM_ENABLE:-0}" == "1" ]] && (( PRM_GPUS_TOTAL <= 0 )); then + echo "PRM_ENABLE=1 but PRM_GPUS_TOTAL=${PRM_GPUS_TOTAL} <= 0" + exit 1 + fi +} + +cleanup_prev() { + log "cleanup previous processes" + pkill -9 sglang || true + sleep 3 + ray stop --force || true + pkill -9 ray || true + pkill -9 python || true + sleep 3 + pkill -9 ray || true + pkill -9 python || true +} + +start_router() { + if [[ ${MLP_ROLE_INDEX} -ne 0 ]]; then + log "worker node ${MLP_ROLE_INDEX}, skip start_router" + return 0 + fi + + require_cmd curl + mkdir -p "${ROUTER_PROJECT_DIR}/logs" + local logf="${ROUTER_PROJECT_DIR}/logs/router_${ROUTER_PORT}_2nodes.log" + + "${ROUTER_CONDA_ENV_PATH}/bin/python" -m terminal-rl.router_server \ + --host "${ROUTER_HOST}" --port "${ROUTER_PORT}" --workers "${WORKER_URLS}" \ + > "${logf}" 2>&1 & + + export ROUTER_PID=$! + log "router started pid=${ROUTER_PID}, log=${logf}" + + trap 'set +e; kill "${ROUTER_PID}" 2>/dev/null || true' EXIT INT TERM + + sleep 1 + tail -n 50 "${logf}" || true +} + +check_router() { + if [[ ${MLP_ROLE_INDEX} -ne 0 ]]; then + log "worker node ${MLP_ROLE_INDEX}, skip check_router" + return 0 + fi + + require_cmd curl + local base_url="http://${CHECK_HOST}:${ROUTER_PORT}" + + log "wait router healthz up to ${CHECK_WAIT_SECS}s: ${base_url}/healthz" + for ((i=1; i<=CHECK_WAIT_SECS; i++)); do + if curl -fsS "${base_url}/healthz" >/dev/null 2>&1; then + log "router is up" + break + fi + sleep 1 + done + + log "curl ${base_url}/status" + curl -sS "${base_url}/status" + echo + log "curl ${base_url}/healthz" + curl -sS "${base_url}/healthz" + echo +} + +detect_nvlink() { + local count + count="$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l || true)" + if [[ "${count:-0}" -gt 0 ]]; then + export HAS_NVLINK=1 + else + export HAS_NVLINK=0 + fi + log "HAS_NVLINK=${HAS_NVLINK} (detected ${count} NVLink references)" +} + +maybe_fill_env_server_url() { + if [[ "${USE_REMOTE_ENV}" == "1" && -z "${ENV_SERVER_URL}" ]]; then + export ENV_SERVER_URL="http://${ENV_SERVER_HOST}:${ENV_SERVER_PORT}" + if [[ "${START_ENV_POOL_SERVER}" == "0" ]]; then + export START_ENV_POOL_SERVER=1 + fi + fi + log "ENV_SERVER_URL=${ENV_SERVER_URL} START_ENV_POOL_SERVER=${START_ENV_POOL_SERVER}" +} + +start_ray_head() { + require_cmd ray + mkdir -p "${RAY_TMPDIR}" + + if [[ ${MLP_ROLE_INDEX} -eq 0 ]]; then + log "start ray head on node 0" + ray start --head \ + --node-ip-address "${NODE_IP}" \ + --num-gpus "${NUM_GPUS_PER_NODE}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 \ + --temp-dir "${RAY_TMPDIR}" + else + log "worker node ${MLP_ROLE_INDEX}: wait 30s then join ray cluster at ${HEAD_ADDR}:6379" + sleep 30 + ray start \ + --address="${HEAD_ADDR}:6379" \ + --num-gpus "${NUM_GPUS_PER_NODE}" \ + --node-ip-address "${NODE_IP}" \ + --temp-dir "${RAY_TMPDIR}" + fi +} + +build_runtime_env_json() { + python3 - <<'PY' +import json, os + +parts = [ + os.environ.get("REPO_ROOT",""), + os.environ.get("SLIME_PKG_DIR",""), + os.environ.get("MEGATRON_DIR",""), + os.environ.get("SCRIPT_DIR",""), +] +pythonpath = ":".join([p for p in parts if p]) + +env_vars = { + "PYTHONPATH": pythonpath, + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": os.environ.get("HAS_NVLINK","0"), + "PYTORCH_CUDA_ALLOC_CONF": os.environ.get("PYTORCH_CUDA_ALLOC_CONF",""), + "USE_REMOTE_ENV": os.environ.get("USE_REMOTE_ENV","0"), + "ENV_SERVER_URL": os.environ.get("ENV_SERVER_URL",""), + "PRM_SGLANG_URL": os.environ.get("PRM_SGLANG_URL",""), +} +print(json.dumps({"env_vars": env_vars})) +PY +} + +submit_job() { + require_env HF_CKPT + require_env REF_LOAD + require_env SAVE_CKPT + require_env ROLLOUT_PROMPT_DATA + if [[ "${PRM_ENABLE:-0}" == "1" ]]; then + require_env PRM_MODEL_PATH + fi + + if [[ ${MLP_ROLE_INDEX} -eq 0 ]]; then + log "submit ray job (head node)" + local runtime_env_json + runtime_env_json="$(build_runtime_env_json)" + local submission_id="${RAY_JOB_SUBMISSION_ID:-terminal_qwen3_8b_prm_2nodes_$(date +%Y%m%d_%H%M%S)}" + + ray job submit --address="http://${HEAD_ADDR}:8265" \ + --submission-id "${submission_id}" \ + --no-wait \ + --runtime-env-json="${runtime_env_json}" \ + -- python3 -u ${SLIME_DIR}/train_async.py \ + --actor-num-nodes "${ACTOR_NUM_NODES}" \ + --actor-num-gpus-per-node "${ACTOR_GPUS_PER_NODE}" \ + --rollout-num-gpus "${ROLLOUT_GPUS_TOTAL}" \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${CUSTOM_ARGS[@]}" \ + "${PRM_ARGS[@]}" + + log "Following live Ray logs for ${submission_id}" + set +e + ray job logs --address="http://${HEAD_ADDR}:8265" "${submission_id}" -f --log-style=record + local ray_log_exit=$? + local ray_status + ray_status=$(ray job status --address="http://${HEAD_ADDR}:8265" "${submission_id}" --log-style=record 2>&1) + echo "${ray_status}" + set -e + + if [[ "${ray_status}" == *"SUCCEEDED"* ]]; then + exit 0 + fi + + echo "Ray job failed (submission id: ${submission_id}, logs exit: ${ray_log_exit})" + exit 1 + else + log "Worker node ${MLP_ROLE_INDEX} joined the cluster. Waiting for job to finish..." + while ray status > /dev/null 2>&1; do + sleep 60 + done + log "Ray cluster stopped. Worker node ${MLP_ROLE_INDEX} exiting." + fi +} + +cleanup_prev + +start_router +check_router + +check_gpus +detect_nvlink +maybe_fill_env_server_url +export SCRIPT_DIR +start_ray_head +submit_job diff --git a/terminal-rl/terminal_qwen3_8b_rl.sh b/terminal-rl/terminal_qwen3_8b_rl.sh new file mode 100644 index 0000000..6dd43b5 --- /dev/null +++ b/terminal-rl/terminal_qwen3_8b_rl.sh @@ -0,0 +1,309 @@ +#!/usr/bin/env bash +set -euo pipefail +set -x + +log() { echo "[$(date +'%F %T')] $*"; } + +require_cmd() { command -v "$1" >/dev/null 2>&1 || { echo "[ERROR] missing cmd: $1"; exit 1; }; } + +export PYTHONBUFFERED=16 + +NUM_GPUS="${NUM_GPUS:-8}" +ACTOR_GPUS="${ACTOR_GPUS:-4}" +ROLLOUT_GPUS="${ROLLOUT_GPUS:-4}" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}" + +export REPO_ROOT +export SLIME_DIR="${REPO_ROOT}/slime" +export MEGATRON_DIR="${MEGATRON_DIR:-${REPO_ROOT}/Megatron-LM}" + +source "${SLIME_DIR}/scripts/models/qwen3-8B.sh" + +# Paths: set/export before running (no built-in defaults). +HF_HOME="${HF_HOME:-}" +HF_CKPT="${HF_CKPT:-}" +REF_LOAD="${REF_LOAD:-}" +SAVE_CKPT="${SAVE_CKPT:-}" +RESUME_LOAD="${RESUME_LOAD:-${SAVE_CKPT}}" +ROLLOUT_PROMPT_DATA="${ROLLOUT_PROMPT_DATA:-}" + +export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-max_split_size_mb:2048,expandable_segments:True}" +export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}" + +export USE_REMOTE_ENV="${USE_REMOTE_ENV:-1}" +export PROVIDER_NAME="${PROVIDER_NAME:-pull}" +export ENV_SERVER_BIND_HOST="${ENV_SERVER_BIND_HOST:-0.0.0.0}" +export ENV_SERVER_PORT="${ENV_SERVER_PORT:-18080}" +export ENV_SERVER_HOST="${ENV_SERVER_HOST:-${MASTER_ADDR}}" +export ENV_SERVER_URL="${ENV_SERVER_URL:-}" +export START_ENV_POOL_SERVER="${START_ENV_POOL_SERVER:-0}" + +export RAY_TMPDIR="${RAY_TMPDIR:-}" + +export WORKER_URLS="${WORKER_URLS:-}" + +ROUTER_SESSION_NAME="${ROUTER_SESSION_NAME:-terminal_router}" +ROUTER_CONDA_ENV_PATH="${ROUTER_CONDA_ENV_PATH:-}" +ROUTER_PROJECT_DIR="${ROUTER_PROJECT_DIR:-${REPO_ROOT}}" +export ROUTER_CONDA_ENV_PATH +CONDA_PYTHON_VERSION="${CONDA_PYTHON_VERSION:-3.12}" +export CONDA_PYTHON_VERSION +ROUTER_HOST="${ROUTER_HOST:-0.0.0.0}" +ROUTER_PORT="${ROUTER_PORT:-${ENV_SERVER_PORT}}" + +CHECK_HOST="${CHECK_HOST:-127.0.0.1}" +CHECK_WAIT_SECS="${CHECK_WAIT_SECS:-60}" +ROUTER_RESTART="${ROUTER_RESTART:-1}" + +CKPT_ARGS=( + --hf-checkpoint "${HF_CKPT}" + --ref-load "${REF_LOAD}" + --load "${RESUME_LOAD}" + --save "${SAVE_CKPT}" + --save-interval 3 + --rotary-base 1000000 +) + +ROLLOUT_ARGS=( + --prompt-data "${ROLLOUT_PROMPT_DATA}" + --input-key task + --rollout-shuffle + --reward-key score + --num-rollout 2000 + --rollout-batch-size 16 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-max-context-len 16384 + --rollout-temperature 1 + + --num-steps-per-rollout 2 + --balance-data +) + +EVAL_ARGS=( + --n-samples-per-eval-prompt 16 + --eval-max-response-len 16384 + --eval-top-p 1 +) + + +PERF_ARGS=( + --tensor-model-parallel-size 4 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 + --log-probs-chunk-size 1024 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --dynamic_history + --use-kl-loss + --kl-loss-coef 0.01 + --kl-loss-type k3 +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + --use-wandb + --wandb-project slime + --wandb-group qwen3-8B-rl_terminal + --wandb-key ${WANDB_KEY} +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 2 + --sglang-mem-fraction-static 0.6 +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + +CUSTOM_ARGS=( + --custom-generate-function-path generate.generate + --custom-rollout-log-function-path rollout_log.rollout_log +) + +check_gpus() { + if (( ACTOR_GPUS + ROLLOUT_GPUS > NUM_GPUS )); then + echo "ACTOR_GPUS + ROLLOUT_GPUS must be <= NUM_GPUS" + echo "ACTOR_GPUS=${ACTOR_GPUS}, ROLLOUT_GPUS=${ROLLOUT_GPUS}, NUM_GPUS=${NUM_GPUS}" + exit 1 + fi +} + +cleanup_prev() { + log "cleanup previous processes" + pkill -9 sglang || true + sleep 3 + ray stop --force || true + pkill -9 ray || true + pkill -9 python || true + sleep 3 + pkill -9 ray || true + pkill -9 python || true +} + +start_router() { + require_cmd curl + mkdir -p "${ROUTER_PROJECT_DIR}/logs" + local logf="${ROUTER_PROJECT_DIR}/logs/router_${ROUTER_PORT}.log" + + "${ROUTER_CONDA_ENV_PATH}/bin/python" -m terminal-rl.router_server \ + --host "${ROUTER_HOST}" --port "${ROUTER_PORT}" --workers "${WORKER_URLS}" \ + > "${logf}" 2>&1 & + + export ROUTER_PID=$! + log "router started pid=${ROUTER_PID}, log=${logf}" + + sleep 1 + tail -n 50 "${logf}" || true +} + +check_router() { + require_cmd curl + local base_url="http://${CHECK_HOST}:${ROUTER_PORT}" + + log "wait router healthz up to ${CHECK_WAIT_SECS}s: ${base_url}/healthz" + for ((i=1; i<=CHECK_WAIT_SECS; i++)); do + if curl -fsS "${base_url}/healthz" >/dev/null 2>&1; then + log "router is up" + break + fi + sleep 1 + done + + log "curl ${base_url}/status" + curl -sS "${base_url}/status" + echo + log "curl ${base_url}/healthz" + curl -sS "${base_url}/healthz" + echo +} + +detect_nvlink() { + local count + count="$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l || true)" + if [[ "${count:-0}" -gt 0 ]]; then + export HAS_NVLINK=1 + else + export HAS_NVLINK=0 + fi + log "HAS_NVLINK=${HAS_NVLINK} (detected ${count} NVLink references)" +} + +maybe_fill_env_server_url() { + if [[ "${USE_REMOTE_ENV}" == "1" && -z "${ENV_SERVER_URL}" ]]; then + export ENV_SERVER_URL="http://${ENV_SERVER_HOST}:${ENV_SERVER_PORT}" + if [[ "${START_ENV_POOL_SERVER}" == "0" ]]; then + export START_ENV_POOL_SERVER=1 + fi + fi + log "ENV_SERVER_URL=${ENV_SERVER_URL} START_ENV_POOL_SERVER=${START_ENV_POOL_SERVER}" +} + +start_ray_head() { + require_cmd ray + log "start ray head" + mkdir -p "${RAY_TMPDIR}" + ray start --head \ + --node-ip-address "${MASTER_ADDR}" \ + --num-gpus "${NUM_GPUS}" \ + --disable-usage-stats \ + --dashboard-host=0.0.0.0 \ + --dashboard-port=8265 \ + --temp-dir "${RAY_TMPDIR}" +} + + +build_runtime_env_json() { + python3 - <<'PY' +import json, os + +conda_env = os.environ.get("ROUTER_CONDA_ENV_PATH", "") +py_ver = os.environ.get("CONDA_PYTHON_VERSION", "3.12") +site_packages = f"{conda_env}/lib/python{py_ver}/site-packages" if conda_env else "" + +parts = [ + os.environ.get("REPO_ROOT",""), + os.environ.get("SLIME_PKG_DIR",""), + os.environ.get("MEGATRON_DIR",""), + os.environ.get("SCRIPT_DIR",""), + site_packages, +] +pythonpath = ":".join([p for p in parts if p]) + +env_vars = { + "PYTHONPATH": pythonpath, + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": os.environ.get("HAS_NVLINK","0"), + "PYTORCH_CUDA_ALLOC_CONF": os.environ.get("PYTORCH_CUDA_ALLOC_CONF",""), + "USE_REMOTE_ENV": os.environ.get("USE_REMOTE_ENV","0"), + "ENV_SERVER_URL": os.environ.get("ENV_SERVER_URL",""), +} +print(json.dumps({"env_vars": env_vars})) +PY +} + +submit_job() { + log "submit ray job" + local runtime_env_json + runtime_env_json="$(build_runtime_env_json)" + + ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${runtime_env_json}" \ + -- python3 ${SLIME_DIR}/train_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node "${ACTOR_GPUS}" \ + --rollout-num-gpus "${ROLLOUT_GPUS}" \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${ROLLOUT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${GRPO_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${SGLANG_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${CUSTOM_ARGS[@]}" +} + +cleanup_prev + +start_router +check_router + +check_gpus +detect_nvlink +maybe_fill_env_server_url +export SCRIPT_DIR +start_ray_head +submit_job