From 162b110c8b66f23df5597c9835297a4410b4a371 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:04 -0700 Subject: [PATCH 01/33] eval: factor run_eval_datasets out; add eval-fleet retargeting helpers run_eval_datasets moves next to eval_rollout_single_dataset so three consumers (InferenceRolloutFn, FullyAsyncRolloutFn, the checkpoint eval service) share one eval driver. checkpoint_eval.py retargets a shallow-copied args namespace at the eval fleet's router with eval GPU sizing, which fixes both the hardcoded router URL in the generate path and the concurrency semaphore in one place. eval_rollout_single_dataset drops ABORTED/reward-less samples (engine died mid-eval) with a failed_samples count instead of corrupting dataset metrics. --- miles/rollout/checkpoint_eval.py | 68 +++++++++++++++++++ .../inference_rollout_common.py | 22 +++--- .../inference_rollout_eval.py | 24 +++++++ 3 files changed, 105 insertions(+), 9 deletions(-) create mode 100644 miles/rollout/checkpoint_eval.py diff --git a/miles/rollout/checkpoint_eval.py b/miles/rollout/checkpoint_eval.py new file mode 100644 index 0000000000..9d3bf1d06b --- /dev/null +++ b/miles/rollout/checkpoint_eval.py @@ -0,0 +1,68 @@ +"""Eval against a dedicated eval fleet pinned to checkpoint snapshots. + +The eval fleet is a second model entry (``name="eval"``, ``update_weights=False``) +behind its own router; it never joins the training weight-update group and is never +paused or aborted by training. Weights reach it exclusively through +``update_weights_from_disk`` on an HF snapshot exported for a specific rollout_id, +so every eval point measures one well-defined weight version. + +The helpers here retarget an args namespace at the eval fleet so the existing +generate machinery (which reads ``args.sglang_router_ip/port`` and sizes its +concurrency semaphore off ``args.rollout_num_gpus``) works against it unchanged. +""" + +import copy +from argparse import Namespace + +from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnEvalOutput +from miles.rollout.inference_rollout.inference_rollout_common import GenerateState + +__all__ = ["retarget_args", "make_eval_args", "make_eval_generate_state", "CheckpointEvalRolloutFn"] + + +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 is all that is needed to run + the standard eval path against a different set of engines. The original ``args`` + is not modified. + """ + 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: + """Retarget ``args`` at the in-job eval fleet's router (requires servers started).""" + 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 CheckpointEvalRolloutFn: + """Eval-only rollout function running against the dedicated eval fleet. + + Addressable via ``--eval-function-path`` when the train rollout function should + not serve eval itself. + """ + + def __init__(self, input: RolloutFnConstructorInput): + self.args = input.args + self._state: GenerateState | None = None + self._prompt_dataset_cache = {} + + async def __call__(self, input: RolloutFnEvalInput) -> RolloutFnEvalOutput: + from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets + + assert input.evaluation, "CheckpointEvalRolloutFn only serves eval" + if self._state is None: + self._state = make_eval_generate_state(self.args) + results = await run_eval_datasets(self._state, self._prompt_dataset_cache) + return RolloutFnEvalOutput(data=results) diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py index 9f1cc603b0..cb0f306289 100644 --- a/miles/rollout/inference_rollout/inference_rollout_common.py +++ b/miles/rollout/inference_rollout/inference_rollout_common.py @@ -174,6 +174,7 @@ def __init__(self, input: RolloutFnConstructorInput): self.data_source = input.data_source self.state = GenerateState(input.args) self.eval_prompt_dataset_cache = {} + self._eval_state: GenerateState | None = None async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: if input.evaluation: @@ -190,13 +191,16 @@ 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" - - 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()} + from miles.rollout.checkpoint_eval import make_eval_generate_state + from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets + + if getattr(self.state.args, "eval_num_gpus", 0) > 0: + # A dedicated eval fleet exists: run eval against its router with a + # state sized for it, instead of contending with train generation. + if self._eval_state is None: + self._eval_state = make_eval_generate_state(self.state.args) + state = self._eval_state + else: + state = self.state + results = await run_eval_datasets(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..df9eedec97 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -19,6 +19,21 @@ logger = logging.getLogger(__name__) +async def run_eval_datasets( + state: GenerateState, + prompt_dataset_cache: dict[Any, Dataset], +) -> dict[str, dict[str, Any]]: + """Run every configured eval dataset against the engines behind ``state.args``'s router.""" + 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, @@ -104,11 +119,20 @@ async def eval_rollout_single_dataset( data.sort(key=lambda sample: sample.index) + # A sample can come back ABORTED or reward-less if an engine died mid-eval; + # drop it and report the count instead of corrupting the dataset metrics. + 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, } } From 9fbe31c4e9fea964006d833c657e2f745fa2ef47 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:04 -0700 Subject: [PATCH 02/33] fully-async: serve eval from the dedicated eval fleet FullyAsyncRolloutFn now answers evaluation=True itself by running the standard eval datasets against the eval fleet's router (lazy GenerateState, cached eval prompt datasets), instead of raising unconditionally. Without --eval-num-gpus the raise remains, now pointing at the fleet flag and the external service. --- miles/rollout/fully_async_rollout.py | 35 ++++++++++++---- .../fast/rollout/test_fully_async_rollout.py | 41 ++++++++++++++++++- 2 files changed, 67 insertions(+), 9 deletions(-) diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py index fa7f0b5023..0cde09dcc7 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 requires a dedicated eval fleet (``--eval-num-gpus``): the rollout fleet +never has a quiet window, so eval runs on separate engines pinned per-eval to an HF +checkpoint snapshot (see ``miles/rollout/checkpoint_eval.py``). """ 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 make_eval_generate_state 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,31 @@ def __init__(self, input: RolloutFnConstructorInput): self._weight_version = _CachedWeightVersion() self._worker: asyncio.Task | None = None self._output: asyncio.Queue[Group] | None = None + self._eval_state: GenerateState | None = None + self._eval_prompt_dataset_cache: dict = {} 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 getattr(self.args, "eval_num_gpus", 0) <= 0: + raise ValueError( + "fully-async eval requires a dedicated eval fleet: set --eval-num-gpus > 0 " + "(or run tools/checkpoint_eval_service.py against --save-hf checkpoints)" + ) + # The eval fleet has its own router and engines, so eval coroutines coexist + # with the producer task on the shared loop without contending for capacity. + if self._eval_state is None: + self._eval_state = make_eval_generate_state(self.args) + results = await run_eval_datasets(self._eval_state, self._eval_prompt_dataset_cache) + return RolloutFnEvalOutput(data=results) + # -------------------------- producer -------------------------- def _max_in_flight_groups(self) -> int: diff --git a/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py index 8f1453ece9..e39fe666e7 100644 --- a/tests/fast/rollout/test_fully_async_rollout.py +++ b/tests/fast/rollout/test_fully_async_rollout.py @@ -105,13 +105,50 @@ async def test_drain_collects_batch_sorted_with_metrics(monkeypatch): assert len(output2.samples) == 3 -async def test_eval_raises(monkeypatch): +async def test_eval_without_fleet_raises(monkeypatch): fn = make_fn(monkeypatch, make_args(), FakeDataSource()) - with pytest.raises(ValueError, match="does not serve eval"): + with pytest.raises(ValueError, match="requires a dedicated eval fleet"): await fn(RolloutFnEvalInput(rollout_id=0)) assert fn._worker is None +async def test_eval_runs_on_dedicated_fleet(monkeypatch): + 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(fully_async, "make_eval_generate_state", fake_make_eval_generate_state) + monkeypatch.setattr(fully_async, "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): aborted = make_group(1, status=Sample.Status.ABORTED) data_source = FakeDataSource(scripted=[aborted]) From 2c81c8a167482d97683bc60d422710c3255f6b0a Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:23 -0700 Subject: [PATCH 03/33] args: dedicated eval fleet flags and validation --eval-num-gpus/--eval-num-gpus-per-engine size the fleet; --eval-hf-dir stages per-eval HF snapshots (tmpfs for the no-disk debug case) with --save-hf reuse as the zero-export production mode; --eval-dispatch/--eval-max-in-flight/ --eval-overflow-policy bound the fire-and-forget pipeline; --eval-keep-snapshots rings the staging dir and is validated >= max-in-flight so a pending eval's snapshot can never be GC'd. --- miles/utils/arguments.py | 108 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 858b5c0efd..b0203ae3a3 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -972,6 +972,80 @@ 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) + # Dedicated eval fleet (checkpoint-interfaced eval; required for fully-async eval). + 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): @@ -2452,6 +2526,40 @@ 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: + # Reuse mode: every eval-due step must coincide with a save-due step. + 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." From 3058a58c8ec4e52d89cf766459512d986a18dee3 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:23 -0700 Subject: [PATCH 04/33] ray: eval fleet placement, server config synthesis, http client sizing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The eval fleet is a second ModelConfig (name=eval, update_weights=False — the default inference would mark it updatable and trip the exactly-one-updatable assert) appended in _resolve_sglang_config; the placement group grows by eval_num_gpus at the tail of the rollout slice, past megatron's range, so the offload choreography skips it for free. rollout_num_gpus keeps meaning training rollout GPUs everywhere; the shared http client is sized for both fleets. --- miles/ray/placement_group.py | 3 +- miles/ray/rollout/rollout_server.py | 48 +++++++++++++++++++++++------ miles/utils/http_utils.py | 3 ++ 3 files changed, 44 insertions(+), 10 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index a9e7122c45..5e24dc5ee4 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -107,7 +107,8 @@ 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 + # eval_num_gpus adds a dedicated eval fleet at the tail of the rollout slice. + 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/rollout_server.py b/miles/ray/rollout/rollout_server.py index 55d230b0f5..07412dcf46 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/miles/ray/rollout/rollout_server.py @@ -104,24 +104,54 @@ 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 = getattr(args, "eval_num_gpus", 0) + 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: + # Dedicated eval fleet: own router, never receives training weight updates + # (weights are pinned per-eval via update_weights_from_disk). update_weights + # must be explicit — resolve() would infer True from the matching model_path. + 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/utils/http_utils.py b/miles/utils/http_utils.py index 9b27fdb613..96a5022aef 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -231,6 +231,9 @@ def init_http_client(args): return _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine + if getattr(args, "eval_num_gpus", 0) > 0: + # The eval fleet is served by the same client; size for both. + _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), From e29a79f01bd88ba2d493f8d2a82cc059a7424e35 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:23 -0700 Subject: [PATCH 05/33] megatron: on-demand HF export with completeness marker save_hf_model gains an explicit path, a raise_on_error mode for the on-demand eval snapshot path (the periodic --save-hf path keeps swallow-and-log), a rank-0 .complete marker written after a barrier so consumers can distinguish finished exports from partial ones, and a cached AutoBridge (it reloaded the HF config on every call). export_hf is exposed through MegatronTrainRayActor and RayTrainGroup. --- miles/backends/megatron_utils/actor.py | 18 +++++++++- miles/backends/megatron_utils/model.py | 46 +++++++++++++++++++++----- miles/ray/actor_group.py | 4 +++ miles/ray/train_actor.py | 4 +++ miles/utils/hf_export.py | 22 ++++++++++++ 5 files changed, 85 insertions(+), 9 deletions(-) create mode 100644 miles/utils/hf_export.py diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index b922417598..6a0ff727b0 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,22 @@ 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). + + 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() + if self.args.debug_rollout_only: + return + + 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..076ba2735b 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_export import HF_EXPORT_COMPLETE_MARKER 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,24 @@ def save( enable_forward_pre_hook(model) -def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: +_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,26 +852,28 @@ 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 so consumers (eval snapshot verification, the + external eval service) can distinguish finished exports from partial ones. 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 (used by the + on-demand eval snapshot path, where the caller wants to skip that eval). """ 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) + bridge = _get_hf_bridge(args.hf_checkpoint) path.mkdir(parents=True, exist_ok=True) @@ -862,22 +882,32 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: # adapter weights into base weights for a standalone HF model. bridge.save_hf_pretrained(model, path=path) + # save_hf_pretrained is collective; make sure every rank is done writing + # before the marker claims the snapshot is complete. + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + (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/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/utils/hf_export.py b/miles/utils/hf_export.py new file mode 100644 index 0000000000..6eaee146e0 --- /dev/null +++ b/miles/utils/hf_export.py @@ -0,0 +1,22 @@ +"""Shared constants and helpers for HF checkpoint export and consumption. + +Kept dependency-free: imported by the megatron export path, the rollout manager's +eval controller, and the standalone checkpoint eval service. +""" + +from pathlib import Path + +HF_EXPORT_COMPLETE_MARKER = ".complete" + + +def is_complete_hf_export(path: str | Path) -> bool: + """Whether ``path`` holds a finished HF export (the completeness marker exists).""" + return (Path(path) / HF_EXPORT_COMPLETE_MARKER).exists() + + +def looks_like_hf_checkpoint(path: str | Path) -> bool: + """Marker-less fallback heuristic 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")) From 6aead46b2886004624136b4d3b11dbae32159977 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:51 -0700 Subject: [PATCH 06/33] rollout manager: checkpoint-pinned eval controller RolloutManager.eval(rollout_id, hf_dir) pins the eval fleet to the snapshot before generating: an eval lock serializes load -> verify -> generate; every engine loads via update_weights_from_disk with weight_version=str(rollout_id) and the version is verified (one retry, then the point is skipped with eval/pin_violation). Fleet recovery is self-healing at eval time; a missing completeness marker or unhealthy fleet degrades to a skipped point logged at that rollout_id (eval/skipped_*), never a stall. Consumed staging snapshots are GC'd beyond the keep ring; --save-hf checkpoints are never touched. lag/duration/ export-time metrics ride the eval payload; log_eval_skip makes curve gaps attributable from the dashboard. --- miles/ray/rollout/metrics.py | 14 ++++ miles/ray/rollout/rollout_manager.py | 100 ++++++++++++++++++++++++++- 2 files changed, 112 insertions(+), 2 deletions(-) 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..2c61d51223 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,6 +29,7 @@ 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.hf_export import is_complete_hf_export from miles.utils.http_utils import init_http_client from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker @@ -87,6 +90,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 +140,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 getattr(self.args, "eval_num_gpus", 0) > 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 +164,93 @@ 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's weights to the snapshot for ``rollout_id``, then run eval. + + Every failure mode degrades to a skipped point logged at ``rollout_id`` with a + reason counter — never wrong data, never a training stall. + """ + start_time = time.time() + # Serializes load -> verify -> generate per fleet: no snapshot load can start + # while an eval is generating, and nothing dispatches before the version is + # confirmed on every engine. This is the version-pinning enforcement. + async with self._eval_lock: + srv = self.servers["eval"] + try: + 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._log_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._log_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): + await asyncio.gather( + *[e.update_weights_from_disk.remote(hf_dir, weight_version=weight_version) for e in engines] + ) + versions = await asyncio.gather(*[e.get_weight_version.remote() for e in engines]) + if 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._log_eval_skip(rollout_id, "pin_violation") + 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) + + def report_eval_skip(self, rollout_id: int, reason: str) -> None: + """Log a skipped eval point at ``rollout_id`` (called by the driver, e.g. on export failure).""" + self._log_eval_skip(rollout_id, reason) + + def _log_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 staging snapshots beyond the keep ring. + + Only dirs under --eval-hf-dir are ever deleted; --save-hf checkpoints and the + base checkpoint are never touched. Pending evals reference unconsumed dirs, + which are never GC candidates. + """ + 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) From 3acecf698230432c7e7ccceabbedf1ab4e2f5a4c Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:51 -0700 Subject: [PATCH 07/33] train_async: fire-and-forget eval dispatch EvalDispatcher exports the HF snapshot for an eval-due step (or reuses the --save-hf checkpoint), fires eval.remote without awaiting, and bounds pending evals via --eval-max-in-flight with backpressure or skip overflow policy. Finished refs are reaped nonblocking each step; all pending evals are drained before dispose so the last points land and MetricChecker sees completed evals. Export failure degrades to a skipped point. The final rollout now always gets an eval point (num_rollout passed to should_run_periodic_action), and the pre-train eval uses the base checkpoint as its snapshot. --- tests/fast/rollout/test_checkpoint_eval.py | 427 +++++++++++++++++++++ train_async.py | 94 ++++- 2 files changed, 518 insertions(+), 3 deletions(-) create mode 100644 tests/fast/rollout/test_checkpoint_eval.py diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py new file mode 100644 index 0000000000..ac81fdbb40 --- /dev/null +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -0,0 +1,427 @@ +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"} + + +# ---------------- controller (RolloutManager._eval_on_dedicated_fleet) ---------------- + + +class FakeRemoteMethod: + """Mimics a Ray actor method: .remote(...) returns an awaitable.""" + + 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 FakeEvalServer: + def __init__(self, engines): + self._engines = engines + self.recover_calls = 0 + + @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) + 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_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 train_async + + # ray.wait/ray.get over asyncio futures: done iff the future is resolved. + monkeypatch.setattr(train_async.ray, "wait", lambda refs, timeout=0: (refs, []) if refs[0].done() else ([], refs)) + monkeypatch.setattr(train_async.ray, "get", lambda ref: ref.result()) + return train_async + + +def make_dispatcher(train_async, 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 train_async.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_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 diff --git a/train_async.py b/train_async.py index 8f72c10693..8c4c2838b7 100644 --- a/train_async.py +++ b/train_async.py @@ -1,5 +1,10 @@ import asyncio import logging +import os +import time +from collections import deque + +import ray from miles.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models from miles.utils.arguments import parse_args @@ -15,6 +20,85 @@ logger = logging.getLogger(__name__) +class EvalDispatcher: + """Dispatch evals against the dedicated eval fleet without stalling training. + + With ``--eval-num-gpus 0`` this degrades to today's blocking ``eval.remote`` call. + Otherwise each due eval gets an HF snapshot (exported to ``--eval-hf-dir``, or the + ``--save-hf`` checkpoint in reuse mode) and is fired fire-and-forget; the pending + set is bounded by ``--eval-max-in-flight`` via the ``--eval-overflow-policy``. + Failures degrade to a skipped point logged at that rollout_id, 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) -> 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": + 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: + """Await every pending eval (before dispose, so the last points always land).""" + 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: + # Reuse mode: the --save-hf checkpoint for this step was just written + # (eval_interval is validated to be a multiple of save_interval). + 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: + # The manager's eval lock serializes evals, so pending refs finish in order. + 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") + + 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") + + # The framework supports other asynchronous approaches such as fully async (see miles/rollout/fully_async_rollout.py). async def train(args): assert not args.colocate, "Colocation is not supported for async training." @@ -52,8 +136,11 @@ 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) + # The base checkpoint is the snapshot for the pre-train eval; no export needed. + 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 +179,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) if ( args.debug_exit_after_rollout is not None @@ -106,6 +193,7 @@ async def train(args): ) break + await eval_dispatcher.drain() await rollout_manager.dispose.remote() From b2a176465063c85ec1aecee90c96b90595f60f13 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:51 -0700 Subject: [PATCH 08/33] tools: standalone checkpoint eval service MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Watches --save-hf output (marker-gated, with an mtime-quiescence fallback for pre-marker checkpoints), pins its own sglang server per snapshot via /update_weights_from_disk with weight_version=str(rollout_id), runs the shared run_eval_datasets core, and logs at the snapshot's step — late points land at the right x because the step travels in the payload. A ledger file makes restarts idempotent and backfills missed snapshots; --catchup latest skips a backlog. No Ray, no GPU carve-out from the training job; works for colocate runs and after the run ends. --- tools/checkpoint_eval_service.py | 300 +++++++++++++++++++++++++++++++ 1 file changed, 300 insertions(+) create mode 100644 tools/checkpoint_eval_service.py diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py new file mode 100644 index 0000000000..8dd1e7ee51 --- /dev/null +++ b/tools/checkpoint_eval_service.py @@ -0,0 +1,300 @@ +"""Standalone checkpoint eval service. + +Watches a directory of HF checkpoint exports (``--save-hf`` output), and for each new +complete snapshot: pins an sglang server to it via ``/update_weights_from_disk`` +(stamping ``weight_version=str(rollout_id)``), runs the standard miles eval datasets +against it, and logs the metrics at the snapshot's ``rollout_id`` — hours-late points +land at the right x because the step travels inside the payload (``eval/step``). + +Runs anywhere with sglang + miles importable — another node, spot capacity, or after +the training run. No Ray, no training-job membership. A ledger file next to the watch +dir makes restarts idempotent and backfills missed snapshots. + +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 \\ + --wandb-mode separate + +The watch dir is scanned for subdirectories whose name contains the rollout id +(``step_{id}`` or any ``{rollout_id}``-formatted layout); a snapshot is consumed once +it has the ``.complete`` marker (or, for checkpoints written before the marker +existed, once it looks like a finished HF export and has been quiescent). +""" + +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_export import is_complete_hf_export, looks_like_hf_checkpoint +from miles.utils.http_utils import init_http_client, post + +logger = logging.getLogger("checkpoint_eval_service") + +QUIESCENCE_SECS = 120.0 + +# Fields the eval path consumes that are not exposed as service flags; values match +# the miles argument defaults. tests/fast/rollout/test_checkpoint_eval_service.py pins +# this table against the eval path's consumption. +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_seed=42, + sglang_enable_deterministic_inference=False, + 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, + wandb_always_use_train_step=False, + eval_num_gpus=0, # the service talks to its own server; no in-job fleet + ci_test=False, +) + + +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("--use-wandb", action="store_true", default=False) + 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: + """Compose the namespace the miles eval path consumes.""" + 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: + """Persistent record of consumed snapshots, next to the watch dir.""" + + 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]]: + """Unconsumed snapshot dirs ready for eval, ordered by rollout_id.""" + 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): + # Pre-marker checkpoint: accept once quiescent. + 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: + from miles.utils.http_utils import get + + deadline = time.time() + timeout + while time.time() < deadline: + try: + await get(f"http://{ip}:{port}/health_generate", max_retries=1) + return + except Exception: + await asyncio.sleep(5) + raise TimeoutError(f"sglang server at {ip}:{port} not healthy after {timeout}s") + + +async def eval_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") + assert ( + str(info.get("weight_version")) == weight_version + ), 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": + assert args.wandb_run_id, "--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) + assert watch_dir.is_dir(), f"--watch-dir {watch_dir} does not exist" + + proc, server_ip, server_port = launch_server(service_args) + try: + await wait_server_healthy(server_ip, server_port) + + args = build_eval_namespace(service_args, server_ip, server_port) + init_http_client(args) + 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 eval_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()) From c31aa5555dcb08f65c8afdeed6da4585882251bb Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 01:28:51 -0700 Subject: [PATCH 09/33] docs/examples: fully-async eval Qwen3.5-4B launch script with a 4+3+1 GPU split (actor/rollout/eval) and a fully-async docs section covering the two postures (tmpfs staging for on-time debug curves; --save-hf reuse and the external service for production), the weight_version namespace, and the skip counters. --- docs/user-guide/cli-reference.md | 8 + docs/user-guide/fully-async.md | 37 ++++ examples/fully_async/README.md | 6 +- .../run-qwen3.5-4b-fully_async-eval.sh | 168 ++++++++++++++++++ 4 files changed, 218 insertions(+), 1 deletion(-) create mode 100644 examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh 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..56f1d7d64a 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -100,6 +100,43 @@ 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 + +The rollout fleet never has a quiet window, so fully-async eval runs on a **dedicated +eval fleet** synced through HF checkpoint snapshots — never by joining training weight +updates. Enable it with: + +```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. + ## 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..7181ae2725 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,11 @@ 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 +Eval runs on a dedicated eval fleet synced via HF checkpoint snapshots (`--eval-num-gpus`, +`--eval-hf-dir`); see `run-qwen3.5-4b-fully_async-eval.sh` and the fully-async docs page. + ## 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..22f414d8c2 --- /dev/null +++ b/examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh @@ -0,0 +1,168 @@ +#!/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. Carving one +# GPU out of rollout costs ~25% rollout throughput vs the 4-GPU baseline; in exchange +# eval runs continuously without ever pausing training. Eval weights are pinned +# per-eval to an HF snapshot exported to tmpfs (--eval-hf-dir on /dev/shm), so eval is +# never blocked by disk and every point measures exactly one weight version. + +# 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[@]} From 63068ca70fe0023849dfce0de4addbe00a09f2e4 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 02:01:21 -0700 Subject: [PATCH 10/33] megatron: export eval snapshots via the direct converters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found in GPU verification: AutoBridge auto-selects Qwen35VLBridge for Qwen3.5, whose mapping matches none of the text spec's megatron params — it exports config/tokenizer but zero weight files while reporting success. export_hf now uses miles' own megatron->HF converters (the weight updater's machinery, so export coverage matches weight-sync coverage exactly), streaming safetensors shards from the gathered chunks with an index and metadata copied from the base checkpoint; bridge mode falls back to save_hf_model. The bridge path gains a no-weight-files check before writing the completeness marker, and the eval controller downgrades weight-load failures to a skipped point (pin_violation) instead of raising out of the fire-and-forget ref. --- miles/backends/megatron_utils/actor.py | 6 +- miles/backends/megatron_utils/model.py | 126 +++++++++++++++++++++---- miles/ray/rollout/rollout_manager.py | 14 ++- 3 files changed, 123 insertions(+), 23 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 6a0ff727b0..327f102bb7 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -568,8 +568,10 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None: def export_hf(self, rollout_id: int, path: str) -> None: """Export current weights as an HF checkpoint to ``path`` (collective). - 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. + 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() if self.args.debug_rollout_only: diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 076ba2735b..500f681b79 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 load_hf_config from miles.utils.hf_export import HF_EXPORT_COMPLETE_MARKER from miles.utils.memory_utils import clear_memory from miles.utils.test_utils.ft_test_actions import FTTestActionActorExecutor @@ -826,6 +827,78 @@ def save( enable_forward_pre_hook(model) +HF_METADATA_SKIP_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".gguf") + + +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. + + This is the same conversion machinery the weight updater uses, so it covers exactly + the model families weight sync supports — unlike the bridge-based ``save_hf_model``, + which silently exports zero weights for specs it has no mapping for (e.g. the + qwen3.5 attention-output-gate layout). + + Collective — all ranks must call it (the iterator gathers PP/EP/TP shards + internally and every rank materializes the full tensors); global rank 0 writes + safetensors shards streamingly plus the index, copies tokenizer/config metadata + from the base checkpoint, and stamps the completeness marker. + """ + 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" + for meta_file in Path(args.hf_checkpoint).iterdir(): + # Copy tokenizer/config metadata only — never base weight files or the + # base checkpoint's safetensors index (ours is written below). + if meta_file.is_file() and not any(s in meta_file.name for s in HF_METADATA_SKIP_SUFFIXES): + shutil.copy2(meta_file, path / meta_file.name) + index = {"metadata": {"total_size": total_size}, "weight_map": weight_map} + (path / "model.safetensors.index.json").write_text(json.dumps(index, indent=2)) + + # Everyone waits for the writer before the marker claims the snapshot complete. + torch.distributed.barrier() + if is_writer: + (path / HF_EXPORT_COMPLETE_MARKER).touch() + + _hf_bridge_cache: dict = {} @@ -868,25 +941,46 @@ def save_hf_model( path = Path(path if path is not None else args.save_hf.format(rollout_id=rollout_id)) try: - from miles.utils.megatron_bridge_utils import patch_megatron_model - if should_log: logger.info(f"Saving model in HuggingFace format to {path}") - 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) - - # save_hf_pretrained is collective; make sure every rank is done writing - # before the marker claims the snapshot is complete. - torch.distributed.barrier() - if torch.distributed.get_rank() == 0: - (path / HF_EXPORT_COMPLETE_MARKER).touch() + if args.megatron_to_hf_mode == "raw" and not is_lora_model(model): + # miles' own converters cover exactly what weight sync covers; the + # bridge silently exports zero weights for specs it has no mapping for + # (e.g. qwen3.5). 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) + + # save_hf_pretrained is collective; make sure every rank is done writing + # before checking the result and claiming the snapshot complete. + torch.distributed.barrier() + if torch.distributed.get_rank() == 0: + # The bridge silently exports zero weights for specs it has no mapping + # for; a marked-but-weightless snapshot must never exist. + 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}") diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 2c61d51223..49128bbd6b 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -192,11 +192,15 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti engines = [e.actor_handle for e in srv.engines] weight_version = str(rollout_id) for _attempt in range(2): - await asyncio.gather( - *[e.update_weights_from_disk.remote(hf_dir, weight_version=weight_version) for e in engines] - ) - versions = await asyncio.gather(*[e.get_weight_version.remote() for e in engines]) - if all(str(v) == weight_version for v in versions): + try: + await asyncio.gather( + *[e.update_weights_from_disk.remote(hf_dir, weight_version=weight_version) for e in engines] + ) + versions = await asyncio.gather(*[e.get_weight_version.remote() for e in engines]) + 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( From 1ce152d5dc505256f72dc4c815c06b88d4f8e730 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 03:16:22 -0700 Subject: [PATCH 11/33] tools: fix eval service arg contract and health probe (found in GPU test) Running the service against real snapshots surfaced missing namespace fields the eval path consumes (rollout_max_context_len, sglang_speculative_algorithm, rm_url, log_reward_category, advantage_estimator, load_debug_rollout_data) and two health-probe bugs: miles' get() helper takes no kwargs and json-decodes the response body, but /health_generate returns a non-JSON 200 (probe now uses httpx directly), and the http client must be initialized before the probe. --- tools/checkpoint_eval_service.py | 25 +++++++++++++++++-------- 1 file changed, 17 insertions(+), 8 deletions(-) diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index 8dd1e7ee51..d79d957b16 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -60,8 +60,11 @@ 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, @@ -70,6 +73,9 @@ 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, @@ -209,14 +215,18 @@ def launch_server(service_args: Namespace) -> tuple[subprocess.Popen | None, str async def wait_server_healthy(ip: str, port: int, timeout: float = 1800.0) -> None: - from miles.utils.http_utils import get + import httpx deadline = time.time() + timeout - while time.time() < deadline: - try: - await get(f"http://{ip}:{port}/health_generate", max_retries=1) - return - except Exception: + async with httpx.AsyncClient() as client: + while time.time() < deadline: + try: + # /health_generate runs a tiny generation; 200 body is not JSON. + response = await client.get(f"http://{ip}:{port}/health_generate", timeout=60) + if response.status_code == 200: + return + except httpx.HTTPError: + pass await asyncio.sleep(5) raise TimeoutError(f"sglang server at {ip}:{port} not healthy after {timeout}s") @@ -263,10 +273,9 @@ async def main() -> None: proc, server_ip, server_port = launch_server(service_args) try: - await wait_server_healthy(server_ip, server_port) - 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 From 7b688456a3c5af6a930cb90d21996aa1b08ee0a7 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 03:45:19 -0700 Subject: [PATCH 12/33] rollout manager: bound eval weight loads against zombie engines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Kill-testing surfaced a hang the design missed: killing an engine's sglang scheduler subprocess leaves the Ray actor alive but its backend dead — the update_weights_from_disk call is accepted and never answered (the shared http client has no timeout), holding the eval lock forever and eventually stalling the training loop through the dispatcher's backpressure await. Each load attempt is now bounded by a 600s timeout that degrades the point to eval/pin_violation and releases the lock. --- miles/ray/rollout/rollout_manager.py | 24 +++++++++++++-- tests/fast/rollout/test_checkpoint_eval.py | 35 ++++++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 49128bbd6b..36d916272b 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -43,6 +43,11 @@ logger = logging.getLogger(__name__) +# Bounds each eval-fleet weight load (per attempt). Generous enough for large +# checkpoints from network storage; small enough that a zombie engine costs a +# skipped point, not a wedged eval pipeline. +EVAL_WEIGHT_LOAD_TIMEOUT_SECS = 600.0 + @ray.remote class RolloutManager: @@ -193,10 +198,23 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti weight_version = str(rollout_id) for _attempt in range(2): try: - await asyncio.gather( - *[e.update_weights_from_disk.remote(hf_dir, weight_version=weight_version) for e in engines] + # The timeout bounds the zombie-engine case: an actor whose sglang + # backend died accepts the call but never answers (the shared http + # client has no timeout), which would otherwise hold the eval lock + # forever and eventually stall the driver via backpressure. + 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, ) - versions = await asyncio.gather(*[e.get_weight_version.remote() for e in engines]) except Exception as e: logger.warning(f"Eval fleet weight load from {hf_dir} failed: {e}") versions = [] diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index ac81fdbb40..465fa8c46d 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -221,6 +221,41 @@ async def test_controller_pin_violation_skips_after_retry(controller_env, tmp_pa assert ("generate", 5) not in controller_env.log +async def test_controller_zombie_engine_times_out_to_skip(controller_env, tmp_path, monkeypatch): + """An engine whose backend died accepts the call but never answers; the load + timeout must convert that into a skipped point instead of holding the eval + lock forever (which would eventually stall the driver via backpressure).""" + 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_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) From eb74bb4d07d04ba3b0fbbba0f4395c039e19e2b5 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 04:14:16 -0700 Subject: [PATCH 13/33] rollout manager: probe eval engines so recover() actually self-heals MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The actor-kill test showed recover() never revives a dead eval engine unless fault tolerance's health monitor marked it stopped — without FT every subsequent eval skipped with pin_violation instead of self-healing. The controller now probes each allocated engine before recovery (60s bound) and marks unreachable ones stopped itself, which covers both the dead-actor and zombie-backend cases; the eval that observed the death still degrades to a skipped point, and the next eval runs on the revived engine. --- miles/ray/rollout/rollout_manager.py | 22 +++++++++++ tests/fast/rollout/test_checkpoint_eval.py | 45 ++++++++++++++++++++++ 2 files changed, 67 insertions(+) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 36d916272b..09fc9445ff 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -182,6 +182,7 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti 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: @@ -243,6 +244,27 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti self._gc_eval_snapshots(hf_dir) + async def _mark_unreachable_eval_engines(self, srv) -> None: + """Probe eval engines and mark unreachable ones stopped so recover() revives them. + + Without fault tolerance nothing records an engine death (recover() only + restarts engines already marked stopped), so a dead or zombie eval engine + would otherwise fail every future eval instead of self-healing. + """ + 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 a skipped eval point at ``rollout_id`` (called by the driver, e.g. on export failure).""" self._log_eval_skip(rollout_id, reason) diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index 465fa8c46d..d7082ca7e2 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -100,11 +100,31 @@ def __getattr__(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] @@ -256,6 +276,31 @@ def remote(self, *args, **kwargs): assert not mgr._eval_lock.locked() +async def test_controller_marks_dead_engine_for_recovery(controller_env, tmp_path): + """A dead engine actor must be marked stopped so recover() revives it; the + eval itself degrades to a skipped point (the fake recover cannot really + replace the actor), never a raise out of the controller.""" + 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_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) From 75ff11fe1310eef84f7109fa12b2a07e21a2df93 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 04:54:14 -0700 Subject: [PATCH 14/33] eval: survive the post-revival router window and per-sample failures MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The revival test showed the remaining gap: after an eval engine is revived, the router needs a health-check cycle to evict the dead worker, and a generate dispatched inside that window 503s — one raised request killed the whole eval point. The controller now proves the route end-to-end with a retried one-token probe before dispatching (skipping with eval/skipped_unhealthy if it never comes up), and eval sample generation tolerates individual raises (counted into failed_samples; the point is only abandoned if every sample failed). --- miles/ray/rollout/rollout_manager.py | 32 +++++++++++++++++++ .../inference_rollout_eval.py | 16 ++++++++-- tests/fast/rollout/test_checkpoint_eval.py | 29 +++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 09fc9445ff..16811db7ae 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -228,6 +228,13 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti self._log_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._log_eval_skip(rollout_id, "unhealthy") + return + result = await asyncio.to_thread( call_rollout_function, self.eval_generate_rollout, RolloutFnEvalInput(rollout_id=rollout_id) ) @@ -244,6 +251,31 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti self._gc_eval_snapshots(hf_dir) + async def _wait_eval_router_ready(self, srv, timeout: float = 180.0) -> None: + """Probe end-to-end generation through the eval router before dispatching. + + After an engine revival the router needs a health-check cycle to evict the + dead worker and pick up the new one; a generate dispatched inside that window + gets 503s. A one-token probe (retried) proves the route is actually usable. + """ + import httpx + + url = f"http://{srv.router_ip}:{srv.router_port}/generate" + payload = {"input_ids": [0], "sampling_params": {"max_new_tokens": 1, "temperature": 0}} + deadline = time.time() + timeout + async with httpx.AsyncClient() as client: + while True: + try: + response = await client.post(url, json=payload, timeout=60) + 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"eval router at {url} not ready after {timeout}s: {last_error}") + await asyncio.sleep(5) + async def _mark_unreachable_eval_engines(self, srv) -> None: """Probe eval engines and mark unreachable ones stopped so recover() revives them. diff --git a/miles/rollout/inference_rollout/inference_rollout_eval.py b/miles/rollout/inference_rollout/inference_rollout_eval.py index df9eedec97..09f30ed197 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -12,7 +12,6 @@ ) 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 @@ -97,9 +96,19 @@ 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: + # One failed request (engine crash mid-eval, transient router error) + # costs one sample, not the whole eval point. + 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 @@ -117,6 +126,9 @@ 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) # A sample can come back ABORTED or reward-less if an engine died mid-eval; diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index d7082ca7e2..4ea2538d53 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -158,6 +158,13 @@ def fake_call_rollout_function(fn, input): 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, @@ -301,6 +308,28 @@ def dead(*args, **kwargs): 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) From 370624378334f13c1ebf10294ed14cad49090cb2 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 05:59:53 -0700 Subject: [PATCH 15/33] tests: e2e for fully-async eval on a dedicated fleet Qwen2.5-0.5B on gsm8k, 2 actor + 5 rollout + 1 eval GPUs (FEW_GPU: 1+2+1), FullyAsyncRolloutFn with eval-interval 2 and tmpfs snapshot staging via the bridge export path (HF-format ref-load requires the bridge loader). Verified green on 8xH200: eval points at 0/1/3 with pinned weight versions, job succeeded. --- .../test_qwen2.5_0.5B_fully_async_eval.py | 132 ++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 tests/e2e/short/test_qwen2.5_0.5B_fully_async_eval.py 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() From d58d8721244d07245c7059f4f58b5fb1fe7b7798 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 06:19:54 -0700 Subject: [PATCH 16/33] =?UTF-8?q?eval:=20address=20review=20=E2=80=94=20hu?= =?UTF-8?q?b-id=20metadata=20guard,=20service=20errors,=20formatting?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The direct HF export no longer crashes when hf_checkpoint is a hub model id rather than a local dir (metadata copy is skipped with a warning); the eval service raises typed errors instead of asserts for its runtime checks; black formatting for the last round of edits. --- miles/backends/megatron_utils/model.py | 20 ++++++++++++-------- tools/checkpoint_eval_service.py | 14 ++++++++------ 2 files changed, 20 insertions(+), 14 deletions(-) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 500f681b79..c1fc984d35 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -863,9 +863,7 @@ def export_hf_model_direct( if is_writer: path.mkdir(parents=True, exist_ok=True) - iterator = HfWeightIteratorDirect( - args, model, model_name=model_name, quantization_config=quantization_config - ) + iterator = HfWeightIteratorDirect(args, model, model_name=model_name, quantization_config=quantization_config) weight_map: dict[str, str] = {} total_size = 0 @@ -885,11 +883,17 @@ def export_hf_model_direct( if is_writer: assert weight_map, f"HF export to {path} produced no weights" - for meta_file in Path(args.hf_checkpoint).iterdir(): - # Copy tokenizer/config metadata only — never base weight files or the - # base checkpoint's safetensors index (ours is written below). - if meta_file.is_file() and not any(s in meta_file.name for s in HF_METADATA_SKIP_SUFFIXES): - shutil.copy2(meta_file, path / meta_file.name) + base_checkpoint = Path(args.hf_checkpoint) + if base_checkpoint.is_dir(): + for meta_file in base_checkpoint.iterdir(): + # Copy tokenizer/config metadata only — never base weight files or the + # base checkpoint's safetensors index (ours is written below). + if meta_file.is_file() and not any(s in meta_file.name for s in HF_METADATA_SKIP_SUFFIXES): + shutil.copy2(meta_file, path / meta_file.name) + else: + # hf_checkpoint can be a hub model id; the snapshot then lacks + # tokenizer/config metadata and consumers must point at the base. + 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)) diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index d79d957b16..0290a2f981 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -244,9 +244,10 @@ async def eval_snapshot(args: Namespace, state, cache: dict, rollout_id: int, sn {"model_path": str(snapshot), "weight_version": weight_version}, ) info = await get(f"{url}/model_info") - assert ( - str(info.get("weight_version")) == weight_version - ), f"weight_version pin failed: engine reports {info.get('weight_version')}, expected {weight_version}" + 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} @@ -260,8 +261,8 @@ def init_service_tracking(args: Namespace) -> None: from miles.utils.tracking_utils.tracking import init_tracking args.use_wandb = True - if args.wandb_mode == "shared": - assert args.wandb_run_id, "--wandb-mode shared requires --wandb-run-id of the training run" + 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") @@ -269,7 +270,8 @@ 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) - assert watch_dir.is_dir(), f"--watch-dir {watch_dir} does not exist" + 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: From 4779ca8638a1b5ff4b2d6da2797fae0fc91a5fb6 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 10:55:12 -0700 Subject: [PATCH 17/33] eval: fold hf_export helpers into hf_config miles/utils/hf_export.py was a three-symbol fragment whose name suggested it did exporting; the marker constant and the two HF-checkpoint-dir predicates belong with the existing HF checkpoint utilities in hf_config.py. --- miles/backends/megatron_utils/model.py | 3 +-- miles/ray/rollout/rollout_manager.py | 2 +- miles/utils/hf_config.py | 20 ++++++++++++++++++++ miles/utils/hf_export.py | 22 ---------------------- tools/checkpoint_eval_service.py | 2 +- 5 files changed, 23 insertions(+), 26 deletions(-) delete mode 100644 miles/utils/hf_export.py diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index c1fc984d35..9bde9b3063 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -30,8 +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 load_hf_config -from miles.utils.hf_export import HF_EXPORT_COMPLETE_MARKER +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 diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 16811db7ae..f1a021ea9e 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -29,7 +29,7 @@ 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.hf_export import is_complete_hf_export +from miles.utils.hf_config import is_complete_hf_export from miles.utils.http_utils import init_http_client from miles.utils.logging_utils import configure_logger from miles.utils.metric_checker import MetricChecker diff --git a/miles/utils/hf_config.py b/miles/utils/hf_config.py index cd081798fa..3bcb49af35 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,22 @@ def load_hf_config( def is_dsa(hf_config) -> bool: return getattr(hf_config, "model_type", None) in ("deepseek_v32", "glm_moe_dsa") + + +# Marker written by HF checkpoint exports (save_hf_model / export_hf) once every rank +# has finished writing, so consumers (the eval snapshot verification, the external +# checkpoint eval service) can distinguish finished exports from partial ones. +HF_EXPORT_COMPLETE_MARKER = ".complete" + + +def is_complete_hf_export(path: str | Path) -> bool: + """Whether ``path`` holds a finished HF export (the completeness marker exists).""" + return (Path(path) / HF_EXPORT_COMPLETE_MARKER).exists() + + +def looks_like_hf_checkpoint(path: str | Path) -> bool: + """Marker-less fallback heuristic 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/hf_export.py b/miles/utils/hf_export.py deleted file mode 100644 index 6eaee146e0..0000000000 --- a/miles/utils/hf_export.py +++ /dev/null @@ -1,22 +0,0 @@ -"""Shared constants and helpers for HF checkpoint export and consumption. - -Kept dependency-free: imported by the megatron export path, the rollout manager's -eval controller, and the standalone checkpoint eval service. -""" - -from pathlib import Path - -HF_EXPORT_COMPLETE_MARKER = ".complete" - - -def is_complete_hf_export(path: str | Path) -> bool: - """Whether ``path`` holds a finished HF export (the completeness marker exists).""" - return (Path(path) / HF_EXPORT_COMPLETE_MARKER).exists() - - -def looks_like_hf_checkpoint(path: str | Path) -> bool: - """Marker-less fallback heuristic 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/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index 0290a2f981..736bb4bbf1 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -35,7 +35,7 @@ from pathlib import Path from miles.rollout.checkpoint_eval import retarget_args -from miles.utils.hf_export import is_complete_hf_export, looks_like_hf_checkpoint +from miles.utils.hf_config import is_complete_hf_export, looks_like_hf_checkpoint from miles.utils.http_utils import init_http_client, post logger = logging.getLogger("checkpoint_eval_service") From fc5d58f1602a06e558bd19fd32b3a23c36222925 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 11:01:40 -0700 Subject: [PATCH 18/33] eval: trim comments to the necessary ones --- .../run-qwen3.5-4b-fully_async-eval.sh | 8 +--- miles/backends/megatron_utils/model.py | 35 +++++---------- miles/ray/rollout/rollout_manager.py | 44 +++++-------------- miles/rollout/checkpoint_eval.py | 25 +++-------- miles/rollout/fully_async_rollout.py | 7 +-- .../inference_rollout_common.py | 3 +- .../inference_rollout_eval.py | 7 +-- miles/utils/hf_config.py | 7 +-- tests/fast/rollout/test_checkpoint_eval.py | 11 ++--- tools/checkpoint_eval_service.py | 30 +++---------- train_async.py | 14 ++---- 11 files changed, 48 insertions(+), 143 deletions(-) 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 index 22f414d8c2..e074dcb4f6 100644 --- a/examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh +++ b/examples/fully_async/run-qwen3.5-4b-fully_async-eval.sh @@ -1,11 +1,7 @@ #!/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. Carving one -# GPU out of rollout costs ~25% rollout throughput vs the 4-GPU baseline; in exchange -# eval runs continuously without ever pausing training. Eval weights are pinned -# per-eval to an HF snapshot exported to tmpfs (--eval-hf-dir on /dev/shm), so eval is -# never blocked by disk and every point measures exactly one weight version. +# 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 diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 9bde9b3063..21fac08214 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -840,15 +840,9 @@ def export_hf_model_direct( ) -> None: """Export current weights as an HF checkpoint via miles' own megatron->HF converters. - This is the same conversion machinery the weight updater uses, so it covers exactly - the model families weight sync supports — unlike the bridge-based ``save_hf_model``, - which silently exports zero weights for specs it has no mapping for (e.g. the - qwen3.5 attention-output-gate layout). - - Collective — all ranks must call it (the iterator gathers PP/EP/TP shards - internally and every rank materializes the full tensors); global rank 0 writes - safetensors shards streamingly plus the index, copies tokenizer/config metadata - from the base checkpoint, and stamps the completeness marker. + 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 @@ -884,19 +878,16 @@ def export_hf_model_direct( assert weight_map, f"HF export to {path} produced no weights" base_checkpoint = Path(args.hf_checkpoint) if base_checkpoint.is_dir(): + # Tokenizer/config metadata only — the skip list also excludes the base + # checkpoint's safetensors index, which would clobber ours below. for meta_file in base_checkpoint.iterdir(): - # Copy tokenizer/config metadata only — never base weight files or the - # base checkpoint's safetensors index (ours is written below). if meta_file.is_file() and not any(s in meta_file.name for s in HF_METADATA_SKIP_SUFFIXES): shutil.copy2(meta_file, path / meta_file.name) else: - # hf_checkpoint can be a hub model id; the snapshot then lacks - # tokenizer/config metadata and consumers must point at the base. 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)) - # Everyone waits for the writer before the marker claims the snapshot complete. torch.distributed.barrier() if is_writer: (path / HF_EXPORT_COMPLETE_MARKER).touch() @@ -929,16 +920,14 @@ def save_hf_model( so it can be loaded with ``PeftModel.from_pretrained``. This function is collective — all ranks must call it. On success, global rank 0 - writes a ``.complete`` marker file so consumers (eval snapshot verification, the - external eval service) can distinguish finished exports from partial ones. + 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 (used by the - on-demand eval snapshot path, where the caller wants to skip that eval). + 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)) @@ -948,9 +937,7 @@ def save_hf_model( logger.info(f"Saving model in HuggingFace format to {path}") if args.megatron_to_hf_mode == "raw" and not is_lora_model(model): - # miles' own converters cover exactly what weight sync covers; the - # bridge silently exports zero weights for specs it has no mapping for - # (e.g. qwen3.5). LoRA keeps the bridge (adapter merging). + # LoRA keeps the bridge (adapter merging). from .update_weight.common import named_params_and_buffers hf_config = load_hf_config(args.hf_checkpoint) @@ -972,12 +959,10 @@ def save_hf_model( # adapter weights into base weights for a standalone HF model. bridge.save_hf_pretrained(model, path=path) - # save_hf_pretrained is collective; make sure every rank is done writing - # before checking the result and claiming the snapshot complete. torch.distributed.barrier() if torch.distributed.get_rank() == 0: - # The bridge silently exports zero weights for specs it has no mapping - # for; a marked-but-weightless snapshot must never exist. + # The bridge exports zero weights for specs it has no mapping for; + # a marked-but-weightless snapshot must never exist. 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 " diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index f1a021ea9e..98453c8bcb 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -43,9 +43,7 @@ logger = logging.getLogger(__name__) -# Bounds each eval-fleet weight load (per attempt). Generous enough for large -# checkpoints from network storage; small enough that a zombie engine costs a -# skipped point, not a wedged eval pipeline. +# Per-attempt bound on eval-fleet weight loads; a zombie engine costs a skipped point. EVAL_WEIGHT_LOAD_TIMEOUT_SECS = 600.0 @@ -170,15 +168,12 @@ async def eval(self, rollout_id, hf_dir: str | None = None, export_time_seconds: 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's weights to the snapshot for ``rollout_id``, then run eval. + """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`` with a - reason counter — never wrong data, never a training stall. + Every failure mode degrades to a skipped point logged at ``rollout_id``. """ start_time = time.time() - # Serializes load -> verify -> generate per fleet: no snapshot load can start - # while an eval is generating, and nothing dispatches before the version is - # confirmed on every engine. This is the version-pinning enforcement. + # The lock holds across load -> verify -> generate; this is the pinning enforcement. async with self._eval_lock: srv = self.servers["eval"] try: @@ -199,10 +194,8 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti weight_version = str(rollout_id) for _attempt in range(2): try: - # The timeout bounds the zombie-engine case: an actor whose sglang - # backend died accepts the call but never answers (the shared http - # client has no timeout), which would otherwise hold the eval lock - # forever and eventually stall the driver via backpressure. + # A zombie engine (backend dead, actor alive) accepts the call and + # never answers; unbounded, it would hold the eval lock forever. await asyncio.wait_for( asyncio.gather( *[ @@ -252,12 +245,8 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti self._gc_eval_snapshots(hf_dir) async def _wait_eval_router_ready(self, srv, timeout: float = 180.0) -> None: - """Probe end-to-end generation through the eval router before dispatching. - - After an engine revival the router needs a health-check cycle to evict the - dead worker and pick up the new one; a generate dispatched inside that window - gets 503s. A one-token probe (retried) proves the route is actually usable. - """ + """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.""" import httpx url = f"http://{srv.router_ip}:{srv.router_port}/generate" @@ -277,12 +266,8 @@ async def _wait_eval_router_ready(self, srv, timeout: float = 180.0) -> None: await asyncio.sleep(5) async def _mark_unreachable_eval_engines(self, srv) -> None: - """Probe eval engines and mark unreachable ones stopped so recover() revives them. - - Without fault tolerance nothing records an engine death (recover() only - restarts engines already marked stopped), so a dead or zombie eval engine - would otherwise fail every future eval instead of self-healing. - """ + """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: @@ -298,19 +283,14 @@ async def _mark_unreachable_eval_engines(self, srv) -> None: engine.mark_stopped() def report_eval_skip(self, rollout_id: int, reason: str) -> None: - """Log a skipped eval point at ``rollout_id`` (called by the driver, e.g. on export failure).""" self._log_eval_skip(rollout_id, reason) def _log_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 staging snapshots beyond the keep ring. - - Only dirs under --eval-hf-dir are ever deleted; --save-hf checkpoints and the - base checkpoint are never touched. Pending evals reference unconsumed dirs, - which are never GC candidates. - """ + """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 diff --git a/miles/rollout/checkpoint_eval.py b/miles/rollout/checkpoint_eval.py index 9d3bf1d06b..bacda96dd9 100644 --- a/miles/rollout/checkpoint_eval.py +++ b/miles/rollout/checkpoint_eval.py @@ -1,14 +1,7 @@ -"""Eval against a dedicated eval fleet pinned to checkpoint snapshots. +"""Eval against a dedicated eval fleet pinned to HF checkpoint snapshots. -The eval fleet is a second model entry (``name="eval"``, ``update_weights=False``) -behind its own router; it never joins the training weight-update group and is never -paused or aborted by training. Weights reach it exclusively through -``update_weights_from_disk`` on an HF snapshot exported for a specific rollout_id, -so every eval point measures one well-defined weight version. - -The helpers here retarget an args namespace at the eval fleet so the existing -generate machinery (which reads ``args.sglang_router_ip/port`` and sizes its -concurrency semaphore off ``args.rollout_num_gpus``) works against it unchanged. +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 @@ -24,9 +17,8 @@ def retarget_args(args: Namespace, router_ip, router_port, num_gpus: int, num_gp """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 is all that is needed to run - the standard eval path against a different set of engines. The original ``args`` - is not modified. + 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 @@ -37,7 +29,6 @@ def retarget_args(args: Namespace, router_ip, router_port, num_gpus: int, num_gp def make_eval_args(args: Namespace) -> Namespace: - """Retarget ``args`` at the in-job eval fleet's router (requires servers started).""" 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) @@ -47,11 +38,7 @@ def make_eval_generate_state(args: Namespace) -> GenerateState: class CheckpointEvalRolloutFn: - """Eval-only rollout function running against the dedicated eval fleet. - - Addressable via ``--eval-function-path`` when the train rollout function should - not serve eval itself. - """ + """Eval-only rollout function for the dedicated eval fleet (via --eval-function-path).""" def __init__(self, input: RolloutFnConstructorInput): self.args = input.args diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py index 0cde09dcc7..d3f2d5fdfb 100644 --- a/miles/rollout/fully_async_rollout.py +++ b/miles/rollout/fully_async_rollout.py @@ -10,9 +10,8 @@ --rollout-function-path miles.rollout.fully_async_rollout.FullyAsyncRolloutFn -Evaluation requires a dedicated eval fleet (``--eval-num-gpus``): the rollout fleet -never has a quiet window, so eval runs on separate engines pinned per-eval to an HF -checkpoint snapshot (see ``miles/rollout/checkpoint_eval.py``). +Evaluation requires a dedicated eval fleet (``--eval-num-gpus``); see +``miles/rollout/checkpoint_eval.py``. """ import asyncio @@ -118,8 +117,6 @@ async def _call_eval(self, input: RolloutFnInput) -> RolloutFnOutput: "fully-async eval requires a dedicated eval fleet: set --eval-num-gpus > 0 " "(or run tools/checkpoint_eval_service.py against --save-hf checkpoints)" ) - # The eval fleet has its own router and engines, so eval coroutines coexist - # with the producer task on the shared loop without contending for capacity. if self._eval_state is None: self._eval_state = make_eval_generate_state(self.args) results = await run_eval_datasets(self._eval_state, self._eval_prompt_dataset_cache) diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py index cb0f306289..b50922875a 100644 --- a/miles/rollout/inference_rollout/inference_rollout_common.py +++ b/miles/rollout/inference_rollout/inference_rollout_common.py @@ -195,8 +195,7 @@ async def _call_eval(self, input: RolloutFnEvalInput) -> RolloutFnEvalOutput: from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets if getattr(self.state.args, "eval_num_gpus", 0) > 0: - # A dedicated eval fleet exists: run eval against its router with a - # state sized for it, instead of contending with train generation. + # Run against the dedicated eval fleet instead of the train engines. if self._eval_state is None: self._eval_state = make_eval_generate_state(self.state.args) state = self._eval_state diff --git a/miles/rollout/inference_rollout/inference_rollout_eval.py b/miles/rollout/inference_rollout/inference_rollout_eval.py index 09f30ed197..26ff4b6350 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -22,7 +22,6 @@ async def run_eval_datasets( state: GenerateState, prompt_dataset_cache: dict[Any, Dataset], ) -> dict[str, dict[str, Any]]: - """Run every configured eval dataset against the engines behind ``state.args``'s router.""" args = state.args assert not args.group_rm, "Group RM is not supported for eval rollout" @@ -103,8 +102,7 @@ async def eval_rollout_single_dataset( try: sample = await future except Exception as e: - # One failed request (engine crash mid-eval, transient router error) - # costs one sample, not the whole eval point. + # One failed request costs one sample, not the whole eval point. logger.warning(f"Eval {dataset_cfg.name}: sample generation raised {e!r}") num_raised += 1 pbar.update(1) @@ -131,8 +129,7 @@ async def eval_rollout_single_dataset( data.sort(key=lambda sample: sample.index) - # A sample can come back ABORTED or reward-less if an engine died mid-eval; - # drop it and report the count instead of corrupting the dataset metrics. + # Drop ABORTED/reward-less samples (engine died mid-eval) and report the count. 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: diff --git a/miles/utils/hf_config.py b/miles/utils/hf_config.py index 3bcb49af35..b730859a13 100644 --- a/miles/utils/hf_config.py +++ b/miles/utils/hf_config.py @@ -110,19 +110,16 @@ def is_dsa(hf_config) -> bool: return getattr(hf_config, "model_type", None) in ("deepseek_v32", "glm_moe_dsa") -# Marker written by HF checkpoint exports (save_hf_model / export_hf) once every rank -# has finished writing, so consumers (the eval snapshot verification, the external -# checkpoint eval service) can distinguish finished exports from partial ones. +# 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: - """Whether ``path`` holds a finished HF export (the completeness marker exists).""" return (Path(path) / HF_EXPORT_COMPLETE_MARKER).exists() def looks_like_hf_checkpoint(path: str | Path) -> bool: - """Marker-less fallback heuristic for checkpoints written before the marker existed.""" + """Fallback for checkpoints written before the marker existed.""" path = Path(path) if not (path / "config.json").exists(): return False diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index 4ea2538d53..489444b3ee 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -66,8 +66,6 @@ async def fake_single_dataset(state, cfg, cache): class FakeRemoteMethod: - """Mimics a Ray actor method: .remote(...) returns an awaitable.""" - def __init__(self, engine, name): self.engine = engine self.name = name @@ -249,9 +247,8 @@ async def test_controller_pin_violation_skips_after_retry(controller_env, tmp_pa async def test_controller_zombie_engine_times_out_to_skip(controller_env, tmp_path, monkeypatch): - """An engine whose backend died accepts the call but never answers; the load - timeout must convert that into a skipped point instead of holding the eval - lock forever (which would eventually stall the driver via backpressure).""" + """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() @@ -284,9 +281,7 @@ def remote(self, *args, **kwargs): async def test_controller_marks_dead_engine_for_recovery(controller_env, tmp_path): - """A dead engine actor must be marked stopped so recover() revives it; the - eval itself degrades to a skipped point (the fake recover cannot really - replace the actor), never a raise out of the controller.""" + """A dead actor is marked stopped for revival; the eval degrades to a skip.""" snapshot = tmp_path / "step_5" snapshot.mkdir() (snapshot / ".complete").touch() diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index 736bb4bbf1..f4e2851489 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -1,26 +1,12 @@ -"""Standalone checkpoint eval service. - -Watches a directory of HF checkpoint exports (``--save-hf`` output), and for each new -complete snapshot: pins an sglang server to it via ``/update_weights_from_disk`` -(stamping ``weight_version=str(rollout_id)``), runs the standard miles eval datasets -against it, and logs the metrics at the snapshot's ``rollout_id`` — hours-late points -land at the right x because the step travels inside the payload (``eval/step``). - -Runs anywhere with sglang + miles importable — another node, spot capacity, or after -the training run. No Ray, no training-job membership. A ledger file next to the watch -dir makes restarts idempotent and backfills missed snapshots. +"""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 \\ - --wandb-mode separate - -The watch dir is scanned for subdirectories whose name contains the rollout id -(``step_{id}`` or any ``{rollout_id}``-formatted layout); a snapshot is consumed once -it has the ``.complete`` marker (or, for checkpoints written before the marker -existed, once it looks like a finished HF export and has been quiescent). + --eval-prompt-data aime /data/aime-2024.jsonl --rm-type dapo --reward-key score """ import argparse @@ -42,9 +28,7 @@ QUIESCENCE_SECS = 120.0 -# Fields the eval path consumes that are not exposed as service flags; values match -# the miles argument defaults. tests/fast/rollout/test_checkpoint_eval_service.py pins -# this table against the eval path's consumption. +# 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, @@ -141,7 +125,6 @@ def parse_service_args() -> Namespace: def build_eval_namespace(service_args: Namespace, server_ip: str, server_port: int) -> Namespace: - """Compose the namespace the miles eval path consumes.""" from miles.utils.arguments import _resolve_eval_datasets args = Namespace(**EVAL_ARG_DEFAULTS) @@ -152,8 +135,6 @@ def build_eval_namespace(service_args: Namespace, server_ip: str, server_port: i class SnapshotLedger: - """Persistent record of consumed snapshots, next to the watch dir.""" - def __init__(self, watch_dir: Path): self.path = watch_dir / ".eval_service_state.json" self.consumed: set[int] = set() @@ -166,7 +147,6 @@ def mark(self, rollout_id: int) -> None: def find_ready_snapshots(watch_dir: Path, min_rollout_id: int, consumed: set[int]) -> list[tuple[int, Path]]: - """Unconsumed snapshot dirs ready for eval, ordered by rollout_id.""" ready = [] for child in watch_dir.iterdir(): if not child.is_dir(): diff --git a/train_async.py b/train_async.py index 8c4c2838b7..1df3c9c51a 100644 --- a/train_async.py +++ b/train_async.py @@ -21,14 +21,8 @@ class EvalDispatcher: - """Dispatch evals against the dedicated eval fleet without stalling training. - - With ``--eval-num-gpus 0`` this degrades to today's blocking ``eval.remote`` call. - Otherwise each due eval gets an HF snapshot (exported to ``--eval-hf-dir``, or the - ``--save-hf`` checkpoint in reuse mode) and is fired fire-and-forget; the pending - set is bounded by ``--eval-max-in-flight`` via the ``--eval-overflow-policy``. - Failures degrade to a skipped point logged at that rollout_id, never a crash. - """ + """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 @@ -65,15 +59,13 @@ async def dispatch(self, rollout_id: int, hf_dir: str | None = None) -> None: self.pending.append((rollout_id, ref)) async def drain(self) -> None: - """Await every pending eval (before dispose, so the last points always land).""" 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: - # Reuse mode: the --save-hf checkpoint for this step was just written - # (eval_interval is validated to be a multiple of save_interval). + # Reuse mode: the --save-hf checkpoint for this step was just written. 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() From e04500015612341c2efc4ec90fa831bf02786aa2 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 11:08:16 -0700 Subject: [PATCH 19/33] eval: move EvalDispatcher into miles/ray/rollout train_async.py goes back to being a thin procedural driver; the dispatcher is Ray orchestration and lives with RolloutManager. --- miles/ray/rollout/eval_dispatch.py | 79 ++++++++++++++++++++++ tests/fast/rollout/test_checkpoint_eval.py | 14 ++-- train_async.py | 77 +-------------------- 3 files changed, 88 insertions(+), 82 deletions(-) create mode 100644 miles/ray/rollout/eval_dispatch.py diff --git a/miles/ray/rollout/eval_dispatch.py b/miles/ray/rollout/eval_dispatch.py new file mode 100644 index 0000000000..883b08e640 --- /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) -> 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": + 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: + # Reuse mode: the --save-hf checkpoint for this step was just written. + 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: + # The manager's eval lock serializes evals, so pending refs finish in order. + 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") + + 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") diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index 489444b3ee..6333efef7d 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -406,15 +406,17 @@ async def export_hf(self, rollout_id, path): @pytest.fixture def dispatcher_env(monkeypatch): - import train_async + import miles.ray.rollout.eval_dispatch as eval_dispatch # ray.wait/ray.get over asyncio futures: done iff the future is resolved. - monkeypatch.setattr(train_async.ray, "wait", lambda refs, timeout=0: (refs, []) if refs[0].done() else ([], refs)) - monkeypatch.setattr(train_async.ray, "get", lambda ref: ref.result()) - return train_async + 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(train_async, manager, actor_model, **arg_overrides): +def make_dispatcher(eval_dispatch, manager, actor_model, **arg_overrides): dispatcher_defaults = dict( eval_hf_dir="/dev/shm/eval_hf", eval_dispatch="async", @@ -423,7 +425,7 @@ def make_dispatcher(train_async, manager, actor_model, **arg_overrides): ) dispatcher_defaults.update(arg_overrides) args = make_args(**dispatcher_defaults) - return train_async.EvalDispatcher(args, actor_model, manager), args + return eval_dispatch.EvalDispatcher(args, actor_model, manager), args async def test_dispatcher_exports_and_fires(dispatcher_env): diff --git a/train_async.py b/train_async.py index 1df3c9c51a..894db129c3 100644 --- a/train_async.py +++ b/train_async.py @@ -1,12 +1,8 @@ import asyncio import logging -import os -import time -from collections import deque - -import ray 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 @@ -20,77 +16,6 @@ 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) -> 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": - 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: - # Reuse mode: the --save-hf checkpoint for this step was just written. - 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: - # The manager's eval lock serializes evals, so pending refs finish in order. - 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") - - 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") - - # The framework supports other asynchronous approaches such as fully async (see miles/rollout/fully_async_rollout.py). async def train(args): assert not args.colocate, "Colocation is not supported for async training." From e499cfc6d7113121333f69f8c748e45ec15e92bd Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 11:14:29 -0700 Subject: [PATCH 20/33] fully-async: shared-engine eval via producer pause when no fleet is set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With --eval-num-gpus unset, eval no longer raises: the producer pauses new submissions for the duration of the eval and resumes after. The weight version stays pinned without any snapshot machinery because the driver awaits eval, so no update_weights interleaves — in-flight groups finish and buffer, eval takes over the freed capacity, and rollout production resumes when the point lands. Costs roughly the eval duration in rollout stall; the dedicated fleet remains the zero-stall option. --- docs/user-guide/fully-async.md | 11 +++-- examples/fully_async/README.md | 5 ++- miles/rollout/fully_async_rollout.py | 33 +++++++++----- .../fast/rollout/test_fully_async_rollout.py | 45 ++++++++++++++++--- 4 files changed, 74 insertions(+), 20 deletions(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 56f1d7d64a..8b2f3c710b 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -102,9 +102,14 @@ rollout engines closer to the latest actor weights so fewer groups get recycled ## Evaluation -The rollout fleet never has a quiet window, so fully-async eval runs on a **dedicated -eval fleet** synced through HF checkpoint snapshots — never by joining training weight -updates. Enable it with: +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. The weight version stays pinned because no update interleaves while the driver +awaits 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) diff --git a/examples/fully_async/README.md b/examples/fully_async/README.md index 7181ae2725..bf58fb5f50 100644 --- a/examples/fully_async/README.md +++ b/examples/fully_async/README.md @@ -28,8 +28,9 @@ Started fully-async rollout worker * Aborted or too-stale groups are recycled back into the data source. ## Evaluation -Eval runs on a dedicated eval fleet synced via HF checkpoint snapshots (`--eval-num-gpus`, -`--eval-hf-dir`); see `run-qwen3.5-4b-fully_async-eval.sh` and the fully-async docs page. +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 * Ordering is best effort (sorted at the end by index). diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py index d3f2d5fdfb..5a46a9ba1b 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 requires a dedicated eval fleet (``--eval-num-gpus``); see -``miles/rollout/checkpoint_eval.py``. +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 @@ -101,6 +102,8 @@ def __init__(self, input: RolloutFnConstructorInput): self._output: asyncio.Queue[Group] | None = None self._eval_state: GenerateState | None = None 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: @@ -112,14 +115,23 @@ async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: return await self._drain(input.rollout_id) async def _call_eval(self, input: RolloutFnInput) -> RolloutFnOutput: - if getattr(self.args, "eval_num_gpus", 0) <= 0: - raise ValueError( - "fully-async eval requires a dedicated eval fleet: set --eval-num-gpus > 0 " - "(or run tools/checkpoint_eval_service.py against --save-hf checkpoints)" - ) - if self._eval_state is None: - self._eval_state = make_eval_generate_state(self.args) - results = await run_eval_datasets(self._eval_state, self._eval_prompt_dataset_cache) + if getattr(self.args, "eval_num_gpus", 0) > 0: + if self._eval_state is None: + self._eval_state = make_eval_generate_state(self.args) + results = await run_eval_datasets(self._eval_state, self._eval_prompt_dataset_cache) + return RolloutFnEvalOutput(data=results) + + # No dedicated fleet: eval shares the rollout engines. The producer pauses + # new submissions while eval runs (in-flight groups finish and buffer); the + # driver awaits eval, so no weight update interleaves and the weight version + # stays pinned for both eval and the buffered train samples. + 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 -------------------------- @@ -141,6 +153,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/tests/fast/rollout/test_fully_async_rollout.py b/tests/fast/rollout/test_fully_async_rollout.py index e39fe666e7..9f30388d17 100644 --- a/tests/fast/rollout/test_fully_async_rollout.py +++ b/tests/fast/rollout/test_fully_async_rollout.py @@ -105,11 +105,46 @@ async def test_drain_collects_batch_sorted_with_metrics(monkeypatch): assert len(output2.samples) == 3 -async def test_eval_without_fleet_raises(monkeypatch): - fn = make_fn(monkeypatch, make_args(), FakeDataSource()) - with pytest.raises(ValueError, match="requires a dedicated eval fleet"): - await fn(RolloutFnEvalInput(rollout_id=0)) - assert fn._worker is None +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): From 460e0af8b4584b4198916ae7eed9faef11d9fdb9 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 13:46:22 -0700 Subject: [PATCH 21/33] docs: shared-engine eval version semantics --- docs/user-guide/fully-async.md | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 8b2f3c710b..3cda90b022 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -105,8 +105,10 @@ rollout engines closer to the latest actor weights so fewer groups get recycled 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. The weight version stays pinned because no update interleaves while the driver -awaits eval — the cost is that rollout production stalls for roughly the eval duration, -which is fine for small debug eval sets. +awaits eval — 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). 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: From 9d39aa8af92fbf0d6ee86be29981fabf0eacdd54 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 17:18:46 -0700 Subject: [PATCH 22/33] eval: apply reward_key only to dict rewards A per-dataset rm override (sample metadata rm_type) may return scalar rewards while the training rm returns dicts keyed by --reward-key. --- .../inference_rollout_eval.py | 7 +++- tests/fast/rollout/test_checkpoint_eval.py | 41 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/miles/rollout/inference_rollout/inference_rollout_eval.py b/miles/rollout/inference_rollout/inference_rollout_eval.py index 26ff4b6350..7b5f1ee4a3 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -139,7 +139,12 @@ async def eval_rollout_single_dataset( 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], + # reward_key only applies to dict rewards; a per-dataset rm override + # (sample metadata rm_type) may return plain scalars. + "rewards": [ + sample.reward[reward_key] if reward_key and isinstance(sample.reward, dict) else sample.reward + for sample in data + ], "truncated": [sample.status == Sample.Status.TRUNCATED for sample in data], "samples": data, "failed_samples": num_failed, diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index 6333efef7d..73c6b37f3b 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -62,6 +62,47 @@ async def fake_single_dataset(state, cfg, cache): assert set(results.keys()) == {"a", "b"} +async def test_single_dataset_reward_key_only_applies_to_dict_rewards(monkeypatch): + """A per-dataset rm override (sample metadata rm_type) may return scalar rewards + even when --reward-key is set for the training rm's dict rewards.""" + import miles.rollout.inference_rollout.inference_rollout_eval as eval_mod + from miles.utils.types import Sample + + def make_sample(index, reward): + return Sample(index=index, prompt="p", response="r", label="l", reward=reward) + + async def fake_generate_and_rm(state, sample, sampling_params, evaluation): + return sample + + monkeypatch.setattr(eval_mod, "generate_and_rm", fake_generate_and_rm) + 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="score", + eval_reward_key=None, + ) + dataset_cfg = SimpleNamespace( + name="mixed", + cache_key=("mixed",), + 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, + ) + scalar, dict_reward = make_sample(0, 1), make_sample(1, {"score": 0.5}) + cache = {dataset_cfg.cache_key + ("hf", False, None): SimpleNamespace(samples=[scalar, dict_reward])} + + result = await eval_mod.eval_rollout_single_dataset(SimpleNamespace(args=args), dataset_cfg, cache) + + assert result["mixed"]["rewards"] == [1, 0.5] + + # ---------------- controller (RolloutManager._eval_on_dedicated_fleet) ---------------- From 4e2092d3408c53224b4bd4fbf3a955f8e6cae129 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 17:28:00 -0700 Subject: [PATCH 23/33] types: get_reward_value tolerates scalar rewards when reward_key is set --- miles/utils/types.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/miles/utils/types.py b/miles/utils/types.py index cd7637d4c4..5931d33f88 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -152,7 +152,11 @@ def from_dict(data: dict): return sample def get_reward_value(self, args) -> float: - return self.reward if not args.reward_key else self.reward[args.reward_key] + # reward_key only applies to dict rewards; a per-dataset rm override + # (sample metadata rm_type) may return plain scalars. + if args.reward_key and isinstance(self.reward, dict): + return self.reward[args.reward_key] + return self.reward @property def effective_response_length(self): From f49a6e85b13f7dd5a5de4470884ed406e3a281f4 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 18:23:25 -0700 Subject: [PATCH 24/33] docs: present eval postures as run-type x standalone-resources matrix --- docs/user-guide/fully-async.md | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 3cda90b022..3db5208fa5 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -102,6 +102,14 @@ rollout engines closer to the latest actor weights so fewer groups get recycled ## 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 | +| **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. The weight version stays pinned because no update interleaves while the driver From 774f2acf7b7abb3477f38f5177b3c014600e23b7 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 18:24:14 -0700 Subject: [PATCH 25/33] docs: note pause-the-world as the strictly-on-time option for real runs --- docs/user-guide/fully-async.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 3db5208fa5..77b1ed0b77 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -107,7 +107,7 @@ 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 | +| **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) | | **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**: From e4484444e92920e7d1e56f1fac4e1dffb0fd0b8d Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 21:10:46 -0700 Subject: [PATCH 26/33] eval: never skip the final eval point At the last rollout training is already over, so overflow backpressure costs nothing; dropping the most valuable point under skip policy was pure loss. --- miles/ray/rollout/eval_dispatch.py | 4 ++-- tests/fast/rollout/test_checkpoint_eval.py | 22 ++++++++++++++++++++++ train_async.py | 4 +++- 3 files changed, 27 insertions(+), 3 deletions(-) diff --git a/miles/ray/rollout/eval_dispatch.py b/miles/ray/rollout/eval_dispatch.py index 883b08e640..80de3cb660 100644 --- a/miles/ray/rollout/eval_dispatch.py +++ b/miles/ray/rollout/eval_dispatch.py @@ -18,14 +18,14 @@ def __init__(self, args, actor_model, rollout_manager): 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) -> None: + 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": + 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() diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index 73c6b37f3b..a277e6c854 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -508,6 +508,28 @@ async def test_dispatcher_skip_policy_drops_before_export(dispatcher_env): 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) diff --git a/train_async.py b/train_async.py index 894db129c3..e1689cb72d 100644 --- a/train_async.py +++ b/train_async.py @@ -97,7 +97,9 @@ async def train(args): await actor_model.update_weights(rollout_id=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) + # The final point is never skipped: training is over, so overflow + # backpressure no longer costs any training time. + await eval_dispatcher.dispatch(rollout_id, force=rollout_id == args.num_rollout - 1) if ( args.debug_exit_after_rollout is not None From 3fbe28b999612012eaad86051356a44f45827d73 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 21:50:59 -0700 Subject: [PATCH 27/33] eval: code-quality pass - fold the triplicated lazy-eval-state block into EvalFleetSession; drop the unused CheckpointEvalRolloutFn - report crashed async evals as eval/skipped_crashed so every missing point stays attributable - anchor snapshot rollout-id extraction to the trailing number - drop the service's dead --use-wandb flag; make resolve_eval_datasets public (cross-module use) - match HF export metadata skips by suffix, not substring - keyword-only path/raise_on_error on save_hf_model - share one wait_http_ok helper for router/server readiness probes - count raised samples in failed_samples; failure-path tests for eval, snapshot discovery, ledger, and force dispatch --- miles/backends/megatron_utils/actor.py | 3 - miles/backends/megatron_utils/model.py | 13 ++- miles/ray/rollout/eval_dispatch.py | 2 + miles/ray/rollout/rollout_manager.py | 37 ++---- miles/ray/rollout/rollout_server.py | 2 +- miles/rollout/checkpoint_eval.py | 23 ++-- miles/rollout/fully_async_rollout.py | 11 +- .../inference_rollout_common.py | 16 +-- .../inference_rollout_eval.py | 9 +- miles/utils/arguments.py | 4 +- miles/utils/http_utils.py | 24 +++- miles/utils/types.py | 6 +- tests/fast/rollout/test_checkpoint_eval.py | 105 +++++++++++++++--- .../fast/rollout/test_fully_async_rollout.py | 12 +- tools/checkpoint_eval_service.py | 26 ++--- 15 files changed, 177 insertions(+), 116 deletions(-) diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 327f102bb7..a91cd9eb2b 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -574,9 +574,6 @@ def export_hf(self, rollout_id: int, path: str) -> None: that failed to export can be skipped loudly. """ self._heartbeat.bump() - if self.args.debug_rollout_only: - return - 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) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 21fac08214..1a55b586fa 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -826,7 +826,13 @@ def save( enable_forward_pre_hook(model) -HF_METADATA_SKIP_SUFFIXES = (".safetensors", ".bin", ".pt", ".pth", ".gguf") +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( @@ -878,10 +884,8 @@ def export_hf_model_direct( assert weight_map, f"HF export to {path} produced no weights" base_checkpoint = Path(args.hf_checkpoint) if base_checkpoint.is_dir(): - # Tokenizer/config metadata only — the skip list also excludes the base - # checkpoint's safetensors index, which would clobber ours below. for meta_file in base_checkpoint.iterdir(): - if meta_file.is_file() and not any(s in meta_file.name for s in HF_METADATA_SKIP_SUFFIXES): + 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}") @@ -908,6 +912,7 @@ def save_hf_model( args, rollout_id: int, model: Sequence[DDP], + *, path: str | Path | None = None, raise_on_error: bool = False, ) -> None: diff --git a/miles/ray/rollout/eval_dispatch.py b/miles/ray/rollout/eval_dispatch.py index 80de3cb660..98ec6959df 100644 --- a/miles/ray/rollout/eval_dispatch.py +++ b/miles/ray/rollout/eval_dispatch.py @@ -71,9 +71,11 @@ def _reap_finished(self) -> None: 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/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 98453c8bcb..5300f44d0d 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -30,7 +30,7 @@ from miles.utils.environ import enable_experimental_rollout_refactor from miles.utils.health_monitor import RolloutHealthMonitor from miles.utils.hf_config import is_complete_hf_export -from miles.utils.http_utils import init_http_client +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 @@ -149,7 +149,7 @@ async def eval(self, rollout_id, hf_dir: str | None = None, export_time_seconds: return self._health_monitoring_resume() - if getattr(self.args, "eval_num_gpus", 0) > 0: + 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) @@ -182,12 +182,12 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti await srv.wait_all_engines_alive() except Exception as e: logger.warning(f"Eval fleet unhealthy at rollout {rollout_id}, skipping eval: {e}") - self._log_eval_skip(rollout_id, "unhealthy") + 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._log_eval_skip(rollout_id, "ckpt_missing") + self.report_eval_skip(rollout_id, "ckpt_missing") return engines = [e.actor_handle for e in srv.engines] @@ -218,14 +218,14 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti logger.warning( f"Eval fleet failed to pin weight_version={weight_version} (got {versions}), skipping eval" ) - self._log_eval_skip(rollout_id, "pin_violation") + 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._log_eval_skip(rollout_id, "unhealthy") + self.report_eval_skip(rollout_id, "unhealthy") return result = await asyncio.to_thread( @@ -247,23 +247,11 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti 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.""" - import httpx - - url = f"http://{srv.router_ip}:{srv.router_port}/generate" - payload = {"input_ids": [0], "sampling_params": {"max_new_tokens": 1, "temperature": 0}} - deadline = time.time() + timeout - async with httpx.AsyncClient() as client: - while True: - try: - response = await client.post(url, json=payload, timeout=60) - 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"eval router at {url} not ready after {timeout}s: {last_error}") - await asyncio.sleep(5) + 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 @@ -283,9 +271,6 @@ async def _mark_unreachable_eval_engines(self, srv) -> None: engine.mark_stopped() def report_eval_skip(self, rollout_id: int, reason: str) -> None: - self._log_eval_skip(rollout_id, reason) - - def _log_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: diff --git a/miles/ray/rollout/rollout_server.py b/miles/ray/rollout/rollout_server.py index 07412dcf46..79f4d51b7a 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/miles/ray/rollout/rollout_server.py @@ -104,7 +104,7 @@ 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 = getattr(args, "eval_num_gpus", 0) + eval_num_gpus = args.eval_num_gpus if getattr(args, "sglang_config", None) is not None: config = SglangConfig.from_yaml(args.sglang_config) diff --git a/miles/rollout/checkpoint_eval.py b/miles/rollout/checkpoint_eval.py index bacda96dd9..ea1fe42c75 100644 --- a/miles/rollout/checkpoint_eval.py +++ b/miles/rollout/checkpoint_eval.py @@ -7,10 +7,9 @@ import copy from argparse import Namespace -from miles.rollout.base_types import RolloutFnConstructorInput, RolloutFnEvalInput, RolloutFnEvalOutput from miles.rollout.inference_rollout.inference_rollout_common import GenerateState -__all__ = ["retarget_args", "make_eval_args", "make_eval_generate_state", "CheckpointEvalRolloutFn"] +__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: @@ -37,19 +36,21 @@ def make_eval_generate_state(args: Namespace) -> GenerateState: return GenerateState(make_eval_args(args)) -class CheckpointEvalRolloutFn: - """Eval-only rollout function for the dedicated eval fleet (via --eval-function-path).""" +class EvalFleetSession: + """Eval-fleet ``GenerateState`` + prompt cache, built lazily on first use. - def __init__(self, input: RolloutFnConstructorInput): - self.args = input.args + 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 = {} + self._prompt_dataset_cache: dict = {} - async def __call__(self, input: RolloutFnEvalInput) -> RolloutFnEvalOutput: + async def run(self) -> dict: from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets - assert input.evaluation, "CheckpointEvalRolloutFn only serves eval" if self._state is None: self._state = make_eval_generate_state(self.args) - results = await run_eval_datasets(self._state, self._prompt_dataset_cache) - return RolloutFnEvalOutput(data=results) + 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 5a46a9ba1b..7c691185b0 100644 --- a/miles/rollout/fully_async_rollout.py +++ b/miles/rollout/fully_async_rollout.py @@ -29,7 +29,7 @@ RolloutFnOutput, RolloutFnTrainOutput, ) -from miles.rollout.checkpoint_eval import make_eval_generate_state +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 @@ -100,7 +100,7 @@ def __init__(self, input: RolloutFnConstructorInput): self._weight_version = _CachedWeightVersion() self._worker: asyncio.Task | None = None self._output: asyncio.Queue[Group] | None = None - self._eval_state: GenerateState | None = None + self._eval_fleet = EvalFleetSession(input.args) self._eval_prompt_dataset_cache: dict = {} self._producer_resumed = asyncio.Event() self._producer_resumed.set() @@ -115,11 +115,8 @@ async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: return await self._drain(input.rollout_id) async def _call_eval(self, input: RolloutFnInput) -> RolloutFnOutput: - if getattr(self.args, "eval_num_gpus", 0) > 0: - if self._eval_state is None: - self._eval_state = make_eval_generate_state(self.args) - results = await run_eval_datasets(self._eval_state, self._eval_prompt_dataset_cache) - return RolloutFnEvalOutput(data=results) + if self.args.eval_num_gpus > 0: + return RolloutFnEvalOutput(data=await self._eval_fleet.run()) # No dedicated fleet: eval shares the rollout engines. The producer pauses # new submissions while eval runs (in-flight groups finish and buffer); the diff --git a/miles/rollout/inference_rollout/inference_rollout_common.py b/miles/rollout/inference_rollout/inference_rollout_common.py index b50922875a..7a02922c32 100644 --- a/miles/rollout/inference_rollout/inference_rollout_common.py +++ b/miles/rollout/inference_rollout/inference_rollout_common.py @@ -171,10 +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_state: GenerateState | None = None + self._eval_fleet = EvalFleetSession(input.args) async def __call__(self, input: RolloutFnInput) -> RolloutFnOutput: if input.evaluation: @@ -191,15 +193,9 @@ async def _call_train(self, input: RolloutFnTrainInput) -> RolloutFnTrainOutput: return output async def _call_eval(self, input: RolloutFnEvalInput) -> RolloutFnEvalOutput: - from miles.rollout.checkpoint_eval import make_eval_generate_state from miles.rollout.inference_rollout.inference_rollout_eval import run_eval_datasets - if getattr(self.state.args, "eval_num_gpus", 0) > 0: - # Run against the dedicated eval fleet instead of the train engines. - if self._eval_state is None: - self._eval_state = make_eval_generate_state(self.state.args) - state = self._eval_state - else: - state = self.state - results = await run_eval_datasets(state, self.eval_prompt_dataset_cache) + 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 7b5f1ee4a3..2bc46325de 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -139,14 +139,9 @@ async def eval_rollout_single_dataset( reward_key = args.eval_reward_key or args.reward_key return { dataset_cfg.name: { - # reward_key only applies to dict rewards; a per-dataset rm override - # (sample metadata rm_type) may return plain scalars. - "rewards": [ - sample.reward[reward_key] if reward_key and isinstance(sample.reward, dict) else sample.reward - for sample in data - ], + "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, + "failed_samples": num_failed + num_raised, } } diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index b0203ae3a3..4c7015700c 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2280,7 +2280,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. """ @@ -2338,7 +2338,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)") diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 96a5022aef..3349d09207 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,7 +253,7 @@ def init_http_client(args): return _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine - if getattr(args, "eval_num_gpus", 0) > 0: + if args.eval_num_gpus > 0: # The eval fleet is served by the same client; size for both. _client_concurrency += args.sglang_server_concurrency * args.eval_num_gpus // args.eval_num_gpus_per_engine if _http_client is None: diff --git a/miles/utils/types.py b/miles/utils/types.py index 5931d33f88..cd7637d4c4 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -152,11 +152,7 @@ def from_dict(data: dict): return sample def get_reward_value(self, args) -> float: - # reward_key only applies to dict rewards; a per-dataset rm override - # (sample metadata rm_type) may return plain scalars. - if args.reward_key and isinstance(self.reward, dict): - return self.reward[args.reward_key] - return self.reward + return self.reward if not args.reward_key else self.reward[args.reward_key] @property def effective_response_length(self): diff --git a/tests/fast/rollout/test_checkpoint_eval.py b/tests/fast/rollout/test_checkpoint_eval.py index a277e6c854..019ec98dae 100644 --- a/tests/fast/rollout/test_checkpoint_eval.py +++ b/tests/fast/rollout/test_checkpoint_eval.py @@ -62,32 +62,23 @@ async def fake_single_dataset(state, cfg, cache): assert set(results.keys()) == {"a", "b"} -async def test_single_dataset_reward_key_only_applies_to_dict_rewards(monkeypatch): - """A per-dataset rm override (sample metadata rm_type) may return scalar rewards - even when --reward-key is set for the training rm's dict rewards.""" +def _eval_dataset_env(monkeypatch, generate): import miles.rollout.inference_rollout.inference_rollout_eval as eval_mod from miles.utils.types import Sample - def make_sample(index, reward): - return Sample(index=index, prompt="p", response="r", label="l", reward=reward) - - async def fake_generate_and_rm(state, sample, sampling_params, evaluation): - return sample - - monkeypatch.setattr(eval_mod, "generate_and_rm", fake_generate_and_rm) + 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="score", + reward_key=None, eval_reward_key=None, ) dataset_cfg = SimpleNamespace( - name="mixed", - cache_key=("mixed",), + name="ds", + cache_key=("ds",), n_samples_per_eval_prompt=1, temperature=1.0, top_p=1.0, @@ -95,12 +86,35 @@ async def fake_generate_and_rm(state, sample, sampling_params, evaluation): max_response_len=16, inject_metadata=lambda md: md, ) - scalar, dict_reward = make_sample(0, 1), make_sample(1, {"score": 0.5}) - cache = {dataset_cfg.cache_key + ("hf", False, None): SimpleNamespace(samples=[scalar, dict_reward])} + 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["mixed"]["rewards"] == [1, 0.5] + 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) ---------------- @@ -594,3 +608,60 @@ def remote(self, rollout_id): 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 9f30388d17..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) @@ -114,7 +115,9 @@ async def blocking_generate(state, group, sampling_params, evaluation=False): return group data_source = FakeDataSource() - fn = make_fn(monkeypatch, make_args(rollout_batch_size=2, eval_num_gpus=0), data_source, generate=blocking_generate) + 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() @@ -148,6 +151,9 @@ async def fake_run_eval_datasets(state, cache): 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, @@ -169,8 +175,8 @@ async def fake_run_eval_datasets(state, cache): assert state in seen_states return eval_results - monkeypatch.setattr(fully_async, "make_eval_generate_state", fake_make_eval_generate_state) - monkeypatch.setattr(fully_async, "run_eval_datasets", fake_run_eval_datasets) + 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)) diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index f4e2851489..66b1b844d2 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -22,7 +22,7 @@ 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 +from miles.utils.http_utils import init_http_client, post, wait_http_ok logger = logging.getLogger("checkpoint_eval_service") @@ -63,6 +63,7 @@ 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 ) @@ -115,7 +116,6 @@ def parse_service_args() -> Namespace: 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("--use-wandb", action="store_true", default=False) 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) @@ -125,12 +125,12 @@ def parse_service_args() -> Namespace: def build_eval_namespace(service_args: Namespace, server_ip: str, server_port: int) -> Namespace: - from miles.utils.arguments import _resolve_eval_datasets + 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) + args.eval_datasets = resolve_eval_datasets(args) return retarget_args(args, server_ip, server_port, service_args.num_gpus, service_args.tp) @@ -151,7 +151,7 @@ def find_ready_snapshots(watch_dir: Path, min_rollout_id: int, consumed: set[int for child in watch_dir.iterdir(): if not child.is_dir(): continue - match = re.search(r"(\d+)", child.name) + match = re.search(r"(\d+)$", child.name) if match is None: continue rollout_id = int(match.group(1)) @@ -195,20 +195,8 @@ def launch_server(service_args: Namespace) -> tuple[subprocess.Popen | None, str async def wait_server_healthy(ip: str, port: int, timeout: float = 1800.0) -> None: - import httpx - - deadline = time.time() + timeout - async with httpx.AsyncClient() as client: - while time.time() < deadline: - try: - # /health_generate runs a tiny generation; 200 body is not JSON. - response = await client.get(f"http://{ip}:{port}/health_generate", timeout=60) - if response.status_code == 200: - return - except httpx.HTTPError: - pass - await asyncio.sleep(5) - raise TimeoutError(f"sglang server at {ip}:{port} not healthy after {timeout}s") + # /health_generate runs a tiny generation; its 200 body is not JSON. + await wait_http_ok(f"http://{ip}:{port}/health_generate", timeout=timeout) async def eval_snapshot(args: Namespace, state, cache: dict, rollout_id: int, snapshot: Path) -> None: From 97e7e04d52f520f90a7d44dc2f3a69813d972c00 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 22:11:14 -0700 Subject: [PATCH 28/33] tests: add eval fleet fields to the ray rollout args fixture --- tests/fast/ray/rollout/conftest.py | 2 ++ 1 file changed, 2 insertions(+) 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, From 0abdd375851f41fc19fcb5a923fad6dcfce138b0 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 22:46:56 -0700 Subject: [PATCH 29/33] eval: trim narrative comments --- miles/ray/placement_group.py | 1 - miles/ray/rollout/eval_dispatch.py | 1 - miles/ray/rollout/rollout_server.py | 5 ++--- miles/rollout/fully_async_rollout.py | 6 ++---- miles/rollout/inference_rollout/inference_rollout_eval.py | 2 -- miles/utils/arguments.py | 1 - miles/utils/http_utils.py | 1 - tools/checkpoint_eval_service.py | 1 - train_async.py | 3 --- 9 files changed, 4 insertions(+), 17 deletions(-) diff --git a/miles/ray/placement_group.py b/miles/ray/placement_group.py index 5e24dc5ee4..f1a6067e13 100644 --- a/miles/ray/placement_group.py +++ b/miles/ray/placement_group.py @@ -107,7 +107,6 @@ 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: - # eval_num_gpus adds a dedicated eval fleet at the tail of the rollout slice. 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: diff --git a/miles/ray/rollout/eval_dispatch.py b/miles/ray/rollout/eval_dispatch.py index 98ec6959df..b961ad3752 100644 --- a/miles/ray/rollout/eval_dispatch.py +++ b/miles/ray/rollout/eval_dispatch.py @@ -53,7 +53,6 @@ async def drain(self) -> None: async def _ensure_snapshot(self, rollout_id: int) -> tuple[str, float | None]: if self.args.eval_hf_dir is None: - # Reuse mode: the --save-hf checkpoint for this step was just written. 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() diff --git a/miles/ray/rollout/rollout_server.py b/miles/ray/rollout/rollout_server.py index 79f4d51b7a..bb77d7b34e 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/miles/ray/rollout/rollout_server.py @@ -134,9 +134,8 @@ def _resolve_sglang_config(args) -> SglangConfig: ) if eval_num_gpus > 0: - # Dedicated eval fleet: own router, never receives training weight updates - # (weights are pinned per-eval via update_weights_from_disk). update_weights - # must be explicit — resolve() would infer True from the matching model_path. + # update_weights=False must be explicit — resolve() would infer True from + # the matching model_path. config.models.append( ModelConfig( name="eval", diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py index 7c691185b0..a0593c27a6 100644 --- a/miles/rollout/fully_async_rollout.py +++ b/miles/rollout/fully_async_rollout.py @@ -118,10 +118,8 @@ async def _call_eval(self, input: RolloutFnInput) -> RolloutFnOutput: if self.args.eval_num_gpus > 0: return RolloutFnEvalOutput(data=await self._eval_fleet.run()) - # No dedicated fleet: eval shares the rollout engines. The producer pauses - # new submissions while eval runs (in-flight groups finish and buffer); the - # driver awaits eval, so no weight update interleaves and the weight version - # stays pinned for both eval and the buffered train samples. + # Shared-engine eval: the driver awaits it, so no weight update interleaves + # and the version stays pinned while submissions are paused. logger.info("Pausing fully-async producer submissions for shared-engine eval") self._producer_resumed.clear() try: diff --git a/miles/rollout/inference_rollout/inference_rollout_eval.py b/miles/rollout/inference_rollout/inference_rollout_eval.py index 2bc46325de..82ec439438 100644 --- a/miles/rollout/inference_rollout/inference_rollout_eval.py +++ b/miles/rollout/inference_rollout/inference_rollout_eval.py @@ -102,7 +102,6 @@ async def eval_rollout_single_dataset( try: sample = await future except Exception as e: - # One failed request costs one sample, not the whole eval point. logger.warning(f"Eval {dataset_cfg.name}: sample generation raised {e!r}") num_raised += 1 pbar.update(1) @@ -129,7 +128,6 @@ async def eval_rollout_single_dataset( data.sort(key=lambda sample: sample.index) - # Drop ABORTED/reward-less samples (engine died mid-eval) and report the count. 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: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index 4c7015700c..c8650c9d79 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -2551,7 +2551,6 @@ def miles_validate_args(args): f"({args.eval_max_in_flight}), otherwise a pending eval's snapshot could be GC'd." ) if args.eval_hf_dir is None: - # Reuse mode: every eval-due step must coincide with a save-due step. 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}, " diff --git a/miles/utils/http_utils.py b/miles/utils/http_utils.py index 3349d09207..a9ae32d836 100644 --- a/miles/utils/http_utils.py +++ b/miles/utils/http_utils.py @@ -254,7 +254,6 @@ def init_http_client(args): _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine if args.eval_num_gpus > 0: - # The eval fleet is served by the same client; size for both. _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( diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index 66b1b844d2..a7d3407a98 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -160,7 +160,6 @@ def find_ready_snapshots(watch_dir: Path, min_rollout_id: int, consumed: set[int if is_complete_hf_export(child): ready.append((rollout_id, child)) elif looks_like_hf_checkpoint(child): - # Pre-marker checkpoint: accept once quiescent. newest_mtime = max(p.stat().st_mtime for p in child.iterdir()) if time.time() - newest_mtime > QUIESCENCE_SECS: ready.append((rollout_id, child)) diff --git a/train_async.py b/train_async.py index e1689cb72d..a158ce6b3d 100644 --- a/train_async.py +++ b/train_async.py @@ -56,7 +56,6 @@ async def train(args): 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: - # The base checkpoint is the snapshot for the pre-train eval; no export needed. await eval_dispatcher.dispatch(0, hf_dir=args.hf_checkpoint) # async train loop. @@ -97,8 +96,6 @@ async def train(args): await actor_model.update_weights(rollout_id=rollout_id) if should_run_periodic_action(rollout_id, args.eval_interval, num_rollout_per_epoch, args.num_rollout): - # The final point is never skipped: training is over, so overflow - # backpressure no longer costs any training time. await eval_dispatcher.dispatch(rollout_id, force=rollout_id == args.num_rollout - 1) if ( From ff9b694783e1d4cc7147d15e2023919cadf3a7a8 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 22:51:34 -0700 Subject: [PATCH 30/33] eval: drop remaining explanatory comments --- miles/backends/megatron_utils/model.py | 2 -- miles/ray/rollout/eval_dispatch.py | 1 - miles/ray/rollout/rollout_manager.py | 4 ---- miles/ray/rollout/rollout_server.py | 2 -- miles/rollout/fully_async_rollout.py | 2 -- miles/utils/arguments.py | 1 - tools/checkpoint_eval_service.py | 1 - 7 files changed, 13 deletions(-) diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 1a55b586fa..4080af8c32 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -966,8 +966,6 @@ def save_hf_model( torch.distributed.barrier() if torch.distributed.get_rank() == 0: - # The bridge exports zero weights for specs it has no mapping for; - # a marked-but-weightless snapshot must never exist. 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 " diff --git a/miles/ray/rollout/eval_dispatch.py b/miles/ray/rollout/eval_dispatch.py index b961ad3752..de1d97f439 100644 --- a/miles/ray/rollout/eval_dispatch.py +++ b/miles/ray/rollout/eval_dispatch.py @@ -60,7 +60,6 @@ async def _ensure_snapshot(self, rollout_id: int) -> tuple[str, float | None]: return hf_dir, time.time() - start def _reap_finished(self) -> None: - # The manager's eval lock serializes evals, so pending refs finish in order. while self.pending: done, _ = ray.wait([self.pending[0][1]], timeout=0) if not done: diff --git a/miles/ray/rollout/rollout_manager.py b/miles/ray/rollout/rollout_manager.py index 5300f44d0d..4941c6a1c2 100644 --- a/miles/ray/rollout/rollout_manager.py +++ b/miles/ray/rollout/rollout_manager.py @@ -43,7 +43,6 @@ logger = logging.getLogger(__name__) -# Per-attempt bound on eval-fleet weight loads; a zombie engine costs a skipped point. EVAL_WEIGHT_LOAD_TIMEOUT_SECS = 600.0 @@ -173,7 +172,6 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti Every failure mode degrades to a skipped point logged at ``rollout_id``. """ start_time = time.time() - # The lock holds across load -> verify -> generate; this is the pinning enforcement. async with self._eval_lock: srv = self.servers["eval"] try: @@ -194,8 +192,6 @@ async def _eval_on_dedicated_fleet(self, rollout_id: int, hf_dir: str, export_ti weight_version = str(rollout_id) for _attempt in range(2): try: - # A zombie engine (backend dead, actor alive) accepts the call and - # never answers; unbounded, it would hold the eval lock forever. await asyncio.wait_for( asyncio.gather( *[ diff --git a/miles/ray/rollout/rollout_server.py b/miles/ray/rollout/rollout_server.py index bb77d7b34e..5e97a6415d 100644 --- a/miles/ray/rollout/rollout_server.py +++ b/miles/ray/rollout/rollout_server.py @@ -134,8 +134,6 @@ def _resolve_sglang_config(args) -> SglangConfig: ) if eval_num_gpus > 0: - # update_weights=False must be explicit — resolve() would infer True from - # the matching model_path. config.models.append( ModelConfig( name="eval", diff --git a/miles/rollout/fully_async_rollout.py b/miles/rollout/fully_async_rollout.py index a0593c27a6..d5b6f43de8 100644 --- a/miles/rollout/fully_async_rollout.py +++ b/miles/rollout/fully_async_rollout.py @@ -118,8 +118,6 @@ async def _call_eval(self, input: RolloutFnInput) -> RolloutFnOutput: if self.args.eval_num_gpus > 0: return RolloutFnEvalOutput(data=await self._eval_fleet.run()) - # Shared-engine eval: the driver awaits it, so no weight update interleaves - # and the version stays pinned while submissions are paused. logger.info("Pausing fully-async producer submissions for shared-engine eval") self._producer_resumed.clear() try: diff --git a/miles/utils/arguments.py b/miles/utils/arguments.py index c8650c9d79..917480552f 100644 --- a/miles/utils/arguments.py +++ b/miles/utils/arguments.py @@ -972,7 +972,6 @@ 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) - # Dedicated eval fleet (checkpoint-interfaced eval; required for fully-async eval). parser.add_argument( "--eval-num-gpus", type=int, diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index a7d3407a98..bca8341676 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -194,7 +194,6 @@ def launch_server(service_args: Namespace) -> tuple[subprocess.Popen | None, str async def wait_server_healthy(ip: str, port: int, timeout: float = 1800.0) -> None: - # /health_generate runs a tiny generation; its 200 body is not JSON. await wait_http_ok(f"http://{ip}:{port}/health_generate", timeout=timeout) From e7eb153a2e9a3ec470c65e48d2a710193cb984a4 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 22:55:42 -0700 Subject: [PATCH 31/33] service: rename eval_snapshot to evaluate_snapshot --- tools/checkpoint_eval_service.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/checkpoint_eval_service.py b/tools/checkpoint_eval_service.py index bca8341676..d3fea57a6f 100644 --- a/tools/checkpoint_eval_service.py +++ b/tools/checkpoint_eval_service.py @@ -197,7 +197,7 @@ async def wait_server_healthy(ip: str, port: int, timeout: float = 1800.0) -> No await wait_http_ok(f"http://{ip}:{port}/health_generate", timeout=timeout) -async def eval_snapshot(args: Namespace, state, cache: dict, rollout_id: int, snapshot: Path) -> None: +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 @@ -261,7 +261,7 @@ async def main() -> None: for rollout_id, snapshot in ready: logger.info(f"Evaluating snapshot {snapshot} (rollout_id={rollout_id})") try: - await eval_snapshot(args, state, cache, rollout_id, snapshot) + 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") From 09826bb43e77ca141df872169af7df759f00b9c5 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 23:27:05 -0700 Subject: [PATCH 32/33] docs: clarify weight delivery and pause semantics per eval posture --- docs/user-guide/fully-async.md | 33 ++++++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index 77b1ed0b77..f4cc272e10 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -112,11 +112,22 @@ persisted anyway), and does eval get standalone GPU resources? 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. The weight version stays pinned because no update interleaves while the driver -awaits eval — 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). The cost is that rollout production stalls for -roughly the eval duration, which is fine for small debug eval sets. +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: @@ -152,6 +163,18 @@ are logged at the affected step. `eval/{ds}/weight_version/mean == eval/step` an 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 From df1130eec5afecc153691f490574f9be00a34679 Mon Sep 17 00:00:00 2001 From: Zhichenzzz Date: Tue, 21 Jul 2026 23:40:55 -0700 Subject: [PATCH 33/33] docs: note pause-the-world and external service compose for real runs --- docs/user-guide/fully-async.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/user-guide/fully-async.md b/docs/user-guide/fully-async.md index f4cc272e10..69c23ba8f9 100644 --- a/docs/user-guide/fully-async.md +++ b/docs/user-guide/fully-async.md @@ -107,7 +107,7 @@ 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) | +| **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**: