diff --git a/examples/experimental/swe-agent-v2-amd/README.md b/examples/experimental/swe-agent-v2-amd/README.md new file mode 100644 index 0000000000..104f10d084 --- /dev/null +++ b/examples/experimental/swe-agent-v2-amd/README.md @@ -0,0 +1,47 @@ +# Agent V2 — Qwen3 / Qwen3-Coder on ROCm/AMD + +ROCm/AMD-oriented variant of [`../swe-agent-v2`](../swe-agent-v2). It is **self-contained** — +it carries its own copies of `run.py`, `generate.py`, and `swe_agent_function.py` so it runs +without modifying the shared example. See [`../swe-agent-v2/README.md`](../swe-agent-v2/README.md) +for the full pipeline/architecture; this doc covers only the Qwen3 launcher and the ROCm notes. + +## Qwen3 / Qwen3-Coder + +`run-qwen3-swe.py` runs the same pipeline with Qwen3 instead of GLM-4.7-Flash. It reuses +`run.py` and only swaps the model + SGLang/TITO parsers (`qwen25` / `qwen3` / `tito-model qwen3`). + +```bash +# General model +python run-qwen3-swe.py --prompt-data /root/swe_train.jsonl + +# Coding-specialised base (recommended for SWE-bench — see "Reward signal" below) +python run-qwen3-swe.py --coder --prompt-data /root/swe_train.jsonl +``` + +`--coder` uses **Qwen3-Coder-30B-A3B-Instruct**, which is arch-identical to Qwen3-30B-A3B +(same converter and `scripts/models/qwen3-30B-A3B.sh`) but uses `rope_theta=1e7`. The launcher +sets `MODEL_ARGS_ROTARY_BASE=10000000` for the Megatron side automatically; SGLang reads the +value from the HF config. + +### Reward signal (why the coder) + +GRPO learns only from **within-group reward variance**. A cold general model that solves ~0% +of SWE-bench Verified produces all-zero groups → zero advantage → no learning. Use a base that +lands in the "some pass / some fail" band (the coder), an easier task set (e.g. SWE-Gym), or a +warm-started checkpoint. With `rollout_batch_size=2`, also consider +`--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std` +plus over-sampling so every step contains variance-bearing groups. + +### ROCm / AMD prerequisites + +- **SGLang build with the `return_meta_info` chat-completion patch.** The TITO session server + requires `choice.meta_info.output_token_logprobs` from `/v1/chat/completions`; a build + without it returns HTTP 502 on every agent call. (This is a miles-SGLang-version + requirement, not AMD-specific.) +- **GPU visibility:** handled automatically — `run.py` sets + `RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1` only on ROCm. +- **agent-env container:** needs the Docker Compose v2 plugin (`docker compose`) — the Debian + `docker.io` package does not bundle it — and, for `mini-swe-agent`, a litellm install that + includes its proxy deps (recent litellm eagerly imports them): install with + `--with 'litellm[proxy]'` or pin litellm, ideally baked into the task base image so it is not + reinstalled per trial. diff --git a/examples/experimental/swe-agent-v2-amd/generate.py b/examples/experimental/swe-agent-v2-amd/generate.py new file mode 100644 index 0000000000..878fba3282 --- /dev/null +++ b/examples/experimental/swe-agent-v2-amd/generate.py @@ -0,0 +1,111 @@ +""" +Agent V2: reward, metrics, and rollout class. + +The generate function is provided by: + miles.rollout.generate_hub.agentic_tool_call.generate +with --custom-agent-function-path pointing to swe_agent_function.run + +Task-type agnostic — reward is pre-computed by the Harbor environment +and stored in sample.metadata["reward"] regardless of task type. + +Dynamic filter uses the general-purpose ``check_no_aborted`` from +``miles.rollout.filter_hub.dynamic_sampling_filters``. + +Components: + - reward_func: reads pre-computed reward from sample metadata + - aggregate_agent_metrics: aggregates agent timing/count metrics + - RolloutFn: InferenceRolloutFn subclass that logs agent metrics +""" + +import logging + +from miles.rollout.base_types import RolloutFnTrainInput, RolloutFnTrainOutput +from miles.rollout.inference_rollout.inference_rollout_common import InferenceRolloutFn +from miles.utils.types import Sample + +logger = logging.getLogger(__name__) + + +# -- Reward -- + + +async def reward_func(args, samples: Sample | list[Sample], **kwargs) -> float | list[float]: + """Reward is pre-computed by the agent environment during generate(). + + Handles both single-sample calls (from ``async_rm``) and batched calls + (from ``batched_async_rm`` when ``--custom-rm-path`` is set). + """ + if isinstance(samples, list): + return [s.metadata.get("reward", 0.0) for s in samples] + return samples.metadata.get("reward", 0.0) + + +# -- Agent Metrics Aggregation -- + + +def _collect_values(all_metrics: list[dict], key: str) -> list[float]: + return [m.get(key, 0) for m in all_metrics] + + +def _agg_mean(metrics: dict, all_metrics: list[dict], keys: list[str], prefix: str = "agent/", suffix: str = "_mean"): + for key in keys: + values = _collect_values(all_metrics, key) + if values: + metrics[f"{prefix}{key}{suffix}"] = sum(values) / len(values) + + +def aggregate_agent_metrics(samples: list[Sample]) -> dict: + """Aggregate agent metrics across samples for logging.""" + all_metrics = [ + s.metadata.get("agent_metrics", {}) + for s in samples + if hasattr(s, "metadata") and s.metadata and s.metadata.get("agent_metrics") + ] + if not all_metrics: + return {} + + metrics = {} + + for key in ["turns", "tool_calls"]: + values = _collect_values(all_metrics, key) + if values: + metrics[f"agent/{key}_mean"] = sum(values) / len(values) + metrics[f"agent/{key}_sum"] = sum(values) + + _agg_mean(metrics, all_metrics, ["model_query_time_sum", "env_execution_time_sum", "eval_time", "agent_run_time"]) + _agg_mean(metrics, all_metrics, ["time_per_turn", "model_query_time_avg", "env_execution_time_avg"], suffix="") + _agg_mean(metrics, all_metrics, ["model_time_ratio", "env_time_ratio", "eval_time_ratio"], suffix="") + + values = _collect_values(all_metrics, "total_time") + if values: + metrics["agent/total_time_mean"] = sum(values) / len(values) + metrics["agent/total_time_max"] = max(values) + metrics["agent/total_time_min"] = min(values) + + return metrics + + +# -- Rollout Function -- + + +class RolloutFn(InferenceRolloutFn): + """Rollout function with agent metrics aggregation.""" + + async def _call_train(self, input: RolloutFnTrainInput) -> RolloutFnTrainOutput: + output = await super()._call_train(input) + + all_samples = [] + for group in output.samples: + if isinstance(group, list): + all_samples.extend(group) + else: + all_samples.append(group) + + agent_metrics = aggregate_agent_metrics(all_samples) + if agent_metrics: + metrics = output.metrics or {} + metrics.update(agent_metrics) + output.metrics = metrics + logger.info(f"Agent metrics for rollout {input.rollout_id}: {agent_metrics}") + + return output diff --git a/examples/experimental/swe-agent-v2-amd/run-qwen3-swe.py b/examples/experimental/swe-agent-v2-amd/run-qwen3-swe.py new file mode 100644 index 0000000000..1907ca0d59 --- /dev/null +++ b/examples/experimental/swe-agent-v2-amd/run-qwen3-swe.py @@ -0,0 +1,74 @@ +"""Agent V2 launcher (Qwen3 / Qwen3-Coder): Miles <-> Harbor agent orchestration. + +Thin wrapper over ``run.py`` — reuses its ``cleanup()`` / ``prepare()`` / +``execute()`` and only overrides the model identity and the SGLang/TITO +parsers, which are the sole differences from the GLM-4.7-Flash default. + +Usage: + # Qwen3-30B-A3B (general model) + python run-qwen3-swe.py --prompt-data /root/swe_train.jsonl + + # Qwen3-Coder-30B-A3B-Instruct (coding-specialised; recommended base for + # SWE-bench so rollouts produce non-zero reward variance) + python run-qwen3-swe.py --coder --prompt-data /root/swe_train.jsonl + + # quick pipeline check (rollout only, no weight updates) + python run-qwen3-swe.py --coder --mode debug_rollout_only +""" + +import os +from dataclasses import dataclass + +import typer + +# run.py lives in the same directory (added to PYTHONPATH by execute()). +from run import ScriptArgs as BaseScriptArgs +from run import cleanup, execute, prepare + +import miles.utils.external_utils.command_utils as U + + +@dataclass +class ScriptArgs(BaseScriptArgs): + # Model identity (Qwen3-30B-A3B). ``--coder`` swaps these to the coder below. + megatron_model_type: str = "qwen3-30B-A3B" + model_name: str = "Qwen3-30B-A3B" + hf_checkpoint: str = "Qwen/Qwen3-30B-A3B" + ref_load: str = "/root/Qwen3-30B-A3B_torch_dist" + save_dir: str = "/root/Qwen3-30B-A3B_agent_v2/" + + # SGLang / TITO parsers for Qwen3 (GLM defaults are glm47 / glm45 / glm47). + sglang_tool_call_parser: str = "qwen25" + sglang_reasoning_parser: str = "qwen3" + tito_model: str = "qwen3" + + wandb_project: str = os.environ.get("WANDB_PROJECT", "qwen3-swe-agentic") + wandb_run_name: str = "qwen3-swe-tito" + prometheus_run_name: str = "qwen3-swe-tito" + + # Use Qwen3-Coder-30B-A3B-Instruct instead of the general model. The coder + # is arch-identical (same converter/model script) but uses rope_theta=1e7 + # (vs 1e6), so we override MODEL_ARGS_ROTARY_BASE for the Megatron side; + # SGLang picks up the correct value from the HF config automatically. + coder: bool = False + + +@U.dataclass_cli +def main(args: ScriptArgs): + if args.coder: + args.model_name = "Qwen3-Coder-30B-A3B-Instruct" + args.hf_checkpoint = "Qwen/Qwen3-Coder-30B-A3B-Instruct" + args.ref_load = "/root/Qwen3-Coder-30B-A3B-Instruct_torch_dist" + args.save_dir = "/root/Qwen3-Coder-30B-A3B-Instruct_agent_v2/" + # Read by scripts/models/qwen3-30B-A3B.sh (${MODEL_ARGS_ROTARY_BASE:-1000000}) + # when the model script is sourced for both conversion and training. + os.environ["MODEL_ARGS_ROTARY_BASE"] = "10000000" + + cleanup() + if not args.skip_prepare: + prepare(args) + execute(args) + + +if __name__ == "__main__": + typer.run(main) diff --git a/examples/experimental/swe-agent-v2-amd/run.py b/examples/experimental/swe-agent-v2-amd/run.py new file mode 100644 index 0000000000..e7e1b86bd9 --- /dev/null +++ b/examples/experimental/swe-agent-v2-amd/run.py @@ -0,0 +1,274 @@ +"""Agent V2 launcher (GLM-4.7-Flash): Miles <-> Harbor agent orchestration. + +Supports any task type (SWE-bench, Terminal-Bench, custom) via Harbor. + +Python equivalent of run.sh. Usage: + python run.py + python run.py --mode normal + python run.py --base-dir /my/models --prompt-data /my/data.jsonl +""" + +import os +import socket +import subprocess +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +import typer + +import miles.utils.external_utils.command_utils as U + +SCRIPT_DIR = Path(__file__).resolve().parent + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + mode: Literal["normal", "debug_rollout_only"] = "normal" + run_id: str = U.create_run_id() + megatron_model_type: str = "glm4.7-flash" + num_gpus_per_node: int = 8 + megatron_path: str = "/root/Megatron-LM" + + # Paths + skip_prepare: bool = False + base_dir: str = "/root" + model_name: str = "GLM-4.7-Flash" + hf_checkpoint: str = "zai-org/GLM-4.7-Flash" + ref_load: str = "/root/GLM-4.7-Flash_torch_dist" + save_dir: str = "/root/GLM-4.7-Flash_agent_v2/" + prompt_data: str = "/root/swe_train.jsonl" + + # Training settings + max_seq_len: int = 16384 + rollout_batch_size: int = 2 + n_samples_per_prompt: int = 4 + global_batch_size: int = 8 + + # SGLang / TITO parsers (model-family specific; GLM-4.7-Flash defaults). + # Subclasses (e.g. run-qwen3-swe.py) override these for other models. + sglang_tool_call_parser: str = "glm47" + sglang_reasoning_parser: str = "glm45" + tito_model: str = "glm47" + + # Agent settings + agent_server_url: str = os.environ.get( + "AGENT_SERVER_URL", os.environ.get("SWE_AGENT_URL", "http://agent_env:11000") + ) + agent_model_name: str = os.environ.get("AGENT_MODEL_NAME", "model") + harbor_tasks_dir: str = os.environ.get("HARBOR_TASKS_DIR", "/root/harbor_tasks") + router_external_host: str = os.environ.get("MILES_ROUTER_EXTERNAL_HOST", socket.gethostname()) # public IP + miles_host_ip: str = os.environ.get("MILES_HOST_IP", socket.gethostname()) # cluster/pod IP + + # W&B settings + wandb_key: str = os.environ.get("WANDB_KEY", os.environ.get("WANDB_API_KEY", "")) + wandb_project: str = os.environ.get("WANDB_PROJECT", "glm47-flash-agentic") + wandb_team: str = os.environ.get("WANDB_TEAM", "") + wandb_run_name: str = "glm47-flash-swe-tito" + + # Prometheus settings + use_prometheus: bool = True + prometheus_port: int = 9090 + prometheus_run_name: str = "glm47-flash-swe-tito" + + +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( + model_name=args.model_name, + megatron_model_type=args.megatron_model_type, + num_gpus_per_node=args.num_gpus_per_node, + dir_dst=args.base_dir, + hf_checkpoint=args.hf_checkpoint, + megatron_path=args.megatron_path, + ) + + +def execute(args: ScriptArgs): + ckpt_args = ( + f"--hf-checkpoint {args.hf_checkpoint} " + f"--ref-load {args.ref_load} " + f"--save {args.save_dir} " + "--save-interval 100 " + ) + + rollout_args = ( + f"--prompt-data {args.prompt_data} " + "--input-key prompt " + "--metadata-key metadata " + "--rollout-shuffle " + "--num-rollout 3000 " + 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 " + ) + + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 16384 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--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 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 1 " + "--sglang-mem-fraction-static 0.7 " + f"--sglang-tool-call-parser {args.sglang_tool_call_parser} " + f"--sglang-reasoning-parser {args.sglang_reasoning_parser} " + "--use-miles-router " + "--sglang-router-port 31000 " + # TODO: speculative decoding has issue, need to fix later + ) + + agent_args = ( + "--custom-generate-function-path miles.rollout.generate_hub.agentic_tool_call.generate " + "--custom-agent-function-path swe_agent_function.run " + "--custom-rm-path generate.reward_func " + "--rollout-function-path generate.RolloutFn " + "--dynamic-sampling-filter-path miles.rollout.filter_hub.dynamic_sampling_filters.check_no_aborted " + f"--tito-model {args.tito_model} " + "--use-session-server " + "--session-server-port 30000 " + # This is required by terminus-2 harness + "--tito-allowed-append-roles user tool " + ) + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--colocate " + f"--actor-num-nodes {args.num_nodes} " + f"--actor-num-gpus-per-node {args.num_gpus_per_node} " + f"--rollout-num-gpus {args.num_gpus_per_node} " + ) + + debug_args = "--debug-rollout-only " if args.mode == "debug_rollout_only" 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} " + ) + + train_args = ( + f"{ckpt_args}" + f"{rollout_args}" + f"{optimizer_args}" + f"{grpo_args}" + f"{wandb_args}" + f"{prometheus_args}" + f"{perf_args}" + f"{sglang_args}" + f"{agent_args}" + f"{misc_args}" + f"{debug_args}" + ) + + miles_root = U.repo_base_dir + + extra_env_vars = { + "PYTHONPATH": f"{args.megatron_path}:{SCRIPT_DIR}:{miles_root}", + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "AGENT_SERVER_URL": args.agent_server_url, + "AGENT_MODEL_NAME": args.agent_model_name, + "MILES_ROUTER_EXTERNAL_HOST": args.router_external_host, + "HARBOR_TASKS_DIR": args.harbor_tasks_dir, + "MILES_HOST_IP": args.miles_host_ip, + } + + # On ROCm/AMD, Ray blanks the Ray-job driver's HIP_VISIBLE_DEVICES, which + # makes SGLang's import-time GPU probe fail ("No HIP GPUs are available"). + # Tell Ray not to touch device visibility so miles manages placement itself. + # No-op on NVIDIA (``torch.version.hip`` is None) — CUDA device management is + # left entirely to Ray there, so this does not change NVIDIA behaviour. + import torch + + if getattr(torch.version, "hip", None): + extra_env_vars["RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES"] = "1" + + U.execute_train( + train_args=train_args, + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type=args.megatron_model_type, + megatron_path=args.megatron_path, + extra_env_vars=extra_env_vars, + ) + + +@U.dataclass_cli +def main(args: ScriptArgs): + cleanup() + if not args.skip_prepare: + prepare(args) + execute(args) + + +if __name__ == "__main__": + typer.run(main) diff --git a/examples/experimental/swe-agent-v2-amd/server.py b/examples/experimental/swe-agent-v2-amd/server.py new file mode 100644 index 0000000000..fb2ffe224e --- /dev/null +++ b/examples/experimental/swe-agent-v2-amd/server.py @@ -0,0 +1,352 @@ +""" +FastAPI server wrapping Harbor for generalized agent-environment orchestration. + +Provides a single ``/run`` endpoint that handles any task type (SWE-bench, +Terminal-Bench, custom datasets, etc.) through Harbor's unified Trial API. +Harbor handles Docker orchestration, agent execution, and grading — the +server is task-type agnostic. + +Requires: + - Harbor installed: pip install harbor-framework + - Prepared task dirs under HARBOR_TASKS_DIR (via adapters or prepare_harbor_tasks.py) + +Usage: + python server.py --port 11000 --max-concurrent 8 +""" + +import argparse +import asyncio +import logging +import os +import re +import traceback +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +import uvicorn +from fastapi import FastAPI +from pydantic import BaseModel + +logger = logging.getLogger(__name__) + + +_semaphore: asyncio.Semaphore | None = None + + +@asynccontextmanager +async def _lifespan(app: FastAPI) -> AsyncIterator[None]: + global _semaphore + max_concurrent = int(os.getenv("AGENT_MAX_CONCURRENT", os.getenv("SWE_AGENT_MAX_CONCURRENT", "8"))) + _semaphore = asyncio.Semaphore(max_concurrent) + logger.info(f"Initialized semaphore with max_concurrent={max_concurrent}") + yield + + +app = FastAPI(title="Agent Environment Server (Harbor)", lifespan=_lifespan) + + +class RunRequest(BaseModel): + base_url: str + model: str + sampling_params: dict[str, Any] = {} + api_key: str = "dummy" + + instance_id: str = "" + agent_name: str = "mini-swe-agent" + max_seq_len: int | None = None + + model_config = {"extra": "allow"} + + +class RunResponse(BaseModel): + reward: float = 0.0 + exit_status: str = "" + agent_metrics: dict[str, Any] = {} + eval_report: dict[str, Any] = {} + + +def get_semaphore() -> asyncio.Semaphore: + assert _semaphore is not None, "Semaphore not initialized — server not started?" + return _semaphore + + +_TIMEOUT_EXCEPTIONS = {"AgentTimeoutError", "VerifierTimeoutError", "EnvironmentStartTimeoutError"} +_OUTPUT_LIMIT_EXCEPTIONS = {"MaxSeqLenExceededError"} + +_HOST_PROCESS_AGENTS = {"terminus-2", "terminus-1", "terminus"} + +_SAFE_INSTANCE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def _extract_exit_status(result) -> str: + """Derive exit status from Harbor TrialResult.""" + exc = getattr(result, "exception_info", None) + if exc is not None: + exc_type = getattr(exc, "exception_type", "") + if exc_type in _TIMEOUT_EXCEPTIONS: + return "TimeLimitExceeded" + if exc_type in _OUTPUT_LIMIT_EXCEPTIONS: + return "SequenceLengthLimitExceeded" + return "AgentError" + if getattr(result, "verifier_result", None) is not None: + return "Submitted" + return "Unknown" + + +def _timing_duration_sec(timing) -> float | None: + started = getattr(timing, "started_at", None) + finished = getattr(timing, "finished_at", None) + if started and finished: + return (finished - started).total_seconds() + return None + + +def _extract_reward(result) -> tuple[float, dict[str, Any]]: + """Extract scalar reward and full eval report from Harbor TrialResult. + + Looks for the ``"reward"`` key first, then falls back to the first value + in the rewards dict. Works with both ``reward.txt`` and ``reward.json``. + """ + vr = getattr(result, "verifier_result", None) + if vr is None: + return 0.0, {} + rewards = getattr(vr, "rewards", None) or {} + reward_val = rewards.get("reward") + if reward_val is None: + reward_val = next(iter(rewards.values()), 0.0) + try: + reward = float(reward_val) if reward_val is not None else 0.0 + except (TypeError, ValueError): + reward = 0.0 + return reward, dict(rewards) + + +def _extract_metrics(result) -> dict[str, Any]: + """Extract agent metrics from Harbor TrialResult.""" + metrics: dict[str, Any] = {} + try: + ar = getattr(result, "agent_result", None) + if ar is not None: + for field in ("n_input_tokens", "n_output_tokens", "cost_usd"): + val = getattr(ar, field, None) + if val is not None: + metrics[field] = val + agent_meta = getattr(ar, "metadata", None) + if isinstance(agent_meta, dict): + metrics.update(agent_meta) + + # Turn / tool-call counts, so miles' aggregate_agent_metrics (which + # reads the "turns" / "tool_calls" keys) can populate agent/turns_mean + # etc. Derived from the ATIF trajectory's agent-authored steps. + # Best-effort: never let a missing/renamed field fail the trial. + traj = getattr(ar, "trajectory", None) + steps = getattr(traj, "steps", None) if traj is not None else None + if steps: + agent_steps = [s for s in steps if getattr(s, "source", None) == "agent"] + if agent_steps: + metrics["turns"] = len(agent_steps) + metrics["tool_calls"] = sum(len(getattr(s, "tool_calls", None) or []) for s in agent_steps) + + agent_timing = getattr(result, "agent_execution", None) + if agent_timing is not None: + dur = _timing_duration_sec(agent_timing) + if dur is not None: + metrics["agent_run_time"] = dur + + verifier_timing = getattr(result, "verifier", None) + if verifier_timing is not None: + dur = _timing_duration_sec(verifier_timing) + if dur is not None: + metrics["eval_time"] = dur + except Exception as e: + logger.warning(f"Failed to extract metrics: {e}", exc_info=True) + return metrics + + +def _error_response(exit_status: str) -> dict[str, Any]: + return {"reward": 0.0, "exit_status": exit_status, "agent_metrics": {}, "eval_report": {}} + + +async def _run_trial(request: RunRequest) -> dict[str, Any]: + """Run a Harbor trial for a single task instance. + + Task-type agnostic — all differentiation (environment, grading harness) + is encoded in the Harbor task directory's 4 files. + """ + try: + from harbor.models.trial.config import AgentConfig, EnvironmentConfig, TaskConfig, TrialConfig + + # Harbor v0.13.x: ``Trial`` is an abstract base; use the concrete + # SingleStepTrial and the async ``create`` factory instead of direct + # construction. + from harbor.trial.single_step import SingleStepTrial + except ImportError: + logger.error("Harbor not installed. Install with: pip install harbor") + return _error_response("ImportError") + + try: + tasks_dir = Path( + os.getenv("HARBOR_TASKS_DIR", "/root/harbor_tasks"), + ).resolve() + + if not request.instance_id: + logger.error("Empty instance_id") + return _error_response("InvalidInstanceId") + + raw_id = request.instance_id + if not _SAFE_INSTANCE_ID.match(raw_id): + logger.error(f"Invalid instance_id rejected: {raw_id!r}") + return _error_response("InvalidInstanceId") + + # Resolve and verify the path stays within tasks_dir. Path.is_relative_to + # is used instead of a string startswith() check, which is vulnerable to a + # partial-prefix bypass (e.g. "/root/harbor_tasks_secret" starts with + # "/root/harbor_tasks"). + def _resolve_within(name: str) -> Path | None: + candidate = (tasks_dir / name).resolve() + return candidate if candidate.is_relative_to(tasks_dir) else None + + task_path = _resolve_within(raw_id) + if task_path is None: + logger.error(f"Path traversal blocked: {raw_id!r}") + return _error_response("InvalidInstanceId") + + # Case-insensitive fallback: the SWE-Gym adapter lowercases task dir names + # (Docker repo names must be lowercase), but dataset instance_ids keep the + # original case (e.g. Project-MONAI__MONAI-1030). Try the lowercased dir if + # the exact-case one is absent so those instances aren't TaskNotFound. + if not task_path.exists() and raw_id != raw_id.lower(): + lower_path = _resolve_within(raw_id.lower()) + if lower_path is not None and lower_path.exists(): + logger.info(f"Resolved {raw_id!r} -> lowercase task dir {raw_id.lower()!r}") + task_path = lower_path + + if not task_path.exists(): + logger.error(f"Task directory not found: {task_path}") + return _error_response("TaskNotFound") + agent_kwargs: dict[str, Any] = {} + agent_env: dict[str, str] = {} + + is_host_agent = request.agent_name in _HOST_PROCESS_AGENTS + + if "hosted_vllm" in request.model or "openai" in request.model: + agent_kwargs["model_info"] = { + "max_input_tokens": int(os.getenv("AGENT_MAX_INPUT_TOKENS", "32768")), + "max_output_tokens": int(os.getenv("AGENT_MAX_OUTPUT_TOKENS", "8192")), + "input_cost_per_token": 0.0, + "output_cost_per_token": 0.0, + } + + if request.max_seq_len is not None: + agent_kwargs["max_seq_len"] = request.max_seq_len + + # Optional: point mini-swe-agent at a specific config file (e.g. the + # merged swebench_local.yaml with environment_class=local). Gowtham 8b + # equivalent, routed through the Harbor adapter kwargs. + config_file = os.getenv("MSWEA_CONFIG_FILE") + if config_file: + agent_kwargs["config_file"] = config_file + + if is_host_agent: + agent_kwargs["api_base"] = request.base_url + agent_kwargs["api_key"] = request.api_key or "dummy" + agent_kwargs["enable_summarize"] = False + agent_env = { + "OPENAI_API_KEY": request.api_key or "dummy", + "OPENAI_API_BASE": request.base_url, + } + else: + agent_env = { + "OPENAI_API_BASE": request.base_url, + "OPENAI_API_KEY": request.api_key, + "HOSTED_VLLM_API_BASE": request.base_url, + "HOSTED_VLLM_API_KEY": request.api_key, + "MSWEA_COST_TRACKING": "ignore_errors", + } + + env_config_kwargs: dict[str, Any] = { + "type": "docker", + "delete": os.getenv("HARBOR_DELETE_CONTAINERS", "false").lower() in ("true", "1", "t"), + } + # Optional: extra docker-compose override(s) so Harbor task containers + # join swe-net (and can reach the Miles router/session server). + extra_compose_env = os.getenv("HARBOR_EXTRA_DOCKER_COMPOSE") + if extra_compose_env: + extra_compose = [Path(p) for p in extra_compose_env.split(os.pathsep) if p] + if extra_compose: + env_config_kwargs["extra_docker_compose"] = extra_compose + + config = TrialConfig( + task=TaskConfig(path=task_path), + agent=AgentConfig( + name=request.agent_name, + model_name=request.model, + env=agent_env, + kwargs=agent_kwargs, + ), + environment=EnvironmentConfig(**env_config_kwargs), + # Trials must live on a path that resolves identically inside agent_env + # and on the host, else the host Docker daemon can't bind-mount the + # verifier/agent dirs (DinD path mismatch). Default to the identity mount. + trials_dir=Path(os.environ.get("HARBOR_TRIALS_DIR", "trials")), + ) + + trial = await SingleStepTrial.create(config) + result = await trial.run() + + reward, eval_report = _extract_reward(result) + exit_status = _extract_exit_status(result) + agent_metrics = _extract_metrics(result) + + return { + "reward": reward, + "exit_status": exit_status, + "agent_metrics": agent_metrics, + "eval_report": eval_report, + } + + except Exception as e: + logger.error(f"Harbor trial failed: {e}\n{traceback.format_exc()}") + return _error_response(f"Error: {type(e).__name__}") + + +@app.post("/run") +async def run_instance(request: RunRequest) -> RunResponse: + """Run an agent on a single task instance via Harbor.""" + logger.info(f"Running instance: {request.instance_id}") + async with get_semaphore(): + result = await _run_trial(request) + logger.info( + f"Instance {request.instance_id} finished: exit_status={result['exit_status']}, reward={result['reward']}" + ) + return RunResponse(**result) + + +@app.get("/health") +async def health(): + return {"status": "ok"} + + +def main(): + parser = argparse.ArgumentParser(description="Agent Environment Server (Harbor)") + parser.add_argument("--host", type=str, default="0.0.0.0") + parser.add_argument("--port", type=int, default=11000) + parser.add_argument("--max-concurrent", type=int, default=8) + args = parser.parse_args() + + os.environ["AGENT_MAX_CONCURRENT"] = str(args.max_concurrent) + + os.environ.setdefault("MSWEA_API_KEY", "dummy") + os.environ.setdefault("HOSTED_VLLM_API_KEY", "dummy") + + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(name)s %(levelname)s %(message)s", + ) + uvicorn.run(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/examples/experimental/swe-agent-v2-amd/swe_agent_function.py b/examples/experimental/swe-agent-v2-amd/swe_agent_function.py new file mode 100644 index 0000000000..d1460fbe06 --- /dev/null +++ b/examples/experimental/swe-agent-v2-amd/swe_agent_function.py @@ -0,0 +1,94 @@ +""" +Custom agent function for agentic_tool_call.generate. + +Dispatches to a Harbor-based agent server and returns env metadata +as a plain dict. The generate layer merges this into sample.metadata so +downstream reward models (--custom-rm-path) can extract reward, eval +reports, etc. + +Task-type agnostic — the server + Harbor task directory handle all +differentiation (environment, grading harness, agent selection). +""" + +import asyncio +import logging +import os +from typing import Any +from urllib.parse import urlparse, urlsplit, urlunparse + +from miles.utils.http_utils import post + +logger = logging.getLogger(__name__) + + +async def run( + base_url: str, + prompt: Any, + request_kwargs: dict[str, Any] | None = None, + metadata: dict[str, Any] | None = None, + **kwargs, +) -> dict[str, Any] | None: + """Run a single task instance via the Harbor agent server.""" + metadata = metadata or {} + request_kwargs = request_kwargs or {} + + agent_server_url = os.getenv( + "AGENT_SERVER_URL", + os.getenv("SWE_AGENT_URL", "http://localhost:11000"), + ) + model_name = os.getenv( + "AGENT_MODEL_NAME", + os.getenv("SWE_AGENT_MODEL_NAME", "model"), + ) + + session_url = f"{base_url}/v1" + external_host = os.getenv("MILES_ROUTER_EXTERNAL_HOST") + if external_host: + parsed = urlparse(session_url) + port = parsed.port + netloc = f"{external_host}:{port}" if port else external_host + session_url = urlunparse(parsed._replace(netloc=netloc)) + + request: dict[str, Any] = { + **metadata, + "base_url": session_url, + "model": f"openai/{model_name}", + "sampling_params": request_kwargs, + } + + max_seq_len = metadata.get("max_seq_len") + if max_seq_len is not None: + request["max_seq_len"] = int(max_seq_len) + + session_server_id = metadata.get("session_server_id") + if session_server_id is not None: + if external_host: + port = urlsplit(f"http://{session_server_id}").port + session_server_id = f"{external_host}:{port}" + request["session_server_id"] = session_server_id + + session_server_instance_id = metadata.get("session_server_instance_id") + if session_server_instance_id is not None: + request["session_server_instance_id"] = session_server_instance_id + + try: + response = await asyncio.wait_for( + post(f"{agent_server_url}/run", request), + timeout=3600, # 1 hour max per trial + ) + except asyncio.TimeoutError: + logger.error("Agent server call timed out after 3600s") + return None + except asyncio.CancelledError: + logger.warning("Agent server call cancelled (sibling task failure?)") + return None + except Exception as e: + logger.error(f"Agent server call failed: {e}") + return None + + return { + "reward": response.get("reward", 0.0), + "exit_status": response.get("exit_status", ""), + "eval_report": response.get("eval_report", {}), + "agent_metrics": response.get("agent_metrics", {}), + }