Skip to content

refactor: move AF3 model running out of the backend - #660

Merged
DimaMolod merged 2 commits into
mainfrom
refactor/af3-runner
Oct 13, 2026
Merged

DimaMolod merged 2 commits into
mainfrom
refactor/af3-runner

Conversation

@DimaMolod

@DimaMolod DimaMolod commented Oct 13, 2026 •

Copy link
Copy Markdown
Collaborator

Last of three behaviour-preserving PRs that split alphafold3_backend.py; stacked on the input-preparation PR (#659). The backend goes from 2,455 lines on main to 300, holding only AlphaFold3Backend (setup, prepare_input, predict, postprocess).

Layout after the three PRs

Module Holds
alphafold3_backend.py AlphaFold3Backend; re-exports process_fold_input, predict_structure, ResultsForSeed
alphafold3_inputs.py AlphaPulldown objects and AF3 JSON → one fold input (FoldInputBuilder, PreparedFold)
alphafold3_runner.py Model running, adapted from upstream run_alphafold.py
alphafold3_outputs.py Output directory: per-seed and per-fold writers, ModelCIF metadata, resume
alphafold3_debug.py Debug artifacts: translated MSAs and their summary, final complex MSA
utils/output_paths.py String naming rules, shared with the AF2 path, no AF3 import

No extracted module imports the backend. alphafold3_outputs names ResultsForSeed only for type checking.

Commits

  1. The move. ConfigurableModel, ModelRunner, ResultsForSeed, predict_structure, process_fold_input, _record_inference_kernels and _initialise_tokamax_trace_context go unchanged to alphafold3_runner.py, so the code adapted from upstream can be diffed against it on its own. The backend re-exports what callers import from it. Unit tests of runner code patch and call it in the runner.
  2. setup tidied:
    • Its nested make_model_config becomes a runner function, without the leftover "new code approach" comments.
    • Its docstring says once why every field is set behind hasattr: AF3 releases differ in which fields their config has. return_distogram came to the AF3 history later than the others (2025-05).
    • setup looks up model.Model once per call and passes it to both the config and the runner, as the nested import did.
    • predict's buckets is typed as what it accepts: an int or a sequence.

Verification

  • CPU equivalence harness (the 15 inputs of the previous PR), after each commit: byte-identical to main.
  • Unit + integration, both release images: the same results as main plus the 6 builder tests of the previous PR. AF2 image: 937 passed, 26 skipped. AF3 image: 1008 passed, 15 skipped. The same pre-existing environmental failures.
  • GPU, full AF3 suite with RUN_EXPENSIVE_GPU_TESTS=1: 39/39 passed, including TestAlphaFold3MetadataEndToEnd (real feature generation through inference and ModelCIF), resume, fast kernels and the MMseqs/regression checks.

🤖 Generated with Claude Code

DimaMolod and others added 2 commits October 13, 2026 19:57
ConfigurableModel, ModelRunner, ResultsForSeed, predict_structure,
process_fold_input, _record_inference_kernels and
_initialise_tokamax_trace_context move unchanged to alphafold3_runner.py, the
part adapted from upstream run_alphafold.py, so it can be diffed against it on
its own. The runner imports the output writers and the debug writer; the
output module names ResultsForSeed only for type checking, from the runner.

The backend keeps AlphaFold3Backend and re-exports process_fold_input,
predict_structure and ResultsForSeed for callers that import them from it.
Unit tests of runner code patch and call it in the runner.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
setup's nested make_model_config becomes a runner function, without the
leftover "new code approach" comments; its docstring says once why every field
is set behind hasattr: AF3 releases differ in which fields their config has
(return_distogram is newer than the rest). setup looks up model.Model once per
call and passes it to both the config and the runner, as the nested import
did. predict's buckets is typed as what it accepts, an int or a sequence.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 13, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-10-13T18:22:53.901459Z 11cde8f PR opened
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@DimaMolod
DimaMolod changed the base branch from refactor/af3-inputs to main October 13, 2026 19:18
@DimaMolod
DimaMolod merged commit 9b6d369 into main Oct 13, 2026
6 checks passed
@DimaMolod
DimaMolod deleted the refactor/af3-runner branch October 13, 2026 19:19
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