diff --git a/.gitignore b/.gitignore index c534206037..35ab4a1f0b 100644 --- a/.gitignore +++ b/.gitignore @@ -115,6 +115,13 @@ outputs/ # Local artifacts tinker.db +# Large local Toolathlon run data. Ray's runtime-environment packager honors +# .gitignore, so these must stay out of the distributed working-directory zip. +/jobs/ +/toolathlon-tasks/ +/toolathlon-tasks.zip +/data/toolathlon-harbor/ + # Alembic - don't track pycache tx/tinker/alembic/__pycache__/ diff --git a/examples/train/toolathlon_harbor/README.md b/examples/train/toolathlon_harbor/README.md index 91566a22f2..f57f0f4696 100644 --- a/examples/train/toolathlon_harbor/README.md +++ b/examples/train/toolathlon_harbor/README.md @@ -18,14 +18,14 @@ export TOOLATHLON_API_KEY=dummy export TOOLATHLON_MODEL=Qwen3-8B bash examples/train/toolathlon_harbor/run_eval.sh \ - -i 401k-watchlist-recency-window-refresh + -i q3-eta-leaderboard-definition-aware-reconcile ``` `run_eval.sh` uses Harbor's local Docker environment by default. Use the separate `run_compute_eval.sh` wrapper only when task containers should run on AfterQuery Compute. -The default task root is `toolathlon-tasks/tasks`. Override it with +The default evaluation root is `toolathlon-tasks/eval_tasks`, which contains the 100 held-out tasks listed in `toolathlon-tasks/samples/pass-at-4-100.txt`. Override it with `TOOLATHLON_TASKS_DIR`. The launcher builds `toolathlon-json-runtime:v1` from the bundled runtime archive as `linux/amd64` when the image is missing. Additional arguments are forwarded to `harbor run`, so omitting `-i` runs the @@ -56,7 +56,7 @@ echo "$TOKEN" | docker login us-docker.pkg.dev \ -u oauth2accesstoken --password-stdin bash examples/train/toolathlon_harbor/run_compute_eval.sh \ - -i 401k-watchlist-recency-window-refresh + -i q3-eta-leaderboard-definition-aware-reconcile ``` `run_compute_eval.sh` exits immediately with an explanatory error when @@ -74,8 +74,8 @@ The reference and no-op paths do not call a model: ```bash cd toolathlon-tasks -harbor run --path tasks -i 401k-watchlist-recency-window-refresh --agent oracle -harbor run --path tasks -i 401k-watchlist-recency-window-refresh --agent nop +harbor run --path eval_tasks -i q3-eta-leaderboard-definition-aware-reconcile --agent oracle +harbor run --path eval_tasks -i q3-eta-leaderboard-definition-aware-reconcile --agent nop ``` Expect rewards `1.0` and `0.0`, respectively. @@ -87,8 +87,8 @@ Registry as above and run: COMPUTE_API_URL=${COMPUTE_API_URL:-https://compute-api.afterquery.com} \ COMPUTE_IMAGE_REGISTRY=${COMPUTE_IMAGE_REGISTRY:-us-docker.pkg.dev/afterquery-compute/compute-images/} \ uv run --extra harbor harbor run \ - --path toolathlon-tasks/tasks \ - -i 401k-watchlist-recency-window-refresh \ + --path toolathlon-tasks/eval_tasks \ + -i q3-eta-leaderboard-definition-aware-reconcile \ --agent oracle \ --env compute \ --n-concurrent 1 @@ -99,6 +99,39 @@ replay, and verification without making a model request. ## SkyRL generation and training +### Shared Daytona runtime (no local Docker required) + +For Daytona evaluation or GRPO, build the runtime once as a named snapshot on +Daytona itself. The build happens remotely and does not require a Docker daemon +on the SkyRL host: + +```bash +export DAYTONA_API_KEY=... +uv run --extra harbor python \ + examples/train/toolathlon_harbor/create_daytona_runtime_snapshot.py \ + --archive toolathlon-tasks/runtime/toolathlon-json-runtime-src.tar.gz \ + --name toolathlon-json-runtime-v1 +``` + +Stage the task dataset once on storage visible to the Ray workers. Staged tasks +contain `environment/task/` and `environment/mcp.json`, but no Dockerfile. Their +`task.toml` selects the shared runtime image, `/opt` upload workdir, and the +task-specific `T3_SERVERS` value: + +```bash +uv run python examples/train/toolathlon_harbor/prepare_daytona_tasks.py \ + --source toolathlon-tasks/tasks \ + --output "$HOME/data/toolathlon-harbor/tasks" \ + --runtime-image us-docker.pkg.dev/afterquery-compute/compute-images/toolathlon-json-runtime:v1 +``` + +The named Daytona snapshot is the normal startup path. The registry image is a +fallback and must be built and pushed by CI or another Docker-capable machine; +this host does not need to build it. The two-node launcher stages the dataset +automatically when `TOOLATHLON_TASKS_DIR` is absent. Override +`TOOLATHLON_RUNTIME_IMAGE` or `DAYTONA_SNAPSHOT_TEMPLATE` to select another +version. + Use the existing Harbor entrypoint with this adapter's trial configuration and the task directory as the dataset: diff --git a/examples/train/toolathlon_harbor/create_daytona_runtime_snapshot.py b/examples/train/toolathlon_harbor/create_daytona_runtime_snapshot.py new file mode 100644 index 0000000000..070368171b --- /dev/null +++ b/examples/train/toolathlon_harbor/create_daytona_runtime_snapshot.py @@ -0,0 +1,69 @@ +"""Build the shared Toolathlon runtime remotely as a named Daytona snapshot.""" + +from __future__ import annotations + +import argparse +import asyncio +import tarfile +import tempfile +from pathlib import Path + + +async def create_snapshot(archive: Path, name: str, target: str | None) -> None: + from daytona import ( + AsyncDaytona, + CreateSnapshotParams, + DaytonaConfig, + Image, + Resources, + ) + + config = DaytonaConfig(target=target) if target else None + daytona = AsyncDaytona(config) if config else AsyncDaytona() + try: + try: + existing = await daytona.snapshot.get(name) + except Exception as error: + if ( + "not found" not in str(error).lower() + and "notfound" not in type(error).__name__.lower() + ): + raise + else: + print( + f"Snapshot {name!r} already exists in state {existing.state}; nothing to do" + ) + return + + with tempfile.TemporaryDirectory(prefix="toolathlon-runtime-") as temporary: + root = Path(temporary) + with tarfile.open(archive, "r:gz") as bundle: + bundle.extractall(root, filter="data") + dockerfiles = list(root.glob("*/Dockerfile")) + if len(dockerfiles) != 1: + raise RuntimeError( + f"expected one runtime Dockerfile in {archive}, found {len(dockerfiles)}" + ) + await daytona.snapshot.create( + CreateSnapshotParams( + name=name, + image=Image.from_dockerfile(str(dockerfiles[0])), + resources=Resources(cpu=2, memory=4, disk=20), + ) + ) + print(f"Submitted remote Daytona build for snapshot {name!r}") + finally: + await daytona.close() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--archive", type=Path, required=True) + parser.add_argument("--name", required=True) + parser.add_argument("--target", help="optional Daytona target, for example us") + args = parser.parse_args() + asyncio.run(create_snapshot(args.archive.resolve(), args.name, args.target)) + + +if __name__ == "__main__": + main() diff --git a/examples/train/toolathlon_harbor/harbor_compute_trial_config.yaml b/examples/train/toolathlon_harbor/harbor_compute_trial_config.yaml index f53e8f217b..5fd452f942 100644 --- a/examples/train/toolathlon_harbor/harbor_compute_trial_config.yaml +++ b/examples/train/toolathlon_harbor/harbor_compute_trial_config.yaml @@ -18,6 +18,8 @@ agent: deadline_sec: 7000 request_timeout_sec: 900 max_tokens: 8192 + max_context_tokens: 131072 + context_manager_factory: examples.train.toolathlon_harbor.toolathlon_context_manager:create_managed_context temperature: 1.0 collect_rollout_details: true strict_rollout_details: true diff --git a/examples/train/toolathlon_harbor/harbor_daytona_training_config.yaml b/examples/train/toolathlon_harbor/harbor_daytona_training_config.yaml new file mode 100644 index 0000000000..77385856cc --- /dev/null +++ b/examples/train/toolathlon_harbor/harbor_daytona_training_config.yaml @@ -0,0 +1,42 @@ +# Harbor TrialConfig for context-managed Toolathlon GRPO rollouts on Daytona. +# SkyRL injects task.path, agent.model_name, api_base, and session_id per trial. + +trials_dir: ~/toolathlon_grpo_qwen38_27b/trials + +agent: + name: null + import_path: examples.train_integrations.harbor.mcp_agent:HarborMCPAgent + override_timeout_sec: 7200 + mcp_servers: + - name: toolathlon + transport: stdio + command: python + args: ["-m", "t3.mcp_server"] + env: + OPENAI_API_KEY: "${TOOLATHLON_API_KEY:-dummy}" + kwargs: + max_turns: 64 + deadline_sec: 7000 + request_timeout_sec: 900 + max_tokens: 8192 + max_context_tokens: 131072 + context_manager_factory: examples.train.toolathlon_harbor.toolathlon_context_manager:create_managed_context + temperature: 0.7 + collect_rollout_details: true + strict_rollout_details: true + +environment: + type: daytona + force_build: false + delete: true + override_cpus: null + override_memory_mb: null + override_storage_mb: null + kwargs: + # One runtime snapshot is built remotely; task files are uploaded into /opt. + snapshot_template_name: toolathlon-json-runtime-v1 + auto_snapshot: false + auto_stop_interval_mins: 30 + +verifier: + disable: false diff --git a/examples/train/toolathlon_harbor/harbor_trial_config.yaml b/examples/train/toolathlon_harbor/harbor_trial_config.yaml index 14e59bbee9..01c005997e 100644 --- a/examples/train/toolathlon_harbor/harbor_trial_config.yaml +++ b/examples/train/toolathlon_harbor/harbor_trial_config.yaml @@ -20,6 +20,8 @@ agent: deadline_sec: 7000 request_timeout_sec: 900 max_tokens: 8192 + max_context_tokens: 131072 + context_manager_factory: examples.train.toolathlon_harbor.toolathlon_context_manager:create_managed_context temperature: 1.0 collect_rollout_details: true strict_rollout_details: true diff --git a/examples/train/toolathlon_harbor/launch_vllm_slurm.sh b/examples/train/toolathlon_harbor/launch_vllm_slurm.sh new file mode 100755 index 0000000000..dc073015ab --- /dev/null +++ b/examples/train/toolathlon_harbor/launch_vllm_slurm.sh @@ -0,0 +1,222 @@ +#!/usr/bin/env bash +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +LOCAL_REPO_ROOT="$(cd "$HERE/../../.." && pwd)" +SHARED_REPO_ROOT="/workspace/users/${USER}/$(basename "$LOCAL_REPO_ROOT")" +if [[ -d "$SHARED_REPO_ROOT" ]]; then + REPO_ROOT=${REPO_ROOT:-$SHARED_REPO_ROOT} +else + REPO_ROOT=${REPO_ROOT:-$LOCAL_REPO_ROOT} +fi + +MODEL=${MODEL:-Qwen/Qwen3.8-27B} +SERVED_MODEL=${SERVED_MODEL:-Qwen3.8-27B} +LORA_PATH=${LORA_PATH:-} +LORA_NAME=${LORA_NAME:-} +MAX_LORA_RANK=${MAX_LORA_RANK:-32} +SERVER_COUNT=${SERVER_COUNT:-8} +GPUS=${GPUS:-8} +BASE_PORT=${BASE_PORT:-18000} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-262144} +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.9} +WORKER_START_TIMEOUT=${WORKER_START_TIMEOUT:-900} + +PARTITION=${PARTITION:-gpu} +WCKEY=${WCKEY:-afterquery_research} +EXCLUDE_NODES=${EXCLUDE_NODES:-node-2} +NODELIST=${NODELIST:-} +CPUS_PER_TASK=${CPUS_PER_TASK:-64} +MEMORY=${MEMORY:-768G} +TIME_LIMIT=${TIME_LIMIT:-12:00:00} +JOB_NAME=${JOB_NAME:-qwen38-27b-8x-262k} +LOG_ROOT=${LOG_ROOT:-} +DRY_RUN=0 + +usage() { + cat < 0)); do + case "$1" in + --model) MODEL=$2; shift 2 ;; + --served-model-name) SERVED_MODEL=$2; shift 2 ;; + --lora-path) LORA_PATH=$2; shift 2 ;; + --lora-name) LORA_NAME=$2; shift 2 ;; + --max-lora-rank) MAX_LORA_RANK=$2; shift 2 ;; + --servers) SERVER_COUNT=$2; shift 2 ;; + --gpus) GPUS=$2; shift 2 ;; + --base-port) BASE_PORT=$2; shift 2 ;; + --max-model-len) MAX_MODEL_LEN=$2; shift 2 ;; + --gpu-memory-utilization) GPU_MEMORY_UTILIZATION=$2; shift 2 ;; + --worker-start-timeout) WORKER_START_TIMEOUT=$2; shift 2 ;; + --time) TIME_LIMIT=$2; shift 2 ;; + --partition) PARTITION=$2; shift 2 ;; + --wckey) WCKEY=$2; shift 2 ;; + --exclude) EXCLUDE_NODES=$2; shift 2 ;; + --nodelist) NODELIST=$2; shift 2 ;; + --cpus) CPUS_PER_TASK=$2; shift 2 ;; + --mem) MEMORY=$2; shift 2 ;; + --job-name) JOB_NAME=$2; shift 2 ;; + --repo-root) REPO_ROOT=$2; shift 2 ;; + --log-root) LOG_ROOT=$2; shift 2 ;; + --dry-run) DRY_RUN=1; shift ;; + -h|--help) usage; exit 0 ;; + *) echo "Unknown argument: $1" >&2; usage >&2; exit 2 ;; + esac +done + +LOG_ROOT=${LOG_ROOT:-$REPO_ROOT/jobs/model_server} + +for value in "$SERVER_COUNT" "$GPUS" "$BASE_PORT" "$MAX_MODEL_LEN" "$MAX_LORA_RANK" "$WORKER_START_TIMEOUT"; do + [[ "$value" =~ ^[0-9]+$ ]] || { echo "Expected a non-negative integer, got: $value" >&2; exit 2; } +done +((SERVER_COUNT > 0)) || { echo "--servers must be greater than zero" >&2; exit 2; } +((GPUS > 0)) || { echo "--gpus must be greater than zero" >&2; exit 2; } +((SERVER_COUNT <= GPUS)) || { echo "--servers cannot exceed --gpus" >&2; exit 2; } +if [[ -n "$LORA_PATH" ]]; then + [[ -d "$LORA_PATH" ]] || { echo "LoRA adapter directory not found: $LORA_PATH" >&2; exit 2; } + [[ -n "$LORA_NAME" ]] || { echo "--lora-name is required with --lora-path" >&2; exit 2; } +elif [[ -n "$LORA_NAME" ]]; then + echo "--lora-path is required with --lora-name" >&2 + exit 2 +fi + +mkdir -p "$LOG_ROOT/slurm" "$LOG_ROOT/servers" + +# Do not use sbatch --export here. On this cluster, explicit --export settings +# have caused allocations to fail with NODE_FAIL. Slurm's default environment +# propagation carries these exported variables into the job. +export SERVER_REPO_ROOT="$REPO_ROOT" +export SERVER_LOG_ROOT="$LOG_ROOT/servers" +export MODEL SERVED_MODEL LORA_PATH LORA_NAME MAX_LORA_RANK +export SERVER_COUNT BASE_PORT MAX_MODEL_LEN +export GPU_MEMORY_UTILIZATION WORKER_START_TIMEOUT + +SBATCH_ARGS=( + --parsable + --partition="$PARTITION" + --wckey="$WCKEY" + --job-name="$JOB_NAME" + --nodes=1 + --ntasks=1 + --gres="gpu:$GPUS" + --cpus-per-task="$CPUS_PER_TASK" + --mem="$MEMORY" + --time="$TIME_LIMIT" + --output="$LOG_ROOT/slurm/$JOB_NAME-%j.out" + --chdir="$REPO_ROOT" +) +[[ -n "$EXCLUDE_NODES" ]] && SBATCH_ARGS+=(--exclude="$EXCLUDE_NODES") +[[ -n "$NODELIST" ]] && SBATCH_ARGS+=(--nodelist="$NODELIST") + +JOB_BODY='set -euo pipefail +cd "$SERVER_REPO_ROOT" +log_dir="$SERVER_LOG_ROOT/job-$SLURM_JOB_ID" +mkdir -p "$log_dir" + +pids=() +ports=() +cleanup() { + if ((${#pids[@]} > 0)); then + kill "${pids[@]}" 2>/dev/null || true + fi +} +trap cleanup EXIT INT TERM + +for ((gpu = 0; gpu < SERVER_COUNT; gpu++)); do + port=$((BASE_PORT + gpu)) + lora_args=() + if [[ -n "$LORA_PATH" ]]; then + lora_args=(--enable-lora --max-lora-rank "$MAX_LORA_RANK" --lora-modules "$LORA_NAME=$LORA_PATH") + fi + CUDA_VISIBLE_DEVICES=$gpu uv run --isolated --extra fsdp \ + python -m vllm.entrypoints.openai.api_server \ + --model "$MODEL" \ + --served-model-name "$SERVED_MODEL" \ + --enable-auto-tool-choice \ + --tool-call-parser qwen3_xml \ + --max-model-len "$MAX_MODEL_LEN" \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \ + "${lora_args[@]}" \ + --host 0.0.0.0 \ + --port "$port" \ + >"$log_dir/server-$port.log" 2>&1 & + pids+=("$!") + ports+=("$port") +done + +echo "Started $SERVER_COUNT vLLM processes in parallel on $(hostname)." +for ((i = 0; i < SERVER_COUNT; i++)); do + elapsed=0 + until curl --fail --silent --max-time 2 "http://127.0.0.1:${ports[$i]}/health" >/dev/null 2>&1; do + if ! kill -0 "${pids[$i]}" 2>/dev/null; then + echo "Worker on port ${ports[$i]} exited during startup; see $log_dir/server-${ports[$i]}.log" >&2 + exit 1 + fi + if ((elapsed >= WORKER_START_TIMEOUT)); then + echo "Worker on port ${ports[$i]} did not become healthy within $WORKER_START_TIMEOUT seconds" >&2 + exit 1 + fi + sleep 5 + ((elapsed += 5)) + done + echo "Worker on port ${ports[$i]} is healthy." +done + +last_port=$((BASE_PORT + SERVER_COUNT - 1)) +echo "All $SERVER_COUNT vLLM servers are healthy on $(hostname), ports $BASE_PORT-$last_port." +wait +' +export JOB_BODY + +if ((DRY_RUN)); then + printf 'Environment:\n' + printf ' MODEL=%q SERVED_MODEL=%q SERVER_COUNT=%q BASE_PORT=%q MAX_MODEL_LEN=%q GPU_MEMORY_UTILIZATION=%q\n' \ + "$MODEL" "$SERVED_MODEL" "$SERVER_COUNT" "$BASE_PORT" "$MAX_MODEL_LEN" "$GPU_MEMORY_UTILIZATION" + if [[ -n "$LORA_PATH" ]]; then + printf ' LORA_PATH=%q LORA_NAME=%q MAX_LORA_RANK=%q\n' "$LORA_PATH" "$LORA_NAME" "$MAX_LORA_RANK" + fi + printf 'Command:\n sbatch' + printf ' %q' "${SBATCH_ARGS[@]}" + printf ' --wrap=%q\n' 'exec bash -c "$JOB_BODY"' + exit 0 +fi + +job_id=$(sbatch "${SBATCH_ARGS[@]}" --wrap='exec bash -c "$JOB_BODY"') +echo "Submitted vLLM cluster as Slurm job $job_id." +echo "Status: squeue -j $job_id -o '%.10i %.28j %.8T %.12M %.20R'" +echo "Slurm log: $LOG_ROOT/slurm/$JOB_NAME-$job_id.out" +echo "Worker logs: $LOG_ROOT/servers/job-$job_id/server-.log" diff --git a/examples/train/toolathlon_harbor/prepare_daytona_tasks.py b/examples/train/toolathlon_harbor/prepare_daytona_tasks.py new file mode 100644 index 0000000000..f36c22b614 --- /dev/null +++ b/examples/train/toolathlon_harbor/prepare_daytona_tasks.py @@ -0,0 +1,138 @@ +"""Stage Toolathlon tasks for a shared Daytona runtime snapshot/image.""" + +from __future__ import annotations + +import argparse +import json +import shutil +import tomllib +from pathlib import Path + + +def _task_ids(source: Path, task_list: Path | None) -> list[str]: + if task_list is None: + return sorted(path.name for path in source.iterdir() if path.is_dir()) + ids = [ + line.strip() + for line in task_list.read_text().splitlines() + if line.strip() and not line.lstrip().startswith("#") + ] + if len(ids) != len(set(ids)): + raise ValueError(f"duplicate task IDs in {task_list}") + return ids + + +def _rewrite_task_toml(path: Path, runtime_image: str) -> None: + document = tomllib.loads(path.read_text()) + servers = document.get("metadata", {}).get("mcp_servers") + if ( + not isinstance(servers, list) + or not servers + or not all(isinstance(item, str) for item in servers) + ): + raise ValueError( + f"{path}: metadata.mcp_servers must be a non-empty string list" + ) + + lines = path.read_text().splitlines() + start = next( + (i for i, line in enumerate(lines) if line.strip() == "[environment]"), None + ) + if start is None: + raise ValueError(f"{path}: missing [environment] section") + end = next( + (i for i in range(start + 1, len(lines)) if lines[i].lstrip().startswith("[")), + len(lines), + ) + retained = [ + line + for line in lines[start + 1 : end] + if not line.lstrip().startswith(("docker_image =", "workdir =", "env =")) + ] + while retained and not retained[-1].strip(): + retained.pop() + + q = json.dumps + environment_vars = ( + "env = { " + f"T3_BUNDLE_DIR = {q('/opt/task')}, " + f"T3_SERVERS = {q(','.join(servers))}, " + f"T3_WORLD_DUMP = {q('/logs/world_after.json')} " + "}" + ) + runtime = [ + f"docker_image = {q(runtime_image)}", + 'workdir = "/opt"', + environment_vars, + ] + rewritten = lines[: start + 1] + retained + runtime + [""] + lines[end:] + path.write_text("\n".join(rewritten).rstrip() + "\n") + + +def stage_tasks( + source: Path, + output: Path, + runtime_image: str, + *, + task_list: Path | None = None, + force: bool = False, +) -> int: + if output.exists(): + if not force: + raise FileExistsError( + f"output already exists: {output}; pass --force to replace it" + ) + shutil.rmtree(output) + output.mkdir(parents=True) + + ids = _task_ids(source, task_list) + for task_id in ids: + source_task = source / task_id + if not source_task.is_dir(): + raise FileNotFoundError(f"task not found: {source_task}") + destination = output / task_id + shutil.copytree( + source_task, + destination, + ignore=shutil.ignore_patterns("Dockerfile", "Dockerfile.*"), + ) + environment = destination / "environment" + if not (environment / "task" / "initial_state.json").is_file(): + raise FileNotFoundError(f"missing task bundle: {environment / 'task'}") + if not (environment / "mcp.json").is_file(): + raise FileNotFoundError(f"missing MCP config: {environment / 'mcp.json'}") + _rewrite_task_toml(destination / "task.toml", runtime_image) + + return len(ids) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--source", type=Path, required=True, help="source tasks directory" + ) + parser.add_argument( + "--output", type=Path, required=True, help="persistent staged tasks directory" + ) + parser.add_argument( + "--runtime-image", required=True, help="pullable fallback runtime image" + ) + parser.add_argument( + "--task-list", type=Path, help="optional file containing one task ID per line" + ) + parser.add_argument( + "--force", action="store_true", help="replace an existing output directory" + ) + args = parser.parse_args() + count = stage_tasks( + args.source.resolve(), + args.output.resolve(), + args.runtime_image, + task_list=args.task_list.resolve() if args.task_list else None, + force=args.force, + ) + print(f"Staged {count} upload-only tasks at {args.output.resolve()}") + + +if __name__ == "__main__": + main() diff --git a/examples/train/toolathlon_harbor/run_eval.sh b/examples/train/toolathlon_harbor/run_eval.sh index f60a136bc6..948bae2476 100755 --- a/examples/train/toolathlon_harbor/run_eval.sh +++ b/examples/train/toolathlon_harbor/run_eval.sh @@ -4,7 +4,7 @@ set -euo pipefail HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "$HERE/../../.." && pwd)" -TASKS_DIR="${TOOLATHLON_TASKS_DIR:-$REPO_ROOT/toolathlon-tasks/tasks}" +TASKS_DIR="${TOOLATHLON_TASKS_DIR:-$REPO_ROOT/toolathlon-tasks/eval_tasks}" RUNTIME_ARCHIVE="${TOOLATHLON_RUNTIME_ARCHIVE:-$REPO_ROOT/toolathlon-tasks/runtime/toolathlon-json-runtime-src.tar.gz}" RUNTIME_IMAGE="${TOOLATHLON_RUNTIME_IMAGE:-toolathlon-json-runtime:v1}" API_BASE="${TOOLATHLON_API_BASE:?Set TOOLATHLON_API_BASE to an OpenAI-compatible /v1 endpoint}" diff --git a/examples/train/toolathlon_harbor/run_grpo_qwen38_27b_2node.sh b/examples/train/toolathlon_harbor/run_grpo_qwen38_27b_2node.sh new file mode 100755 index 0000000000..2e3681e716 --- /dev/null +++ b/examples/train/toolathlon_harbor/run_grpo_qwen38_27b_2node.sh @@ -0,0 +1,155 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Two-node (16 GPU), colocated GRPO for Qwen3.8-27B on Harbor-native +# Toolathlon tasks. This expects a two-node Ray cluster to already be running +# and RAY_ADDRESS to point at it (normally: export RAY_ADDRESS=auto). +# +# Required: +# DAYTONA_API_KEY Daytona credential visible to Ray workers +# WANDB_API_KEY unless TRAINER_LOGGER=console +# TOOLATHLON_TASKS_DIR persistent, upload-only Daytona task directory +# +# No local Docker daemon is required. If TOOLATHLON_TASKS_DIR does not exist, +# the launcher stages it from the source tasks using TOOLATHLON_RUNTIME_IMAGE. + +script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P) +repo_root=$(cd -- "$script_dir/../../.." && pwd -P) + +: "${DAYTONA_API_KEY:?Set DAYTONA_API_KEY before launching training}" +: "${RAY_ADDRESS:=auto}" + +model_path=${MODEL_PATH:-Qwen/Qwen3.8-27B} +served_model_name=${SERVED_MODEL_NAME:-Qwen3.8-27B} +tasks_dir=${TOOLATHLON_TASKS_DIR:-$HOME/data/toolathlon-harbor/tasks} +source_tasks_dir=${TOOLATHLON_SOURCE_TASKS_DIR:-$repo_root/toolathlon-tasks/tasks} +runtime_image=${TOOLATHLON_RUNTIME_IMAGE:-us-docker.pkg.dev/afterquery-compute/compute-images/toolathlon-json-runtime:v1} +snapshot_template=${DAYTONA_SNAPSHOT_TEMPLATE:-toolathlon-json-runtime-v1} +storage_root=${STORAGE_ROOT:-$HOME/toolathlon_grpo_qwen38_27b} +lora_sync_path=${LORA_SYNC_PATH:-$repo_root/jobs/rl/runs/toolathlon-grpo-qwen38-27b-${SLURM_JOB_ID:-local}/lora_sync} +trial_config=${HARBOR_TRIAL_CONFIG:-$script_dir/harbor_daytona_training_config.yaml} + +run_name=${RUN_NAME:-toolathlon-grpo-qwen38-27b-lora-2node-130k-g8} +trainer_logger=${TRAINER_LOGGER:-wandb} + +# Four prompts per optimizer step, with eight samples each, yields 32 complete +# agent trajectories per update. Start here because a single trajectory may +# contain many 131k-token model turns; increase only after measuring one update. +group_size=${GROUP_SIZE:-8} +train_batch_size=${TRAIN_BATCH_SIZE:-4} +policy_mini_batch_size=${POLICY_MINI_BATCH_SIZE:-4} +max_model_len=${MAX_MODEL_LEN:-131072} + +num_nodes=${NUM_NODES:-2} +gpus_per_node=${GPUS_PER_NODE:-8} +num_inference_engines=${NUM_INFERENCE_ENGINES:-16} +tensor_parallel_size=${TENSOR_PARALLEL_SIZE:-1} +max_concurrency=${MAX_CONCURRENCY:-128} +sequence_parallel_size=${SEQUENCE_PARALLEL_SIZE:-4} +num_logger_train_samples=${NUM_LOGGER_TRAIN_SAMPLES:-2} + +lora_rank=${LORA_RANK:-32} +lora_alpha=${LORA_ALPHA:-64} +lora_targets=${LORA_TARGETS:-'[q_proj,k_proj,v_proj,o_proj]'} +learning_rate=${LEARNING_RATE:-1.0e-6} + +if [[ ! -d "$tasks_dir" ]]; then + echo "Staging upload-only Toolathlon tasks at $tasks_dir" >&2 + uv run python "$script_dir/prepare_daytona_tasks.py" \ + --source "$source_tasks_dir" \ + --output "$tasks_dir" \ + --runtime-image "$runtime_image" +fi +if [[ ! -f "$trial_config" ]]; then + echo "Harbor trial config not found: $trial_config" >&2 + exit 2 +fi +if [[ "$trainer_logger" == wandb && -z "${WANDB_API_KEY:-}" ]]; then + echo "WANDB_API_KEY is required when TRAINER_LOGGER=wandb." >&2 + exit 2 +fi + +mkdir -p \ + "$storage_root/trials" \ + "$storage_root/ckpts" \ + "$storage_root/exports" \ + "$storage_root/logs" \ + "$lora_sync_path" + +cd "$repo_root" +export RAY_ADDRESS +export TILELANG_CLEANUP_TEMP_FILES=${TILELANG_CLEANUP_TEMP_FILES:-1} + +python_bin=${SKYRL_PYTHON_BIN:-$repo_root/.venv/bin/python} +if [[ ! -x "$python_bin" ]]; then + echo "Shared SkyRL Python is missing: $python_bin" >&2 + exit 2 +fi + +"$python_bin" -m examples.train_integrations.harbor.entrypoints.main_harbor \ + data.train_data="['$tasks_dir']" \ + trainer.policy.model.path="$model_path" \ + trainer.policy.language_model_only=true \ + trainer.ref.language_model_only=true \ + generator.inference_engine.served_model_name="$served_model_name" \ + generator.inference_engine.language_model_only=true \ + harbor_trial_config_path="$trial_config" \ + harbor_trial_config.trials_dir="$storage_root/trials" \ + harbor_trial_config.environment.kwargs.snapshot_template_name="$snapshot_template" \ + trainer.export_path="$storage_root/exports" \ + trainer.ckpt_path="$storage_root/ckpts" \ + trainer.log_path="$storage_root/logs" \ + trainer.algorithm.advantage_estimator=grpo \ + trainer.algorithm.loss_reduction=token_mean \ + trainer.algorithm.grpo_norm_by_std=false \ + trainer.algorithm.dynamic_sampling.type=filter \ + trainer.algorithm.dynamic_sampling.max_sample_batches=60 \ + trainer.algorithm.use_kl_loss=false \ + trainer.algorithm.max_seq_len="$max_model_len" \ + trainer.policy.model.lora.rank="$lora_rank" \ + trainer.policy.model.lora.alpha="$lora_alpha" \ + trainer.policy.model.lora.target_modules="$lora_targets" \ + trainer.policy.model.lora.lora_sync_path="$lora_sync_path" \ + trainer.policy.optimizer_config.lr="$learning_rate" \ + trainer.placement.colocate_all=true \ + trainer.strategy=fsdp \ + trainer.placement.policy_num_nodes="$num_nodes" \ + trainer.placement.ref_num_nodes="$num_nodes" \ + trainer.placement.policy_num_gpus_per_node="$gpus_per_node" \ + trainer.placement.ref_num_gpus_per_node="$gpus_per_node" \ + trainer.policy.sequence_parallel_size="$sequence_parallel_size" \ + generator.inference_engine.num_engines="$num_inference_engines" \ + generator.inference_engine.tensor_parallel_size="$tensor_parallel_size" \ + generator.inference_engine.engine_init_kwargs.max_model_len="$max_model_len" \ + generator.inference_engine.engine_init_kwargs.enable_auto_tool_choice=true \ + generator.inference_engine.engine_init_kwargs.tool_call_parser=qwen3_xml \ + generator.inference_engine.engine_init_kwargs.enable_log_requests=false \ + generator.inference_engine.gpu_memory_utilization=0.8 \ + generator.inference_engine.backend=vllm \ + generator.inference_engine.run_engines_locally=true \ + generator.inference_engine.weight_sync_backend=nccl \ + generator.inference_engine.enforce_eager=false \ + generator.batched=false \ + generator.step_wise_trajectories=true \ + generator.merge_stepwise_output=true \ + generator.n_samples_per_prompt="$group_size" \ + generator.apply_overlong_filtering=true \ + generator.rate_limit.enabled=true \ + generator.rate_limit.trajectories_per_second=4 \ + generator.rate_limit.max_concurrency="$max_concurrency" \ + trainer.epochs=3 \ + trainer.train_batch_size="$train_batch_size" \ + trainer.policy_mini_batch_size="$policy_mini_batch_size" \ + trainer.micro_forward_batch_size_per_gpu=4 \ + trainer.micro_train_batch_size_per_gpu=4 \ + trainer.update_epochs_per_batch=1 \ + trainer.num_logger_train_samples="$num_logger_train_samples" \ + trainer.eval_interval=-1 \ + trainer.ckpt_interval=5 \ + trainer.max_ckpts_to_keep=3 \ + trainer.hf_save_interval=5 \ + trainer.logger="$trainer_logger" \ + trainer.project_name=toolathlon-harbor \ + trainer.run_name="$run_name" \ + trainer.resume_mode=latest \ + "$@" diff --git a/examples/train/toolathlon_harbor/run_grpo_qwen38_27b_2node_slurm.sh b/examples/train/toolathlon_harbor/run_grpo_qwen38_27b_2node_slurm.sh new file mode 100755 index 0000000000..9631955da8 --- /dev/null +++ b/examples/train/toolathlon_harbor/run_grpo_qwen38_27b_2node_slurm.sh @@ -0,0 +1,186 @@ +#!/usr/bin/env bash +#SBATCH --job-name=toolathlon-grpo-qwen38-27b +#SBATCH --partition=gpu +#SBATCH --wckey=afterquery_research +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --cpus-per-task=128 +#SBATCH --mem=768G +#SBATCH --time=48:00:00 +#SBATCH --exclude=node-2 + +set -euo pipefail + +if [[ -z "${SLURM_JOB_ID:-}" ]]; then + script_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P) + repo_root=$(cd -- "$script_dir/../../.." && pwd -P) +else + repo_root=$(pwd -P) + script_dir="$repo_root/examples/train/toolathlon_harbor" +fi + +# Make the same script convenient to invoke from a login shell. Any arguments +# are forwarded to the inner SkyRL launcher as Hydra overrides. +if [[ -z "${SLURM_JOB_ID:-}" ]]; then + log_dir=${SLURM_LOG_DIR:-$repo_root/jobs/rl/slurm} + mkdir -p "$log_dir" + exec sbatch --parsable \ + --output="$log_dir/toolathlon-grpo-qwen38-27b-%j.out" \ + --chdir="$repo_root" \ + "$0" "$@" +fi + +mapfile -t nodes < <(scontrol show hostnames "$SLURM_JOB_NODELIST") +if [[ ${#nodes[@]} -ne 2 ]]; then + echo "Expected exactly two allocated nodes; got ${#nodes[@]}: ${nodes[*]}" >&2 + exit 2 +fi + +head_node=${nodes[0]} +worker_node=${nodes[1]} +ray_port=${RAY_PORT:-6379} +ray_dashboard_port=${RAY_DASHBOARD_PORT:-8265} +ray_start_timeout=${RAY_START_TIMEOUT:-300} +ray_num_cpus=${RAY_NUM_CPUS_PER_NODE:-${SLURM_CPUS_PER_TASK:-128}} +ray_num_gpus=${RAY_NUM_GPUS_PER_NODE:-8} +ray_tmp_root=${RAY_TMP_ROOT:-/tmp/skyrl-ray-$SLURM_JOB_ID} +cuda_home=${SKYRL_CUDA_HOME:-/usr/local/cuda-12.8} +shared_tasks_dir=${TOOLATHLON_SHARED_TASKS_DIR:-$repo_root/data/toolathlon-harbor/tasks} +local_tasks_dir=${TOOLATHLON_LOCAL_TASKS_DIR:-$ray_tmp_root/toolathlon-harbor/tasks} + +python_bin=${SKYRL_PYTHON_BIN:-$repo_root/.venv/bin/python} +ray_bin=${SKYRL_RAY_BIN:-$repo_root/.venv/bin/ray} +if [[ ! -x "$python_bin" || ! -x "$ray_bin" ]]; then + echo "Shared SkyRL environment is incomplete under $repo_root/.venv" >&2 + exit 2 +fi +if [[ ! -x "$cuda_home/bin/nvcc" ]]; then + echo "CUDA toolkit is incomplete: $cuda_home/bin/nvcc is not executable" >&2 + exit 2 +fi +export CUDA_HOME="$cuda_home" +export TILELANG_CLEANUP_TEMP_FILES=${TILELANG_CLEANUP_TEMP_FILES:-1} +export PATH="$CUDA_HOME/bin:$PATH" +ray_cmd=("$ray_bin") +head_ip=$(srun --overlap --nodes=1 --ntasks=1 --cpus-per-task=1 \ + --gres=gpu:0 --nodelist="$head_node" \ + bash -lc "hostname -I | awk '{for (i=1; i<=NF; i++) if (\$i ~ /^10\.65\.0\./) {print \$i; found=1; break} if (!found) print \$1}'") +head_ip=${head_ip//$'\n'/} +if [[ -z "$head_ip" ]]; then + echo "Could not determine the Ray head IP for $head_node." >&2 + exit 2 +fi +ray_address="$head_ip:$ray_port" + +head_step_pid="" +worker_step_pid="" +cleanup() { + status=$? + trap - EXIT INT TERM + for pid in "$worker_step_pid" "$head_step_pid"; do + if [[ -n "$pid" ]]; then + kill "$pid" 2>/dev/null || true + fi + done + wait 2>/dev/null || true + exit "$status" +} +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +echo "Staging Toolathlon tasks from shared storage onto both nodes" +srun --overlap --nodes=2 --ntasks=2 --ntasks-per-node=1 --cpus-per-task=4 \ + --gres=gpu:0 \ + "$script_dir/stage_tasks_local.sh" "$shared_tasks_dir" "$local_tasks_dir" +export TOOLATHLON_TASKS_DIR="$local_tasks_dir" +# Ray rendezvous stays on the routable 10.65.0.x network. Use the eight +# matching 400 Gb/s rails for both Gloo's CPU model-state distribution and +# NCCL's GPU collectives; pinning either backend to ens1 bottlenecks it at +# 10 Gb/s. +export GLOO_SOCKET_IFNAME=${GLOO_SOCKET_IFNAME:-ens2,ens3,ens4,ens5,ens6,ens7,ens8,ens9} +export NCCL_SOCKET_IFNAME=${NCCL_SOCKET_IFNAME:-ens2,ens3,ens4,ens5,ens6,ens7,ens8,ens9} +export NCCL_IB_DISABLE=${NCCL_IB_DISABLE:-0} +export NCCL_IB_HCA=${NCCL_IB_HCA:-=rocep9s0,rocep23s0,rocep64s0,rocep73s0,rocep134s0,rocep143s0,rocep189s0,rocep198s0} +export NCCL_DEBUG=${NCCL_DEBUG:-INFO} +export NCCL_DEBUG_SUBSYS=${NCCL_DEBUG_SUBSYS:-INIT,NET} + +echo "Distributed networking: GLOO_SOCKET_IFNAME=$GLOO_SOCKET_IFNAME" +echo "Distributed networking: NCCL_SOCKET_IFNAME=$NCCL_SOCKET_IFNAME NCCL_IB_HCA=$NCCL_IB_HCA" + +echo "Starting Ray head on $head_node ($head_ip)" +srun --overlap --nodes=1 --ntasks=1 --cpus-per-task="${SLURM_CPUS_PER_TASK:-128}" \ + --gres="gpu:$ray_num_gpus" --nodelist="$head_node" \ + bash -lc " + set -euo pipefail + cd '$repo_root' + mkdir -p '$ray_tmp_root/head' + '$repo_root/.venv/bin/ray' stop --force >/dev/null 2>&1 || true + exec '$repo_root/.venv/bin/ray' start \ + --head --block --disable-usage-stats \ + --node-ip-address='$head_ip' --port='$ray_port' \ + --dashboard-host=0.0.0.0 --dashboard-port='$ray_dashboard_port' \ + --num-cpus='$ray_num_cpus' --num-gpus='$ray_num_gpus' \ + --temp-dir='$ray_tmp_root/head' + " & +head_step_pid=$! + +deadline=$((SECONDS + ray_start_timeout)) +until "${ray_cmd[@]}" status --address="$ray_address" >/dev/null 2>&1; do + if ! kill -0 "$head_step_pid" 2>/dev/null; then + echo "Ray head step exited during startup." >&2 + wait "$head_step_pid" + fi + if (( SECONDS >= deadline )); then + echo "Ray head did not become ready within $ray_start_timeout seconds." >&2 + exit 1 + fi + sleep 2 +done + +echo "Starting Ray worker on $worker_node" +srun --overlap --nodes=1 --ntasks=1 --cpus-per-task="${SLURM_CPUS_PER_TASK:-128}" \ + --gres="gpu:$ray_num_gpus" --nodelist="$worker_node" \ + bash -lc " + set -euo pipefail + cd '$repo_root' + mkdir -p '$ray_tmp_root/worker' + worker_ip=\$(hostname -I | awk '{for (i=1; i<=NF; i++) if (\$i ~ /^10\.65\.0\./) {print \$i; found=1; break} if (!found) print \$1}') + '$repo_root/.venv/bin/ray' stop --force >/dev/null 2>&1 || true + exec '$repo_root/.venv/bin/ray' start \ + --block --disable-usage-stats --address='$ray_address' \ + --node-ip-address=\"\$worker_ip\" \ + --num-cpus='$ray_num_cpus' --num-gpus='$ray_num_gpus' \ + --temp-dir='$ray_tmp_root/worker' + " & +worker_step_pid=$! + +echo "Waiting for Ray to report 2 nodes and 16 GPUs" +deadline=$((SECONDS + ray_start_timeout)) +until RAY_ADDRESS="$ray_address" "$python_bin" - <<'PY' +import ray + +ray.init(address="auto", logging_level="ERROR") +alive = [node for node in ray.nodes() if node["Alive"]] +resources = ray.cluster_resources() +raise SystemExit(0 if len(alive) == 2 and resources.get("GPU", 0) >= 16 else 1) +PY +do + if ! kill -0 "$worker_step_pid" 2>/dev/null; then + echo "Ray worker step exited during startup." >&2 + wait "$worker_step_pid" + fi + if (( SECONDS >= deadline )); then + echo "The two-node Ray cluster did not become ready within $ray_start_timeout seconds." >&2 + "${ray_cmd[@]}" status --address="$ray_address" || true + exit 1 + fi + sleep 3 +done + +echo "Ray cluster ready at $ray_address; launching Toolathlon GRPO" +export RAY_ADDRESS="$ray_address" +srun --overlap --nodes=1 --ntasks=1 --cpus-per-task=4 --gres=gpu:0 \ + --nodelist="$head_node" --chdir="$repo_root" \ + "$script_dir/run_grpo_qwen38_27b_2node.sh" "$@" diff --git a/examples/train/toolathlon_harbor/stage_tasks_local.sh b/examples/train/toolathlon_harbor/stage_tasks_local.sh new file mode 100755 index 0000000000..159d9444fa --- /dev/null +++ b/examples/train/toolathlon_harbor/stage_tasks_local.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Stage Toolathlon tasks from shared storage onto node-local storage before a +# Slurm training run. Local staging avoids repeatedly reading thousands of small +# task files from the shared filesystem. The source and copied destination must +# each contain TOOLATHLON_EXPECTED_TASKS task.toml files (1,900 by default). +# +# rsync intentionally updates the destination in place without --delete. Files +# left over from an earlier staging run are therefore retained; extra task.toml +# files are detected by the destination count, but other stale files are not. + +if [[ $# -ne 2 ]]; then + echo "Usage: $0 SHARED_TASKS_DIR LOCAL_TASKS_DIR" >&2 + exit 2 +fi + +source_dir=$1 +destination_dir=$2 +expected_tasks=${TOOLATHLON_EXPECTED_TASKS:-1900} + +if [[ ! -d "$source_dir" ]]; then + echo "Shared Toolathlon task directory is missing: $source_dir" >&2 + exit 2 +fi + +source_count=$(find "$source_dir" -mindepth 2 -maxdepth 2 -name task.toml -type f | wc -l) +if [[ "$source_count" -ne "$expected_tasks" ]]; then + echo "Expected $expected_tasks shared tasks, found $source_count in $source_dir" >&2 + exit 2 +fi + +mkdir -p "$destination_dir" +rsync -a "$source_dir/" "$destination_dir/" + +destination_count=$(find "$destination_dir" -mindepth 2 -maxdepth 2 -name task.toml -type f | wc -l) +if [[ "$destination_count" -ne "$expected_tasks" ]]; then + echo "Expected $expected_tasks staged tasks, found $destination_count in $destination_dir" >&2 + exit 1 +fi + +echo "node=$(hostname) staged_tasks=$destination_count destination=$destination_dir" diff --git a/examples/train/toolathlon_harbor/toolathlon_context_manager.py b/examples/train/toolathlon_harbor/toolathlon_context_manager.py new file mode 100644 index 0000000000..0a2d825ad1 --- /dev/null +++ b/examples/train/toolathlon_harbor/toolathlon_context_manager.py @@ -0,0 +1,577 @@ +"""Bounded active context with durable, searchable Harbor trajectory history.""" + +from __future__ import annotations + +import hashlib +import json +import re +import uuid +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +LOCAL_TOOL_SCHEMAS = [ + { + "name": "context_check_status", + "description": "Report active context usage and prior compaction events.", + "inputSchema": {"type": "object", "properties": {}, "additionalProperties": False}, + }, + { + "name": "context_manage", + "description": "Schedule deterministic removal of older conversation exchanges.", + "inputSchema": { + "type": "object", + "properties": { + "method": { + "type": "string", + "enum": ["keep_recent_turns", "keep_recent_percent", "delete_first_turns", "delete_first_percent"], + }, + "value": {"type": "number", "exclusiveMinimum": 0}, + }, + "required": ["method", "value"], + "additionalProperties": False, + }, + }, + { + "name": "history_search", + "description": "Search the complete durable conversation history, including compacted turns.", + "inputSchema": { + "type": "object", + "properties": { + "query": {"type": "string"}, + "max_results": {"type": "integer", "minimum": 1, "maximum": 20}, + }, + "required": ["query"], + "additionalProperties": False, + }, + }, + { + "name": "history_view", + "description": "View bounded records from the complete durable conversation history.", + "inputSchema": { + "type": "object", + "properties": { + "start_seq": {"type": "integer", "minimum": 0}, + "count": {"type": "integer", "minimum": 1, "maximum": 20}, + }, + "required": ["start_seq"], + "additionalProperties": False, + }, + }, + { + "name": "tool_output_search", + "description": "Regex-search an offloaded tool output and start a paginated search session.", + "inputSchema": { + "type": "object", + "properties": { + "artifact_id": {"type": "string"}, + "pattern": {"type": "string"}, + "page_size": {"type": "integer", "minimum": 1, "maximum": 50}, + "context_size": {"type": "integer", "minimum": 1}, + }, + "required": ["artifact_id", "pattern"], + "additionalProperties": False, + }, + }, + { + "name": "tool_output_search_navigate", + "description": "Navigate a search session created by tool_output_search.", + "inputSchema": { + "type": "object", + "properties": { + "search_session_id": {"type": "string"}, + "action": { + "type": "string", + "enum": ["next_page", "prev_page", "jump_to_page", "first_page", "last_page"], + }, + "target_page": {"type": "integer", "minimum": 1}, + }, + "required": ["search_session_id", "action"], + "additionalProperties": False, + }, + }, + { + "name": "tool_output_view", + "description": "View the first character page of an offloaded tool output.", + "inputSchema": { + "type": "object", + "properties": { + "artifact_id": {"type": "string"}, + "page_size": {"type": "integer", "minimum": 1, "maximum": 100000}, + }, + "required": ["artifact_id"], + "additionalProperties": False, + }, + }, + { + "name": "tool_output_view_navigate", + "description": "Navigate a view session created by tool_output_view.", + "inputSchema": { + "type": "object", + "properties": { + "view_session_id": {"type": "string"}, + "action": { + "type": "string", + "enum": ["next_page", "prev_page", "jump_to_page", "first_page", "last_page"], + }, + "target_page": {"type": "integer", "minimum": 1}, + }, + "required": ["view_session_id", "action"], + "additionalProperties": False, + }, + }, +] + + +@dataclass +class ContextPolicy: + max_context_tokens: int + max_output_tokens: int + safety_tokens: int = 2048 + warning_ratio: float = 0.75 + compact_ratio: float = 0.85 + target_ratio: float = 0.70 + keep_recent_turns: int = 4 + keep_reasoning_turns: int = 1 + inline_tool_output_chars: int = 12_000 + preview_chars: int = 2_000 + max_resets: int = 2 + + @property + def prompt_budget(self) -> int: + return self.max_context_tokens - self.max_output_tokens - self.safety_tokens + + def validate(self) -> None: + if self.prompt_budget <= 0: + raise ValueError("max_context_tokens must exceed max_tokens + context_safety_tokens") + if not 0 < self.target_ratio < self.warning_ratio < self.compact_ratio < 1: + raise ValueError("context ratios must satisfy target < warning < compact < 1") + + +@dataclass +class ManagedContext: + instruction: str + artifact_dir: Path + policy: ContextPolicy + full_messages: list[dict[str, Any]] = field(default_factory=list) + active_messages: list[dict[str, Any]] = field(default_factory=list) + events: list[dict[str, Any]] = field(default_factory=list) + latest_prompt_tokens: int = 0 + additions_since_prompt: list[dict[str, Any]] = field(default_factory=list) + pending_compaction: dict[str, Any] | None = None + reset_count: int = 0 + search_sessions: dict[str, dict[str, Any]] = field(default_factory=dict) + view_sessions: dict[str, dict[str, Any]] = field(default_factory=dict) + + def __post_init__(self) -> None: + self.policy.validate() + self.artifact_dir.mkdir(parents=True, exist_ok=True) + (self.artifact_dir / "tool_outputs").mkdir(exist_ok=True) + initial = {"role": "user", "content": self.instruction} + self.full_messages.append(initial.copy()) + self.active_messages.append(initial.copy()) + self._append_history(initial, active=True) + + @property + def tool_schemas(self) -> list[dict[str, Any]]: + return LOCAL_TOOL_SCHEMAS + + @property + def history_path(self) -> Path: + return self.artifact_dir / "context_history.jsonl" + + def _append_history(self, message: dict[str, Any], *, active: bool) -> None: + record = {"seq": self._history_count(), "active_at_write": active, "message": message} + with self.history_path.open("a", encoding="utf-8") as stream: + stream.write(json.dumps(record, ensure_ascii=False) + "\n") + + def _history_count(self) -> int: + if not self.history_path.exists(): + return 0 + with self.history_path.open("r", encoding="utf-8") as stream: + return sum(1 for _ in stream) + + def append(self, message: dict[str, Any]) -> None: + full = json.loads(json.dumps(message, ensure_ascii=False)) + active = json.loads(json.dumps(message, ensure_ascii=False)) + self.full_messages.append(full) + self.active_messages.append(active) + self.additions_since_prompt.append(active) + self._append_history(full, active=True) + + def append_tool_result(self, message: dict[str, Any]) -> dict[str, Any]: + content = str(message.get("content") or "") + if len(content) <= self.policy.inline_tool_output_chars: + self.append(message) + return message + digest = hashlib.sha256(content.encode("utf-8", errors="replace")).hexdigest()[:16] + artifact_id = f"tool-{digest}" + artifact_path = self.artifact_dir / "tool_outputs" / f"{artifact_id}.json" + artifact_path.write_text(json.dumps({"artifact_id": artifact_id, "content": content}, ensure_ascii=False), encoding="utf-8") + n = self.policy.preview_chars + preview = ( + f"[Tool output offloaded as {artifact_id}; {len(content)} characters. " + "Use tool_output_search or tool_output_view for complete content.]\n" + f"{content[:n]}\n... [offloaded] ...\n{content[-n:]}" + ) + active_message = dict(message) + active_message["content"] = preview + durable_message = dict(active_message) + durable_message["artifact_id"] = artifact_id + self.full_messages.append(durable_message) + self.active_messages.append(active_message) + self.additions_since_prompt.append(active_message) + self._append_history(durable_message, active=True) + self._event("tool_output_offloaded", artifact_id=artifact_id, original_chars=len(content)) + return durable_message + + @staticmethod + def _estimate(messages: list[dict[str, Any]]) -> int: + # Conservative tokenizer-independent estimate; calibrated by real prompt usage after each request. + return sum(max(1, len(json.dumps(message, ensure_ascii=False)) // 3) for message in messages) + + def estimated_prompt_tokens(self) -> int: + if not self.latest_prompt_tokens: + return self._estimate(self.active_messages) + return self.latest_prompt_tokens + self._estimate(self.additions_since_prompt) + + def observe_prompt_tokens(self, tokens: int) -> None: + self.latest_prompt_tokens = max(0, tokens) + self.additions_since_prompt.clear() + + def _event(self, kind: str, **details: Any) -> None: + self.events.append({"type": kind, **details}) + + def prepare_for_request(self) -> None: + estimate = self.estimated_prompt_tokens() + ratio = estimate / self.policy.prompt_budget + if ratio >= self.policy.warning_ratio: + self._event("context_warning", estimated_tokens=estimate, ratio=ratio) + if ratio >= self.policy.compact_ratio: + self.compact_automatic(estimate) + + def _exchange_ranges(self) -> list[tuple[int, int]]: + ranges: list[tuple[int, int]] = [] + start: int | None = None + for index, message in enumerate(self.active_messages): + if message.get("role") == "assistant": + if start is not None: + ranges.append((start, index)) + start = index + if start is not None: + ranges.append((start, len(self.active_messages))) + return ranges + + def _prune_old_reasoning(self) -> int: + assistant_indices = [i for i, m in enumerate(self.active_messages) if m.get("role") == "assistant"] + eligible = assistant_indices[:-self.policy.keep_reasoning_turns] if self.policy.keep_reasoning_turns else assistant_indices + removed = 0 + for index in eligible: + reasoning = self.active_messages[index].pop("reasoning_content", None) + if reasoning is not None: + removed += max(1, len(str(reasoning)) // 3) + if removed: + self._event("reasoning_pruned", estimated_tokens_removed=removed, assistant_turns=len(eligible)) + return removed + + def compact_automatic(self, estimate_before: int) -> None: + current_estimate = max(0, estimate_before - self._prune_old_reasoning()) + target = int(self.policy.prompt_budget * self.policy.target_ratio) + ranges = self._exchange_ranges() + removable = max(0, len(ranges) - self.policy.keep_recent_turns) + removed = 0 + while removable > 0 and current_estimate > target: + ranges = self._exchange_ranges() + start, end = ranges[0] + current_estimate = max( + 0, current_estimate - self._estimate(self.active_messages[start:end]) + ) + del self.active_messages[start:end] + removed += 1 + removable -= 1 + self.latest_prompt_tokens = 0 + self.additions_since_prompt.clear() + self._event( + "automatic_compaction", + estimated_tokens_before=estimate_before, + estimated_tokens_after=current_estimate, + exchanges_removed=removed, + ) + + def schedule_manual(self, args: dict[str, Any]) -> dict[str, Any]: + method = str(args.get("method") or "") + value = args.get("value") + valid = {"keep_recent_turns", "keep_recent_percent", "delete_first_turns", "delete_first_percent"} + if method not in valid or not isinstance(value, (int, float)) or value <= 0: + return {"status": "error", "message": "invalid context compaction method or value"} + if "percent" in method and value >= 100: + return {"status": "error", "message": "percentage must be less than 100"} + self.pending_compaction = {"method": method, "value": value} + return {"status": "scheduled", "method": method, "value": value} + + def apply_pending(self) -> None: + if not self.pending_compaction: + return + request = self.pending_compaction + self.pending_compaction = None + ranges = self._exchange_ranges() + eligible = max(0, len(ranges) - self.policy.keep_recent_turns) + method, value = request["method"], request["value"] + if method == "keep_recent_turns": + delete = max(0, len(ranges) - max(self.policy.keep_recent_turns, int(value))) + elif method == "keep_recent_percent": + keep = max(self.policy.keep_recent_turns, int(len(ranges) * float(value) / 100)) + delete = max(0, len(ranges) - keep) + elif method == "delete_first_turns": + delete = min(eligible, int(value)) + else: + delete = min(eligible, int(eligible * float(value) / 100)) + for _ in range(delete): + start, end = self._exchange_ranges()[0] + del self.active_messages[start:end] + self.latest_prompt_tokens = 0 + self.additions_since_prompt.clear() + self._event("manual_compaction", method=method, value=value, exchanges_removed=delete) + + def status(self) -> dict[str, Any]: + estimate = self.estimated_prompt_tokens() + return { + "estimated_prompt_tokens": estimate, + "prompt_budget": self.policy.prompt_budget, + "usage_ratio": estimate / self.policy.prompt_budget, + "active_messages": len(self.active_messages), + "full_messages": len(self.full_messages), + "reset_count": self.reset_count, + "recent_events": self.events[-10:], + } + + def _read_history(self) -> list[dict[str, Any]]: + if not self.history_path.exists(): + return [] + return [json.loads(line) for line in self.history_path.read_text(encoding="utf-8").splitlines() if line] + + def _artifact_content(self, artifact_id: str) -> str: + if not re.fullmatch(r"tool-[0-9a-f]{16}", artifact_id): + raise ValueError("invalid artifact_id") + path = self.artifact_dir / "tool_outputs" / f"{artifact_id}.json" + return str(json.loads(path.read_text(encoding="utf-8"))["content"]) + + @staticmethod + def _target_page(session: dict[str, Any], args: dict[str, Any], total_pages: int) -> int: + current = int(session.get("current_page", 1)) + action = str(args.get("action", "next_page")) + if action == "next_page": + return min(current + 1, total_pages) + if action == "prev_page": + return max(current - 1, 1) + if action == "first_page": + return 1 + if action == "last_page": + return total_pages + if action == "jump_to_page": + target = int(args.get("target_page", 0)) + if not 1 <= target <= total_pages: + raise ValueError(f"target_page must be between 1 and {total_pages}") + return target + raise ValueError(f"invalid navigation action: {action}") + + @staticmethod + def _format_search_page(session_id: str, session: dict[str, Any], page: int) -> dict[str, Any]: + matches = session["matches"] + page_size = session["page_size"] + total_pages = max(1, (len(matches) + page_size - 1) // page_size) + start = (page - 1) * page_size + results = [] + for match in matches[start : start + page_size]: + results.append( + { + "match_text": match["match_text"], + "start_pos": match["start_pos"], + "end_pos": match["end_pos"], + "line_num": match["line_num"], + "context": ( + match["before_context"] + + f">>>{match['match_text']}<<<" + + match["after_context"] + ), + } + ) + return { + "artifact_id": session["artifact_id"], + "pattern": session["pattern"], + "search_session_id": session_id, + "total_matches": len(matches), + "current_page": page, + "total_pages": total_pages, + "page_size": page_size, + "file_size_chars": session["content_length"], + "results": results, + } + + def _search_output(self, args: dict[str, Any]) -> dict[str, Any]: + artifact_id = str(args["artifact_id"]) + pattern = str(args["pattern"]).strip() + page_size = int(args.get("page_size", 10)) + context_size = int(args.get("context_size", 1000)) + if not pattern: + raise ValueError("pattern is required") + if not 1 <= page_size <= 50: + raise ValueError("page_size must be between 1 and 50") + if context_size < 1: + raise ValueError("context_size must be positive") + try: + regex = re.compile(pattern, re.IGNORECASE | re.MULTILINE | re.DOTALL) + except re.error as exc: + raise ValueError(f"invalid regex pattern: {exc}") from exc + content = self._artifact_content(artifact_id) + matches = [] + for match in regex.finditer(content): + start, end = match.span() + context_start = max(0, start - context_size // 2) + context_end = min(len(content), end + context_size // 2) + matches.append( + { + "match_text": match.group(0), + "start_pos": start, + "end_pos": end, + "line_num": content[:start].count("\n") + 1, + "before_context": content[context_start:start], + "after_context": content[end:context_end], + } + ) + session_id = uuid.uuid4().hex[:8] + session = { + "artifact_id": artifact_id, + "pattern": pattern, + "matches": matches, + "page_size": page_size, + "context_size": context_size, + "content_length": len(content), + "current_page": 1, + } + self.search_sessions[session_id] = session + return self._format_search_page(session_id, session, 1) + + def _navigate_search(self, args: dict[str, Any]) -> dict[str, Any]: + session_id = str(args["search_session_id"]) + if session_id not in self.search_sessions: + raise ValueError("invalid or expired search_session_id") + session = self.search_sessions[session_id] + total_pages = max(1, (len(session["matches"]) + session["page_size"] - 1) // session["page_size"]) + page = self._target_page(session, args, total_pages) + session["current_page"] = page + return self._format_search_page(session_id, session, page) + + @staticmethod + def _format_view_page(session_id: str, session: dict[str, Any], content: str, page: int) -> dict[str, Any]: + page_size = session["page_size"] + total_pages = max(1, (len(content) + page_size - 1) // page_size) + start = (page - 1) * page_size + end = min(start + page_size, len(content)) + return { + "artifact_id": session["artifact_id"], + "view_session_id": session_id, + "current_page": page, + "total_pages": total_pages, + "page_size": page_size, + "start_pos": start, + "end_pos": end, + "file_size_chars": len(content), + "content": content[start:end], + } + + def _view_output(self, args: dict[str, Any]) -> dict[str, Any]: + artifact_id = str(args["artifact_id"]) + page_size = int(args.get("page_size", 10_000)) + if not 1 <= page_size <= 100_000: + raise ValueError("page_size must be between 1 and 100000") + content = self._artifact_content(artifact_id) + session_id = uuid.uuid4().hex[:8] + session = {"artifact_id": artifact_id, "page_size": page_size, "current_page": 1} + self.view_sessions[session_id] = session + return self._format_view_page(session_id, session, content, 1) + + def _navigate_view(self, args: dict[str, Any]) -> dict[str, Any]: + session_id = str(args["view_session_id"]) + if session_id not in self.view_sessions: + raise ValueError("invalid or expired view_session_id") + session = self.view_sessions[session_id] + content = self._artifact_content(session["artifact_id"]) + total_pages = max(1, (len(content) + session["page_size"] - 1) // session["page_size"]) + page = self._target_page(session, args, total_pages) + session["current_page"] = page + return self._format_view_page(session_id, session, content, page) + + def call_local_tool(self, name: str, args: dict[str, Any]) -> dict[str, Any]: + if name == "context_check_status": + return self.status() + if name == "context_manage": + return self.schedule_manual(args) + if name == "history_view": + start, count = int(args["start_seq"]), min(20, int(args.get("count", 5))) + return {"records": self._read_history()[start : start + count]} + if name == "history_search": + query = str(args["query"]).casefold() + limit = min(20, int(args.get("max_results", 10))) + found = [] + for record in self._read_history(): + text = json.dumps(record.get("message"), ensure_ascii=False) + pos = text.casefold().find(query) + if pos >= 0: + found.append({"seq": record["seq"], "snippet": text[max(0, pos - 500) : pos + len(query) + 500]}) + if len(found) >= limit: + break + return {"results": found} + if name == "tool_output_search": + return self._search_output(args) + if name == "tool_output_search_navigate": + return self._navigate_search(args) + if name == "tool_output_view": + return self._view_output(args) + if name == "tool_output_view_navigate": + return self._navigate_view(args) + raise ValueError(f"unknown local tool: {name}") + + def emergency_reset(self) -> bool: + if self.reset_count >= self.policy.max_resets: + return False + self.reset_count += 1 + recent = self.full_messages[-12:] + overview = json.dumps(recent, ensure_ascii=False) + overview = overview[:12_000] + recovery = { + "role": "user", + "content": ( + "[Context reset] The previous active context exceeded the model limit. " + "Continue the original task using this recent structural overview. Use history_search, " + "history_view, tool_output_search, or tool_output_view for omitted details.\n\n" + f"Recent history:\n{overview}" + ), + } + self.active_messages = [{"role": "user", "content": self.instruction}, recovery] + self.full_messages.append(recovery.copy()) + self._append_history(recovery, active=True) + self.latest_prompt_tokens = 0 + self.additions_since_prompt.clear() + self._event("emergency_reset", reset_count=self.reset_count) + return True + + +def create_managed_context( + *, + instruction: str, + artifact_dir: Path, + max_context_tokens: int, + max_output_tokens: int, + **policy_kwargs: Any, +) -> ManagedContext: + """Create the Toolathlon-inspired managed context for the generic runner.""" + return ManagedContext( + instruction=instruction, + artifact_dir=artifact_dir, + policy=ContextPolicy( + max_context_tokens=max_context_tokens, + max_output_tokens=max_output_tokens, + **policy_kwargs, + ), + ) diff --git a/examples/train_integrations/harbor/entrypoints/main_harbor.py b/examples/train_integrations/harbor/entrypoints/main_harbor.py index 128c9b4f9c..9a29f7d2b8 100644 --- a/examples/train_integrations/harbor/entrypoints/main_harbor.py +++ b/examples/train_integrations/harbor/entrypoints/main_harbor.py @@ -3,20 +3,21 @@ """ import sys +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any import ray import yaml -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any, Dict +from skyrl.train.config import GeneratorConfig, SkyRLTrainConfig from skyrl.train.entrypoints.main_base import BasePPOExp -from skyrl.train.config import SkyRLTrainConfig, GeneratorConfig, get_config_as_yaml_str from skyrl.train.utils import validate_cfg -from skyrl.train.utils.utils import initialize_ray from skyrl.train.utils.rate_limiter import RateLimiterConfig -from ..harbor_generator import HarborGenerator +from skyrl.train.utils.utils import initialize_ray + from ..dataset import HarborTaskDataset +from ..harbor_generator import HarborGenerator # NOTE (sumanthrh): We use a YAML to store the defaults for the Harbor trial configuration # TODO: Convert to a dataclass @@ -33,6 +34,32 @@ def _deep_merge(base: dict, overrides: dict) -> dict: return base +def _load_yaml_mapping(path: Path, description: str) -> dict: + """Load a YAML mapping, with a useful error for invalid trial configs.""" + with path.open() as f: + value = yaml.safe_load(f) + if value is None: + return {} + if not isinstance(value, dict): + raise TypeError(f"{description} must contain a YAML mapping, got {type(value).__name__}: {path}") + return value + + +def load_harbor_trial_config(cfg: "HarborSkyRLConfig") -> None: + """Merge Harbor defaults, an optional config file, and CLI overrides.""" + if not isinstance(cfg.harbor_trial_config, dict): + raise TypeError( + "harbor_trial_config must be a mapping of CLI overrides. " + "Use harbor_trial_config_path=/path/to/config.yaml to load a YAML file." + ) + + merged = _load_yaml_mapping(HARBOR_DEFAULT_CONFIG, "Harbor default config") + if cfg.harbor_trial_config_path: + config_path = Path(cfg.harbor_trial_config_path).expanduser() + merged = _deep_merge(merged, _load_yaml_mapping(config_path, "Harbor trial config")) + cfg.harbor_trial_config = _deep_merge(merged, cfg.harbor_trial_config) + + @dataclass class HarborGeneratorConfig(GeneratorConfig): """GeneratorConfig with Harbor-specific rate limiting.""" @@ -45,7 +72,8 @@ class HarborGeneratorConfig(GeneratorConfig): class HarborSkyRLConfig(SkyRLTrainConfig): """SkyRLTrainConfig with Harbor trial configuration.""" - harbor_trial_config: Dict[str, Any] = field(default_factory=dict) + harbor_trial_config: dict[str, Any] = field(default_factory=dict) + harbor_trial_config_path: str | None = None generator: HarborGeneratorConfig = field(default_factory=HarborGeneratorConfig) @@ -100,10 +128,7 @@ def skyrl_entrypoint(cfg): def main() -> None: cfg = HarborSkyRLConfig.from_cli_overrides(sys.argv[1:]) - # Load harbor defaults and merge CLI overrides on top - with open(HARBOR_DEFAULT_CONFIG) as f: - defaults = yaml.safe_load(f) - cfg.harbor_trial_config = _deep_merge(defaults, cfg.harbor_trial_config) + load_harbor_trial_config(cfg) validate_cfg(cfg) if cfg.trainer.algorithm.max_seq_len is None: diff --git a/examples/train_integrations/harbor/entrypoints/main_harbor_fully_async.py b/examples/train_integrations/harbor/entrypoints/main_harbor_fully_async.py index 11a46a5435..fbb7f37b83 100644 --- a/examples/train_integrations/harbor/entrypoints/main_harbor_fully_async.py +++ b/examples/train_integrations/harbor/entrypoints/main_harbor_fully_async.py @@ -6,17 +6,15 @@ ``examples/train/fully_async/main_fully_async.py`` for harbor. """ -import asyncio import sys import ray -import yaml from skyrl.train.fully_async_trainer import FullyAsyncRayPPOTrainer from skyrl.train.utils import validate_cfg from skyrl.train.utils.utils import initialize_ray -from .main_harbor import HARBOR_DEFAULT_CONFIG, HarborExp, HarborSkyRLConfig, _deep_merge +from .main_harbor import HarborExp, HarborSkyRLConfig, load_harbor_trial_config class HarborFullyAsyncExp(HarborExp): @@ -52,9 +50,7 @@ def skyrl_entrypoint(cfg): def main() -> None: cfg = HarborSkyRLConfig.from_cli_overrides(sys.argv[1:]) - with open(HARBOR_DEFAULT_CONFIG) as f: - defaults = yaml.safe_load(f) - cfg.harbor_trial_config = _deep_merge(defaults, cfg.harbor_trial_config) + load_harbor_trial_config(cfg) validate_cfg(cfg) if cfg.trainer.algorithm.max_seq_len is None: diff --git a/examples/train_integrations/harbor/entrypoints/main_harbor_generate.py b/examples/train_integrations/harbor/entrypoints/main_harbor_generate.py index 3bc5d534d2..eea52e8879 100644 --- a/examples/train_integrations/harbor/entrypoints/main_harbor_generate.py +++ b/examples/train_integrations/harbor/entrypoints/main_harbor_generate.py @@ -2,26 +2,24 @@ Main entrypoint for generating rollouts on Harbor tasks. For debugging purposes. """ +import asyncio import sys import ray -import asyncio -import yaml from loguru import logger -from skyrl.train.utils import validate_cfg -from skyrl.train.utils.utils import initialize_ray from skyrl.train.entrypoints.main_base import BasePPOExp from skyrl.train.generators.base import GeneratorInput, TrajectoryID -from ..harbor_generator import HarborGenerator +from skyrl.train.utils import validate_cfg +from skyrl.train.utils.utils import initialize_ray + from ..dataset import HarborTaskDataset +from ..harbor_generator import HarborGenerator from .main_harbor import ( HarborSkyRLConfig, - HARBOR_DEFAULT_CONFIG, - _deep_merge, + load_harbor_trial_config, ) - # For debugging purposes, we only generate a few samples. NUM_SAMPLES_TO_TEST = 10 @@ -106,10 +104,7 @@ def skyrl_entrypoint(cfg): def main() -> None: cfg = HarborSkyRLConfig.from_cli_overrides(sys.argv[1:]) - # Load harbor defaults and merge CLI overrides on top - with open(HARBOR_DEFAULT_CONFIG) as f: - defaults = yaml.safe_load(f) - cfg.harbor_trial_config = _deep_merge(defaults, cfg.harbor_trial_config) + load_harbor_trial_config(cfg) validate_cfg(cfg) if cfg.trainer.algorithm.max_seq_len is None: diff --git a/examples/train_integrations/harbor/harbor_generator.py b/examples/train_integrations/harbor/harbor_generator.py index 6bd325a4e0..7c9db93eec 100644 --- a/examples/train_integrations/harbor/harbor_generator.py +++ b/examples/train_integrations/harbor/harbor_generator.py @@ -2,10 +2,8 @@ import logging import os import time -from pathlib import Path from copy import deepcopy from dataclasses import dataclass -from typing import List, Optional from uuid import uuid4 # Suppress LiteLLM verbose logging @@ -17,7 +15,10 @@ from harbor.models.agent.rollout_detail import RolloutDetail from harbor.models.trial.config import TrialConfig from harbor.trial.trial import Trial -from skyrl.backends.skyrl_train.inference_servers.base import ConversationType, InferenceEngineInterface +from skyrl.backends.skyrl_train.inference_servers.base import ( + ConversationType, + InferenceEngineInterface, +) from skyrl.train.generators.base import ( GeneratorInput, GeneratorInterface, @@ -59,7 +60,7 @@ class HarborTrajectoryOutput: trajectory_id: TrajectoryID # Entire rollout_details list as returned by harbor's agent_result. None for failed trajectories # (agent_timeout / error) that we will mask in `build_step_wise_generator_output`. - rollout_details: Optional[List[RolloutDetail]] = None + rollout_details: list[RolloutDetail] | None = None reward: float = 0.0 num_turns: int = 0 # One of: "complete", "context_length", "agent_timeout", "error". Used by @@ -67,20 +68,20 @@ class HarborTrajectoryOutput: stop_reason: str = "complete" # End-to-end wall-clock time (seconds) to generate this trajectory. Optional: left as None if # timing was not recorded. - e2e_time: Optional[float] = None + e2e_time: float | None = None # Agent-reported detail, when the agent puts it in agent_result.metadata. Optional because # not every agent does (terminus-2 reports only n_episodes), so metrics below are emitted # only for the trajectories that carry them. - n_tool_calls: Optional[int] = None - agent_stop_reason: Optional[str] = None + n_tool_calls: int | None = None + agent_stop_reason: str | None = None # Target-tool recall, and the raw verifier coverage before any shaping. Both kept so the # unshaped number stays reportable and comparable to the benchmark. - tool_recall: Optional[float] = None - raw_reward: Optional[float] = None + tool_recall: float | None = None + raw_reward: float | None = None def build_step_wise_generator_output( - trajectory_outputs: List[HarborTrajectoryOutput], overlong_filtering: bool + trajectory_outputs: list[HarborTrajectoryOutput], overlong_filtering: bool ) -> GeneratorOutput: """Flatten per-trajectory rollout details into one entry per LLM turn. @@ -108,24 +109,24 @@ def build_step_wise_generator_output( masked_instance_ids = timeout_instance_ids | error_instance_ids # 2. Walk trajectories and emit one entry of GeneratorOutput per step. - prompt_token_ids: List[List[int]] = [] - response_ids: List[List[int]] = [] - rewards: List[float] = [] - loss_masks: List[List[int]] = [] - stop_reasons: List[str] = [] - is_last_step_list: List[bool] = [] - out_trajectory_ids: List[TrajectoryID] = [] - rollout_logprobs_list: List[List[float]] = [] - - successful_trajectories: List[HarborTrajectoryOutput] = [] - response_ids_for_metrics: List[List[int]] = [] - rewards_for_metrics: List[float] = [] + prompt_token_ids: list[list[int]] = [] + response_ids: list[list[int]] = [] + rewards: list[float] = [] + loss_masks: list[list[int]] = [] + stop_reasons: list[str] = [] + is_last_step_list: list[bool] = [] + out_trajectory_ids: list[TrajectoryID] = [] + rollout_logprobs_list: list[list[float]] = [] + + successful_trajectories: list[HarborTrajectoryOutput] = [] + response_ids_for_metrics: list[list[int]] = [] + rewards_for_metrics: list[float] = [] # One generation time per successful trajectory; used for completion-time metrics (avoids the # duplicate per-step entries below inflating the stats). - trajectory_generation_times_per_prompt: List[Optional[float]] = [] + trajectory_generation_times_per_prompt: list[float | None] = [] # One generation time per emitted step, aligned 1:1 with the flattened per-step arrays above. # Per trajectory we replicate its trajectory-level e2e_time across all of its steps. - out_trajectory_generation_times: List[Optional[float]] = [] + out_trajectory_generation_times: list[float | None] = [] for traj in trajectory_outputs: tid = traj.trajectory_id @@ -146,8 +147,9 @@ def build_step_wise_generator_output( successful_trajectories.append(traj) # 2.3. Check rollout_details expected format. - # Expect no summarization; rollout_details is a single linear chat segment from the main agent. - # TODO(Charlie): Support summarization. + # One segment contains the exact prompt sampled at each step. Prompts need not be + # prefix extensions: a context-managed agent may compact later prompts while the + # already sampled completion IDs and logprobs remain immutable. assert len(traj.rollout_details) == 1, f"Expected exactly one rollout segment, got {len(traj.rollout_details)}." rollout_detail = traj.rollout_details[0] prompt_token_ids_per_turn = rollout_detail["prompt_token_ids"] @@ -345,7 +347,7 @@ def __init__( rate_limit_config = getattr(generator_cfg, "rate_limit", None) self._rate_limiter = create_rate_limiter(rate_limit_config) - def _compute_cache_salt(self) -> Optional[str]: + def _compute_cache_salt(self) -> str | None: """Derive a prefix-cache salt from the current policy version. Mirrors ``SkyRLGymGenerator._compute_cache_salt``: keyed on the engine's ``weight_version`` and @@ -374,7 +376,7 @@ async def generate(self, input_batch: GeneratorInput, disable_tqdm: bool = False # Captured once so every trajectory shares the policy version at the start of the batch. cache_salt = self._compute_cache_salt() - all_outputs: List[HarborTrajectoryOutput] = [None] * len(prompts) # type: ignore[list-item] + all_outputs: list[HarborTrajectoryOutput] = [None] * len(prompts) # type: ignore[list-item] progress = tqdm( disable=disable_tqdm, # disable for fully async training total=len(prompts), @@ -403,7 +405,7 @@ async def _harbor_agent_loop( self, prompt: ConversationType, trajectory_id: TrajectoryID, - cache_salt: Optional[str] = None, + cache_salt: str | None = None, ) -> HarborTrajectoryOutput: """Run a single Harbor trial and return the rollout details plus a trajectory-level reward. Retries on unknown errors; context length errors train with reward=0; agent timeouts mask the trajectory. @@ -496,7 +498,7 @@ async def _harbor_agent_loop( break else: logger.warning(f"{prefix} failed: empty/missing rollout_details. Results: {results}") - except (asyncio.TimeoutError, TimeoutError): + except TimeoutError: # Must precede the generic handler, whose `continue` would retry and gamble # another full TRIAL_RUN_TIMEOUT_S while (in sync training) the entire batch # waits. wait_for has already cancelled the trial task; its container, if @@ -509,7 +511,7 @@ async def _harbor_agent_loop( f"wedged docker exec there is the usual cause." ) break - except Exception as e: + except Exception as e: # noqa: BLE001 - retry arbitrary Harbor/provider failures logger.warning(f"{prefix} failed: Error running trial: {e}. Results: {results}") continue finally: diff --git a/examples/train_integrations/harbor/mcp_agent.py b/examples/train_integrations/harbor/mcp_agent.py index d3458d41cb..7805f2932d 100644 --- a/examples/train_integrations/harbor/mcp_agent.py +++ b/examples/train_integrations/harbor/mcp_agent.py @@ -5,29 +5,37 @@ import asyncio import json import os -import shlex -from pathlib import Path, PurePosixPath +from importlib import import_module from typing import Any -from uuid import uuid4 from harbor.agents.base import BaseAgent from harbor.environments.base import BaseEnvironment from harbor.models.agent.context import AgentContext -from harbor.models.trial.paths import EnvironmentPaths -from .mcp_runner import run_loop +from .mcp_runner import ContextManagerFactory, run_loop +from .remote_mcp import RemoteMCPBridge + + +class ContextLengthExceededError(RuntimeError): + """Harbor-visible terminal context overflow with usable rollout details.""" + + +def _load_context_manager_factory(import_path: str | None) -> ContextManagerFactory | None: + if import_path is None: + return None + module_name, separator, attribute = import_path.partition(":") + if not separator or not module_name or not attribute: + raise ValueError("context_manager_factory must use module.path:attribute syntax") + factory = getattr(import_module(module_name), attribute) + if not callable(factory): + raise TypeError(f"context manager factory is not callable: {import_path}") + return factory class HarborMCPAgent(BaseAgent): """Run the model locally and dispatch configured stdio MCP tools remotely.""" SUPPORTS_ATIF = False - _REMOTE_DIR = PurePosixPath("/opt/harbor-mcp-agent") - _REMOTE_BRIDGE = _REMOTE_DIR / "bridge.py" - _REMOTE_SOCKET = _REMOTE_DIR / "bridge.sock" - _REMOTE_PID = _REMOTE_DIR / "bridge.pid" - _REMOTE_RPC_DIR = _REMOTE_DIR / "rpc" - _BRIDGE_LOG = EnvironmentPaths.agent_dir / "bridge.log" _MAX_TRAJECTORY_BYTES = 100_000_000 def __init__( @@ -42,6 +50,16 @@ def __init__( temperature: float = 0.0, collect_rollout_details: bool = False, strict_rollout_details: bool = False, + context_manager_factory: str | None = None, + max_context_tokens: int | None = None, + context_safety_tokens: int = 2048, + context_warning_ratio: float = 0.75, + context_compact_ratio: float = 0.85, + context_target_ratio: float = 0.70, + context_keep_recent_turns: int = 4, + context_keep_reasoning_turns: int = 1, + inline_tool_output_chars: int = 12_000, + max_context_resets: int = 2, extra_env: dict[str, str] | None = None, *args: Any, **kwargs: Any, @@ -57,15 +75,24 @@ def __init__( self.temperature = temperature self.collect_rollout_details = collect_rollout_details self.strict_rollout_details = strict_rollout_details + self.context_manager_factory = _load_context_manager_factory(context_manager_factory) + self.max_context_tokens = max_context_tokens + self.context_safety_tokens = context_safety_tokens + self.context_warning_ratio = context_warning_ratio + self.context_compact_ratio = context_compact_ratio + self.context_target_ratio = context_target_ratio + self.context_keep_recent_turns = context_keep_recent_turns + self.context_keep_reasoning_turns = context_keep_reasoning_turns + self.inline_tool_output_chars = inline_tool_output_chars + self.max_context_resets = max_context_resets self._bridge_tools: list[dict[str, Any]] = [] + self._remote_bridge = RemoteMCPBridge(self.logs_dir) if not self.api_base: raise ValueError("HarborMCPAgent requires api_base or OPENAI_BASE_URL") if not self.model_name: raise ValueError("HarborMCPAgent requires agent.model_name") if len(self.mcp_servers) != 1 or self.mcp_servers[0].transport != "stdio": - raise ValueError( - "HarborMCPAgent currently requires exactly one stdio MCP server" - ) + raise ValueError("HarborMCPAgent currently requires exactly one stdio MCP server") if not self.mcp_servers[0].command: raise ValueError("HarborMCPAgent requires an MCP server command") @@ -74,76 +101,17 @@ def name() -> str: return "harbor-mcp" def version(self) -> str | None: - return "2" + return "3" def _request_model(self) -> str: assert self.model_name is not None return self.model_name.split("/", 1)[-1] async def setup(self, environment: BaseEnvironment) -> None: - local_bridge = Path(__file__).with_name("mcp_bridge.py") - directories = [ - self._REMOTE_DIR.as_posix(), - self._REMOTE_RPC_DIR.as_posix(), - EnvironmentPaths.agent_dir.as_posix(), - ] - result = await environment.exec( - "mkdir -p " + " ".join(shlex.quote(path) for path in directories), - user="root", - timeout_sec=30, - ) - if result.return_code != 0: - raise RuntimeError( - result.stderr - or result.stdout - or "failed to create MCP bridge directory" - ) - await environment.upload_file(local_bridge, self._REMOTE_BRIDGE.as_posix()) - - server = self.mcp_servers[0] - command = [ - "python3", - self._REMOTE_BRIDGE.as_posix(), - "serve", - "--socket", - self._REMOTE_SOCKET.as_posix(), - "--mcp-command", - server.command or "", - ] - for arg in server.args: - command.append(f"--mcp-arg={arg}") - shell_command = " ".join(shlex.quote(part) for part in command) - start = await environment.exec( - f"rm -f {shlex.quote(self._REMOTE_SOCKET.as_posix())}; " - f"nohup {shell_command} > {shlex.quote(self._BRIDGE_LOG.as_posix())} 2>&1 " - f"< /dev/null & echo $! > {shlex.quote(self._REMOTE_PID.as_posix())}", - user="root", - timeout_sec=30, - ) - if start.return_code != 0: - raise RuntimeError( - start.stderr or start.stdout or "failed to start MCP bridge" - ) + await self._remote_bridge.setup(environment, self.mcp_servers[0]) + self._bridge_tools = self._remote_bridge.tools - last_error: Exception | None = None - for _ in range(30): - try: - response = await self._bridge_rpc( - environment, {"op": "list_tools"}, timeout_sec=30 - ) - tools = response.get("tools") - if not isinstance(tools, list): - raise TypeError("MCP bridge list_tools response omitted tools") - self._bridge_tools = tools - return - except Exception as exc: # noqa: BLE001 - readiness retries include provider errors - last_error = exc - await asyncio.sleep(1) - raise RuntimeError(f"MCP bridge did not become ready: {last_error}") - - async def run( - self, instruction: str, environment: BaseEnvironment, context: AgentContext - ) -> None: + async def run(self, instruction: str, environment: BaseEnvironment, context: AgentContext) -> None: trajectory: dict[str, Any] | None = None async def call_tool(name: str, arguments: dict[str, Any]) -> dict[str, Any]: @@ -174,6 +142,17 @@ async def checkpoint(value: dict[str, Any]) -> None: collect_rollout_details=self.collect_rollout_details, strict_rollout_details=self.strict_rollout_details, session_id=self.session_id, + context_manager_factory=self.context_manager_factory, + artifact_dir=self.logs_dir / "context", + max_context_tokens=self.max_context_tokens, + context_safety_tokens=self.context_safety_tokens, + context_warning_ratio=self.context_warning_ratio, + context_compact_ratio=self.context_compact_ratio, + context_target_ratio=self.context_target_ratio, + context_keep_recent_turns=self.context_keep_recent_turns, + context_keep_reasoning_turns=self.context_keep_reasoning_turns, + inline_tool_output_chars=self.inline_tool_output_chars, + max_context_resets=self.max_context_resets, ) finally: shutdown = asyncio.create_task(self._shutdown_bridge(environment)) @@ -186,22 +165,13 @@ async def checkpoint(value: dict[str, Any]) -> None: exc, ) + if trajectory and trajectory.get("stop_reason") == "context_length": + raise ContextLengthExceededError(str(trajectory.get("error") or "context recovery exhausted")) if trajectory and trajectory.get("error"): raise RuntimeError(str(trajectory["error"])) async def _shutdown_bridge(self, environment: BaseEnvironment) -> None: - await self._bridge_rpc(environment, {"op": "shutdown"}, timeout_sec=30) - pid_file = shlex.quote(self._REMOTE_PID.as_posix()) - wait = await environment.exec( - f"pid=$(cat {pid_file}); " - 'while kill -0 "$pid" 2>/dev/null; do sleep 0.1; done', - user="root", - timeout_sec=30, - ) - if wait.return_code != 0: - raise RuntimeError( - wait.stderr or wait.stdout or "MCP bridge did not stop cleanly" - ) + await self._remote_bridge.shutdown(environment) async def _bridge_rpc( self, @@ -210,50 +180,7 @@ async def _bridge_rpc( *, timeout_sec: int, ) -> dict[str, Any]: - rpc_id = uuid4().hex - local_request = self.logs_dir / f"bridge-request-{rpc_id}.json" - local_response = self.logs_dir / f"bridge-response-{rpc_id}.json" - remote_request = self._REMOTE_RPC_DIR / f"{rpc_id}.request.json" - remote_response = self._REMOTE_RPC_DIR / f"{rpc_id}.response.json" - local_request.write_text(json.dumps(request, ensure_ascii=False)) - try: - await environment.upload_file(local_request, remote_request.as_posix()) - command = " ".join( - shlex.quote(part) - for part in [ - "python3", - self._REMOTE_BRIDGE.as_posix(), - "request", - "--socket", - self._REMOTE_SOCKET.as_posix(), - "--request-file", - remote_request.as_posix(), - "--response-file", - remote_response.as_posix(), - ] - ) - result = await environment.exec( - command, user="root", timeout_sec=timeout_sec - ) - if result.return_code != 0: - raise RuntimeError( - result.stderr or result.stdout or "MCP bridge request failed" - ) - await environment.download_file(remote_response.as_posix(), local_response) - if local_response.stat().st_size > self._MAX_TRAJECTORY_BYTES: - raise RuntimeError("MCP bridge response exceeds 100 MB") - envelope = json.loads(local_response.read_text()) - if envelope.get("status") != "ok": - raise RuntimeError( - str(envelope.get("error") or "MCP bridge returned an error") - ) - payload = envelope.get("result") - if not isinstance(payload, dict): - raise TypeError("MCP bridge returned an invalid result") - return payload - finally: - local_request.unlink(missing_ok=True) - local_response.unlink(missing_ok=True) + return await self._remote_bridge.rpc(environment, request, timeout_sec=timeout_sec) def _write_trajectory(self, trajectory: dict[str, Any]) -> None: local = self.logs_dir / "trajectory.json" @@ -276,6 +203,9 @@ def _apply_trajectory(trajectory: dict[str, Any], context: AgentContext) -> None "agent_stop_reason": trajectory.get("stop_reason", "error"), "all_messages": trajectory.get("messages") or [], "elapsed_seconds": trajectory.get("elapsed_seconds"), + "context_events": trajectory.get("context_events") or [], + "active_messages": trajectory.get("active_messages") or [], + "context_artifact_dir": trajectory.get("context_artifact_dir"), } details = trajectory.get("rollout_details") if details: diff --git a/examples/train_integrations/harbor/mcp_bridge.py b/examples/train_integrations/harbor/mcp_bridge.py index 093f932d76..b0634dbcff 100644 --- a/examples/train_integrations/harbor/mcp_bridge.py +++ b/examples/train_integrations/harbor/mcp_bridge.py @@ -17,7 +17,7 @@ def _content_json(block: Any) -> dict[str, Any]: if hasattr(block, "model_dump"): - return block.model_dump(mode="json") + return block.model_dump(mode="json", exclude_none=True) text = getattr(block, "text", None) return {"type": "text", "text": str(text if text is not None else block)} @@ -41,9 +41,7 @@ async def serve(args: argparse.Namespace) -> None: ): await session.initialize() - async def handle( - reader: asyncio.StreamReader, writer: asyncio.StreamWriter - ) -> None: + async def handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: is_shutdown = False try: raw = await reader.readline() @@ -66,8 +64,7 @@ async def handle( { "name": tool.name, "description": tool.description or "", - "inputSchema": tool.inputSchema - or {"type": "object", "properties": {}}, + "inputSchema": tool.inputSchema or {"type": "object", "properties": {}}, } for tool in listed.tools ] @@ -77,18 +74,13 @@ async def handle( name = request_data.get("name") arguments = request_data.get("arguments") if not isinstance(name, str) or not isinstance(arguments, dict): - raise ValueError( - "call_tool requires string name and object arguments" - ) + raise ValueError("call_tool requires string name and object arguments") async with session_lock: tool_result = await session.call_tool(name, arguments) result = { "status": "ok", "result": { - "content": [ - _content_json(block) - for block in tool_result.content or [] - ], + "content": [_content_json(block) for block in tool_result.content or []], "is_error": bool(tool_result.isError), }, } @@ -110,9 +102,7 @@ async def handle( if is_shutdown: shutdown_reply_sent.set() - server = await asyncio.start_unix_server( - handle, path=socket_path, limit=_MAX_MESSAGE_BYTES - ) + server = await asyncio.start_unix_server(handle, path=socket_path, limit=_MAX_MESSAGE_BYTES) await stop.wait() server.close() await server.wait_closed() @@ -127,9 +117,7 @@ async def request(args: argparse.Namespace) -> None: raw = request_path.read_bytes() if len(raw) > _MAX_MESSAGE_BYTES: raise ValueError("bridge request exceeds 100 MB") - reader, writer = await asyncio.open_unix_connection( - args.socket, limit=_MAX_MESSAGE_BYTES - ) + reader, writer = await asyncio.open_unix_connection(args.socket, limit=_MAX_MESSAGE_BYTES) writer.write(raw.rstrip(b"\n") + b"\n") await writer.drain() response = await reader.readline() diff --git a/examples/train_integrations/harbor/mcp_runner.py b/examples/train_integrations/harbor/mcp_runner.py index fefcca68ff..110fa5f597 100644 --- a/examples/train_integrations/harbor/mcp_runner.py +++ b/examples/train_integrations/harbor/mcp_runner.py @@ -6,7 +6,8 @@ import json import time from collections.abc import Awaitable, Callable -from typing import Any +from pathlib import Path +from typing import Any, Protocol import httpx @@ -14,6 +15,50 @@ Checkpoint = Callable[[dict[str, Any]], Awaitable[None]] +class ManagedContext(Protocol): + """Context-management contract optionally supplied by an integration.""" + + artifact_dir: Path + full_messages: list[dict[str, Any]] + active_messages: list[dict[str, Any]] + events: list[dict[str, Any]] + tool_schemas: list[dict[str, Any]] + + def prepare_for_request(self) -> None: ... + + def observe_prompt_tokens(self, tokens: int) -> None: ... + + def append(self, message: dict[str, Any]) -> None: ... + + def append_tool_result(self, message: dict[str, Any]) -> dict[str, Any]: ... + + def call_local_tool(self, name: str, args: dict[str, Any]) -> dict[str, Any]: ... + + def apply_pending(self) -> None: ... + + def emergency_reset(self) -> bool: ... + + +ContextManagerFactory = Callable[..., ManagedContext] + + +class ContextLengthExceeded(RuntimeError): + """The model endpoint rejected a request because its context was too long.""" + + +def _is_context_length_response(response: httpx.Response) -> bool: + if response.status_code not in {400, 413, 422}: + return False + text = response.text.casefold() + return any( + marker in text + for marker in ( + "context length", "context_length", "maximum context", + "max_model_len", "maximum number of tokens", "too many tokens", + ) + ) + + def openai_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: """Convert bridge MCP tool descriptions to OpenAI function tools.""" return [ @@ -100,6 +145,10 @@ async def completion( response.status_code not in {408, 409, 429} and response.status_code < 500 ): + if _is_context_length_response(response): + raise ContextLengthExceeded( + f"model API context limit: {response.text[:1000]}" + ) response.raise_for_status() return response.json() last_error = RuntimeError( @@ -124,12 +173,16 @@ def _trajectory( prompt_ids: list[list[int]], completion_ids: list[list[int]], logprobs: list[list[float]], + full_messages: list[dict[str, Any]] | None = None, + active_messages: list[dict[str, Any]] | None = None, + context_events: list[dict[str, Any]] | None = None, + context_artifact_dir: str | None = None, error: str | None = None, ) -> dict[str, Any]: trajectory: dict[str, Any] = { - "schema_version": 2, + "schema_version": 3, "model": model, - "messages": messages, + "messages": full_messages if full_messages is not None else messages, "tools": tools, "tool_calls": calls_log, "usage": usage, @@ -137,6 +190,12 @@ def _trajectory( "stop_reason": stop_reason, "elapsed_seconds": time.monotonic() - started, } + if active_messages is not None: + trajectory["active_messages"] = active_messages + if context_events is not None: + trajectory["context_events"] = context_events + if context_artifact_dir is not None: + trajectory["context_artifact_dir"] = context_artifact_dir if error: trajectory["error"] = error if prompt_ids or completion_ids or logprobs: @@ -165,15 +224,53 @@ async def run_loop( collect_rollout_details: bool, strict_rollout_details: bool, session_id: str | None = None, + context_manager_factory: ContextManagerFactory | None = None, + artifact_dir: Path | None = None, + max_context_tokens: int | None = None, + context_safety_tokens: int = 2048, + context_warning_ratio: float = 0.75, + context_compact_ratio: float = 0.85, + context_target_ratio: float = 0.70, + context_keep_recent_turns: int = 4, + context_keep_reasoning_turns: int = 1, + inline_tool_output_chars: int = 12_000, + max_context_resets: int = 2, ) -> dict[str, Any]: """Run the model locally while dispatching tool actions through a bridge.""" + managed: ManagedContext | None = None messages: list[dict[str, Any]] = [{"role": "user", "content": instruction}] + if max_context_tokens is not None: + if context_manager_factory is None: + raise ValueError("context_manager_factory is required when context management is enabled") + if artifact_dir is None: + raise ValueError("artifact_dir is required when context management is enabled") + managed = context_manager_factory( + instruction=instruction, + artifact_dir=artifact_dir, + max_context_tokens=max_context_tokens, + max_output_tokens=max_tokens, + safety_tokens=context_safety_tokens, + warning_ratio=context_warning_ratio, + compact_ratio=context_compact_ratio, + target_ratio=context_target_ratio, + keep_recent_turns=context_keep_recent_turns, + keep_reasoning_turns=context_keep_reasoning_turns, + inline_tool_output_chars=inline_tool_output_chars, + max_resets=max_context_resets, + ) + messages = managed.active_messages calls_log: list[dict[str, Any]] = [] prompt_ids_per_turn: list[list[int]] = [] completion_ids_per_turn: list[list[int]] = [] logprobs_per_turn: list[list[float]] = [] usage = {"prompt_tokens": 0, "completion_tokens": 0, "cached_tokens": 0} - tools = openai_tools(bridge_tools) + remote_names = {tool["name"] for tool in bridge_tools} + local_schemas = managed.tool_schemas if managed else [] + local_names = {tool["name"] for tool in local_schemas} + collisions = remote_names & local_names + if collisions: + raise ValueError(f"MCP tools collide with context tools: {sorted(collisions)}") + tools = openai_tools(bridge_tools + local_schemas) tool_names = {tool["function"]["name"] for tool in tools} started = time.monotonic() stop_reason = "running" @@ -191,6 +288,10 @@ def snapshot() -> dict[str, Any]: prompt_ids=prompt_ids_per_turn, completion_ids=completion_ids_per_turn, logprobs=logprobs_per_turn, + full_messages=managed.full_messages if managed else None, + active_messages=managed.active_messages if managed else None, + context_events=managed.events if managed else None, + context_artifact_dir=str(managed.artifact_dir) if managed else None, error=error, ) @@ -200,18 +301,30 @@ def snapshot() -> dict[str, Any]: if time.monotonic() - started >= deadline_sec: stop_reason = "deadline" break - response = await completion( - client, - api_base=api_base, - api_key=api_key, - model=model, - messages=messages, - tools=tools, - max_tokens=max_tokens, - temperature=temperature, - collect_rollout_details=collect_rollout_details, - session_id=session_id, - ) + if managed: + managed.prepare_for_request() + messages = managed.active_messages + try: + response = await completion( + client, + api_base=api_base, + api_key=api_key, + model=model, + messages=messages, + tools=tools, + max_tokens=max_tokens, + temperature=temperature, + collect_rollout_details=collect_rollout_details, + session_id=session_id, + ) + except ContextLengthExceeded: + if managed and managed.emergency_reset(): + messages = managed.active_messages + await checkpoint(snapshot()) + continue + stop_reason = "context_length" + error = "ContextLengthExceededError: context recovery exhausted" + break choice = (response.get("choices") or [{}])[0] message = choice.get("message") or {} assistant: dict[str, Any] = { @@ -223,16 +336,22 @@ def snapshot() -> dict[str, Any]: tool_calls = message.get("tool_calls") or [] if tool_calls: assistant["tool_calls"] = tool_calls - messages.append(assistant) - turn_usage = response.get("usage") or {} + if managed: + # The reported prompt usage describes the request before this assistant + # message; calibrate first, then track the assistant as new context. + managed.observe_prompt_tokens(int(turn_usage.get("prompt_tokens") or 0)) + managed.append(assistant) + messages = managed.active_messages + else: + messages.append(assistant) + usage["prompt_tokens"] += int(turn_usage.get("prompt_tokens") or 0) usage["completion_tokens"] += int( turn_usage.get("completion_tokens") or 0 ) details = turn_usage.get("prompt_tokens_details") or {} usage["cached_tokens"] += int(details.get("cached_tokens") or 0) - if collect_rollout_details: prompt_ids, completion_ids, logprobs = token_data(response) if strict_rollout_details and ( @@ -281,22 +400,35 @@ def snapshot() -> dict[str, Any]: raise TypeError("tool arguments must decode to an object") if name not in tool_names: raise ValueError(f"unknown tool: {name}") - result = await call_tool(name, parsed_args) - result_text = tool_result_text(result) - record["is_error"] = bool(result.get("is_error")) + if managed and name in local_names: + result_text = json.dumps(managed.call_local_tool(name, parsed_args), ensure_ascii=False) + record["is_error"] = False + else: + result = await call_tool(name, parsed_args) + result_text = tool_result_text(result) + record["is_error"] = bool(result.get("is_error")) except Exception as exc: # noqa: BLE001 - tool errors are model-visible result_text = f"Error: {type(exc).__name__}: {exc}" record["is_error"] = True record["result"] = result_text calls_log.append(record) - messages.append( - { - "role": "tool", - "tool_call_id": call_id, - "name": name, - "content": result_text, - } - ) + tool_message = { + "role": "tool", + "tool_call_id": call_id, + "name": name, + "content": result_text, + } + if managed: + compact = managed.append_tool_result(tool_message) + record["result"] = compact["content"] + if compact.get("artifact_id"): + record["artifact_id"] = compact["artifact_id"] + messages = managed.active_messages + else: + messages.append(tool_message) + if managed: + managed.apply_pending() + messages = managed.active_messages await checkpoint(snapshot()) else: stop_reason = "max_turns" diff --git a/examples/train_integrations/harbor/remote_mcp.py b/examples/train_integrations/harbor/remote_mcp.py new file mode 100644 index 0000000000..dc3275b943 --- /dev/null +++ b/examples/train_integrations/harbor/remote_mcp.py @@ -0,0 +1,145 @@ +"""Host-side client for a stateful stdio MCP server in a Harbor environment.""" + +from __future__ import annotations + +import json +import shlex +from pathlib import Path, PurePosixPath +from typing import Any +from uuid import uuid4 + +from harbor.environments.base import BaseEnvironment +from harbor.models.task.config import MCPServerConfig +from harbor.models.trial.paths import EnvironmentPaths + + +class RemoteMCPBridge: + """Own the remote MCP bridge lifecycle and its file-backed RPC protocol.""" + + _REMOTE_DIR = PurePosixPath("/opt/harbor-mcp-agent") + _REMOTE_BRIDGE = _REMOTE_DIR / "bridge.py" + _REMOTE_SOCKET = _REMOTE_DIR / "bridge.sock" + _REMOTE_PID = _REMOTE_DIR / "bridge.pid" + _REMOTE_RPC_DIR = _REMOTE_DIR / "rpc" + _BRIDGE_LOG = EnvironmentPaths.agent_dir / "bridge.log" + _MAX_RESPONSE_BYTES = 100_000_000 + + def __init__(self, logs_dir: Path) -> None: + self.logs_dir = logs_dir + self.tools: list[dict[str, Any]] = [] + self._started = False + + async def setup(self, environment: BaseEnvironment, server: MCPServerConfig) -> None: + local_bridge = Path(__file__).with_name("mcp_bridge.py") + directories = [ + self._REMOTE_DIR.as_posix(), + self._REMOTE_RPC_DIR.as_posix(), + EnvironmentPaths.agent_dir.as_posix(), + ] + result = await environment.exec( + "mkdir -p " + " ".join(shlex.quote(path) for path in directories), + user="root", + timeout_sec=30, + ) + if result.return_code != 0: + raise RuntimeError(result.stderr or result.stdout or "failed to create MCP bridge directory") + await environment.upload_file(local_bridge, self._REMOTE_BRIDGE.as_posix()) + + command = [ + "python3", + self._REMOTE_BRIDGE.as_posix(), + "serve", + "--socket", + self._REMOTE_SOCKET.as_posix(), + "--mcp-command", + server.command or "", + ] + for arg in server.args: + command.append(f"--mcp-arg={arg}") + shell_command = " ".join(shlex.quote(part) for part in command) + start = await environment.exec( + f"rm -f {shlex.quote(self._REMOTE_SOCKET.as_posix())}; " + f"nohup {shell_command} > {shlex.quote(self._BRIDGE_LOG.as_posix())} 2>&1 " + f"< /dev/null & echo $! > {shlex.quote(self._REMOTE_PID.as_posix())}", + user="root", + timeout_sec=30, + ) + if start.return_code != 0: + raise RuntimeError(start.stderr or start.stdout or "failed to start MCP bridge") + self._started = True + + last_error: Exception | None = None + for _ in range(30): + try: + response = await self.rpc(environment, {"op": "list_tools"}, timeout_sec=30) + tools = response.get("tools") + if not isinstance(tools, list): + raise TypeError("MCP bridge list_tools response omitted tools") + self.tools = tools + return + except Exception as exc: # noqa: BLE001 - readiness includes remote errors + last_error = exc + import asyncio + + await asyncio.sleep(1) + raise RuntimeError(f"MCP bridge did not become ready: {last_error}") + + async def shutdown(self, environment: BaseEnvironment) -> None: + if not self._started: + return + self._started = False + await self.rpc(environment, {"op": "shutdown"}, timeout_sec=30) + pid_file = shlex.quote(self._REMOTE_PID.as_posix()) + wait = await environment.exec( + f"pid=$(cat {pid_file}); " 'while kill -0 "$pid" 2>/dev/null; do sleep 0.1; done', + user="root", + timeout_sec=30, + ) + if wait.return_code != 0: + raise RuntimeError(wait.stderr or wait.stdout or "MCP bridge did not stop cleanly") + + async def rpc( + self, + environment: BaseEnvironment, + request: dict[str, Any], + *, + timeout_sec: int, + ) -> dict[str, Any]: + rpc_id = uuid4().hex + local_request = self.logs_dir / f"bridge-request-{rpc_id}.json" + local_response = self.logs_dir / f"bridge-response-{rpc_id}.json" + remote_request = self._REMOTE_RPC_DIR / f"{rpc_id}.request.json" + remote_response = self._REMOTE_RPC_DIR / f"{rpc_id}.response.json" + local_request.write_text(json.dumps(request, ensure_ascii=False)) + try: + await environment.upload_file(local_request, remote_request.as_posix()) + command = " ".join( + shlex.quote(part) + for part in [ + "python3", + self._REMOTE_BRIDGE.as_posix(), + "request", + "--socket", + self._REMOTE_SOCKET.as_posix(), + "--request-file", + remote_request.as_posix(), + "--response-file", + remote_response.as_posix(), + ] + ) + result = await environment.exec(command, user="root", timeout_sec=timeout_sec) + if result.return_code != 0: + raise RuntimeError(result.stderr or result.stdout or "MCP bridge request failed") + await environment.download_file(remote_response.as_posix(), local_response) + if local_response.stat().st_size > self._MAX_RESPONSE_BYTES: + raise RuntimeError("MCP bridge response exceeds 100 MB") + envelope = json.loads(local_response.read_text()) + if envelope.get("status") != "ok": + raise RuntimeError(str(envelope.get("error") or "MCP bridge returned an error")) + payload = envelope.get("result") + if not isinstance(payload, dict): + raise TypeError("MCP bridge returned an invalid result") + return payload + finally: + local_request.unlink(missing_ok=True) + local_response.unlink(missing_ok=True) diff --git a/pyproject.toml b/pyproject.toml index ba0fbbe484..98653a84fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -322,7 +322,7 @@ torchvision = [ ] # AfterQuery's Harbor fork carries the RL rollout extensions and existing # raw-HTTP Compute provider used by the training integrations in this repository. -harbor = { git = "ssh://git@github.com/AfterQuery/harbor-aq.git", rev = "ff8961ece97f744e3c7e373c0b410d87b01cf8f7" } +harbor = { git = "ssh://git@github.com/AfterQuery/harbor-aq.git", rev = "1627cfb2295fa9b42d70e19dffbecc6f3243eaad" } megatron-bridge = {git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge", rev = "91a15142a4b4442a8d46ab539d1b923bd08570d0", marker = "sys_platform == 'linux'"} # megatron-core main branch: https://github.com/NVIDIA/Megatron-LM/tree/main latest as of 6/8/26 megatron-core = {git = "https://github.com/NVIDIA/Megatron-LM", rev = "71e418ea7d7b3a6c9a53238c543c3e0b43e11026", marker = "sys_platform == 'linux'"} diff --git a/tests/train/test_harbor_mcp_agent.py b/tests/train/test_harbor_mcp_agent.py index 3454d2a428..0975f29427 100644 --- a/tests/train/test_harbor_mcp_agent.py +++ b/tests/train/test_harbor_mcp_agent.py @@ -5,6 +5,10 @@ import httpx import pytest +from examples.train.toolathlon_harbor.toolathlon_context_manager import ( + create_managed_context, +) + RUNNER = Path(__file__).parents[2] / "examples/train_integrations/harbor/mcp_runner.py" SPEC = importlib.util.spec_from_file_location("harbor_mcp_runner", RUNNER) assert SPEC and SPEC.loader @@ -89,9 +93,89 @@ def handler(request: httpx.Request) -> httpx.Response: assert seen["payload"]["return_token_ids"] is True +@pytest.mark.asyncio +async def test_completion_classifies_vllm_context_rejection(): + def handler(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + json={"error": {"message": "maximum context length is 32768 tokens"}}, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as client: + with pytest.raises(runner.ContextLengthExceeded): + await runner.completion( + client, + api_base="http://model/v1", + api_key="dummy", + model="model", + messages=[{"role": "user", "content": "task"}], + tools=[], + max_tokens=32, + temperature=0.0, + collect_rollout_details=False, + session_id=None, + ) + + +@pytest.mark.asyncio +async def test_run_loop_recovers_from_context_rejection(monkeypatch, tmp_path): + requests = [] + + async def fake_completion(_client, **kwargs): + requests.append(kwargs["messages"]) + if len(requests) == 1: + raise runner.ContextLengthExceeded("too long") + return { + "choices": [{"message": {"role": "assistant", "content": "done"}}], + "usage": {"prompt_tokens": 20, "completion_tokens": 1}, + } + + checkpoints = [] + + async def checkpoint(value): + checkpoints.append(value) + + async def call_tool(_name, _arguments): + raise AssertionError("no tool call expected") + + monkeypatch.setattr(runner, "completion", fake_completion) + result = await runner.run_loop( + instruction="task", + bridge_tools=[], + call_tool=call_tool, + checkpoint=checkpoint, + api_base="http://model/v1", + api_key="dummy", + model="model", + max_turns=3, + deadline_sec=60, + request_timeout_sec=30, + max_tokens=32, + temperature=0.0, + collect_rollout_details=False, + strict_rollout_details=False, + artifact_dir=tmp_path / "context", + context_manager_factory=create_managed_context, + max_context_tokens=256, + context_safety_tokens=16, + ) + + assert result["stop_reason"] == "complete" + assert any(event["type"] == "emergency_reset" for event in result["context_events"]) + assert "Context reset" in requests[1][1]["content"] + assert checkpoints + + def test_agent_uses_environment_rpc_without_container_model_credentials(): agent_source = (RUNNER.parent / "mcp_agent.py").read_text() - assert "environment.upload_file" in agent_source - assert "environment.download_file" in agent_source + bridge_source = (RUNNER.parent / "remote_mcp.py").read_text() + assert "environment.upload_file" in bridge_source + assert "environment.download_file" in bridge_source + assert "RemoteMCPBridge" in agent_source assert "container_reachable_url" not in agent_source assert 'env={"OPENAI_BASE_URL"' not in agent_source + + +def test_remote_bridge_omits_null_mcp_content_fields(): + bridge_source = (RUNNER.parent / "mcp_bridge.py").read_text() + assert 'model_dump(mode="json", exclude_none=True)' in bridge_source diff --git a/tests/train/test_toolathlon_context_manager.py b/tests/train/test_toolathlon_context_manager.py new file mode 100644 index 0000000000..fee1805a49 --- /dev/null +++ b/tests/train/test_toolathlon_context_manager.py @@ -0,0 +1,148 @@ +import json + +from examples.train.toolathlon_harbor.toolathlon_context_manager import ( + ContextPolicy, + ManagedContext, +) + + +def _manager(tmp_path, **overrides): + values = { + "max_context_tokens": 10_000, + "max_output_tokens": 1_000, + "safety_tokens": 0, + "warning_ratio": 0.50, + "compact_ratio": 0.60, + "target_ratio": 0.40, + "keep_recent_turns": 2, + "keep_reasoning_turns": 1, + "inline_tool_output_chars": 100, + "preview_chars": 20, + } + values.update(overrides) + return ManagedContext("do the task", tmp_path / "context", ContextPolicy(**values)) + + +def _append_exchange(manager, index, size=600): + manager.append( + { + "role": "assistant", + "content": f"turn {index}", + "reasoning_content": "r" * size, + "tool_calls": [{"id": f"call-{index}", "function": {"name": "work", "arguments": "{}"}}], + } + ) + manager.append_tool_result( + {"role": "tool", "tool_call_id": f"call-{index}", "name": "work", "content": "x" * size} + ) + + +def test_automatic_compaction_prunes_reasoning_and_atomic_exchanges(tmp_path): + manager = _manager(tmp_path) + for index in range(6): + _append_exchange(manager, index) + + full_before = json.loads(json.dumps(manager.full_messages)) + manager.latest_prompt_tokens = 8_000 + manager.prepare_for_request() + + assert manager.full_messages == full_before + active_assistants = [m for m in manager.active_messages if m["role"] == "assistant"] + assert len(active_assistants) >= 2 + assert "reasoning_content" not in active_assistants[0] + calls = {call["id"] for m in active_assistants for call in m.get("tool_calls", [])} + results = {m["tool_call_id"] for m in manager.active_messages if m["role"] == "tool"} + assert calls == results + assert any(event["type"] == "automatic_compaction" for event in manager.events) + + +def test_oversized_tool_output_is_offloaded_and_searchable(tmp_path): + manager = _manager(tmp_path) + content = "prefix NEEDLE " + "z" * 200 + " NEEDLE " + "q" * 200 + compact = manager.append_tool_result( + {"role": "tool", "tool_call_id": "call-1", "name": "work", "content": content} + ) + + artifact_id = compact["artifact_id"] + assert "artifact_id" not in manager.active_messages[-1] + assert len(compact["content"]) < len(content) + 200 + found = manager.call_local_tool( + "tool_output_search", + {"artifact_id": artifact_id, "pattern": "needle", "page_size": 1, "context_size": 20}, + ) + assert found["results"] + assert found["total_matches"] == 2 + assert found["current_page"] == 1 + next_search = manager.call_local_tool( + "tool_output_search_navigate", + {"search_session_id": found["search_session_id"], "action": "next_page"}, + ) + assert next_search["current_page"] == 2 + viewed = manager.call_local_tool( + "tool_output_view", {"artifact_id": artifact_id, "page_size": 100} + ) + assert viewed["content"] == content[:100] + next_view = manager.call_local_tool( + "tool_output_view_navigate", + {"view_session_id": viewed["view_session_id"], "action": "next_page"}, + ) + assert next_view["content"] == content[100:200] + + +def test_tool_output_search_validates_toolathlon_limits(tmp_path): + manager = _manager(tmp_path) + compact = manager.append_tool_result( + {"role": "tool", "tool_call_id": "call-1", "name": "work", "content": "x" * 200} + ) + artifact_id = compact["artifact_id"] + + for args in ( + {"artifact_id": artifact_id, "pattern": "["}, + {"artifact_id": artifact_id, "pattern": "x", "page_size": 51}, + {"artifact_id": artifact_id, "pattern": "x", "context_size": 0}, + ): + try: + manager.call_local_tool("tool_output_search", args) + except ValueError: + pass + else: + raise AssertionError(f"invalid search arguments accepted: {args}") + + +def test_manual_compaction_is_scheduled_until_exchange_finishes(tmp_path): + manager = _manager(tmp_path) + for index in range(5): + _append_exchange(manager, index, size=20) + before = list(manager.active_messages) + + result = manager.call_local_tool( + "context_manage", {"method": "delete_first_turns", "value": 2} + ) + assert result["status"] == "scheduled" + assert manager.active_messages == before + manager.apply_pending() + assert len([m for m in manager.active_messages if m["role"] == "assistant"]) == 3 + + +def test_emergency_reset_is_bounded_and_preserves_durable_history(tmp_path): + manager = _manager(tmp_path, max_resets=2) + _append_exchange(manager, 1, size=20) + full_before = list(manager.full_messages) + + assert manager.emergency_reset() is True + assert manager.emergency_reset() is True + assert manager.emergency_reset() is False + assert manager.full_messages[: len(full_before)] == full_before + assert len(manager.full_messages) == len(full_before) + 2 + assert manager.active_messages[0] == {"role": "user", "content": "do the task"} + assert "Context reset" in manager.active_messages[1]["content"] + + +def test_invalid_context_policy_is_rejected(tmp_path): + policy = ContextPolicy(max_context_tokens=1_000, max_output_tokens=900, safety_tokens=100) + try: + ManagedContext("task", tmp_path / "context", policy) + except ValueError as exc: + assert "must exceed" in str(exc) + else: + raise AssertionError("invalid context budget was accepted") diff --git a/tests/train/test_toolathlon_daytona_staging.py b/tests/train/test_toolathlon_daytona_staging.py new file mode 100644 index 0000000000..1ac096cdb0 --- /dev/null +++ b/tests/train/test_toolathlon_daytona_staging.py @@ -0,0 +1,63 @@ +import json +import tomllib +from pathlib import Path + +import pytest + +from examples.train.toolathlon_harbor.prepare_daytona_tasks import stage_tasks + + +def _write_task(root: Path, task_id: str = "sample") -> Path: + task = root / task_id + environment = task / "environment" + (environment / "task").mkdir(parents=True) + (environment / "task" / "initial_state.json").write_text("{}") + (environment / "mcp.json").write_text(json.dumps({"mcpServers": {}})) + (environment / "Dockerfile").write_text("FROM local-runtime\n") + task.joinpath("instruction.md").write_text("Do the task") + task.joinpath("task.toml").write_text( + """schema_version = "1.4" + +[metadata] +mcp_servers = ["emails", "word"] + +[environment] +cpus = 2 +docker_image = "old:image" +""" + ) + return task + + +def test_stages_upload_only_task_with_shared_runtime(tmp_path: Path) -> None: + source = tmp_path / "source" + output = tmp_path / "output" + _write_task(source) + + assert stage_tasks(source, output, "registry/runtime:v1") == 1 + + staged = output / "sample" + config = tomllib.loads((staged / "task.toml").read_text()) + assert config["environment"] == { + "cpus": 2, + "docker_image": "registry/runtime:v1", + "workdir": "/opt", + "env": { + "T3_BUNDLE_DIR": "/opt/task", + "T3_SERVERS": "emails,word", + "T3_WORLD_DUMP": "/logs/world_after.json", + }, + } + assert not (staged / "environment" / "Dockerfile").exists() + assert (staged / "environment" / "task" / "initial_state.json").exists() + assert (staged / "environment" / "mcp.json").exists() + + +def test_existing_output_requires_force(tmp_path: Path) -> None: + source = tmp_path / "source" + output = tmp_path / "output" + _write_task(source) + output.mkdir() + + with pytest.raises(FileExistsError, match="--force"): + stage_tasks(source, output, "registry/runtime:v1") diff --git a/tests/train/test_toolathlon_harbor_adapter.py b/tests/train/test_toolathlon_harbor_adapter.py index ba8d9d7db1..e7016296c9 100644 --- a/tests/train/test_toolathlon_harbor_adapter.py +++ b/tests/train/test_toolathlon_harbor_adapter.py @@ -47,7 +47,7 @@ def test_compute_config_keeps_agent_on_host_and_selects_compute(): def test_launchers_use_restored_bundle_layout_and_compute_environment(): local_launcher = (ADAPTER / "run_eval.sh").read_text() compute_launcher = (ADAPTER / "run_compute_eval.sh").read_text() - assert "toolathlon-tasks/tasks}" in local_launcher + assert "toolathlon-tasks/eval_tasks}" in local_launcher assert "toolathlon-tasks/runtime/" in local_launcher assert "--platform linux/amd64 --load" in local_launcher assert '"$HERE/run_eval.sh" --env compute "$@"' in compute_launcher @@ -55,15 +55,25 @@ def test_launchers_use_restored_bundle_layout_and_compute_environment(): def test_launcher_keeps_toolathlon_out_of_generic_agent(): - generic = ( - (ROOT / "examples/train_integrations/harbor/mcp_agent.py").read_text().lower() - ) - runner = ( - (ROOT / "examples/train_integrations/harbor/mcp_runner.py").read_text().lower() - ) - bridge = ( - (ROOT / "examples/train_integrations/harbor/mcp_bridge.py").read_text().lower() - ) + generic = (ROOT / "examples/train_integrations/harbor/mcp_agent.py").read_text().lower() + runner = (ROOT / "examples/train_integrations/harbor/mcp_runner.py").read_text().lower() + bridge = (ROOT / "examples/train_integrations/harbor/mcp_bridge.py").read_text().lower() assert "toolathlon" not in generic assert "toolathlon" not in runner assert "toolathlon" not in bridge + + + +def test_daytona_training_uses_shared_runtime_snapshot(): + config = yaml.safe_load( + (ADAPTER / "harbor_daytona_training_config.yaml").read_text() + ) + kwargs = config["environment"]["kwargs"] + assert kwargs["snapshot_template_name"] == "toolathlon-json-runtime-v1" + assert kwargs["auto_snapshot"] is False + + launcher = (ADAPTER / "run_grpo_qwen38_27b_2node.sh").read_text() + assert "prepare_daytona_tasks.py" in launcher + assert "TOOLATHLON_RUNTIME_IMAGE" in launcher + assert "DAYTONA_SNAPSHOT_TEMPLATE" in launcher + assert "Dockerfile" not in launcher diff --git a/uv.lock b/uv.lock index d42e53b833..3725aeb5c7 100644 --- a/uv.lock +++ b/uv.lock @@ -3033,7 +3033,7 @@ wheels = [ [[package]] name = "harbor" version = "0.18.0" -source = { git = "ssh://git@github.com/AfterQuery/harbor-aq.git?rev=ff8961ece97f744e3c7e373c0b410d87b01cf8f7#ff8961ece97f744e3c7e373c0b410d87b01cf8f7" } +source = { git = "ssh://git@github.com/AfterQuery/harbor-aq.git?rev=1627cfb2295fa9b42d70e19dffbecc6f3243eaad#1627cfb2295fa9b42d70e19dffbecc6f3243eaad" } dependencies = [ { name = "dirhash", marker = "python_full_version >= '3.12' or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, { name = "fastapi", marker = "python_full_version >= '3.12' or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-jax') or (extra == 'extra-5-skyrl-fsdp' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-gpu' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-jax' and extra == 'extra-5-skyrl-megatron') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-miniswe') or (extra == 'extra-5-skyrl-megatron' and extra == 'extra-5-skyrl-tpu') or (extra == 'extra-5-skyrl-miniswe' and extra == 'extra-5-skyrl-tpu')" }, @@ -8951,7 +8951,7 @@ requires-dist = [ { name = "func-timeout", marker = "extra == 'skyrl-train'" }, { name = "gcsfs", marker = "extra == 'skyrl-train'" }, { name = "griffe2md", marker = "extra == 'dev'" }, - { name = "harbor", extras = ["daytona", "modal"], marker = "python_full_version >= '3.12' and extra == 'harbor'", git = "ssh://git@github.com/AfterQuery/harbor-aq.git?rev=ff8961ece97f744e3c7e373c0b410d87b01cf8f7" }, + { name = "harbor", extras = ["daytona", "modal"], marker = "python_full_version >= '3.12' and extra == 'harbor'", git = "ssh://git@github.com/AfterQuery/harbor-aq.git?rev=1627cfb2295fa9b42d70e19dffbecc6f3243eaad" }, { name = "hf-transfer" }, { name = "hf-transfer", marker = "extra == 'skyrl-train'" }, { name = "hydra-core", marker = "extra == 'skyrl-train'", specifier = "==1.3.2" },