diff --git a/docs/user-guide/cli-reference.md b/docs/user-guide/cli-reference.md index 1ca81e8519..dfa4e3cb09 100644 --- a/docs/user-guide/cli-reference.md +++ b/docs/user-guide/cli-reference.md @@ -200,6 +200,14 @@ Sections mirror the launch-script argument groups. | `--eval-max-response-len` | int | – | Max eval response length. Inherits from rollout if unset. | | `--eval-temperature` | float | – | Eval temperature. Inherits from rollout if unset. | | `--eval-top-p` | float | – | Eval top-p. Inherits from rollout if unset. | +| `--eval-num-gpus` | int | `0` | Dedicated eval fleet size (required for fully-async eval). `0` = shared-engine eval. | +| `--eval-num-gpus-per-engine` | int | `1` | Eval engine TP, independent of rollout TP. | +| `--eval-hf-dir` | str | – | Staging dir for per-eval HF snapshots (tmpfs recommended). Unset + `--save-hf` = reuse mode. | +| `--eval-model-path` | str | – | Boot checkpoint for eval engines. Defaults to `--hf-checkpoint`. | +| `--eval-dispatch` | str | `async` | `async` fires eval without blocking training; `blocking` awaits inline. | +| `--eval-max-in-flight` | int | `2` | Max concurrently pending async evals. | +| `--eval-overflow-policy` | str | `backpressure` | At the cap: await the oldest eval, or `skip` the new point (logged as `eval/skipped_busy`). | +| `--eval-keep-snapshots` | int | `2` | GC ring for `--eval-hf-dir`; `--save-hf` output is never deleted. | ### Performance diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index c0e35d57d3..69c23ba8f9 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -100,6 +100,81 @@ signal. Fast [P2P weight transfer](/advanced/p2p-weight-transfer) keeps the rollout engines closer to the latest actor weights so fewer groups get recycled by `--max-weight-staleness`. +## Evaluation + +Pick a posture by two questions: is this a test run or a real run (are checkpoints +persisted anyway), and does eval get standalone GPU resources? + +| | Test run | Real run | +|---|---|---| +| **No standalone eval** | Pause-the-world (shared engines) | External service over `--save-hf` output; pause-the-world if eval must be strictly on-time (costs ~eval duration of rollout production per point). The two compose: pause-the-world on a small sanity set for on-time points, the service on the full set for delayed high-fidelity points | +| **Standalone eval fleet** | Fleet + tmpfs snapshot (`--eval-hf-dir /dev/shm/...`) | Fleet + checkpoint reuse (no `--eval-hf-dir`) | + +Without extra GPUs (`--eval-num-gpus` unset), eval **shares the rollout engines**: +the producer pauses new submissions for the duration of the blocking eval and resumes +after. This is a gate, not a retract — in-flight rollout requests finish and buffer, +nothing is aborted, and the `pause_generation` API is never involved; eval requests +simply share engine capacity with the draining tail. + +No extra weight movement happens either: the engines already carry the weights the +step's `update_weights` broadcast just pushed (in fully-async, only *generation* is +continuous — weight updates are still driver-scheduled per step, each with its own +generation pause per `--pause-generation-mode`; eval adds no pause of its own on top). +Pinning comes from ordering: the driver awaits the eval, so the next +`update_weights` cannot interleave — +expect `mixed_version_ratio == 0` with the training fleet's update counter as the +version label (a constant offset from `eval/step`, unlike the dedicated fleet which +stamps the rollout_id). This ordering is also why shared-engine eval must stay +blocking: fired-and-forgotten, the next weight update would rewrite the engines +mid-eval. The cost is that rollout production stalls for roughly the eval duration, +which is fine for small debug eval sets. + +For eval that never touches training capacity, use a **dedicated eval fleet** synced +through HF checkpoint snapshots — never by joining training weight updates: + +```bash +--eval-num-gpus 1 # dedicated eval engines (own router) +--eval-interval K +--eval-hf-dir /dev/shm/miles_eval_hf # snapshot staging; tmpfs = no disk dependency +--eval-prompt-data aime /path/to/aime.jsonl +``` + +Per eval-due step the trainer exports an HF snapshot (seconds to tmpfs), fires the +eval **fire-and-forget**, and keeps training; the eval fleet pins its weights to the +snapshot (`weight_version = str(rollout_id)`), runs the standard eval datasets, and the +point lands at the right x-axis step even when it completes a few steps later +(`eval/lag_steps` reports how late). + +Two production-oriented variants: + +- **Reuse mode**: with `--save-hf` set and `--eval-hf-dir` unset, eval reuses the + periodic HF checkpoints (requires `eval_interval % save_interval == 0`) — zero extra + export cost. Pair with `--eval-overflow-policy skip` so a slow eval set can never + stall training. +- **External service**: `tools/checkpoint_eval_service.py` watches `--save-hf` output + with its own sglang server — no GPU carve-out from the training job, restartable, + backfills missed points from its ledger. Works for `--colocate` runs too. + +Every skipped point is attributable from the dashboard: `eval/skipped_busy`, +`eval/skipped_ckpt_missing`, `eval/skipped_unhealthy`, `eval/skipped_export_failed` +are logged at the affected step. `eval/{ds}/weight_version/mean == eval/step` and +`mixed_version_ratio == 0` confirm every point measured exactly the intended weights. + +The eval-engine `weight_version` namespace is the snapshot's `rollout_id` — deliberately +different from the training fleet's job-local update counter; the two fleets never mix. + +Where each posture's weights come from — and what "the weights at step R" means: + +| Posture | Weight delivery | Measures | +|---|---|---| +| Pause-the-world | none needed — training's own `update_weights` broadcast already put them on the shared engines | the engines' last-broadcast version (equals the actor's current weights when `update_weights_interval` is 1) | +| Dedicated fleet | `update_weights_from_disk` on a snapshot exported **directly from the actor** | the actor's exact step-R weights, regardless of broadcast schedule | +| External service | loads `--save-hf` checkpoints itself | the actor's exact step-R weights | + +Eval engines are never added to the training broadcast group: collectives cannot skip +members, so a fleet inside the group would have its weights rewritten by every update +and asynchronous points could not be pinned. + ## Example implementation For a complete Qwen3 launch script and worker implementation, see the diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index d65011d824..bf58fb5f50 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -6,6 +6,7 @@ The implementation lives in the core library at `miles/rollout/fully_async_rollo ## Files * `run-qwen3-4b-fully_async.sh`: example launch script with Qwen3‑4B. +* `run-qwen3.5-4b-fully_async-eval.sh`: Qwen3.5‑4B with a dedicated eval fleet (fully-async eval). ## Prerequisite First set up model & environment following the Qwen3-4B example. @@ -26,8 +27,12 @@ Started fully-async rollout worker * Completed groups are pushed into a queue; each step drains until it has `--rollout-batch-size` groups. * Aborted or too-stale groups are recycled back into the data source. +## Evaluation +Without extra GPUs, eval shares the rollout engines (producer pauses during the blocking +eval). With `--eval-num-gpus`/`--eval-hf-dir`, eval runs on a dedicated fleet synced via +HF checkpoint snapshots; see `run-qwen3.5-4b-fully_async-eval.sh` and the fully-async docs. + ## Limitations -* No evaluation mode (`--eval-function-path` must point at the standard `InferenceRolloutFn`). * Ordering is best effort (sorted at the end by index). ## Config Differences (3 Key Points) diff --git a/examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh b/examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh new file mode 100644 index 0000000000..e074dcb4f6 --- /dev/null +++ b/examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh @@ -0,0 +1,164 @@ +#!/bin/bash +# Fully-async training with a dedicated eval fleet (Qwen3.5-4B, one 8-GPU node). +# GPU split: 4 actor (TP=2) + 3 rollout engines (TP=1) + 1 eval engine; eval weights +# are pinned per-eval to an HF snapshot on tmpfs, so eval never pauses training. + +# for rerun the task +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +# will prevent ray from buffering stdout/stderr +export PYTHONBUFFERED=16 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +source "${SCRIPT_DIR}/../../scripts/models/qwen3.5-4B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3.5-4B + --ref-load /root/Qwen3.5-4B_torch_dist + --load /root/Qwen3.5-4B_miles/ + --save /root/Qwen3.5-4B_miles/ + --save-interval 20 +) + +PROMPT_SET=/path/to/dapo-math-17k.jsonl + +ROLLOUT_ARGS=( + --rollout-function-path miles.rollout.fully_async_rollout.FullyAsyncRolloutFn + --prompt-data ${PROMPT_SET} + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + + --rm-type dapo + --reward-key score + + --num-rollout 3000 + --rollout-batch-size 32 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --global-batch-size 256 + --balance-data + + # retract (default) can deadlock flush_cache in fully_async under load + --pause-generation-mode in_place + + # for staleness control + #--max-weight-staleness 2 +) + +EVAL_ARGS=( + --eval-interval 5 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 16384 + --eval-top-p 1 + + # dedicated eval fleet: 1 engine, pinned per-eval to a tmpfs HF snapshot + --eval-num-gpus 1 + --eval-num-gpus-per-engine 1 + --eval-hf-dir /dev/shm/miles_eval_hf + --eval-keep-snapshots 2 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 9216 +) + +GRPO_ARGS=( + --advantage-estimator grpo + --use-kl-loss + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --entropy-coef 0.00 + --eps-clip 0.2 + --eps-clip-high 0.28 + + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.7 +) + +MISC_ARGS=( + # default dropout in megatron is 0.1 + --attention-dropout 0.0 + --hidden-dropout 0.0 + # should be good for model performance + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + # need to comment this when using model with MLA + --attention-backend flash +) + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats + +RUNTIME_ENV_JSON="{ + \"env_vars\": { + \"PYTHONPATH\": \"/root/Megatron-LM/\", + \"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR\": \"1\", + \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\", + \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\" + } +}" + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 train_async.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 4 \ + --rollout-num-gpus 3 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${GRPO_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${SGLANG_ARGS[@]} \ + ${MISC_ARGS[@]} diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b922417598..a91cd9eb2b 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -162,7 +162,7 @@ def init( dict(no_load_optim=False, no_load_rng=False, finetune=False) if recv_ckpt_src_rank is not None else {} ) with inplace_modify_args(args, heal_load_overrides): - (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer( + self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer( args, role, checkpointing_context=checkpointing_context ) @@ -563,6 +563,21 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: if self.args.offload_train: destroy_process_groups() + @with_logs + @timer + def export_hf(self, rollout_id: int, path: str) -> None: + """Export current weights as an HF checkpoint to ``path`` (collective). + + Uses the direct megatron->HF converters (the weight updater's machinery), so + export coverage matches weight-sync coverage. Unlike the periodic --save-hf + path inside save_model, failures propagate to the caller so an eval snapshot + that failed to export can be skipped loudly. + """ + self._heartbeat.bump() + from miles.backends.megatron_utils.model import save_hf_model + + save_hf_model(self.args, rollout_id, self.model, path=path, raise_on_error=True) + @with_logs @timer def update_weights(self, info: "EnginesAndLock") -> None: diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 61a49fe395..4080af8c32 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -30,6 +30,7 @@ from miles.utils.audit_utils.witness.allocator import WitnessInfo from miles.utils.audit_utils.witness.module import witness_dump_and_clear_stale from miles.utils.dumper_utils import DumperMegatronUtil, DumperPhase +from miles.utils.hf_config import HF_EXPORT_COMPLETE_MARKER, load_hf_config from miles.utils.memory_utils import clear_memory from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor from miles.utils.tracking_utils.structured_log import log_structured @@ -825,7 +826,96 @@ def save( enable_forward_pre_hook(model) -def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: +HF_WEIGHT_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".gguf") + + +def _is_hf_metadata_file(path: Path) -> bool: + """Tokenizer/config files worth copying into an export — not weights, and not the + base checkpoint's weight index, which would clobber the one the export writes.""" + return path.is_file() and path.suffix not in HF_WEIGHT_SUFFIXES and not path.name.endswith(".index.json") + + +def export_hf_model_direct( + args, + model: Sequence[DDP], + path: str | Path, + *, + model_name: str, + quantization_config, + megatron_local_weights, +) -> None: + """Export current weights as an HF checkpoint via miles' own megatron->HF converters. + + Same conversion machinery as the weight updater, so export coverage matches + weight-sync coverage (the bridge silently exports zero weights for specs it has + no mapping for, e.g. qwen3.5). Collective — all ranks must call it; rank 0 writes. + """ + import json + import shutil + + import safetensors.torch + + from miles.backends.megatron_utils.update_weight.hf_weight_iterator_direct import HfWeightIteratorDirect + + path = Path(path) + is_writer = torch.distributed.get_rank() == 0 + if is_writer: + path.mkdir(parents=True, exist_ok=True) + + iterator = HfWeightIteratorDirect(args, model, model_name=model_name, quantization_config=quantization_config) + + weight_map: dict[str, str] = {} + total_size = 0 + shard_index = 0 + for hf_named_tensors in iterator.get_hf_weight_chunks(megatron_local_weights): + if not is_writer: + continue + shard_index += 1 + shard_name = f"model-{shard_index:05d}.safetensors" + shard_tensors = {} + for name, tensor in hf_named_tensors: + shard_tensors[name] = tensor.detach().to("cpu").contiguous() + weight_map[name] = shard_name + total_size += shard_tensors[name].numel() * shard_tensors[name].element_size() + safetensors.torch.save_file(shard_tensors, path / shard_name) + del shard_tensors + + if is_writer: + assert weight_map, f"HF export to {path} produced no weights" + base_checkpoint = Path(args.hf_checkpoint) + if base_checkpoint.is_dir(): + for meta_file in base_checkpoint.iterdir(): + if _is_hf_metadata_file(meta_file): + shutil.copy2(meta_file, path / meta_file.name) + else: + logger.warning(f"hf_checkpoint {args.hf_checkpoint} is not a local dir; metadata not copied to {path}") + index = {"metadata": {"total_size": total_size}, "weight_map": weight_map} + (path / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + + torch.distributed.barrier() + if is_writer: + (path / HF_EXPORT_COMPLETE_MARKER).touch() + + +_hf_bridge_cache: dict = {} + + +def _get_hf_bridge(hf_checkpoint: str): + from megatron.bridge import AutoBridge + + if hf_checkpoint not in _hf_bridge_cache: + _hf_bridge_cache[hf_checkpoint] = AutoBridge.from_hf_pretrained(hf_checkpoint, trust_remote_code=True) + return _hf_bridge_cache[hf_checkpoint] + + +def save_hf_model( + args, + rollout_id: int, + model: Sequence[DDP], + *, + path: str | Path | None = None, + raise_on_error: bool = False, +) -> None: """Save Megatron model in HuggingFace format. For LoRA models this saves both: @@ -834,50 +924,75 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: - An **adapter-only** HF PEFT checkpoint at ``{path}/adapter/`` so it can be loaded with ``PeftModel.from_pretrained``. - This function is collective — all ranks must call it. + This function is collective — all ranks must call it. On success, global rank 0 + writes a ``.complete`` marker file. Args: args: Runtime arguments. model (Sequence[DDP]): Sequence of DDP-wrapped model chunks. rollout_id (int): Rollout ID for path formatting. + path: Destination directory; defaults to ``args.save_hf.format(rollout_id)``. + raise_on_error: Re-raise export failures instead of logging them. """ should_log = get_parallel_state().effective_dp_cp.rank == 0 and get_parallel_state().tp.rank == 0 + path = Path(path if path is not None else args.save_hf.format(rollout_id=rollout_id)) try: - from megatron.bridge import AutoBridge - - from miles.utils.megatron_bridge_utils import patch_megatron_model - - path = Path(args.save_hf.format(rollout_id=rollout_id)) - if should_log: logger.info(f"Saving model in HuggingFace format to {path}") - bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True) - - path.mkdir(parents=True, exist_ok=True) - - with patch_megatron_model(model): - # For LoRA models, merge_adapter_weights=True (default) merges - # adapter weights into base weights for a standalone HF model. - bridge.save_hf_pretrained(model, path=path) + if args.megatron_to_hf_mode == "raw" and not is_lora_model(model): + # LoRA keeps the bridge (adapter merging). + from .update_weight.common import named_params_and_buffers + + hf_config = load_hf_config(args.hf_checkpoint) + export_hf_model_direct( + args, + model, + path, + model_name=type(hf_config).__name__.lower() if args.model_name is None else args.model_name, + quantization_config=getattr(hf_config, "quantization_config", None), + megatron_local_weights=dict(named_params_and_buffers(args, model, convert_to_global_name=True)), + ) + else: + from miles.utils.megatron_bridge_utils import patch_megatron_model + + bridge = _get_hf_bridge(args.hf_checkpoint) + path.mkdir(parents=True, exist_ok=True) + with patch_megatron_model(model): + # For LoRA models, merge_adapter_weights=True (default) merges + # adapter weights into base weights for a standalone HF model. + bridge.save_hf_pretrained(model, path=path) + + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + if not any(path.glob("*.safetensors")) and not any(path.glob("*.bin")): + raise RuntimeError( + f"HF export to {path} produced no weight files — the megatron " + f"bridge likely has no mapping for this model architecture." + ) + (path / HF_EXPORT_COMPLETE_MARKER).touch() if should_log: logger.info(f"Successfully saved merged HuggingFace model to {path}") except Exception as e: + if raise_on_error: + raise if should_log: logger.error(f"Failed to save HuggingFace format: {e}") # Additionally save adapter-only checkpoint for LoRA models if is_lora_model(model): try: - adapter_path = Path(args.save_hf.format(rollout_id=rollout_id)) / "adapter" + adapter_path = path / "adapter" if should_log: logger.info(f"Saving LoRA adapter (HF PEFT format) to {adapter_path}") save_lora_checkpoint(model, args, str(adapter_path)) if should_log: logger.info(f"Successfully saved LoRA adapter to {adapter_path}") except Exception as e: + if raise_on_error: + raise if should_log: logger.error(f"Failed to save LoRA adapter: {e}") diff --git a/miles/ray/actor_group.py b/miles/ray/actor_group.py index ca47ddff6a..837836840c 100644 --- a/miles/ray/actor_group.py +++ b/miles/ray/actor_group.py @@ -88,6 +88,10 @@ async def save_model(self, rollout_id, force_sync=False): """Save actor model""" await self._broadcast("save_model", rollout_id, force_sync=force_sync) + async def export_hf(self, rollout_id: int, path: str): + """Export current weights as an HF checkpoint (collective across all ranks).""" + await self._broadcast("export_hf", rollout_id, path) + async def update_weights(self, rollout_id: int | None = None): """Broadcast weights from rank 0 to all other ranks.""" if self.args.debug_train_only or self.args.debug_rollout_only: diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index a9e7122c45..f1a6067e13 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -107,7 +107,7 @@ def create_placement_groups(args): num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node else: - num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + args.rollout_num_gpus + num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + args.rollout_num_gpus + args.eval_num_gpus rollout_offset = args.actor_num_nodes * args.actor_num_gpus_per_node if args.use_critic: num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node diff --git a/miles/ray/rollout/eval_dispatch.py b/miles/ray/rollout/eval_dispatch.py new file mode 100644 index 0000000000..de1d97f439 --- /dev/null +++ b/miles/ray/rollout/eval_dispatch.py @@ -0,0 +1,79 @@ +import logging +import os +import time +from collections import deque + +import ray + +logger = logging.getLogger(__name__) + + +class EvalDispatcher: + """Fire-and-forget evals against the dedicated eval fleet; blocking legacy call + when ``--eval-num-gpus`` is 0. Failures degrade to a skipped point, never a crash.""" + + def __init__(self, args, actor_model, rollout_manager): + self.args = args + self.actor_model = actor_model + self.rollout_manager = rollout_manager + self.pending: deque[tuple[int, ray.ObjectRef]] = deque() + + async def dispatch(self, rollout_id: int, hf_dir: str | None = None, force: bool = False) -> None: + if self.args.eval_num_gpus <= 0: + await self.rollout_manager.eval.remote(rollout_id) + return + + self._reap_finished() + if len(self.pending) >= self.args.eval_max_in_flight: + if self.args.eval_overflow_policy == "skip" and not force: + await self.rollout_manager.report_eval_skip.remote(rollout_id, "busy") + return + oldest_id, oldest_ref = self.pending.popleft() + await self._await_ref(oldest_id, oldest_ref) + + export_time = None + if hf_dir is None: + try: + hf_dir, export_time = await self._ensure_snapshot(rollout_id) + except Exception as e: + logger.error(f"HF snapshot export for eval {rollout_id} failed: {e}") + await self.rollout_manager.report_eval_skip.remote(rollout_id, "export_failed") + return + + ref = self.rollout_manager.eval.remote(rollout_id, hf_dir=hf_dir, export_time_seconds=export_time) + if self.args.eval_dispatch == "blocking": + await self._await_ref(rollout_id, ref) + else: + self.pending.append((rollout_id, ref)) + + async def drain(self) -> None: + while self.pending: + rollout_id, ref = self.pending.popleft() + await self._await_ref(rollout_id, ref) + + async def _ensure_snapshot(self, rollout_id: int) -> tuple[str, float | None]: + if self.args.eval_hf_dir is None: + return self.args.save_hf.format(rollout_id=rollout_id), None + hf_dir = os.path.join(self.args.eval_hf_dir, f"step_{rollout_id}") + start = time.time() + await self.actor_model.export_hf(rollout_id, hf_dir) + return hf_dir, time.time() - start + + def _reap_finished(self) -> None: + while self.pending: + done, _ = ray.wait([self.pending[0][1]], timeout=0) + if not done: + break + rollout_id, ref = self.pending.popleft() + try: + ray.get(ref) + except Exception: + logger.exception(f"Async eval for rollout {rollout_id} raised") + self.rollout_manager.report_eval_skip.remote(rollout_id, "crashed") + + async def _await_ref(self, rollout_id: int, ref) -> None: + try: + await ref + except Exception: + logger.exception(f"Async eval for rollout {rollout_id} raised") + await self.rollout_manager.report_eval_skip.remote(rollout_id, "crashed") diff --git a/miles/ray/rollout/metrics.py b/miles/ray/rollout/metrics.py index 9ee7ae41e9..aee01b4cb7 100644 --- a/miles/ray/rollout/metrics.py +++ b/miles/ray/rollout/metrics.py @@ -26,7 +26,11 @@ def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] log_dict = extra_metrics or {} for key in data.keys(): + if (num_failed := data[key].get("failed_samples")) is not None and num_failed > 0: + log_dict[f"eval/{key}/failed_samples"] = num_failed rewards = data[key]["rewards"] + if not rewards: + continue log_dict[f"eval/{key}"] = sum(rewards) / len(rewards) if (samples := data[key].get("samples")) is not None: log_dict |= dict_add_prefix(_compute_metrics_from_samples(args, samples), f"eval/{key}/") @@ -51,6 +55,16 @@ def log_eval_rollout_data(rollout_id, args, data, extra_metrics: dict[str, Any] return log_dict +def log_eval_skip(rollout_id, args, reason: str): + """Log a skipped eval point at ``rollout_id`` so curve gaps are attributable.""" + log_dict = { + f"eval/skipped_{reason}": 1, + "eval/step": compute_rollout_step(args, rollout_id), + } + logger.warning(f"eval {rollout_id} skipped: {reason}") + tracking.log(args, log_dict, step_key="eval/step") + + def log_rollout_data(rollout_id, args, samples, rollout_extra_metrics, rollout_time): if (x := args.custom_rollout_log_function_path) is not None: custom_log_func = load_function(x) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 4d98fbd9fa..4941c6a1c2 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -1,14 +1,16 @@ import asyncio import logging +import shutil import time from dataclasses import dataclass +from pathlib import Path import ray from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS from miles.ray.rollout.addr_allocator import PortCursors from miles.ray.rollout.debug_data import RolloutDataInjectionUtil, load_debug_rollout_data, save_debug_rollout_data -from miles.ray.rollout.metrics import log_eval_rollout_data, log_rollout_data +from miles.ray.rollout.metrics import log_eval_rollout_data, log_eval_skip, log_rollout_data from miles.ray.rollout.rollout_data_conversion import postprocess_rollout_data from miles.ray.rollout.rollout_server import RolloutServer, start_rollout_servers from miles.ray.rollout.router_manager import start_session_server @@ -27,7 +29,8 @@ from miles.utils.audit_utils.process_identity import RolloutManagerProcessIdentity from miles.utils.environ import enable_experimental_rollout_refactor from miles.utils.health_monitor import RolloutHealthMonitor -from miles.utils.http_utils import init_http_client +from miles.utils.hf_config import is_complete_hf_export +from miles.utils.http_utils import init_http_client, wait_http_ok from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker from miles.utils.misc import load_function @@ -40,6 +43,8 @@ logger = logging.getLogger(__name__) +EVAL_WEIGHT_LOAD_TIMEOUT_SECS = 600.0 + @ray.remote class RolloutManager: @@ -87,6 +92,8 @@ def __init__(self, args, pg): start_session_server(args) self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote() self.rollout_id = -1 + self._eval_lock = asyncio.Lock() + self._eval_consumed_snapshots: list[str] = [] self._metric_checker = MetricChecker.maybe_create(args) @@ -135,12 +142,16 @@ async def generate(self, rollout_id): data_ref = split_train_data_by_dp(self.args, data, self.train_parallel_config["dp_size"]) return dict(sample_indices=sample_indices, data_ref=data_ref) - async def eval(self, rollout_id): + async def eval(self, rollout_id, hf_dir: str | None = None, export_time_seconds: float | None = None): if self.args.debug_train_only: # if debug train only, we don't generate evaluation data return self._health_monitoring_resume() + if self.args.eval_num_gpus > 0: + assert hf_dir is not None, "eval with a dedicated fleet requires an HF snapshot dir" + return await self._eval_on_dedicated_fleet(rollout_id, hf_dir, export_time_seconds) + if self.use_experimental_refactor: result = await asyncio.to_thread( call_rollout_function, self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id) @@ -155,6 +166,128 @@ async def eval(self, rollout_id): if self._metric_checker is not None: self._metric_checker.on_eval(metrics) + async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_time_seconds: float | None): + """Pin the eval fleet to the snapshot for ``rollout_id``, then run eval. + + Every failure mode degrades to a skipped point logged at ``rollout_id``. + """ + start_time = time.time() + async with self._eval_lock: + srv = self.servers["eval"] + try: + await self._mark_unreachable_eval_engines(srv) + await srv.recover() + await srv.wait_all_engines_alive() + except Exception as e: + logger.warning(f"Eval fleet unhealthy at rollout {rollout_id}, skipping eval: {e}") + self.report_eval_skip(rollout_id, "unhealthy") + return + + if hf_dir != self.args.hf_checkpoint and not is_complete_hf_export(hf_dir): + logger.warning(f"Eval snapshot {hf_dir} missing or incomplete, skipping eval {rollout_id}") + self.report_eval_skip(rollout_id, "ckpt_missing") + return + + engines = [e.actor_handle for e in srv.engines] + weight_version = str(rollout_id) + for _attempt in range(2): + try: + await asyncio.wait_for( + asyncio.gather( + *[ + e.update_weights_from_disk.remote(hf_dir, weight_version=weight_version) + for e in engines + ] + ), + timeout=EVAL_WEIGHT_LOAD_TIMEOUT_SECS, + ) + versions = await asyncio.wait_for( + asyncio.gather(*[e.get_weight_version.remote() for e in engines]), + timeout=EVAL_WEIGHT_LOAD_TIMEOUT_SECS, + ) + except Exception as e: + logger.warning(f"Eval fleet weight load from {hf_dir} failed: {e}") + versions = [] + if versions and all(str(v) == weight_version for v in versions): + break + else: + logger.warning( + f"Eval fleet failed to pin weight_version={weight_version} (got {versions}), skipping eval" + ) + self.report_eval_skip(rollout_id, "pin_violation") + return + + try: + await self._wait_eval_router_ready(srv) + except Exception as e: + logger.warning(f"Eval router not ready at rollout {rollout_id}, skipping eval: {e}") + self.report_eval_skip(rollout_id, "unhealthy") + return + + result = await asyncio.to_thread( + call_rollout_function, self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id) + ) + data = result.data + save_debug_rollout_data(self.args, data, rollout_id=rollout_id, evaluation=True) + extra_metrics = dict(result.metrics or {}) + extra_metrics["eval/lag_steps"] = max(self.rollout_id - rollout_id, 0) + extra_metrics["eval/duration_seconds"] = time.time() - start_time + if export_time_seconds is not None: + extra_metrics["eval/export_time_seconds"] = export_time_seconds + metrics = log_eval_rollout_data(rollout_id, self.args, data, extra_metrics) + if self._metric_checker is not None: + self._metric_checker.on_eval(metrics) + + self._gc_eval_snapshots(hf_dir) + + async def _wait_eval_router_ready(self, srv, timeout: float = 180.0) -> None: + """After a revival the router 503s until its health cycle evicts the dead + worker; a retried one-token probe proves the route is usable before dispatch.""" + await wait_http_ok( + f"http://{srv.router_ip}:{srv.router_port}/generate", + json_payload={"input_ids": [0], "sampling_params": {"max_new_tokens": 1, "temperature": 0}}, + timeout=timeout, + ) + + async def _mark_unreachable_eval_engines(self, srv) -> None: + """Without fault tolerance nothing records an engine death (recover() only + restarts engines already marked stopped), so the controller probes itself.""" + for group in srv.server_groups: + for engine in group.all_engines: + if not engine.is_allocated: + continue + try: + await asyncio.wait_for(engine.actor_handle.get_weight_version.remote(), timeout=60) + except Exception as e: + logger.warning(f"Eval engine unreachable ({e!r}); marking stopped for recovery") + try: + ray.kill(engine.actor_handle) + except Exception: + pass + engine.mark_stopped() + + def report_eval_skip(self, rollout_id: int, reason: str) -> None: + log_eval_skip(rollout_id, self.args, reason) + + def _gc_eval_snapshots(self, consumed_dir: str) -> None: + """Delete consumed --eval-hf-dir snapshots beyond the keep ring; nothing else + is ever deleted (pending evals reference unconsumed dirs).""" + staging = getattr(self.args, "eval_hf_dir", None) + if staging is None: + return + staging_root = Path(staging).resolve() + consumed = Path(consumed_dir).resolve() + if staging_root not in consumed.parents: + return + consumed = str(consumed) + if consumed in self._eval_consumed_snapshots: + self._eval_consumed_snapshots.remove(consumed) + self._eval_consumed_snapshots.append(consumed) + while len(self._eval_consumed_snapshots) > self.args.eval_keep_snapshots: + victim = self._eval_consumed_snapshots.pop(0) + shutil.rmtree(victim, ignore_errors=True) + logger.info(f"GC'd consumed eval snapshot {victim}") + async def _get_rollout_data(self, rollout_id): if self.args.load_debug_rollout_data: data, metadata = load_debug_rollout_data(self.args, rollout_id=rollout_id) diff --git a/miles/ray/rollout/rollout_server.py b/miles/ray/rollout/rollout_server.py index 55d230b0f5..5e97a6415d 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/miles/ray/rollout/rollout_server.py @@ -104,24 +104,51 @@ def start_rollout_servers(args, pg) -> dict[str, "RolloutServer"]: def _resolve_sglang_config(args) -> SglangConfig: """Build a SglangConfig from args, choosing the right source.""" + eval_num_gpus = args.eval_num_gpus + if getattr(args, "sglang_config", None) is not None: config = SglangConfig.from_yaml(args.sglang_config) - expected = args.rollout_num_gpus + expected = args.rollout_num_gpus + eval_num_gpus actual = config.total_num_gpus - assert actual == expected, f"sglang_config total GPUs ({actual}) != rollout_num_gpus ({expected})" + assert ( + actual == expected + ), f"sglang_config total GPUs ({actual}) != rollout_num_gpus + eval_num_gpus ({expected})" + if eval_num_gpus > 0: + eval_models = [m for m in config.models if m.name == "eval"] + assert len(eval_models) == 1 and eval_models[0].total_num_gpus == eval_num_gpus, ( + f"--eval-num-gpus {eval_num_gpus} requires the sglang_config YAML to contain " + f"exactly one model named 'eval' with that many GPUs." + ) return config if args.prefill_num_servers is not None: - return SglangConfig.from_prefill_num_servers(args) + config = SglangConfig.from_prefill_num_servers(args) + else: + config = SglangConfig( + models=[ + ModelConfig( + name="default", + server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)], + ) + ] + ) - return SglangConfig( - models=[ + if eval_num_gpus > 0: + config.models.append( ModelConfig( - name="default", - server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)], + name="eval", + model_path=args.eval_model_path, + update_weights=False, + server_groups=[ + ServerGroupConfig( + worker_type="regular", + num_gpus=eval_num_gpus, + num_gpus_per_engine=args.eval_num_gpus_per_engine, + ) + ], ) - ] - ) + ) + return config def _compute_rollout_offset(args) -> int: diff --git a/miles/ray/train_actor.py b/miles/ray/train_actor.py index 782d192b23..2b7d3b36dc 100644 --- a/miles/ray/train_actor.py +++ b/miles/ray/train_actor.py @@ -163,6 +163,10 @@ def train(self, rollout_id, rollout_data_ref): def save_model(self, rollout_id, force_sync=False): raise NotImplementedError + def export_hf(self, rollout_id: int, path: str) -> None: + """Export current weights as an HF checkpoint to ``path`` (eval snapshots).""" + raise NotImplementedError(f"{type(self).__name__} does not support HF export") + @abc.abstractmethod def update_weights(self, info: "EnginesAndLock") -> None: raise NotImplementedError diff --git a/miles/rollout/checkpoint_eval.py b/miles/rollout/checkpoint_eval.py new file mode 100644 index 0000000000..ea1fe42c75 --- /dev/null +++ b/miles/rollout/checkpoint_eval.py @@ -0,0 +1,56 @@ +"""Eval against a dedicated eval fleet pinned to HF checkpoint snapshots. + +The eval fleet never joins training weight updates; weights reach it only through +``update_weights_from_disk`` on a snapshot exported for a specific rollout_id. +""" + +import copy +from argparse import Namespace + +from miles.rollout.inference_rollout.inference_rollout_common import GenerateState + +__all__ = ["retarget_args", "make_eval_args", "make_eval_generate_state", "EvalFleetSession"] + + +def retarget_args(args: Namespace, router_ip, router_port, num_gpus: int, num_gpus_per_engine: int) -> Namespace: + """Shallow-copy ``args`` with the router address and GPU sizing swapped for eval. + + Generate functions read the router from ``args`` and ``GenerateState`` sizes its + semaphore off the GPU counts, so a retargeted copy runs the standard eval path + against a different set of engines unchanged. + """ + eval_args = copy.copy(args) + eval_args.sglang_router_ip = router_ip + eval_args.sglang_router_port = router_port + eval_args.rollout_num_gpus = num_gpus + eval_args.rollout_num_gpus_per_engine = num_gpus_per_engine + return eval_args + + +def make_eval_args(args: Namespace) -> Namespace: + router_ip, router_port = args.sglang_model_routers["eval"] + return retarget_args(args, router_ip, router_port, args.eval_num_gpus, args.eval_num_gpus_per_engine) + + +def make_eval_generate_state(args: Namespace) -> GenerateState: + return GenerateState(make_eval_args(args)) + + +class EvalFleetSession: + """Eval-fleet ``GenerateState`` + prompt cache, built lazily on first use. + + Lazy because the eval router only exists once the servers are up, while rollout + functions are constructed earlier. + """ + + def __init__(self, args: Namespace): + self.args = args + self._state: GenerateState | None = None + self._prompt_dataset_cache: dict = {} + + async def run(self) -> dict: + from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets + + if self._state is None: + self._state = make_eval_generate_state(self.args) + return await run_eval_datasets(self._state, self._prompt_dataset_cache) diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py index fa7f0b5023..d5b6f43de8 100644 --- a/miles/rollout/fully_async_rollout.py +++ b/miles/rollout/fully_async_rollout.py @@ -10,8 +10,9 @@ --rollout-function-path miles.rollout.fully_async_rollout.FullyAsyncRolloutFn -Evaluation is not served by this function; point ``--eval-function-path`` at -``miles.rollout.inference_rollout.inference_rollout_common.InferenceRolloutFn``. +Evaluation runs on a dedicated eval fleet when ``--eval-num-gpus`` is set (see +``miles/rollout/checkpoint_eval.py``); otherwise it shares the rollout engines, +pausing producer submissions for the duration of the (blocking) eval. """ import asyncio @@ -21,8 +22,16 @@ import httpx -from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnInput, RolloutFnOutput, RolloutFnTrainOutput +from miles.rollout.base_types import ( + RolloutFnConstructorInput, + RolloutFnEvalOutput, + RolloutFnInput, + RolloutFnOutput, + RolloutFnTrainOutput, +) +from miles.rollout.checkpoint_eval import EvalFleetSession from miles.rollout.inference_rollout.inference_rollout_common import GenerateState, generate_and_rm_group +from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets from miles.utils.http_utils import get from miles.utils.types import Sample @@ -91,19 +100,33 @@ def __init__(self, input: RolloutFnConstructorInput): self._weight_version = _CachedWeightVersion() self._worker: asyncio.Task | None = None self._output: asyncio.Queue[Group] | None = None + self._eval_fleet = EvalFleetSession(input.args) + self._eval_prompt_dataset_cache: dict = {} + self._producer_resumed = asyncio.Event() + self._producer_resumed.set() async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: if input.evaluation: - raise ValueError( - "FullyAsyncRolloutFn does not serve eval; set --eval-function-path to " - "miles.rollout.inference_rollout.inference_rollout_common.InferenceRolloutFn" - ) + return await self._call_eval(input) if self._worker is None: self._output = asyncio.Queue(maxsize=OUTPUT_QUEUE_MAX_GROUPS) self._worker = asyncio.create_task(self._worker_loop()) logger.info("Started fully-async rollout worker") return await self._drain(input.rollout_id) + async def _call_eval(self, input: RolloutFnInput) -> RolloutFnOutput: + if self.args.eval_num_gpus > 0: + return RolloutFnEvalOutput(data=await self._eval_fleet.run()) + + logger.info("Pausing fully-async producer submissions for shared-engine eval") + self._producer_resumed.clear() + try: + results = await run_eval_datasets(self.state, self._eval_prompt_dataset_cache) + finally: + self._producer_resumed.set() + logger.info("Resumed fully-async producer submissions after eval") + return RolloutFnEvalOutput(data=results) + # -------------------------- producer -------------------------- def _max_in_flight_groups(self) -> int: @@ -123,6 +146,7 @@ def _submit_one_group(self) -> asyncio.Task: async def _worker_loop(self): active: set[asyncio.Task] = set() while True: + await self._producer_resumed.wait() while len(active) < self._max_in_flight_groups(): active.add(self._submit_one_group()) done, active = await asyncio.wait(active, return_when=asyncio.FIRST_COMPLETED) diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py index 9f1cc603b0..7a02922c32 100644 --- a/miles/rollout/inference_rollout/inference_rollout_common.py +++ b/miles/rollout/inference_rollout/inference_rollout_common.py @@ -171,9 +171,12 @@ def compute_sampling_params( class InferenceRolloutFn: def __init__(self, input: RolloutFnConstructorInput): + from miles.rollout.checkpoint_eval import EvalFleetSession + self.data_source = input.data_source self.state = GenerateState(input.args) self.eval_prompt_dataset_cache = {} + self._eval_fleet = EvalFleetSession(input.args) async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: if input.evaluation: @@ -190,13 +193,9 @@ async def _call_train(self, input: RolloutFnTrainInput) -> RolloutFnTrainOutput: return output async def _call_eval(self, input: RolloutFnEvalInput) -> RolloutFnEvalOutput: - from miles.rollout.inference_rollout.inference_rollout_eval import eval_rollout_single_dataset - - assert not self.state.args.group_rm, "Group RM is not supported for eval rollout" + from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets - coros = [] - for dataset_cfg in getattr(self.state.args, "eval_datasets", []) or []: - coros.append(eval_rollout_single_dataset(self.state, dataset_cfg, self.eval_prompt_dataset_cache)) - results_list = await asyncio.gather(*coros) - results = {k: v for r in results_list for k, v in r.items()} + if self.state.args.eval_num_gpus > 0: + return RolloutFnEvalOutput(data=await self._eval_fleet.run()) + results = await run_eval_datasets(self.state, self.eval_prompt_dataset_cache) return RolloutFnEvalOutput(data=results) diff --git a/miles/rollout/inference_rollout/inference_rollout_eval.py b/miles/rollout/inference_rollout/inference_rollout_eval.py index 2747776791..82ec439438 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -12,13 +12,26 @@ ) from miles.utils.data import Dataset from miles.utils.eval_config import EvalDatasetConfig -from miles.utils.misc import as_completed_async from miles.utils.processing_utils import load_processor, load_tokenizer from miles.utils.types import Sample logger = logging.getLogger(__name__) +async def run_eval_datasets( + state: GenerateState, + prompt_dataset_cache: dict[Any, Dataset], +) -> dict[str, dict[str, Any]]: + args = state.args + assert not args.group_rm, "Group RM is not supported for eval rollout" + + coros = [] + for dataset_cfg in getattr(args, "eval_datasets", []) or []: + coros.append(eval_rollout_single_dataset(state, dataset_cfg, prompt_dataset_cache)) + results_list = await asyncio.gather(*coros) + return {k: v for r in results_list for k, v in r.items()} + + async def eval_rollout_single_dataset( state: GenerateState, dataset_cfg: EvalDatasetConfig, @@ -82,9 +95,17 @@ async def eval_rollout_single_dataset( ) data = [] + num_raised = 0 do_print = True pbar = tqdm(total=len(tasks), desc=f"Eval {dataset_cfg.name}", disable=not do_print) - async for sample in as_completed_async(tasks): + for future in asyncio.as_completed(tasks): + try: + sample = await future + except Exception as e: + logger.warning(f"Eval {dataset_cfg.name}: sample generation raised {e!r}") + num_raised += 1 + pbar.update(1) + continue if do_print: # TODO improve this after enhancing samples' type s = (sample[0] if len(sample) > 0 else None) if isinstance(sample, list) else sample @@ -102,13 +123,23 @@ async def eval_rollout_single_dataset( pbar.update(1) pbar.close() + if num_raised == len(tasks): + raise RuntimeError(f"Eval {dataset_cfg.name}: all {num_raised} sample generations failed") + data.sort(key=lambda sample: sample.index) + kept = [s for s in data if s.status != Sample.Status.ABORTED and s.reward is not None] + num_failed = len(data) - len(kept) + if num_failed: + logger.warning(f"Eval {dataset_cfg.name}: dropping {num_failed} aborted/reward-less samples") + data = kept + reward_key = args.eval_reward_key or args.reward_key return { dataset_cfg.name: { "rewards": [sample.reward if not reward_key else sample.reward[reward_key] for sample in data], "truncated": [sample.status == Sample.Status.TRUNCATED for sample in data], "samples": data, + "failed_samples": num_failed + num_raised, } } diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 858b5c0efd..917480552f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -972,6 +972,79 @@ def add_eval_arguments(parser): parser.add_argument("--eval-min-new-tokens", type=int, default=None) parser.add_argument("--eval-max-context-len", type=int, default=None) + parser.add_argument( + "--eval-num-gpus", + type=int, + default=0, + help=( + "Number of GPUs for a dedicated eval engine fleet. When > 0, eval runs on " + "its own engines behind its own router, synced by loading HF checkpoint " + "snapshots (never by joining training weight updates). 0 disables the " + "fleet and keeps today's shared-engine eval behavior." + ), + ) + parser.add_argument( + "--eval-num-gpus-per-engine", + type=int, + default=1, + help="GPUs per eval engine (TP size), independent of --rollout-num-gpus-per-engine.", + ) + parser.add_argument( + "--eval-hf-dir", + type=str, + default=None, + help=( + "Staging directory for per-eval HF snapshots (written to " + "`{eval_hf_dir}/step_{rollout_id}`). Point at tmpfs (e.g. /dev/shm/...) to " + "avoid disk. When unset and --save-hf is set, eval reuses the --save-hf " + "checkpoints instead of exporting its own snapshots." + ), + ) + parser.add_argument( + "--eval-model-path", + type=str, + default=None, + help="Boot checkpoint for eval engines. Defaults to --hf-checkpoint.", + ) + parser.add_argument( + "--eval-dispatch", + type=str, + choices=["async", "blocking"], + default="async", + help=( + "With a dedicated eval fleet, whether the training loop fires eval " + "fire-and-forget (async) or awaits it inline (blocking). Ignored when " + "--eval-num-gpus is 0." + ), + ) + parser.add_argument( + "--eval-max-in-flight", + type=int, + default=2, + help="Maximum number of concurrently pending async evals.", + ) + parser.add_argument( + "--eval-overflow-policy", + type=str, + choices=["backpressure", "skip"], + default="backpressure", + help=( + "What to do when an eval is due but --eval-max-in-flight evals are pending: " + "'backpressure' awaits the oldest pending eval (deterministic curve, bounded " + "stall); 'skip' drops the new eval point and logs eval/skipped_busy at that " + "step (training cadence is never stalled)." + ), + ) + parser.add_argument( + "--eval-keep-snapshots", + type=int, + default=2, + help=( + "How many snapshot dirs to keep under --eval-hf-dir (consumed snapshots " + "beyond this are deleted). --save-hf checkpoints are never deleted." + ), + ) + return parser def add_algo_arguments(parser): @@ -2206,7 +2279,7 @@ def parse_args_train_backend(): return args_partial.train_backend -def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]: +def resolve_eval_datasets(args) -> list[EvalDatasetConfig]: """ Build evaluation dataset configurations from either --eval-config or --eval-prompt-data. """ @@ -2264,7 +2337,7 @@ def _resolve_ft_components(args: argparse.Namespace) -> list[str]: def miles_validate_args(args): args.ft_components = _resolve_ft_components(args) - args.eval_datasets = _resolve_eval_datasets(args) + args.eval_datasets = resolve_eval_datasets(args) if args.mini_ft_controller_enable and args.control_server_port == 0: raise ValueError("--mini-ft-controller-enable requires --control-server-port to be set (non-zero)") @@ -2452,6 +2525,39 @@ def miles_validate_args(args): if args.eval_interval is not None: assert args.eval_datasets, "Evaluation datasets must be configured when eval_interval is set." + if args.eval_num_gpus > 0: + assert ( + enable_experimental_rollout_refactor() + ), "--eval-num-gpus requires the class-based rollout API (MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1)." + assert args.eval_interval is not None, "--eval-num-gpus requires --eval-interval." + assert args.eval_hf_dir is not None or args.save_hf is not None, ( + "--eval-num-gpus requires a snapshot source: set --eval-hf-dir (staging exports) " + "or --save-hf (reuse periodic HF checkpoints)." + ) + assert not args.colocate, ( + "--eval-num-gpus is not supported with --colocate; " + "use tools/checkpoint_eval_service.py against --save-hf checkpoints instead." + ) + assert ( + not args.debug_train_only and not args.debug_rollout_only + ), "--eval-num-gpus is not supported with debug_train_only/debug_rollout_only." + assert args.eval_num_gpus % args.eval_num_gpus_per_engine == 0, ( + f"eval_num_gpus ({args.eval_num_gpus}) must be divisible by " + f"eval_num_gpus_per_engine ({args.eval_num_gpus_per_engine})." + ) + assert args.eval_keep_snapshots >= args.eval_max_in_flight, ( + f"--eval-keep-snapshots ({args.eval_keep_snapshots}) must be >= --eval-max-in-flight " + f"({args.eval_max_in_flight}), otherwise a pending eval's snapshot could be GC'd." + ) + if args.eval_hf_dir is None: + assert args.save_interval is not None and args.eval_interval % args.save_interval == 0, ( + "Reusing --save-hf checkpoints for eval requires eval_interval to be a " + f"multiple of save_interval (got eval_interval={args.eval_interval}, " + f"save_interval={args.save_interval}). Set --eval-hf-dir for independent snapshots." + ) + if args.eval_model_path is None: + args.eval_model_path = args.hf_checkpoint + if args.save_interval is not None: assert args.save is not None, "'--save' is required when save_interval is set." diff --git a/miles/utils/hf_config.py b/miles/utils/hf_config.py index cd081798fa..b730859a13 100644 --- a/miles/utils/hf_config.py +++ b/miles/utils/hf_config.py @@ -13,6 +13,7 @@ import importlib from dataclasses import dataclass +from pathlib import Path from transformers import AutoConfig, AutoModelForCausalLM from transformers.models.auto.configuration_auto import CONFIG_MAPPING_NAMES @@ -107,3 +108,19 @@ def load_hf_config( def is_dsa(hf_config) -> bool: return getattr(hf_config, "model_type", None) in ("deepseek_v32", "glm_moe_dsa") + + +# Written by HF exports after all ranks finish, so consumers can tell finished from partial. +HF_EXPORT_COMPLETE_MARKER = ".complete" + + +def is_complete_hf_export(path: str | Path) -> bool: + return (Path(path) / HF_EXPORT_COMPLETE_MARKER).exists() + + +def looks_like_hf_checkpoint(path: str | Path) -> bool: + """Fallback for checkpoints written before the marker existed.""" + path = Path(path) + if not (path / "config.json").exists(): + return False + return any(path.glob("*.safetensors")) or any(path.glob("*.bin")) diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 9b27fdb613..a9ae32d836 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -224,6 +224,28 @@ async def _post(client, url, payload, max_retries=60, action="post", headers=Non return output +async def wait_http_ok(url: str, *, json_payload=None, timeout: float = 180.0, request_timeout: float = 60.0) -> None: + """Poll ``url`` until it answers HTTP 200 (POST when ``json_payload`` is given, + else GET); raise ``TimeoutError`` past the deadline.""" + deadline = time.time() + timeout + last_error = "no attempt made" + async with httpx.AsyncClient() as client: + while True: + try: + if json_payload is not None: + response = await client.post(url, json=json_payload, timeout=request_timeout) + else: + response = await client.get(url, timeout=request_timeout) + if response.status_code == 200: + return + last_error = f"HTTP {response.status_code}" + except httpx.HTTPError as e: + last_error = repr(e) + if time.time() > deadline: + raise TimeoutError(f"{url} not ready after {timeout}s: {last_error}") + await asyncio.sleep(5) + + def init_http_client(args): """Initialize HTTP client and optionally enable distributed POST via Ray.""" global _http_client, _client_concurrency, _distributed_post_enabled @@ -231,6 +253,8 @@ def init_http_client(args): return _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + if args.eval_num_gpus > 0: + _client_concurrency += args.sglang_server_concurrency * args.eval_num_gpus // args.eval_num_gpus_per_engine if _http_client is None: _http_client = httpx.AsyncClient( limits=httpx.Limits(max_connections=_client_concurrency), diff --git a/tests/e2e/short/test_qwen2.5_0.5B_fully_async_eval.py b/tests/e2e/short/test_qwen2.5_0.5B_fully_async_eval.py new file mode 100644 index 0000000000..f98f3bfbab --- /dev/null +++ b/tests/e2e/short/test_qwen2.5_0.5B_fully_async_eval.py @@ -0,0 +1,132 @@ +import os + +from tests.ci.ci_register import register_cuda_ci + +import miles.utils.external_utils.command_utils as U + +register_cuda_ci(est_time=400, suite="stage-c-8-gpu-h100", labels=["short"]) + +FEW_GPU = U.get_bool_env_var("MILES_TEST_FEW_GPU", "0") + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 4 if FEW_GPU else 8 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/gsm8k") + + +def execute(): + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + + rollout_args = ( + "--rollout-function-path miles.rollout.fully_async_rollout.FullyAsyncRolloutFn " + "--prompt-data /root/datasets/gsm8k/train.parquet " + "--input-key messages " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type math " + "--num-rollout 4 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 1024 " + "--rollout-temperature 0.8 " + "--global-batch-size 32 " + # retract (default) can deadlock flush_cache in fully_async under load + "--pause-generation-mode in_place " + ) + + # Dedicated eval fleet pinned to tmpfs HF snapshots: every eval point must + # report weight_version == its rollout_id (asserted by the pin verify in the + # eval controller; visible in eval/gsm8k/weight_version metrics). + eval_args = ( + "--eval-interval 2 " + "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet " + "--n-samples-per-eval-prompt 1 " + "--eval-max-response-len 1024 " + "--eval-top-k 1 " + "--eval-num-gpus 1 " + "--eval-num-gpus-per-engine 1 " + "--eval-hf-dir /dev/shm/miles_e2e_eval_hf " + "--eval-keep-snapshots 2 " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--use-kl-loss " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = "--rollout-num-gpus-per-engine 1 " "--sglang-mem-fraction-static 0.65 " "--sglang-enable-metrics " + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {1 if FEW_GPU else 2} " + f"--rollout-num-gpus {2 if FEW_GPU else 5} " + # HF-format --ref-load requires the bridge loader; eval snapshots are + # exported through the bridge path as well (marker-gated). + "--megatron-to-hf-mode bridge " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + train_script="train_async.py", + extra_env_vars={"MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1"}, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/fast/ray/rollout/conftest.py b/tests/fast/ray/rollout/conftest.py index 7e5e94fdbd..5eff8e2d98 100644 --- a/tests/fast/ray/rollout/conftest.py +++ b/tests/fast/ray/rollout/conftest.py @@ -30,6 +30,8 @@ def make_args(**overrides: Any) -> Namespace: # rollout core rollout_num_gpus=8, rollout_num_gpus_per_engine=1, + eval_num_gpus=0, + eval_num_gpus_per_engine=1, num_gpus_per_node=8, rollout_batch_size=8, n_samples_per_prompt=4, diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py new file mode 100644 index 0000000000..019ec98dae --- /dev/null +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -0,0 +1,667 @@ +from tests.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=60, suite="stage-a-cpu", labels=[]) + +import asyncio +from argparse import Namespace +from types import SimpleNamespace + +import pytest + +import miles.ray.rollout.rollout_manager as rollout_manager_mod +from miles.rollout.checkpoint_eval import make_eval_args, retarget_args + + +def make_args(**overrides) -> Namespace: + defaults = dict( + sglang_router_ip="10.0.0.1", + sglang_router_port=30000, + rollout_num_gpus=4, + rollout_num_gpus_per_engine=2, + eval_num_gpus=1, + eval_num_gpus_per_engine=1, + sglang_model_routers={"default": ("10.0.0.1", 30000), "eval": ("10.0.0.2", 31000)}, + ) + defaults.update(overrides) + return Namespace(**defaults) + + +def test_retarget_args_swaps_router_and_sizing(): + args = make_args() + eval_args = retarget_args(args, "10.0.0.9", 39000, num_gpus=2, num_gpus_per_engine=2) + + assert (eval_args.sglang_router_ip, eval_args.sglang_router_port) == ("10.0.0.9", 39000) + assert eval_args.rollout_num_gpus == 2 + assert eval_args.rollout_num_gpus_per_engine == 2 + # The original namespace is untouched. + assert (args.sglang_router_ip, args.sglang_router_port) == ("10.0.0.1", 30000) + assert args.rollout_num_gpus == 4 + + +def test_make_eval_args_reads_router_registry(): + args = make_args() + eval_args = make_eval_args(args) + + assert (eval_args.sglang_router_ip, eval_args.sglang_router_port) == ("10.0.0.2", 31000) + assert eval_args.rollout_num_gpus == args.eval_num_gpus + assert eval_args.rollout_num_gpus_per_engine == args.eval_num_gpus_per_engine + + +async def test_run_eval_datasets_merges_datasets(monkeypatch): + import miles.rollout.inference_rollout.inference_rollout_eval as eval_mod + + async def fake_single_dataset(state, cfg, cache): + return {cfg.name: {"rewards": [1.0], "truncated": [False], "samples": []}} + + monkeypatch.setattr(eval_mod, "eval_rollout_single_dataset", fake_single_dataset) + + state = SimpleNamespace( + args=Namespace(group_rm=False, eval_datasets=[SimpleNamespace(name="a"), SimpleNamespace(name="b")]) + ) + results = await eval_mod.run_eval_datasets(state, {}) + assert set(results.keys()) == {"a", "b"} + + +def _eval_dataset_env(monkeypatch, generate): + import miles.rollout.inference_rollout.inference_rollout_eval as eval_mod + from miles.utils.types import Sample + + monkeypatch.setattr(eval_mod, "generate_and_rm", generate) + monkeypatch.setattr(eval_mod, "compute_sampling_params", lambda args, **kw: {}) + args = Namespace( + group_rm=False, + hf_checkpoint="hf", + apply_chat_template=False, + chat_template_path=None, + reward_key=None, + eval_reward_key=None, + ) + dataset_cfg = SimpleNamespace( + name="ds", + cache_key=("ds",), + n_samples_per_eval_prompt=1, + temperature=1.0, + top_p=1.0, + top_k=-1, + max_response_len=16, + inject_metadata=lambda md: md, + ) + samples = [Sample(index=i, prompt="p", response="r", label="l", reward=1) for i in range(4)] + cache = {dataset_cfg.cache_key + ("hf", False, None): SimpleNamespace(samples=samples)} + return eval_mod, args, dataset_cfg, cache + + +async def test_single_dataset_tolerates_partial_failures(monkeypatch): + from miles.utils.types import Sample + + async def generate(state, sample, sampling_params, evaluation): + if sample.index == 0: + raise RuntimeError("engine died") + if sample.index == 1: + sample.status = Sample.Status.ABORTED + return sample + + eval_mod, args, dataset_cfg, cache = _eval_dataset_env(monkeypatch, generate) + result = await eval_mod.eval_rollout_single_dataset(SimpleNamespace(args=args), dataset_cfg, cache) + + assert result["ds"]["rewards"] == [1, 1] + assert result["ds"]["failed_samples"] == 2 # one raised + one aborted + + +async def test_single_dataset_all_failures_raise(monkeypatch): + async def generate(state, sample, sampling_params, evaluation): + raise RuntimeError("engine died") + + eval_mod, args, dataset_cfg, cache = _eval_dataset_env(monkeypatch, generate) + with pytest.raises(RuntimeError, match="all 4 sample generations failed"): + await eval_mod.eval_rollout_single_dataset(SimpleNamespace(args=args), dataset_cfg, cache) + + +# ---------------- controller (RolloutManager._eval_on_dedicated_fleet) ---------------- + + +class FakeRemoteMethod: + def __init__(self, engine, name): + self.engine = engine + self.name = name + + def remote(self, *args, **kwargs): + self.engine.log.append((self.name, args, kwargs)) + result = self.engine.responses[self.name](*args, **kwargs) + fut = asyncio.get_event_loop().create_future() + fut.set_result(result) + return fut + + +class FakeEngine: + def __init__(self, log): + self.log = log + self.weight_version = None + + def load(model_path, weight_version=None): + self.weight_version = weight_version + return None + + self.responses = { + "update_weights_from_disk": load, + "get_weight_version": lambda: self.weight_version, + } + + def __getattr__(self, name): + if name in ("update_weights_from_disk", "get_weight_version"): + return FakeRemoteMethod(self, name) + raise AttributeError(name) + + +class FakeServerEngineWrapper: + def __init__(self, actor): + self._actor = actor + self.is_allocated = True + self.stopped = False + + @property + def actor_handle(self): + return self._actor + + def mark_stopped(self): + self.stopped = True + self.is_allocated = False + + +class FakeEvalServer: + def __init__(self, engines): + self._engines = engines + self.wrappers = [FakeServerEngineWrapper(e) for e in engines] + self.recover_calls = 0 + + @property + def server_groups(self): + return [SimpleNamespace(all_engines=self.wrappers)] + + @property + def engines(self): + return [SimpleNamespace(actor_handle=e) for e in self._engines] + + async def recover(self): + self.recover_calls += 1 + + async def wait_all_engines_alive(self): + pass + + +def make_manager(args, engines, eval_fn_result=None): + mgr = object.__new__(rollout_manager_mod.RolloutManager.__ray_actor_class__) + mgr.args = args + mgr.rollout_id = 7 + mgr._eval_lock = asyncio.Lock() + mgr._eval_consumed_snapshots = [] + mgr.servers = {"eval": FakeEvalServer(engines)} + mgr._metric_checker = None + mgr.eval_generate_rollout = lambda input: eval_fn_result + return mgr + + +@pytest.fixture +def controller_env(monkeypatch, tmp_path): + log = [] + logged = {} + + def fake_call_rollout_function(fn, input): + log.append(("generate", input.rollout_id)) + return fn(input) + + monkeypatch.setattr(rollout_manager_mod, "call_rollout_function", fake_call_rollout_function) + + async def noop_router_ready(self, srv, timeout=180.0): + return None + + monkeypatch.setattr( + rollout_manager_mod.RolloutManager.__ray_actor_class__, "_wait_eval_router_ready", noop_router_ready + ) + monkeypatch.setattr(rollout_manager_mod, "save_debug_rollout_data", lambda *a, **k: None) + monkeypatch.setattr( + rollout_manager_mod, + "log_eval_rollout_data", + lambda rollout_id, args, data, extra: logged.setdefault("eval", (rollout_id, data, extra)) or {}, + ) + monkeypatch.setattr( + rollout_manager_mod, + "log_eval_skip", + lambda rollout_id, args, reason: logged.setdefault("skip", (rollout_id, reason)), + ) + return SimpleNamespace(log=log, logged=logged, tmp_path=tmp_path) + + +async def test_controller_pins_all_engines_before_generate(controller_env, tmp_path): + snapshot = tmp_path / "step_5" + snapshot.mkdir() + (snapshot / ".complete").touch() + + log = controller_env.log + engines = [FakeEngine(log), FakeEngine(log)] + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + result = SimpleNamespace(data={"ds": {"rewards": [1.0]}}, metrics=None) + mgr = make_manager(args, engines, eval_fn_result=result) + + await mgr._eval_on_dedicated_fleet(5, str(snapshot), export_time_seconds=1.5) + + load_events = [e for e in log if e[0] == "update_weights_from_disk"] + assert len(load_events) == 2 + assert all(e[2]["weight_version"] == "5" for e in load_events) + # Every load strictly precedes generation. + assert log.index(("generate", 5)) > max(i for i, e in enumerate(log) if e[0] == "update_weights_from_disk") + rollout_id, _data, extra = controller_env.logged["eval"] + assert rollout_id == 5 + assert extra["eval/lag_steps"] == 2 + assert extra["eval/export_time_seconds"] == 1.5 + + +async def test_controller_skips_on_missing_marker(controller_env, tmp_path): + snapshot = tmp_path / "step_5" + snapshot.mkdir() # no .complete marker + + engines = [FakeEngine(controller_env.log)] + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + mgr = make_manager(args, engines) + + await mgr._eval_on_dedicated_fleet(5, str(snapshot), export_time_seconds=None) + + assert controller_env.logged["skip"] == (5, "ckpt_missing") + assert not [e for e in controller_env.log if e[0] == "update_weights_from_disk"] + + +async def test_controller_base_checkpoint_needs_no_marker(controller_env, tmp_path): + engines = [FakeEngine(controller_env.log)] + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + result = SimpleNamespace(data={}, metrics=None) + mgr = make_manager(args, engines, eval_fn_result=result) + + await mgr._eval_on_dedicated_fleet(0, "/base", export_time_seconds=None) + + assert "eval" in controller_env.logged + load_events = [e for e in controller_env.log if e[0] == "update_weights_from_disk"] + assert len(load_events) == 1 and load_events[0][2]["weight_version"] == "0" + + +async def test_controller_pin_violation_skips_after_retry(controller_env, tmp_path): + snapshot = tmp_path / "step_5" + snapshot.mkdir() + (snapshot / ".complete").touch() + + engine = FakeEngine(controller_env.log) + engine.responses["get_weight_version"] = lambda: "999" # never matches + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + mgr = make_manager(args, [engine]) + + await mgr._eval_on_dedicated_fleet(5, str(snapshot), export_time_seconds=None) + + assert controller_env.logged["skip"] == (5, "pin_violation") + load_events = [e for e in controller_env.log if e[0] == "update_weights_from_disk"] + assert len(load_events) == 2 # one retry + assert ("generate", 5) not in controller_env.log + + +async def test_controller_zombie_engine_times_out_to_skip(controller_env, tmp_path, monkeypatch): + """A zombie engine never answers; the load timeout must skip the point + instead of holding the eval lock forever.""" + snapshot = tmp_path / "step_5" + snapshot.mkdir() + (snapshot / ".complete").touch() + + engine = FakeEngine(controller_env.log) + + class _NeverResolves: + def remote(self, *args, **kwargs): + return asyncio.get_event_loop().create_future() # never resolved + + engine.update_weights_from_disk_override = _NeverResolves() + monkeypatch.setattr( + type(engine), + "__getattr__", + lambda self, name: ( + self.update_weights_from_disk_override + if name == "update_weights_from_disk" + else FakeRemoteMethod(self, name) + ), + ) + monkeypatch.setattr(rollout_manager_mod, "EVAL_WEIGHT_LOAD_TIMEOUT_SECS", 0.05) + + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + mgr = make_manager(args, [engine]) + + await asyncio.wait_for(mgr._eval_on_dedicated_fleet(5, str(snapshot), export_time_seconds=None), timeout=5) + + assert controller_env.logged["skip"] == (5, "pin_violation") + assert not mgr._eval_lock.locked() + + +async def test_controller_marks_dead_engine_for_recovery(controller_env, tmp_path): + """A dead actor is marked stopped for revival; the eval degrades to a skip.""" + snapshot = tmp_path / "step_5" + snapshot.mkdir() + (snapshot / ".complete").touch() + + engine = FakeEngine(controller_env.log) + + def dead(*args, **kwargs): + raise RuntimeError("actor died") + + engine.responses["get_weight_version"] = dead + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + mgr = make_manager(args, [engine]) + + await mgr._eval_on_dedicated_fleet(5, str(snapshot), export_time_seconds=None) + + srv = mgr.servers["eval"] + assert srv.wrappers[0].stopped # probed, found unreachable, marked for revival + assert srv.recover_calls == 1 + assert controller_env.logged["skip"] == (5, "pin_violation") + + +async def test_controller_router_not_ready_skips(controller_env, tmp_path, monkeypatch): + snapshot = tmp_path / "step_5" + snapshot.mkdir() + (snapshot / ".complete").touch() + + async def router_never_ready(self, srv, timeout=180.0): + raise TimeoutError("router not ready") + + monkeypatch.setattr( + rollout_manager_mod.RolloutManager.__ray_actor_class__, "_wait_eval_router_ready", router_never_ready + ) + + engines = [FakeEngine(controller_env.log)] + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + mgr = make_manager(args, engines) + + await mgr._eval_on_dedicated_fleet(5, str(snapshot), export_time_seconds=None) + + assert controller_env.logged["skip"] == (5, "unhealthy") + assert ("generate", 5) not in controller_env.log + + +async def test_controller_gc_keeps_ring(controller_env, tmp_path): + engines = [FakeEngine(controller_env.log)] + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(tmp_path), eval_keep_snapshots=2) + result = SimpleNamespace(data={}, metrics=None) + mgr = make_manager(args, engines, eval_fn_result=result) + + dirs = [] + for rollout_id in (1, 2, 3): + snapshot = tmp_path / f"step_{rollout_id}" + snapshot.mkdir() + (snapshot / ".complete").touch() + dirs.append(snapshot) + controller_env.logged.pop("eval", None) + await mgr._eval_on_dedicated_fleet(rollout_id, str(snapshot), export_time_seconds=None) + + assert not dirs[0].exists() # oldest consumed snapshot beyond keep-2 is deleted + assert dirs[1].exists() and dirs[2].exists() + + +async def test_controller_never_deletes_outside_staging(controller_env, tmp_path): + engines = [FakeEngine(controller_env.log)] + save_hf = tmp_path / "save_hf" / "step_1" + save_hf.mkdir(parents=True) + (save_hf / ".complete").touch() + staging = tmp_path / "staging" + staging.mkdir() + args = make_args(hf_checkpoint="/base", eval_hf_dir=str(staging), eval_keep_snapshots=2) + result = SimpleNamespace(data={}, metrics=None) + mgr = make_manager(args, engines, eval_fn_result=result) + + await mgr._eval_on_dedicated_fleet(1, str(save_hf), export_time_seconds=None) + + assert save_hf.exists() + assert mgr._eval_consumed_snapshots == [] + + +# ---------------- driver (train_async.EvalDispatcher) ---------------- + + +class FakeManagerActor: + def __init__(self): + self.eval_calls = [] + self.skip_calls = [] + self._futures = [] + + outer = self + + class _Eval: + def remote(self, rollout_id, hf_dir=None, export_time_seconds=None): + outer.eval_calls.append((rollout_id, hf_dir, export_time_seconds)) + fut = asyncio.get_event_loop().create_future() + outer._futures.append(fut) + return fut + + class _Skip: + def remote(self, rollout_id, reason): + outer.skip_calls.append((rollout_id, reason)) + fut = asyncio.get_event_loop().create_future() + fut.set_result(None) + return fut + + self.eval = _Eval() + self.report_eval_skip = _Skip() + + def finish(self, index=0): + self._futures[index].set_result(None) + + +class FakeActorModel: + def __init__(self, fail=False): + self.exports = [] + self.fail = fail + + async def export_hf(self, rollout_id, path): + if self.fail: + raise RuntimeError("export boom") + self.exports.append((rollout_id, path)) + + +@pytest.fixture +def dispatcher_env(monkeypatch): + import miles.ray.rollout.eval_dispatch as eval_dispatch + + # ray.wait/ray.get over asyncio futures: done iff the future is resolved. + monkeypatch.setattr( + eval_dispatch.ray, "wait", lambda refs, timeout=0: (refs, []) if refs[0].done() else ([], refs) + ) + monkeypatch.setattr(eval_dispatch.ray, "get", lambda ref: ref.result()) + return eval_dispatch + + +def make_dispatcher(eval_dispatch, manager, actor_model, **arg_overrides): + dispatcher_defaults = dict( + eval_hf_dir="/dev/shm/eval_hf", + eval_dispatch="async", + eval_max_in_flight=2, + eval_overflow_policy="backpressure", + ) + dispatcher_defaults.update(arg_overrides) + args = make_args(**dispatcher_defaults) + return eval_dispatch.EvalDispatcher(args, actor_model, manager), args + + +async def test_dispatcher_exports_and_fires(dispatcher_env): + manager = FakeManagerActor() + actor_model = FakeActorModel() + dispatcher, _ = make_dispatcher(dispatcher_env, manager, actor_model) + + await dispatcher.dispatch(4) + + assert actor_model.exports == [(4, "/dev/shm/eval_hf/step_4")] + assert len(manager.eval_calls) == 1 + rollout_id, hf_dir, export_time = manager.eval_calls[0] + assert (rollout_id, hf_dir) == (4, "/dev/shm/eval_hf/step_4") + assert export_time is not None + assert len(dispatcher.pending) == 1 + + +async def test_dispatcher_export_failure_skips(dispatcher_env): + manager = FakeManagerActor() + dispatcher, _ = make_dispatcher(dispatcher_env, manager, FakeActorModel(fail=True)) + + await dispatcher.dispatch(4) + + assert manager.eval_calls == [] + assert manager.skip_calls == [(4, "export_failed")] + + +async def test_dispatcher_skip_policy_drops_before_export(dispatcher_env): + manager = FakeManagerActor() + actor_model = FakeActorModel() + dispatcher, _ = make_dispatcher( + dispatcher_env, manager, actor_model, eval_max_in_flight=1, eval_overflow_policy="skip" + ) + + await dispatcher.dispatch(1) + await dispatcher.dispatch(2) # at cap: dropped, no export + + assert manager.skip_calls == [(2, "busy")] + assert actor_model.exports == [(1, "/dev/shm/eval_hf/step_1")] + + +async def test_dispatcher_force_overrides_skip_policy(dispatcher_env): + """The final eval point must never be dropped: training is already over.""" + manager = FakeManagerActor() + actor_model = FakeActorModel() + dispatcher, _ = make_dispatcher( + dispatcher_env, manager, actor_model, eval_max_in_flight=1, eval_overflow_policy="skip" + ) + + await dispatcher.dispatch(1) + + async def finish_soon(): + await asyncio.sleep(0.01) + manager.finish(0) + + finisher = asyncio.create_task(finish_soon()) + await dispatcher.dispatch(2, force=True) # at cap: waits instead of dropping + await finisher + + assert manager.skip_calls == [] + assert [c[0] for c in manager.eval_calls] == [1, 2] + + +async def test_dispatcher_backpressure_awaits_oldest(dispatcher_env): + manager = FakeManagerActor() + dispatcher, _ = make_dispatcher(dispatcher_env, manager, FakeActorModel(), eval_max_in_flight=1) + + await dispatcher.dispatch(1) + + async def finish_soon(): + await asyncio.sleep(0.01) + manager.finish(0) + + finisher = asyncio.create_task(finish_soon()) + await dispatcher.dispatch(2) # must wait for eval 1 to finish + await finisher + + assert [c[0] for c in manager.eval_calls] == [1, 2] + assert len(dispatcher.pending) == 1 # only eval 2 pending + + +async def test_dispatcher_reuse_mode_uses_save_hf(dispatcher_env): + manager = FakeManagerActor() + actor_model = FakeActorModel() + dispatcher, _ = make_dispatcher( + dispatcher_env, manager, actor_model, eval_hf_dir=None, save_hf="/ckpt/hf/{rollout_id}" + ) + + await dispatcher.dispatch(10) + + assert actor_model.exports == [] # no extra export in reuse mode + assert manager.eval_calls[0][:2] == (10, "/ckpt/hf/10") + + +async def test_dispatcher_drain_awaits_all(dispatcher_env): + manager = FakeManagerActor() + dispatcher, _ = make_dispatcher(dispatcher_env, manager, FakeActorModel()) + + await dispatcher.dispatch(1) + await dispatcher.dispatch(2) + assert len(dispatcher.pending) == 2 + + manager.finish(0) + manager.finish(1) + await dispatcher.drain() + assert len(dispatcher.pending) == 0 + + +async def test_dispatcher_without_fleet_blocks_like_today(dispatcher_env): + manager = FakeManagerActor() + + class _LegacyEval: + def __init__(self): + self.calls = [] + + def remote(self, rollout_id): + self.calls.append(rollout_id) + fut = asyncio.get_event_loop().create_future() + fut.set_result(None) + return fut + + manager.eval = _LegacyEval() + dispatcher, _ = make_dispatcher(dispatcher_env, manager, FakeActorModel(), eval_num_gpus=0) + + await dispatcher.dispatch(3) + assert manager.eval.calls == [3] + assert len(dispatcher.pending) == 0 + + +# ---------------- external service (tools/checkpoint_eval_service.py) ---------------- + + +@pytest.fixture +def service_mod(): + import importlib.util + from pathlib import Path as _Path + + path = _Path(__file__).resolve().parents[3] / "tools" / "checkpoint_eval_service.py" + spec = importlib.util.spec_from_file_location("checkpoint_eval_service", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_find_ready_snapshots(service_mod, tmp_path): + from miles.utils.hf_config import HF_EXPORT_COMPLETE_MARKER + + def snapshot(name, *, marker=False, config=False, old=False): + d = tmp_path / name + d.mkdir() + if marker: + (d / HF_EXPORT_COMPLETE_MARKER).touch() + if config: + (d / "config.json").touch() + (d / "model.safetensors").touch() + if old: + import os + + stale = time.time() - 2 * service_mod.QUIESCENCE_SECS + for p in [d] + list(d.iterdir()): + os.utime(p, (stale, stale)) + return d + + import time + + snapshot("step_3", marker=True) + snapshot("step_7", config=True, old=True) # pre-marker, quiescent + snapshot("step_9", config=True) # still being written + snapshot("qwen2.5-step_12", marker=True) # id must come from the trailing number + snapshot("step_1", marker=True) # below min_rollout_id + snapshot("step_5", marker=True) # already consumed + (tmp_path / "notes.txt").touch() + + ready = service_mod.find_ready_snapshots(tmp_path, min_rollout_id=2, consumed={5}) + + assert [(rid, p.name) for rid, p in ready] == [(3, "step_3"), (7, "step_7"), (12, "qwen2.5-step_12")] + + +def test_snapshot_ledger_roundtrip(service_mod, tmp_path): + ledger = service_mod.SnapshotLedger(tmp_path) + ledger.mark(3) + ledger.mark(7) + + assert service_mod.SnapshotLedger(tmp_path).consumed == {3, 7} diff --git a/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py index 8f1453ece9..55d96ba77c 100644 --- a/tests/fast/rollout/test_fully_async_rollout.py +++ b/tests/fast/rollout/test_fully_async_rollout.py @@ -72,6 +72,7 @@ def make_args(**overrides) -> Namespace: max_weight_staleness=None, sglang_router_ip="127.0.0.1", sglang_router_port=30000, + eval_num_gpus=0, ) defaults.update(overrides) return Namespace(**defaults) @@ -105,11 +106,88 @@ async def test_drain_collects_batch_sorted_with_metrics(monkeypatch): assert len(output2.samples) == 3 -async def test_eval_raises(monkeypatch): - fn = make_fn(monkeypatch, make_args(), FakeDataSource()) - with pytest.raises(ValueError, match="does not serve eval"): - await fn(RolloutFnEvalInput(rollout_id=0)) +async def test_eval_without_fleet_pauses_producer(monkeypatch): + """Shared-engine eval: producer submissions pause during eval and resume after.""" + release = asyncio.Event() + + async def blocking_generate(state, group, sampling_params, evaluation=False): + await release.wait() + return group + + data_source = FakeDataSource() + fn = make_fn( + monkeypatch, make_args(rollout_batch_size=2, eval_num_gpus=0), data_source, generate=blocking_generate + ) + + eval_started = asyncio.Event() + eval_release = asyncio.Event() + eval_results = {"fake_ds": {"rewards": [1.0], "truncated": [False], "samples": []}} + + async def fake_run_eval_datasets(state, cache): + assert state is fn.state # shared-engine eval uses the train state + eval_started.set() + await eval_release.wait() + return eval_results + + monkeypatch.setattr(fully_async, "run_eval_datasets", fake_run_eval_datasets) + + # Start the producer via a train call, then run eval concurrently. + drain = asyncio.create_task(fn(RolloutFnTrainInput(rollout_id=0))) + await asyncio.sleep(0.05) + submitted_before_eval = data_source.num_get_calls + + eval_task = asyncio.create_task(fn(RolloutFnEvalInput(rollout_id=0))) + await eval_started.wait() + release.set() # in-flight groups finish and buffer, but no NEW submissions + await asyncio.sleep(0.05) + assert data_source.num_get_calls == submitted_before_eval + + eval_release.set() + output = await eval_task + assert output.data == eval_results + + # Producer resumes and the train drain completes. + assert (await drain).samples + + +async def test_eval_runs_on_dedicated_fleet(monkeypatch): + import miles.rollout.checkpoint_eval as checkpoint_eval + import miles.rollout.inference_rollout.inference_rollout_eval as eval_mod + + args = make_args( + eval_num_gpus=1, + eval_num_gpus_per_engine=1, + sglang_model_routers={"eval": ("127.0.0.1", 31000)}, + ) + data_source = FakeDataSource() + fn = make_fn(monkeypatch, args, data_source) + + eval_results = {"fake_ds": {"rewards": [1.0], "truncated": [False], "samples": []}} + seen_states = [] + + def fake_make_eval_generate_state(a): + assert a.sglang_model_routers["eval"] == ("127.0.0.1", 31000) + state = FakeGenerateState(a) + seen_states.append(state) + return state + + async def fake_run_eval_datasets(state, cache): + assert state in seen_states + return eval_results + + monkeypatch.setattr(checkpoint_eval, "make_eval_generate_state", fake_make_eval_generate_state) + monkeypatch.setattr(eval_mod, "run_eval_datasets", fake_run_eval_datasets) + + output = await fn(RolloutFnEvalInput(rollout_id=0)) + + assert output.data == eval_results + # Eval must not start the producer or consume training prompts. assert fn._worker is None + assert data_source.num_get_calls == 0 + + # The eval state is created once and reused across evals. + await fn(RolloutFnEvalInput(rollout_id=1)) + assert len(seen_states) == 1 async def test_aborted_group_recycled(monkeypatch): diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py new file mode 100644 index 0000000000..d3fea57a6f --- /dev/null +++ b/tools/checkpoint_eval_service.py @@ -0,0 +1,277 @@ +"""Standalone checkpoint eval service: watch --save-hf output, pin an sglang server +to each complete snapshot, run the miles eval datasets, log at the snapshot's step. +No Ray; restarts resume from a ledger file next to the watch dir. + +Example:: + + python tools/checkpoint_eval_service.py \\ + --watch-dir /ckpt/exp/hf --hf-checkpoint /models/Qwen3.5-4B --tp 1 \\ + --eval-prompt-data aime /data/aime-2024.jsonl --rm-type dapo --reward-key score +""" + +import argparse +import asyncio +import json +import logging +import re +import subprocess +import sys +import time +from argparse import Namespace +from pathlib import Path + +from miles.rollout.checkpoint_eval import retarget_args +from miles.utils.hf_config import is_complete_hf_export, looks_like_hf_checkpoint +from miles.utils.http_utils import init_http_client, post, wait_http_ok + +logger = logging.getLogger("checkpoint_eval_service") + +QUIESCENCE_SECS = 120.0 + +# Fields the eval path consumes that are not service flags; values match miles defaults. +EVAL_ARG_DEFAULTS = dict( + chat_template_path=None, + custom_generate_function_path=None, + custom_eval_rollout_log_function_path=None, + custom_rm_path=None, + group_rm=False, + partial_rollout=False, + mask_offpolicy_in_partial_rollout=False, + multimodal_keys=None, + metadata_key="metadata", + tool_key=None, + apply_chat_template_kwargs=None, + rollout_stop=None, + rollout_stop_token_ids=None, + rollout_skip_special_tokens=True, + rollout_max_context_len=None, + rollout_seed=42, + rm_url=None, + sglang_enable_deterministic_inference=False, + sglang_speculative_algorithm=None, + sglang_server_concurrency=512, + use_distributed_post=False, + use_rollout_routing_replay=False, + use_rollout_indexer_replay=False, + use_opd=False, + lora_rank=0, + lora_adapter_path=None, + log_passrate=False, + log_reward_category=None, + advantage_estimator="grpo", + load_debug_rollout_data=None, + wandb_always_use_train_step=False, + eval_num_gpus=0, # the service talks to its own server; no in-job fleet + ci_test=False, + use_wandb=False, # set by init_service_tracking from --wandb-mode +) + + +def parse_service_args() -> Namespace: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + # snapshot source + parser.add_argument("--watch-dir", type=str, required=True, help="Directory containing HF snapshot subdirs.") + parser.add_argument("--poll-interval", type=float, default=60.0) + parser.add_argument("--min-rollout-id", type=int, default=0) + parser.add_argument( + "--catchup", + type=str, + choices=["all", "latest"], + default="all", + help="On startup/backlog: eval every unconsumed snapshot, or skip to the newest.", + ) + parser.add_argument("--once", action="store_true", help="Process the current backlog and exit.") + # server + parser.add_argument("--hf-checkpoint", type=str, required=True, help="Base checkpoint (tokenizer/arch source).") + parser.add_argument("--server-url", type=str, default=None, help="Attach to a running sglang server.") + parser.add_argument("--num-gpus", type=int, default=1) + parser.add_argument("--tp", type=int, default=1) + parser.add_argument("--server-port", type=int, default=31000) + parser.add_argument("--sglang-mem-fraction-static", type=float, default=0.8) + # eval datasets (mirrors the miles eval surface) + parser.add_argument("--eval-prompt-data", type=str, nargs="+", default=None) + parser.add_argument("--eval-config", type=str, default=None) + parser.add_argument("--eval-input-key", type=str, default=None) + parser.add_argument("--eval-label-key", type=str, default=None) + parser.add_argument("--eval-tool-key", type=str, default=None) + parser.add_argument("--n-samples-per-eval-prompt", type=int, default=1) + parser.add_argument("--eval-temperature", type=float, default=None) + parser.add_argument("--eval-top-p", type=float, default=None) + parser.add_argument("--eval-top-k", type=int, default=None) + parser.add_argument("--eval-max-response-len", type=int, default=None) + parser.add_argument("--eval-max-prompt-len", type=int, default=None) + parser.add_argument("--rollout-temperature", type=float, default=1.0) + parser.add_argument("--rollout-top-p", type=float, default=1.0) + parser.add_argument("--rollout-top-k", type=int, default=-1) + parser.add_argument("--rollout-max-response-len", type=int, default=8192) + parser.add_argument("--input-key", type=str, default="prompt") + parser.add_argument("--label-key", type=str, default=None) + parser.add_argument("--apply-chat-template", action="store_true", default=False) + parser.add_argument("--rm-type", type=str, default=None) + parser.add_argument("--reward-key", type=str, default=None) + parser.add_argument("--eval-reward-key", type=str, default=None) + # step mapping (only needed with --wandb-always-use-train-step trainers) + parser.add_argument("--rollout-batch-size", type=int, default=None) + parser.add_argument("--n-samples-per-prompt", type=int, default=None) + parser.add_argument("--global-batch-size", type=int, default=None) + # tracking + parser.add_argument("--wandb-mode", type=str, choices=["shared", "separate", "off"], default="off") + parser.add_argument("--wandb-project", type=str, default=None) + parser.add_argument("--wandb-group", type=str, default=None) + parser.add_argument("--wandb-run-id", type=str, default=None) + parser.add_argument("--wandb-key", type=str, default=None) + parser.add_argument("--wandb-host", type=str, default=None) + return parser.parse_args() + + +def build_eval_namespace(service_args: Namespace, server_ip: str, server_port: int) -> Namespace: + from miles.utils.arguments import resolve_eval_datasets + + args = Namespace(**EVAL_ARG_DEFAULTS) + for key, value in vars(service_args).items(): + setattr(args, key, value) + args.eval_datasets = resolve_eval_datasets(args) + return retarget_args(args, server_ip, server_port, service_args.num_gpus, service_args.tp) + + +class SnapshotLedger: + def __init__(self, watch_dir: Path): + self.path = watch_dir / ".eval_service_state.json" + self.consumed: set[int] = set() + if self.path.exists(): + self.consumed = set(json.loads(self.path.read_text()).get("consumed", [])) + + def mark(self, rollout_id: int) -> None: + self.consumed.add(rollout_id) + self.path.write_text(json.dumps({"consumed": sorted(self.consumed)})) + + +def find_ready_snapshots(watch_dir: Path, min_rollout_id: int, consumed: set[int]) -> list[tuple[int, Path]]: + ready = [] + for child in watch_dir.iterdir(): + if not child.is_dir(): + continue + match = re.search(r"(\d+)$", child.name) + if match is None: + continue + rollout_id = int(match.group(1)) + if rollout_id < min_rollout_id or rollout_id in consumed: + continue + if is_complete_hf_export(child): + ready.append((rollout_id, child)) + elif looks_like_hf_checkpoint(child): + newest_mtime = max(p.stat().st_mtime for p in child.iterdir()) + if time.time() - newest_mtime > QUIESCENCE_SECS: + ready.append((rollout_id, child)) + return sorted(ready) + + +def launch_server(service_args: Namespace) -> tuple[subprocess.Popen | None, str, int]: + if service_args.server_url is not None: + url = service_args.server_url.removeprefix("http://") + ip, port = url.split(":") + return None, ip, int(port) + + cmd = [ + sys.executable, + "-m", + "sglang.launch_server", + "--model-path", + service_args.hf_checkpoint, + "--tp", + str(service_args.tp), + "--host", + "127.0.0.1", + "--port", + str(service_args.server_port), + "--mem-fraction-static", + str(service_args.sglang_mem_fraction_static), + "--trust-remote-code", + ] + logger.info(f"Launching sglang server: {' '.join(cmd)}") + proc = subprocess.Popen(cmd) + return proc, "127.0.0.1", service_args.server_port + + +async def wait_server_healthy(ip: str, port: int, timeout: float = 1800.0) -> None: + await wait_http_ok(f"http://{ip}:{port}/health_generate", timeout=timeout) + + +async def evaluate_snapshot(args: Namespace, state, cache: dict, rollout_id: int, snapshot: Path) -> None: + from miles.ray.rollout.metrics import log_eval_rollout_data + from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets + from miles.utils.http_utils import get + + url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}" + start = time.time() + weight_version = str(rollout_id) + await post( + f"{url}/update_weights_from_disk", + {"model_path": str(snapshot), "weight_version": weight_version}, + ) + info = await get(f"{url}/model_info") + if str(info.get("weight_version")) != weight_version: + raise RuntimeError( + f"weight_version pin failed: engine reports {info.get('weight_version')}, expected {weight_version}" + ) + + results = await run_eval_datasets(state, cache) + extra = {"eval/duration_seconds": time.time() - start} + log_eval_rollout_data(rollout_id, args, results, extra) + + +def init_service_tracking(args: Namespace) -> None: + if args.wandb_mode == "off": + args.use_wandb = False + return + from miles.utils.tracking_utils.tracking import init_tracking + + args.use_wandb = True + if args.wandb_mode == "shared" and not args.wandb_run_id: + raise ValueError("--wandb-mode shared requires --wandb-run-id of the training run") + init_tracking(args, primary=args.wandb_mode == "separate") + + +async def main() -> None: + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s") + service_args = parse_service_args() + watch_dir = Path(service_args.watch_dir) + if not watch_dir.is_dir(): + raise FileNotFoundError(f"--watch-dir {watch_dir} does not exist") + + proc, server_ip, server_port = launch_server(service_args) + try: + args = build_eval_namespace(service_args, server_ip, server_port) + init_http_client(args) + await wait_server_healthy(server_ip, server_port) + init_service_tracking(args) + + from miles.rollout.inference_rollout.inference_rollout_common import GenerateState + + state = GenerateState(args) + cache: dict = {} + ledger = SnapshotLedger(watch_dir) + + while True: + ready = find_ready_snapshots(watch_dir, service_args.min_rollout_id, ledger.consumed) + if ready and service_args.catchup == "latest": + for rollout_id, _ in ready[:-1]: + ledger.mark(rollout_id) + ready = ready[-1:] + for rollout_id, snapshot in ready: + logger.info(f"Evaluating snapshot {snapshot} (rollout_id={rollout_id})") + try: + await evaluate_snapshot(args, state, cache, rollout_id, snapshot) + ledger.mark(rollout_id) + except Exception: + logger.exception(f"Eval of {snapshot} failed; will retry next scan") + if service_args.once: + break + await asyncio.sleep(service_args.poll_interval) + finally: + if proc is not None: + proc.terminate() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/train_async.py b/train_async.py index 8f72c10693..a158ce6b3d 100644 --- a/train_async.py +++ b/train_async.py @@ -2,6 +2,7 @@ import logging from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models +from miles.ray.rollout.eval_dispatch import EvalDispatcher from miles.utils.arguments import parse_args from miles.utils.async_utils import eager_create_task from miles.utils.audit_utils.process_identity import MainProcessIdentity @@ -52,8 +53,10 @@ async def train(args): skip_list=args.check_weight_update_skip_list, ) + eval_dispatcher = EvalDispatcher(args, actor_model, rollout_manager) + if args.eval_interval is not None and args.start_rollout_id == 0 and not args.skip_eval_before_train: - await rollout_manager.eval.remote(0) + await eval_dispatcher.dispatch(0, hf_dir=args.hf_checkpoint) # async train loop. rollout_data_next_future = rollout_manager.generate.remote(args.start_rollout_id) @@ -92,8 +95,8 @@ async def train(args): rollout_data_next_future = None await actor_model.update_weights(rollout_id=rollout_id) - if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch): - await rollout_manager.eval.remote(rollout_id) + if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch, args.num_rollout): + await eval_dispatcher.dispatch(rollout_id, force=rollout_id == args.num_rollout - 1) if ( args.debug_exit_after_rollout is not None @@ -106,6 +109,7 @@ async def train(args): ) break + await eval_dispatcher.drain() await rollout_manager.dispose.remote()