openenv/tbench2: adapter-driven TB2 fidelity (no OpenEnv fork/vendor)

examples/experimental/openenv/openenv_agent_function.py

@@ -21,6 +21,11 @@
                      stragglers that would otherwise stall the whole rollout batch).
   AGENT_MODEL_NAME   model name sent to the policy (default: "model")
   MILES_ROUTER_EXTERNAL_HOST  optional host rewrite for off-cluster agents
+  OPENENV_TASK_WORKDIR  container dir every agent command + eval runs in (default:
+                     /app, the TB2 task image WORKDIR). Empty string disables the
+                     prefix. Needed because upstream OpenEnv defaults to /task.
+  OPENENV_TB2_TESTS_SRC  where the upstream env stages the task's tests inside the
+                     container (default: /task/tests); copied to /tests for test.sh.
 
 Daytona pool (optional; takes precedence over OPENENV_ENV_URL when set):
   OPENENV_DAYTONA_SNAPSHOT      Daytona snapshot name to provision sandboxes from
@@ -65,6 +70,55 @@
 # Max chars of command output fed back to the policy per turn (keeps context bounded).
 _OBS_CHAR_CAP = 4000
 
+# --- Adapter-driven Terminal-Bench-2 fidelity --------------------------------
+# Upstream OpenEnv's Tbench2DockerEnvironment runs the task container with workdir
+# /task (a copy of the task *source*) and scores via bare `pytest tests/` there.
+# Real TB2 tasks live at /app (the task image's WORKDIR) and are scored by the
+# task's canonical tests/test.sh, which pins the pytest toolchain, copies test.py
+# into /app, runs test_outputs.py, and writes the binary result to
+# /logs/verifier/reward.txt. We reproduce that faithfully from the adapter --
+# without patching OpenEnv or vendoring it -- by (a) running every agent command
+# in _TASK_WORKDIR and (b) driving the canonical harness through a plain `exec`
+# step instead of the env's built-in (non-canonical) `evaluate` action.
+#
+# This assumes the env server is UNMODIFIED upstream, which copies the task dir
+# (tests included) into the container at _TB2_TESTS_SRC.
+_TASK_WORKDIR = os.getenv("OPENENV_TASK_WORKDIR", "/app")
+_TB2_TESTS_SRC = os.getenv("OPENENV_TB2_TESTS_SRC", "/task/tests")
+
+# The eval exec echoes reward.txt on this marker so we can parse it out of stdout.
+_REWARD_MARKER = "__TB2_REWARD__:"
+# Honor an empty _TASK_WORKDIR (workdir prefix disabled) the same way
+# _apply_workdir does, instead of silently forcing /app.
+_EVAL_CD_CMD = f"cd {_TASK_WORKDIR} && " if _TASK_WORKDIR else ""
+_CANONICAL_EVAL_CMD = (
+    # rm the reward file first so a stale one can never be read back if test.sh
+    # fails to run (e.g. in a reused sandbox where /logs survives across episodes).
+    "mkdir -p /tests /logs/verifier && rm -f /logs/verifier/reward.txt && "
+    f"cp -a {_TB2_TESTS_SRC}/. /tests/ 2>/dev/null || true; "
+    f"{_EVAL_CD_CMD}bash /tests/test.sh > /tmp/tb2_testsh.log 2>&1; "
+    f"echo {_REWARD_MARKER}$(cat /logs/verifier/reward.txt 2>/dev/null)"
+)
+
+
+def _apply_workdir(command: str) -> str:
+    """Prefix an agent command so it runs in the real task workdir (/app)."""
+    if not _TASK_WORKDIR:
+        return command
+    return f"cd {_TASK_WORKDIR} && {command}"
+
+
+def _parse_reward_marker(output: str) -> float:
+    """Parse the reward.txt value the canonical-eval exec echoed on its marker line."""
+    for line in output.splitlines()[::-1]:
+        if _REWARD_MARKER in line:
+            raw = line.split(_REWARD_MARKER, 1)[1].strip()
+            try:
+                return float(raw) if raw else 0.0
+            except ValueError:
+                return 0.0
+    return 0.0
+
 # Per-message WS recv timeout. Docker-mode tbench2 reset (container create),
 # exec, and evaluate (pytest) each routinely exceed the EnvClient default of 60s.
 _MESSAGE_TIMEOUT_S = float(os.getenv("OPENENV_MESSAGE_TIMEOUT_S", "600"))
@@ -361,9 +415,11 @@ async def _multi_turn(
     """Agentic loop: reset(task) -> {policy -> exec -> feed output back} -> evaluate (tbench2).
 
     The policy emits one shell command per turn (a ```bash block or the bare
-    reply); the loop ends when the policy stops emitting a command, says
-    TASK_COMPLETE, or hits OPENENV_MAX_TURNS. The final ``evaluate`` action runs
-    the task's pytest suite and returns the binary reward.
+    reply), executed in the real task workdir (_TASK_WORKDIR, /app); the loop ends
+    when the policy stops emitting a command, says TASK_COMPLETE, or hits
+    OPENENV_MAX_TURNS. Scoring runs the task's canonical tests/test.sh via an
+    ``exec`` step and parses /logs/verifier/reward.txt for the binary reward
+    (faithful to Terminal-Bench-2, and needs no OpenEnv-side changes).
     """
     action_cls = classes["action"]
     task_id = metadata.get("task_id") or metadata.get("task_name")
@@ -409,7 +465,9 @@ async def body(env: Any) -> tuple[float, int, list[float], list[float], float, f
                 break
 
             t0 = time.monotonic()
-            step_result = await env.step(action_cls(action_type="exec", command=command))
+            step_result = await env.step(
+                action_cls(action_type="exec", command=_apply_workdir(command))
+            )
             tool_times.append(time.monotonic() - t0)
             output = _obs_field(step_result, "output")
             # Feed the command output back as a user turn, not a tool turn. GLM
@@ -425,9 +483,11 @@ async def body(env: Any) -> tuple[float, int, list[float], list[float], float, f
             convo.append({"role": "user", "content": content})
 
         t0 = time.monotonic()
-        eval_result = await env.step(action_cls(action_type="evaluate"))
+        eval_result = await env.step(
+            action_cls(action_type="exec", command=_CANONICAL_EVAL_CMD)
+        )
         eval_time = time.monotonic() - t0
-        reward = float(getattr(eval_result, "reward", 0.0) or 0.0)
+        reward = _parse_reward_marker(_obs_field(eval_result, "output"))
 
         # rm-hack: the tbench2 env server (TB2_OUTPUT_DIR=/tmp/tbench2_env_runs)
         # leaves a per-episode trial dir under that path after every episode, which
@@ -512,6 +572,8 @@ async def run(
             f"OpenEnv tbench2 episode exceeded {_MAX_ROLLOUT_TIME_S:.0f}s; "
             "terminating with reward 0"
         )
+        # eval_report empty: the episode was cancelled before the canonical
+        # eval ever ran, so there is no pytest report to surface.
         return {
             "reward": 0.0,
             "exit_status": "timeout",
@@ -522,6 +584,11 @@ async def run(
         logger.error(f"OpenEnv tbench2 episode failed: {e}", exc_info=True)
         return None
 
+    # eval_report is intentionally empty: the canonical-eval marker protocol
+    # (see _REWARD_MARKER) echoes back only the scalar reward. The detailed
+    # pytest CTRF report is written inside the sandbox at
+    # /logs/verifier/ctrf.json and is deliberately not captured back to the
+    # trainer, which consumes only `reward`.
     return {
         "reward": reward,
         "exit_status": "completed",

examples/experimental/openenv/openenv_launch_common.py

@@ -0,0 +1,164 @@
+"""Shared launch helpers for the OpenEnv tbench2 learning launchers.
+
+``run-openenv-tbench2.py`` (GLM-4.7-Flash) and ``run-openenv-tbench2-dsv4.py``
+(DeepSeek-V4-Flash) drive the same agentic adapter and differ only in the
+model-family serving/training profile. The model-agnostic fragments (process
+cleanup, GRPO/optimizer/rollout/agent flags, W&B + Prometheus wiring, and the
+OpenEnv/Daytona env-var plumbing) live here so the two launchers cannot silently
+drift apart. Each launcher keeps only its own perf/sglang/misc profile and its
+``ScriptArgs`` defaults.
+"""
+
+import os
+import subprocess
+import time
+from typing import Protocol
+
+
+class LaunchArgs(Protocol):
+    """The config fields the shared helpers read (satisfied by each launcher's ScriptArgs)."""
+
+    prompt_data: str
+    rollout_batch_size: int
+    n_samples_per_prompt: int
+    max_seq_len: int
+    global_batch_size: int
+
+    openenv_env_url: str
+    agent_model_name: str
+    openenv_max_turns: int
+    openenv_max_rollout_time_seconds: int
+    openenv_daytona_snapshot: str
+    openenv_daytona_pool_size: int
+    openenv_daytona_port: int
+    daytona_api_key: str
+    router_external_host: str
+    miles_host_ip: str
+
+    wandb_key: str
+    wandb_project: str
+    wandb_team: str
+    wandb_run_name: str
+
+    use_prometheus: bool
+    prometheus_port: int
+    prometheus_run_name: str
+
+
+def cleanup() -> None:
+    """Kill old Ray jobs and stale processes to free GPU resources."""
+    my_pid = os.getpid()
+    ppid = os.getppid()
+    print(f"Cleanup starting (pid={my_pid}, ppid={ppid})")
+    targets = ["sglang", "train.py", "MegatronTrain"]
+    exclude = f"grep -v '^{my_pid}$' | grep -v '^{ppid}$'"
+    for t in targets:
+        subprocess.run(
+            f"pgrep -f '{t}' | {exclude} | xargs -r kill 2>/dev/null || true",
+            shell=True,
+        )
+    time.sleep(5)
+    print(f"Cleanup complete (pid={my_pid}) — old processes killed.")
+
+
+def rollout_args(args: LaunchArgs) -> str:
+    return (
+        f"--prompt-data {args.prompt_data} "
+        "--input-key prompt "
+        "--metadata-key metadata "
+        "--rollout-shuffle "
+        "--num-rollout 40 "
+        f"--rollout-batch-size {args.rollout_batch_size} "
+        f"--n-samples-per-prompt {args.n_samples_per_prompt} "
+        "--rollout-temperature 0.8 "
+        "--rollout-max-response-len 8192 "
+        f"--max-seq-len {args.max_seq_len} "
+        f"--global-batch-size {args.global_batch_size} "
+        "--balance-data "
+    )
+
+
+def grpo_args() -> str:
+    return (
+        "--advantage-estimator grpo "
+        "--use-kl-loss "
+        "--kl-loss-coef 0.01 "
+        "--kl-loss-type low_var_kl "
+        "--entropy-coef 0.0 "
+        "--eps-clip 0.2 "
+        "--eps-clip-high 0.28 "
+    )
+
+
+def optimizer_args() -> str:
+    return (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+
+def agent_args(tito_model: str) -> str:
+    """Agentic-rollout wiring. Only the TITO surface differs across models."""
+    return (
+        "--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate "
+        "--custom-agent-function-path openenv_agent_function.run "
+        "--custom-rm-path openenv_generate.reward_func "
+        "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_aborted "
+        f"--tito-model {tito_model} "
+        "--use-session-server "
+        "--session-server-port 30000 "
+        "--tito-allowed-append-roles user tool "
+    )
+
+
+def wandb_args(args: LaunchArgs) -> str:
+    if not args.wandb_key:
+        return ""
+    out = (
+        "--use-wandb "
+        f"--wandb-project {args.wandb_project} "
+        f"--wandb-group {args.wandb_run_name} "
+        f"--wandb-key {args.wandb_key} "
+    )
+    if args.wandb_team:
+        out += f"--wandb-team {args.wandb_team} "
+    return out
+
+
+def prometheus_args(args: LaunchArgs) -> str:
+    if not args.use_prometheus:
+        return ""
+    return (
+        "--use-prometheus "
+        f"--prometheus-port {args.prometheus_port} "
+        f"--prometheus-run-name {args.prometheus_run_name} "
+    )
+
+
+def base_env_vars(args: LaunchArgs, script_dir: str, megatron_path: str, miles_root: str) -> dict[str, str]:
+    return {
+        "PYTHONPATH": f"{megatron_path}:{script_dir}:{miles_root}",
+        "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1",
+        "OPENENV_ENV_URL": args.openenv_env_url,
+        "OPENENV_MAX_TURNS": str(args.openenv_max_turns),
+        "OPENENV_MAX_ROLLOUT_TIME_SECONDS": str(args.openenv_max_rollout_time_seconds),
+        "AGENT_MODEL_NAME": args.agent_model_name,
+    }
+
+
+def apply_optional_env_vars(env: dict[str, str], args: LaunchArgs) -> None:
+    """Add host-rewrite / Daytona-pool env vars when the args request them."""
+    if args.miles_host_ip:
+        env["MILES_HOST_IP"] = args.miles_host_ip
+    if args.router_external_host:
+        env["MILES_ROUTER_EXTERNAL_HOST"] = args.router_external_host
+    if args.openenv_daytona_snapshot:
+        assert args.daytona_api_key, "DAYTONA_API_KEY required when openenv_daytona_snapshot is set"
+        env["OPENENV_DAYTONA_SNAPSHOT"] = args.openenv_daytona_snapshot
+        env["OPENENV_DAYTONA_POOL_SIZE"] = str(args.openenv_daytona_pool_size)
+        env["OPENENV_DAYTONA_PORT"] = str(args.openenv_daytona_port)
+        env["DAYTONA_API_KEY"] = args.daytona_api_key