From de56ffd49055df2d2010d79f55e8c4183f35b93f Mon Sep 17 00:00:00 2001 From: "coderabbitai[bot]" <136622811+coderabbitai[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 10:29:23 +0000 Subject: [PATCH] fix: apply CodeRabbit auto-fixes Fixed 5 file(s) based on 9 unresolved review comments. Co-authored-by: CodeRabbit --- modelopt/torch/puzzletron/post_mip/runner.py | 41 ++++++++++++------- modelopt/torch/puzzletron/stages/pipeline.py | 41 ++++++++----------- .../subblock_stats/calc_subblock_stats.py | 13 +++--- puzzletron_setup/wizard.py | 9 +++- .../puzzletron/test_sparse_runtime_stats.py | 5 ++- 5 files changed, 63 insertions(+), 46 deletions(-) diff --git a/modelopt/torch/puzzletron/post_mip/runner.py b/modelopt/torch/puzzletron/post_mip/runner.py index 3d23e0d05b0..0511091b87e 100644 --- a/modelopt/torch/puzzletron/post_mip/runner.py +++ b/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( diff --git a/modelopt/torch/puzzletron/stages/pipeline.py b/modelopt/torch/puzzletron/stages/pipeline.py index b723526cc1b..8561b1cc454 100644 --- a/modelopt/torch/puzzletron/stages/pipeline.py +++ b/modelopt/torch/puzzletron/stages/pipeline.py @@ -275,16 +275,18 @@ 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): @@ -292,34 +294,27 @@ def _has_runtime_measurement( 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 diff --git a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py b/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py index 9677b5c55c1..127f4b2d8c4 100644 --- a/modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py +++ b/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, ) diff --git a/puzzletron_setup/wizard.py b/puzzletron_setup/wizard.py index c6140623757..4b092569c34 100644 --- a/puzzletron_setup/wizard.py +++ b/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": diff --git a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py b/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py index 827d1b635ff..019f489cab0 100644 --- a/tests/unit/torch/puzzletron/test_sparse_runtime_stats.py +++ b/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():