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 <jrausch@nvidia.com>

Author: j-rausch

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
 

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
 

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
 
 

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)

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

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}