diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ad0d4acdfac..7383b652554 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -6,6 +6,7 @@ Changelog **New Features** +- Add Puzzletron dynamic post-MIP downstream evaluation through ``lmms-eval`` with vLLM-backed checkpoint evaluation, setup-wizard topology/resource prompts, and an opt-in Nemotron-3 Nano 30B A3B BF16 example flow. - Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. - Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index d5149969d13..b1be0e708e1 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -268,7 +268,13 @@ import vllm with open(os.environ["PUZZLETRON_CI_ENVIRONMENT"], encoding="utf-8") as stream: ci_environment = json.load(stream) -for package in ("torch", "vllm", "nemo-automodel", "aiperf", "nvidia-modelopt"): +for package in ( + "torch", + "vllm", + "nemo-automodel", + "aiperf", + "nvidia-modelopt", +): print(package, metadata.version(package)) print("torch CUDA", torch.version.cuda) @@ -501,6 +507,26 @@ After `mip`, prepare one deduplicated online-evaluation plan. Repeat `--profile-id` for every configured profile; aliases ensure that an identical architecture is evaluated once while remaining visible in every profile. +Downstream `lmms-eval` nodes use a separate evaluator Python. The reproducible +example path is pinned to `lmms-eval==0.7.2` by +`examples/puzzletron/requirements-lmms-eval.txt` and recorded in +`ci_environment.json`. Keep this separate from the Puzzletron runtime +environment because `lmms-eval==0.7.2` pins `wandb==0.25.0`, while the pinned +AutoModel build requires a newer `wandb`. + +```bash +python3 -m venv /workspace/.venv-lmms-eval +source /workspace/.venv-lmms-eval/bin/activate +python -m pip install --upgrade pip "setuptools>=80,<81" wheel packaging +VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \ + python -m pip install --no-build-isolation -e "${VLLM_ROOT}" +python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements-lmms-eval.txt" +python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"' +deactivate + +export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python +``` + ```bash python examples/puzzletron/run_profile_online_evaluation.py \ --puzzle-dir "$PUZZLE_DIR" --prepare \ diff --git a/examples/puzzletron/ci_environment.json b/examples/puzzletron/ci_environment.json index d7520a69109..1e7d0f5be5f 100644 --- a/examples/puzzletron/ci_environment.json +++ b/examples/puzzletron/ci_environment.json @@ -5,6 +5,7 @@ "torch": "2.11.0", "torchvision": "0.26.0", "transformers": "5.8.1", + "lmms_eval": "0.7.2", "nemo_automodel": { "base_version": "0.5.0", "repository": "https://github.com/Separius/Automodel.git", diff --git a/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml new file mode 100644 index 00000000000..a5e8d8c77b4 --- /dev/null +++ b/examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml @@ -0,0 +1,68 @@ +# @package _global_ + +defaults: + - default + - _self_ + +# Opt-in downstream lmms-eval workflow for the realized runtime-075 candidate. +# Use PUZZLETRON_LMMS_EVAL_PYTHON to point at the isolated pinned evaluator env. +# Non-empty post_mip.flows replaces the legacy post-MIP tail in the v2 orchestrator. +zero_shot_evaluation: + enabled: false +aiperf: + enabled: false +global_distillation_sanity: + enabled: false +global_distillation: + enabled: false +post_distillation_evaluation: + enabled: false + +post_mip: + flows: + runtime-075-lmms-eval: + source: + run: runtime-075 + variants: all + objectives: all + nodes: + best_mip: + type: filter + mode: top_k + metric: mip.score + direction: minimize + top_k: 1 + materialized: + type: materialize + input: best_mip + lmms_eval: + type: downstream_evaluation + input: materialized + config: + command_prefix: + - ${oc.env:PUZZLETRON_LMMS_EVAL_PYTHON} + - -m + - lmms_eval + model: vllm + checkpoint_arg: model + tasks: + - ifeval + - gsm8k + limit: 128 + batch_size: 1 + log_samples: true + timeout_seconds: 7200 + topology: + tensor_parallel_size: 8 + pipeline_parallel_size: 1 + data_parallel_size: 1 + prefill_context_parallel_size: 1 + decode_context_parallel_size: 1 + enable_expert_parallel: false + distributed_executor_backend: mp + gpu_group_size: 8 + model_args: + dtype: bfloat16 + gpu_memory_utilization: 0.85 + max_model_len: 262144 + trust_remote_code: ${model.trust_remote_code} diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index 0b28d0b0e99..fd13f55ab38 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -105,9 +105,11 @@ metric lists or cases. Later filters reference metrics as `mip.` or - `evaluation`: evaluates either a config-only candidate or a checkpoint and publishes all result metrics. - `aiperf`: benchmarks a checkpoint and publishes all result metrics. +- `downstream_evaluation`: runs `lmms-eval` against a materialized checkpoint + and publishes task metrics. - `global_kd`: produces a new checkpoint revision. -- `ptq` and `downstream_evaluation`: reserved interfaces; configuring either - currently fails plan compilation with a clear not-implemented error. +- `ptq`: reserved interface; configuring it currently fails plan compilation + with a clear not-implemented error. Nodes that require checkpoints never materialize implicitly. Add a `materialize` node where the transition is needed. @@ -129,6 +131,28 @@ Selection still follows `input`; `model_source` only chooses the artifact operat on. This supports a long KD run selected using short-KD/PTQ results but restarted from the original candidate. +## Downstream evaluation + +`downstream_evaluation` shells out to `python -m lmms_eval` through +`command_prefix`. Install the pinned evaluator into an isolated environment +rather than the Puzzletron runtime environment, because `lmms-eval==0.7.2` pins +`wandb==0.25.0` and the pinned AutoModel build requires a newer `wandb`: + +```bash +python3 -m venv /workspace/.venv-lmms-eval +source /workspace/.venv-lmms-eval/bin/activate +python -m pip install -r examples/puzzletron/requirements-lmms-eval.txt +python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"' +deactivate + +export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python +``` + +The runner derives the realized checkpoint path, vLLM topology arguments, task +list, and output path from the campaign config. Use `model_args` only for +non-derived model options such as dtype or maximum model length, and `extra_args` +only for non-reserved `lmms-eval` flags. + ## Filters `top_k` accepts one integer or separate homogeneous/heterogeneous quotas. diff --git a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py index ba46a5b8178..90f1f745f9a 100644 --- a/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py +++ b/modelopt/torch/puzzletron/orchestration/adapters/post_mip.py @@ -111,7 +111,7 @@ def plan(self, plan: CampaignPlan, node: StagePlanNode) -> WorkPlan: config = _node_config(plan, node.stage_id) node_type = str(config.get("type")) count = 1 if node_type in {"filter", "manual_filter"} else node.instances - if node_type == "evaluation": + if node_type in {"evaluation", "downstream_evaluation"}: available = _available_evaluation_candidates(plan, node.stage_id, config) if available is not None: if available < 1: diff --git a/modelopt/torch/puzzletron/orchestration/compiler.py b/modelopt/torch/puzzletron/orchestration/compiler.py index 924716ba946..1f493aeec52 100644 --- a/modelopt/torch/puzzletron/orchestration/compiler.py +++ b/modelopt/torch/puzzletron/orchestration/compiler.py @@ -57,6 +57,7 @@ "aiperf": ExecutionStrategy.SHARDED, } + def _mapping(value: Any) -> dict[str, Any]: return dict(value) if isinstance(value, Mapping) else {} @@ -85,7 +86,6 @@ def _mapping(value: Any) -> dict[str, Any]: "downstream_evaluation": { "kind": "evaluator", "accepts": {"checkpoint"}, - "implemented": False, }, } @@ -516,7 +516,7 @@ def compile_campaign_plan( parallel[key] = node_config[key] elif key in global_kd and key not in parallel: parallel[key] = global_kd[key] - if dynamic["node_type"] == "aiperf": + if dynamic["node_type"] in {"aiperf", "downstream_evaluation"}: topology = _mapping(node_config.get("topology")) topology_mesh = vllm_topology_to_mesh(topology) if override: @@ -525,7 +525,7 @@ def compile_campaign_plan( if ParallelMesh.from_mapping(overridden) != topology_mesh: raise ValueError( f"{stage_id} execution parallel override conflicts with " - "its AIPerf topology" + "its vLLM topology" ) mesh = topology_mesh else: diff --git a/modelopt/torch/puzzletron/orchestration/controller.py b/modelopt/torch/puzzletron/orchestration/controller.py index b2a567a3fd0..335066e7100 100644 --- a/modelopt/torch/puzzletron/orchestration/controller.py +++ b/modelopt/torch/puzzletron/orchestration/controller.py @@ -74,6 +74,29 @@ def create_executor(plan: CampaignPlan, *, local: bool = False) -> Executor: raise ValueError(f"Unsupported runner kind: {plan.runner.kind}") +def _stage_dashboard_display_name( + config: Mapping[str, Any], + stage_id: str, + *, + granularity: str | None = None, +) -> str: + if stage_id.startswith("post."): + parts = stage_id.split(".", 2) + if len(parts) == 3: + _prefix, flow_id, node_id = parts + node = ( + (config.get("post_mip") or {}) + .get("flows", {}) + .get(flow_id, {}) + .get("nodes", {}) + .get(node_id, {}) + ) + node_type = str(node.get("type") or "") if isinstance(node, Mapping) else "" + if node_type == "downstream_evaluation": + return "Downstream Evaluation" + return stage_display_name(stage_id, granularity=granularity) + + @dataclass class DryRunSubmission: stage_id: str @@ -748,7 +771,8 @@ def _stage_views(self) -> list[StageView]: views.append( StageView( stage_id=node.stage_id, - display_name=stage_display_name( + display_name=_stage_dashboard_display_name( + self.plan.experiment_config, node.stage_id, granularity=str(granularity) if granularity is not None else None, ), diff --git a/modelopt/torch/puzzletron/orchestration/progress.py b/modelopt/torch/puzzletron/orchestration/progress.py index e395b9477ea..a9c943357e6 100644 --- a/modelopt/torch/puzzletron/orchestration/progress.py +++ b/modelopt/torch/puzzletron/orchestration/progress.py @@ -434,6 +434,7 @@ def _post_mip_progress( labels = { "evaluation": "evaluated", + "downstream_evaluation": "evaluated", "aiperf": "benchmarked", "global_kd": "distilled", "materialize": "materialized", diff --git a/modelopt/torch/puzzletron/post_mip/builtin.py b/modelopt/torch/puzzletron/post_mip/builtin.py index 8c62a1d5a11..050352af445 100644 --- a/modelopt/torch/puzzletron/post_mip/builtin.py +++ b/modelopt/torch/puzzletron/post_mip/builtin.py @@ -10,7 +10,12 @@ from .base import NodeCapabilities, NodeKind, PostMIPNode, post_mip_node from .filters import filter_metric_references, validate_filter_config from .records import ArtifactKind -from .reporting import render_aiperf_report, render_evaluation_report, render_global_kd_report +from .reporting import ( + render_aiperf_report, + render_downstream_evaluation_report, + render_evaluation_report, + render_global_kd_report, +) if TYPE_CHECKING: from collections.abc import Mapping @@ -115,6 +120,9 @@ class DownstreamEvaluationNode(PostMIPNode): NodeKind.EVALUATOR, frozenset({ArtifactKind.CHECKPOINT}), distributed=True, - implemented=False, default_strategy="sharded", ) + + @classmethod + def render_report(cls, node, payload): + return render_downstream_evaluation_report(str(payload["section_id"]), payload) diff --git a/modelopt/torch/puzzletron/post_mip/reporting.py b/modelopt/torch/puzzletron/post_mip/reporting.py index 86e3c11f44e..bc287990805 100644 --- a/modelopt/torch/puzzletron/post_mip/reporting.py +++ b/modelopt/torch/puzzletron/post_mip/reporting.py @@ -16,6 +16,7 @@ __all__ = [ "build_post_mip_report_payloads", "render_aiperf_report", + "render_downstream_evaluation_report", "render_evaluation_report", "render_global_kd_report", ] @@ -356,6 +357,16 @@ def render_aiperf_report(section_id: str, payload: Mapping[str, Any]) -> str: ) +def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str: + """Render lmms-eval task metrics for downstream-evaluation nodes.""" + + return render_evaluation_report(section_id, payload).replace( + "

Candidate evaluation

", + "

Downstream evaluation

", + 1, + ) + + def render_global_kd_report(section_id: str, payload: Mapping[str, Any]) -> str: """Render several candidate KD histories on shared, lineage-colored plots.""" diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 819cbbce010..4de690a60b8 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -24,7 +24,10 @@ import json import math import os +import shlex +import signal import subprocess +import sys import traceback import uuid from contextlib import contextmanager @@ -33,6 +36,7 @@ from typing import Any, Iterator, Mapping, Sequence from ..identity import canonicalize, stable_hash +from ..orchestration.mesh import normalize_vllm_topology from .base import CompiledPostMIPNode, NodeKind, compile_post_mip_flows from .filters import apply_filter from .records import ArtifactKind, CandidateLedger, CandidateSet, NodeObservation @@ -523,6 +527,610 @@ def _aiperf( } +_LMMS_EVAL_MODEL_ARG_FIELDS = frozenset( + { + "dtype", + "gpu_memory_utilization", + "max_model_len", + "trust_remote_code", + "tokenizer", + "tokenizer_mode", + "enforce_eager", + "limit_mm_per_prompt", + "reasoning_parser", + } +) +_LMMS_EVAL_RESERVED_TOPOLOGY_MODEL_ARG_FIELDS = frozenset( + { + "tensor_parallel_size", + "pipeline_parallel_size", + "data_parallel_size", + "prefill_context_parallel_size", + "decode_context_parallel_size", + "enable_expert_parallel", + "distributed_executor_backend", + "expert_parallel_size", + "gpu_group_size", + "tp", + "pp", + "dp", + "prefill_cp", + "decode_cp", + "ep", + } +) +_LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS = frozenset( + { + "--model_args", + "--model-args", + "--output_path", + "--output-path", + "--tasks", + } +) +_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 + + +def _as_cli_bool(value: bool) -> str: + return "True" if value else "False" + + +def _as_lmms_eval_arg(value: Any) -> str: + if isinstance(value, bool): + return _as_cli_bool(value) + if isinstance(value, (int, float)) and not isinstance(value, bool): + return str(value) + if isinstance(value, (list, tuple, dict)): + return json.dumps(value, sort_keys=True, separators=(",", ":")) + return str(value) + + +def _join_cli_values(value: Any, *, path: str) -> str: + if isinstance(value, str): + text = value.strip() + if not text: + raise ValueError(f"{path} must not be empty") + return text + if not isinstance(value, Sequence): + raise TypeError(f"{path} must be a string or sequence") + values = [str(item).strip() for item in value] + if not values or any(not item for item in values): + raise ValueError(f"{path} must contain at least one non-empty value") + return ",".join(values) + + +def _lmms_eval_model_arg_keys(value: str) -> tuple[str, ...]: + keys: list[str] = [] + start = 0 + depth = 0 + quote: str | None = None + escaped = False + + def append(segment: str) -> None: + key, separator, _ = segment.strip().partition("=") + if separator and key.strip(): + keys.append(key.strip()) + + for index, char in enumerate(value): + if escaped: + escaped = False + continue + if quote: + if char == "\\": + escaped = True + elif char == quote: + quote = None + continue + if char in {"'", '"'}: + quote = char + elif char in "([{": + depth += 1 + elif char in ")]}" and depth: + depth -= 1 + elif char == "," and depth == 0: + append(value[start:index]) + start = index + 1 + append(value[start:]) + return tuple(keys) + + +def _lmms_eval_reserved_model_arg_fields(checkpoint_arg: str) -> frozenset[str]: + return frozenset( + key + for key in ( + str(checkpoint_arg).strip(), + *_LMMS_EVAL_RESERVED_TOPOLOGY_MODEL_ARG_FIELDS, + ) + if key + ) + + +def _reject_reserved_lmms_eval_model_args( + keys: Sequence[Any], reserved_fields: frozenset[str] +) -> None: + reserved = sorted({str(key).strip() for key in keys} & reserved_fields) + if reserved: + raise ValueError( + "downstream_evaluation.config.model_args must not set reserved " + f"lmms-eval model arguments: {', '.join(reserved)}" + ) + + +def _configured_lmms_eval_tasks(settings: Mapping[str, Any]) -> tuple[str, ...]: + tasks = _join_cli_values(settings.get("tasks"), path="downstream_evaluation.config.tasks") + values = tuple(task.strip() for task in tasks.split(",")) + if not values or any(not task for task in values): + raise ValueError("downstream_evaluation.config.tasks must contain non-empty task names") + return values + + +def _model_arg_string(values: Mapping[str, Any]) -> str: + parts = [] + for key, value in values.items(): + if value is None: + continue + key_text = str(key).strip() + if not key_text or "," in key_text or "=" in key_text: + raise ValueError(f"invalid lmms-eval model_args key: {key!r}") + rendered = _as_lmms_eval_arg(value) + if "," in rendered: + raise ValueError( + f"lmms-eval model_args value for {key_text!r} contains a comma; " + "provide model_args as a preformatted string instead" + ) + parts.append(f"{key_text}={rendered}") + if not parts: + raise ValueError("lmms-eval model_args must contain at least the checkpoint path") + return ",".join(parts) + + +def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) -> str: + raw = settings.get("model_args") + checkpoint_arg = str(settings.get("checkpoint_arg", "model")) + topology = dict(settings.get("topology") or {}) + canonical_topology = normalize_vllm_topology(topology) if topology else {} + reserved_fields = _lmms_eval_reserved_model_arg_fields(checkpoint_arg) + derived = { + checkpoint_arg: checkpoint, + } + if canonical_topology: + derived.update( + { + "tensor_parallel_size": canonical_topology["tp"], + "pipeline_parallel_size": canonical_topology["pp"], + "data_parallel_size": canonical_topology["dp"], + "enable_expert_parallel": canonical_topology["enable_expert_parallel"], + "distributed_executor_backend": canonical_topology[ + "distributed_executor_backend" + ], + } + ) + for key in _LMMS_EVAL_MODEL_ARG_FIELDS: + if key in settings: + derived[key] = settings[key] + + if isinstance(raw, str): + _reject_reserved_lmms_eval_model_args( + _lmms_eval_model_arg_keys(raw), reserved_fields + ) + prefix = raw.strip().strip(",") + suffix = _model_arg_string(derived) + return ",".join(part for part in (prefix, suffix) if part) + if raw is not None and not isinstance(raw, Mapping): + raise TypeError("downstream_evaluation.config.model_args must be a mapping or string") + _reject_reserved_lmms_eval_model_args(tuple((raw or {}).keys()), reserved_fields) + merged = dict(raw or {}) + merged.update(derived) + return _model_arg_string(merged) + + +def _command_prefix(settings: Mapping[str, Any]) -> list[str]: + raw = settings.get("command_prefix") + if raw is None: + return [sys.executable, "-m", "lmms_eval"] + if isinstance(raw, str): + values = [raw] + else: + values = [str(item) for item in raw] + if not values or any(not value for value in values): + raise ValueError("downstream_evaluation.config.command_prefix must not be empty") + return values + + +def _lmms_eval_extra_args(settings: Mapping[str, Any]) -> list[str]: + raw = settings.get("extra_args") + if raw is None: + return [] + if isinstance(raw, str): + values = shlex.split(raw) + elif isinstance(raw, Sequence): + values = [str(item) for item in raw] + else: + raise TypeError("downstream_evaluation.config.extra_args must be a string or sequence") + if any(not value for value in values): + raise ValueError("downstream_evaluation.config.extra_args must not contain empty values") + reserved = sorted( + { + value.split("=", 1)[0] + for value in values + if value.split("=", 1)[0] in _LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS + } + ) + if reserved: + raise ValueError( + "downstream_evaluation.config.extra_args must not set reserved " + f"lmms-eval flags: {', '.join(reserved)}" + ) + return values + + +def _lmms_eval_command( + settings: Mapping[str, Any], + *, + checkpoint: str, + output_path: Path, +) -> tuple[list[str], dict[str, str], float | None]: + """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" + + tasks = ",".join(_configured_lmms_eval_tasks(settings)) + argv = [ + *_command_prefix(settings), + "--model", + str(settings.get("model", "vllm")), + "--model_args", + _merge_lmms_eval_model_args(settings, checkpoint), + "--tasks", + tasks, + "--batch_size", + str(settings.get("batch_size", 1)), + "--output_path", + str(output_path), + ] + optional_fields = { + "limit": "--limit", + "num_fewshot": "--num_fewshot", + "seed": "--seed", + "verbosity": "--verbosity", + "device": "--device", + "use_cache": "--use_cache", + } + for key, flag in optional_fields.items(): + value = settings.get(key) + if value is not None: + argv.extend([flag, str(value)]) + if settings.get("gen_kwargs") is not None: + argv.extend( + [ + "--gen_kwargs", + ( + settings["gen_kwargs"] + if isinstance(settings["gen_kwargs"], str) + else _model_arg_string(dict(settings["gen_kwargs"])) + ), + ] + ) + if bool(settings.get("log_samples", False)): + argv.append("--log_samples") + argv.extend(_lmms_eval_extra_args(settings)) + + env = os.environ.copy() + for key, value in dict(settings.get("env") or {}).items(): + if value is not None: + env[str(key)] = str(value) + if settings.get("cache_dir") is not None: + env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"])) + timeout = settings.get("timeout_seconds", settings.get("timeout")) + return argv, env, (float(timeout) if timeout is not None else None) + + +def _metric_key(value: Any) -> str: + return ( + str(value) + .strip() + .replace(" ", "_") + .replace(",", "_") + .replace("/", "_") + .replace("\\", "_") + ) + + +def _numeric_metrics(task_payload: Mapping[str, Any]) -> dict[str, float]: + metrics = {} + for metric_name, value in task_payload.items(): + if ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and math.isfinite(value) + ): + metrics[str(metric_name)] = float(value) + return metrics + + +def _flatten_lmms_eval_metrics(payload: Mapping[str, Any]) -> dict[str, float]: + results = payload.get("results") + if not isinstance(results, Mapping): + return {} + metrics = {} + for task_name, task_payload in results.items(): + if not isinstance(task_payload, Mapping): + continue + for metric_name, value in _numeric_metrics(task_payload).items(): + metrics[f"{_metric_key(task_name)}.{_metric_key(metric_name)}"] = value + return metrics + + +def _resolved_lmms_eval_tasks( + payload: Mapping[str, Any], configured_tasks: Sequence[str] +) -> tuple[str, ...]: + group_subtasks = payload.get("group_subtasks") + if not isinstance(group_subtasks, Mapping): + group_subtasks = {} + + def expand(task: str, seen: frozenset[str]) -> tuple[str, ...]: + raw_subtasks = group_subtasks.get(task) + if ( + isinstance(raw_subtasks, Sequence) + and not isinstance(raw_subtasks, str) + and raw_subtasks + and task not in seen + ): + expanded = [] + for raw_subtask in raw_subtasks: + expanded.extend(expand(str(raw_subtask), seen | {task})) + return tuple(dict.fromkeys(expanded)) + return (task,) + + resolved = [] + for task in configured_tasks: + resolved.extend(expand(task, frozenset())) + return tuple(dict.fromkeys(resolved)) + + +def _sample_count(payload: Mapping[str, Any], task: str) -> float | None: + samples = payload.get("n-samples", payload.get("n_samples")) + if not isinstance(samples, Mapping): + return None + value = samples.get(task) + if isinstance(value, Mapping): + if "effective" in value: + value = value["effective"] + elif "original" in value: + value = value["original"] + else: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value): + return None + return float(value) + + +def _validate_lmms_eval_completion( + payload: Mapping[str, Any], configured_tasks: Sequence[str] +) -> dict[str, float]: + results = payload.get("results") + if not isinstance(results, Mapping): + raise RuntimeError("lmms-eval result is missing the results mapping") + + expected_tasks = _resolved_lmms_eval_tasks(payload, configured_tasks) + missing_results = [task for task in expected_tasks if task not in results] + if missing_results: + raise RuntimeError( + "lmms-eval result is missing configured task results: " + f"{sorted(missing_results)}" + ) + + missing_metrics = [ + task + for task in expected_tasks + if not isinstance(results[task], Mapping) or not _numeric_metrics(results[task]) + ] + if missing_metrics: + raise RuntimeError( + "lmms-eval result has no numeric metrics for configured tasks: " + f"{sorted(missing_metrics)}" + ) + + sample_counts = {} + missing_samples = [] + zero_samples = [] + for task in expected_tasks: + sample_count = _sample_count(payload, task) + if sample_count is None: + missing_samples.append(task) + elif sample_count <= 0: + zero_samples.append(task) + else: + sample_counts[task] = sample_count + if missing_samples: + raise RuntimeError( + "lmms-eval result is missing sample counts for configured tasks: " + f"{sorted(missing_samples)}" + ) + if zero_samples: + raise RuntimeError( + "lmms-eval result has zero effective samples for configured tasks: " + f"{sorted(zero_samples)}" + ) + return sample_counts + + +def _lmms_eval_result_payload(output_path: Path) -> tuple[dict[str, Any], Path]: + candidates = [] + for path in sorted(output_path.rglob("*.json")): + try: + payload = json.loads(path.read_text()) + except (OSError, ValueError): + continue + if isinstance(payload, Mapping) and isinstance(payload.get("results"), Mapping): + candidates.append((path.stat().st_mtime_ns, path, dict(payload))) + if not candidates: + raise FileNotFoundError(f"lmms-eval wrote no JSON results below {output_path}") + _mtime, path, payload = max(candidates, key=lambda item: item[0]) + return payload, path + + +def _write_lmms_eval_streams( + output_path: Path, result: subprocess.CompletedProcess[str] +) -> dict[str, str]: + stream_paths = {} + for stream_name, text in (("stdout", result.stdout), ("stderr", result.stderr)): + if not text: + continue + stream_path = output_path / f"{stream_name}.txt" + stream_path.write_text(text) + stream_paths[f"{stream_name}_path"] = str(stream_path) + return stream_paths + + +def _lmms_eval_output_tail(result: subprocess.CompletedProcess[str], *, max_lines: int = 20) -> str: + sections = [] + for stream_name, text in (("stderr", result.stderr), ("stdout", result.stdout)): + lines = (text or "").strip().splitlines() + if lines: + sections.append(f"{stream_name} tail:") + sections.extend(lines[-max_lines:]) + return "\n".join(sections) + + +def _signal_lmms_eval_process_group( + process: subprocess.Popen[str], signal_number: int +) -> None: + try: + if os.name == "posix": + os.killpg(process.pid, signal_number) + else: + process.send_signal(signal_number) + except ProcessLookupError: + pass + + +def _lmms_eval_process_group_exists(process: subprocess.Popen[str]) -> bool: + if os.name != "posix": + return process.poll() is None + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _run_lmms_eval_process( + argv: list[str], + *, + cwd: str, + env: Mapping[str, str], + timeout: float | None, +) -> subprocess.CompletedProcess[str]: + process = subprocess.Popen( + argv, + cwd=cwd, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=os.name == "posix", + ) + try: + stdout, stderr = process.communicate(timeout=timeout) + except subprocess.TimeoutExpired as error: + _signal_lmms_eval_process_group(process, signal.SIGTERM) + try: + stdout, stderr = process.communicate( + timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS + ) + except subprocess.TimeoutExpired: + _signal_lmms_eval_process_group(process, signal.SIGKILL) + stdout, stderr = process.communicate() + else: + if _lmms_eval_process_group_exists(process): + _signal_lmms_eval_process_group(process, signal.SIGKILL) + raise subprocess.TimeoutExpired( + argv, + timeout, + output=stdout if stdout is not None else error.output, + stderr=stderr if stderr is not None else error.stderr, + ) from error + return subprocess.CompletedProcess(argv, process.returncode, stdout, stderr) + + +def _downstream_evaluation( + config: dict[str, Any], + node: CompiledPostMIPNode, + source, + execution_identity: str, +) -> dict[str, Any]: + if source.artifact_kind is not ArtifactKind.CHECKPOINT: + raise ValueError("downstream_evaluation requires materialized checkpoint artifacts") + settings = dict(node.config.get("config") or {}) + output_root = ( + _execution_root(config, node, execution_identity) + / "raw" + / source.architecture_id + / "lmms_eval" + ) + output = output_root / f"attempt_{uuid.uuid4().hex}" + output.mkdir(parents=True, exist_ok=True) + argv, env, timeout = _lmms_eval_command( + settings, + checkpoint=str(source.artifact["checkpoint"]), + output_path=output, + ) + command_path = output / "command.json" + _atomic_json( + command_path, + { + "argv": argv, + "env_overrides": sorted(str(key) for key in dict(settings.get("env") or {})), + "timeout": timeout, + }, + ) + # Campaign config controls the executable and arguments, but subprocess receives + # an argv list directly; no shell parsing is involved. + result = _run_lmms_eval_process( + argv, + cwd=str(output), + env=env, + timeout=timeout, + ) + stream_paths = _write_lmms_eval_streams(output, result) + if result.returncode: + tail = _lmms_eval_output_tail(result) + raise RuntimeError( + f"lmms-eval failed with exit code {result.returncode}" + + (f": {tail}" if tail else "") + ) + try: + payload, result_path = _lmms_eval_result_payload(output) + except FileNotFoundError as error: + tail = _lmms_eval_output_tail(result) + raise FileNotFoundError(str(error) + (f": {tail}" if tail else "")) from error + sample_counts = _validate_lmms_eval_completion( + payload, _configured_lmms_eval_tasks(settings) + ) + metrics = _flatten_lmms_eval_metrics(payload) + if not metrics: + raise RuntimeError(f"lmms-eval result has no numeric task metrics: {result_path}") + summary_path = output / "summary.json" + _atomic_json( + summary_path, + { + "architecture_id": source.architecture_id, + "checkpoint": source.artifact["checkpoint"], + "metrics": metrics, + "result_path": str(result_path), + "sample_counts": sample_counts, + }, + ) + return { + "metrics": metrics, + "result_path": str(summary_path), + "raw_result_path": str(result_path), + "command_path": str(command_path), + **stream_paths, + } + + def _post_mip_kd_settings( config: Mapping[str, Any], node_settings: Mapping[str, Any], @@ -582,6 +1190,8 @@ def _run_candidate( result = _evaluate(config, node, source, execution_identity) elif node.node_type == "aiperf": result = _aiperf(config, node, source, execution_identity) + elif node.node_type == "downstream_evaluation": + result = _downstream_evaluation(config, node, source, execution_identity) elif node.node_type == "global_kd": result = _global_kd(config, node, source, execution_identity) else: @@ -649,7 +1259,7 @@ def run_post_mip_node_shard( row = _run_candidate(config, node, ledger, revision_id, execution_identity) row = {**row, "execution_identity": execution_identity} except Exception as error: - timed_out = node.node_type == "aiperf" and isinstance( + timed_out = node.node_type in {"aiperf", "downstream_evaluation"} and isinstance( error, (subprocess.TimeoutExpired, TimeoutError) ) row = { @@ -661,12 +1271,14 @@ def run_post_mip_node_shard( **_exception_diagnostics(error), } if timed_out: - timeout_field = ( - "benchmark_timeout" - if isinstance(error, subprocess.TimeoutExpired) - else "readiness_timeout" + timeout_field = "benchmark_timeout" + if node.node_type == "downstream_evaluation": + timeout_field = "timeout_seconds" + elif not isinstance(error, subprocess.TimeoutExpired): + timeout_field = "readiness_timeout" + default_timeout = 3600 if node.node_type == "downstream_evaluation" else ( + 600 if timeout_field == "benchmark_timeout" else 1200 ) - default_timeout = 600 if timeout_field == "benchmark_timeout" else 1200 row["timeout_seconds"] = float( getattr(error, "timeout", None) or (node.config.get("config") or {}).get(timeout_field, default_timeout) diff --git a/puzzletron_setup/bundle.py b/puzzletron_setup/bundle.py index ba091af39a8..3ee3e8f0936 100644 --- a/puzzletron_setup/bundle.py +++ b/puzzletron_setup/bundle.py @@ -346,6 +346,13 @@ def _post_mip_flows( if smoke: config["minimum_request_count"] = 4 config["requests_per_concurrency"] = 1 + elif node_type == "downstream_evaluation": + config.setdefault("model", "vllm") + config.setdefault("batch_size", 1) + config.setdefault("log_samples", True) + config.setdefault("topology", deepcopy(dict(default_serving_topology))) + if smoke: + config["limit"] = min(int(config.get("limit", 8) or 8), 8) elif node_type == "global_kd": config.setdefault("automodel", {})["parallel"] = _parallel(global_kd_mesh) config["local_batch_size"] = _aligned_batch_size( @@ -816,7 +823,7 @@ def _dynamic_stage_entries( entry.update(resource="cpu", partition=cpu_partition) if node_type == "evaluation": entry["parallel"] = dict(common) - elif node_type == "aiperf": + elif node_type in {"aiperf", "downstream_evaluation"}: config = _mapping(node.get("config")) entry["parallel"] = _serving_parallel(_mapping(config.get("topology"))) elif node_type == "materialize": diff --git a/puzzletron_setup/v2/parallel_validation.py b/puzzletron_setup/v2/parallel_validation.py index 0ce29e95297..c22058fd9f5 100644 --- a/puzzletron_setup/v2/parallel_validation.py +++ b/puzzletron_setup/v2/parallel_validation.py @@ -50,7 +50,9 @@ "vllm_stats", } ) -_CANDIDATE_POST_MIP_TYPES = frozenset({"evaluation", "global_kd", "aiperf"}) +_CANDIDATE_POST_MIP_TYPES = frozenset( + {"evaluation", "global_kd", "aiperf", "downstream_evaluation"} +) @dataclass(frozen=True) diff --git a/puzzletron_setup/v2/post_mip.py b/puzzletron_setup/v2/post_mip.py index 14d34f6b8dc..067dc18b459 100644 --- a/puzzletron_setup/v2/post_mip.py +++ b/puzzletron_setup/v2/post_mip.py @@ -29,9 +29,10 @@ "materialize", "evaluation", "aiperf", + "downstream_evaluation", "global_kd", ) -RESERVED_NODE_TYPES = ("ptq", "downstream_evaluation") +RESERVED_NODE_TYPES = ("ptq",) @dataclass(frozen=True) diff --git a/puzzletron_setup/v2/validation.py b/puzzletron_setup/v2/validation.py index f5c8dff0b4a..d0cf216d693 100644 --- a/puzzletron_setup/v2/validation.py +++ b/puzzletron_setup/v2/validation.py @@ -298,7 +298,7 @@ def validate_state(state: WizardState) -> tuple[ValidationIssue, ...]: for stage_id, node in post_mip_nodes.items(): node_type = str(node.get("type", "")) config = _mapping(node.get("config")) - if node_type == "aiperf": + if node_type in {"aiperf", "downstream_evaluation"}: topology = _mapping(config.get("topology")) if topology: issues.extend( diff --git a/puzzletron_setup/v2/wizard.py b/puzzletron_setup/v2/wizard.py index c26372ffa3b..cca7b77206f 100644 --- a/puzzletron_setup/v2/wizard.py +++ b/puzzletron_setup/v2/wizard.py @@ -3668,7 +3668,7 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context strategy=_post_mip_strategy(node), batch=1, ) - elif node.node_type == "aiperf": + elif node.node_type in {"aiperf", "downstream_evaluation"}: node_preview["resources"] = { "instances": int(session.state.get_field("infrastructure.gpus_per_node", 8)), "topology": node.config.get("topology", {}), @@ -3765,7 +3765,7 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context "aiperf", "global_kd", ("PTQ — unavailable", "unavailable"), - ("Downstream evaluation — unavailable", "unavailable"), + "downstream_evaluation", ], default="evaluation", ) @@ -3839,6 +3839,18 @@ def post_mip_section(session: WizardSession, resolver: DefaultsResolver, context "concurrency": list(configured["concurrency"]), "benchmark_timeout": 900, } + elif node_type == "downstream_evaluation": + configured = _downstream_evaluation_setting_prompt( + session, + f"post_mip.{run_id}.{node_id}", + {}, + inventory=context["model"].inventory, + pruning=_mapping_copy(_pruning_payload(session.state)), + stage_id=f"post.{run_id}.{node_id}", + ) + if configured is BACK: + return False + config = configured elif node_type == "global_kd": max_steps = session.integer( f"post_mip.{run_id}.{node_id}.max_steps", @@ -4106,6 +4118,83 @@ def validate_concurrency(value: str) -> bool | str: return values +def _downstream_evaluation_setting_prompt( + session: WizardSession, + prefix: str, + defaults: Mapping[str, Any], + *, + inventory: Any, + pruning: Mapping[str, Any], + stage_id: str, +) -> Any: + """Ask lmms-eval task settings and the vLLM topology used to run them.""" + + def validate_tasks(value: str) -> bool | str: + tasks = [item.strip() for item in value.split(",") if item.strip()] + return True if tasks else "Enter at least one lmms-eval task." + + raw_default_tasks = defaults.get("tasks", ("ifeval", "gsm8k")) + default_tasks = ( + str(raw_default_tasks) + if isinstance(raw_default_tasks, str) + else ",".join(str(item) for item in raw_default_tasks) + ) + default_model_args = _mapping_copy(defaults.get("model_args")) + tasks = session.text( + f"{prefix}.tasks", + "lmms-eval tasks (comma-separated):", + default=default_tasks, + validate=validate_tasks, + ) + if tasks is BACK: + return BACK + limit = session.integer( + f"{prefix}.limit", + "lmms-eval sample limit:", + default=int(defaults.get("limit", 128)), + minimum=1, + ) + batch_size = session.integer( + f"{prefix}.batch_size", + "lmms-eval batch size:", + default=int(defaults.get("batch_size", 1)), + minimum=1, + ) + timeout = session.integer( + f"{prefix}.timeout_seconds", + "Per-candidate lmms-eval timeout (seconds):", + default=int(defaults.get("timeout_seconds", 3600)), + minimum=1, + ) + if BACK in (limit, batch_size, timeout): + return BACK + topology = _vllm_topology_prompt( + session, + f"{prefix}.topology", + _mapping_copy(defaults.get("topology")), + inventory=inventory, + pruning=pruning, + stage_id=stage_id, + label_prefix="lmms-eval vLLM", + ) + if topology is BACK: + return BACK + return { + "model": str(defaults.get("model", "vllm")), + "tasks": [item.strip() for item in str(tasks).split(",") if item.strip()], + "limit": int(limit), + "batch_size": int(batch_size), + "log_samples": bool(defaults.get("log_samples", True)), + "topology": topology, + "model_args": { + **default_model_args, + "dtype": default_model_args.get("dtype", "bfloat16"), + "gpu_memory_utilization": default_model_args.get("gpu_memory_utilization", 0.85), + }, + "timeout_seconds": int(timeout), + } + + def _configure_dynamic_resources( session: WizardSession, editor: PostMIPFlowEditor, @@ -4158,7 +4247,7 @@ def _configure_dynamic_resources( "resource": "gpu", "gpus_per_node": gpus_per_node, } - if node.node_type == "aiperf": + if node.node_type in {"aiperf", "downstream_evaluation"}: topology = _mapping_copy(node.config.get("topology")) allocation_mesh = vllm_topology_to_mesh(topology) entry["parallel"] = { diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index a9e153e5260..90dba4cd26f 100644 --- a/puzzletron_setup/wizard.py +++ b/puzzletron_setup/wizard.py @@ -39,6 +39,10 @@ _MESH_KEYS = ("tp", "cp", "pp", "dp_shard", "dp_replicate", "ep") _DEFAULT_MIP_SOLUTION_COUNT = 3 _DEFAULT_HOMOGENEOUS_SOLUTIONS_PER_SCENARIO = 8 +_DOWNSTREAM_EVALUATION_METRICS_BY_TASK = { + "gsm8k": ("exact_match_strict-match",), + "ifeval": ("prompt_level_strict_acc_none",), +} def _default(state: AnswerState, section: str, key: str, fallback: Any) -> Any: @@ -733,6 +737,128 @@ def _ask_aiperf_config( return config +def _ask_downstream_evaluation_config( + prompts: PromptSession, + *, + detailed: bool, + moe: bool, + defaults: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + """Ask for lmms-eval task and vLLM settings.""" + + defaults = defaults or {} + tasks = prompts.text( + "lmms-eval tasks (comma-separated):", + default=str(defaults.get("tasks", "ifeval,gsm8k")), + ) + limit = prompts.integer( + "lmms-eval sample limit:", + default=int(defaults.get("limit", 128)), + ) + batch_size = prompts.integer( + "lmms-eval batch size:", + default=int(defaults.get("batch_size", 1)), + ) + topology_defaults = dict(defaults.get("topology") or {}) + checkpoint = prompts.checkpoint() + while True: + tp = prompts.integer( + "lmms-eval vLLM tensor parallel (TP):", + default=int(topology_defaults.get("tensor_parallel_size", 1)), + ) + pp = prompts.integer( + "lmms-eval vLLM pipeline parallel (PP):", + default=int(topology_defaults.get("pipeline_parallel_size", 1)), + ) + dp = prompts.integer( + "lmms-eval vLLM data parallel (DP):", + default=int(topology_defaults.get("data_parallel_size", 1)), + ) + prefill_cp = prompts.integer( + "lmms-eval vLLM prefill context parallel (CP):", + default=int(topology_defaults.get("prefill_context_parallel_size", 1)), + ) + decode_cp = prompts.integer( + "lmms-eval vLLM decode context parallel (CP):", + default=int(topology_defaults.get("decode_context_parallel_size", 1)), + ) + enable_expert_parallel = ( + prompts.confirm( + ( + "Enable lmms-eval vLLM expert parallelism? " + "vLLM effective EP is TP * DP." + ), + default=bool( + topology_defaults.get("enable_expert_parallel") + or int(topology_defaults.get("expert_parallel_size", 1)) > 1 + ), + ) + if moe + else False + ) + dimensions = { + "TP": tp, + "PP": pp, + "DP": dp, + "prefill CP": prefill_cp, + "decode CP": decode_cp, + } + error = None + if any(int(value) < 1 for value in dimensions.values()): + error = f"lmms-eval vLLM parallel dimensions must be positive: {dimensions}" + elif decode_cp > tp or tp % decode_cp: + error = f"lmms-eval vLLM decode CP={decode_cp} must divide TP={tp}." + if error is None: + break + print(error) + prompts.rewind(checkpoint) + topology = { + "tensor_parallel_size": tp, + "pipeline_parallel_size": pp, + "prefill_context_parallel_size": prefill_cp, + "decode_context_parallel_size": decode_cp, + "data_parallel_size": dp, + "enable_expert_parallel": bool(enable_expert_parallel), + "distributed_executor_backend": "mp", + "gpu_group_size": tp * pp * prefill_cp * dp, + } + timeout = ( + prompts.integer( + "Per-candidate lmms-eval timeout (seconds):", + default=int(defaults.get("timeout_seconds", 3600)), + ) + if detailed + else int(defaults.get("timeout_seconds", 3600)) + ) + return { + "model": "vllm", + "tasks": [item.strip() for item in str(tasks).split(",") if item.strip()], + "limit": int(limit), + "batch_size": int(batch_size), + "log_samples": True, + "topology": topology, + "model_args": { + "dtype": "bfloat16", + "gpu_memory_utilization": 0.85, + }, + "timeout_seconds": int(timeout), + } + + +def _downstream_evaluation_metric_suggestions(node_id: str, config: Mapping[str, Any]) -> list[str]: + """Return filter metric names produced by the downstream-evaluation runner.""" + + suggestions = [] + tasks = config.get("tasks") or () + if isinstance(tasks, str): + tasks = [item.strip() for item in tasks.split(",") if item.strip()] + for task in tasks: + task_name = str(task).strip() + for metric in _DOWNSTREAM_EVALUATION_METRICS_BY_TASK.get(task_name, ()): + suggestions.append(f"{node_id}.{task_name}.{metric}") + return suggestions + + def _default_flow( run_id: str, run: Mapping[str, Any], @@ -916,7 +1042,7 @@ def _custom_flow( "global_kd", "manual_filter", ("PTQ (reserved; not executable yet)", "ptq"), - ("Downstream evaluation (reserved; not executable yet)", "downstream_evaluation"), + "downstream_evaluation", ], default="filter", ) @@ -963,9 +1089,18 @@ def _custom_flow( runtime=runtime, ) available_metrics.append(f"{node_id}.request_throughput") + elif node_type == "downstream_evaluation": + node["config"] = _ask_downstream_evaluation_config( + prompts, + detailed=detailed, + moe=moe, + ) + available_metrics.extend( + _downstream_evaluation_metric_suggestions(node_id, node["config"]) + ) elif node_type == "global_kd": node["config"] = {"max_steps": prompts.integer("Global KD steps:", default=128)} - elif node_type in {"ptq", "downstream_evaluation"}: + elif node_type == "ptq": print( f"{node_type} records the reserved interface, but current orchestration " "validation will report it as unimplemented." @@ -1164,6 +1299,11 @@ def post_mip_gpus_per_instance(node_type: str, default: int) -> int: post_mip_gpus_per_instance("aiperf", 1), post_mip_instances("aiperf", sharded_workers, sharded_workers), ), + ( + "downstream eval", + post_mip_gpus_per_instance("downstream_evaluation", 1), + post_mip_instances("downstream_evaluation", sharded_workers, sharded_workers), + ), ( "evaluation", _mesh_product(common), diff --git a/tests/unit/torch/puzzletron/test_orchestration_compiler.py b/tests/unit/torch/puzzletron/test_orchestration_compiler.py index 1a0c1c085cc..1d23353dd4e 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_compiler.py +++ b/tests/unit/torch/puzzletron/test_orchestration_compiler.py @@ -175,3 +175,62 @@ def test_post_mip_compiler_topologically_orders_serialized_nodes() -> None: stages = _post_mip_stage_metadata(config) assert [stage["node_id"] for stage in stages] == ["initial", "final_eval", "best"] + + +def test_compile_campaign_plan_allocates_downstream_evaluation_from_vllm_topology( + tmp_configs, +) -> None: + experiment_path, runner_path, execution_path = tmp_configs + experiment = yaml.safe_load(experiment_path.read_text()) + experiment.update( + { + "mip": {"runs": {"runtime": {}}}, + "post_mip": { + "flows": { + "runtime": { + "source": {"run": "runtime"}, + "nodes": { + "materialized": {"type": "materialize"}, + "lmms_eval": { + "type": "downstream_evaluation", + "input": "materialized", + "config": { + "tasks": ["ifeval"], + "topology": { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 2, + "data_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": False, + "gpu_group_size": 8, + }, + }, + }, + }, + } + } + }, + } + ) + experiment_path.write_text(yaml.safe_dump(experiment)) + execution = yaml.safe_load(execution_path.read_text()) + execution["execution"]["stages"]["post.runtime.lmms_eval"] = { + "strategy": "sharded", + "instances": 2, + } + execution_path.write_text(yaml.safe_dump(execution)) + + plan = compile_campaign_plan( + experiment_config_path=experiment_path, + runner=load_runner_config(runner_path), + execution=load_execution_config(execution_path), + stage_filter="post.runtime.lmms_eval", + ) + node = plan.stages[0] + + assert node.stage_id == "post.runtime.lmms_eval" + assert node.parents == ("post.runtime.materialized",) + assert node.gpus_per_instance == 8 + assert node.instances == 2 + assert node.nodes == 2 diff --git a/tests/unit/torch/puzzletron/test_orchestration_controller.py b/tests/unit/torch/puzzletron/test_orchestration_controller.py index a0390c30fb6..2a42ca5f181 100644 --- a/tests/unit/torch/puzzletron/test_orchestration_controller.py +++ b/tests/unit/torch/puzzletron/test_orchestration_controller.py @@ -14,7 +14,7 @@ load_execution_config, load_runner_config, ) -from puzzletron_orchestrator.controller import dry_run_plan +from puzzletron_orchestrator.controller import _stage_dashboard_display_name, dry_run_plan from puzzletron_orchestrator.executors.baremetal import GpuLeaseManager from puzzletron_orchestrator.schema import ( AttemptSpec, diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index aa12508ecb6..7b650284a51 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -16,12 +16,16 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import json +import signal +import subprocess from pathlib import Path from types import SimpleNamespace from omegaconf import OmegaConf from modelopt.torch.puzzletron.post_mip import runner +from modelopt.torch.puzzletron.post_mip.records import ArtifactKind from modelopt.torch.puzzletron.post_mip.runner import ( _exception_diagnostics, _needs_puzzletron_process_group, @@ -184,3 +188,479 @@ def fake_run_aiperf_sweep(checkpoint, **settings): assert "requests_per_concurrency" not in captured assert "best_selection_mode" not in captured assert result["metrics"] == {} + + +def test_lmms_eval_command_maps_checkpoint_and_vllm_topology(tmp_path): + argv, env, timeout = runner._lmms_eval_command( + { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "batch_size": 2, + "limit": 8, + "cache_dir": tmp_path / "cache", + "timeout_seconds": 123, + "topology": { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 2, + "data_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": False, + "gpu_group_size": 8, + }, + "model_args": {"dtype": "bfloat16"}, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + model_args = argv[argv.index("--model_args") + 1] + assert argv[:5] == ["python", "-m", "lmms_eval", "--model", "vllm"] + assert argv[argv.index("--tasks") + 1] == "ifeval,gsm8k" + assert argv[argv.index("--batch_size") + 1] == "2" + assert argv[argv.index("--limit") + 1] == "8" + assert "model=/ckpts/candidate" in model_args + assert "tensor_parallel_size=4" in model_args + assert "pipeline_parallel_size=2" in model_args + assert "gpu_group_size" not in model_args + assert env["LMMS_EVAL_HOME"] == str(tmp_path / "cache") + assert timeout == 123 + + +def test_lmms_eval_command_rejects_reserved_model_args(tmp_path): + cases = ( + ({"model": "/ckpts/wrong"}, "model"), + ("dtype=bfloat16,tensor_parallel_size=1", "tensor_parallel_size"), + ) + for model_args, expected in cases: + try: + runner._lmms_eval_command( + { + "tasks": ["ifeval"], + "topology": {"gpu_group_size": 1}, + "model_args": model_args, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + except ValueError as error: + message = str(error) + else: + raise AssertionError("expected reserved lmms-eval model_args to fail") + + assert "reserved lmms-eval model arguments" in message + assert expected in message + + +def test_lmms_eval_command_rejects_reserved_extra_args(tmp_path): + cases = ( + (["--tasks", "gsm8k"], "--tasks"), + ("--output_path /tmp/other", "--output_path"), + (["--model_args=model=/ckpts/wrong"], "--model_args"), + ) + for extra_args, expected in cases: + try: + runner._lmms_eval_command( + { + "tasks": ["ifeval"], + "topology": {"gpu_group_size": 1}, + "extra_args": extra_args, + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + except ValueError as error: + message = str(error) + else: + raise AssertionError("expected reserved lmms-eval extra_args to fail") + + assert "reserved lmms-eval flags" in message + assert expected in message + + +def test_lmms_eval_timeout_terminates_process_group(monkeypatch, tmp_path): + created = [] + signals = [] + + class FakeProcess: + pid = 1234 + returncode = None + + def __init__(self): + self.communicate_timeouts = [] + + def communicate(self, timeout=None): + self.communicate_timeouts.append(timeout) + if len(self.communicate_timeouts) == 1: + raise subprocess.TimeoutExpired( + ["python", "-m", "lmms_eval"], + timeout, + output="partial stdout", + stderr="partial stderr", + ) + self.returncode = -signal.SIGTERM + return "partial stdout", "partial stderr" + + def fake_popen(argv, **kwargs): + process = FakeProcess() + created.append((argv, kwargs, process)) + return process + + def fake_killpg(pid, signal_number): + if signal_number == 0: + raise ProcessLookupError + signals.append((pid, signal_number)) + + monkeypatch.setattr(runner.subprocess, "Popen", fake_popen) + monkeypatch.setattr(runner.os, "killpg", fake_killpg) + + try: + runner._run_lmms_eval_process( + ["python", "-m", "lmms_eval"], + cwd=str(tmp_path), + env={}, + timeout=7.0, + ) + except subprocess.TimeoutExpired as error: + assert error.timeout == 7.0 + assert error.output == "partial stdout" + assert error.stderr == "partial stderr" + else: + raise AssertionError("expected lmms-eval timeout to be raised") + + argv, kwargs, process = created[0] + assert argv == ["python", "-m", "lmms_eval"] + assert kwargs["start_new_session"] is True + assert process.communicate_timeouts == [ + 7.0, + runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, + ] + assert signals == [(1234, signal.SIGTERM)] + + +def test_lmms_eval_timeout_kills_remaining_process_group(monkeypatch, tmp_path): + signals = [] + + class FakeProcess: + pid = 3456 + returncode = None + + def communicate(self, timeout=None): + if timeout == 7.0: + raise subprocess.TimeoutExpired( + ["python", "-m", "lmms_eval"], + timeout, + output="partial stdout", + stderr="partial stderr", + ) + self.returncode = -signal.SIGTERM + return "partial stdout", "partial stderr" + + def fake_killpg(pid, signal_number): + if signal_number != 0: + signals.append((pid, signal_number)) + + monkeypatch.setattr(runner.subprocess, "Popen", lambda *args, **kwargs: FakeProcess()) + monkeypatch.setattr(runner.os, "killpg", fake_killpg) + + try: + runner._run_lmms_eval_process( + ["python", "-m", "lmms_eval"], + cwd=str(tmp_path), + env={}, + timeout=7.0, + ) + except subprocess.TimeoutExpired: + pass + else: + raise AssertionError("expected lmms-eval timeout to be raised") + + assert signals == [(3456, signal.SIGTERM), (3456, signal.SIGKILL)] + + +def test_lmms_eval_timeout_kills_stubborn_process_group(monkeypatch, tmp_path): + signals = [] + + class FakeProcess: + pid = 5678 + returncode = None + + def __init__(self): + self.communicate_timeouts = [] + + def communicate(self, timeout=None): + self.communicate_timeouts.append(timeout) + if len(self.communicate_timeouts) < 3: + raise subprocess.TimeoutExpired( + ["python", "-m", "lmms_eval"], + timeout, + output="partial stdout", + stderr="partial stderr", + ) + self.returncode = -signal.SIGKILL + return "partial stdout", "partial stderr" + + process = FakeProcess() + monkeypatch.setattr(runner.subprocess, "Popen", lambda *args, **kwargs: process) + monkeypatch.setattr( + runner.os, + "killpg", + lambda pid, signal_number: signals.append((pid, signal_number)), + ) + + try: + runner._run_lmms_eval_process( + ["python", "-m", "lmms_eval"], + cwd=str(tmp_path), + env={}, + timeout=7.0, + ) + except subprocess.TimeoutExpired as error: + assert error.output == "partial stdout" + assert error.stderr == "partial stderr" + else: + raise AssertionError("expected lmms-eval timeout to be raised") + + assert process.communicate_timeouts == [ + 7.0, + runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, + None, + ] + assert signals == [(5678, signal.SIGTERM), (5678, signal.SIGKILL)] + + +def test_downstream_evaluation_runs_lmms_eval_and_flattens_metrics(monkeypatch, tmp_path): + captured = {} + + def fake_run(argv, *, cwd, env, timeout): + del env, timeout + captured["argv"] = argv + output = Path(cwd) / "nested" + output.mkdir(parents=True) + (output / "results.json").write_text( + json.dumps( + { + "results": { + "ifeval": {"prompt_level_strict_acc,none": 0.5}, + "gsm8k": {"exact_match,strict-match": 0.75}, + }, + "group_subtasks": {"ifeval": [], "gsm8k": []}, + "n-samples": { + "ifeval": {"original": 541, "effective": 4}, + "gsm8k": {"original": 1319, "effective": 4}, + }, + } + ) + ) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "limit": 4, + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + result = runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, + node, + source, + "execution", + ) + + assert captured["argv"][:3] == ["python", "-m", "lmms_eval"] + assert result["metrics"] == { + "gsm8k.exact_match_strict-match": 0.75, + "ifeval.prompt_level_strict_acc_none": 0.5, + } + assert Path(result["result_path"]).is_file() + assert Path(result["raw_result_path"]).name == "results.json" + summary = json.loads(Path(result["result_path"]).read_text()) + assert summary["sample_counts"] == {"gsm8k": 4.0, "ifeval": 4.0} + + +def test_lmms_eval_completion_validates_resolved_task_expansion(): + sample_counts = runner._validate_lmms_eval_completion( + { + "results": { + "arc_challenge": {"acc,none": 0.25}, + "hellaswag": {"acc_norm,none": 0.5}, + }, + "group_subtasks": { + "leaderboard": ["arc_challenge", "hellaswag"], + "arc_challenge": [], + "hellaswag": [], + }, + "n-samples": { + "arc_challenge": {"original": 1172, "effective": 8}, + "hellaswag": {"original": 10042, "effective": 8}, + }, + }, + ("leaderboard",), + ) + + assert sample_counts == {"arc_challenge": 8.0, "hellaswag": 8.0} + + +def test_downstream_evaluation_rejects_missing_configured_task(monkeypatch, tmp_path): + def fake_run(argv, *, cwd, env, timeout): + del env, timeout + output = Path(cwd) + (output / "results.json").write_text( + json.dumps( + { + "results": {"ifeval": {"prompt_level_strict_acc,none": 0.5}}, + "group_subtasks": {"ifeval": []}, + "n-samples": {"ifeval": {"original": 541, "effective": 4}}, + } + ) + ) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + try: + runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, node, source, "execution" + ) + except RuntimeError as error: + message = str(error) + else: + raise AssertionError("expected incomplete lmms-eval result to fail") + + assert "missing configured task results" in message + assert "gsm8k" in message + + +def test_downstream_evaluation_rejects_zero_sample_task(monkeypatch, tmp_path): + def fake_run(argv, *, cwd, env, timeout): + del env, timeout + output = Path(cwd) + (output / "results.json").write_text( + json.dumps( + { + "results": { + "ifeval": {"prompt_level_strict_acc,none": 0.5}, + "gsm8k": {"exact_match,strict-match": 0.75}, + }, + "group_subtasks": {"ifeval": [], "gsm8k": []}, + "n-samples": { + "ifeval": {"original": 541, "effective": 4}, + "gsm8k": {"original": 1319, "effective": 0}, + }, + } + ) + ) + return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval", "gsm8k"], + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + try: + runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, node, source, "execution" + ) + except RuntimeError as error: + message = str(error) + else: + raise AssertionError("expected zero-sample lmms-eval result to fail") + + assert "zero effective samples" in message + assert "gsm8k" in message + + +def test_downstream_evaluation_reports_lmms_eval_output_when_results_are_missing( + monkeypatch, tmp_path +): + def fake_run(argv, *, cwd, env, timeout): + del cwd, env, timeout + return subprocess.CompletedProcess( + argv, + 0, + stdout="Saving results aggregated\nCould not save results aggregated\n", + stderr="", + ) + + monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) + node = SimpleNamespace( + node_id="lmms_eval", + flow_id="runtime", + stage_id="post.runtime.lmms_eval", + config={ + "config": { + "command_prefix": ["python", "-m", "lmms_eval"], + "tasks": ["ifeval"], + "limit": 1, + "topology": {"gpu_group_size": 1}, + } + }, + ) + source = SimpleNamespace( + architecture_id="architecture", + artifact_kind=ArtifactKind.CHECKPOINT, + artifact={"checkpoint": str(tmp_path / "checkpoint")}, + ) + + try: + runner._downstream_evaluation( + {"puzzle_dir": str(tmp_path)}, node, source, "execution" + ) + except FileNotFoundError as error: + message = str(error) + else: + raise AssertionError("expected missing lmms-eval results to fail") + + assert "lmms-eval wrote no JSON results" in message + assert "stdout tail:" in message + assert "Could not save results aggregated" in message + stream_root = ( + tmp_path + / "artifacts/post_mip/nodes/lmms_eval/executions/execution/raw/architecture/lmms_eval" + ) + assert list(stream_root.glob("attempt_*/stdout.txt")) diff --git a/tests/unit/torch/puzzletron/test_setup_bundle.py b/tests/unit/torch/puzzletron/test_setup_bundle.py index f453aa55d49..5f4f264ea65 100644 --- a/tests/unit/torch/puzzletron/test_setup_bundle.py +++ b/tests/unit/torch/puzzletron/test_setup_bundle.py @@ -34,6 +34,7 @@ _ask_mesh, _ask_mip, _default_flow, + _downstream_evaluation_metric_suggestions, _resource_rows, ) @@ -549,6 +550,82 @@ def test_render_execution_uses_common_mesh_for_post_mip_evaluation_only() -> Non assert execution["post.run.materialized"]["instances"] == 1 +def test_downstream_evaluation_metric_suggestions_match_runner_keys() -> None: + assert _downstream_evaluation_metric_suggestions( + "lmms_eval", + {"tasks": ["ifeval", "gsm8k", "custom_task"]}, + ) == [ + "lmms_eval.ifeval.prompt_level_strict_acc_none", + "lmms_eval.gsm8k.exact_match_strict-match", + ] + assert _downstream_evaluation_metric_suggestions( + "lmms_eval", + {"tasks": "gsm8k,ifeval"}, + ) == [ + "lmms_eval.gsm8k.exact_match_strict-match", + "lmms_eval.ifeval.prompt_level_strict_acc_none", + ] + + +def test_render_execution_uses_vllm_mesh_for_post_mip_downstream_evaluation() -> None: + state = { + "answers": { + "infrastructure": { + "gpus_per_node": 8, + "workers": {"pool": 8, "sharded": 8}, + "runner": {"slurm": {}}, + "meshes": { + "common": {"tp": 1, "cp": 1, "pp": 1, "dp_shard": 2, "ep": 1}, + "bypass": {"tp": 1, "cp": 1, "pp": 1, "dp_shard": 1, "ep": 1}, + "global_kd": {"tp": 1, "cp": 1, "pp": 1, "dp_shard": 1, "ep": 1}, + }, + } + }, + } + experiment = { + "embedding_pruning": {"widths": []}, + "vllm_stats": {"runtime_stats": {"topology": {"gpu_group_size": 1}}}, + "post_mip": { + "flows": { + "run": { + "nodes": { + "materialized": {"type": "materialize"}, + "lmms_eval": { + "type": "downstream_evaluation", + "input": "materialized", + "config": { + "topology": { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 2, + "data_parallel_size": 1, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": False, + "gpu_group_size": 8, + } + }, + }, + } + } + } + }, + } + + stages = render_execution(state, experiment, "production")["execution"]["stages"] + + assert stages["post.run.lmms_eval"]["strategy"] == "sharded" + assert stages["post.run.lmms_eval"]["instances"] == 8 + assert stages["post.run.lmms_eval"]["parallel"] == { + "tp": 4, + "cp": 1, + "pp": 2, + "ep": 1, + "dp_shard": 1, + "dp_replicate": 1, + "sequence_parallel": False, + } + + def test_render_execution_caps_post_mip_workers_at_upstream_top_k() -> None: common = { "tp": 1, diff --git a/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py b/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py index 459011a3bc3..247f0e85ebf 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_post_mip.py @@ -1,7 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from puzzletron_setup.v2.post_mip import recommended_flow +from collections import OrderedDict + +from puzzletron_setup.v2.post_mip import FlowDraft, NodeDraft, PostMIPFlowEditor, recommended_flow def test_recommended_flow_propagates_aiperf_sweep_selection_mode(): @@ -45,3 +47,32 @@ def test_recommended_flow_accepts_a_single_concurrency_value(): assert flow.nodes["serving"].config["concurrency"] == [2] assert flow.nodes["fastest"].selector["best_selection_mode"] == "individual_best" + + +def test_downstream_evaluation_node_is_configurable_after_materialization(): + flow = FlowDraft( + "runtime", + "runtime", + nodes=OrderedDict( + ( + ("materialized", NodeDraft("materialized", "materialize")), + ( + "lmms_eval", + NodeDraft( + "lmms_eval", + "downstream_evaluation", + input_id="materialized", + config={"tasks": ["ifeval"]}, + ), + ), + ) + ), + ) + editor = PostMIPFlowEditor({"runtime": {}}) + editor.add_flow(flow) + + review = editor.review("runtime") + + assert review.node_order == ("materialized", "lmms_eval") + assert review.parents["lmms_eval"] == ("materialized",) + assert review.artifacts["lmms_eval"] == "checkpoint" diff --git a/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py b/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py index 049905630a8..78be65d9916 100644 --- a/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py +++ b/tests/unit/torch/puzzletron/test_setup_v2_state_validation.py @@ -181,6 +181,38 @@ def test_persisted_post_mip_aiperf_topology_is_candidate_checked(tmp_path): assert "expert counts [48, 64]" in messages +def test_persisted_post_mip_downstream_eval_topology_is_candidate_checked(tmp_path): + state = _state(tmp_path) + topology = { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 1, + "data_parallel_size": 8, + "prefill_context_parallel_size": 1, + "decode_context_parallel_size": 1, + "enable_expert_parallel": True, + "gpu_group_size": 32, + } + state.set_collection( + "post_mip_flows", + { + "run": { + "source": {"run": "run"}, + "nodes": { + "lmms_eval": { + "type": "downstream_evaluation", + "config": {"tasks": ["ifeval"], "topology": topology}, + } + }, + } + }, + ) + + messages = _messages(state) + + assert "effective EP=32 (TP * DP)" in messages + assert "expert counts [48, 64]" in messages + + def test_persisted_vllm_measurement_topology_is_candidate_checked(tmp_path): state = _state(tmp_path) state.set_collection(