Skip to content

protein_binding: switch s4 to Boltz-2, remove hardcoded paths, pin GPUs - #61

Open
drawadiagram wants to merge 1 commit into
update/core-manager-0926from
update/protein-binding-boltz-0926
Open

protein_binding: switch s4 to Boltz-2, remove hardcoded paths, pin GPUs#61
drawadiagram wants to merge 1 commit into
update/core-manager-0926from
update/protein-binding-boltz-0926

Conversation

@drawadiagram

Copy link
Copy Markdown
Contributor

Splits the protein-binding workflow changes out of #59. Second of four.

Based on #60 (needs find_gpus and IMPRESS_SESSION_DIR). GitHub will retarget this to main when #60 merges — review only the commit on this branch.

Structure prediction: AlphaFold multimer → Boltz-2

  • s4_boltz.sh resolves its interpreter from BOLTZ_VENV/VIRTUAL_ENV instead of sourcing a hardcoded /anvil venv, and makes --use_msa_server opt-in via BOLTZ_USE_MSA_SERVER=1 because compute nodes have no internet.
  • CCD cache extraction race fixed. This killed 13 of 16 pipelines in a production run. boltz's own download_boltz2() decides to skip extraction based on whether the mols/ directory exists, but tarfile.extractall() creates that directory entry immediately — so a second concurrent task saw mols/ present, skipped extraction, read a half-populated cache, and died with CCD component not found. The lock is now held across a completeness check of mols/ against mols.tar, with a .mols_complete marker so later runs skip the O(45k) recount once warmed.
  • s3 embeds pre-computed MSA paths in the FASTA header, falling back to single-sequence mode when the cache is cold.
  • af2_multimer_reduced.sh and scripts/s4_alphafold.sh are deleted, along with the commented-out s4_alphafold task block that was their only remaining call site. That also resolves @AymenFJA's "remove any Dragon reference" comment on s4_alphafold.sh:10 — the file is gone.

⚠️ GPU scheduling — the change most worth reviewing

task_description={"gpus_per_rank": 1} is removed from the s1 and s4 signatures. GPU assignment now happens by pinning: a gpu_id kwarg is threaded into s4_boltz.sh as CUDA_VISIBLE_DEVICES, propagated to child pipelines, with the runner round-robining find_gpus() across pipelines. A module-level per-GPU asyncio.Semaphore(2) caps concurrent Boltz launches on each GPU across every pipeline pinned to it.

@AymenFJA — this is what your protein_binding.py:203/:212 comments were pointing at. To be precise about one thing: no task decorator changed. s4_post_exec was already local_task=True on main (I diffed them), and s4 already had capture_stdio=True. The real change is the resource-request removal above, and it deserves your eye.

Portability

MPNN_PATH comes from the environment and fails loudly when unset; base_path splits into input_base_path/output_base_path (IMPRESS_BASE_DIR / IMPRESS_OUTPUT_DIR / IMPRESS_SCRIPTS_DIR); every scripts/*.sh re-activates $VIRTUAL_ENV instead of a hardcoded cluster path; plddt_extract_pipeline.py writes its CSV under --path rather than cwd.

Adds delta_env_setup.sh and delta_gpu_run.sh (with IMPRESS_BACKEND=dragon|local and IMPRESS_TEST_MODE switches). The leaked --mail-user address @AymenFJA flagged is now a <your e-mail> placeholder.

Robustness

s1, s4 and s5 each check for their real output before re-raising, because the execution backend intermittently reports failure for a task that completed. run() now aborts the pass if every s4 failed instead of proceeding to pLDDT extraction on nothing.

Beyond a straight extraction

Deleting s4_alphafold.sh left dangling references in CLAUDE.md and README.md; both now describe Boltz-2 as the only s4 path.

Still open from #59

@AymenFJA's comment on delta_env_setup.sh:105 (the Dragon/Delta comment being misleading) is carried over unchanged — I'd rather you tell me what it should say than guess. On installing Dragon twice (:107): it's installed exactly once, via rhapsody-py[dragon,telemetry], which is what your comment asks for.

Replaces part of #59.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Lo8DwSbyvdWZRkkkka6gA2

Protein-binding workflow changes extracted from the impress_fixes branch.
Depends on the ImpressManager/find_gpus changes in the preceding commit.

Structure prediction (s4): AlphaFold multimer is replaced by Boltz-2.
- s4_boltz.sh resolves its interpreter from BOLTZ_VENV/VIRTUAL_ENV rather
  than sourcing a hardcoded /anvil venv, and makes --use_msa_server opt-in
  (BOLTZ_USE_MSA_SERVER=1) because compute nodes have no internet.
- Fix a CCD cache extraction race that killed 13/16 pipelines in a
  production run: boltz's download_boltz2() skips extraction when the
  mols/ directory merely exists, but tarfile.extractall() creates that
  directory entry immediately, so a second concurrent task read a
  half-populated cache and failed with "CCD component not found". The
  lock is now held across a completeness check of mols/ against
  mols.tar, with a .mols_complete marker to skip the O(45k) recount once
  warmed.
- s3 embeds pre-computed MSA paths in the FASTA header, falling back to
  single-sequence mode when the cache is cold.
- af2_multimer_reduced.sh and the commented-out s4_alphafold task are
  removed along with scripts/s4_alphafold.sh; the AF2 path had been dead
  for some time and its test-mode stub emitted Boltz-shaped output.

Portability: MPNN_PATH now comes from the environment and fails loudly
when unset; base_path is split into input_base_path/output_base_path
(IMPRESS_BASE_DIR / IMPRESS_OUTPUT_DIR / IMPRESS_SCRIPTS_DIR) so inputs,
outputs and scripts can live apart; every scripts/*.sh re-activates
$VIRTUAL_ENV instead of a hardcoded cluster path; plddt_extract_pipeline
writes its CSV under --path rather than cwd.

GPU scheduling: a gpu_id kwarg is threaded into s4_boltz.sh as
CUDA_VISIBLE_DEVICES and propagated to child pipelines, with the runner
round-robining find_gpus() across pipelines. A per-GPU semaphore caps
concurrent Boltz launches at 2 for each GPU, shared across every
pipeline pinned to it.

Failure handling: s1, s4 and s5 each check for their real output before
re-raising, because the execution backend intermittently reports a
failure for a task that completed; run() now aborts the pass if every s4
failed instead of proceeding to pLDDT extraction on nothing.

Adds delta_env_setup.sh and delta_gpu_run.sh for Delta HPC, with
IMPRESS_BACKEND=dragon|local and IMPRESS_TEST_MODE switches.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Lo8DwSbyvdWZRkkkka6gA2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant