Experimental claude skill for puzzletron algoritgm#1769
Experimental claude skill for puzzletron algoritgm#1769danielkorzekwa wants to merge 13 commits into
Conversation
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
Signed-off-by: Daniel Korzekwa <dkorzekwa@nvidia.com>
📝 WalkthroughWalkthroughAdds an experimental Claude Code agent skill for Puzzletron under ChangesPuzzletron Agent Skill
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1769 +/- ##
==========================================
+ Coverage 58.45% 64.78% +6.32%
==========================================
Files 510 511 +1
Lines 56271 56792 +521
==========================================
+ Hits 32896 36791 +3895
+ Misses 23375 20001 -3374
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:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Warning
CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.
Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.
Actionable comments posted: 4
🤖 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 @.agents/skills/puzzletron/all_progress.py:
- Around line 80-84: The variables `cur_b` and `total_b` are only defined inside
the elif block when `batch_matches` is truthy, but they are used later in the
code (around line 100) regardless of which conditional branch executes. When the
if condition on line 80 evaluates to true (sol_done is not None and sol_total is
truthy), the elif block is skipped entirely, leaving `cur_b` and `total_b`
undefined. Extract the batch data unpacking logic (extracting pct, cur_b, and
total_b from batch_matches[-1]) before the if-elif conditional block to ensure
these variables are always defined when batch_matches is non-empty, preventing
NameError when they are referenced later in the code.
In @.agents/skills/puzzletron/mip_progress.py:
- Around line 53-59: Replace the hardcoded source line number markers with
content-based semantic markers to make detection robust to code refactoring. In
the completion detection block around line 57, replace the condition checking
for "sweep.py:292" with a check for "Results written to:" which is the actual
completion message. In the related detection block around lines 109-114 that
currently guards on "sweep.py:258", remove the line number check entirely and
instead use unconditional regex matching on the "compression_rate=" pattern
which is already a proven approach used at line 99 for results detection.
In @.agents/skills/puzzletron/SKILL.md:
- Around line 34-40: The specification lacks numeric validation for the
nproc_per_node parameter before it is interpolated into shell commands, creating
a security vulnerability for shell injection attacks. Add an explicit validation
rule to both the "all" and "local" command sections in the skill specification
that checks whether nproc_per_node matches the pattern of a positive integer
(^[0-9]+$). Insert this validation check after the "value not found" check and
before the "Otherwise use the parsed value" instruction in both sections. If the
value is not strictly numeric, the specification should instruct to ask the user
"nproc_per_node must be a positive integer." and STOP before any shell command
execution occurs.
- Around line 46-53: The shell pipeline using torchrun piped to tee piped to
grep does not properly propagate exit codes because without pipefail, the
pipeline only returns the exit code of the rightmost command (grep). When
torchrun fails but grep successfully finds the "Puzzletron Progress" pattern,
the pipeline reports success even though the actual torchrun command failed. To
fix this, add set -o pipefail before or at the beginning of the script block
containing the torchrun command to ensure that the pipeline returns a non-zero
exit code when any command in the pipeline fails, allowing accurate exit code
reporting as mentioned in the instructions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f507f804-2357-44dd-934e-633f88d0cd06
📒 Files selected for processing (7)
.agents/skills/puzzletron/README.md.agents/skills/puzzletron/SKILL.md.agents/skills/puzzletron/all_progress.py.agents/skills/puzzletron/mip_progress.py.claude/skills/puzzletronCHANGELOG.rstexamples/puzzletron/README.md
| if sol_done is not None and sol_total: | ||
| cur_detail = f" ({sol_done}/{sol_total} solutions)" | ||
| elif batch_matches: | ||
| pct, cur_b, total_b = batch_matches[-1] | ||
| cur_detail = f" ({cur_b}/{total_b} batches)" |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
find . -type f -name "all_progress.py" | head -20Repository: NVIDIA/Model-Optimizer
Length of output: 110
🏁 Script executed:
cat -n "./.agents/skills/puzzletron/all_progress.py"Repository: NVIDIA/Model-Optimizer
Length of output: 7918
cur_b/total_b can be undefined and crash at line 100.
Lines 83 defines these variables only when the elif batch_matches: condition (line 82) executes. However, when sol_done = 0 and sol_total is non-zero, the condition at line 80 (if sol_done is not None and sol_total:) evaluates to True, causing lines 82-84 to be skipped. Later, line 100 attempts to use cur_b and total_b in the condition, raising NameError when batch_matches is truthy.
Extract the batch data early as suggested:
Suggested fix
-batch_matches = re.findall(r"calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)", text)
+batch_matches = re.findall(r"calculate_losses_pipeline[^:]*:\s*(\d+)%.*?(\d+)/(\d+)", text)
+latest_batch = batch_matches[-1] if batch_matches else None
@@
-elif batch_matches:
- pct, cur_b, total_b = batch_matches[-1]
+elif latest_batch:
+ pct, cur_b, total_b = latest_batch
cur_detail = f" ({cur_b}/{total_b} batches)"
@@
-elif batch_matches and int(cur_b) > 0 and int(cur_b) < int(total_b):
- rate_per_batch = cur_step_elapsed / int(cur_b)
- step_remaining = rate_per_batch * (int(total_b) - int(cur_b))
+elif latest_batch:
+ _, cur_b, total_b = latest_batch
+ if int(cur_b) > 0 and int(cur_b) < int(total_b):
+ rate_per_batch = cur_step_elapsed / int(cur_b)
+ step_remaining = rate_per_batch * (int(total_b) - int(cur_b))🤖 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 @.agents/skills/puzzletron/all_progress.py around lines 80 - 84, The
variables `cur_b` and `total_b` are only defined inside the elif block when
`batch_matches` is truthy, but they are used later in the code (around line 100)
regardless of which conditional branch executes. When the if condition on line
80 evaluates to true (sol_done is not None and sol_total is truthy), the elif
block is skipped entirely, leaving `cur_b` and `total_b` undefined. Extract the
batch data unpacking logic (extracting pct, cur_b, and total_b from
batch_matches[-1]) before the if-elif conditional block to ensure these
variables are always defined when batch_matches is non-empty, preventing
NameError when they are referenced later in the code.
| # Detect completion via step 8 marker or sweep.py:292 | ||
| complete_ts = None | ||
| for line in lines: | ||
| ts = get_ts(line) | ||
| if ts and ("sweep.py:292" in line or "Puzzletron Progress 8/8" in line): | ||
| complete_ts = ts | ||
| break |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
find . -name "mip_progress.py" -type fRepository: NVIDIA/Model-Optimizer
Length of output: 110
🏁 Script executed:
cat -n ".agents/skills/puzzletron/mip_progress.py" | sed -n '50,120p'Repository: NVIDIA/Model-Optimizer
Length of output: 3163
🏁 Script executed:
find . -name "sweep.py" -type f | head -5Repository: NVIDIA/Model-Optimizer
Length of output: 107
🏁 Script executed:
cat -n "./modelopt/torch/puzzletron/mip/sweep.py" | sed -n '255,265p'Repository: NVIDIA/Model-Optimizer
Length of output: 607
🏁 Script executed:
cat -n "./modelopt/torch/puzzletron/mip/sweep.py" | sed -n '290,300p'Repository: NVIDIA/Model-Optimizer
Length of output: 390
🏁 Script executed:
rg "Results written to:" "./modelopt/torch/puzzletron/mip/" -A 2 -B 2Repository: NVIDIA/Model-Optimizer
Length of output: 357
🏁 Script executed:
rg "sweep.py:292" "./.agents/skills/puzzletron/" -A 3 -B 3Repository: NVIDIA/Model-Optimizer
Length of output: 939
Replace hardcoded source line number markers with semantic markers to improve robustness.
Lines 57 and 109 rely on exact source line numbers (sweep.py:292, sweep.py:258) embedded in logging output. These become stale when sweep.py is refactored, causing false [RUNNING] states even when runs complete.
Replace with stable, content-based detection:
- Line 57: Replace
"sweep.py:292"with"Results written to:"(the actual completion message) - Lines 109–114: Replace the
"sweep.py:258"guard with unconditional regex matching oncompression_rate=pattern (no line dependency needed)
This approach is already partially in use (line 99 detects results via regex on "Results written to:"), so the pattern is established and proven.
🤖 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 @.agents/skills/puzzletron/mip_progress.py around lines 53 - 59, Replace the
hardcoded source line number markers with content-based semantic markers to make
detection robust to code refactoring. In the completion detection block around
line 57, replace the condition checking for "sweep.py:292" with a check for
"Results written to:" which is the actual completion message. In the related
detection block around lines 109-114 that currently guards on "sweep.py:258",
remove the line number check entirely and instead use unconditional regex
matching on the "compression_rate=" pattern which is already a proven approach
used at line 99 for results detection.
| Parse `nproc_per_node` from args using either positional or flag syntax: | ||
| - Positional: second word is a number, e.g. `all 2` | ||
| - Flag: `--nproc_per_node <value>` anywhere in args, e.g. `all --nproc_per_node 2` | ||
|
|
||
| - If the second word is exactly `progress`, execute the **all progress** sub-command below. | ||
| - If no `nproc_per_node` value can be found, ask the user: "Please provide the number of GPUs per node (nproc_per_node)." and **STOP**. | ||
| - Otherwise use the parsed value and run the full pipeline. |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
cd /tmp && find . -name "SKILL.md" 2>/dev/null | head -5Repository: NVIDIA/Model-Optimizer
Length of output: 48
🏁 Script executed:
git ls-files | grep -i "skill.md"Repository: NVIDIA/Model-Optimizer
Length of output: 669
🏁 Script executed:
cat -n ".agents/skills/puzzletron/SKILL.md" | sed -n '30,85p'Repository: NVIDIA/Model-Optimizer
Length of output: 2840
Add explicit numeric validation rule to prevent shell injection in nproc_per_node.
The specification at lines 34-40 and 65-71 describes parsing nproc_per_node but does not mandate validation that the value is numeric before shell interpolation at lines 44-50 and 75-81. Malformed input (e.g., --nproc_per_node "$(malicious)") can reach the Bash torchrun command.
Add validation rule to both command sections:
- If
nproc_per_nodeis not strictly numeric (^[0-9]+$), ask the user: "nproc_per_node must be a positive integer." and STOP.
This should be added after the "value not found" check and before "Otherwise use the parsed value..." in both command specifications.
🤖 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 @.agents/skills/puzzletron/SKILL.md around lines 34 - 40, The specification
lacks numeric validation for the nproc_per_node parameter before it is
interpolated into shell commands, creating a security vulnerability for shell
injection attacks. Add an explicit validation rule to both the "all" and "local"
command sections in the skill specification that checks whether nproc_per_node
matches the pattern of a positive integer (^[0-9]+$). Insert this validation
check after the "value not found" check and before the "Otherwise use the parsed
value" instruction in both sections. If the value is not strictly numeric, the
specification should instruct to ask the user "nproc_per_node must be a positive
integer." and STOP before any shell command execution occurs.
Source: Coding guidelines
| ```bash | ||
| export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \ | ||
| torchrun --nproc_per_node <nproc_per_node> examples/puzzletron/main.py \ | ||
| --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \ | ||
| 2>&1 | tee ./log.txt | grep "Puzzletron Progress" | ||
| ``` | ||
|
|
||
| Stream output to the user as it arrives. When the command finishes, report the exit code. |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
# First, check if the file exists and view the relevant lines
if [ -f ".agents/skills/puzzletron/SKILL.md" ]; then
echo "=== Lines 40-55 (around line 46-53) ==="
sed -n '40,55p' ".agents/skills/puzzletron/SKILL.md"
echo ""
echo "=== Lines 75-90 (around line 77-84) ==="
sed -n '75,90p' ".agents/skills/puzzletron/SKILL.md"
else
echo "File not found"
fiRepository: NVIDIA/Model-Optimizer
Length of output: 1368
🏁 Script executed:
# Verify bash pipeline exit code behavior with a test
python3 - <<'PY'
import subprocess
import sys
# Test 1: torchrun fails but grep succeeds (simulating the issue)
print("Test 1: Command fails, grep succeeds")
result = subprocess.run(
"false | tee /tmp/test.txt | grep 'x'",
shell=True,
capture_output=True
)
print(f"Exit code without pipefail: {result.returncode}")
# Test 2: With set -o pipefail
print("\nTest 2: With set -o pipefail")
result = subprocess.run(
"set -o pipefail; false | tee /tmp/test.txt | grep 'x'",
shell=True,
capture_output=True
)
print(f"Exit code with pipefail: {result.returncode}")
# Test 3: PIPESTATUS array
print("\nTest 3: PIPESTATUS array")
result = subprocess.run(
"set -o pipefail; false | tee /tmp/test.txt | grep 'x'; echo \"PIPESTATUS: ${PIPESTATUS[@]}\"",
shell=True,
capture_output=True,
text=True
)
print(f"Output: {result.stdout}")
print(f"Exit code: {result.returncode}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 238
Pipeline exit code can mask torchrun failures.
Lines 46-53 and 77-84 use torchrun | tee | grep. In shell pipelines without pipefail, the exit code defaults to the rightmost command. If torchrun fails but grep finds the "Puzzletron Progress" pattern, the pipeline reports success (exit 0), masking the actual failure while the instruction says to report exit code.
Apply set -o pipefail to make the pipeline return non-zero when any command fails:
Suggested update
+ set -o pipefail && \
export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \
torchrun --nproc_per_node <nproc_per_node> examples/puzzletron/main.py \
--config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \
2>&1 | tee ./log.txt | grep "Puzzletron Progress"📝 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.
| ```bash | |
| export PYTHONPATH=$PYTHONPATH:/workspace/Model-Optimizer && \ | |
| torchrun --nproc_per_node <nproc_per_node> examples/puzzletron/main.py \ | |
| --config examples/puzzletron/configs/llama-3_1-8B_pruneffn_memory/llama-3_1-8B_pruneffn_memory.yaml \ | |
| 2>&1 | tee ./log.txt | grep "Puzzletron Progress" | |
| ``` | |
| Stream output to the user as it arrives. When the command finishes, report the exit code. |
🤖 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 @.agents/skills/puzzletron/SKILL.md around lines 46 - 53, The shell pipeline
using torchrun piped to tee piped to grep does not properly propagate exit codes
because without pipefail, the pipeline only returns the exit code of the
rightmost command (grep). When torchrun fails but grep successfully finds the
"Puzzletron Progress" pattern, the pipeline reports success even though the
actual torchrun command failed. To fix this, add set -o pipefail before or at
the beginning of the script block containing the torchrun command to ensure that
the pipeline returns a non-zero exit code when any command in the pipeline
fails, allowing accurate exit code reporting as mentioned in the instructions.
1 participant
What does this PR do?
Type of change: new feature
Experimental claude skill for puzzletron compression algorithm. See
.agents/skills/puzzletron/README.mdfor detailsUsage
see
.agents/skills/puzzletron/README.mdTesting
Before your PR is "Ready for review"
Summary by CodeRabbit
Release Notes
New Features
/puzzletron mipand/puzzletron allcommands to run model optimization pipelines./puzzletron progresscommand to track step status, elapsed time, and estimated remaining time.Documentation