Add Puzzletron v2 GPU quality baseline - #2166
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (5)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughThe PR updates Puzzletron v2 setup, configuration provenance, execution identities, distributed launch topology, replacement-scoring finalization, AIPerf policies, CI validation, and unit and GPU integration coverage. ChangesPuzzletron v2 workflow
Estimated code review effort: 5 (Critical) | ~120 minutes Mergeability Score: 🔵 Low · up to The PR is mergeable with owner follow-up to approve the security-scan suppressions and clarify the security-option behavior and required setup defaults so users do not run the workflow with unintended policies or incomplete configuration. Sequence Diagram(s)sequenceDiagram
participant CampaignPlan
participant Controller
participant Worker
participant ArtifactStore
CampaignPlan->>Controller: compiled overrides and execution identity
Controller->>Worker: submit identity-bound attempt
Worker->>ArtifactStore: publish completion artifacts
Controller->>ArtifactStore: validate settled artifacts
Controller->>CampaignPlan: persist completion or stage failure
Possibly related PRs
Suggested reviewers: Estimated code review effort: 5 (Critical) | ~120 minutes Caution Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional.
❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 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: 11
🧹 Nitpick comments (7)
tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py (1)
49-81: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFail loudly if the wizard prompt text drifts.
_DefaultsBackend.selectand_DefaultsBackend.textmatch exact prompt strings such as"Model:","Dataset:", and"Campaign directory:". Ifrun_wizard_v2renames a prompt, the methods fall through todefault. The fixture then builds a campaign against a different model or data source, or against the wrong campaign directory, and the failure appears later as an unrelated stage error.Record which expected prompts were answered, then assert the set in
build_tiny_qwen_campaignafterrun_wizard_v2returns.♻️ Proposed change
class _DefaultsBackend: """Select resolved guided defaults while supplying the campaign directory.""" def __init__(self, campaign_dir: Path) -> None: self.campaign_dir = campaign_dir + self.answered: set[str] = set() def text(self, message: str, default: str) -> Any: if message == "Campaign directory:": + self.answered.add(message) return str(self.campaign_dir) return default def select( self, message: str, choices: Sequence[PromptChoice], default: Any, ) -> Any: if message == "Model:": + self.answered.add(message) return _DEFAULT_MODEL_SOURCE if message == "Dataset:": + self.answered.add(message) return _DEFAULT_DATA_SOURCE if default is not None: return default return next(choice.value for choice in choices if choice.disabled is None)Then verify the contract after the wizard runs:
backend = _DefaultsBackend(campaign_dir) generated = run_wizard_v2( resume=None, defaults_path=defaults_path, backend=backend, ) expected_prompts = {"Campaign directory:", "Model:", "Dataset:"} if backend.answered != expected_prompts: raise AssertionError( f"wizard prompt contract changed; answered {sorted(backend.answered)}" )🤖 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/_test_utils/torch/puzzletron/tiny_qwen_campaign.py` around lines 49 - 81, Track every recognized prompt answered by _DefaultsBackend.text and _DefaultsBackend.select, exposing the collected names through an answered set. In build_tiny_qwen_campaign, immediately after run_wizard_v2 returns, compare that set with {"Campaign directory:", "Model:", "Dataset:"} and raise an AssertionError including the sorted answered prompts when they differ.tests/gpu/torch/puzzletron/test_puzzletron.py (1)
269-287: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAvoid depending on a tie-break order that production may not share.
Lines 276-282 rebuild the expected
best_lmorder by sorting(loss, revision_id)pairs. Lines 342-347 do the same forfastestwith throughput. If two revisions produce an equal loss or an equal throughput, the assertion requires the production node to break the tie by revision id in the same direction. This tiny two-layer model can produce equal metrics across candidate configurations, which would make an expensive GPU test flaky.Assert the selected metric values instead of the exact identifier order. For example, check that the selected losses equal the three smallest losses, and that the selected throughputs equal the two largest throughputs.
Also applies to: 340-347
🤖 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/gpu/torch/puzzletron/test_puzzletron.py` around lines 269 - 287, Update the best_lm and fastest assertions to compare selected metric values rather than exact revision-id ordering. For best_lm, assert its selected losses equal the three smallest values in online_losses; for fastest, assert its selected throughputs equal the two largest values. Retain the candidate counts and selected-observation membership checks, but remove revision-id-based sorting from both assertions.tests/unit/torch/puzzletron/test_width_scenarios.py (1)
111-151: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSplit the marker-currency phases into separate tests.
Lines 111-151 assert five independent invalidation rules against one shared mutable state: missing summary, changed summary, missing manifest, changed
semantic_identity, and marker-to-manifest binding. A failure in an early phase hides every later phase.Consider extracting the marker rules into a parametrized test that builds its own manifest and summary fixtures. The manifest-publication assertions at Lines 88-109 can then stay in this test.
Line 123 also writes
summary_payloadand Line 124 overwrites it immediately, so Line 123 has no effect.🤖 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_width_scenarios.py` around lines 111 - 151, Split the independent marker-currency invalidation checks from the existing test into separate parametrized cases, with each case creating its own manifest, summary, and marker fixtures; keep the manifest-publication assertions in the current test. Cover missing summary, changed summary, missing manifest, changed semantic_identity, and marker-to-manifest binding independently, and remove the redundant summary.write_text call that is immediately overwritten.tests/unit/torch/puzzletron/test_orchestration_executors.py (1)
912-939: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winAdd a width-variation case for the completion directory.
This case changes
replacement_scoring.default_metric, which is areplacement_scoringsemantic section, so the completion identity changes as expected. It does not cover a change toembedding_pruning.widths, which changesFINALIZE_EXPECTED_COMPLETIONSand the set ofwidth-*.donemarker names.Add a case that changes only the width list and asserts the completion directory changes. That case pins the marker-isolation contract described in the
modelopt/torch/puzzletron/orchestration/adapters/pool.pyLines 147-161 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 `@tests/unit/torch/puzzletron/test_orchestration_executors.py` around lines 912 - 939, Extend the orchestration executor test around changed_plan and changed_attempt with a case that modifies only embedding_pruning.widths, then assert the resulting FINALIZE_COMPLETION_DIR differs from the original attempt. Keep replacement_scoring unchanged and preserve the existing marker-isolation behavior, including the width-specific FINALIZE_EXPECTED_COMPLETIONS and width-*.done markers.tests/unit/torch/puzzletron/test_orchestration_task_topology.py (1)
246-252: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a timeout to this subprocess call.
test_nonzero_group_rank_does_not_own_pool_control_pathpassestimeout=10, and this call has no timeout. Ifrun_worker.shblocks, the unit test hangs instead of failing. Unit tests target a few-seconds budget.💚 Proposed fix
result = subprocess.run( ["bash", str(script)], env=env, check=True, capture_output=True, text=True, + timeout=10, )As per path instructions: "Tests placed in the wrong directory for their cost profile (e.g., multi-minute tests under tests/unit, which targets a few-seconds budget)".
🤖 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_orchestration_task_topology.py` around lines 246 - 252, Update the subprocess.run call in test_orchestration_task_topology.py to include a short timeout consistent with test_nonzero_group_rank_does_not_own_pool_control_path, ensuring a blocked run_worker.sh fails promptly while preserving the existing subprocess options.Source: Path instructions
noxfile.py (1)
240-248: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider pinning GPU visibility instead of failing on multi-GPU hosts.
The session fails when
torch.cuda.device_count() != 1. A runner with more than one GPU cannot run the session, even though the test needs only one GPU. SetCUDA_VISIBLE_DEVICES=0for the session and keep the assertion as a post-condition. The CUDA 12.9 check stays unchanged.♻️ Proposed change
def gpu_puzzletron(session): """Run the focused Puzzletron suite in its pinned one-GPU image.""" + session.env["CUDA_VISIBLE_DEVICES"] = os.environ.get("CUDA_VISIBLE_DEVICES", "0") _verify_puzzletron_v2_environment(session)🤖 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 `@noxfile.py` around lines 240 - 248, Update the GPU CI session configuration around the torch validation command to set CUDA_VISIBLE_DEVICES=0, restricting the session to the first GPU before execution. Keep the torch.cuda.device_count() == 1 assertion as a post-condition and leave the CUDA 12.9 version check unchanged.modelopt/torch/puzzletron/orchestration/controller.py (1)
367-399: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConsider caching the stage execution identity.
_stage_execution_identitycallsplan_to_dict(self.plan), which serializes every stage in the campaign, and it callsadapter_for_stage(node).plan(...)when no work plan is passed. The run loop calls it for each stage on every poll iteration through_required_work_is_completed, and again in_persisted_stage_attempts,_recover_failed_stages, and_fail_stage_if_artifacts_did_not_settle.
CampaignPlanis frozen, so the compiled-node projection is stable for the controller's lifetime. Some adapterplan()implementations also read artifacts from disk, for example the post-MIP evaluation candidate count. Memoize the serialized stage nodes, and memoize the identity perstage_idwhen the work plan is supplied by the caller.♻️ Proposed memoization of the compiled stage nodes
+ def _compiled_stage_nodes(self) -> dict[str, Any]: + if self._compiled_stage_nodes_cache is None: + self._compiled_stage_nodes_cache = { + stage["stage_id"]: stage for stage in plan_to_dict(self.plan)["stages"] + } + return self._compiled_stage_nodes_cache + def _stage_execution_identity( self, node: StagePlanNode, work_plan: WorkPlan | None = None, ) -> str: work_plan = work_plan or adapter_for_stage(node).plan(self.plan, node) - compiled_node = next( - stage - for stage in plan_to_dict(self.plan)["stages"] - if stage["stage_id"] == node.stage_id - ) + compiled_node = self._compiled_stage_nodes()[node.stage_id]Initialize
self._compiled_stage_nodes_cache: dict[str, Any] | None = Nonein__init__.🤖 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/orchestration/controller.py` around lines 367 - 399, Optimize _stage_execution_identity by memoizing the serialized stage-node projection: add the proposed _compiled_stage_nodes_cache field in __init__, populate it once from plan_to_dict(self.plan)["stages"], and reuse the matching stage_id entry on subsequent calls. Also cache the computed identity by stage_id only when work_plan is supplied by the caller; continue invoking adapter_for_stage(node).plan(...) and recomputing the identity when it is omitted so artifact-dependent planning remains current.
🤖 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/finalize_replacement_scoring.py`:
- Around line 36-44: Update _successful_manifest_identity to validate that the
parsed manifest is a mapping immediately after json.loads; return None for valid
JSON lists, scalars, or other non-object values before calling get, while
preserving the existing stage, status, and semantic_identity checks for mapping
manifests.
- Around line 31-33: Define the module public API in
finalize_replacement_scoring.py by adding an __all__ declaration listing its
reusable helper symbols, using the module’s existing helper functions and
classes rather than imported dependencies.
In `@examples/puzzletron/README.md`:
- Line 245: Update the documented Python invocation around
verify_installed_vcs_source so it runs with ${MODEL_OPT_ROOT} on PYTHONPATH or
changes into ${MODEL_OPT_ROOT} first, ensuring the
examples.puzzletron.ci_environment import works from the container’s default
/workspace directory.
In `@modelopt/torch/puzzletron/distillation/global_kd_recipe.py`:
- Around line 639-645: Update refresh_realized_checkpoint_config to accept a
caller-configurable trust_remote_code parameter defaulting to False, and pass it
through to AutoConfig.from_pretrained instead of hardcoding True. In the
checkpoint refresh call near the block_configs check, preserve the safe default
and only enable True through a documented, explicitly validated exception path.
In `@modelopt/torch/puzzletron/distributed_eval/automodel_executor.py`:
- Around line 272-273: Update the local torch import in _score: move it to
module scope to match the module’s import convention, unless deferred loading is
required; if it is required, retain the local import and add a brief comment
explaining that requirement.
In `@modelopt/torch/puzzletron/orchestration/adapters/pool.py`:
- Around line 147-161: Update _replacement_completion_identity in
modelopt/torch/puzzletron/orchestration/adapters/pool.py (lines 147-161) to
include the resolved widths from _replacement_widths(plan), ensuring width-list
changes produce a distinct completion directory. Add a test in
tests/unit/torch/puzzletron/test_orchestration_executors.py (lines 912-939) that
changes only embedding_pruning.widths and asserts FINALIZE_COMPLETION_DIR
changes.
In `@modelopt/torch/puzzletron/orchestration/controller.py`:
- Around line 726-737: Update the nested aggregate handling around
adapter.aggregate and _record_stage_aggregation_failure so ManualInputRequired
is not converted into an aggregation failure on the second call. Re-raise
ManualInputRequired to the caller or preserve the existing manual-waiting
behavior by setting self._manual_waiting and returning False, while retaining
failure recording for other OSError, ValueError, and RuntimeError exceptions.
In `@tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py`:
- Around line 97-127: Align TinyQwenCampaign.run’s default subprocess timeout
with the 2400-second caller test budget so two run() calls and their assertions
can complete before pytest times out. Lower the default timeout or provide
explicit smaller timeouts at both call sites, while preserving diagnostic output
from require_success.
In `@tests/gpu/torch/puzzletron/test_puzzletron.py`:
- Around line 135-141: Update the torch.load call in the score_tensors
comprehension to use weights_only=True, since _tensor_values handles only
supported tensor/container payloads. Preserve the existing CPU mapping and
tensor validation behavior; only retain weights_only=False if the payload
requires it, with an inline comment documenting that the checkpoint is generated
inside tmp_path and is not user-supplied.
In `@tests/unit/torch/puzzletron/test_global_kd_canonical.py`:
- Line 744: Move the _WeightedObjectiveMixin import out of the test-local scope
and place it at module scope in test_global_kd_canonical.py; only retain the
lazy import if you document a concrete optional, circular, or heavy-import
requirement.
In `@tests/unit/torch/puzzletron/test_orchestration_lightweight.py`:
- Line 373: Add a concise comment immediately above the in-function import of
stage_is_complete explaining why it remains local, such as avoiding an optional
or heavy puzzletron_orchestrator dependency during test collection. Apply the
same justification to the matching pre-existing local import if needed, or move
both imports to module scope if no such constraint exists.
---
Nitpick comments:
In `@modelopt/torch/puzzletron/orchestration/controller.py`:
- Around line 367-399: Optimize _stage_execution_identity by memoizing the
serialized stage-node projection: add the proposed _compiled_stage_nodes_cache
field in __init__, populate it once from plan_to_dict(self.plan)["stages"], and
reuse the matching stage_id entry on subsequent calls. Also cache the computed
identity by stage_id only when work_plan is supplied by the caller; continue
invoking adapter_for_stage(node).plan(...) and recomputing the identity when it
is omitted so artifact-dependent planning remains current.
In `@noxfile.py`:
- Around line 240-248: Update the GPU CI session configuration around the torch
validation command to set CUDA_VISIBLE_DEVICES=0, restricting the session to the
first GPU before execution. Keep the torch.cuda.device_count() == 1 assertion as
a post-condition and leave the CUDA 12.9 version check unchanged.
In `@tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py`:
- Around line 49-81: Track every recognized prompt answered by
_DefaultsBackend.text and _DefaultsBackend.select, exposing the collected names
through an answered set. In build_tiny_qwen_campaign, immediately after
run_wizard_v2 returns, compare that set with {"Campaign directory:", "Model:",
"Dataset:"} and raise an AssertionError including the sorted answered prompts
when they differ.
In `@tests/gpu/torch/puzzletron/test_puzzletron.py`:
- Around line 269-287: Update the best_lm and fastest assertions to compare
selected metric values rather than exact revision-id ordering. For best_lm,
assert its selected losses equal the three smallest values in online_losses; for
fastest, assert its selected throughputs equal the two largest values. Retain
the candidate counts and selected-observation membership checks, but remove
revision-id-based sorting from both assertions.
In `@tests/unit/torch/puzzletron/test_orchestration_executors.py`:
- Around line 912-939: Extend the orchestration executor test around
changed_plan and changed_attempt with a case that modifies only
embedding_pruning.widths, then assert the resulting FINALIZE_COMPLETION_DIR
differs from the original attempt. Keep replacement_scoring unchanged and
preserve the existing marker-isolation behavior, including the width-specific
FINALIZE_EXPECTED_COMPLETIONS and width-*.done markers.
In `@tests/unit/torch/puzzletron/test_orchestration_task_topology.py`:
- Around line 246-252: Update the subprocess.run call in
test_orchestration_task_topology.py to include a short timeout consistent with
test_nonzero_group_rank_does_not_own_pool_control_path, ensuring a blocked
run_worker.sh fails promptly while preserving the existing subprocess options.
In `@tests/unit/torch/puzzletron/test_width_scenarios.py`:
- Around line 111-151: Split the independent marker-currency invalidation checks
from the existing test into separate parametrized cases, with each case creating
its own manifest, summary, and marker fixtures; keep the manifest-publication
assertions in the current test. Cover missing summary, changed summary, missing
manifest, changed semantic_identity, and marker-to-manifest binding
independently, and remove the redundant summary.write_text call that is
immediately overwritten.
🪄 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: 4cdead70-f756-4c0f-a848-37ab87eb91c9
📒 Files selected for processing (50)
examples/puzzletron/README.mdexamples/puzzletron/ci_environment.jsonexamples/puzzletron/ci_environment.pyexamples/puzzletron/distributed_eval/run_coordinator.shexamples/puzzletron/distributed_eval/run_depth_pool.shexamples/puzzletron/distributed_eval/run_replacement_pool.shexamples/puzzletron/distributed_eval/run_worker.shexamples/puzzletron/embedding_pipeline.pyexamples/puzzletron/finalize_replacement_scoring.pyexamples/puzzletron/main.pyexamples/puzzletron/run_axis_diagnostic_worker.pyexamples/puzzletron/tokenize_data.pymodelopt/torch/puzzletron/benchmarks/aiperf.pymodelopt/torch/puzzletron/distillation/global_kd_recipe.pymodelopt/torch/puzzletron/distributed_eval/automodel_executor.pymodelopt/torch/puzzletron/manifest.pymodelopt/torch/puzzletron/orchestration/adapters/pool.pymodelopt/torch/puzzletron/orchestration/adapters/post_mip.pymodelopt/torch/puzzletron/orchestration/compiler.pymodelopt/torch/puzzletron/orchestration/config.pymodelopt/torch/puzzletron/orchestration/controller.pymodelopt/torch/puzzletron/orchestration/schema.pymodelopt/torch/puzzletron/orchestration/task_launcher.pymodelopt/torch/puzzletron/pipeline_config.pymodelopt/torch/puzzletron/post_mip/runner.pymodelopt/torch/puzzletron/stage_runner.pymodelopt/torch/puzzletron/stages/graph.pynoxfile.pypuzzletron_setup/bundle.pytests/_test_utils/torch/puzzletron/tiny_qwen_campaign.pytests/_test_utils/torch/puzzletron/utils.pytests/gpu/torch/puzzletron/test_puzzletron.pytests/unit/torch/puzzletron/conftest.pytests/unit/torch/puzzletron/test_aiperf_context_capacity.pytests/unit/torch/puzzletron/test_automodel_solution_scoring.pytests/unit/torch/puzzletron/test_ci_environment.pytests/unit/torch/puzzletron/test_diagnostic_scoring_config.pytests/unit/torch/puzzletron/test_example_runner.pytests/unit/torch/puzzletron/test_global_kd_canonical.pytests/unit/torch/puzzletron/test_orchestration_executors.pytests/unit/torch/puzzletron/test_orchestration_lightweight.pytests/unit/torch/puzzletron/test_orchestration_shutdown_progress.pytests/unit/torch/puzzletron/test_orchestration_task_topology.pytests/unit/torch/puzzletron/test_post_mip_adapter.pytests/unit/torch/puzzletron/test_post_mip_runner.pytests/unit/torch/puzzletron/test_setup_bundle.pytests/unit/torch/puzzletron/test_stage_graph.pytests/unit/torch/puzzletron/test_tokenize_data.pytests/unit/torch/puzzletron/test_width_scenarios.pytests/unit/torch/puzzletron/test_width_slice_equivalence.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## feature/puzzletron_v2 #2166 +/- ##
=========================================================
+ Coverage 53.12% 62.59% +9.47%
=========================================================
Files 706 709 +3
Lines 91565 91929 +364
=========================================================
+ Hits 48640 57546 +8906
+ Misses 42925 34383 -8542
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
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: 1
🧹 Nitpick comments (2)
tests/unit/torch/puzzletron/test_orchestration_executors.py (1)
642-666: 🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick winAdd a negative case for the default-off policy.
The test proves both flags appear when the configuration enables them. No test proves they are absent when the configuration omits them.
Falseis the secure default for both policies, so a regression that always appends the flags would still pass. Add a second assertion path with an emptyexperiment_config.✅ Proposed additional coverage for the default-off policy
assert "--trust-remote-code" in attempt.command.argv assert "--allow-aiperf-v011-online-tokenizer-resolution" in attempt.command.argv + + default_plan = replace(plan, experiment_config={}) + default_adapter = adapter_for_stage(node) + default_attempt = default_adapter.command( + plan=default_plan, + node=node, + item=default_adapter.plan(default_plan, node).items[0], + attempt_id="a2", + runner=runner, + ) + + assert "--trust-remote-code" not in default_attempt.command.argv + assert "--allow-aiperf-v011-online-tokenizer-resolution" not in default_attempt.command.argv
replacecomes fromdataclasses. If it is not already imported in this module, addfrom dataclasses import replaceat the top, or construct a secondCampaignPlanexplicitly.🤖 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_orchestration_executors.py` around lines 642 - 666, Add a default-off test path alongside the existing positive assertions in the campaign plan orchestration test, using dataclasses.replace or an equivalent plan reconstruction to set experiment_config to an empty mapping. Re-plan and build the command through adapter.plan and adapter.command, then assert both --trust-remote-code and --allow-aiperf-v011-online-tokenizer-resolution are absent from the resulting argv.modelopt/torch/puzzletron/orchestration/adapters/post_mip.py (1)
44-51: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated dual-path import of the post-MIP identity module. Both files select the identity module with the same
__package__.startswith("puzzletron_orchestrator.")branch, so both depend on the identity-API surface existing on both import paths. Confirm every consumed symbol resolves on thepuzzletron_orchestratorpath, then consider centralizing the resolver so only one file owns the branch.
modelopt/torch/puzzletron/orchestration/adapters/post_mip.py#L44-L51: confirmprepare_post_mip_candidate_ledger,expected_post_mip_execution_contract,expected_post_mip_candidate_count, andPostMIPExecutionContractUnavailableresolve on both paths, and export_post_mip_identity_apias the single resolver.modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py#L466-L478: confirmexpected_post_mip_execution_identityresolves on both paths, and reuse the shared resolver instead of repeating the__package__branch.🤖 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/orchestration/adapters/post_mip.py` around lines 44 - 51, The post-MIP identity resolver is duplicated and must be centralized. In modelopt/torch/puzzletron/orchestration/adapters/post_mip.py:44-51, verify that prepare_post_mip_candidate_ledger, expected_post_mip_execution_contract, expected_post_mip_candidate_count, and PostMIPExecutionContractUnavailable resolve through both import paths, then expose _post_mip_identity_api as the single resolver. In modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py:466-478, verify expected_post_mip_execution_identity resolves through both paths and replace the repeated __package__ branch with reuse of _post_mip_identity_api.
🤖 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 `@tests/unit/torch/puzzletron/test_width_scenarios.py`:
- Around line 153-170: Update finalization_marker_is_current to validate that
the decoded manifest is an object before any manifest.get(...) calls. Return
False immediately for arrays, strings, integers, null, or other non-object
values, while preserving the existing identity and report checks for valid
object manifests.
---
Nitpick comments:
In `@modelopt/torch/puzzletron/orchestration/adapters/post_mip.py`:
- Around line 44-51: The post-MIP identity resolver is duplicated and must be
centralized. In
modelopt/torch/puzzletron/orchestration/adapters/post_mip.py:44-51, verify that
prepare_post_mip_candidate_ledger, expected_post_mip_execution_contract,
expected_post_mip_candidate_count, and PostMIPExecutionContractUnavailable
resolve through both import paths, then expose _post_mip_identity_api as the
single resolver. In
modelopt/torch/puzzletron/orchestration/adapters/stage_compat.py:466-478, verify
expected_post_mip_execution_identity resolves through both paths and replace the
repeated __package__ branch with reuse of _post_mip_identity_api.
In `@tests/unit/torch/puzzletron/test_orchestration_executors.py`:
- Around line 642-666: Add a default-off test path alongside the existing
positive assertions in the campaign plan orchestration test, using
dataclasses.replace or an equivalent plan reconstruction to set
experiment_config to an empty mapping. Re-plan and build the command through
adapter.plan and adapter.command, then assert both --trust-remote-code and
--allow-aiperf-v011-online-tokenizer-resolution are absent from the resulting
argv.
🪄 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: 2ba6374c-a6be-4a89-aefb-909deb005d5b
📒 Files selected for processing (29)
examples/puzzletron/README.mdexamples/puzzletron/finalize_replacement_scoring.pyexamples/puzzletron/run_profile_aiperf_worker.pymodelopt/torch/puzzletron/benchmarks/aiperf.pymodelopt/torch/puzzletron/distillation/global_kd_recipe.pymodelopt/torch/puzzletron/distributed_eval/automodel_executor.pymodelopt/torch/puzzletron/orchestration/adapters/base.pymodelopt/torch/puzzletron/orchestration/adapters/post_mip.pymodelopt/torch/puzzletron/orchestration/adapters/sharded.pymodelopt/torch/puzzletron/orchestration/adapters/stage_compat.pymodelopt/torch/puzzletron/orchestration/controller.pymodelopt/torch/puzzletron/post_mip/identity.pymodelopt/torch/puzzletron/post_mip/runner.pymodelopt/torch/puzzletron/stages/future.pymodelopt/torch/puzzletron/utils/vllm_adapter.pynoxfile.pytests/_test_utils/torch/puzzletron/tiny_qwen_campaign.pytests/unit/torch/puzzletron/test_aiperf_context_capacity.pytests/unit/torch/puzzletron/test_ci_environment.pytests/unit/torch/puzzletron/test_global_kd_canonical.pytests/unit/torch/puzzletron/test_orchestration_executors.pytests/unit/torch/puzzletron/test_orchestration_lightweight.pytests/unit/torch/puzzletron/test_orchestration_shutdown_progress.pytests/unit/torch/puzzletron/test_orchestration_task_topology.pytests/unit/torch/puzzletron/test_post_mip_adapter.pytests/unit/torch/puzzletron/test_post_mip_execution_identity.pytests/unit/torch/puzzletron/test_post_mip_runner.pytests/unit/torch/puzzletron/test_vllm_axis_contract.pytests/unit/torch/puzzletron/test_width_scenarios.py
🚧 Files skipped from review as they are similar to previous changes (11)
- tests/unit/torch/puzzletron/test_post_mip_adapter.py
- examples/puzzletron/README.md
- modelopt/torch/puzzletron/distributed_eval/automodel_executor.py
- tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py
- examples/puzzletron/finalize_replacement_scoring.py
- modelopt/torch/puzzletron/distillation/global_kd_recipe.py
- tests/unit/torch/puzzletron/test_post_mip_runner.py
- noxfile.py
- modelopt/torch/puzzletron/orchestration/controller.py
- tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py
- tests/unit/torch/puzzletron/test_global_kd_canonical.py
55478b8 to
de15733
Compare
There was a problem hiding this comment.
♻️ Duplicate comments (1)
tests/gpu/torch/puzzletron/test_puzzletron.py (1)
135-141: 🔒 Security & Privacy | 🔴 Critical | ⚡ Quick winReplace
weights_only=Falseon line 138.The coding guidelines prohibit
torch.load(..., weights_only=False)without a documented exception. Static analysis flags the same line._tensor_valuesonly traverses tensors, dicts, lists, and tuples, whichweights_only=Truesupports. If pickle is genuinely required, keepweights_only=Falseand add an inline comment that states the file is produced by this test insidetmp_path.🔒️ Proposed fix
score_tensors = [ tensor for score_file in score_files - for tensor in _tensor_values(torch.load(score_file, map_location="cpu", weights_only=False)) + for tensor in _tensor_values(torch.load(score_file, map_location="cpu", weights_only=True)) ]As per coding guidelines: "Do not use
torch.load(..., weights_only=False)unless a documented exception is provided."🤖 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/gpu/torch/puzzletron/test_puzzletron.py` around lines 135 - 141, Replace weights_only=False with weights_only=True in the score-loading comprehension around _tensor_values and preserve the existing tensor validation. Only retain the unsafe setting if loading requires pickle, in which case add an inline comment documenting that the file is produced by this test within tmp_path.Sources: Coding guidelines, Linters/SAST tools
🤖 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 `@tests/gpu/torch/puzzletron/test_puzzletron.py`:
- Around line 135-141: Replace weights_only=False with weights_only=True in the
score-loading comprehension around _tensor_values and preserve the existing
tensor validation. Only retain the unsafe setting if loading requires pickle, in
which case add an inline comment documenting that the file is produced by this
test within tmp_path.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fa954247-b55d-4fdf-a7e0-c478fe152399
📒 Files selected for processing (4)
examples/puzzletron/README.mdnoxfile.pytests/gpu/torch/puzzletron/test_puzzletron.pytests/unit/torch/puzzletron/test_stage_graph.py
🚧 Files skipped from review as they are similar to previous changes (2)
- examples/puzzletron/README.md
- noxfile.py
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: 2
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/stages/future.py (1)
388-394: 🔒 Security & Privacy | 🟠 Major | ⚡ Quick winReject non-boolean policy values before
bool()coercion.
aiperf_stage()receives a plain dictionary without schema validation. A string such as"false"therefore becomesTrue, which can enabletrust_remote_codeor online tokenizer resolution. Reject non-boolean values at the configuration boundary.🤖 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/future.py` around lines 388 - 394, Update aiperf_stage() configuration handling to validate trust_remote_code and allow_aiperf_v011_online_tokenizer_resolution before coercion: accept only boolean values, reject non-boolean inputs such as strings, and preserve the existing default and precedence behavior for omitted settings.Source: Coding guidelines
🧹 Nitpick comments (1)
examples/puzzletron/main.py (1)
418-426: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument why these imports remain local.
The helper uses local imports to support both package and standalone entry points. Add a brief comment that explains why the imports cannot move to module scope. Match the explanation used by
_run_embedding_stage.As per coding guidelines, keep Python imports at the top unless a local import is justified and documented.
🤖 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/main.py` around lines 418 - 426, Add a brief comment in _run_tokenize_data_stage before the conditional imports, matching the rationale used by _run_embedding_stage: the imports must remain local to support both package and standalone entry points. Keep the existing import behavior unchanged.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 `@modelopt/torch/puzzletron/orchestration/adapters/post_mip.py`:
- Around line 266-267: Remove the # nosec B603 suppression from subprocess.run
in modelopt/torch/puzzletron/orchestration/adapters/post_mip.py (lines 266-267)
and the # nosec B404 suppression from its subprocess import (line 23). Also
remove the corresponding B603 suppression from subprocess.Popen in
modelopt/torch/puzzletron/post_mip/runner.py (lines 997-998) and B404
suppression from its subprocess import (lines 30-31); leave the subprocess
behavior unchanged.
In `@modelopt/torch/puzzletron/stages/future.py`:
- Around line 470-484: Validate the configured checkpoint value before the
comprehension in the checkpoint-selection flow, ensuring it is a list or tuple
and rejecting scalar strings (including truthy strings) with the existing
configuration error behavior. Keep the configured-entry conversion unchanged for
valid collections, and preserve the fallback branches in the surrounding
checkpoint resolution logic.
---
Outside diff comments:
In `@modelopt/torch/puzzletron/stages/future.py`:
- Around line 388-394: Update aiperf_stage() configuration handling to validate
trust_remote_code and allow_aiperf_v011_online_tokenizer_resolution before
coercion: accept only boolean values, reject non-boolean inputs such as strings,
and preserve the existing default and precedence behavior for omitted settings.
---
Nitpick comments:
In `@examples/puzzletron/main.py`:
- Around line 418-426: Add a brief comment in _run_tokenize_data_stage before
the conditional imports, matching the rationale used by _run_embedding_stage:
the imports must remain local to support both package and standalone entry
points. Keep the existing import behavior unchanged.
🪄 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: f6eaaff3-173c-41c8-91d5-938f003b16dc
📒 Files selected for processing (20)
examples/puzzletron/finalize_replacement_scoring.pyexamples/puzzletron/main.pymodelopt/torch/puzzletron/distillation/global_kd_recipe.pymodelopt/torch/puzzletron/distributed_eval/automodel_executor.pymodelopt/torch/puzzletron/manifest.pymodelopt/torch/puzzletron/orchestration/adapters/base.pymodelopt/torch/puzzletron/orchestration/adapters/pool.pymodelopt/torch/puzzletron/orchestration/adapters/post_mip.pymodelopt/torch/puzzletron/orchestration/adapters/sharded.pymodelopt/torch/puzzletron/orchestration/compiler.pymodelopt/torch/puzzletron/orchestration/task_launcher.pymodelopt/torch/puzzletron/post_mip/identity.pymodelopt/torch/puzzletron/post_mip/runner.pymodelopt/torch/puzzletron/stages/future.pypuzzletron_setup/bundle.pypyproject.tomltests/gpu/torch/puzzletron/test_puzzletron.pytests/unit/torch/puzzletron/test_global_kd_canonical.pytests/unit/torch/puzzletron/test_post_mip_execution_identity.pytests/unit/torch/puzzletron/test_vllm_axis_contract.py
🚧 Files skipped from review as they are similar to previous changes (14)
- puzzletron_setup/bundle.py
- tests/gpu/torch/puzzletron/test_puzzletron.py
- modelopt/torch/puzzletron/orchestration/adapters/base.py
- modelopt/torch/puzzletron/post_mip/identity.py
- modelopt/torch/puzzletron/orchestration/task_launcher.py
- modelopt/torch/puzzletron/manifest.py
- examples/puzzletron/finalize_replacement_scoring.py
- modelopt/torch/puzzletron/orchestration/adapters/pool.py
- modelopt/torch/puzzletron/orchestration/adapters/sharded.py
- tests/unit/torch/puzzletron/test_vllm_axis_contract.py
- modelopt/torch/puzzletron/distributed_eval/automodel_executor.py
- tests/unit/torch/puzzletron/test_post_mip_execution_identity.py
- modelopt/torch/puzzletron/distillation/global_kd_recipe.py
- tests/unit/torch/puzzletron/test_global_kd_canonical.py
e013fe8 to
e3f78c2
Compare
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: 2
🤖 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 29-31: Remove the # nosec B404 suppression from the subprocess
import at modelopt/torch/puzzletron/post_mip/runner.py lines 29-31. At lines
1004-1005, remove # nosec B603 from the subprocess.Popen call and update that
implementation so Bandit passes without suppression, preserving the existing
subprocess behavior; document any necessary security exception and obtain the
required approval if no compliant implementation exists.
In `@tests/unit/torch/puzzletron/test_future_stages.py`:
- Around line 55-60: Move the future import to module scope in
tests/unit/torch/puzzletron/test_future_stages.py, alongside the existing
imports. Remove the local future imports from lines 55-60 and 68-81, leaving
both tests to use the module-level import.
🪄 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: 62d16e1c-e559-4917-8f64-b00b704bfaa9
📒 Files selected for processing (9)
examples/puzzletron/main.pymodelopt/torch/puzzletron/orchestration/adapters/sharded.pymodelopt/torch/puzzletron/post_mip/runner.pymodelopt/torch/puzzletron/security_policy.pymodelopt/torch/puzzletron/stages/future.pytests/unit/torch/puzzletron/test_future_stages.pytests/unit/torch/puzzletron/test_orchestration_executors.pytests/unit/torch/puzzletron/test_orchestration_lightweight.pytests/unit/torch/puzzletron/test_post_mip_runner.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/puzzletron/main.py
- modelopt/torch/puzzletron/orchestration/adapters/sharded.py
- modelopt/torch/puzzletron/stages/future.py
6a6b3d0 to
79eb6cd
Compare
|
@coderabbitai review |
✅ Action performedReview finished.
|
|
@coderabbitai review |
|
79eb6cd to
a34b939
Compare
|
@coderabbitai full review |
✅ Action performedFull review finished. |
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: 6
🧹 Nitpick comments (5)
tests/unit/torch/puzzletron/test_orchestration_task_topology.py (1)
261-285: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCapture script output so a failure is diagnosable.
subprocess.run(..., check=True)without capture sends the script output to the test runner's streams and raisesCalledProcessErrorwithout the script message attached. Capture the output and assert that no pool-control command ran. The current test proves only a zero exit code.♻️ Proposed change
- subprocess.run(["bash", str(script)], env=env, check=True, timeout=10) + result = subprocess.run( + ["bash", str(script)], + env=env, + check=True, + capture_output=True, + text=True, + timeout=10, + ) + + assert "drain" not in result.stdoutAdjust the assertion to the marker that group rank 0 prints when it owns the control path.
🤖 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_orchestration_task_topology.py` around lines 261 - 285, Update test_nonzero_group_rank_does_not_own_pool_control_path to capture the subprocess output while retaining check=True, then assert that the output does not contain the group-rank-0 pool-control marker. Preserve the existing environment and zero-exit validation, but make the test verify that no pool-control command ran.modelopt/torch/puzzletron/orchestration/controller.py (1)
371-407: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the stage execution identity per stage.
_stage_execution_identitycallsadapter.plan(...),plan_to_dict(self.plan),semantic_stage_config(...), andadapter.execution_identity_projection(...)on every invocation. The run loop reaches it through_required_completed_attempts,_persisted_stage_attempts,_fail_stage_if_artifacts_did_not_settle, and_bind_attempt_to_stage_execution, so one poll iteration recomputes it several times for every stage.PostMIPAdapter.planandexecution_identity_projectionread campaign artifacts from disk, so the cost is filesystem I/O on the controller thread.The inputs are the immutable plan plus adapter projections, so memoize per
stage_idfor the duration of one loop iteration, or cacheplan_to_dict(self.plan)once in__init__.🤖 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/orchestration/controller.py` around lines 371 - 407, Cache the immutable plan-derived data and memoize the stage execution identity per stage within each controller loop iteration. Update _stage_execution_identity and its callers (_required_completed_attempts, _persisted_stage_attempts, _fail_stage_if_artifacts_did_not_settle, and _bind_attempt_to_stage_execution) to reuse the cached value, while preserving recalculation across loop iterations so filesystem-backed adapter projections remain current.examples/puzzletron/ci_environment.py (1)
33-55: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winTranslate git and metadata failures into a clear verification error.
Users run this helper as the documented environment check. Three failure paths currently surface as raw tracebacks:
metadata.distribution(package)raisesPackageNotFoundErrorwhen the package is not installed.read_text("direct_url.json")returnsNonefor a plain index install; the function then reports a(None, None)mismatch instead of stating that the package was not installed from a repository.subprocess.check_outputraisesCalledProcessErrorwhenrootis not a git checkout.Add a timeout to the git calls and wrap these cases in
RuntimeErrorwith the package name.♻️ Proposed error translation
def _installed_vcs_source(package: str) -> tuple[str | None, str | None]: - payload = json.loads(metadata.distribution(package).read_text("direct_url.json") or "{}") + try: + distribution = metadata.distribution(package) + except metadata.PackageNotFoundError as error: + raise RuntimeError(f"Pinned Puzzletron dependency {package!r} is not installed") from error + payload = json.loads(distribution.read_text("direct_url.json") or "{}") vcs_info = payload.get("vcs_info") or {} if vcs_info.get("commit_id"): return payload.get("url"), vcs_info["commit_id"] if (payload.get("dir_info") or {}).get("editable") and str(payload.get("url", "")).startswith( "file:" ): root = unquote(urlparse(payload["url"]).path) - repository = subprocess.check_output( - ["git", "-C", root, "remote", "get-url", "origin"], text=True - ).strip() - commit = subprocess.check_output( - ["git", "-C", root, "rev-parse", "HEAD"], text=True - ).strip() - dirty = subprocess.check_output( - ["git", "-C", root, "status", "--porcelain", "--untracked-files=all"], - text=True, - ).strip() + def _git(*args: str) -> str: + try: + # Fixed git argv; no shell is used. + return subprocess.check_output( # nosec B603 + ["git", "-C", root, *args], text=True, timeout=60 + ).strip() + except (subprocess.SubprocessError, OSError) as error: + raise RuntimeError( + f"Pinned Puzzletron dependency {package!r} git query failed in {root!r}" + ) from error + + repository = _git("remote", "get-url", "origin") + commit = _git("rev-parse", "HEAD") + dirty = _git("status", "--porcelain", "--untracked-files=all") if dirty: raise RuntimeError(f"Pinned Puzzletron dependency {package!r} is dirty: {dirty}") return repository, commitThe
# nosecmarker is unrelated to the security-guideline restriction only if the repository already accepts inline markers for fixed argv lists. Confirm the project Bandit policy before adding it; the coding guidelines forbid# nosecas a bypass.🤖 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/ci_environment.py` around lines 33 - 55, Update _installed_vcs_source to translate missing distributions, absent direct_url.json metadata, and CalledProcessError failures from git commands into RuntimeError messages that include the package name and clearly describe the verification failure. Add a timeout to each subprocess.check_output call, preserving the existing dirty-check behavior and repository/commit return values; do not add a # nosec marker unless it is already permitted by the project’s Bandit policy.Source: Coding guidelines
modelopt/torch/puzzletron/distillation/global_kd_recipe.py (1)
666-667: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove
Pathto the module imports.Line 666 adds a local standard-library import without a circular, optional, or heavy-import justification. Import
Pathat module scope.As per 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/distillation/global_kd_recipe.py` around lines 666 - 667, Move the pathlib Path import from the local scope to the module-level imports in global_kd_recipe.py, alongside the other standard-library imports, and remove the local import while preserving all existing Path usage.Source: Coding guidelines
modelopt/torch/puzzletron/orchestration/adapters/post_mip.py (1)
46-55: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueGuard the package check against a
None__package__.Line 49 calls
__package__.startswith(...).__package__isNonewhen a module is executed in some non-package contexts, which raisesAttributeErrorinstead of falling back to the relative import. Use a defensive default so the loader always resolves one of the two import paths.♻️ Proposed change
- if __package__.startswith("puzzletron_orchestrator."): + if (__package__ or "").startswith("puzzletron_orchestrator."):🤖 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/orchestration/adapters/post_mip.py` around lines 46 - 55, Update _post_mip_identity_api to handle a None __package__ before calling startswith, using a defensive default that selects the relative import path when no package is defined while preserving the existing orchestrator-package branch.
🤖 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/run_profile_aiperf_worker.py`:
- Around line 344-348: Update the argparse definitions near trust-remote-code
and allow-aiperf-v011-online-tokenizer-resolution with explicit help text
describing the security opt-ins: trust remote code only for trusted model
sources, and the tokenizer option enabling online resolution. Add matching
documentation to run_worker and the required SECURITY.md guidance for these
security-sensitive exceptions.
In `@modelopt/torch/puzzletron/distillation/global_kd_recipe.py`:
- Around line 674-678: Update the trust_remote_code argument in the global KD
checkpoint refresh flow to default missing values to False while rejecting any
configured value that is not a Boolean. Remove the bool() coercion around
_config_value(model_config, "trust_remote_code"), preserve explicit Boolean
opt-in, and pass the validated value to refresh_realized_checkpoint_config.
In `@modelopt/torch/puzzletron/orchestration/adapters/sharded.py`:
- Around line 260-263: Update the trust_remote_code validation around
require_boolean_policy to pass the configuration path corresponding to the
value’s source: use aiperf.trust_remote_code when that key is present, otherwise
model.trust_remote_code. Update the parametrized expectation in
test_orchestration_executors.py to assert the model key is reported for invalid
model-sourced values.
- Around line 20-21: Address the inline Bandit suppression approvals for the
subprocess import in sharded.py (lines 20-21), the aggregation subprocess.run
call in sharded.py (line 389), and the os.execvpe call in task_launcher.py
(lines 275-276): either document an explicit security justification for each #
nosec marker in the PR description with required codeowner approval, or remove
the inline markers and configure the approved skips centrally in Bandit
configuration.
Apply the same fix in `@modelopt/torch/puzzletron/post_mip/runner.py` around lines
29 - 31: Same prohibited subprocess suppression.
Apply the same fix in
`@modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` around lines 22 -
25: Same prohibited subprocess-call suppression.
In `@modelopt/torch/puzzletron/orchestration/controller.py`:
- Around line 448-462: Update _completed_work_artifact_settling_elapsed and
controller initialization to track the first time each completed work item is
observed, using self._first_completion_observed keyed by the node/work
identifier. Record the current time on first observation, then compute settling
elapsed from max(completion_time, first_observed_time) so resumed work retains a
full settling window while normal completion timing remains unchanged.
In `@puzzletron_setup/bundle.py`:
- Line 470: Remove the "sequence_length" entry from the rendered data mapping
used to construct PuzzletronDataSpec, leaving sequence_length available only
through its derived property and preserving the remaining constructor inputs.
---
Nitpick comments:
In `@examples/puzzletron/ci_environment.py`:
- Around line 33-55: Update _installed_vcs_source to translate missing
distributions, absent direct_url.json metadata, and CalledProcessError failures
from git commands into RuntimeError messages that include the package name and
clearly describe the verification failure. Add a timeout to each
subprocess.check_output call, preserving the existing dirty-check behavior and
repository/commit return values; do not add a # nosec marker unless it is
already permitted by the project’s Bandit policy.
In `@modelopt/torch/puzzletron/distillation/global_kd_recipe.py`:
- Around line 666-667: Move the pathlib Path import from the local scope to the
module-level imports in global_kd_recipe.py, alongside the other
standard-library imports, and remove the local import while preserving all
existing Path usage.
In `@modelopt/torch/puzzletron/orchestration/adapters/post_mip.py`:
- Around line 46-55: Update _post_mip_identity_api to handle a None __package__
before calling startswith, using a defensive default that selects the relative
import path when no package is defined while preserving the existing
orchestrator-package branch.
In `@modelopt/torch/puzzletron/orchestration/controller.py`:
- Around line 371-407: Cache the immutable plan-derived data and memoize the
stage execution identity per stage within each controller loop iteration. Update
_stage_execution_identity and its callers (_required_completed_attempts,
_persisted_stage_attempts, _fail_stage_if_artifacts_did_not_settle, and
_bind_attempt_to_stage_execution) to reuse the cached value, while preserving
recalculation across loop iterations so filesystem-backed adapter projections
remain current.
In `@tests/unit/torch/puzzletron/test_orchestration_task_topology.py`:
- Around line 261-285: Update
test_nonzero_group_rank_does_not_own_pool_control_path to capture the subprocess
output while retaining check=True, then assert that the output does not contain
the group-rank-0 pool-control marker. Preserve the existing environment and
zero-exit validation, but make the test verify that no pool-control command ran.
🪄 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: 46642561-e51b-42b7-b7a3-1ba23e4eb356
📒 Files selected for processing (62)
examples/puzzletron/README.mdexamples/puzzletron/ci_environment.jsonexamples/puzzletron/ci_environment.pyexamples/puzzletron/distributed_eval/run_coordinator.shexamples/puzzletron/distributed_eval/run_depth_pool.shexamples/puzzletron/distributed_eval/run_replacement_pool.shexamples/puzzletron/distributed_eval/run_worker.shexamples/puzzletron/embedding_pipeline.pyexamples/puzzletron/finalize_replacement_scoring.pyexamples/puzzletron/main.pyexamples/puzzletron/run_axis_diagnostic_worker.pyexamples/puzzletron/run_profile_aiperf_worker.pyexamples/puzzletron/tokenize_data.pymodelopt/torch/puzzletron/benchmarks/aiperf.pymodelopt/torch/puzzletron/distillation/global_kd_recipe.pymodelopt/torch/puzzletron/distributed_eval/automodel_executor.pymodelopt/torch/puzzletron/manifest.pymodelopt/torch/puzzletron/orchestration/adapters/base.pymodelopt/torch/puzzletron/orchestration/adapters/pool.pymodelopt/torch/puzzletron/orchestration/adapters/post_mip.pymodelopt/torch/puzzletron/orchestration/adapters/sharded.pymodelopt/torch/puzzletron/orchestration/adapters/stage_compat.pymodelopt/torch/puzzletron/orchestration/compiler.pymodelopt/torch/puzzletron/orchestration/config.pymodelopt/torch/puzzletron/orchestration/controller.pymodelopt/torch/puzzletron/orchestration/schema.pymodelopt/torch/puzzletron/orchestration/task_launcher.pymodelopt/torch/puzzletron/pipeline_config.pymodelopt/torch/puzzletron/post_mip/identity.pymodelopt/torch/puzzletron/post_mip/runner.pymodelopt/torch/puzzletron/security_policy.pymodelopt/torch/puzzletron/stage_runner.pymodelopt/torch/puzzletron/stages/future.pymodelopt/torch/puzzletron/stages/graph.pymodelopt/torch/puzzletron/utils/vllm_adapter.pynoxfile.pypuzzletron_setup/bundle.pypyproject.tomltests/_test_utils/torch/puzzletron/tiny_qwen_campaign.pytests/_test_utils/torch/puzzletron/utils.pytests/gpu/torch/puzzletron/test_puzzletron.pytests/unit/torch/puzzletron/conftest.pytests/unit/torch/puzzletron/test_aiperf_context_capacity.pytests/unit/torch/puzzletron/test_automodel_solution_scoring.pytests/unit/torch/puzzletron/test_ci_environment.pytests/unit/torch/puzzletron/test_diagnostic_scoring_config.pytests/unit/torch/puzzletron/test_example_runner.pytests/unit/torch/puzzletron/test_future_stages.pytests/unit/torch/puzzletron/test_global_kd_canonical.pytests/unit/torch/puzzletron/test_orchestration_executors.pytests/unit/torch/puzzletron/test_orchestration_lightweight.pytests/unit/torch/puzzletron/test_orchestration_shutdown_progress.pytests/unit/torch/puzzletron/test_orchestration_task_topology.pytests/unit/torch/puzzletron/test_post_mip_adapter.pytests/unit/torch/puzzletron/test_post_mip_execution_identity.pytests/unit/torch/puzzletron/test_post_mip_runner.pytests/unit/torch/puzzletron/test_setup_bundle.pytests/unit/torch/puzzletron/test_stage_graph.pytests/unit/torch/puzzletron/test_tokenize_data.pytests/unit/torch/puzzletron/test_vllm_axis_contract.pytests/unit/torch/puzzletron/test_width_scenarios.pytests/unit/torch/puzzletron/test_width_slice_equivalence.py
| parser.add_argument("--trust-remote-code", action="store_true") | ||
| parser.add_argument( | ||
| "--allow-aiperf-v011-online-tokenizer-resolution", | ||
| action="store_true", | ||
| ) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Document the security opt-ins.
Add explicit CLI help for both options. State that --trust-remote-code is only for trusted model sources. State that the tokenizer flag permits online resolution. Document the same behavior in run_worker.
As per coding guidelines, “Document public and higher-level APIs with docstrings.” As per path instructions, SECURITY.md requires documentation for security-sensitive exceptions.
Proposed CLI documentation
- parser.add_argument("--trust-remote-code", action="store_true")
+ parser.add_argument(
+ "--trust-remote-code",
+ action="store_true",
+ help="Allow remote model code. Use only with trusted model sources.",
+ )
parser.add_argument(
"--allow-aiperf-v011-online-tokenizer-resolution",
action="store_true",
+ help="Allow AIPerf v0.11 to resolve the tokenizer online.",
)🤖 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/run_profile_aiperf_worker.py` around lines 344 - 348,
Update the argparse definitions near trust-remote-code and
allow-aiperf-v011-online-tokenizer-resolution with explicit help text describing
the security opt-ins: trust remote code only for trusted model sources, and the
tokenizer option enabling online resolution. Add matching documentation to
run_worker and the required SECURITY.md guidance for these security-sensitive
exceptions.
Sources: Coding guidelines, Path instructions
| # Commands are compiled argv lists and never use a shell. | ||
| import subprocess # nosec B404 |
There was a problem hiding this comment.
🔒 Security & Privacy | 🟡 Minor | ⚡ Quick win
Remove the inline # nosec suppressions or obtain the required codeowner approval and justification. The changed subprocess and exec calls use fixed argument vectors, so the suppressions should not remain as unreviewed Bandit bypasses. This applies to the subprocess import and aggregation call in modelopt/torch/puzzletron/orchestration/adapters/sharded.py, the os.execvpe call in modelopt/torch/puzzletron/orchestration/task_launcher.py, and the corresponding subprocess suppressions in modelopt/torch/puzzletron/post_mip/runner.py and modelopt/torch/puzzletron/orchestration/adapters/post_mip.py.
📍 Affects 3 files
modelopt/torch/puzzletron/orchestration/adapters/sharded.py#L20-L21(this comment)modelopt/torch/puzzletron/post_mip/runner.py#L29-L31modelopt/torch/puzzletron/orchestration/adapters/post_mip.py#L22-L25
🤖 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/orchestration/adapters/sharded.py` around lines 20
- 21, Address the inline Bandit suppression approvals for the subprocess import
in sharded.py (lines 20-21), the aggregation subprocess.run call in sharded.py
(line 389), and the os.execvpe call in task_launcher.py (lines 275-276): either
document an explicit security justification for each # nosec marker in the PR
description with required codeowner approval, or remove the inline markers and
configure the approved skips centrally in Bandit configuration.
Apply the same fix in `@modelopt/torch/puzzletron/post_mip/runner.py` around lines
29 - 31: Same prohibited subprocess suppression.
Apply the same fix in
`@modelopt/torch/puzzletron/orchestration/adapters/post_mip.py` around lines 22 -
25: Same prohibited subprocess-call suppression.
Source: Coding guidelines
| "modality": data["modality"], | ||
| "layout": data["layout"], | ||
| "max_sample_length": sequence_length, | ||
| "sequence_length": sequence_length, |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Inspect the dataset config type that consumes the rendered `data` mapping.
set -euo pipefail
fd -t f 'config.py' modelopt/torch/puzzletron/dataset --exec cat -n {}
# Find every consumer that reads data.sequence_length or constructs the dataset config.
rg -nP -C4 '\bsequence_length\b' --type=py modelopt/torch/puzzletron | head -80Repository: NVIDIA/Model-Optimizer
Length of output: 2064
🏁 Script executed:
set -euo pipefail
for path in \
puzzletron_setup/bundle.py \
modelopt/torch/puzzletron/dataset/config.py
do
echo "===== $path ====="
curl -fsSL "https://raw.githubusercontent.com/NVIDIA/Model-Optimizer/main/$path" \
| nl -ba | sed -n '1,560p'
done
echo "===== sequence_length references ====="
curl -fsSL "https://api.github.com/repos/NVIDIA/Model-Optimizer/contents/modelopt/torch/puzzletron/dataset" \
| jq -r '.[].download_url' \
| while read -r url; do
curl -fsSL "$url"
done \
| nl -ba \
| rg -n -C4 '\bsequence_length\b'Repository: NVIDIA/Model-Optimizer
Length of output: 2035
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer "class" "sequence_length" "puzzletron" dataset config.py
💡 Result:
In the context of the NVIDIA Model Optimizer Puzzletron framework, the sequence length is typically managed via a parameter named block_size when configuring dataset validation or preprocessing [1]. While Puzzletron configuration files (such as those for memory pruning, e.g., llama-3_1-8B_pruneffn_memory.yaml) primarily focus on parameters like puzzle_dir, input_hf_model_path, dataset_path, intermediate_size_list, and target_memory, the underlying validation tools define block_size as the maximum sequence length used for tokenization [1][2][3]. Additionally, when preparing datasets for Megatron-based workflows in the same repository, users specify the sequence length using the --max_sequence_length argument (commonly set to 256,000 for specific post-training datasets) [4]. To verify or adjust these settings, check: 1. The validation scripts (e.g., validate_puzzle_with_multi_replacements.py), where block_size is explicitly documented [1]. 2. Any relevant YAML configuration files within the examples/puzzletron/configs/ directory, which govern the specific pruning or memory optimization experiments [2][5][3].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/modelopt/torch/puzzletron/tools/validate_puzzle_with_multi_replacements.py
- 2: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/examples/puzzletron/README.md
- 4: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/dataset/MEGATRON_DATA_PREP.md
- 5: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml
🌐 Web query:
https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt/torch/puzzletron/dataset/config.py
💡 Result:
Puzzletron is an experimental model compression and Neural Architecture Search (NAS) algorithm integrated into the NVIDIA Model Optimizer framework [1][2]. It is based on the Puzzle research paper, which aims to produce smaller, more efficient, and often heterogeneous versions of large language models (LLMs) [3][2][4]. Key technical aspects of Puzzletron include: Functionality: It optimizes model architectures by searching for the most effective modifications to MLP (Multi-Layer Perceptron) and attention layers [4]. It uses Mixed-Integer Programming (MIP) to identify combinations of layer modifications that satisfy specific constraints, such as target memory usage or parameter counts, while aiming to maximize model accuracy [3][4]. Workflow: The algorithm typically involves several stages, including scoring pruning activations, performing a MIP-based NAS search, and optionally conducting knowledge distillation [2][4][5]. It is designed to handle heterogeneous architectures, meaning it can prune different parts of the model to different extents to meet global constraints [6][4]. Integration: Puzzletron is accessible within the NVIDIA Model Optimizer repository under the modelopt.torch.puzzletron package [1]. Official documentation and tutorials are provided within the examples/puzzletron directory of the repository [3][4]. It is recommended to use the algorithm within an NVIDIA NeMo container environment, although users may need to perform an editable installation of the Model Optimizer from source because standard container versions may not include the necessary Puzzletron components [7][4]. The specific file you referenced, modelopt/torch/puzzletron/dataset/config.py, is part of the internal configuration logic for the Puzzletron NAS plugin, which facilitates the conversion of models from Hugging Face format to the heterogeneous formats used by the optimization process [2].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/tree/main/modelopt/torch/puzzletron
- 2: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.puzzletron_nas_plugin.html
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
- 4: https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/puzzletron
- 5: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/examples/puzzletron/README.md
- 6: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/pruning/puzzletron/README.md
- 7: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/examples/puzzletron/README.md
🌐 Web query:
https://github.com/NVIDIA/Model-Optimizer/blob/main/puzzletron_setup/bundle.py
💡 Result:
The file path provided, https://github.com/NVIDIA/Model-Optimizer/blob/main/puzzletron_setup/bundle.py, does not appear to exist in the main branch of the NVIDIA Model Optimizer repository as of August 2026 [1][2]. Based on the official documentation for the Puzzletron project within the NVIDIA Model Optimizer repository, the relevant script for setting up and managing Puzzletron experiments is typically named puzzletron_setup.py, which is located in the examples/puzzletron directory [3][1]. Puzzletron is an experimental model compression algorithm in the Model Optimizer library based on Local Neural Architecture Search (NAS) and the Puzzle paper [2][4]. The setup wizard, executed via puzzletron_setup.py, is designed to inspect model configurations and generate self-contained smoke and production experiment, runner, and execution bundles [3]. These bundles encapsulate the necessary components for running the Puzzletron pipeline, which includes converting models from Hugging Face format to Puzzletron-compatible formats, performing NAS-based pruning, and validating checkpoints [1][4]. If you are encountering issues with Puzzletron setup, it is recommended to review the README.md in the examples/puzzletron directory for current usage instructions [1][5], or check the repository's issue tracker for known setup challenges [5].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
- 2: https://github.com/NVIDIA/Model-Optimizer/tree/main/modelopt/torch/puzzletron
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/examples/puzzletron/README.md
- 4: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.puzzletron_nas_plugin.html
- 5: Puzzletron README initial setup fails with number of issues #1637
🌐 Web query:
https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py
💡 Result:
The file modelopt/torch/puzzletron/dataset/config.py in the puzzletron_v2 branch of the NVIDIA Model-Optimizer repository is part of Puzzletron, an experimental model compression framework [1][2]. Puzzletron is a tool designed to find smaller, faster variants of pre-trained large language models [2]. It utilizes a combination of techniques, including local neural architecture search, physical slicing, mixed-integer programming (MIP), and optional knowledge distillation [3][2]. The framework aims to optimize models to satisfy specific constraints, such as target parameter counts or memory usage requirements [3]. The dataset module within the puzzletron directory typically handles data configuration for these compression workflows [4]. In the context of the broader Puzzletron project, such configuration files define the parameters for preparing and processing datasets required during the model search and evaluation phases [2]. For the most accurate and up-to-date implementation details of config.py, it is recommended to review the file directly within the puzzletron_v2 branch or compare it against the main branch of the repository, as Puzzletron is under active development [1][2].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/tree/main/modelopt/torch/puzzletron
- 2: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/examples/puzzletron/README.md
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
- 4: https://github.com/NVIDIA/Model-Optimizer/blob/main/modelopt/torch/puzzletron/__init__.py
🌐 Web query:
https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/puzzletron_setup/bundle.py
💡 Result:
The file at https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/puzzletron_setup/bundle.py is part of the Puzzletron experimental branch within the NVIDIA Model Optimizer repository [1][2]. Puzzletron is an advanced model compression algorithm designed for heterogeneous pruning and Neural Architecture Search (NAS) of Large Language Models (LLMs) and Vision Language Models (VLMs) [3][4][5]. It utilizes Mixed Integer Programming (MIP) to identify optimal model variants by searching across dimensions such as FFN (Feed-Forward Network) intermediate sizes per layer, attention layer removal, and GQA (Grouped Query Attention) head configurations, all to meet specific constraints like target memory, latency, or parameter counts [4][6]. The specific file path puzzletron_setup/bundle.py appears to be associated with the setup or bundling process for this experimental functionality [1]. In the context of the broader Model Optimizer toolkit, Puzzletron is intended to produce heterogeneous architectures that maintain accuracy while significantly reducing model size or serving costs [1][4][6]. Users are generally directed to documentation in examples/puzzletron/README.md for implementation details, including dataset preparation and configuration, though it has been noted in developer discussions that the experimental setup process can encounter dependency and configuration challenges [7][8][9].
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/examples/puzzletron/README.md
- 2: https://github.com/NVIDIA/Model-Optimizer/tree/main/modelopt/torch/puzzletron
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst
- 4: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/pruning/README.md
- 5: https://github.com/NVIDIA/Model-Optimizer
- 6: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.puzzletron_nas_plugin.html
- 7: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
- 8: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/examples/puzzletron/README.md
- 9: Puzzletron README initial setup fails with number of issues #1637
🌐 Web query:
"sequence_length" "pack_size" "max_sample_length" "modelopt/torch/puzzletron/dataset/config.py"
💡 Result:
In the context of the NVIDIA Model Optimizer library, specifically within the Puzzletron module (modelopt/torch/puzzletron), the parameters sequence_length, pack_size, and max_sample_length are commonly used to configure data handling for Large Language Model (LLM) training, pruning, and optimization tasks [1][2][3]. While these parameters are often used together in NVIDIA's broader ecosystem (such as in NeMo and Megatron-Bridge) to manage data efficiency [4][5][6], they serve distinct roles: 1. max_sample_length: Generally refers to the maximum length to which individual sequences are truncated before any packing occurs [7][5]. It defines the upper bound for the length of a single input sample [7]. 2. sequence_length: In many configurations, this refers to the target sequence length for the model input [5][6]. When sequence packing is used, the individual samples are truncated to a specific length, and then concatenated into a larger "packed" sequence [5]. 3. pack_size: Refers to the total length of the packed sequence [4][5]. By concatenating multiple shorter samples into a single sequence of pack_size, training efficiency can be improved by filling GPU memory more effectively without the wasted computation often associated with padding [5]. In the Model Optimizer (ModelOpt) framework, these configurations are typically managed within dataset preparation or NAS (Neural Architecture Search) plugins [8][9][10]. Users often specify these in YAML configuration files (e.g., for pruneffn tasks) to align data dimensions with model architecture requirements and available GPU memory [1][2]. For precise implementation details, refer to the following sources in the Model Optimizer repository: - Dataset preparation utilities: modelopt.torch.puzzletron.dataset [8] - Puzzletron NAS configuration: modelopt.torch.puzzletron.puzzletron_nas_plugin [9] - General dataset utilities: modelopt/torch/utils/dataset_utils.py [7]
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
- 2: https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/puzzletron
- 3: https://github.com/NVIDIA/Model-Optimizer/tree/main/modelopt/torch/puzzletron
- 4: https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.packed_sequence.html
- 5: https://docs.nvidia.com/nemo-framework/user-guide/24.09/nemotoolkit/features/optimizations/sequence_packing.html
- 6: https://docs.nvidia.com/nemo/megatron-bridge/latest/skills/nemo-mbridge-perf-sequence-packing/SKILL.html
- 7: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/modelopt/torch/utils/dataset_utils.py
- 8: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.dataset.html
- 9: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.puzzletron_nas_plugin.html
- 10: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.html
🌐 Web query:
"sequence_length" "packed_varlen" "pack_size" Puzzletron
💡 Result:
The terms mentioned refer to different concepts in high-performance machine learning workflows, specifically within the NVIDIA ecosystem (Megatron, NeMo, and Model Optimizer). There is no single "Puzzletron" object that combines sequence_length, packed_varlen, and pack_size; rather, these terms are associated with distinct, though sometimes overlapping, optimization techniques. 1. Sequence Packing (packed_varlen/pack_size): This technique is used to improve training efficiency by concatenating multiple variable-length sequences into a single "packed" sequence, significantly reducing padding waste [1][2][3]. - packed_sequence_size (or pack_size): This parameter defines the maximum size (length) of the packed sequences [4][1][5]. - packed_varlen (variable length packing): This refers to the underlying processing logic that handles sequences of differing lengths by passing cumulative sequence length metadata (often called cu_seqlens) to attention kernels (like FlashAttention) instead of relying on inefficient custom attention masks [2][5][3][6]. 2. Puzzletron: Puzzletron is a tool within the NVIDIA Model Optimizer suite used for compressing large language models [7][8]. It implements a Mixed-Integer Programming (MIP) based Neural Architecture Search (NAS) algorithm to prune and modify model architectures (such as FFN intermediate sizes and attention layers) to meet specific memory and accuracy targets [7]. It is unrelated to the sequence packing parameters used in training data pipelines [9][7]. In summary, sequence_length, pack_size, and concepts related to packed_varlen belong to the data loading and training pipeline (e.g., Megatron Bridge, NeMo), while Puzzletron is a separate model compression/pruning framework [4][1][7][5].
Citations:
- 1: https://docs.nvidia.com/nemo/megatron-bridge/latest/training/packed-sequences.html
- 2: https://docs.nvidia.com/nemo/rl/latest/design-docs/sequence-packing-and-dynamic-batching.html
- 3: https://docs.nvidia.com/nemo-framework/user-guide/25.07/sft_peft/packed_sequence.html
- 4: https://docs.nvidia.com/nemo/megatron-bridge/latest/apidocs/bridge/bridge.data.datasets.packed_sequence.html
- 5: https://docs.nvidia.com/nemo/megatron-bridge/0.3.1/training/packed-sequences.html
- 6: https://github.com/NVIDIA/Megatron-LM/blob/40b1fd37/megatron/rl/sequence_packing_utils.py
- 7: https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/puzzletron
- 8: https://nvidia.github.io/Model-Optimizer/reference/generated/modelopt.torch.puzzletron.html
- 9: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/puzzletron/README.md
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py "def sequence_length"
💡 Result:
Found it in config.py at lines 520–524:
`@property`
def sequence_length(self) -> int:
return self.packing.pack_size if self.packing is not None else self.max_sample_length[1] (github.com)
Citations:
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer "max_sample_length" "packed_varlen"
💡 Result:
In the context of the NVIDIA Model Optimizer library, max_sample_length and packing (often referred to in the context of packed_varlen or pack=True) are parameters used to control how calibration data is prepared for model optimization, such as quantization or pruning [1][2]. max_sample_length This parameter defines the maximum number of tokens for a single sequence or sample during data preparation [3][4][5]. It ensures that calibration inputs adhere to specific length constraints, which helps maintain predictable calibration behavior, reduces peak memory usage, and ensures consistency during preprocessing [5]. Packed Variable Length Data (pack=True) The library supports a pack=True mode in its dataset utilities (e.g., get_dataset_dataloader), which implements Megatron-LM-style global-stream document packing [1][2]. - When enabled, raw samples are concatenated into a single continuous token stream, separated by end-of-sequence (EOS) tokens [1][2]. - This stream is then sliced into uniform sequences of length max_sample_length [1][2]. - This approach is designed to match the distribution seen by models during pretraining, which can lead to more representative activation behavior during the calibration phase [1]. These settings are typically configured within the dataloader setup scripts (e.g., hf_ptq.py) to control how data is ingested and processed before it is fed through the model for calibration [6].
Citations:
- 1: 5d0441a
- 2: https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst
- 3: https://github.com/NVIDIA/Model-Optimizer/blob/f34f488a/modelopt/torch/utils/dataset_utils.py
- 4: https://github.com/NVIDIA/Model-Optimizer/blob/7f5fd650/modelopt/torch/utils/dataset_utils.py
- 5: 8663678
- 6: https://github.com/NVIDIA/Model-Optimizer/blob/main/examples/llm_ptq/hf_ptq.py
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py "class DatasetConfig"
💡 Result:
The file exists, but it does not contain class DatasetConfig. It defines:
PackingSpecPuzzletronDataSpec
The relevant class is PuzzletronDataSpec at approximately line 422. [1]
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset "DatasetConfig(" "sequence_length"
💡 Result:
No matching DatasetConfig("sequence_length" occurrence was found in the specified GitHub search. The puzzletron_v2 dataset module instead exposes process_and_save_dataset(...) in prepare_dataset.py; it has no DatasetConfig or sequence_length field. [1] (raw.githubusercontent.com)
Citations:
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py "class PuzzletronDataSpec" "model_config"
💡 Result:
The referenced PuzzletronDataSpec class does not contain a model_config field. Its fields are:
modalitylayoutmax_sample_lengthpacking
The class is defined at approximately lines 420–530 in the puzzletron_v2 branch. (github.com)
Citations:
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py "PuzzletronDataSpec(" "sequence_length"
💡 Result:
PuzzletronDataSpec does not define a sequence_length field. It exposes sequence_length as a property:
`@property`
def sequence_length(self) -> int:
return self.packing.pack_size if self.packing is not None else self.max_sample_lengthThus, it returns packing.pack_size when packing is configured; otherwise, it returns max_sample_length. [1]
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py "PuzzletronDataSpec" "BaseModel"
💡 Result:
The file is available here:
modelopt/torch/puzzletron/dataset/config.py on the puzzletron_v2 branch. (github.com)
Key details:
PuzzletronDataSpecis a frozen dataclass.- It contains:
modality: Modalitylayout: DataLayoutmax_sample_length: intpacking: PackingSpec | None
- It does not inherit from
BaseModel; it uses Python’sdataclasses.dataclass. (raw.githubusercontent.com)
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py
- 2: https://raw.githubusercontent.com/NVIDIA/Model-Optimizer/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/config.py "PuzzletronDataSpec" "dataclass"
💡 Result:
Found it: NVIDIA/Model-Optimizer config.py.
PuzzletronDataSpec is defined as a frozen dataclass with:
`@dataclass`(frozen=True)
class PuzzletronDataSpec:
modality: Modality
layout: DataLayout
max_sample_length: int
packing: PackingSpec | None = NoneIt appears at lines 420–430. (github.com)
Citations:
🌐 Web query:
site:github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/prepare_dataset.py "PuzzletronDataSpec"
💡 Result:
The file exists at:
modelopt/torch/puzzletron/dataset/prepare_dataset.py on the puzzletron_v2 branch. (github.com)
However, it does not contain PuzzletronDataSpec; GitHub reports no matching text. (github.com)
The file defines process_and_save_dataset() and PREBUILT_KD_DATASET.
Citations:
- 1: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/prepare_dataset.py
- 2: https://github.com/NVIDIA/Model-Optimizer/blob/puzzletron_v2/modelopt/torch/puzzletron/dataset/prepare_dataset.py
Remove sequence_length from the rendered data mapping.
PuzzletronDataSpec is a frozen dataclass with no sequence_length input field. It exposes sequence_length only as a derived property. Passing this key to its constructor raises an unexpected-keyword error.
🤖 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/bundle.py` at line 470, Remove the "sequence_length" entry
from the rendered data mapping used to construct PuzzletronDataSpec, leaving
sequence_length available only through its derived property and preserving the
remaining constructor inputs.
Replace the legacy model matrix with one hermetic current-route campaign, and fix the configuration and orchestration contracts it exposes. Keep the dedicated GPU target separate from generic GPU coverage. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Keep controller-compatible authored configuration separate from normalized worker settings so equivalent stage manifests remain resumable. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Use the authored configuration for semantic compatibility checks while preserving normalized worker settings in execution records and synthesized post-MIP candidates. This keeps worker self-validation aligned with controller resume checks. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
The current-route baseline exposed generic rendezvous, finalization, provenance, and AIPerf integration defects. Preserve compiled execution identity across workers, recover aggregation publication failures, verify exact source dependencies, and keep local tokenizer loading compatible with the pinned AIPerf environment. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Replacement scoring artifacts are nested beneath both width and depth scenario directories. Include the depth directory when validating the distributed evaluation results. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Apply the repository formatter and full license headers across the affected files. Normalize parameterized test values and document the fused-kernel compatibility names required by the test double. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Accept either configured pruned FFN width in the selected architecture. The 90 percent parameter constraint does not guarantee that the optimizer chooses the most aggressive candidate. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Disable KV caching only for the synthetic all-full-attention Qwen 3.5 forward. This preserves the CUDA reload check without relying on a linear-attention cache layer that the hermetic fixture intentionally omits. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Share the hermetic tiny-Qwen campaign fixture and validate the selected KD checkpoint, post-MIP report, and no-op resume through the public orchestrator. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Avoid fingerprinting the canonical stage pointer as one of its own immutable outputs, and keep test manifests aligned with the execution-record schema. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Keep shard workers and aggregation on the same campaign configuration so they derive one execution identity and consume the same results. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Refresh AnyModel interchange metadata before publishing a consolidated checkpoint, and reload selected heterogeneous artifacts through the descriptor-aware ModelOpt path. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Bind post-MIP attempts to the producer contract and preserve explicit AIPerf security policies so resume resubmits stale work without weakening offline execution. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Load post-MIP execution identity helpers only when dynamic-stage currentness or planning needs them, keeping normal Puzzletron package initialization acyclic. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Signed-off-by: Johannes Rausch <jrausch@nvidia.com> (cherry picked from commit 9bf94f22640b155d0f71a78a8afca9b255e582e4)
Signed-off-by: Johannes Rausch <jrausch@nvidia.com> (cherry picked from commit 6065146820736e918069019b8b41945926d889c1)
Reject truthy string and numeric values across AIPerf execution routes, validate checkpoint collections, and keep the cold-import regression independent of subprocess coverage. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Keep the new configuration-boundary tests consistent with the repository import convention without changing their behavior. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Reject ambiguous security-policy values and give resumed controllers a fresh artifact-settling window before recording terminal failure. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Let automation accept resolved setup defaults through the public CLI so integration tests and unattended campaigns share the same bundle-generation path. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
a34b939 to
57c87e1
Compare
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: 3
🤖 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/README.md`:
- Around line 88-90: Update the automation setup description in the README to
state that the defaults file must provide every required value without a
resolved default, rather than only values lacking a built-in default. Align the
wording with the layered resolution behavior implemented by
puzzletron_setup/v2/defaults.py.
In `@puzzletron_setup/v2/prompts.py`:
- Around line 209-213: Update NonInteractiveBackend.text to distinguish an
absent required default from an explicitly provided empty string, using a
sentinel or explicit requiredness parameter, while still rejecting missing
required defaults. Ensure wizard.py accepts the empty string produced for an
empty prerun_commands list, and add a regression test covering --full
--non-interactive setup with an empty prerun_commands default.
In `@tests/_test_utils/torch/puzzletron/tiny_qwen_campaign.py`:
- Around line 246-262: Add a bounded timeout to the subprocess.run call in the
setup flow, catch subprocess.TimeoutExpired, and raise an assertion containing
any available stdout and stderr so setup failures are reported before the outer
test timeout.
🪄 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: 974eb0f7-ad6c-45ef-8556-3e20dadb1e47
📒 Files selected for processing (16)
examples/puzzletron/README.mdexamples/puzzletron/run_profile_aiperf_worker.pymodelopt/torch/puzzletron/distillation/global_kd_recipe.pymodelopt/torch/puzzletron/orchestration/adapters/post_mip.pymodelopt/torch/puzzletron/orchestration/adapters/sharded.pymodelopt/torch/puzzletron/orchestration/controller.pymodelopt/torch/puzzletron/post_mip/runner.pypuzzletron_setup/v2/cli.pypuzzletron_setup/v2/prompts.pypuzzletron_setup/v2/wizard.pytests/_test_utils/torch/puzzletron/tiny_qwen_campaign.pytests/unit/torch/puzzletron/test_global_kd_canonical.pytests/unit/torch/puzzletron/test_orchestration_executors.pytests/unit/torch/puzzletron/test_orchestration_shutdown_progress.pytests/unit/torch/puzzletron/test_post_mip_runner.pytests/unit/torch/puzzletron/test_setup_v2_quick.py
🚧 Files skipped from review as they are similar to previous changes (8)
- tests/unit/torch/puzzletron/test_orchestration_executors.py
- modelopt/torch/puzzletron/orchestration/adapters/sharded.py
- modelopt/torch/puzzletron/orchestration/adapters/post_mip.py
- modelopt/torch/puzzletron/orchestration/controller.py
- examples/puzzletron/run_profile_aiperf_worker.py
- modelopt/torch/puzzletron/distillation/global_kd_recipe.py
- tests/unit/torch/puzzletron/test_orchestration_shutdown_progress.py
- modelopt/torch/puzzletron/post_mip/runner.py
| Automation can use the same setup entry point without answering prompts. The | ||
| defaults file must provide every required value that has no built-in default: | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Describe resolved defaults accurately.
puzzletron_setup/v2/defaults.py resolves values from built-in, model-derived, preset, model-profile, and defaults-file layers. A required value can therefore lack a built-in default but still have a resolved default. Replace “no built-in default” with “no resolved default” or list the supported default layers.
Suggested wording
-Automation can use the same setup entry point without answering prompts. The
-defaults file must provide every required value that has no built-in default:
+Automation can use the same setup entry point without answering prompts. The
+defaults file must provide every required value that has no resolved default:📝 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.
| Automation can use the same setup entry point without answering prompts. The | |
| defaults file must provide every required value that has no built-in default: | |
| Automation can use the same setup entry point without answering prompts. The | |
| defaults file must provide every required value that has no resolved default: |
🤖 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/README.md` around lines 88 - 90, Update the automation
setup description in the README to state that the defaults file must provide
every required value without a resolved default, rather than only values lacking
a built-in default. Align the wording with the layered resolution behavior
implemented by puzzletron_setup/v2/defaults.py.
Preserve explicit empty defaults while failing closed on missing required values, and bound setup subprocess execution so failures surface within the test timeout. Signed-off-by: Johannes Rausch <jrausch@nvidia.com>
1 participant
What does this PR do?
Type of change: New feature, new tests, and bug fixes.
Problem / situation
Puzzletron v2 has focused unit coverage, but it does not have one small GPU baseline that proves the current public route from setup through final model selection and resume. The missing integration coverage allowed configuration identity, aggregation, artifact publication, and checkpoint portability defects to remain invisible until a real campaign reached those boundaries.
Automation also had no supported way to run setup without answering prompts. The test therefore depended on prompt labels and ordering, which made it unnecessarily sensitive to setup UI refactors.
What changed
This change adds one reusable tiny Qwen 3.5 campaign fixture and one focused GPU end-to-end test. The test covers setup-v2 bundle generation, compilation, automated public orchestration, activation and replacement scoring, MIP solving, evaluation, AIPerf serving measurements, short global distillation, final evaluation and selection, descriptor-aware reload of the selected heterogeneous checkpoint, report generation, and no-op resume.
Setup v2 now provides a public non-interactive mode. It accepts an explicit defaults file, campaign directory, and setup profile, then runs the same wizard sections, resolved configuration, bundle rendering, and validation used by interactive setup. It fails instead of guessing when a required answer has no resolved default. The shared E2E fixture invokes this CLI and the public orchestrator CLI, so setup and pipeline refactors are exercised through production entry points rather than duplicated test configuration.
The campaign exposed generic current-route defects rather than test-only exceptions. The corresponding fixes:
Focused CPU tests cover the same setup, configuration, graph, controller, manifest, topology, aggregation, checkpoint, security-policy, and resume contracts. The dedicated sessions also verify the installed lmms-eval and NeMo AutoModel repositories and pinned commits before collection.
Existing interactive setup and campaign configurations with boolean security-policy values and list-valued checkpoint overrides remain accepted. The non-interactive CLI options are additive. Security-sensitive AIPerf settings reject strings and numbers rather than interpreting them by truthiness. AIPerf no longer enables remote code implicitly; callers that need it must opt in. Persisted attempts or post-MIP publications that lack or mismatch the current execution identity are rerun rather than reused.
How it works
The shared fixture creates its model, tokenizer, and data locally and exposes exactly one CUDA device. It invokes the public setup CLI and orchestrator CLI, then validates canonical stage manifests, post-MIP candidate identity and lineage, real AIPerf artifacts, distillation records, final checkpoint geometry, report verification, and immutable artifacts across resume.
The dedicated one-GPU Puzzletron Nox test session (
gpu_puzzletron) verifies the prepared runtime against the repository's pinned environment contract, requires exactly one visible CUDA 12.9 GPU, and runs only the focused lifecycle test. A Nox session is a named repository automation target defined innoxfile.py; it is not itself a GitHub workflow.Why this approach
A single small current-route campaign keeps the expensive integration surface reviewable while focused CPU tests retain detailed failure localization. Reusing the public setup and orchestration entry points makes configuration and pipeline changes flow through the same contracts used by real campaigns. Follow-on tests can reuse one fixture instead of copying model, dataset, YAML, or orchestration setup.
Usage
Run setup without prompts by supplying a complete defaults file and campaign destination:
Testing
The required Puzzletron CPU workflow covers the setup CLI, resolved configuration, graph, controller, manifest, topology, aggregation, checkpoint, security-policy, and resume contracts. The changed-file quality workflow checks formatting, license headers, security findings, and cache-independent mypy coverage for the affected Python files.
Manual validation in the pinned CUDA 12.9 environment covers the complete one-GPU lifecycle, including real AIPerf measurements, two distillation optimizer steps, final selected-checkpoint reload and CUDA forward, verified report generation, and no-op resume.
The checked GitHub workflow in this PR does not invoke the
gpu_puzzletronNox session. The stacked Puzzletron GPU CI workflow PR proposes the repository-owned execution image and required-check wiring separately.Summary by CodeRabbit
New Features
Bug Fixes
Tests