Add flexible Diffuse params runs#225
Open
xraymemory wants to merge 4 commits intomainfrom
Open
Conversation
Declares sampleworks run profiles (boltz2-xrd, boltz2-md, protenix, rf3) with input schemas, args templates, volume mounts, and metadata fields. Applied to Diffuse via `diffuse apply`.
Contributor
📝 WalkthroughWalkthroughAdds a params-file execution mode and Diffuse contract plus related plumbing: new Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant Diffuse as Diffuse Scheduler
participant Entrypoint as docker-entrypoint.sh
participant RunParams as run_params module
participant RunGrid as run_grid_search.py
participant Container as Sampleworks Container
participant FS as Filesystem
User->>Diffuse: submit job with params_json
Diffuse->>FS: materialize `/diffuse/input/params.json`
Diffuse->>Container: start container with `--params /diffuse/input/params.json --output-dir /data/results`
Container->>Entrypoint: invoke entrypoint with `--params` args
Entrypoint->>RunParams: call infer_env_from_params -> determines pixi env
Entrypoint->>RunGrid: exec run_grid_search.py under inferred env
RunGrid->>RunParams: load and apply_run_params(params.json)
RunParams->>FS: write `params.original.json`, `params.resolved.json`, inline proteins CSV (if present)
RunGrid->>Container: orchestrate jobs (process pool, guidance)
RunGrid->>FS: write `run_summary.json` (status, counts, timestamps, error if any)
Container->>Diffuse: exit status / logs
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Review rate limit: 0/1 reviews remaining, refill in 60 minutes.Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@diffuse.yaml`:
- Around line 243-247: The static_args are using per-model flags
--protenix-checkpoint and --rf3-checkpoint which the script does not recognize;
replace those occurrences in static_args with the single generic flag
--model-checkpoint and keep the same checkpoint path values so the arguments
match the parser in run_grid_search.py (search for --protenix-checkpoint and
--rf3-checkpoint and change them to --model-checkpoint).
- Around line 79-93: The args_template in the diffuse profiles references
incorrect CLI flags for run_grid_search.py: update every profile's flag_args to
use the singular flags the script expects (change --models → --model and
--methods → --method) and replace any profile-specific checkpoint flags in
static_args (e.g., --protenix-checkpoint, --rf3-checkpoint) with the generic
--model-checkpoint used by run_grid_search.py; also correct the input_schema
entry for partial_diffusion_step so its default is a numeric value (not the
string "120") and its type remains number.
- Around line 42-44: The YAML schema sets partial_diffusion_step as type: number
but its default is a quoted string ("120"); update all occurrences of the
partial_diffusion_step default (the four instances) to an unquoted numeric value
(120) so the default matches the declared type; search for the key name
partial_diffusion_step in diffuse.yaml and replace the quoted default "120" with
120 for each occurrence.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a diffuse.yaml contract file for integrating Sampleworks with the Diffuse platform, defining 4 runnable profiles and shared metadata fields so Diffuse can apply/upsert these run configurations.
Changes:
- Introduces
diffuse.yamlwith 4 profiles (boltz2-xrd,boltz2-md,protenix,rf3) including container config, input schemas, and argument templates. - Defines default runtime settings (GPU min, shared memory, pull policy, retry/poll behavior) plus volume mounts for inputs/results/cache/checkpoints.
- Adds a
fields:section describing standardized metadata fields for runs.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+86
to
+89
| ensemble_sizes: --ensemble-sizes | ||
| gradient_weights: --gradient-weights | ||
| methods: --methods | ||
| boolean_args: |
| static_args: | ||
| - --output-dir | ||
| - /data/results | ||
| - --protenix-checkpoint |
| static_args: | ||
| - --output-dir | ||
| - /data/results | ||
| - --rf3-checkpoint |
| default: pure_guidance | ||
| - key: partial_diffusion_step | ||
| type: number | ||
| default: "120" |
Comment on lines
+31
to
+35
| - key: models | ||
| type: enum | ||
| required: true | ||
| allowed_values: [boltz2] | ||
| default: boltz2 |
Comment on lines
+36
to
+38
| - key: methods | ||
| type: text | ||
| default: "X-RAY DIFFRACTION" |
saada
approved these changes
Apr 22, 2026
xraymemory
changed the title
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@GRID_SEARCH.md`:
- Around line 29-30: The continued shell lines for the example flags (--model
boltz2 and --method "X-RAY DIFFRACTION") include trailing inline comments after
the backslash which breaks bash line continuation; remove or relocate those
inline annotations so the backslash is the last character on the line (e.g., put
explanatory comments on their own preceding lines or after the full continued
command), ensuring the --model and --method lines end with a literal "\" and no
trailing text so the block is copy/pasteable.
In `@run_grid_search.py`:
- Around line 292-305: The run summary currently derives overall status only
from the JobResult objects in results, which misses jobs dropped by crashed
workers; update the post-run summary logic (the block using results, successful,
failed, status before calling write_run_summary) to treat any missing JobResult
entries as failures by computing expected = len(filtered_jobs), got =
len(results), missing = max(0, expected - got) and adding missing to failed (or
use an explicit worker-failure count returned by run_grid_search() if
available); then compute status = "dry_run" if args.dry_run else ("failed" if
failed > 0 else "success") and pass the adjusted failed_jobs and successful_jobs
into write_run_summary so the summary reflects dropped jobs.
In `@tests/utils/test_run_params.py`:
- Around line 242-243: The test's regex in the pytest.raises call uses an
unescaped dot so "inputs.proteins or proteins_path" matches any char; update the
match pattern used in the with pytest.raises(...) assertion for apply_run_params
to escape the literal dot (e.g., use "inputs\.proteins or proteins_path") so the
test only passes for the exact error message referencing inputs.proteins; keep
the rest of the assertion and the call to apply_run_params(_namespace(...,
proteins=None)) unchanged.
🪄 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: defaults
Review profile: CHILL
Plan: Pro
Run ID: 0c3c980b-dc60-4184-8cbb-b7928e91f3b2
📒 Files selected for processing (9)
DockerfileGRID_SEARCH.mdREADME.mddiffuse.yamldocker-entrypoint.shrun_grid_search.pysrc/sampleworks/utils/guidance_script_utils.pysrc/sampleworks/utils/run_params.pytests/utils/test_run_params.py
✅ Files skipped from review due to trivial changes (1)
- README.md
Comment on lines
+29
to
+30
| --model boltz2 \ # options: boltz1, boltz2, protenix, rf3 (make sure env aligns!) | ||
| --method "X-RAY DIFFRACTION" \ # only useful for Boltz-2, ignored otherwise |
Contributor
There was a problem hiding this comment.
Move the inline annotations off these continued shell lines.
In bash, \ only continues the command when it is the last character on the line. With the trailing comments here, this example is not copy/pasteable as written.
🔧 Suggested doc fix
- --model boltz2 \ # options: boltz1, boltz2, protenix, rf3 (make sure env aligns!)
- --method "X-RAY DIFFRACTION" \ # only useful for Boltz-2, ignored otherwise
+ # options: boltz1, boltz2, protenix, rf3 (make sure env aligns!)
+ --model boltz2 \
+ # only useful for Boltz-2, ignored otherwise
+ --method "X-RAY DIFFRACTION" \📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| --model boltz2 \ # options: boltz1, boltz2, protenix, rf3 (make sure env aligns!) | |
| --method "X-RAY DIFFRACTION" \ # only useful for Boltz-2, ignored otherwise | |
| # options: boltz1, boltz2, protenix, rf3 (make sure env aligns!) | |
| --model boltz2 \ | |
| # only useful for Boltz-2, ignored otherwise | |
| --method "X-RAY DIFFRACTION" \ |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@GRID_SEARCH.md` around lines 29 - 30, The continued shell lines for the
example flags (--model boltz2 and --method "X-RAY DIFFRACTION") include trailing
inline comments after the backslash which breaks bash line continuation; remove
or relocate those inline annotations so the backslash is the last character on
the line (e.g., put explanatory comments on their own preceding lines or after
the full continued command), ensuring the --model and --method lines end with a
literal "\" and no trailing text so the block is copy/pasteable.
Comment on lines
+292
to
+305
| successful = sum(1 for result in results if result.status == "success") | ||
| failed = sum(1 for result in results if result.status == "failed") | ||
| status = "dry_run" if args.dry_run else "success" | ||
| if failed: | ||
| status = "failed" | ||
| write_run_summary( | ||
| args=args, | ||
| output_dir=args.output_dir, | ||
| status=status, | ||
| started_at=started_at, | ||
| total_jobs=len(filtered_jobs), | ||
| successful_jobs=successful, | ||
| failed_jobs=failed, | ||
| ) |
Contributor
There was a problem hiding this comment.
Don't derive the final run status from results alone.
If a worker crashes before writing its .results.pkl, run_grid_search() records that failure at Line 223-Line 225 but returns no JobResult for the dropped jobs. This block will then write status="success" with failed_jobs=0, which makes the run summary lie about a partially failed execution.
🔧 Suggested fix
- successful = sum(1 for result in results if result.status == "success")
- failed = sum(1 for result in results if result.status == "failed")
+ successful = sum(1 for result in results if result.status == "success")
+ failed = 0 if args.dry_run else max(0, len(filtered_jobs) - successful)
status = "dry_run" if args.dry_run else "success"
if failed:
status = "failed"At minimum, count missing JobResults as failures. Returning explicit worker-failure counts from run_grid_search() would be even better.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| successful = sum(1 for result in results if result.status == "success") | |
| failed = sum(1 for result in results if result.status == "failed") | |
| status = "dry_run" if args.dry_run else "success" | |
| if failed: | |
| status = "failed" | |
| write_run_summary( | |
| args=args, | |
| output_dir=args.output_dir, | |
| status=status, | |
| started_at=started_at, | |
| total_jobs=len(filtered_jobs), | |
| successful_jobs=successful, | |
| failed_jobs=failed, | |
| ) | |
| successful = sum(1 for result in results if result.status == "success") | |
| failed = 0 if args.dry_run else max(0, len(filtered_jobs) - successful) | |
| status = "dry_run" if args.dry_run else "success" | |
| if failed: | |
| status = "failed" | |
| write_run_summary( | |
| args=args, | |
| output_dir=args.output_dir, | |
| status=status, | |
| started_at=started_at, | |
| total_jobs=len(filtered_jobs), | |
| successful_jobs=successful, | |
| failed_jobs=failed, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@run_grid_search.py` around lines 292 - 305, The run summary currently derives
overall status only from the JobResult objects in results, which misses jobs
dropped by crashed workers; update the post-run summary logic (the block using
results, successful, failed, status before calling write_run_summary) to treat
any missing JobResult entries as failures by computing expected =
len(filtered_jobs), got = len(results), missing = max(0, expected - got) and
adding missing to failed (or use an explicit worker-failure count returned by
run_grid_search() if available); then compute status = "dry_run" if args.dry_run
else ("failed" if failed > 0 else "success") and pass the adjusted failed_jobs
and successful_jobs into write_run_summary so the summary reflects dropped jobs.
Comment on lines
+242
to
+243
| with pytest.raises(ValueError, match="inputs.proteins or proteins_path"): | ||
| apply_run_params(_namespace(params=str(params_path), proteins=None)) |
Contributor
There was a problem hiding this comment.
Escape the dot in this match= pattern.
pytest.raises(..., match=...) uses regex syntax, so inputs.proteins currently matches any character between inputs and proteins. That can let the test pass on the wrong error string.
🔧 Suggested fix
- with pytest.raises(ValueError, match="inputs.proteins or proteins_path"):
+ with pytest.raises(ValueError, match=r"inputs\.proteins or proteins_path"):🧰 Tools
🪛 Ruff (0.15.12)
[warning] 242-242: Pattern passed to match= contains metacharacters but is neither escaped nor raw
(RUF043)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/utils/test_run_params.py` around lines 242 - 243, The test's regex in
the pytest.raises call uses an unescaped dot so "inputs.proteins or
proteins_path" matches any char; update the match pattern used in the with
pytest.raises(...) assertion for apply_run_params to escape the literal dot
(e.g., use "inputs\.proteins or proteins_path") so the test only passes for the
exact error message referencing inputs.proteins; keep the rest of the assertion
and the call to apply_run_params(_namespace(..., proteins=None)) unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
sampleworksprofile that accepts a flexibleparams.jsonpayload and runs via--params /diffuse/input/params.json --output-dir ....-e <env> run_grid_search.py ...usage.params.original.json,params.resolved.json,run_summary.json, enrichedjob_metadata.json) and update docs/examples away from stale fixed-profile flags.Validation
bash -n docker-entrypoint.shpython3 -m py_compile run_grid_search.py src/sampleworks/utils/run_params.py src/sampleworks/utils/guidance_script_utils.py tests/utils/test_run_params.pyPYTHONPATH=src python3smoke checks for nested config, inline proteins, unknown-param preservation, and model inferenceNotes
--params <file>for parameter passing; no--params-jsonor--params-envuser path is exposed.Summary by CodeRabbit
New Features
Documentation
Build
Tests