fix: CodeRabbit auto-fixes for PR #2104

modelopt/torch/puzzletron/post_mip/runner.py

@@ -45,6 +45,8 @@
     "run_post_mip_node_shard",
 ]
 
+_DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0
+
 
 def _puzzle_dir(config: Mapping[str, Any]) -> Path:
     return Path(config.get("puzzle_dir") or (config.get("experiment") or {})["dir"])
@@ -672,19 +674,18 @@ def _merge_lmms_eval_model_args(settings: Mapping[str, Any], checkpoint: str) ->
                 ],
             }
         )
-    for key in _LMMS_EVAL_MODEL_ARG_FIELDS:
+    for key in sorted(_LMMS_EVAL_MODEL_ARG_FIELDS):
         if key in settings:
             derived[key] = settings[key]
 
     if isinstance(raw, str):
         prefix = raw.strip().strip(",")
         suffix = _model_arg_string(derived)
-        return ",".join(part for part in (prefix, suffix) if part)
+        return ",".join(part for part in (suffix, prefix) 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")
     merged = dict(raw or {})
-    for key, value in derived.items():
-        merged.setdefault(key, value)
+    merged.update(derived)
     return _model_arg_string(merged)
 
 
@@ -779,7 +780,7 @@ def _lmms_eval_command(
     if settings.get("cache_dir") is not None:
         env.setdefault("LMMS_EVAL_HOME", str(settings["cache_dir"]))
     timeout = settings.get("timeout_seconds", settings.get("timeout"))
-    return argv, env, (float(timeout) if timeout is not None else None)
+    return argv, env, (float(timeout) if timeout is not None else _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS)
 
 
 def _metric_key(value: Any) -> str:
@@ -949,15 +950,25 @@ def _downstream_evaluation(
     )
     # Campaign config controls the executable and arguments, but subprocess receives
     # an argv list directly; no shell parsing is involved.
-    result = subprocess.run(
-        argv,
-        cwd=str(output),
-        env=env,
-        capture_output=True,
-        text=True,
-        timeout=timeout,
-        check=False,
-    )
+    try:
+        result = subprocess.run(
+            argv,
+            cwd=str(output),
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+            check=False,
+        )
+    except subprocess.TimeoutExpired as timeout_error:
+        timeout_result = subprocess.CompletedProcess(
+            args=argv,
+            returncode=-1,
+            stdout=timeout_error.stdout.decode("utf-8", errors="replace") if timeout_error.stdout else "",
+            stderr=timeout_error.stderr.decode("utf-8", errors="replace") if timeout_error.stderr else "",
+        )
+        _write_lmms_eval_streams(output, timeout_result)
+        raise
     stream_paths = _write_lmms_eval_streams(output, result)
     if result.returncode:
         tail = _lmms_eval_output_tail(result)
@@ -1165,7 +1176,7 @@ def run_post_mip_node_shard(
                         timeout_field = "timeout_seconds"
                     elif not isinstance(error, subprocess.TimeoutExpired):
                         timeout_field = "readiness_timeout"
-                    default_timeout = 3600 if node.node_type == "downstream_evaluation" else (
+                    default_timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS if node.node_type == "downstream_evaluation" else (
                         600 if timeout_field == "benchmark_timeout" else 1200
                     )
                     row["timeout_seconds"] = float(

modelopt/torch/puzzletron/stages/pipeline.py

@@ -275,51 +275,46 @@ def _has_runtime_measurement(
 ) -> bool:
     """
     Determine whether a statistics file contains a compatible runtime measurement.
-    
+
     Parameters:
     	path (Path): Statistics file to inspect.
     	hidden_width (int): Model hidden width expected by the measurement.
     	measurement (Any): Runtime measurement configuration to match.
     	allow_missing_workload_id (bool): Whether entries without a workload identifier may match.
-    
+
     Returns:
     	bool: `True` if a compatible runtime measurement is present, `False` otherwise.
     """
+    from ..subblock_stats.calc_subblock_stats import _runtime_reuse_key, _runtime_stats_identity
+
     try:
         payload = json.loads(path.read_text())
     except (OSError, ValueError):
         return False
     if not isinstance(payload, list):
         return False
     expected_backend = (measurement.runtime_stats or {}).get("backend")
+    requested_key = _runtime_reuse_key(
+        width=hidden_width,
+        batch_size=measurement.batch_size,
+        prefill_seq_len=measurement.prefill_seq_len,
+        generation_seq_len=measurement.generation_seq_len,
+        runtime_stats_config=measurement.runtime_stats or {},
+    )
     for entry in payload:
         if not isinstance(entry, dict):
             continue
         args = entry.get("args") or {}
         if not isinstance(args, dict) or args.get("runtime_stats") is not True:
             continue
-        if int(args.get("n_embd", -1)) != int(hidden_width):
-            continue
-        if args.get("weights_dtype") != "torch.bfloat16":
-            continue
-        if int(args.get("batch_size", -1)) != int(measurement.batch_size):
-            continue
-        if int(args.get("prefill_seq_len", -1)) != int(measurement.prefill_seq_len):
-            continue
-        if int(args.get("generation_seq_len", -1)) != int(measurement.generation_seq_len):
-            continue
-        if int(args.get("max_num_seqs", -1)) != int(measurement.max_num_seqs):
-            continue
-        if args.get("runtime_granularity", "subblock") != measurement.granularity:
-            continue
-        if expected_backend is not None and args.get("runtime_backend") != expected_backend:
-            continue
-        workload_id = args.get("workload_id")
-        if workload_id is None and not allow_missing_workload_id:
-            continue
-        if workload_id is not None and workload_id != measurement.measurement_id:
+        persisted_key = _runtime_stats_identity(
+            args,
+            fallback_workload_id=measurement.measurement_id if allow_missing_workload_id else None,
+        )
+        if persisted_key is None:
             continue
-        return True
+        if persisted_key == requested_key:
+            return True
     return False
 
 

modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py

@@ -289,20 +289,23 @@ def _runtime_reuse_key_from_args(
 
     if not args.get("runtime_stats") or args.get("weights_dtype") != str(torch.bfloat16):
         return None
+    required_fields = ["n_embd", "batch_size", "prefill_seq_len", "generation_seq_len"]
+    if any(args.get(field) is None for field in required_fields):
+        return None
     workload_id = args.get("workload_id")
     if workload_id is None:
         workload_id = fallback_workload_id
     return (
         int(args["n_embd"]),
         int(args["batch_size"]),
-        int(args.get("prefill_seq_len")),
-        int(args.get("generation_seq_len")),
+        int(args["prefill_seq_len"]),
+        int(args["generation_seq_len"]),
         args.get("max_num_seqs"),
         args.get("runtime_granularity", "subblock"),
         args.get("runtime_backend"),
-        args.get("num_iters"),
-        args.get("num_warmup_iters"),
-        args.get("repeat_block_n_times"),
+        args.get("num_iters", 30),
+        args.get("num_warmup_iters", 10),
+        max(2, int(args.get("repeat_block_n_times", 10))),
         _freeze_stats_args(args.get("vllm_args")),
         workload_id,
     )

puzzletron_setup/wizard.py

@@ -764,9 +764,13 @@ def _ask_downstream_evaluation_config(
     """
 
     defaults = defaults or {}
+    default_tasks = defaults.get("tasks", "ifeval,gsm8k")
+    if isinstance(default_tasks, list):
+        default_tasks = ",".join(default_tasks)
     tasks = prompts.text(
         "lmms-eval tasks (comma-separated):",
-        default=str(defaults.get("tasks", "ifeval,gsm8k")),
+        default=str(default_tasks),
+        validate=lambda value: bool(str(value).strip()) or "Enter at least one task.",
     )
     limit = prompts.integer(
         "lmms-eval sample limit:",
@@ -1127,7 +1131,8 @@ def _custom_flow(
                 detailed=detailed,
                 moe=moe,
             )
-            available_metrics.append(f"{node_id}.gsm8k.exact_match")
+            for task_name in node["config"].get("tasks", []):
+                available_metrics.append(f"{node_id}.{task_name}.strict-match")
         elif node_type == "global_kd":
             node["config"] = {"max_steps": prompts.integer("Global KD steps:", default=128)}
         elif node_type == "ptq":

tests/unit/torch/puzzletron/test_sparse_runtime_stats.py

@@ -1180,6 +1180,10 @@ def test_width_scenario_runtime_stats_reuse_root_measurement(tmp_path, monkeypat
                         "generation_seq_len": 1024,
                         "max_num_seqs": 1,
                         "n_embd": 2688,
+                        "num_iters": 30,
+                        "num_warmup_iters": 10,
+                        "repeat_block_n_times": 10,
+                        "vllm_args": [],
                         "workload_id": "serving-default",
                     },
                     "subblocks": [],
@@ -1327,7 +1331,6 @@ def test_runtime_stats_resume_signature_includes_workload_id():
         **kwargs,
         runtime_workload_id="different-workload",
     )
-    assert hydra_cfg.calc_subblock_stats.merge_with_existing_stats is False
 
 
 def test_sparse_runtime_selection_is_unique_and_layer_independent():