Add Puzzletron downstream evaluation - #2034
Conversation
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds ChangesDownstream evaluation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant SetupWizard
participant CampaignCompiler
participant PostMIPRunner
participant LMMS-Eval
SetupWizard->>CampaignCompiler: configure downstream_evaluation topology
CampaignCompiler->>PostMIPRunner: allocate compiled evaluation stage
PostMIPRunner->>LMMS-Eval: execute checkpoint evaluation
LMMS-Eval-->>PostMIPRunner: return JSON task metrics
PostMIPRunner-->>CampaignCompiler: publish flattened evaluation results
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🧹 Nitpick comments (2)
modelopt/torch/puzzletron/post_mip/reporting.py (1)
360-367: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winParameterize the heading instead of post-rendering string replacement.
The
.replace()silently no-ops ifrender_evaluation_report's<h3>text is ever reworded, leaving downstream nodes labelled "Candidate evaluation". Aheadingkeyword also fixes the stalearia-label='Candidate evaluation metrics'on the plot div.♻️ Proposed refactor
-def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str: - """Render lmms-eval task metrics for downstream-evaluation nodes.""" - - return render_evaluation_report(section_id, payload).replace( - "<h3>Candidate evaluation</h3>", - "<h3>Downstream evaluation</h3>", - 1, - ) +def render_downstream_evaluation_report(section_id: str, payload: Mapping[str, Any]) -> str: + """Render lmms-eval task metrics for downstream-evaluation nodes.""" + + return render_evaluation_report(section_id, payload, heading="Downstream evaluation")with
render_evaluation_reportgainingheading: str = "Candidate evaluation"and using it for both the<h3>and the plotaria-label.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@modelopt/torch/puzzletron/post_mip/reporting.py` around lines 360 - 367, Update render_evaluation_report to accept a heading keyword defaulting to "Candidate evaluation", and use it for both the evaluation section’s h3 text and plot div aria-label. Change render_downstream_evaluation_report to pass "Downstream evaluation" directly and remove the post-render .replace() workaround.tests/unit/torch/puzzletron/test_post_mip_runner.py (1)
229-282: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering the new timeout-classification path.
These tests cover the success path well, but the new
downstream_evaluationbranch inrun_post_mip_node_shard(runner.pyLines 942-965 —status="timed_out",timeout_field="timeout_seconds",3600default) is untested. A fakesubprocess.runraisingsubprocess.TimeoutExpiredwould pin that mapping cheaply.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unit/torch/puzzletron/test_post_mip_runner.py` around lines 229 - 282, Add a unit test covering the TimeoutExpired path in run_post_mip_node_shard when downstream_evaluation invokes subprocess.run. Make the fake subprocess.run raise subprocess.TimeoutExpired, then assert the shard result maps to status="timed_out", uses timeout_field="timeout_seconds", and applies the default timeout_seconds value of 3600.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@modelopt/torch/puzzletron/post_mip/runner.py`:
- Around line 611-613: Update the loop in the model-argument derivation logic
around _LMMS_EVAL_MODEL_ARG_FIELDS to iterate those fields in a deterministic
order, such as a sorted sequence, before populating derived. Preserve the
existing filtering and value-copying behavior so generated --model_args and
command.json argv remain stable.
- Around line 689-696: Update the timeout selection in the environment/command
preparation flow to use a finite default of 3600 seconds when both
timeout_seconds and timeout are unset. Preserve explicit configured timeout
values and continue returning the selected value as a float for subprocess.run.
In `@puzzletron_setup/wizard.py`:
- Around line 741-744: Update the lmms-eval task prompt in the wizard flow to
use a non-empty comma-separated task validator, matching the validator behavior
already used by the v2 wizard. Reject blank input and prevent serialization of
an empty tasks list while preserving the existing default task values.
- Around line 1069-1075: Update the downstream_evaluation handling around
_ask_downstream_evaluation_config so available_metrics only includes metrics for
tasks selected in the node configuration. Do not unconditionally append the
gsm8k.exact_match metric; preserve metric availability for selected tasks such
as gsm8k while excluding it when only ifeval or other tasks are selected.
---
Nitpick comments:
In `@modelopt/torch/puzzletron/post_mip/reporting.py`:
- Around line 360-367: Update render_evaluation_report to accept a heading
keyword defaulting to "Candidate evaluation", and use it for both the evaluation
section’s h3 text and plot div aria-label. Change
render_downstream_evaluation_report to pass "Downstream evaluation" directly and
remove the post-render .replace() workaround.
In `@tests/unit/torch/puzzletron/test_post_mip_runner.py`:
- Around line 229-282: Add a unit test covering the TimeoutExpired path in
run_post_mip_node_shard when downstream_evaluation invokes subprocess.run. Make
the fake subprocess.run raise subprocess.TimeoutExpired, then assert the shard
result maps to status="timed_out", uses timeout_field="timeout_seconds", and
applies the default timeout_seconds value of 3600.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5b5ddf46-03f8-46cd-9016-a735148cad1d
📒 Files selected for processing (19)
CHANGELOG.rstexamples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yamlmodelopt/torch/puzzletron/orchestration/adapters/post_mip.pymodelopt/torch/puzzletron/orchestration/compiler.pymodelopt/torch/puzzletron/orchestration/progress.pymodelopt/torch/puzzletron/post_mip/builtin.pymodelopt/torch/puzzletron/post_mip/reporting.pymodelopt/torch/puzzletron/post_mip/runner.pypuzzletron_setup/bundle.pypuzzletron_setup/v2/parallel_validation.pypuzzletron_setup/v2/post_mip.pypuzzletron_setup/v2/validation.pypuzzletron_setup/v2/wizard.pypuzzletron_setup/wizard.pytests/unit/torch/puzzletron/test_orchestration_compiler.pytests/unit/torch/puzzletron/test_post_mip_runner.pytests/unit/torch/puzzletron/test_setup_bundle.pytests/unit/torch/puzzletron/test_setup_v2_post_mip.pytests/unit/torch/puzzletron/test_setup_v2_state_validation.py
| for key in _LMMS_EVAL_MODEL_ARG_FIELDS: | ||
| if key in settings: | ||
| derived[key] = settings[key] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Iterating a frozenset makes the generated --model_args order non-deterministic.
_LMMS_EVAL_MODEL_ARG_FIELDS is a frozenset of strings, so iteration order varies between processes under hash randomization. That leaks into derived/merged insertion order and therefore into the --model_args string and the recorded command.json argv, which the docstring on _lmms_eval_command claims is deterministic.
♻️ Proposed fix
- 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]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for key in _LMMS_EVAL_MODEL_ARG_FIELDS: | |
| if key in settings: | |
| derived[key] = settings[key] | |
| for key in sorted(_LMMS_EVAL_MODEL_ARG_FIELDS): | |
| if key in settings: | |
| derived[key] = settings[key] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/puzzletron/post_mip/runner.py` around lines 611 - 613, Update
the loop in the model-argument derivation logic around
_LMMS_EVAL_MODEL_ARG_FIELDS to iterate those fields in a deterministic order,
such as a sorted sequence, before populating derived. Preserve the existing
filtering and value-copying behavior so generated --model_args and command.json
argv remain stable.
| 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", settings.get("timeout")) | ||
| return argv, env, (float(timeout) if timeout is not None else None) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether the setup/validation layer always emits a timeout for downstream_evaluation nodes
rg -nP -C3 'timeout_seconds' --type=py | rg -n -C3 'downstream'
fd -t f 'post_mip.py' -x rg -nP -C4 'downstream_evaluation|timeout_seconds' {}Repository: NVIDIA/Model-Optimizer
Length of output: 243
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant runner code and search for timeout handling in the surrounding pipeline.
FILE="modelopt/torch/puzzletron/post_mip/runner.py"
wc -l "$FILE"
sed -n '620,710p' "$FILE"
printf '\n--- search timeout callers ---\n'
rg -n "timeout_seconds|timeout" modelopt/torch/puzzletron -g '*.py'
printf '\n--- search downstream_evaluation ---\n'
rg -n "downstream_evaluation" modelopt/torch/puzzletron -g '*.py'Repository: NVIDIA/Model-Optimizer
Length of output: 16252
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the specific runner and nearby configuration objects for timeout defaults.
FILE="modelopt/torch/puzzletron/post_mip/runner.py"
sed -n '1,260p' "$FILE"
printf '\n--- mid file ---\n'
sed -n '260,620p' "$FILE"
printf '\n--- tail ---\n'
sed -n '710,1040p' "$FILE"
printf '\n--- timeout-related definitions in puzzletron ---\n'
rg -n "timeout_seconds|timeout" modelopt/torch/puzzletron -g '*.py'Repository: NVIDIA/Model-Optimizer
Length of output: 47964
Add a default timeout here. When neither timeout_seconds nor timeout is set, subprocess.run(...) gets None, so a hung lmms-eval process can block the worker indefinitely. The 3600 fallback only applies after a timeout exception is raised, so it does not bound the run itself.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@modelopt/torch/puzzletron/post_mip/runner.py` around lines 689 - 696, Update
the timeout selection in the environment/command preparation flow to use a
finite default of 3600 seconds when both timeout_seconds and timeout are unset.
Preserve explicit configured timeout values and continue returning the selected
value as a float for subprocess.run.
| tasks = prompts.text( | ||
| "lmms-eval tasks (comma-separated):", | ||
| default=str(defaults.get("tasks", "ifeval,gsm8k")), | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Reject empty lmms-eval task selections.
Blank input is serialized as tasks: [], producing a downstream-evaluation node with nothing to run. Add a non-empty comma-separated task validator, as the v2 wizard does.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@puzzletron_setup/wizard.py` around lines 741 - 744, Update the lmms-eval task
prompt in the wizard flow to use a non-empty comma-separated task validator,
matching the validator behavior already used by the v2 wizard. Reject blank
input and prevent serialization of an empty tasks list while preserving the
existing default task values.
Source: Coding guidelines
| elif node_type == "downstream_evaluation": | ||
| node["config"] = _ask_downstream_evaluation_config( | ||
| prompts, | ||
| detailed=detailed, | ||
| moe=moe, | ||
| ) | ||
| available_metrics.append(f"{node_id}.gsm8k.exact_match") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Only expose metrics for selected tasks.
Selecting only ifeval still adds <node>.gsm8k.exact_match to filter choices. A later filter can therefore reference a metric this node will never emit.
Proposed fix
- node["config"] = _ask_downstream_evaluation_config(
+ config = _ask_downstream_evaluation_config(
prompts,
detailed=detailed,
moe=moe,
)
- available_metrics.append(f"{node_id}.gsm8k.exact_match")
+ node["config"] = config
+ if "gsm8k" in config["tasks"]:
+ available_metrics.append(f"{node_id}.gsm8k.exact_match")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| elif node_type == "downstream_evaluation": | |
| node["config"] = _ask_downstream_evaluation_config( | |
| prompts, | |
| detailed=detailed, | |
| moe=moe, | |
| ) | |
| available_metrics.append(f"{node_id}.gsm8k.exact_match") | |
| elif node_type == "downstream_evaluation": | |
| config = _ask_downstream_evaluation_config( | |
| prompts, | |
| detailed=detailed, | |
| moe=moe, | |
| ) | |
| node["config"] = config | |
| if "gsm8k" in config["tasks"]: | |
| available_metrics.append(f"{node_id}.gsm8k.exact_match") |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@puzzletron_setup/wizard.py` around lines 1069 - 1075, Update the
downstream_evaluation handling around _ask_downstream_evaluation_config so
available_metrics only includes metrics for tasks selected in the node
configuration. Do not unconditionally append the gsm8k.exact_match metric;
preserve metric availability for selected tasks such as gsm8k while excluding it
when only ifeval or other tasks are selected.
|
Note Docstrings generation - SUCCESS |
Docstrings generation was requested by @grzegorz-k-karch. * #2034 (comment) The following files were modified: * `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` * `modelopt/torch/puzzletron/orchestration/compiler.py` * `modelopt/torch/puzzletron/orchestration/progress.py` * `modelopt/torch/puzzletron/post_mip/builtin.py` * `modelopt/torch/puzzletron/post_mip/reporting.py` * `modelopt/torch/puzzletron/post_mip/runner.py` * `puzzletron_setup/bundle.py` * `puzzletron_setup/v2/validation.py` * `puzzletron_setup/v2/wizard.py` * `puzzletron_setup/wizard.py`
Docstrings generation was requested by @grzegorz-k-karch. * #2034 (comment) The following files were modified: * `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` * `modelopt/torch/puzzletron/orchestration/compiler.py` * `modelopt/torch/puzzletron/orchestration/progress.py` * `modelopt/torch/puzzletron/post_mip/builtin.py` * `modelopt/torch/puzzletron/post_mip/reporting.py` * `modelopt/torch/puzzletron/post_mip/runner.py` * `puzzletron_setup/bundle.py` * `puzzletron_setup/v2/validation.py` * `puzzletron_setup/v2/wizard.py` * `puzzletron_setup/wizard.py` <details> <summary>These files were kept as they were</summary> * `tests/unit/torch/puzzletron/test_orchestration_compiler.py` * `tests/unit/torch/puzzletron/test_post_mip_runner.py` * `tests/unit/torch/puzzletron/test_setup_bundle.py` * `tests/unit/torch/puzzletron/test_setup_v2_post_mip.py` * `tests/unit/torch/puzzletron/test_setup_v2_state_validation.py` </details> <details> <summary>These file types are not supported</summary> * `CHANGELOG.rst` * `examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml` </details> <details> <summary>ℹ️ Note</summary><blockquote> CodeRabbit cannot perform edits on its own pull requests yet. </blockquote></details> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
1 participant
What does this PR do?
Type of change: new feature, new example, new tests
Adds a Puzzletron dynamic post-MIP
downstream_evaluationnode that evaluates realized checkpoint artifacts withlmms-evalusing the vLLM backend. The stage builds a shell-free lmms-eval argv, derives vLLM model arguments from the node topology, records command metadata, parses lmms-eval JSON results, and publishes flattened task metrics into the post-MIP ledger/reporting path.This also extends the Puzzletron setup/orchestration flow so downstream evaluation can be configured with task names, sample limit, batch size, timeout, and vLLM topology defaults. A new opt-in Nemotron-3 Nano 30B A3B BF16 example flow evaluates the best
runtime-075realized model withifevalandgsm8k.Usage
The lmms-eval CLI shape follows the upstream vLLM examples from https://github.com/EvolvingLMMs-Lab/lmms-eval.
Testing
python -m py_compile modelopt/torch/puzzletron/orchestration/adapters/post_mip.py modelopt/torch/puzzletron/orchestration/compiler.py modelopt/torch/puzzletron/orchestration/progress.py modelopt/torch/puzzletron/post_mip/builtin.py modelopt/torch/puzzletron/post_mip/reporting.py modelopt/torch/puzzletron/post_mip/runner.py puzzletron_setup/bundle.py puzzletron_setup/v2/parallel_validation.py puzzletron_setup/v2/post_mip.py puzzletron_setup/v2/validation.py puzzletron_setup/v2/wizard.py puzzletron_setup/wizard.py tests/unit/torch/puzzletron/test_orchestration_compiler.py tests/unit/torch/puzzletron/test_post_mip_runner.py tests/unit/torch/puzzletron/test_setup_bundle.py tests/unit/torch/puzzletron/test_setup_v2_post_mip.py tests/unit/torch/puzzletron/test_setup_v2_state_validation.pygit diff --checkpytestis not installed.Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Assumes
lmms-evaland vLLM are already installed in the target execution virtual environment.Summary by CodeRabbit