feat (async): support evaluation for fully-async training (dedicated fleet / pause-the-world / external service)

[Zhichenzzz]

Motivation

Fully-async rollout has no evaluation story: the producer never stops, so there is no quiet window to evaluate in, and running eval on the training engines is structurally broken — eval requests get aborted by every weight update, and no well-defined weight version is measured. Heavy eval sets also steal training inference capacity even in sync mode.

Design

Eval has exactly two postures, split by where the weights come from:

1. Against the live training engines — version pinned by blocking call order: the driver awaits the eval, so the next weight update cannot interleave. FullyAsyncRolloutFn serves it itself, pausing new producer submissions for the duration (in-flight requests finish and buffer — a gate, not a retract).
2. Against a checkpoint file — version pinned by the file itself. The only interface between training and eval is a (rollout_id, HF snapshot dir) pair; eval engines never join training weight updates, and weights reach them only through update_weights_from_disk(hf_dir, weight_version=str(rollout_id)).

args.eval_uses_snapshots is the single discriminant, derived from args in _resolve_rollout_functions and read by both the driver's dispatcher and RolloutManager — which never import each other.

Two ways to fill the second posture

They are mutually exclusive (validated), and they are different kinds of thing:

The in-job eval fleet (--eval-num-gpus N) delivers weights and nothing else. EvalFleet.pin(dir, version) -> GenerateState probes every engine, revives dead ones, loads the snapshot, and confirms each engine reports the expected version before returning the state to generate against. Because pin returns the state, generating before pinning is not expressible. Your --eval-function-path fn then generates against that state exactly as it would against the training engines, so custom eval fns work on the fleet unchanged.

A black box (--eval-function-path at a CheckpointEvalFn subclass) takes it from the directory onward:

class CheckpointEvalFn(abc.ABC):
    async def evaluate_checkpoint(self, checkpoint_dir, input) -> RolloutFnEvalOutput: ...
    def dispose(self): ...          # tear down anything launched in __init__
# raise EvalSkip(reason) for an attributable skipped point (eval/skipped_{reason})

Directory in, results out, anything in between — a non-sglang service implements this by submitting checkpoint_dir to its API. The fleet is deliberately not one of these; it is internal plumbing, not a user-facing contract. examples/fully_async/external_eval_fn.py is the reference implementation: it launches its own sglang server on user-named GPUs (MILES_EXTERNAL_EVAL_GPUS) or attaches to one (MILES_EXTERNAL_EVAL_URL).

Because the fn runs in-job, all eval config (datasets, sampling, rm, lora) comes from the real training args — nothing is hand-copied.

Engine configuration

The eval fleet inherits every --sglang-* setting from the rollout engines; override any single field with --eval-sglang-* (argparse.SUPPRESS defaults, so unset really means inherit; booleans take a --no- form). Not inheritable: TP comes from --eval-num-gpus-per-engine, which also places the engines, and dp/pp/ep/attn_cp fall back to 1 when the eval TP differs from the rollout TP — SGLang ties them to tp_size, so inheriting them across a different TP produces an engine that fails its own validation at boot.

Semantics

• Fire-and-forget eval: the point lands at the dispatch-time step even when it completes later (eval/lag_steps reports how late). The export is not fire-and-forget — it is a collective the training loop waits on, eval/export_time_seconds per point; reuse mode has no export at all.
• Version pinning: the engine-stamped weight_version is verified after every load; eval/{ds}/weight_version/mean == eval/step and mixed_version_ratio == 0 prove each point measured exactly the intended weights.
• Snapshot ownership: the dispatcher exports the snapshot and the dispatcher retires it, on every eval outcome, so no failure path can leak one. Staging holds at most --eval-keep-snapshots + --eval-max-in-flight model-sized dirs.
• Training is never derailed: every failure mode — unhealthy fleet, incomplete snapshot, pin violation, crashed eval, overflow — degrades to a skipped point logged at that step (eval/skipped_{busy,export_failed,ckpt_missing,unhealthy,pin_violation,crashed}). Overflow is two orthogonal knobs (--eval-max-in-flight N × --eval-overflow-policy backpressure|skip), and the final eval point is never skipped.
• Pause modes: eval adds no pause, retract, or weight sync of its own; it works under both --pause-generation-mode in_place and retract (retract additionally requires [sglang-miles] Fix flush_cache() no-op after pause_generation in retract sgl-project/sglang#31962 + [Fix] retract-mode flush_cache no-op crash #1750).

What's in the PR

Core:

• miles/rollout/checkpoint_eval.py — the black-box contract: CheckpointEvalFn, EvalSkip, is_checkpoint_eval_fn, and retarget_args (public for external backends)
• miles/ray/rollout/eval_fleet.py — EvalFleet.pin: health probe, load, per-engine version read-back, router probe
• miles/ray/rollout/eval_dispatch.py — driver-side EvalDispatcher: bounded in-flight dispatch, overflow policy, snapshot export and retirement, drain-before-dispose
• miles/ray/rollout/rollout_manager.py — _eval_checkpoint, one call site for both postures: pin first when there is a fleet, then call eval_generate_rollout
• miles/rollout/fully_async_rollout.py — _call_eval: generate on the injected state, or pause the producer and use its own
• miles/backends/megatron_utils/hf_export.py — export_hf_model_direct goes through miles' direct megatron→HF converters, so export coverage always matches weight-sync coverage; save_hf_model moved here from model.py, restoring the 1,000-line file budget
• examples/fully_async/external_eval_fn.py + run_qwen3_5_4b_fully_async_eval.py — both backends behind one flag (--eval-backend fleet|external)
• New args + a validation matrix shared by both snapshot postures, docs (docs/user-guide/fully-async.md, CLI reference)

Tests: tests/e2e/megatron/test_qwen3_4b_fully_async_eval.py covers all three postures as short Qwen3-4B GRPO runs (suite stage-c-8-gpu-h200, labels megatron / eval; apply run-ci-eval to run it on a PR), plus fast tests for the contract, fleet pin/recovery, dispatcher policies and snapshot ownership, the manager snapshot path, and the --eval-sglang-* inheritance mechanism.

Review history

The branch was reviewed and reworked in place; the design above is the reworked one. What changed and why:

• FleetEvalFn → EvalFleet, out of the CheckpointEvalFn contract and into miles/ray/rollout/. The fleet is not a black box — it must hand generation back to the configured eval fn — and forcing it into the contract produced the inner constructor argument no other implementation had, a "privileged, not via the CLI flag" caveat, and --eval-num-gpus silently beating --eval-function-path. resolve_checkpoint_eval_fn disappeared with it.
• Snapshot lifecycle given one owner. Export and deletion sat on opposite sides of a Ray call, and the manager only reached its GC on the success path, so every skip and crash leaked a model-sized directory into --eval-hf-dir — a node-down failure with the documented /dev/shm staging.
• Shared-engine eval never paused the producer. resolve_rollout_function_paths computed the eval path before the --fully-async override, so eval resolved to a second InferenceRolloutFn and FullyAsyncRolloutFn._call_eval was never reached. The GPU verification below therefore did not exercise the pause, despite reporting it.
• --eval-sglang-* replaced "the fleet is whatever rollout is"; the sglang_overrides merge in _compute_server_args moved after the args-derived branches so a per-group override actually wins. That also fixes existing --sglang-config per-group overrides of dtype / LoRA / disaggregation keys, which previously lost to those branches — worth a look from anyone running PD configs.
• Correctness fixes: the export barrier is now reachable when rank 0 fails alone (it used to leave every other rank in torch.distributed.barrier() until the NCCL watchdog); the .complete marker is cleared before a re-export and written after the LoRA adapter; both snapshot postures now get the same validation, including the class-based-API check whose absence crashed the black-box posture at the first eval point.
• Dropped: --eval-model-path (nothing set it, and the boot weights are overwritten before the first eval); the --eval-keep-snapshots >= --eval-max-in-flight assert (the GC ring only ever held retired snapshots).

Verification

CI on fc2a4bf, the first commit:

• Full PR Test run green with the eval e2e actually executing: test_qwen3_4b_fully_async_eval.py PASS in stage-c-8-gpu-h200 (23 min) — shared / fleet / external each completed as a separate successful ray job. Note: without a run-ci-<label> PR label the GPU stages select nothing, so a labelless green does not exercise this test.
• All CPU shards green; tests/fast 3934 passed / 20 skipped / 0 failed.

GPU matrix (8×H200, Qwen3.5-4B, dapo-math training, gsm8k eval), also on fc2a4bf:

• Fleet + tmpfs: all 5 eval points pinned, gsm8k 0.640 → 0.882 over 20 steps, lag_steps ≤ 11, training uninterrupted
• Fleet + ckpt reuse: zero-extra-export points 0.641 → 0.847; overflow under skip observed and attributed
• External service: on the same checkpoint the external score matched the in-job score exactly
• Kill tests: an eval engine killed mid-run is detected by the pre-eval probe, revived by recover(), and the next eval lands pinned — training untouched

These runs predate the rework and need repeating on the head commit. In particular the shared-engine numbers were produced by code in which the producer pause could not run, so that posture is effectively unverified.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces a dedicated evaluation fleet for fully-async training, allowing evaluations to run on separate engines pinned to Hugging Face checkpoint snapshots without stalling training. It adds new CLI configuration options, an EvalDispatcher to manage async evaluation tasks, a standalone checkpoint evaluation service tool, and corresponding tests. The review feedback focuses on improving robustness and error handling, specifically suggesting to verify if the checkpoint path is a local directory before exporting, replacing assert statements with explicit exceptions (like ValueError and RuntimeError) to prevent them from being stripped under Python optimization, and narrowing exception handling in health check probes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

If args.hf_checkpoint is a Hugging Face Hub model ID (e.g., Qwen/Qwen2.5-7B-Instruct) rather than a local directory, Path(args.hf_checkpoint).iterdir() will raise a FileNotFoundError and crash the direct HF export process. Check if the path is a directory before attempting to copy metadata files.

Suggested change

	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)
	if is_writer:
	assert weight_map, f"HF export to {path} produced no weights"
	checkpoint_path = Path(args.hf_checkpoint)
	if not checkpoint_path.is_dir():
	raise FileNotFoundError(
	f"hf_checkpoint '{args.hf_checkpoint}' is not a local directory. "
	f"Direct HF export requires a local directory containing the model metadata files."
	)
	for meta_file in checkpoint_path.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)

[gemini-code-assist]

Use ValueError instead of assert for validating command-line arguments. If Python is run with optimization (-O), assert statements are compiled away, which would silently bypass all validation checks.

    if args.eval_num_gpus > 0:
        if not enable_experimental_rollout_refactor():
            raise ValueError("--eval-num-gpus requires the class-based rollout API (MILES_EXPERIMENTAL_ROLLOUT_REFACTOR=1).")
        if args.eval_interval is None:
            raise ValueError("--eval-num-gpus requires --eval-interval.")
        if args.eval_hf_dir is None and args.save_hf is None:
            raise ValueError(
                "--eval-num-gpus requires a snapshot source: set --eval-hf-dir (staging exports) "
                "or --save-hf (reuse periodic HF checkpoints)."
            )
        if args.colocate:
            raise ValueError(
                "--eval-num-gpus is not supported with --colocate; "
                "use tools/checkpoint_eval_service.py against --save-hf checkpoints instead."
            )
        if args.debug_train_only or args.debug_rollout_only:
            raise ValueError("--eval-num-gpus is not supported with debug_train_only/debug_rollout_only.")
        if args.eval_num_gpus % args.eval_num_gpus_per_engine != 0:
            raise ValueError(
                f"eval_num_gpus ({args.eval_num_gpus}) must be divisible by "
                f"eval_num_gpus_per_engine ({args.eval_num_gpus_per_engine})."
            )
        if args.eval_keep_snapshots < args.eval_max_in_flight:
            raise ValueError(
                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.
            if args.save_interval is None or args.eval_interval % args.save_interval != 0:
                raise ValueError(
                    "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."
                )

References

1. Use ValueError instead of assert for validating function or constructor arguments.

[gemini-code-assist]

Avoid catching broad exceptions like Exception in health check probes. Only catch expected network, timeout, or unavailability exceptions (such as httpx.HTTPError and asyncio.TimeoutError) to prevent swallowing internal programming errors or bugs.

Suggested change

	try:
	await get(f"http://{ip}:{port}/health_generate", max_retries=1)
	return
	except Exception:
	await asyncio.sleep(5)
	import httpx
	try:
	await get(f"http://{ip}:{port}/health_generate", max_retries=1)
	return
	except (httpx.HTTPError, asyncio.TimeoutError):
	await asyncio.sleep(5)

References

1. In health check endpoints or probes, avoid catching broad exceptions (like Exception). Only catch expected network, timeout, or unavailability exceptions.

[gemini-code-assist]

Do not use assert for runtime checks that must always execute, as they are stripped when Python is run with optimization (-O). Use RuntimeError instead.

Suggested change

	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}"
	)

[gemini-code-assist]

Use ValueError instead of assert for validating command-line arguments or configurations to prevent checks from being stripped under Python optimization (-O).

Suggested change

	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")

References

1. Use ValueError instead of assert for validating function or constructor arguments.

[gemini-code-assist]

Use FileNotFoundError instead of assert for validating file or directory existence to prevent checks from being stripped under Python optimization (-O).

Suggested change

	watch_dir = Path(service_args.watch_dir)
	assert watch_dir.is_dir(), f"--watch-dir {watch_dir} does not exist"
	watch_dir = Path(service_args.watch_dir)
	if not watch_dir.is_dir():
	raise FileNotFoundError(f"--watch-dir {watch_dir} does not exist")

References

1. Use ValueError instead of assert for validating function or constructor arguments.

[ashtonchew]

I found one mismatch between the shared-engine eval contract in this PR and the current head (9febdfa9).

resolve_rollout_function_paths() resolves the default eval path before applying the fully-async rollout override. With rollout_function_path=None, eval_function_path=None, and fully_async=True, it currently returns:

(
    "miles.rollout.fully_async_rollout.FullyAsyncRolloutFn",
    "miles.rollout.inference_rollout.inference_rollout_common.InferenceRolloutFn",
)

RolloutManager reuses the rollout instance only when the two paths are equal. The default shared-engine eval therefore constructs a separate InferenceRolloutFn and bypasses FullyAsyncRolloutFn._call_eval(), including its producer pause.

The bounded fix is to resolve eval_path = args.eval_function_path or rollout_path after the fully-async override. An explicit --eval-function-path remains unchanged. I reproduced the failure with a focused resolver test, added a preservation test for an explicit eval backend, and verified both pass after the reorder fix.

[yueming-yuan]

I found one mismatch between the shared-engine eval contract in this PR and the current head (9febdfa9).

resolve_rollout_function_paths() resolves the default eval path before applying the fully-async rollout override. With rollout_function_path=None, eval_function_path=None, and fully_async=True, it currently returns:

(
    "miles.rollout.fully_async_rollout.FullyAsyncRolloutFn",
    "miles.rollout.inference_rollout.inference_rollout_common.InferenceRolloutFn",
)

RolloutManager reuses the rollout instance only when the two paths are equal. The default shared-engine eval therefore constructs a separate InferenceRolloutFn and bypasses FullyAsyncRolloutFn._call_eval(), including its producer pause.

The bounded fix is to resolve eval_path = args.eval_function_path or rollout_path after the fully-async override. An explicit --eval-function-path remains unchanged. I reproduced the failure with a focused resolver test, added a preservation test for an explicit eval backend, and verified both pass after the reorder fix.

Thanks for the comment! This issue has been fixed in b0aa5e5 commit

[Zhichenzzz]

LGTM for the code changes!