Merge latest Puzzletron v2 target

Signed-off-by: Johannes Rausch <jrausch@nvidia.com>

Author: j-rausch

examples/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+"""ModelOpt examples."""

examples/puzzletron/README.md

@@ -284,6 +284,7 @@ import os
 from packaging.version import Version
 
 import aiperf
+import lmms_eval
 import modelopt
 import nemo_automodel
 import torch
@@ -298,6 +299,7 @@ for package in (
     "vllm",
     "nemo-automodel",
     "aiperf",
+    "lmms-eval",
     "nvidia-modelopt",
 ):
     print(package, metadata.version(package))
@@ -312,6 +314,7 @@ assert Version(metadata.version("torchvision")).release == Version(
     ci_environment["torchvision"]
 ).release
 assert transformers.__version__ == ci_environment["transformers"]
+assert metadata.version("lmms-eval") == ci_environment["lmms_eval"]
 assert Version(metadata.version("nemo-automodel")).base_version == (
     ci_environment["nemo_automodel"]["base_version"]
 )
@@ -612,25 +615,12 @@ After `mip`, prepare one deduplicated online-evaluation plan. Repeat
 `--profile-id` for every configured profile; aliases ensure that an identical
 architecture is evaluated once while remaining visible in every profile.
 
-Downstream `lmms-eval` nodes use a separate evaluator Python. The reproducible
-example path is pinned to `lmms-eval==0.7.2` by
-`examples/puzzletron/requirements-lmms-eval.txt` and recorded in
-`ci_environment.json`. Keep this separate from the Puzzletron runtime
-environment because `lmms-eval==0.7.2` pins `wandb==0.25.0`, while the pinned
-AutoModel build requires a newer `wandb`.
-
-```bash
-python3 -m venv /workspace/.venv-lmms-eval
-source /workspace/.venv-lmms-eval/bin/activate
-python -m pip install --upgrade pip "setuptools>=80,<81" wheel packaging
-VLLM_USE_PRECOMPILED=1 VLLM_PRECOMPILED_WHEEL_VARIANT=cu129 \
-  python -m pip install --no-build-isolation -e "${VLLM_ROOT}"
-python -m pip install -r "${MODEL_OPT_ROOT}/examples/puzzletron/requirements-lmms-eval.txt"
-python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"'
-deactivate
-
-export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python
-```
+Downstream `lmms-eval` nodes run in the standard Puzzletron worker environment.
+The reproducible example path is pinned by
+`examples/puzzletron/requirements.txt` to an upstream 0.7.0 snapshot that
+packages the task definitions and accepts the newer `wandb` required by the
+pinned AutoModel build. The expected package version is recorded in
+`ci_environment.json`.
 
 ```bash
 python examples/puzzletron/run_profile_online_evaluation.py \

examples/puzzletron/ci_environment.json

@@ -5,7 +5,7 @@
   "torch": "2.11.0",
   "torchvision": "0.26.0",
   "transformers": "5.8.1",
-  "lmms_eval": "0.7.2",
+  "lmms_eval": "0.7.0",
   "nemo_automodel": {
     "base_version": "0.5.0",
     "repository": "https://github.com/Separius/Automodel.git",

examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml

@@ -5,7 +5,7 @@ defaults:
   - _self_
 
 # Opt-in downstream lmms-eval workflow for the realized runtime-075 candidate.
-# Use PUZZLETRON_LMMS_EVAL_PYTHON to point at the isolated pinned evaluator env.
+# Install examples/puzzletron/requirements.txt to use the pinned lmms-eval version.
 # Non-empty post_mip.flows replaces the legacy post-MIP tail in the v2 orchestrator.
 zero_shot_evaluation:
   enabled: false
@@ -39,10 +39,6 @@ post_mip:
           type: downstream_evaluation
           input: materialized
           config:
-            command_prefix:
-              - ${oc.env:PUZZLETRON_LMMS_EVAL_PYTHON}
-              - -m
-              - lmms_eval
             model: vllm
             checkpoint_arg: model
             tasks:

examples/puzzletron/docs/post_mip_pipeline.md

@@ -91,9 +91,10 @@ post_mip:
             max_steps: 2048
 ```
 
-Evaluation nodes publish every finite metric they produce. They do not declare
-metric lists or cases. Later filters reference metrics as `mip.<metric>` or
-`<node_id>.<metric>`.
+Evaluation nodes publish every finite metric they produce. Downstream lmms-eval
+nodes publish results only after every configured task, or every leaf of a
+configured task group, has numeric metrics and positive effective sample counts.
+Later filters reference metrics as `mip.<metric>` or `<node_id>.<metric>`.
 
 ## Node types
 
@@ -114,6 +115,13 @@ metric lists or cases. Later filters reference metrics as `mip.<metric>` or
 Nodes that require checkpoints never materialize implicitly. Add a `materialize`
 node where the transition is needed.
 
+## Add downstream evaluation to an existing campaign
+
+Keep the campaign's `puzzle_dir` and add a `post_mip.flows` entry whose source
+selects the completed MIP run. See the
+[lmms-eval run configuration](../configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml)
+for a complete filter, materialization, and downstream-evaluation flow.
+
 ## Lineage and model source
 
 Architectures are deduplicated across MIP and homogeneous results, while all MIP
@@ -133,21 +141,15 @@ from the original candidate.
 
 ## Downstream evaluation
 
-`downstream_evaluation` runs `python -m lmms_eval` as a subprocess through
-`command_prefix`. 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. Install the pinned evaluator into an isolated environment
-rather than the Puzzletron runtime environment, because `lmms-eval==0.7.2` pins
-`wandb==0.25.0` and the pinned AutoModel build requires a newer `wandb`:
+`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
-python3 -m venv /workspace/.venv-lmms-eval
-source /workspace/.venv-lmms-eval/bin/activate
-python -m pip install -r examples/puzzletron/requirements-lmms-eval.txt
-python -c 'import importlib.metadata as m; assert m.version("lmms-eval") == "0.7.2"'
-deactivate
-
-export PUZZLETRON_LMMS_EVAL_PYTHON=/workspace/.venv-lmms-eval/bin/python
+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

examples/puzzletron/requirements.txt

@@ -1,4 +1,5 @@
 aiohttp>=3.9,<4
+lmms-eval @ git+https://github.com/EvolvingLMMs-Lab/lmms-eval.git@15c32bfec165df13c269ddd3cda03b2ed9137825
 math-verify
 ray
 # Transformers comes from nvidia-modelopt[hf]; CPU CI overlays the exact

modelopt/torch/puzzletron/post_mip/runner.py

@@ -28,6 +28,7 @@
 import signal
 import subprocess
 import sys
+import time
 import traceback
 import uuid
 from contextlib import contextmanager
@@ -561,6 +562,9 @@ def _aiperf(
 )
 _LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS = frozenset(
     {
+        "--batch-size",
+        "--batch_size",
+        "--model",
         "--model_args",
         "--model-args",
         "--output_path",
@@ -570,6 +574,7 @@ def _aiperf(
 )
 _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:
@@ -688,6 +693,8 @@ def _model_arg_string(values: Mapping[str, Any]) -> str:
 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)
@@ -766,14 +773,17 @@ def _lmms_eval_command(
     *,
     checkpoint: str,
     output_path: Path,
-) -> tuple[list[str], dict[str, str], float | None]:
+) -> 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",
-        str(settings.get("model", "vllm")),
+        model,
         "--model_args",
         _merge_lmms_eval_model_args(settings, checkpoint),
         "--tasks",
@@ -816,10 +826,15 @@ def _lmms_eval_command(
             env[str(key)] = str(value)
     if settings.get("cache_dir") is not None:
         env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"]))
-    timeout = settings.get("timeout_seconds", settings.get("timeout"))
+    timeout = settings.get("timeout_seconds")
+    if timeout is None:
+        timeout = settings.get("timeout")
     if timeout is None:
         timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS
-    return argv, env, float(timeout)
+    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:
@@ -1002,12 +1017,22 @@ def _lmms_eval_process_group_exists(process: subprocess.Popen[str]) -> bool:
     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 | None,
+    timeout: float,
 ) -> subprocess.CompletedProcess[str]:
     process = subprocess.Popen(
         argv,
@@ -1026,15 +1051,21 @@ def _run_lmms_eval_process(
             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,

noxfile.py

@@ -110,6 +110,7 @@ def puzzletron_v2(session):
         "torch": PUZZLETRON_V2_CI_ENVIRONMENT["torch"],
         "torchvision": PUZZLETRON_V2_CI_ENVIRONMENT["torchvision"],
         "transformers": PUZZLETRON_V2_CI_ENVIRONMENT["transformers"],
+        "lmms-eval": PUZZLETRON_V2_CI_ENVIRONMENT["lmms_eval"],
         "nemo-automodel": PUZZLETRON_V2_AUTOMODEL_SOURCE["base_version"],
     }
     session.run(
@@ -125,6 +126,7 @@ def puzzletron_v2(session):
             "'torch': Version(version('torch')).base_version, "
             "'torchvision': Version(version('torchvision')).base_version, "
             "'transformers': Version(version('transformers')).base_version, "
+            "'lmms-eval': Version(version('lmms-eval')).base_version, "
             "'nemo-automodel': Version(version('nemo-automodel')).base_version}; "
             "mismatches = {name: (actual[name], expected_version) "
             "for name, expected_version in expected.items() "

tests/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+
+"""ModelOpt test suite."""