From ec0f281f430eeccc28a783853eafe7b241d081ad Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Wed, 12 Aug 2026 18:46:36 +0200 Subject: [PATCH 1/2] Add standalone checkpoint evaluation for Puzzletron Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 28 +- .../puzzletron/docs/checkpoint_evaluation.md | 60 ++ examples/puzzletron/docs/post_mip_pipeline.md | 21 +- .../puzzletron/evaluate_lmms_checkpoint.py | 325 ++++++++ modelopt/torch/puzzletron/__init__.py | 1 + .../torch/puzzletron/evaluation/__init__.py | 18 + modelopt/torch/puzzletron/evaluation/lmms.py | 714 ++++++++++++++++++ modelopt/torch/puzzletron/post_mip/runner.py | 633 +--------------- .../test_evaluate_lmms_checkpoint_cli.py | 323 ++++++++ .../torch/puzzletron/test_lmms_evaluation.py | 452 +++++++++++ .../torch/puzzletron/test_post_mip_runner.py | 539 +------------ 11 files changed, 1960 insertions(+), 1154 deletions(-) create mode 100644 examples/puzzletron/docs/checkpoint_evaluation.md create mode 100644 examples/puzzletron/evaluate_lmms_checkpoint.py create mode 100644 modelopt/torch/puzzletron/evaluation/__init__.py create mode 100644 modelopt/torch/puzzletron/evaluation/lmms.py create mode 100644 tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py create mode 100644 tests/unit/torch/puzzletron/test_lmms_evaluation.py diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 7cd8f81c143..373cd844640 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -10,6 +10,7 @@ distill the selected model. - [Start here](#start-here) - [Setup wizard](#setup-wizard) - [Installation](#installation) +- [Evaluate a checkpoint](#evaluate-a-checkpoint) - [Run with an agent](#run-with-an-agent) - [Configuration](#configuration) - [Run a campaign](#run-a-campaign) @@ -22,6 +23,8 @@ distill the selected model. smoke and production bundles. - **Generated campaign:** complete the [installation](#installation), then [run the campaign](#run-a-campaign) with its generated bundle. +- **Checkpoint evaluation:** use [Evaluate a checkpoint](#evaluate-a-checkpoint) + for a local model without creating or running a pruning campaign. - **Agent-assisted campaign:** follow [Run with an agent](#run-with-an-agent) with your model, data, compute environment, and deployment goals. - **Existing results:** see [Reports](#reports) to regenerate a campaign report @@ -318,6 +321,23 @@ python -m pip check Record the three source revisions and verification output with the campaign. Re-run verification after pulling either fork or rebuilding a CUDA extension. +## Evaluate a checkpoint + +Basic evaluation is independent of MIP and the campaign DAG. In the Puzzletron +worker environment, run any compatible local Hugging Face checkpoint directly: + +```bash +python examples/puzzletron/evaluate_lmms_checkpoint.py \ + --checkpoint /path/to/checkpoint \ + --output-dir /path/to/results/checkpoint-smoke +``` + +The default one-GPU smoke evaluates eight samples each from IFEval and GSM8K. +Qwen 3.5 checkpoints are configured automatically. See +[checkpoint evaluation](docs/checkpoint_evaluation.md) to choose tasks, run a +full evaluation, find results, or override model detection. For options not +covered by the convenience command, use the native `python -m lmms_eval` CLI. + ## Run with an agent The canonical agent workflow is @@ -467,11 +487,9 @@ the same command for the full campaign. Add `--dry-run` to inspect either plan without submitting work, or select one stage while iterating, for example `--stage mip --dry-run`. -The setup wizard can also add downstream `lmms-eval` nodes that evaluate -materialized candidates through vLLM. They run in the standard Puzzletron -worker environment, whose example requirements pin a compatible `lmms-eval` -snapshot. See [post-MIP pipelines](docs/post_mip_pipeline.md) for configuration -details and for adding downstream evaluation to an existing campaign. +The setup wizard can also add downstream evaluation for materialized campaign +candidates. See [post-MIP pipelines](docs/post_mip_pipeline.md) to configure it +or add it to an existing campaign. ### Legacy checked-in Nano campaign diff --git a/examples/puzzletron/docs/checkpoint_evaluation.md b/examples/puzzletron/docs/checkpoint_evaluation.md new file mode 100644 index 00000000000..cadaa1e628a --- /dev/null +++ b/examples/puzzletron/docs/checkpoint_evaluation.md @@ -0,0 +1,60 @@ +# Checkpoint evaluation + +Use the standalone command to evaluate a compatible local Hugging Face +checkpoint without creating or running a Puzzletron campaign. + +## Quick start + +Install the Puzzletron worker requirements: + +```bash +python -m pip install -r examples/puzzletron/requirements.txt +``` + +Then run the default smoke: + +```bash +python examples/puzzletron/evaluate_lmms_checkpoint.py \ + --checkpoint /path/to/checkpoint \ + --output-dir /path/to/results/checkpoint-smoke +``` + +This evaluates eight samples each from IFEval and GSM8K on one GPU. Results and +logs are written under the output directory. + +## Customize the evaluation + +Choose tasks and common runtime settings with command-line options: + +```bash +python examples/puzzletron/evaluate_lmms_checkpoint.py \ + --checkpoint /path/to/checkpoint \ + --output-dir /path/to/results/custom-smoke \ + --tasks ifeval,gsm8k \ + --limit 32 \ + --tensor-parallel-size 2 \ + --dtype bfloat16 \ + --max-model-len 8192 +``` + +Qwen 3.5 checkpoints are detected from their local `config.json` and configured +automatically. Use `--reasoning-parser` to override the detected parser or +`--model-profile none` to disable model detection. + +Use `--trust-remote-code` only after reviewing the checkpoint-provided Python +code. After the smoke succeeds, use `--full` with a separate output directory +to evaluate the complete task datasets. Use `--timeout-seconds` if the full run +needs a different limit. + +For options not exposed by this convenience command, use +`python -m lmms_eval --help` and the native lmms-eval CLI. + +## Results and troubleshooting + +Each run creates a new `attempt_/` directory. Start with `summary.json` for +metrics. If a run fails, inspect `stderr.txt`; the command and raw evaluator +output are retained in the same directory. Rerunning creates another attempt +without overwriting the earlier one. + +To evaluate candidates as part of a pruning campaign, use +[downstream evaluation](post_mip_pipeline.md#downstream-evaluation) instead. diff --git a/examples/puzzletron/docs/post_mip_pipeline.md b/examples/puzzletron/docs/post_mip_pipeline.md index e325bf1efb4..95a4c6d5c18 100644 --- a/examples/puzzletron/docs/post_mip_pipeline.md +++ b/examples/puzzletron/docs/post_mip_pipeline.md @@ -141,21 +141,12 @@ from the original candidate. ## Downstream evaluation -`downstream_evaluation` runs `python -m lmms_eval` as a subprocess from the -Puzzletron worker environment. The runner passes an argument list directly and -does not invoke a shell. Values in `command_prefix` and `extra_args` are arguments; -shell syntax is not interpreted. The standard example requirements pin a snapshot -compatible with the newer `wandb` required by the pinned AutoModel build: - -```bash -python -m pip install -r examples/puzzletron/requirements.txt -python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.0"' -``` - -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. +`downstream_evaluation` adapts the generic +[checkpoint evaluator](checkpoint_evaluation.md) to materialized campaign +candidates and publishes their task metrics. Add it after a `materialize` node; +the linked example config shows the complete flow. Use the standalone +[checkpoint evaluation](checkpoint_evaluation.md) command when campaign +lineage, filtering, and reports are not needed. ## Filters diff --git a/examples/puzzletron/evaluate_lmms_checkpoint.py b/examples/puzzletron/evaluate_lmms_checkpoint.py new file mode 100644 index 00000000000..f427d60f413 --- /dev/null +++ b/examples/puzzletron/evaluate_lmms_checkpoint.py @@ -0,0 +1,325 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Evaluate a local Hugging Face checkpoint with lmms-eval and vLLM.""" + +from __future__ import annotations + +import argparse +import importlib +import importlib.util +import json +import sys +from contextlib import redirect_stdout +from pathlib import Path +from typing import TYPE_CHECKING, cast + +if TYPE_CHECKING: + from collections.abc import Callable + +__all__ = ["main"] + +REPOSITORY_ROOT = Path(__file__).resolve().parents[2] +if str(REPOSITORY_ROOT) not in sys.path: + sys.path.insert(0, str(REPOSITORY_ROOT)) + + +def _load_runner() -> Callable[..., dict[str, object]]: + """Load ModelOpt while keeping the CLI stdout machine-readable.""" + with redirect_stdout(sys.stderr): + module = importlib.import_module("modelopt.torch.puzzletron.evaluation") + return cast("Callable[..., dict[str, object]]", module.run_lmms_eval_checkpoint) + + +run_lmms_eval_checkpoint = _load_runner() + +DEFAULT_TASKS = "ifeval,gsm8k" +_TASK_ALIASES = {"gsm8k": "modelopt_gsm8k"} +_QWEN_3_5_MODEL_TYPES = frozenset({"qwen3_5", "qwen3_5_text"}) +DEFAULT_SMOKE_TIMEOUT_SECONDS = 3_000.0 +DEFAULT_FULL_TIMEOUT_SECONDS = 24 * 60 * 60.0 + + +def _checkpoint_directory(value: str) -> Path: + checkpoint = Path(value).expanduser().resolve() + if not checkpoint.is_dir(): + raise argparse.ArgumentTypeError(f"checkpoint is not a local directory: {checkpoint}") + return checkpoint + + +def _positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be a positive integer") + return parsed + + +def _positive_float(value: str) -> float: + parsed = float(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def _gpu_memory_utilization(value: str) -> float: + parsed = float(value) + if not 0 < parsed <= 1: + raise argparse.ArgumentTypeError("GPU memory utilization must be in (0, 1]") + return parsed + + +def _task_selection(value: str) -> str: + tasks = [task.strip() for task in value.split(",")] + if not tasks or any(not task for task in tasks): + raise argparse.ArgumentTypeError("tasks must be a comma-separated list of names") + return ",".join(tasks) + + +def _resolved_tasks(value: str) -> str: + return ",".join(_TASK_ALIASES.get(task, task) for task in value.split(",")) + + +def _automatic_model_args(checkpoint: Path) -> dict[str, object]: + """Return narrow compatibility defaults inferred from local checkpoint metadata.""" + try: + config = json.loads((checkpoint / "config.json").read_text()) + except (OSError, json.JSONDecodeError): + return {} + if not isinstance(config, dict): + return {} + model_types = set() + if isinstance(config.get("model_type"), str): + model_types.add(config["model_type"]) + text_config = config.get("text_config") + if isinstance(text_config, dict) and isinstance(text_config.get("model_type"), str): + model_types.add(text_config["model_type"]) + if model_types & _QWEN_3_5_MODEL_TYPES: + return {"reasoning_parser": "qwen3"} + return {} + + +def _lmms_eval_gsm8k_config() -> Path: + spec = importlib.util.find_spec("lmms_eval") + locations = spec.submodule_search_locations if spec is not None else None + if not locations: + raise RuntimeError( + "lmms_eval is not installed; install examples/puzzletron/requirements.txt" + ) + config = Path(next(iter(locations))) / "tasks/gsm8k/gsm8k.yaml" + if not config.is_file(): + raise RuntimeError(f"installed lmms_eval has no GSM8K task config: {config}") + return config.resolve() + + +def _prepare_compatibility_tasks(output_root: Path, tasks: str) -> Path | None: + """Materialize narrow task overrides while inheriting pinned lmms-eval behavior.""" + if "modelopt_gsm8k" not in tasks.split(","): + return None + tasks_root = output_root.expanduser().resolve() / "task_configs" + tasks_root.mkdir(parents=True, exist_ok=True) + config = { + "include": str(_lmms_eval_gsm8k_config()), + "task": "modelopt_gsm8k", + "dataset_path": "openai/gsm8k", + "fewshot_config": {"sampler": "default"}, + } + (tasks_root / "modelopt_gsm8k.yaml").write_text( + json.dumps(config, indent=2, sort_keys=True) + "\n" + ) + return tasks_root + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + epilog="For the full native interface, run: python -m lmms_eval --help", + ) + parser.add_argument( + "--checkpoint", + required=True, + type=_checkpoint_directory, + help="Local Hugging Face checkpoint directory to evaluate.", + ) + parser.add_argument( + "--output-dir", + required=True, + type=Path, + help="Root directory for isolated per-attempt artifacts.", + ) + parser.add_argument( + "--tasks", + default=DEFAULT_TASKS, + type=_task_selection, + help=( + f"Comma-separated text tasks (default: {DEFAULT_TASKS}); gsm8k uses a generated " + "namespaced compatibility task that inherits the pinned evaluator config." + ), + ) + limit = parser.add_mutually_exclusive_group() + limit.add_argument( + "--limit", + type=_positive_int, + default=8, + help="Maximum samples per task; the default 8 is a wiring smoke.", + ) + limit.add_argument( + "--full", + dest="limit", + action="store_const", + const=None, + help="Run every sample instead of the default wiring smoke.", + ) + parser.add_argument("--batch-size", type=_positive_int, default=1) + parser.add_argument("--tensor-parallel-size", type=_positive_int, default=1) + parser.add_argument("--dtype", default="bfloat16") + parser.add_argument( + "--gpu-memory-utilization", + type=_gpu_memory_utilization, + default=0.85, + ) + parser.add_argument("--max-model-len", type=_positive_int, default=8192) + parser.add_argument( + "--model-profile", + choices=("auto", "none"), + default="auto", + help=( + "Apply narrow vLLM compatibility defaults inferred from config.json; " + "auto currently maps Qwen 3.5 to reasoning_parser=qwen3, while none " + "leaves all model-specific arguments explicit." + ), + ) + parser.add_argument( + "--reasoning-parser", + default=None, + help="Explicit vLLM reasoning parser; overrides the detected model profile.", + ) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--timeout-seconds", + type=_positive_float, + default=None, + help=( + "Subprocess timeout; defaults to 3000 seconds for a limited smoke " + "and 86400 seconds with --full." + ), + ) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Allow reviewed checkpoint-provided Python code (disabled by default).", + ) + return parser + + +def _settings( + args: argparse.Namespace, + *, + compatibility_tasks_root: Path | None = None, + automatic_model_args: dict[str, object] | None = None, +) -> dict[str, object]: + timeout_seconds = args.timeout_seconds + if timeout_seconds is None: + timeout_seconds = ( + DEFAULT_FULL_TIMEOUT_SECONDS if args.limit is None else DEFAULT_SMOKE_TIMEOUT_SECONDS + ) + model_args = { + "dtype": args.dtype, + "gpu_memory_utilization": args.gpu_memory_utilization, + "max_model_len": args.max_model_len, + "trust_remote_code": args.trust_remote_code, + } + if args.model_profile == "auto": + if automatic_model_args is None: + automatic_model_args = _automatic_model_args(args.checkpoint) + model_args.update(automatic_model_args) + if args.reasoning_parser is not None: + model_args["reasoning_parser"] = args.reasoning_parser + settings = { + "tasks": _resolved_tasks(args.tasks), + "limit": args.limit, + "batch_size": args.batch_size, + "seed": args.seed, + "timeout_seconds": timeout_seconds, + "topology": { + "tensor_parallel_size": args.tensor_parallel_size, + "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": args.tensor_parallel_size, + }, + "model_args": model_args, + } + if compatibility_tasks_root is not None: + settings["extra_args"] = ["--include_path", str(compatibility_tasks_root)] + return settings + + +def main(argv: list[str] | None = None) -> int: + """Evaluate a checkpoint and print results or failure diagnostics as JSON. + + Args: + argv: Command-line arguments. Uses the process arguments when omitted. + + Returns: + Zero on success and one on failure. Results are written to stdout and + failure diagnostics are written to stderr. + """ + args = _build_parser().parse_args(argv) + try: + tasks = _resolved_tasks(args.tasks) + compatibility_tasks_root = _prepare_compatibility_tasks(args.output_dir, tasks) + automatic_model_args = ( + _automatic_model_args(args.checkpoint) if args.model_profile == "auto" else {} + ) + if args.reasoning_parser is None and automatic_model_args: + rendered = ",".join( + f"{key}={value}" for key, value in sorted(automatic_model_args.items()) + ) + print( + f"Detected Qwen 3.5 checkpoint; applying vLLM model argument {rendered}. " + "Override with --reasoning-parser or disable with --model-profile none.", + file=sys.stderr, + ) + result = run_lmms_eval_checkpoint( + args.checkpoint, + output_root=args.output_dir, + settings=_settings( + args, + compatibility_tasks_root=compatibility_tasks_root, + automatic_model_args=automatic_model_args, + ), + ) + except Exception as error: + payload = { + "error": type(error).__name__, + "message": str(error), + **{ + name: getattr(error, name) + for name in ("command_path", "stdout_path", "stderr_path") + if getattr(error, name, None) + }, + } + print(json.dumps(payload, indent=2, sort_keys=True), file=sys.stderr) + return 1 + print(json.dumps(result, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/modelopt/torch/puzzletron/__init__.py b/modelopt/torch/puzzletron/__init__.py index 512b7cabca5..5e4c5f7489f 100644 --- a/modelopt/torch/puzzletron/__init__.py +++ b/modelopt/torch/puzzletron/__init__.py @@ -23,6 +23,7 @@ candidates, dataset, distillation, + evaluation, export, mip, pipeline_config, diff --git a/modelopt/torch/puzzletron/evaluation/__init__.py b/modelopt/torch/puzzletron/evaluation/__init__.py new file mode 100644 index 00000000000..c1151379390 --- /dev/null +++ b/modelopt/torch/puzzletron/evaluation/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Checkpoint evaluation backends shared by direct and campaign workflows.""" + +from .lmms import * diff --git a/modelopt/torch/puzzletron/evaluation/lmms.py b/modelopt/torch/puzzletron/evaluation/lmms.py new file mode 100644 index 00000000000..a99147c89a6 --- /dev/null +++ b/modelopt/torch/puzzletron/evaluation/lmms.py @@ -0,0 +1,714 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Run lmms-eval for a local checkpoint with durable attempt artifacts.""" + +from __future__ import annotations + +import asyncio +import json +import math +import os +import shlex +import signal +import sys +import tempfile +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping, Sequence + +from ..orchestration.mesh import normalize_vllm_topology + +__all__ = [ + "DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS", + "LmmsEvalTimeoutError", + "run_lmms_eval_checkpoint", +] + +_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", + } +) +_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", + } +) +_RESERVED_EXTRA_ARG_FLAGS = frozenset( + { + "--batch-size", + "--batch_size", + "--model", + "--model_args", + "--model-args", + "--output_path", + "--output-path", + "--tasks", + } +) +DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0 +_PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 +_PROCESS_GROUP_POLL_INTERVAL_SECONDS = 0.1 + + +class LmmsEvalTimeoutError(TimeoutError): + """Report a timed-out lmms-eval process with its captured output.""" + + def __init__(self, argv: Sequence[str], timeout: float, *, output: str, stderr: str): + super().__init__(f"lmms-eval exceeded its {timeout:g}-second timeout") + self.cmd = list(argv) + self.timeout = timeout + self.output = output + self.stderr = stderr + + +@dataclass(frozen=True) +class _ProcessResult: + args: list[str] + returncode: int + stdout: str + stderr: str + + +def _atomic_json(path: Path, payload: Any) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + try: + temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + temporary.replace(path) + finally: + temporary.unlink(missing_ok=True) + return path + + +def _as_lmms_eval_arg(value: Any) -> str: + if isinstance(value, bool): + return "True" if value else "False" + 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 _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 _reject_reserved_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( + "evaluation settings.model_args must not set reserved lmms-eval model arguments: " + f"{', '.join(reserved)}" + ) + + +def _configured_tasks(settings: Mapping[str, Any]) -> tuple[str, ...]: + tasks = _join_cli_values(settings.get("tasks"), path="evaluation settings.tasks") + values = tuple(task.strip() for task in tasks.split(",")) + if not values or any(not task for task in values): + raise ValueError("evaluation settings.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_model_args(settings: Mapping[str, Any], checkpoint: str) -> str: + raw = settings.get("model_args") + checkpoint_arg = str(settings.get("checkpoint_arg", "model")) + if checkpoint_arg != "model": + raise ValueError("evaluation settings.checkpoint_arg must be 'model'") + topology = dict(settings.get("topology") or {}) + canonical_topology = normalize_vllm_topology(topology) if topology else {} + reserved_fields = frozenset( + key for key in (checkpoint_arg, *_RESERVED_TOPOLOGY_MODEL_ARG_FIELDS) if key + ) + derived: dict[str, Any] = {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 sorted(_MODEL_ARG_FIELDS): + if key in settings: + derived[key] = settings[key] + + if isinstance(raw, str): + _reject_reserved_model_args(_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("evaluation settings.model_args must be a mapping or string") + _reject_reserved_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 = shlex.split(raw) + elif isinstance(raw, Sequence) and not isinstance(raw, (bytes, bytearray)): + values = [str(item) for item in raw] + else: + raise TypeError("evaluation settings.command_prefix must be a string or sequence") + if not values or any(not value for value in values): + raise ValueError("evaluation settings.command_prefix must not be empty") + return values + + +def _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("evaluation settings.extra_args must be a string or sequence") + if any(not value for value in values): + raise ValueError("evaluation settings.extra_args must not contain empty values") + reserved = sorted( + { + value.split("=", 1)[0] + for value in values + if value.split("=", 1)[0] in _RESERVED_EXTRA_ARG_FLAGS + } + ) + if reserved: + raise ValueError( + "evaluation settings.extra_args must not set reserved lmms-eval flags: " + f"{', '.join(reserved)}" + ) + return values + + +def _build_command( + settings: Mapping[str, Any], + *, + checkpoint: str, + output_path: Path, +) -> tuple[list[str], dict[str, str], float]: + """Build a deterministic lmms-eval CLI invocation for one local checkpoint.""" + + model = str(settings.get("model", "vllm")) + if model != "vllm": + raise ValueError("evaluation settings.model must be 'vllm'") + argv = [ + *_command_prefix(settings), + "--model", + model, + "--model_args", + _merge_model_args(settings, checkpoint), + "--tasks", + ",".join(_configured_tasks(settings)), + "--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(_extra_args(settings)) + + env = os.environ.copy() + env_overrides = dict(settings.get("env") or {}) + for key, value in env_overrides.items(): + if value is not None: + env[str(key)] = str(value) + if settings.get("cache_dir") is not None and "LMMS_EVAL_HOME" not in env_overrides: + env["LMMS_EVAL_HOME"] = str(settings["cache_dir"]) + timeout = settings.get("timeout_seconds") + if timeout is None: + timeout = settings.get("timeout") + if timeout is None: + timeout = DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS + timeout = float(timeout) + if not math.isfinite(timeout) or timeout <= 0: + raise ValueError("lmms-eval timeout must be a finite positive number") + return argv, env, timeout + + +def _numeric_metrics(task_payload: Mapping[str, Any]) -> dict[str, float]: + return { + str(metric_name): float(value) + for metric_name, value in task_payload.items() + if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value) + } + + +def _metric_key(value: Any) -> str: + return ( + str(value).strip().replace(" ", "_").replace(",", "_").replace("/", "_").replace("\\", "_") + ) + + +def _flatten_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 isinstance(task_payload, Mapping): + metrics.update( + { + f"{_metric_key(task_name)}.{_metric_key(metric_name)}": value + for metric_name, value in _numeric_metrics(task_payload).items() + } + ) + return metrics + + +def _resolved_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_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_tasks(payload, configured_tasks) + missing_results = [task for task in expected_tasks if task not in results] + if missing_results: + raise RuntimeError( + f"lmms-eval result is missing configured task results: {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: + count = _sample_count(payload, task) + if count is None: + missing_samples.append(task) + elif count <= 0: + zero_samples.append(task) + else: + sample_counts[task] = 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 _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_streams(output_path: Path, result: _ProcessResult) -> dict[str, str]: + stream_paths = {} + for stream_name, text in (("stdout", result.stdout), ("stderr", result.stderr)): + stream_path = output_path / f"{stream_name}.txt" + stream_path.write_text(text or "") + stream_paths[f"{stream_name}_path"] = str(stream_path) + return stream_paths + + +def _stream_text(value: str | bytes | None) -> str: + return value.decode(errors="replace") if isinstance(value, bytes) else value or "" + + +def _output_tail(result: _ProcessResult, *, 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_process_group(process: asyncio.subprocess.Process, 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 _process_group_exists(process: asyncio.subprocess.Process) -> bool: + if os.name != "posix": + return process.returncode is None + try: + os.killpg(process.pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +async def _wait_for_process_group_exit( + process: asyncio.subprocess.Process, *, deadline: float +) -> None: + loop = asyncio.get_running_loop() + while _process_group_exists(process): + remaining = deadline - loop.time() + if remaining <= 0: + return + await asyncio.sleep(min(_PROCESS_GROUP_POLL_INTERVAL_SECONDS, remaining)) + + +async def _run_process_async( + argv: list[str], + *, + cwd: str, + env: Mapping[str, str], + timeout: float, +) -> _ProcessResult: + # lmms-eval needs process isolation for bounded GPU-worker cleanup. The argument + # vector is passed directly; no shell interprets checkpoint or configuration values. + with tempfile.TemporaryFile() as stdout_file, tempfile.TemporaryFile() as stderr_file: + process = await asyncio.create_subprocess_exec( + *argv, + cwd=cwd, + env=env, + stdout=stdout_file, + stderr=stderr_file, + start_new_session=os.name == "posix", + ) + try: + await asyncio.wait_for(process.wait(), timeout) + except TimeoutError as error: + _signal_process_group(process, signal.SIGTERM) + try: + await asyncio.wait_for(process.wait(), _PROCESS_CLEANUP_TIMEOUT_SECONDS) + except TimeoutError: + _signal_process_group(process, signal.SIGKILL) + try: + await asyncio.wait_for(process.wait(), _PROCESS_CLEANUP_TIMEOUT_SECONDS) + except TimeoutError: + pass + if _process_group_exists(process): + _signal_process_group(process, signal.SIGKILL) + await _wait_for_process_group_exit( + process, + deadline=asyncio.get_running_loop().time() + _PROCESS_CLEANUP_TIMEOUT_SECONDS, + ) + stdout_file.seek(0) + stderr_file.seek(0) + raise LmmsEvalTimeoutError( + argv, + timeout, + output=_stream_text(stdout_file.read()), + stderr=_stream_text(stderr_file.read()), + ) from error + stdout_file.seek(0) + stderr_file.seek(0) + return _ProcessResult( + args=argv, + returncode=int(process.returncode or 0), + stdout=_stream_text(stdout_file.read()), + stderr=_stream_text(stderr_file.read()), + ) + + +def _run_process( + argv: list[str], + *, + cwd: str, + env: Mapping[str, str], + timeout: float, +) -> _ProcessResult: + return asyncio.run(_run_process_async(argv, cwd=cwd, env=env, timeout=timeout)) + + +def _annotate_error( + error: Exception, + *, + command_path: Path, + stream_paths: Mapping[str, str], +) -> None: + setattr(error, "command_path", str(command_path)) + for name in ("stdout_path", "stderr_path"): + if name in stream_paths: + setattr(error, name, stream_paths[name]) + + +def run_lmms_eval_checkpoint( + checkpoint: str | Path, + *, + output_root: str | Path, + settings: Mapping[str, Any], +) -> dict[str, Any]: + """Evaluate one local checkpoint and preserve an isolated lmms-eval attempt. + + Args: + checkpoint: Local Hugging Face checkpoint directory. + output_root: Root under which a unique attempt directory is created. + settings: lmms-eval tasks, vLLM model arguments, topology, and runtime controls. + + Returns: + Flattened metrics and paths to the normalized summary, raw result, command, + stdout, and stderr artifacts. + """ + + checkpoint_path = Path(checkpoint).expanduser().resolve() + if not checkpoint_path.is_dir(): + raise FileNotFoundError(f"checkpoint is not a local directory: {checkpoint_path}") + output = Path(output_root).expanduser().resolve() / f"attempt_{uuid.uuid4().hex}" + settings = dict(settings) + argv, env, timeout = _build_command( + settings, + checkpoint=str(checkpoint_path), + output_path=output, + ) + output.mkdir(parents=True, exist_ok=True) + command_path = _atomic_json( + output / "command.json", + { + "argv": argv, + "env_overrides": sorted(str(key) for key in dict(settings.get("env") or {})), + "timeout": timeout, + }, + ) + try: + result = _run_process(argv, cwd=str(output), env=env, timeout=timeout) + except LmmsEvalTimeoutError as error: + captured = _ProcessResult(argv, -1, error.output, error.stderr) + stream_paths = _write_streams(output, captured) + _annotate_error(error, command_path=command_path, stream_paths=stream_paths) + raise + + stream_paths = _write_streams(output, result) + if result.returncode: + tail = _output_tail(result) + failure = RuntimeError( + f"lmms-eval failed with exit code {result.returncode}" + (f": {tail}" if tail else "") + ) + _annotate_error(failure, command_path=command_path, stream_paths=stream_paths) + raise failure + + try: + payload, result_path = _result_payload(output) + sample_counts = _validate_completion(payload, _configured_tasks(settings)) + metrics = _flatten_metrics(payload) + if not metrics: + raise RuntimeError(f"lmms-eval result has no numeric task metrics: {result_path}") + except FileNotFoundError as error: + tail = _output_tail(result) + failure = FileNotFoundError(f"{error}: {tail}" if tail else str(error)) + _annotate_error(failure, command_path=command_path, stream_paths=stream_paths) + raise failure from error + except RuntimeError as error: + _annotate_error(error, command_path=command_path, stream_paths=stream_paths) + raise + + summary = { + "checkpoint": str(checkpoint_path), + "metrics": metrics, + "result_path": str(result_path), + "sample_counts": sample_counts, + } + summary_path = _atomic_json(output / "summary.json", summary) + return { + "metrics": metrics, + "result_path": str(summary_path), + "raw_result_path": str(result_path), + "command_path": str(command_path), + **stream_paths, + } diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 44b5143d5f3..08aeff370a1 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/modelopt/torch/puzzletron/post_mip/runner.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Execute and aggregate one compiled post-MIP node.""" from __future__ import annotations @@ -24,11 +21,7 @@ import json import math import os -import shlex -import signal -import subprocess import sys -import time import traceback import uuid from contextlib import contextmanager @@ -36,8 +29,8 @@ from pathlib import Path from typing import Any, Iterator, Mapping, Sequence +from ..evaluation import DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS, run_lmms_eval_checkpoint 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 @@ -115,6 +108,13 @@ def _needs_puzzletron_process_group(node_type: str) -> bool: return node_type == "evaluation" +def _is_loaded_subprocess_timeout(error: Exception) -> bool: + """Recognize the timeout type loaded by the optional AIPerf adapter.""" + + timeout_type = getattr(sys.modules.get("subprocess"), "TimeoutExpired", None) + return timeout_type is not None and isinstance(error, timeout_type) + + def _node_root(config: Mapping[str, Any], node: CompiledPostMIPNode) -> Path: return _puzzle_dir(config) / "artifacts" / "post_mip" / "nodes" / node.node_id @@ -528,553 +528,6 @@ 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( - { - "--batch-size", - "--batch_size", - "--model", - "--model_args", - "--model-args", - "--output_path", - "--output-path", - "--tasks", - } -) -_DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0 -_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 -_LMMS_EVAL_PROCESS_GROUP_POLL_INTERVAL_SECONDS = 0.1 - - -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")) - if checkpoint_arg != "model": - raise ValueError("downstream_evaluation.config.checkpoint_arg must be '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 sorted(_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]: - """Build a deterministic lmms-eval CLI invocation for one realized checkpoint.""" - - model = str(settings.get("model", "vllm")) - if model != "vllm": - raise ValueError("downstream_evaluation.config.model must be 'vllm'") - tasks = ",".join(_configured_lmms_eval_tasks(settings)) - argv = [ - *_command_prefix(settings), - "--model", - model, - "--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") - if timeout is None: - timeout = settings.get("timeout") - if timeout is None: - timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS - timeout = float(timeout) - if not math.isfinite(timeout) or timeout <= 0: - raise ValueError("lmms-eval timeout must be a finite positive number") - return argv, env, timeout - - -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( - f"lmms-eval result is missing configured task results: {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 _wait_for_lmms_eval_process_group_exit( - process: subprocess.Popen[str], *, deadline: float -) -> None: - while _lmms_eval_process_group_exists(process): - remaining = deadline - time.monotonic() - if remaining <= 0: - return - time.sleep(min(_LMMS_EVAL_PROCESS_GROUP_POLL_INTERVAL_SECONDS, remaining)) - - -def _run_lmms_eval_process( - argv: list[str], - *, - cwd: str, - env: Mapping[str, str], - timeout: float, -) -> 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) - cleanup_deadline = time.monotonic() + _LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS - try: - stdout, stderr = process.communicate( - timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS - ) - except subprocess.TimeoutExpired as kill_error: - stdout, stderr = kill_error.output, kill_error.stderr - _wait_for_lmms_eval_process_group_exit(process, deadline=cleanup_deadline) - else: - if _lmms_eval_process_group_exists(process): - _signal_lmms_eval_process_group(process, signal.SIGKILL) - _wait_for_lmms_eval_process_group_exit( - process, - deadline=time.monotonic() + _LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, - ) - raise subprocess.TimeoutExpired( - argv, - error.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, @@ -1083,70 +536,17 @@ def _downstream_evaluation( ) -> 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 run_lmms_eval_checkpoint( + source.artifact["checkpoint"], + output_root=output_root, + settings=node.config.get("config") or {}, ) - 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( @@ -1277,8 +677,9 @@ 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 in {"aiperf", "downstream_evaluation"} and isinstance( - error, (subprocess.TimeoutExpired, TimeoutError) + subprocess_timeout = _is_loaded_subprocess_timeout(error) + timed_out = node.node_type in {"aiperf", "downstream_evaluation"} and ( + subprocess_timeout or isinstance(error, TimeoutError) ) row = { "input_revision_id": revision_id, @@ -1292,10 +693,10 @@ def run_post_mip_node_shard( timeout_field = "benchmark_timeout" if node.node_type == "downstream_evaluation": timeout_field = "timeout_seconds" - elif not isinstance(error, subprocess.TimeoutExpired): + elif not subprocess_timeout: timeout_field = "readiness_timeout" default_timeout = ( - _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS + DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS if node.node_type == "downstream_evaluation" else 600 if timeout_field == "benchmark_timeout" diff --git a/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py b/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py new file mode 100644 index 00000000000..9d57eeb6d22 --- /dev/null +++ b/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py @@ -0,0 +1,323 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the direct local-checkpoint lmms-eval CLI.""" + +import importlib.util +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +_REPOSITORY_ROOT = Path(__file__).resolve().parents[4] +_SCRIPT_PATH = _REPOSITORY_ROOT / "examples/puzzletron/evaluate_lmms_checkpoint.py" +_SPEC = importlib.util.spec_from_file_location("evaluate_lmms_checkpoint", _SCRIPT_PATH) +assert _SPEC is not None and _SPEC.loader is not None +evaluate_lmms_checkpoint = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(evaluate_lmms_checkpoint) + + +def test_runner_import_keeps_stdout_machine_readable(monkeypatch, capsys): + runner = object() + + def fake_import_module(name): + assert name == "modelopt.torch.puzzletron.evaluation" + print("import-time diagnostic") + return SimpleNamespace(run_lmms_eval_checkpoint=runner) + + monkeypatch.setattr(evaluate_lmms_checkpoint.importlib, "import_module", fake_import_module) + + assert evaluate_lmms_checkpoint._load_runner() is runner + captured = capsys.readouterr() + assert captured.out == "" + assert captured.err == "import-time diagnostic\n" + + +def test_cli_help_explains_qwen_profile_and_native_escape_hatch(): + help_text = evaluate_lmms_checkpoint._build_parser().format_help() + normalized_help = " ".join(help_text.split()) + + assert "--model-profile {auto,none}" in help_text + assert "Qwen 3.5" in help_text + assert "reasoning_parser=qwen3" in help_text + assert "python -m lmms_eval --help" in normalized_help + + +def test_cli_runs_generic_text_smoke_defaults(monkeypatch, tmp_path, capsys): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + output_root = tmp_path / "results" + upstream_task = tmp_path / "lmms_eval/tasks/gsm8k/gsm8k.yaml" + upstream_task.parent.mkdir(parents=True) + upstream_task.write_text("task: gsm8k\n") + captured = {} + + def fake_run(checkpoint_path, *, output_root, settings): + captured["checkpoint"] = checkpoint_path + captured["output_root"] = output_root + captured["settings"] = settings + return {"result_path": "/results/attempt/summary.json", "metrics": {}} + + monkeypatch.setattr(evaluate_lmms_checkpoint, "run_lmms_eval_checkpoint", fake_run) + monkeypatch.setattr(evaluate_lmms_checkpoint, "_lmms_eval_gsm8k_config", lambda: upstream_task) + + returncode = evaluate_lmms_checkpoint.main( + ["--checkpoint", str(checkpoint), "--output-dir", str(output_root)] + ) + + assert returncode == 0 + assert captured == { + "checkpoint": checkpoint.resolve(), + "output_root": output_root, + "settings": { + "tasks": "ifeval,modelopt_gsm8k", + "limit": 8, + "batch_size": 1, + "seed": 42, + "timeout_seconds": evaluate_lmms_checkpoint.DEFAULT_SMOKE_TIMEOUT_SECONDS, + "topology": { + "tensor_parallel_size": 1, + "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": 1, + }, + "model_args": { + "dtype": "bfloat16", + "gpu_memory_utilization": 0.85, + "max_model_len": 8192, + "trust_remote_code": False, + }, + "extra_args": [ + "--include_path", + str(output_root / "task_configs"), + ], + }, + } + assert json.loads((output_root / "task_configs/modelopt_gsm8k.yaml").read_text()) == { + "dataset_path": "openai/gsm8k", + "fewshot_config": {"sampler": "default"}, + "include": str(upstream_task), + "task": "modelopt_gsm8k", + } + assert json.loads(capsys.readouterr().out) == { + "metrics": {}, + "result_path": "/results/attempt/summary.json", + } + + +def test_cli_auto_configures_qwen_3_5_and_reports_the_choice(monkeypatch, tmp_path, capsys): + checkpoint = tmp_path / "qwen" + checkpoint.mkdir() + (checkpoint / "config.json").write_text( + json.dumps( + { + "model_type": "qwen3_5", + "text_config": {"model_type": "qwen3_5_text"}, + } + ) + ) + captured = {} + + def fake_run(_checkpoint_path, *, output_root, settings): + captured["settings"] = settings + return {"result_path": str(output_root / "summary.json"), "metrics": {}} + + monkeypatch.setattr(evaluate_lmms_checkpoint, "run_lmms_eval_checkpoint", fake_run) + + returncode = evaluate_lmms_checkpoint.main( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + "--tasks", + "ifeval", + ] + ) + + assert returncode == 0 + assert captured["settings"]["model_args"]["reasoning_parser"] == "qwen3" + captured_streams = capsys.readouterr() + assert json.loads(captured_streams.out)["metrics"] == {} + assert "Detected Qwen 3.5 checkpoint" in captured_streams.err + assert "reasoning_parser=qwen3" in captured_streams.err + + +@pytest.mark.parametrize( + ("extra_args", "expected"), + [ + (["--model-profile", "none"], None), + (["--reasoning-parser", "custom"], "custom"), + ], +) +def test_cli_can_disable_or_override_qwen_3_5_profile(tmp_path, extra_args, expected): + checkpoint = tmp_path / "qwen" + checkpoint.mkdir() + (checkpoint / "config.json").write_text('{"model_type": "qwen3_5"}\n') + + args = evaluate_lmms_checkpoint._build_parser().parse_args( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + *extra_args, + ] + ) + + assert ( + evaluate_lmms_checkpoint._settings(args)["model_args"].get("reasoning_parser") == expected + ) + + +def test_cli_ignores_unrecognized_or_malformed_model_metadata(tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text( + '{"model_type": ["not", "a", "string"], "text_config": []}\n' + ) + + assert evaluate_lmms_checkpoint._automatic_model_args(checkpoint) == {} + + +def test_cli_full_run_and_runtime_overrides_are_wired(tmp_path): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + + args = evaluate_lmms_checkpoint._build_parser().parse_args( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + "--full", + "--tensor-parallel-size", + "2", + "--reasoning-parser", + "qwen3", + "--trust-remote-code", + ] + ) + + settings = evaluate_lmms_checkpoint._settings(args) + assert settings["limit"] is None + assert settings["timeout_seconds"] == evaluate_lmms_checkpoint.DEFAULT_FULL_TIMEOUT_SECONDS + assert settings["topology"]["tensor_parallel_size"] == 2 + assert settings["topology"]["gpu_group_size"] == 2 + assert settings["model_args"]["reasoning_parser"] == "qwen3" + assert settings["model_args"]["trust_remote_code"] is True + + +def test_cli_maps_gsm8k_to_namespaced_compatibility_task(tmp_path): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + + args = evaluate_lmms_checkpoint._build_parser().parse_args( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + "--tasks", + " ifeval, gsm8k,custom_task ", + ] + ) + + assert evaluate_lmms_checkpoint._settings(args)["tasks"] == ( + "ifeval,modelopt_gsm8k,custom_task" + ) + + +def test_cli_rejects_empty_task_names(tmp_path): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + + with pytest.raises(SystemExit): + evaluate_lmms_checkpoint._build_parser().parse_args( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + "--tasks", + "ifeval,,gsm8k", + ] + ) + + +def test_compatibility_task_is_not_written_when_gsm8k_is_not_selected(tmp_path): + assert ( + evaluate_lmms_checkpoint._prepare_compatibility_tasks( + tmp_path / "results", "ifeval,custom_task" + ) + is None + ) + assert not (tmp_path / "results").exists() + + +def test_cli_explicit_timeout_overrides_full_default(tmp_path): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + + args = evaluate_lmms_checkpoint._build_parser().parse_args( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + "--full", + "--timeout-seconds", + "123", + ] + ) + + assert evaluate_lmms_checkpoint._settings(args)["timeout_seconds"] == 123 + + +def test_cli_reports_failed_attempt_evidence_payload(monkeypatch, tmp_path, capsys): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + failure = RuntimeError("lmms-eval failed") + failure.command_path = "/results/attempt/command.json" + failure.stdout_path = "/results/attempt/stdout.txt" + failure.stderr_path = "/results/attempt/stderr.txt" + + def fail(*_args, **_kwargs): + raise failure + + monkeypatch.setattr(evaluate_lmms_checkpoint, "run_lmms_eval_checkpoint", fail) + monkeypatch.setattr( + evaluate_lmms_checkpoint, + "_lmms_eval_gsm8k_config", + lambda: tmp_path / "lmms_eval/tasks/gsm8k/gsm8k.yaml", + ) + + returncode = evaluate_lmms_checkpoint.main( + ["--checkpoint", str(checkpoint), "--output-dir", str(tmp_path / "results")] + ) + + assert returncode == 1 + assert json.loads(capsys.readouterr().err) == { + "command_path": failure.command_path, + "error": "RuntimeError", + "message": "lmms-eval failed", + "stderr_path": failure.stderr_path, + "stdout_path": failure.stdout_path, + } diff --git a/tests/unit/torch/puzzletron/test_lmms_evaluation.py b/tests/unit/torch/puzzletron/test_lmms_evaluation.py new file mode 100644 index 00000000000..65d1750443e --- /dev/null +++ b/tests/unit/torch/puzzletron/test_lmms_evaluation.py @@ -0,0 +1,452 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the reusable lmms-eval checkpoint backend.""" + +import json +import os +import sys +from pathlib import Path + +import pytest + +from modelopt.torch.puzzletron.evaluation import lmms + + +def _settings(*tasks: str) -> dict: + return {"tasks": list(tasks), "topology": {"gpu_group_size": 1}} + + +def _write_result( + output: Path, + *, + results: dict, + sample_counts: dict, + group_subtasks: dict | None = None, +) -> None: + output.mkdir(parents=True, exist_ok=True) + (output / "results.json").write_text( + json.dumps( + { + "results": results, + "group_subtasks": group_subtasks or {}, + "n-samples": sample_counts, + } + ) + ) + + +def test_command_maps_checkpoint_and_vllm_topology(tmp_path): + argv, env, timeout = lmms._build_command( + { + "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] == [sys.executable, "-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_command_forwards_unowned_model_and_evaluator_options(tmp_path): + argv, _, _ = lmms._build_command( + { + **_settings("ifeval"), + "model_args": {"new_vllm_option": "enabled"}, + "extra_args": ["--new-lmms-option", "enabled"], + }, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + model_args = argv[argv.index("--model_args") + 1] + evaluator_option = argv.index("--new-lmms-option") + assert "new_vllm_option=enabled" in model_args + assert argv[evaluator_option + 1] == "enabled" + + +def test_command_splits_string_prefix(tmp_path): + argv, _, _ = lmms._build_command( + {**_settings("ifeval"), "command_prefix": "python -m lmms_eval"}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + assert argv[:3] == ["python", "-m", "lmms_eval"] + + +def test_cache_dir_precedence(monkeypatch, tmp_path): + monkeypatch.setenv("LMMS_EVAL_HOME", "/inherited/cache") + settings = {**_settings("ifeval"), "cache_dir": tmp_path / "configured-cache"} + + _, env, _ = lmms._build_command( + settings, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + assert env["LMMS_EVAL_HOME"] == str(tmp_path / "configured-cache") + + settings["env"] = {"LMMS_EVAL_HOME": "/explicit/cache"} + _, env, _ = lmms._build_command( + settings, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + assert env["LMMS_EVAL_HOME"] == "/explicit/cache" + + +def test_command_uses_bounded_default_timeout(tmp_path): + _, _, timeout = lmms._build_command( + _settings("ifeval"), + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + assert timeout == lmms.DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS + + +def test_command_rejects_non_sequence_prefix_setting(tmp_path): + with pytest.raises(TypeError, match="command_prefix must be a string or sequence"): + lmms._build_command( + {**_settings("ifeval"), "command_prefix": {"executable": "lmms_eval"}}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + +@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf")]) +def test_command_rejects_invalid_timeout(tmp_path, timeout): + with pytest.raises(ValueError, match="finite positive number"): + lmms._build_command( + {**_settings("ifeval"), "timeout_seconds": timeout}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + +@pytest.mark.parametrize( + ("model_args", "expected"), + [ + ({"model": "/ckpts/wrong"}, "model"), + ("dtype=bfloat16,tensor_parallel_size=1", "tensor_parallel_size"), + ], +) +def test_command_rejects_reserved_model_args_setting(tmp_path, model_args, expected): + with pytest.raises(ValueError, match="reserved lmms-eval model arguments") as exc_info: + lmms._build_command( + {**_settings("ifeval"), "model_args": model_args}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + assert expected in str(exc_info.value) + + +@pytest.mark.parametrize( + ("settings", "expected"), + [ + ({"model": "hf"}, "settings.model must be 'vllm'"), + ({"checkpoint_arg": "pretrained"}, "settings.checkpoint_arg must be 'model'"), + ], +) +def test_command_rejects_unsupported_backend_contract(tmp_path, settings, expected): + with pytest.raises(ValueError, match=expected): + lmms._build_command( + {**_settings("ifeval"), **settings}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + +@pytest.mark.parametrize( + ("extra_args", "expected"), + [ + pytest.param(["--model", "hf"], "--model", id="model"), + pytest.param(["--tasks", "gsm8k"], "--tasks", id="tasks"), + pytest.param(["--batch_size=99"], "--batch_size", id="batch-size-underscore"), + pytest.param("--batch-size 99", "--batch-size", id="batch-size-hyphen"), + pytest.param("--output_path /tmp/other", "--output_path", id="output-path"), + pytest.param(["--model_args=model=/wrong"], "--model_args", id="model-args"), + ], +) +def test_command_rejects_reserved_extra_args_setting(tmp_path, extra_args, expected): + with pytest.raises(ValueError, match="reserved lmms-eval flags") as exc_info: + lmms._build_command( + {**_settings("ifeval"), "extra_args": extra_args}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + assert expected in str(exc_info.value) + + +@pytest.mark.skipif(os.name != "posix", reason="process groups are POSIX-specific") +def test_timeout_kills_ignored_process_group_members(monkeypatch, tmp_path): + script = ( + "import signal,time; " + "signal.signal(signal.SIGTERM, signal.SIG_IGN); " + "print('partial stdout', flush=True); " + "time.sleep(60)" + ) + monkeypatch.setattr(lmms, "_PROCESS_CLEANUP_TIMEOUT_SECONDS", 0.1) + + with pytest.raises(lmms.LmmsEvalTimeoutError) as exc_info: + lmms._run_process( + [sys.executable, "-c", script], + cwd=str(tmp_path), + env=os.environ.copy(), + timeout=1.0, + ) + + assert exc_info.value.timeout == 1.0 + assert exc_info.value.output == "partial stdout\n" + assert exc_info.value.stderr == "" + + +def test_run_checkpoint_flattens_metrics_and_preserves_artifacts(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + + def fake_run(argv, *, cwd, env, timeout): + del env, timeout + _write_result( + Path(cwd) / "nested", + results={ + "ifeval": {"prompt_level_strict_acc,none": 0.5}, + "gsm8k": {"exact_match,strict-match": 0.75}, + }, + sample_counts={ + "ifeval": {"original": 541, "effective": 4}, + "gsm8k": {"original": 1319, "effective": 4}, + }, + ) + return lmms._ProcessResult(argv, 0, stdout="done\n", stderr="") + + monkeypatch.setattr(lmms, "_run_process", fake_run) + + result = lmms.run_lmms_eval_checkpoint( + checkpoint, + output_root=tmp_path / "results", + settings={**_settings("ifeval", "gsm8k"), "limit": 4}, + ) + + assert result["metrics"] == { + "gsm8k.exact_match_strict-match": 0.75, + "ifeval.prompt_level_strict_acc_none": 0.5, + } + assert Path(result["command_path"]).is_file() + assert Path(result["stdout_path"]).read_text() == "done\n" + assert Path(result["stderr_path"]).read_text() == "" + summary = json.loads(Path(result["result_path"]).read_text()) + assert summary["checkpoint"] == str(checkpoint.resolve()) + assert summary["sample_counts"] == {"gsm8k": 4.0, "ifeval": 4.0} + + +def test_run_checkpoint_executes_real_process(tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + script = """ +import json +from pathlib import Path + +Path("results.json").write_text(json.dumps({ + "results": {"ifeval": {"accuracy": 0.5}}, + "n-samples": {"ifeval": {"effective": 2}}, +})) +print("evaluator stdout") +print("evaluator stderr", file=__import__("sys").stderr) +""" + + result = lmms.run_lmms_eval_checkpoint( + checkpoint, + output_root=tmp_path / "results", + settings={ + **_settings("ifeval"), + "command_prefix": [sys.executable, "-c", script], + }, + ) + + assert result["metrics"] == {"ifeval.accuracy": 0.5} + assert Path(result["stdout_path"]).read_text() == "evaluator stdout\n" + assert Path(result["stderr_path"]).read_text() == "evaluator stderr\n" + assert Path(result["raw_result_path"]).name == "results.json" + command = json.loads(Path(result["command_path"]).read_text()) + assert command["argv"][:2] == [sys.executable, "-c"] + + +def test_run_checkpoint_preserves_failure_artifacts(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + monkeypatch.setattr( + lmms, + "_run_process", + lambda argv, **kwargs: lmms._ProcessResult( + argv, 2, stdout="partial evaluator output\n", stderr="backend failed\n" + ), + ) + + with pytest.raises(RuntimeError, match="lmms-eval failed with exit code 2") as exc_info: + lmms.run_lmms_eval_checkpoint( + checkpoint, + output_root=tmp_path / "results", + settings=_settings("ifeval"), + ) + + error = exc_info.value + assert Path(error.command_path).is_file() + assert Path(error.stdout_path).read_text() == "partial evaluator output\n" + assert Path(error.stderr_path).read_text() == "backend failed\n" + + +def test_run_checkpoint_preserves_timeout_artifacts(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + + def time_out(argv, **_kwargs): + raise lmms.LmmsEvalTimeoutError( + argv, + 7, + output="partial evaluator output\n", + stderr="evaluation timed out\n", + ) + + monkeypatch.setattr(lmms, "_run_process", time_out) + + with pytest.raises(lmms.LmmsEvalTimeoutError) as exc_info: + lmms.run_lmms_eval_checkpoint( + checkpoint, + output_root=tmp_path / "results", + settings=_settings("ifeval"), + ) + + error = exc_info.value + assert Path(error.command_path).is_file() + assert Path(error.stdout_path).read_text() == "partial evaluator output\n" + assert Path(error.stderr_path).read_text() == "evaluation timed out\n" + + +def test_completion_validates_resolved_task_expansion(): + sample_counts = lmms._validate_completion( + { + "results": { + "arc_challenge": {"acc,none": 0.25}, + "hellaswag": {"acc_norm,none": 0.5}, + }, + "group_subtasks": {"leaderboard": ["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} + + +@pytest.mark.parametrize( + ("results", "sample_counts", "expected"), + [ + ( + {"ifeval": {"accuracy": 0.5}}, + {"ifeval": {"effective": 4}}, + "missing configured task results", + ), + ( + {"ifeval": {"accuracy": 0.5}, "gsm8k": {"accuracy": 0.75}}, + {"ifeval": {"effective": 4}, "gsm8k": {"effective": 0}}, + "zero effective samples", + ), + ], +) +def test_run_checkpoint_rejects_incomplete_results( + monkeypatch, tmp_path, results, sample_counts, expected +): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + + def fake_run(argv, *, cwd, env, timeout): + del env, timeout + _write_result(Path(cwd), results=results, sample_counts=sample_counts) + return lmms._ProcessResult(argv, 0, stdout="", stderr="") + + monkeypatch.setattr(lmms, "_run_process", fake_run) + + with pytest.raises(RuntimeError, match=expected): + lmms.run_lmms_eval_checkpoint( + checkpoint, + output_root=tmp_path / "results", + settings=_settings("ifeval", "gsm8k"), + ) + + +def test_run_checkpoint_reports_output_when_results_are_missing(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + monkeypatch.setattr( + lmms, + "_run_process", + lambda argv, **kwargs: lmms._ProcessResult( + argv, + 0, + stdout="Saving results aggregated\nCould not save results aggregated\n", + stderr="", + ), + ) + + with pytest.raises(FileNotFoundError) as exc_info: + lmms.run_lmms_eval_checkpoint( + checkpoint, + output_root=tmp_path / "results", + settings=_settings("ifeval"), + ) + + error = exc_info.value + assert "lmms-eval wrote no JSON results" in str(error) + assert "Could not save results aggregated" in str(error) + assert Path(error.command_path).parent == Path(error.stdout_path).parent + + +def test_run_checkpoint_requires_local_directory(tmp_path): + with pytest.raises(FileNotFoundError, match="checkpoint is not a local directory"): + lmms.run_lmms_eval_checkpoint( + tmp_path / "missing", + output_root=tmp_path / "results", + settings=_settings("ifeval"), + ) diff --git a/tests/unit/torch/puzzletron/test_post_mip_runner.py b/tests/unit/torch/puzzletron/test_post_mip_runner.py index aba8ec44834..47f6e17012e 100644 --- a/tests/unit/torch/puzzletron/test_post_mip_runner.py +++ b/tests/unit/torch/puzzletron/test_post_mip_runner.py @@ -13,19 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - """Tests for post-MIP execution, including managed downstream evaluation.""" -import json -import signal -import subprocess -import sys from pathlib import Path from types import SimpleNamespace -import pytest from omegaconf import OmegaConf from modelopt.torch.puzzletron.post_mip import runner @@ -194,527 +186,38 @@ def fake_run_aiperf_sweep(checkpoint, **settings): assert result["metrics"] == {} -def test_lmms_eval_command_maps_checkpoint_and_vllm_topology(tmp_path): - argv, env, timeout = runner._lmms_eval_command( - { - "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] == [sys.executable, "-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_uses_bounded_default_timeout(tmp_path): - _, _, timeout = runner._lmms_eval_command( - { - "tasks": ["ifeval"], - "topology": {"gpu_group_size": 1}, - }, - checkpoint="/ckpts/candidate", - output_path=tmp_path / "results", - ) - - assert timeout == runner._DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS - - -@pytest.mark.parametrize("timeout", [0, -1, float("nan"), float("inf")]) -def test_lmms_eval_command_rejects_invalid_timeout(tmp_path, timeout): - with pytest.raises(ValueError, match="finite positive number"): - runner._lmms_eval_command( - { - "tasks": ["ifeval"], - "timeout_seconds": timeout, - "topology": {"gpu_group_size": 1}, - }, - checkpoint="/ckpts/candidate", - output_path=tmp_path / "results", - ) - - -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 - - -@pytest.mark.parametrize( - ("settings", "expected"), - [ - ({"model": "hf"}, "config.model must be 'vllm'"), - ({"checkpoint_arg": "pretrained"}, "config.checkpoint_arg must be 'model'"), - ], -) -def test_lmms_eval_command_rejects_non_vllm_managed_settings(tmp_path, settings, expected): - with pytest.raises(ValueError, match=expected): - runner._lmms_eval_command( - { - "tasks": ["ifeval"], - "topology": {"gpu_group_size": 1}, - **settings, - }, - checkpoint="/ckpts/candidate", - output_path=tmp_path / "results", - ) - - -@pytest.mark.parametrize( - ("extra_args", "expected"), - [ - pytest.param(["--model", "hf"], "--model", id="model"), - pytest.param(["--tasks", "gsm8k"], "--tasks", id="tasks"), - pytest.param(["--batch_size=99"], "--batch_size", id="batch-size-underscore"), - pytest.param("--batch-size 99", "--batch-size", id="batch-size-hyphen"), - pytest.param("--output_path /tmp/other", "--output_path", id="output-path"), - pytest.param( - ["--model_args=model=/ckpts/wrong"], - "--model_args", - id="model-args", - ), - ], -) -def test_lmms_eval_command_rejects_reserved_extra_args(tmp_path, extra_args, expected): - with pytest.raises(ValueError, match="reserved lmms-eval flags") as exc_info: - runner._lmms_eval_command( - { - "tasks": ["ifeval"], - "topology": {"gpu_group_size": 1}, - "extra_args": extra_args, - }, - checkpoint="/ckpts/candidate", - output_path=tmp_path / "results", - ) - - assert expected in str(exc_info.value) - - -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) - - with pytest.raises(subprocess.TimeoutExpired) as exc_info: - runner._run_lmms_eval_process( - ["python", "-m", "lmms_eval"], - cwd=str(tmp_path), - env={}, - timeout=7.0, - ) - - assert exc_info.value.timeout == 7.0 - assert exc_info.value.output == "partial stdout" - assert exc_info.value.stderr == "partial stderr" - - 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 = [] - sleep_intervals = [] - clock = [0.0] - - 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) - monkeypatch.setattr(runner, "_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS", 0.25) - monkeypatch.setattr(runner.time, "monotonic", lambda: clock[0]) - - def fake_sleep(interval): - sleep_intervals.append(interval) - clock[0] += interval - - monkeypatch.setattr(runner.time, "sleep", fake_sleep) - - with pytest.raises(subprocess.TimeoutExpired): - runner._run_lmms_eval_process( - ["python", "-m", "lmms_eval"], - cwd=str(tmp_path), - env={}, - timeout=7.0, - ) - - assert signals == [(3456, signal.SIGTERM), (3456, signal.SIGKILL)] - assert sum(sleep_intervals) == pytest.approx(0.25) - assert all( - interval <= runner._LMMS_EVAL_PROCESS_GROUP_POLL_INTERVAL_SECONDS - for interval in sleep_intervals - ) - - -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) - raise subprocess.TimeoutExpired( - ["python", "-m", "lmms_eval"], - timeout, - output="partial stdout", - stderr="partial stderr", - ) - - process = FakeProcess() - monkeypatch.setattr(runner.subprocess, "Popen", lambda *args, **kwargs: process) - - def fake_killpg(pid, signal_number): - if signal_number == 0: - raise ProcessLookupError - signals.append((pid, signal_number)) - - monkeypatch.setattr(runner.os, "killpg", fake_killpg) - - with pytest.raises(subprocess.TimeoutExpired) as exc_info: - runner._run_lmms_eval_process( - ["python", "-m", "lmms_eval"], - cwd=str(tmp_path), - env={}, - timeout=7.0, - ) - - assert exc_info.value.output == "partial stdout" - assert exc_info.value.stderr == "partial stderr" - - assert process.communicate_timeouts == [ - 7.0, - runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, - runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, - ] - assert signals == [(5678, signal.SIGTERM), (5678, signal.SIGKILL)] - - -def test_downstream_evaluation_runs_lmms_eval_and_flattens_metrics(monkeypatch, tmp_path): +def test_downstream_evaluation_delegates_to_generic_checkpoint_evaluator(monkeypatch, tmp_path): + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() 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}, - }, - } - ) + def fake_evaluate(checkpoint_path, *, output_root, settings): + captured.update( + checkpoint=checkpoint_path, + output_root=output_root, + settings=settings, ) - return subprocess.CompletedProcess(argv, 0, stdout="", stderr="") + return {"metrics": {"ifeval.accuracy": 0.5}} - monkeypatch.setattr(runner, "_run_lmms_eval_process", fake_run) + monkeypatch.setattr(runner, "run_lmms_eval_checkpoint", fake_evaluate) 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}, - } - }, + config={"config": {"tasks": ["ifeval"], "limit": 4}}, ) source = SimpleNamespace( architecture_id="architecture", artifact_kind=ArtifactKind.CHECKPOINT, - artifact={"checkpoint": str(tmp_path / "checkpoint")}, + artifact={"checkpoint": str(checkpoint)}, ) - result = runner._downstream_evaluation( - {"puzzle_dir": str(tmp_path)}, - node, - source, - "execution", - ) + 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 result == {"metrics": {"ifeval.accuracy": 0.5}} + assert captured == { + "checkpoint": str(checkpoint), + "output_root": ( + tmp_path + / "artifacts/post_mip/nodes/lmms_eval/executions/execution/raw/architecture/lmms_eval" + ), + "settings": {"tasks": ["ifeval"], "limit": 4}, } - 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")) From 282dfd2cc0872dfb5d33e63795ea3ea71c136ab4 Mon Sep 17 00:00:00 2001 From: Johannes Rausch Date: Thu, 13 Aug 2026 01:22:43 +0200 Subject: [PATCH 2/2] Address checkpoint evaluation review feedback Forward native lmms-eval options through the convenience CLI, preserve timeout compatibility, reject byte-string arguments, and disambiguate persisted result paths. Signed-off-by: Johannes Rausch --- examples/puzzletron/README.md | 3 +- .../puzzletron/docs/checkpoint_evaluation.md | 4 +- .../puzzletron/evaluate_lmms_checkpoint.py | 12 +++- modelopt/torch/puzzletron/evaluation/lmms.py | 11 ++-- .../test_evaluate_lmms_checkpoint_cli.py | 32 ++++++++++ .../torch/puzzletron/test_lmms_evaluation.py | 61 +++++++++++++++++++ 6 files changed, 114 insertions(+), 9 deletions(-) diff --git a/examples/puzzletron/README.md b/examples/puzzletron/README.md index 373cd844640..86af42abd85 100644 --- a/examples/puzzletron/README.md +++ b/examples/puzzletron/README.md @@ -336,7 +336,8 @@ The default one-GPU smoke evaluates eight samples each from IFEval and GSM8K. Qwen 3.5 checkpoints are configured automatically. See [checkpoint evaluation](docs/checkpoint_evaluation.md) to choose tasks, run a full evaluation, find results, or override model detection. For options not -covered by the convenience command, use the native `python -m lmms_eval` CLI. +covered by the convenience command, append `--lmms-eval-args` followed by the +native lmms-eval options. ## Run with an agent diff --git a/examples/puzzletron/docs/checkpoint_evaluation.md b/examples/puzzletron/docs/checkpoint_evaluation.md index cadaa1e628a..27e2b6fab81 100644 --- a/examples/puzzletron/docs/checkpoint_evaluation.md +++ b/examples/puzzletron/docs/checkpoint_evaluation.md @@ -46,8 +46,8 @@ code. After the smoke succeeds, use `--full` with a separate output directory to evaluate the complete task datasets. Use `--timeout-seconds` if the full run needs a different limit. -For options not exposed by this convenience command, use -`python -m lmms_eval --help` and the native lmms-eval CLI. +Pass additional native options after `--lmms-eval-args`, which must be the last +wrapper option. See `python -m lmms_eval --help` for the available options. ## Results and troubleshooting diff --git a/examples/puzzletron/evaluate_lmms_checkpoint.py b/examples/puzzletron/evaluate_lmms_checkpoint.py index f427d60f413..42a84973e34 100644 --- a/examples/puzzletron/evaluate_lmms_checkpoint.py +++ b/examples/puzzletron/evaluate_lmms_checkpoint.py @@ -221,6 +221,13 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Allow reviewed checkpoint-provided Python code (disabled by default).", ) + parser.add_argument( + "--lmms-eval-args", + nargs=argparse.REMAINDER, + default=[], + metavar="ARG", + help="Forward remaining arguments to lmms-eval; this option must be last.", + ) return parser @@ -265,8 +272,11 @@ def _settings( }, "model_args": model_args, } + extra_args = list(args.lmms_eval_args) if compatibility_tasks_root is not None: - settings["extra_args"] = ["--include_path", str(compatibility_tasks_root)] + extra_args.extend(["--include_path", str(compatibility_tasks_root)]) + if extra_args: + settings["extra_args"] = extra_args return settings diff --git a/modelopt/torch/puzzletron/evaluation/lmms.py b/modelopt/torch/puzzletron/evaluation/lmms.py index a99147c89a6..1d8cfc0ed0b 100644 --- a/modelopt/torch/puzzletron/evaluation/lmms.py +++ b/modelopt/torch/puzzletron/evaluation/lmms.py @@ -85,6 +85,7 @@ DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0 _PROCESS_CLEANUP_TIMEOUT_SECONDS = 10.0 _PROCESS_GROUP_POLL_INTERVAL_SECONDS = 0.1 +_TIMEOUT_ERRORS = (TimeoutError, asyncio.TimeoutError) class LmmsEvalTimeoutError(TimeoutError): @@ -272,7 +273,7 @@ def _extra_args(settings: Mapping[str, Any]) -> list[str]: return [] if isinstance(raw, str): values = shlex.split(raw) - elif isinstance(raw, Sequence): + elif isinstance(raw, Sequence) and not isinstance(raw, (bytes, bytearray)): values = [str(item) for item in raw] else: raise TypeError("evaluation settings.extra_args must be a string or sequence") @@ -573,15 +574,15 @@ async def _run_process_async( ) try: await asyncio.wait_for(process.wait(), timeout) - except TimeoutError as error: + except _TIMEOUT_ERRORS as error: _signal_process_group(process, signal.SIGTERM) try: await asyncio.wait_for(process.wait(), _PROCESS_CLEANUP_TIMEOUT_SECONDS) - except TimeoutError: + except _TIMEOUT_ERRORS: _signal_process_group(process, signal.SIGKILL) try: await asyncio.wait_for(process.wait(), _PROCESS_CLEANUP_TIMEOUT_SECONDS) - except TimeoutError: + except _TIMEOUT_ERRORS: pass if _process_group_exists(process): _signal_process_group(process, signal.SIGKILL) @@ -701,7 +702,7 @@ def run_lmms_eval_checkpoint( summary = { "checkpoint": str(checkpoint_path), "metrics": metrics, - "result_path": str(result_path), + "raw_result_path": str(result_path), "sample_counts": sample_counts, } summary_path = _atomic_json(output / "summary.json", summary) diff --git a/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py b/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py index 9d57eeb6d22..6babb30d95b 100644 --- a/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py +++ b/tests/unit/torch/puzzletron/test_evaluate_lmms_checkpoint_cli.py @@ -53,6 +53,7 @@ def test_cli_help_explains_qwen_profile_and_native_escape_hatch(): assert "--model-profile {auto,none}" in help_text assert "Qwen 3.5" in help_text assert "reasoning_parser=qwen3" in help_text + assert "--lmms-eval-args" in help_text assert "python -m lmms_eval --help" in normalized_help @@ -225,6 +226,37 @@ def test_cli_full_run_and_runtime_overrides_are_wired(tmp_path): assert settings["model_args"]["trust_remote_code"] is True +def test_cli_forwards_native_lmms_eval_options_with_compatibility_path(tmp_path): + checkpoint = tmp_path / "teacher" + checkpoint.mkdir() + compatibility_tasks_root = tmp_path / "task-configs" + + args = evaluate_lmms_checkpoint._build_parser().parse_args( + [ + "--checkpoint", + str(checkpoint), + "--output-dir", + str(tmp_path / "results"), + "--lmms-eval-args", + "--verbosity", + "DEBUG", + "--apply_chat_template", + ] + ) + + settings = evaluate_lmms_checkpoint._settings( + args, + compatibility_tasks_root=compatibility_tasks_root, + ) + assert settings["extra_args"] == [ + "--verbosity", + "DEBUG", + "--apply_chat_template", + "--include_path", + str(compatibility_tasks_root), + ] + + def test_cli_maps_gsm8k_to_namespaced_compatibility_task(tmp_path): checkpoint = tmp_path / "teacher" checkpoint.mkdir() diff --git a/tests/unit/torch/puzzletron/test_lmms_evaluation.py b/tests/unit/torch/puzzletron/test_lmms_evaluation.py index 65d1750443e..f334fa7471c 100644 --- a/tests/unit/torch/puzzletron/test_lmms_evaluation.py +++ b/tests/unit/torch/puzzletron/test_lmms_evaluation.py @@ -15,8 +15,10 @@ """Tests for the reusable lmms-eval checkpoint backend.""" +import asyncio import json import os +import signal import sys from pathlib import Path @@ -216,6 +218,16 @@ def test_command_rejects_reserved_extra_args_setting(tmp_path, extra_args, expec assert expected in str(exc_info.value) +@pytest.mark.parametrize("extra_args", [b"--verbosity DEBUG", bytearray(b"--verbosity DEBUG")]) +def test_command_rejects_byte_string_extra_args(tmp_path, extra_args): + with pytest.raises(TypeError, match="extra_args must be a string or sequence"): + lmms._build_command( + {**_settings("ifeval"), "extra_args": extra_args}, + checkpoint="/ckpts/candidate", + output_path=tmp_path / "results", + ) + + @pytest.mark.skipif(os.name != "posix", reason="process groups are POSIX-specific") def test_timeout_kills_ignored_process_group_members(monkeypatch, tmp_path): script = ( @@ -239,6 +251,53 @@ def test_timeout_kills_ignored_process_group_members(monkeypatch, tmp_path): assert exc_info.value.stderr == "" +def test_legacy_asyncio_timeout_is_classified(monkeypatch, tmp_path): + class LegacyAsyncioTimeoutError(Exception): + pass + + class Process: + pid = 123 + returncode = None + + async def wait(self): + return self.returncode + + process = Process() + wait_calls = 0 + + async def create_subprocess_exec(*_args, **_kwargs): + return process + + async def wait_for(awaitable, _timeout): + nonlocal wait_calls + wait_calls += 1 + awaitable.close() + if wait_calls == 1: + raise LegacyAsyncioTimeoutError + process.returncode = -signal.SIGTERM + return process.returncode + + monkeypatch.setattr( + lmms, + "_TIMEOUT_ERRORS", + (TimeoutError, LegacyAsyncioTimeoutError), + ) + monkeypatch.setattr(lmms.asyncio, "create_subprocess_exec", create_subprocess_exec) + monkeypatch.setattr(lmms.asyncio, "wait_for", wait_for) + monkeypatch.setattr(lmms, "_signal_process_group", lambda *_args: None) + monkeypatch.setattr(lmms, "_process_group_exists", lambda _process: False) + + with pytest.raises(lmms.LmmsEvalTimeoutError): + asyncio.run( + lmms._run_process_async( + [sys.executable, "-c", "pass"], + cwd=str(tmp_path), + env=os.environ.copy(), + timeout=1.0, + ) + ) + + def test_run_checkpoint_flattens_metrics_and_preserves_artifacts(monkeypatch, tmp_path): checkpoint = tmp_path / "checkpoint" checkpoint.mkdir() @@ -275,6 +334,8 @@ def fake_run(argv, *, cwd, env, timeout): assert Path(result["stderr_path"]).read_text() == "" summary = json.loads(Path(result["result_path"]).read_text()) assert summary["checkpoint"] == str(checkpoint.resolve()) + assert summary["raw_result_path"] == result["raw_result_path"] + assert "result_path" not in summary assert summary["sample_counts"] == {"gsm8k": 4.0, "ifeval": 4.0}