Skip to content

Add downstream eval to puzzletron v2 - #2104

Open
grzegorz-k-karch wants to merge 15 commits into
feature/puzzletron_v2from
gkarch/add_downstream_eval
Open

Add downstream eval to puzzletron v2#2104
grzegorz-k-karch wants to merge 15 commits into
feature/puzzletron_v2from
gkarch/add_downstream_eval

Conversation

@grzegorz-k-karch

@grzegorz-k-karch grzegorz-k-karch commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: new feature, new example, new tests

Adds a Puzzletron dynamic post-MIP downstream_evaluation node that evaluates realized checkpoint artifacts with lmms-eval through the vLLM backend. The stage b
uilds a deterministic shell-free lmms-eval command, maps the realized checkpoint into vLLM model arguments, derives distributed topology settings from the node c
onfig, captures command metadata/stdout/stderr, parses JSON metrics, and publishes flattened task metrics into the post-MIP ledger and report path.

This also extends the Puzzletron setup and orchestration flow so downstream evaluation can be configured with task names, sample limit, batch size, timeout, and vL
LM topology defaults. The orchestration layer now labels downstream evaluation stages in progress/dashboard output and includes resource allocation handling for th
ose stages.

The PR includes a new opt-in Nemotron-3 Nano 30B A3B BF16 example flow that evaluates the best runtime-075 realized model with ifeval and gsm8k.

Supporting changes include:

  • Adds writable Enroot/Pyxis host-side cache, data, temp, and runtime environment setup for Slurm-launched containers.
  • Sets torch distributed environment variables for direct launcher task payloads.
  • Improves AIPerf vLLM argument handling, including default max-num-seqs.
  • Makes runtime-stat reuse workload-aware and supports width scenario reuse/root measurement.

Usage

defaults:                                                                                                                                                           
  - default                                                                                                                                                         
  - _self_                                                                                                                                                          
                                                                                                                                                                    
post_mip:                                                                                                                                                           
  flows:                                                                                                                                                            
    runtime-075-lmms-eval:                                                                                                                                          
      source:                                                                                                                                                       
        run: runtime-075                                                                                                                                            
        variants: all                                                                                                                                               
        objectives: all                                                                                                                                             
      nodes:                                                                                                                                                        
        best_mip:                                                                                                                                                   
          type: filter                                                                                                                                              
          mode: top_k                                                                                                                                               
          metric: mip.score                                                                                                                                         
          direction: minimize                                                                                                                                       
          top_k: 1                                                                                                                                                  
        materialized:                                                                                                                                               
          type: materialize                                                                                                                                         
          input: best_mip                                                                                                                                           
        lmms_eval:                                                                                                                                                  
          type: downstream_evaluation                                                                                                                               
          input: materialized                                                                                                                                       
          config:                                                                                                                                                   
            model: vllm                                                                                                                                             
            tasks: [ifeval, gsm8k]                                                                                                                                  
            limit: 128                                                                                                                                              
            batch_size: 1                                                                                                                                           
            timeout_seconds: 7200                                                                                                                                   
            topology:                                                                                                                                               
              tensor_parallel_size: 8                                                                                                                               
              pipeline_parallel_size: 1                                                                                                                             
              data_parallel_size: 1                                                                                                                                 
              prefill_context_parallel_size: 1                                                                                                                      
              decode_context_parallel_size: 1                                                                                                                       
              enable_expert_parallel: false                                                                                                                         
              distributed_executor_backend: mp                                                                                                                      
              gpu_group_size: 8                                                                                                                                     
           model_args:                                                                                                                                             
              dtype: bfloat16                                                                                                                                       
              gpu_memory_utilization: 0.85                                                                                                                          
              max_model_len: 262144                                                                                                                                 
              trust_remote_code: ${model.trust_remote_code}                                                                                                         

Testing

  • git diff --check feature/puzzletron_v2...HEAD
  • 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.p\ y puzzletron_setup/bundle.py puzzletron_setup/v2/parallel_validation.py puzzletron_setup/v2/post_mip.py puzzletron_setup/v2/validation.py puzzletron_setup/v2/wizar\ d.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/puz\ zletron/test_setup_bundle.py tests/unit/torch/puzzletron/test_setup_v2_post_mip.py tests/unit/torch/puzzletron/test_setup_v2_state_validation.py
  • Added focused unit coverage for downstream evaluation command construction, result parsing, missing-result diagnostics, orchestration allocation/dashboard labels
    , setup bundle resources, v2 validation, Slurm Enroot/Pyxis paths, direct launcher distributed environment setup, AIPerf vLLM args, and runtime-stat reuse.
  • Manual smoke validation completed after updating the execution environment to include lmms-eval:
    • post.params-75.lmms_eval succeeded on Slurm job 15159279.
    • final_report completed on Slurm job 15159510.
    • Final report artifact: nano/results/smoke/artifacts/campaign_report/campaign_report.html.
  • Focused pytest was attempted but did not run in this local host environment because pytest is not installed.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commi\ t -s -S).

Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contribut\
ors) (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

  • Is this change backward compatible?: Yes
  • If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
  • Did you write any new necessary tests?: Yes
  • Did you update Changelog?: Yes
  • Did you get Claude approval on this PR?: No, not run.

Additional Information

lmms-eval and vLLM must be available in the target execution environment. This PR wires Puzzletron to call lmms-eval; it does not vendor lmms-eval or add it
as a ModelOpt package dependency.

Summary by CodeRabbit

  • New Features

    • Added downstream model evaluation using lmms-eval and vLLM.
    • Setup flows now support configuring evaluation tasks, metrics, batching, timeouts, logging, and hardware topology.
    • Added checkpoint evaluation results, summaries, and dashboard reporting.
    • Added an opt-in Nemotron-3 Nano example for ifeval and gsm8k.
  • Documentation

    • Added setup instructions, environment requirements, usage guidance, and configuration limitations.
  • Bug Fixes

    • Improved resource allocation, topology validation, candidate selection, and progress reporting for downstream evaluations.

Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 7, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 41450810-f950-406d-bf3e-011f6a9aa68f

📥 Commits

Reviewing files that changed from the base of the PR and between 57ccb99 and 7e990d1.

📒 Files selected for processing (3)
  • examples/puzzletron/README.md
  • examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml
  • examples/puzzletron/docs/post_mip_pipeline.md
🚧 Files skipped from review as they are similar to previous changes (2)
  • examples/puzzletron/docs/post_mip_pipeline.md
  • examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml

📝 Walkthrough

Walkthrough

Puzzletron now supports configurable vLLM-backed lmms-eval downstream evaluation after MIP selection. The change adds setup prompts, topology validation, orchestration, subprocess execution, result validation, reporting, tests, documentation, and an opt-in Nemotron-3 Nano example.

Changes

Puzzletron downstream evaluation

Layer / File(s) Summary
Downstream evaluation configuration and validation
puzzletron_setup/..., tests/unit/torch/puzzletron/test_setup_bundle.py, tests/unit/torch/puzzletron/test_setup_v2_*.py
downstream_evaluation is implemented as a post-MIP node. Setup flows collect tasks, limits, model arguments, timeout, logging, and vLLM topology. Resource planning and topology validation include the new node type.
Distributed orchestration and dashboard integration
modelopt/torch/puzzletron/orchestration/..., tests/unit/torch/puzzletron/test_orchestration_*.py
Planning applies candidate and node-count limits. Compilation derives vLLM meshes and GPU allocation. Dashboard names and progress labels identify downstream evaluation stages.
lmms-eval execution and reporting
modelopt/torch/puzzletron/post_mip/..., tests/unit/torch/puzzletron/test_post_mip_runner.py
The runner builds and executes lmms-eval commands, injects checkpoints and topology settings, handles timeouts, persists process output, validates results and metrics, and generates downstream evaluation reports.
Example flow and environment support
examples/puzzletron/..., CHANGELOG.rst
The example environment pins lmms_eval==0.7.2. Documentation describes isolated setup and reserved arguments. The Nemotron-3 Nano example evaluates ifeval and gsm8k through an opt-in post-MIP flow.

Estimated code review effort: 5 (Critical) | ~90 minutes

Sequence Diagram(s)

sequenceDiagram
  participant SetupWizard
  participant CampaignCompiler
  participant DownstreamEvaluationNode
  participant PostMIPRunner
  participant lmms_eval
  participant ArtifactStore

  SetupWizard->>CampaignCompiler: configure downstream_evaluation topology and tasks
  CampaignCompiler->>DownstreamEvaluationNode: materialize checkpoint-dependent node
  DownstreamEvaluationNode->>PostMIPRunner: execute selected checkpoint evaluation
  PostMIPRunner->>lmms_eval: launch vLLM evaluation command
  lmms_eval-->>PostMIPRunner: return evaluation results
  PostMIPRunner->>ArtifactStore: persist results, streams, and summary
Loading

Suggested reviewers: kevalmorabia97, separius


Caution

Pre-merge checks failed

Please resolve all errors before merging. Addressing warnings is optional.

  • Ignore

❌ Failed checks (1 error, 1 warning)

Check name Status Explanation Resolution
Security Anti-Patterns ❌ Error Added code uses # nosec, hardcodes trust_remote_code=True, and adds weights_only=False loads without inline safety comments; SECURITY.md forbids these patterns without an approved exception. Remove new # nosec suppressions; use weights_only=True where possible, document trusted files inline, and expose trust_remote_code with a default of False. Obtain required codeowner approval for exceptions.
Docstring Coverage ⚠️ Warning Docstring coverage is 28.67% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly and concisely summarizes the main change: adding downstream evaluation support to Puzzletron v2.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch gkarch/add_downstream_eval

Comment @coderabbitai help to get the list of available commands.

@coderabbitai

coderabbitai Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Note

Docstrings generation - SUCCESS
Generated docstrings for this pull request at #2105

coderabbitai Bot added a commit that referenced this pull request Aug 7, 2026
Docstrings generation was requested by @grzegorz-k-karch.

* #2104 (comment)

The following files were modified:

* `modelopt/torch/puzzletron/benchmarks/aiperf.py`
* `modelopt/torch/puzzletron/orchestration/adapters/pool.py`
* `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py`
* `modelopt/torch/puzzletron/orchestration/compiler.py`
* `modelopt/torch/puzzletron/orchestration/controller.py`
* `modelopt/torch/puzzletron/orchestration/executors/slurm.py`
* `modelopt/torch/puzzletron/orchestration/progress.py`
* `modelopt/torch/puzzletron/orchestration/task_launcher.py`
* `modelopt/torch/puzzletron/post_mip/builtin.py`
* `modelopt/torch/puzzletron/post_mip/reporting.py`
* `modelopt/torch/puzzletron/post_mip/runner.py`
* `modelopt/torch/puzzletron/stages/pipeline.py`
* `modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py`
* `puzzletron_setup/bundle.py`
* `puzzletron_setup/v2/validation.py`
* `puzzletron_setup/v2/wizard.py`
* `puzzletron_setup/wizard.py`
coderabbitai Bot and others added 2 commits August 7, 2026 12:16
Docstrings generation was requested by @grzegorz-k-karch.

*
#2104 (comment)

The following files were modified:

* `modelopt/torch/puzzletron/benchmarks/aiperf.py`
* `modelopt/torch/puzzletron/orchestration/adapters/pool.py`
* `modelopt/torch/puzzletron/orchestration/adapters/post_mip.py`
* `modelopt/torch/puzzletron/orchestration/compiler.py`
* `modelopt/torch/puzzletron/orchestration/controller.py`
* `modelopt/torch/puzzletron/orchestration/executors/slurm.py`
* `modelopt/torch/puzzletron/orchestration/progress.py`
* `modelopt/torch/puzzletron/orchestration/task_launcher.py`
* `modelopt/torch/puzzletron/post_mip/builtin.py`
* `modelopt/torch/puzzletron/post_mip/reporting.py`
* `modelopt/torch/puzzletron/post_mip/runner.py`
* `modelopt/torch/puzzletron/stages/pipeline.py`
* `modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.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_aiperf_context_capacity.py`
* `tests/unit/torch/puzzletron/test_orchestration_compiler.py`
* `tests/unit/torch/puzzletron/test_orchestration_controller.py`
* `tests/unit/torch/puzzletron/test_orchestration_executors.py`
* `tests/unit/torch/puzzletron/test_orchestration_task_topology.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`
* `tests/unit/torch/puzzletron/test_sparse_runtime_stats.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>

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

👉 Steps to fix this

Actionable comments posted: 9

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
modelopt/torch/puzzletron/benchmarks/aiperf.py (1)

442-456: 🔒 Security & Privacy | 🔴 Critical | ⚡ Quick win

Make vLLM remote-code trust caller-controlled.

run_aiperf_sweep always enables execution of custom model and tokenizer code. Add trust_remote_code: bool = False, append --trust-remote-code only when true, and include the resolved value in cache_identity.

🤖 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/benchmarks/aiperf.py` around lines 442 - 456,
Update run_aiperf_sweep to accept trust_remote_code: bool = False, append
--trust-remote-code to server_cmd only when that option is enabled, and include
the resolved trust_remote_code value in cache_identity so cache entries
distinguish the setting.

Sources: Coding guidelines, Path instructions, MCP tools

🧹 Nitpick comments (4)
modelopt/torch/puzzletron/stages/pipeline.py (2)

221-225: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Optional cleanup: dedented brace and redundant set_struct.

Line 221 closes the workloads literal at 4-space indentation, which does not match the enclosing block. Python accepts this, but it reduces readability. Line 225 calls OmegaConf.set_struct(selected, False), while clone_hydra_config already clears struct mode on the clone.

♻️ Proposed cleanup
-    }
+        }
     for raw_workload in workloads.values():
         workload = dict(raw_workload or {})
         selected = clone_hydra_config(hydra_cfg)
-        OmegaConf.set_struct(selected, False)
         stats_cfg = selected.calc_subblock_stats
🤖 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/stages/pipeline.py` around lines 221 - 225, Clean
up the workload-processing block by aligning the closing workloads-literal brace
with its enclosing indentation and removing the redundant OmegaConf.set_struct
call after clone_hydra_config. Preserve the existing clone and workload
iteration behavior.

359-359: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add a brief comment for the local import.

The new function imports launch_calc_subblock_stats inside the function body. The guideline permits local imports for circular, optional, or heavy dependencies, but requires a brief explanatory comment. Add one line stating the reason, for example the heavy torch/vLLM import chain.

Based on learnings from the coding guidelines: "Keep imports at the top of Python source and test files; use local imports only for justified circular dependencies, optional dependencies, or unusually heavy imports, with a brief explanatory comment."

🤖 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/stages/pipeline.py` at line 359, Add a brief
explanatory comment immediately above the local import of
launch_calc_subblock_stats in the affected function, noting that it is kept
local because its torch/vLLM dependency chain is unusually heavy.

Source: Coding guidelines

modelopt/torch/puzzletron/post_mip/runner.py (1)

728-740: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Result discovery parses every JSON file under the output directory.

rglob("*.json") reads and parses all JSON below output_path. When log_samples is enabled, lmms-eval writes per-sample logs there. The wizard sets log_samples: True by default. For a large limit those files dominate the directory, and each is fully loaded into memory only to be discarded by the results check at line 735.

Narrow the search to the aggregated result files, or skip files above a size threshold before parsing.

🤖 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 728 - 740, The
_lmms_eval_result_payload function currently parses every JSON file, including
large per-sample logs. Narrow candidate discovery to aggregated result files
using their known filename/pattern, or apply an appropriate size threshold
before json.loads, while preserving selection of the newest valid payload
containing a results mapping.
tests/unit/torch/puzzletron/test_post_mip_runner.py (1)

229-333: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add coverage for the model_args precedence and the timeout path.

The two execution tests cover the success path and the missing-results path. Two behaviors that this PR introduces stay uncovered:

  1. A user-supplied model_args mapping that also sets model or tensor_parallel_size. See the finding on modelopt/torch/puzzletron/post_mip/runner.py lines 615-624.
  2. A subprocess.TimeoutExpired raised by the fake. See the finding on modelopt/torch/puzzletron/post_mip/runner.py lines 799-808.

Both are cheap to add with the existing fake_run pattern and no GPU. Add them alongside the fixes.

As per path instructions: "Add lean regression tests for command construction, topology/resource allocation, parsing, failures, and environment behavior."

🤖 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 - 333,
Extend the lmms-eval tests around _downstream_evaluation with two lean
regressions: verify user-supplied model_args takes precedence over generated
model and tensor_parallel_size arguments, and make the subprocess fake raise
TimeoutExpired to assert the timeout failure behavior and persisted output
handling. Reuse the existing node/source fixtures and fake_run pattern without
requiring GPU execution.

Source: Path instructions

🤖 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: Sort _LMMS_EVAL_MODEL_ARG_FIELDS before iterating in the
derived-argument construction so --model_args and the persisted argv remain
deterministic across processes. Preserve the existing settings membership check
and derived[key] assignments.
- Around line 695-696: Define a shared _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS
constant with value 3600.0 in runner.py, use it as the fallback when both
timeout_seconds and timeout are absent in the command-building logic, and reuse
the same constant in the timeout-reporting path around the existing failure
handling instead of duplicating 3600.
- Around line 615-624: Update the model-argument merge logic around
_model_arg_string so derived checkpoint and topology values override
user-supplied values in both branches. For mapping inputs, replace
setdefault-based merging with precedence for derived entries; for string inputs,
ensure duplicate keys resolve to the derived suffix under lmms-eval’s last-value
behavior. Preserve validation and output formatting while guaranteeing the
realized checkpoint and allocated topology are authoritative.
- Around line 799-808: Wrap the subprocess.run call in the relevant runner flow
with a subprocess.TimeoutExpired handler, decode the exception’s stdout and
stderr bytes, and pass them to _write_lmms_eval_streams before re-raising the
timeout. Preserve the existing successful-result handling unchanged.

In `@modelopt/torch/puzzletron/stages/pipeline.py`:
- Around line 254-296: Unify runtime measurement identity across all three
sites: in modelopt/torch/puzzletron/stages/pipeline.py lines 254-296, update
_has_runtime_measurement to use the shared key builder instead of subset
comparisons; in modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py
lines 312-325, apply identical defaults and casts for num_iters,
num_warmup_iters, repeat_block_n_times, max_num_seqs, and granularity on
persisted and requested values; in
tests/unit/torch/puzzletron/test_sparse_runtime_stats.py lines 1159-1179, add
num_iters, num_warmup_iters, repeat_block_n_times, and vllm_args to the root
aggregate fixture.

In `@modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py`:
- Around line 281-299: Update the runtime-stats identity extraction logic around
the tuple construction to validate that n_embd, batch_size, prefill_seq_len, and
generation_seq_len are present before converting them with int. Return None for
entries missing any required identity field, allowing
calculate_subblock_stats_for_puzzle_dir to skip them and report missing
identities through its existing path; preserve current handling for complete
entries.

In `@puzzletron_setup/wizard.py`:
- Around line 741-748: Update the lmms-eval task prompt in the relevant wizard
function to validate that the entered task list is non-empty, matching
_downstream_evaluation_setting_prompt behavior, so cleared input is rejected at
the interface. Render sequence defaults by joining list-like values with commas
while preserving string defaults, ensuring defaults produced by
_ask_aiperf_config display as valid task text before the existing parsing at
line 826.
- Around line 1069-1075: Update the downstream evaluation branch around
_ask_downstream_evaluation_config so available_metrics is populated from the
selected tasks’ emitted metric names rather than the hardcoded gsm8k.exact_match
path. Include each task’s aggregation suffix, such as strict-match, so
registered paths exactly match the keys produced by the runner.

In `@tests/unit/torch/puzzletron/test_sparse_runtime_stats.py`:
- Line 1320: Remove the assertion referencing hydra_cfg from
test_runtime_stats_resume_signature_includes_workload_id, since that symbol is
undefined and unrelated to the test. Leave the remaining test behavior
unchanged.

---

Outside diff comments:
In `@modelopt/torch/puzzletron/benchmarks/aiperf.py`:
- Around line 442-456: Update run_aiperf_sweep to accept trust_remote_code: bool
= False, append --trust-remote-code to server_cmd only when that option is
enabled, and include the resolved trust_remote_code value in cache_identity so
cache entries distinguish the setting.

---

Nitpick comments:
In `@modelopt/torch/puzzletron/post_mip/runner.py`:
- Around line 728-740: The _lmms_eval_result_payload function currently parses
every JSON file, including large per-sample logs. Narrow candidate discovery to
aggregated result files using their known filename/pattern, or apply an
appropriate size threshold before json.loads, while preserving selection of the
newest valid payload containing a results mapping.

In `@modelopt/torch/puzzletron/stages/pipeline.py`:
- Around line 221-225: Clean up the workload-processing block by aligning the
closing workloads-literal brace with its enclosing indentation and removing the
redundant OmegaConf.set_struct call after clone_hydra_config. Preserve the
existing clone and workload iteration behavior.
- Line 359: Add a brief explanatory comment immediately above the local import
of launch_calc_subblock_stats in the affected function, noting that it is kept
local because its torch/vLLM dependency chain is unusually heavy.

In `@tests/unit/torch/puzzletron/test_post_mip_runner.py`:
- Around line 229-333: Extend the lmms-eval tests around _downstream_evaluation
with two lean regressions: verify user-supplied model_args takes precedence over
generated model and tensor_parallel_size arguments, and make the subprocess fake
raise TimeoutExpired to assert the timeout failure behavior and persisted output
handling. Reuse the existing node/source fixtures and fake_run pattern without
requiring GPU execution.
🪄 Autofix

✅ Autofix completed


ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: de470a57-4d2a-4c05-80b0-8cba7c728740

📥 Commits

Reviewing files that changed from the base of the PR and between 3485fd0 and afbb964.

📒 Files selected for processing (35)
  • CHANGELOG.rst
  • examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml
  • examples/puzzletron/distributed_eval/run_coordinator.sh
  • examples/puzzletron/distributed_eval/run_depth_coordinator.sh
  • examples/puzzletron/distributed_eval/run_depth_pool.sh
  • examples/puzzletron/distributed_eval/run_replacement_pool.sh
  • modelopt/torch/puzzletron/benchmarks/aiperf.py
  • modelopt/torch/puzzletron/orchestration/adapters/pool.py
  • modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
  • modelopt/torch/puzzletron/orchestration/compiler.py
  • modelopt/torch/puzzletron/orchestration/controller.py
  • modelopt/torch/puzzletron/orchestration/executors/slurm.py
  • modelopt/torch/puzzletron/orchestration/progress.py
  • modelopt/torch/puzzletron/orchestration/task_launcher.py
  • modelopt/torch/puzzletron/post_mip/builtin.py
  • modelopt/torch/puzzletron/post_mip/reporting.py
  • modelopt/torch/puzzletron/post_mip/runner.py
  • modelopt/torch/puzzletron/stages/pipeline.py
  • modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.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_aiperf_context_capacity.py
  • tests/unit/torch/puzzletron/test_orchestration_compiler.py
  • tests/unit/torch/puzzletron/test_orchestration_controller.py
  • tests/unit/torch/puzzletron/test_orchestration_executors.py
  • tests/unit/torch/puzzletron/test_orchestration_task_topology.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
  • tests/unit/torch/puzzletron/test_sparse_runtime_stats.py

Comment thread modelopt/torch/puzzletron/post_mip/runner.py
Comment thread modelopt/torch/puzzletron/post_mip/runner.py
Comment thread modelopt/torch/puzzletron/post_mip/runner.py
Comment thread modelopt/torch/puzzletron/post_mip/runner.py Outdated
Comment thread modelopt/torch/puzzletron/stages/pipeline.py Outdated
Comment thread modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py Outdated
Comment thread puzzletron_setup/wizard.py
Comment thread puzzletron_setup/wizard.py Outdated
Comment thread tests/unit/torch/puzzletron/test_sparse_runtime_stats.py Outdated
@coderabbitai

coderabbitai Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Fixes Applied Successfully

Fixed 5 file(s) based on 9 unresolved review comments.

A stacked PR containing fixes has been created.

  • Stacked PR: #2106
  • Files modified:
  • modelopt/torch/puzzletron/post_mip/runner.py
  • modelopt/torch/puzzletron/stages/pipeline.py
  • modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py
  • puzzletron_setup/wizard.py
  • tests/unit/torch/puzzletron/test_sparse_runtime_stats.py

Time taken: 5m 49s

@codecov

codecov Bot commented Aug 7, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 1.51515% with 325 lines in your changes missing coverage. Please review.
✅ Project coverage is 54.58%. Comparing base (3485fd0) to head (7e990d1).

Files with missing lines Patch % Lines
modelopt/torch/puzzletron/post_mip/runner.py 0.00% 312 Missing ⚠️
...elopt/torch/puzzletron/orchestration/controller.py 30.00% 7 Missing ⚠️
modelopt/torch/puzzletron/post_mip/builtin.py 0.00% 4 Missing ⚠️
modelopt/torch/puzzletron/post_mip/reporting.py 0.00% 2 Missing ⚠️
Additional details and impacted files
@@                    Coverage Diff                    @@
##           feature/puzzletron_v2    #2104      +/-   ##
=========================================================
- Coverage                  54.76%   54.58%   -0.18%     
=========================================================
  Files                        704      704              
  Lines                      90869    91193     +324     
=========================================================
+ Hits                       49761    49775      +14     
- Misses                     41108    41418     +310     
Flag Coverage Δ
puzzletron 32.77% <1.51%> (-0.11%) ⬇️
unit 29.64% <0.00%> (-0.11%) ⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

This stacked PR contains CodeRabbit auto-fixes for #2104.

**Files modified:**
- `modelopt/torch/puzzletron/post_mip/runner.py`
- `modelopt/torch/puzzletron/stages/pipeline.py`
- `modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py`
- `puzzletron_setup/wizard.py`
- `tests/unit/torch/puzzletron/test_sparse_runtime_stats.py`

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
@grzegorz-k-karch
grzegorz-k-karch marked this pull request as ready for review August 7, 2026 10:34
@grzegorz-k-karch
grzegorz-k-karch requested a review from a team as a code owner August 7, 2026 10:34
@grzegorz-k-karch grzegorz-k-karch self-assigned this Aug 7, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

♻️ Duplicate comments (1)
modelopt/torch/puzzletron/post_mip/runner.py (1)

681-684: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Put derived model arguments last for string model_args.

At Line 684, the raw string follows derived. Repeated keys in raw model_args can override the realized checkpoint and allocated topology. This disagrees with the mapping branch at Lines 687-689.

Proposed fix
-        return ",".join(part for part in (suffix, prefix) if part)
+        return ",".join(part for part in (prefix, suffix) if part)
🤖 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 681 - 684, Update
the string-handling branch in _model_arg_string so the raw model arguments
precede the derived arguments, matching the mapping branch and ensuring derived
checkpoint and topology values take precedence for duplicate keys.
🤖 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.

Duplicate comments:
In `@modelopt/torch/puzzletron/post_mip/runner.py`:
- Around line 681-684: Update the string-handling branch in _model_arg_string so
the raw model arguments precede the derived arguments, matching the mapping
branch and ensuring derived checkpoint and topology values take precedence for
duplicate keys.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d49d16ee-f253-4f33-899e-44f9f7a9bb3d

📥 Commits

Reviewing files that changed from the base of the PR and between afbb964 and 3798b10.

📒 Files selected for processing (21)
  • modelopt/torch/puzzletron/benchmarks/aiperf.py
  • modelopt/torch/puzzletron/orchestration/adapters/pool.py
  • modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
  • modelopt/torch/puzzletron/orchestration/compiler.py
  • modelopt/torch/puzzletron/orchestration/controller.py
  • modelopt/torch/puzzletron/orchestration/executors/slurm.py
  • modelopt/torch/puzzletron/orchestration/progress.py
  • modelopt/torch/puzzletron/orchestration/task_launcher.py
  • modelopt/torch/puzzletron/post_mip/builtin.py
  • modelopt/torch/puzzletron/post_mip/reporting.py
  • modelopt/torch/puzzletron/post_mip/runner.py
  • modelopt/torch/puzzletron/stages/pipeline.py
  • modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py
  • puzzletron_setup/bundle.py
  • puzzletron_setup/v2/validation.py
  • puzzletron_setup/v2/wizard.py
  • puzzletron_setup/wizard.py
  • tests/unit/torch/puzzletron/test_orchestration_controller.py
  • tests/unit/torch/puzzletron/test_orchestration_executors.py
  • tests/unit/torch/puzzletron/test_setup_v2_state_validation.py
  • tests/unit/torch/puzzletron/test_sparse_runtime_stats.py
🚧 Files skipped from review as they are similar to previous changes (19)
  • puzzletron_setup/v2/validation.py
  • tests/unit/torch/puzzletron/test_setup_v2_state_validation.py
  • modelopt/torch/puzzletron/orchestration/executors/slurm.py
  • modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
  • modelopt/torch/puzzletron/orchestration/progress.py
  • modelopt/torch/puzzletron/orchestration/controller.py
  • modelopt/torch/puzzletron/post_mip/builtin.py
  • puzzletron_setup/bundle.py
  • tests/unit/torch/puzzletron/test_orchestration_executors.py
  • modelopt/torch/puzzletron/orchestration/adapters/pool.py
  • modelopt/torch/puzzletron/orchestration/compiler.py
  • modelopt/torch/puzzletron/orchestration/task_launcher.py
  • puzzletron_setup/wizard.py
  • modelopt/torch/puzzletron/benchmarks/aiperf.py
  • modelopt/torch/puzzletron/stages/pipeline.py
  • modelopt/torch/puzzletron/post_mip/reporting.py
  • puzzletron_setup/v2/wizard.py
  • tests/unit/torch/puzzletron/test_sparse_runtime_stats.py
  • modelopt/torch/puzzletron/subblock_stats/calc_subblock_stats.py

Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
@grzegorz-k-karch
grzegorz-k-karch requested a review from a team as a code owner August 8, 2026 18:06
@grzegorz-k-karch
grzegorz-k-karch requested review from kevalmorabia97 and removed request for a team August 8, 2026 18:06

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

👉 Steps to fix this

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
puzzletron_setup/v2/wizard.py (1)

4151-4170: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Return BACK immediately after each prompt.

Lines 4151-4168 run three prompts before line 4169 inspects the results. If the user selects back at the limit prompt, the wizard still asks for batch_size and timeout before it returns BACK. _serving_setting_prompt at lines 4059-4093 checks each answer right after the prompt. Match that behavior.

🛠️ Proposed fix
     limit = session.integer(
         f"{prefix}.limit",
         "lmms-eval sample limit:",
         default=int(defaults.get("limit", 128)),
         minimum=1,
     )
+    if limit is BACK:
+        return BACK
     batch_size = session.integer(
         f"{prefix}.batch_size",
         "lmms-eval batch size:",
         default=int(defaults.get("batch_size", 1)),
         minimum=1,
     )
+    if batch_size is BACK:
+        return BACK
     timeout = session.integer(
         f"{prefix}.timeout_seconds",
         "Per-candidate lmms-eval timeout (seconds):",
         default=int(defaults.get("timeout_seconds", 3600)),
         minimum=1,
     )
-    if BACK in (limit, batch_size, timeout):
+    if timeout is BACK:
         return BACK
🤖 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/v2/wizard.py` around lines 4151 - 4170, Update the prompt
flow around the limit, batch_size, and timeout assignments to check for BACK
immediately after each session.integer call and return it before showing the
next prompt, matching _serving_setting_prompt behavior. Remove the deferred
combined check after all three prompts.
♻️ Duplicate comments (1)
puzzletron_setup/wizard.py (1)

750-753: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

The task prompt still lacks validation and sequence-default rendering.

Line 752 applies str() to the default. If defaults["tasks"] is a list, the prompt renders "['ifeval', 'gsm8k']", and line 835 produces corrupted task names. The prompt also accepts empty input, which yields tasks: [] and fails later in _configured_lmms_eval_tasks. The v2 prompt _downstream_evaluation_setting_prompt handles both cases at puzzletron_setup/v2/wizard.py lines 4132-4147.

🛠️ Proposed fix
     defaults = defaults or {}
+    raw_default_tasks = defaults.get("tasks", "ifeval,gsm8k")
+    default_tasks = (
+        raw_default_tasks
+        if isinstance(raw_default_tasks, str)
+        else ",".join(str(item) for item in raw_default_tasks)
+    )
+
+    def validate_tasks(value: str) -> bool | str:
+        return (
+            True
+            if [item for item in value.split(",") if item.strip()]
+            else "Enter at least one lmms-eval task."
+        )
+
     tasks = prompts.text(
         "lmms-eval tasks (comma-separated):",
-        default=str(defaults.get("tasks", "ifeval,gsm8k")),
+        default=default_tasks,
+        validate=validate_tasks,
     )
🤖 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 750 - 753, Update the lmms-eval
tasks prompt in the wizard flow to render sequence defaults as comma-separated
task names rather than Python list syntax, matching
_downstream_evaluation_setting_prompt. Add validation that rejects empty input,
ensuring the resulting tasks value cannot become an empty list before
_configured_lmms_eval_tasks processes it.
🧹 Nitpick comments (2)
modelopt/torch/puzzletron/post_mip/runner.py (1)

562-570: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick win

Add --model and --batch_size to the reserved extra-argument flags.

The runner always emits --model at line 778 and --batch_size at line 784. lmms-eval uses argparse, so a later duplicate flag in extra_args wins. A config that sets extra_args: ["--model", "hf"] therefore replaces the vLLM backend while --model_args still carries the vLLM topology keys, and the run fails or evaluates a different backend than the recorded topology.

🛡️ Proposed hardening
 _LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS = frozenset(
     {
+        "--model",
         "--model_args",
         "--model-args",
         "--output_path",
         "--output-path",
         "--tasks",
+        "--batch_size",
+        "--batch-size",
     }
 )
🤖 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 562 - 570, Update
the _LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS set to include --model and --batch_size
alongside the existing reserved flags, preventing extra_args from overriding the
runner-emitted backend and batch size arguments.
tests/unit/torch/puzzletron/test_post_mip_runner.py (1)

230-278: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use pytest.raises and pytest.mark.parametrize for these rejection tests.

The manual try/except/else with raise AssertionError reimplements pytest.raises. The in-test loop also hides which case failed. This pattern repeats in the four tests at lines 518, 565, and 618.

♻️ Proposed refactor for one test
-def test_lmms_eval_command_rejects_reserved_model_args(tmp_path):
-    cases = (
-        ({"model": "/ckpts/wrong"}, "model"),
-        ("dtype=bfloat16,tensor_parallel_size=1", "tensor_parallel_size"),
-    )
-    for model_args, expected in cases:
-        try:
-            runner._lmms_eval_command(
-                {
-                    "tasks": ["ifeval"],
-                    "topology": {"gpu_group_size": 1},
-                    "model_args": model_args,
-                },
-                checkpoint="/ckpts/candidate",
-                output_path=tmp_path / "results",
-            )
-        except ValueError as error:
-            message = str(error)
-        else:
-            raise AssertionError("expected reserved lmms-eval model_args to fail")
-
-        assert "reserved lmms-eval model arguments" in message
-        assert expected in message
+@pytest.mark.parametrize(
+    ("model_args", "expected"),
+    [
+        ({"model": "/ckpts/wrong"}, "model"),
+        ("dtype=bfloat16,tensor_parallel_size=1", "tensor_parallel_size"),
+    ],
+)
+def test_lmms_eval_command_rejects_reserved_model_args(tmp_path, model_args, expected):
+    with pytest.raises(ValueError, match="reserved lmms-eval model arguments") as error:
+        runner._lmms_eval_command(
+            {
+                "tasks": ["ifeval"],
+                "topology": {"gpu_group_size": 1},
+                "model_args": model_args,
+            },
+            checkpoint="/ckpts/candidate",
+            output_path=tmp_path / "results",
+        )
+    assert expected in str(error.value)
🤖 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 230 - 278,
Refactor the reserved-argument rejection tests around _lmms_eval_command to use
pytest.mark.parametrize for individual cases and pytest.raises(ValueError) for
the expected failures, asserting the error message within each case. Apply the
same pattern to the repeated rejection tests around lines 518, 565, and 618,
preserving each case’s expected message fragments.

Source: Coding guidelines

🤖 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 `@examples/puzzletron/docs/post_mip_pipeline.md`:
- Around line 136-138: Update the downstream_evaluation description to state
that it invokes python -m lmms_eval via subprocess without a shell, rather than
saying it “shells out.” Clarify that command_prefix and extra_args are passed as
argv elements and must not contain shell syntax.

In `@modelopt/torch/puzzletron/post_mip/runner.py`:
- Around line 708-710: Make the loop over _LMMS_EVAL_MODEL_ARG_FIELDS
deterministic by iterating its keys in a stable order, such as sorted order,
before populating derived. Preserve the existing filtering and value assignment
so _model_arg_string, persisted argv, and the documented command determinism
remain stable for identical configuration.
- Around line 822-823: Define a shared _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS
constant near _LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, set to 3600.0, and use
it as the fallback when both timeout_seconds and timeout are absent in the
timeout resolution logic. Replace the hardcoded 3600 downstream-evaluation
fallback near the existing failure path so both _run_lmms_eval_process behavior
and failure handling use the same default.
- Around line 1042-1044: Update the TimeoutExpired handler around
_signal_lmms_eval_process_group and the final process.communicate() to use a
finite timeout, preserving the existing fallback handling for missing output and
stderr. If the bounded communicate also times out, ensure the runner does not
hang and continues through the established cleanup/error-result path.

In `@tests/unit/torch/puzzletron/test_post_mip_runner.py`:
- Around line 424-429: Update the communicate_timeouts expectation in the
relevant test to use runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS for the
third call instead of None, matching the bounded cleanup communicate behavior in
the post-MIP runner.

---

Outside diff comments:
In `@puzzletron_setup/v2/wizard.py`:
- Around line 4151-4170: Update the prompt flow around the limit, batch_size,
and timeout assignments to check for BACK immediately after each session.integer
call and return it before showing the next prompt, matching
_serving_setting_prompt behavior. Remove the deferred combined check after all
three prompts.

---

Duplicate comments:
In `@puzzletron_setup/wizard.py`:
- Around line 750-753: Update the lmms-eval tasks prompt in the wizard flow to
render sequence defaults as comma-separated task names rather than Python list
syntax, matching _downstream_evaluation_setting_prompt. Add validation that
rejects empty input, ensuring the resulting tasks value cannot become an empty
list before _configured_lmms_eval_tasks processes it.

---

Nitpick comments:
In `@modelopt/torch/puzzletron/post_mip/runner.py`:
- Around line 562-570: Update the _LMMS_EVAL_RESERVED_EXTRA_ARG_FLAGS set to
include --model and --batch_size alongside the existing reserved flags,
preventing extra_args from overriding the runner-emitted backend and batch size
arguments.

In `@tests/unit/torch/puzzletron/test_post_mip_runner.py`:
- Around line 230-278: Refactor the reserved-argument rejection tests around
_lmms_eval_command to use pytest.mark.parametrize for individual cases and
pytest.raises(ValueError) for the expected failures, asserting the error message
within each case. Apply the same pattern to the repeated rejection tests around
lines 518, 565, and 618, preserving each case’s expected message fragments.
🪄 Autofix

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: 10547945-e8ea-4b90-b1a9-cd73eaf798e7

📥 Commits

Reviewing files that changed from the base of the PR and between 3798b10 and 57ccb99.

📒 Files selected for processing (19)
  • examples/puzzletron/README.md
  • examples/puzzletron/ci_environment.json
  • examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml
  • examples/puzzletron/docs/post_mip_pipeline.md
  • examples/puzzletron/requirements.txt
  • modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
  • modelopt/torch/puzzletron/orchestration/compiler.py
  • modelopt/torch/puzzletron/orchestration/controller.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
  • noxfile.py
  • puzzletron_setup/bundle.py
  • puzzletron_setup/v2/validation.py
  • puzzletron_setup/v2/wizard.py
  • puzzletron_setup/wizard.py
  • tests/unit/torch/puzzletron/test_post_mip_runner.py
  • tests/unit/torch/puzzletron/test_setup_bundle.py
💤 Files with no reviewable changes (5)
  • modelopt/torch/puzzletron/orchestration/progress.py
  • modelopt/torch/puzzletron/post_mip/builtin.py
  • modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
  • modelopt/torch/puzzletron/orchestration/controller.py
  • puzzletron_setup/bundle.py
🚧 Files skipped from review as they are similar to previous changes (4)
  • examples/puzzletron/configs/families/nemotron3/nano_30b_a3b_bf16/runs/lmms_eval.yaml
  • puzzletron_setup/v2/validation.py
  • tests/unit/torch/puzzletron/test_setup_bundle.py
  • modelopt/torch/puzzletron/orchestration/compiler.py

Comment on lines +136 to +138
`downstream_evaluation` shells out to `python -m lmms_eval` from the GPU worker
environment. Install `examples/puzzletron/requirements.txt` in that environment;
it pins the evaluator package used by the checked-in example:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Correct the "shells out" wording.

The runner passes an argv list to subprocess.Popen and never invokes a shell. SECURITY.md requires this runner to stay shell-free. The phrase "shells out" states the opposite and could lead a config author to put shell syntax into command_prefix or extra_args, which is not interpreted.

📝 Proposed wording
-`downstream_evaluation` shells out to `python -m lmms_eval` from the GPU worker
-environment. Install `examples/puzzletron/requirements.txt` in that environment;
-it pins the evaluator package used by the checked-in example:
+`downstream_evaluation` runs `python -m lmms_eval` as a subprocess from the GPU
+worker environment. The runner passes an argument list directly and never uses a
+shell. Install `examples/puzzletron/requirements.txt` in that environment; it
+pins the evaluator package used by the checked-in example:
📝 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.

Suggested change
`downstream_evaluation` shells out to `python -m lmms_eval` from the GPU worker
environment. Install `examples/puzzletron/requirements.txt` in that environment;
it pins the evaluator package used by the checked-in example:
`downstream_evaluation` runs `python -m lmms_eval` as a subprocess from the GPU
worker environment. The runner passes an argument list directly and never uses a
shell. Install `examples/puzzletron/requirements.txt` in that environment; it
pins the evaluator package used by the checked-in example:
🤖 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 `@examples/puzzletron/docs/post_mip_pipeline.md` around lines 136 - 138, Update
the downstream_evaluation description to state that it invokes python -m
lmms_eval via subprocess without a shell, rather than saying it “shells out.”
Clarify that command_prefix and extra_args are passed as argv elements and must
not contain shell syntax.

Source: Path instructions

Comment on lines +708 to +710
for key in _LMMS_EVAL_MODEL_ARG_FIELDS:
if key in settings:
derived[key] = settings[key]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

The frozenset iteration still makes --model_args nondeterministic.

_LMMS_EVAL_MODEL_ARG_FIELDS is a frozenset of strings. Its iteration order depends on per-process hash randomization, so the key order in derived and therefore in _model_arg_string changes between runs for identical configuration. The argv persisted to command.json at line 1083 changes too, and the docstring at line 773 promises a deterministic command.

🛠️ 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]
🤖 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 708 - 710, Make
the loop over _LMMS_EVAL_MODEL_ARG_FIELDS deterministic by iterating its keys in
a stable order, such as sorted order, before populating derived. Preserve the
existing filtering and value assignment so _model_arg_string, persisted argv,
and the documented command determinism remain stable for identical
configuration.

Comment on lines +822 to +823
timeout = settings.get("timeout_seconds", settings.get("timeout"))
return argv, env, (float(timeout) if timeout is not None else None)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

A missing timeout_seconds still disables the subprocess timeout.

If the config sets neither timeout_seconds nor timeout, this returns None. _run_lmms_eval_process at line 1035 then calls communicate(timeout=None) and waits forever. A hung vLLM engine blocks the shard worker and holds the GPU allocation with no watchdog.

The failure path at lines 1279-1281 already assumes 3600 as the downstream-evaluation default. Use one shared constant in both places.

🛠️ Proposed fix
-    timeout = settings.get("timeout_seconds", settings.get("timeout"))
-    return argv, env, (float(timeout) if timeout is not None else None)
+    timeout = settings.get("timeout_seconds", settings.get("timeout"))
+    if timeout is None:
+        timeout = _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS
+    return argv, env, float(timeout)

Define the constant once near _LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS and reuse it at line 1279:

_DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS = 3600.0
🤖 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 822 - 823, Define
a shared _DEFAULT_LMMS_EVAL_TIMEOUT_SECONDS constant near
_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS, set to 3600.0, and use it as the
fallback when both timeout_seconds and timeout are absent in the timeout
resolution logic. Replace the hardcoded 3600 downstream-evaluation fallback near
the existing failure path so both _run_lmms_eval_process behavior and failure
handling use the same default.

Comment on lines +1042 to +1044
except subprocess.TimeoutExpired:
_signal_lmms_eval_process_group(process, signal.SIGKILL)
stdout, stderr = process.communicate()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Bound the final communicate() after SIGKILL.

Line 1044 calls communicate() with no timeout. SIGKILL terminates the direct child, but any descendant that escaped the process group with setsid still holds the inherited stdout and stderr pipes. communicate() then blocks until those descendants exit, and the shard worker hangs inside the timeout handler. SECURITY.md requires bounded execution for the subprocess runner.

🛡️ Proposed fix
         except subprocess.TimeoutExpired:
             _signal_lmms_eval_process_group(process, signal.SIGKILL)
-            stdout, stderr = process.communicate()
+            try:
+                stdout, stderr = process.communicate(
+                    timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS
+                )
+            except subprocess.TimeoutExpired:
+                stdout, stderr = None, None

The existing fallback at lines 1051-1052 already substitutes error.output and error.stderr when the values are None.

📝 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.

Suggested change
except subprocess.TimeoutExpired:
_signal_lmms_eval_process_group(process, signal.SIGKILL)
stdout, stderr = process.communicate()
except subprocess.TimeoutExpired:
_signal_lmms_eval_process_group(process, signal.SIGKILL)
try:
stdout, stderr = process.communicate(
timeout=_LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS
)
except subprocess.TimeoutExpired:
stdout, stderr = None, None
🤖 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 1042 - 1044,
Update the TimeoutExpired handler around _signal_lmms_eval_process_group and the
final process.communicate() to use a finite timeout, preserving the existing
fallback handling for missing output and stderr. If the bounded communicate also
times out, ensure the runner does not hang and continues through the established
cleanup/error-result path.

Source: Path instructions

Comment on lines +424 to +429
assert process.communicate_timeouts == [
7.0,
runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS,
None,
]
assert signals == [(5678, signal.SIGTERM), (5678, signal.SIGKILL)]

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

This assertion couples to the unbounded final communicate().

Line 427 asserts the third communicate call receives None. I proposed bounding that call in modelopt/torch/puzzletron/post_mip/runner.py at lines 1042-1044. If you apply that fix, update this expectation to runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS.

🤖 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 424 - 429,
Update the communicate_timeouts expectation in the relevant test to use
runner._LMMS_EVAL_PROCESS_CLEANUP_TIMEOUT_SECONDS for the third call instead of
None, matching the bounded cleanup communicate behavior in the post-MIP runner.

Signed-off-by: Grzegorz Karch <gkarch@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant