Add terminal-rl directory with Terminal RL training scripts and infrastructure

.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

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://<worker-ip>: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://<prm-router-ip>:<prm-router-port>"
+```
+
+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)).

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 = "<tool_call>\n"
+        eot_token = "\n</tool_call>"
+
+        # 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"],
+        )

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] + "...<truncated>")
+
+
+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

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"""
+<role>
+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.
+</role>
+"""
+    TEAM_STRUCTURE_PROMPT = f""
+
+    OPERATING_ENVIRONMENT_PROMPT = f"""
+<operating_environment>
+- **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}.
+</operating_environment>
+"""
+
+    MANDATORY_INSTRUCTIONS_PROMPT = f"""
+<mandatory_instructions>
+- 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.
+</mandatory_instructions>
+"""
+    CAPABILITIES_PROMPT = """
+<capabilities>
+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.
+""" + \
+"""
+</capabilities>
+"""
+
+    PHILOSOPHY_PROMPT = """
+<philosophy>
+- **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. 
+</philosophy>
+"""
+
+    SHELL_TIPS_PROMPT = f"""
+<terminal_tips>
+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.
+    
+</terminal_tips>
+"""
+    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

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,
+    )

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

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: <DATASET_DIR>/<tasks_dir_name>_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,
+    )

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 <dataset_name>")

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

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}")

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,
+                )

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}")

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://<this-machine-ip-or-hostname>: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` → `<prefix>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/<your-org>/<your-image>:task-"
+export COMPOSE_OVERRIDE_PATH="terminal-rl/remote/compose_override.yaml"
+bash terminal-rl/remote/run_pool_server.sh
+```
+

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}

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
+    )

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()

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}"

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 <<EOF
+{
+  "default-address-pools": [
+    {
+      "base": "10.200.0.0/16",
+      "size": 24
+    }
+  ]
+}
+EOF
+# Restart Docker to apply changes
+echo "Restarting Docker daemon..."
+sudo systemctl restart docker
+echo "Docker configuration complete!"
+
+# install uv if not found
+if ! command -v uv &>/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

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")

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 {}

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

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()

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

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