diff --git a/examples/experimental/openenv/openenv_agent_function.py b/examples/experimental/openenv/openenv_agent_function.py index 4b953047e4..772512ea5b 100644 --- a/examples/experimental/openenv/openenv_agent_function.py +++ b/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", diff --git a/examples/experimental/openenv/openenv_launch_common.py b/examples/experimental/openenv/openenv_launch_common.py new file mode 100644 index 0000000000..e30e5fea4d --- /dev/null +++ b/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 diff --git a/examples/experimental/openenv/run-openenv-tbench2-dsv4.py b/examples/experimental/openenv/run-openenv-tbench2-dsv4.py index 82122ab1e4..787fc09a5b 100644 --- a/examples/experimental/openenv/run-openenv-tbench2-dsv4.py +++ b/examples/experimental/openenv/run-openenv-tbench2-dsv4.py @@ -49,8 +49,6 @@ """ import os -import subprocess -import time from dataclasses import dataclass from pathlib import Path from typing import Literal @@ -58,6 +56,7 @@ import typer import miles.utils.external_utils.command_utils as U +import openenv_launch_common as C SCRIPT_DIR = Path(__file__).resolve().parent @@ -154,22 +153,6 @@ def _parallel_config(num_nodes: int, num_gpus_per_node: int) -> str: ) -def cleanup(): - """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 execute(args: ScriptArgs): ckpt_args = ( f"--hf-checkpoint {args.hf_checkpoint} " @@ -178,20 +161,7 @@ def execute(args: ScriptArgs): "--save-interval 20 " ) - rollout_args = ( - 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 " - ) + rollout_args = C.rollout_args(args) # dsv4 perf: bshd + micro-batch-size 1 (NOT --use-dynamic-batch-size). perf_args = _parallel_config(args.num_nodes, args.num_gpus_per_node) @@ -206,24 +176,9 @@ def execute(args: ScriptArgs): "--use-precision-aware-optimizer " ) - grpo_args = ( - "--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 " - ) + grpo_args = C.grpo_args() - optimizer_args = ( - "--optimizer adam " - "--lr 1e-6 " - "--lr-decay-style constant " - "--weight-decay 0.1 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - ) + optimizer_args = C.optimizer_args() # dsv4-flash rollout engines: tp4/dp1/ep4, 4 GPUs/engine. DP attention stays # OFF for Flash. MTP/EAGLE per --enable-mtp. @@ -248,16 +203,7 @@ def execute(args: ScriptArgs): "--sglang-speculative-num-draft-tokens 4 " ) - agent_args = ( - "--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 " - "--tito-model deepseekv4 " - "--use-session-server " - "--session-server-port 30000 " - "--tito-allowed-append-roles user tool " - ) + agent_args = C.agent_args("deepseekv4") misc_args = ( "--attention-dropout 0.0 " @@ -280,13 +226,10 @@ def execute(args: ScriptArgs): if args.enable_r3: misc_args += "--use-rollout-routing-replay " - extra_env_vars = { - "PYTHONPATH": f"{args.megatron_path}:{SCRIPT_DIR}:{U.repo_base_dir}", - "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, + extra_env_vars = C.base_env_vars( + args, str(SCRIPT_DIR), args.megatron_path, U.repo_base_dir + ) + extra_env_vars |= { "SGLANG_SKIP_CHECKPOINT_LOAD_CHECK": "1", "SGLANG_DSV4_FP4_EXPERTS": "0", "SGLANG_HEALTH_CHECK_TIMEOUT": "120", @@ -307,24 +250,9 @@ def execute(args: ScriptArgs): debug_args = "--debug-rollout-only " if args.mode == "debug_rollout_only" else "" dump_args = f"--dump-details {args.dump_details} " if args.dump_details else "" - wandb_args = "" - if args.wandb_key: - wandb_args = ( - "--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: - wandb_args += f"--wandb-team {args.wandb_team} " - - prometheus_args = "" - if args.use_prometheus: - prometheus_args = ( - "--use-prometheus " - f"--prometheus-port {args.prometheus_port} " - f"--prometheus-run-name {args.prometheus_run_name} " - ) + wandb_args = C.wandb_args(args) + + prometheus_args = C.prometheus_args(args) train_args = ( f"{ckpt_args}" @@ -341,16 +269,7 @@ def execute(args: ScriptArgs): f"{dump_args}" ) - if args.miles_host_ip: - extra_env_vars["MILES_HOST_IP"] = args.miles_host_ip - if args.router_external_host: - extra_env_vars["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" - extra_env_vars["OPENENV_DAYTONA_SNAPSHOT"] = args.openenv_daytona_snapshot - extra_env_vars["OPENENV_DAYTONA_POOL_SIZE"] = str(args.openenv_daytona_pool_size) - extra_env_vars["OPENENV_DAYTONA_PORT"] = str(args.openenv_daytona_port) - extra_env_vars["DAYTONA_API_KEY"] = args.daytona_api_key + C.apply_optional_env_vars(extra_env_vars, args) U.execute_train( train_args=train_args, @@ -364,7 +283,7 @@ def execute(args: ScriptArgs): @U.dataclass_cli def main(args: ScriptArgs): - cleanup() + C.cleanup() execute(args) diff --git a/examples/experimental/openenv/run-openenv-tbench2.py b/examples/experimental/openenv/run-openenv-tbench2.py index b5b0da7648..4cce8d23af 100644 --- a/examples/experimental/openenv/run-openenv-tbench2.py +++ b/examples/experimental/openenv/run-openenv-tbench2.py @@ -31,8 +31,6 @@ """ import os -import subprocess -import time from dataclasses import dataclass from pathlib import Path from typing import Literal @@ -40,6 +38,7 @@ import typer import miles.utils.external_utils.command_utils as U +import openenv_launch_common as C SCRIPT_DIR = Path(__file__).resolve().parent @@ -107,22 +106,6 @@ class ScriptArgs(U.ExecuteTrainConfig): prometheus_run_name: str = "openenv-tbench2-learn" -def cleanup(): - """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 prepare(args: ScriptArgs): """Convert HF checkpoint to torch_dist format if not already done.""" U.convert_checkpoint( @@ -143,20 +126,7 @@ def execute(args: ScriptArgs): "--save-interval 100 " ) - rollout_args = ( - 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 " - ) + rollout_args = C.rollout_args(args) perf_args = ( "--tensor-model-parallel-size 4 " @@ -175,24 +145,9 @@ def execute(args: ScriptArgs): "--use-precision-aware-optimizer " ) - grpo_args = ( - "--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 " - ) + grpo_args = C.grpo_args() - optimizer_args = ( - "--optimizer adam " - "--lr 1e-6 " - "--lr-decay-style constant " - "--weight-decay 0.1 " - "--adam-beta1 0.9 " - "--adam-beta2 0.98 " - ) + optimizer_args = C.optimizer_args() sglang_args = ( "--rollout-num-gpus-per-engine 1 " @@ -202,16 +157,7 @@ def execute(args: ScriptArgs): "--sglang-router-port 31000 " ) - agent_args = ( - "--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 " - "--tito-model glm47 " - "--use-session-server " - "--session-server-port 30000 " - "--tito-allowed-append-roles user tool " - ) + agent_args = C.agent_args("glm47") misc_args = ( "--attention-dropout 0.0 " @@ -229,24 +175,9 @@ def execute(args: ScriptArgs): dump_args = f"--dump-details {args.dump_details} " if args.dump_details else "" - wandb_args = "" - if args.wandb_key: - wandb_args = ( - "--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: - wandb_args += f"--wandb-team {args.wandb_team} " - - prometheus_args = "" - if args.use_prometheus: - prometheus_args = ( - "--use-prometheus " - f"--prometheus-port {args.prometheus_port} " - f"--prometheus-run-name {args.prometheus_run_name} " - ) + wandb_args = C.wandb_args(args) + + prometheus_args = C.prometheus_args(args) train_args = ( f"{ckpt_args}" @@ -263,26 +194,10 @@ def execute(args: ScriptArgs): f"{dump_args}" ) - miles_root = U.repo_base_dir - - extra_env_vars = { - "PYTHONPATH": f"{args.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, - } - if args.miles_host_ip: - extra_env_vars["MILES_HOST_IP"] = args.miles_host_ip - if args.router_external_host: - extra_env_vars["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" - extra_env_vars["OPENENV_DAYTONA_SNAPSHOT"] = args.openenv_daytona_snapshot - extra_env_vars["OPENENV_DAYTONA_POOL_SIZE"] = str(args.openenv_daytona_pool_size) - extra_env_vars["OPENENV_DAYTONA_PORT"] = str(args.openenv_daytona_port) - extra_env_vars["DAYTONA_API_KEY"] = args.daytona_api_key + extra_env_vars = C.base_env_vars( + args, str(SCRIPT_DIR), args.megatron_path, U.repo_base_dir + ) + C.apply_optional_env_vars(extra_env_vars, args) U.execute_train( train_args=train_args, @@ -296,7 +211,7 @@ def execute(args: ScriptArgs): @U.dataclass_cli def main(args: ScriptArgs): - cleanup() + C.cleanup() if not args.skip_prepare: prepare(args) execute(args)