perf(rf3): reuse initialized model runtime

[manzuoni-astera]

Summary

• allow RF3Wrapper to reuse a model from an existing wrapper
• reuse the complete originating inference engine, trainer, pipeline, and Fabric context instead of loading a second checkpoint
• reject unsafe bare modules, checkpoint mismatches, and cross-device reuse
• forward pre-loaded RF3 models through get_model_and_device

Closes #291

Why the engine is reused

RF3 preprocessing and model execution are coupled through Foundry's initialized RF3InferenceEngine. Replacing only trainer.state["model"] after initialization would still load the checkpoint and could attach the model to the wrong Fabric runtime. The originating engine is therefore retained with the model and reused as one runtime context.

Validation

• pytest tests/models/test_rf3_device_selection.py tests/utils/test_guidance_script_utils.py -q (21 passed, 1 skipped)
• ty check on touched files in rf3-dev
• ruff check and ruff format --check on touched files
• git diff --check

Summary by CodeRabbit

• New Features

  • RF3 wrappers can now reuse pre-initialized models and inference runtimes.
  • Added validation for checkpoint, device, and model/runtime compatibility when reusing models.
  • Preloaded RF3 models are now correctly passed through model and device setup utilities.

• Bug Fixes

  • Prevented unnecessary inference engine reinitialization for reused RF3 models.
  • Added clear validation for invalid or incomplete model reuse contexts.

[coderabbitai]

Warning

Review limit reached

@manzuoni-astera, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 43 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 11dea641-2581-455d-845b-e1b99b67774f

📥 Commits

Reviewing files that changed from the base of the PR and between cbe1fd0 and 1ab1c4c.

📒 Files selected for processing (1)

• src/sampleworks/models/rf3/wrapper.py

📝 Walkthrough

Walkthrough

RF3Wrapper now accepts and validates preinitialized models with their inference-engine runtime. Freshly initialized models retain the attached engine context, and get_model_and_device forwards preloaded RF3 models to the wrapper. Tests cover reuse, device, binding, and forwarding behavior.

Changes

RF3 model reuse

Layer / File(s)	Summary
RF3Wrapper reuse runtime src/sampleworks/models/rf3/wrapper.py	Adds an optional model parameter, supports fresh or reused inference-engine initialization, validates checkpoint/device/model binding, and stores the engine context on fresh models.
Model reuse regression coverage tests/models/test_rf3_device_selection.py	Tests engine preservation, missing context, device mismatches, and trainer-model binding validation.
Guidance model forwarding src/sampleworks/utils/guidance_script_utils.py, tests/utils/test_guidance_script_utils.py	Forwards preloaded RF3 models to RF3Wrapper and verifies the exact model object is passed through.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Sequence Diagram(s)

sequenceDiagram
  participant GuidanceScript
  participant RF3Wrapper
  participant RF3Model
  participant RF3InferenceEngine
  GuidanceScript->>RF3Wrapper: pass preloaded model
  RF3Wrapper->>RF3Model: retrieve attached engine
  RF3Wrapper->>RF3InferenceEngine: validate checkpoint and device
  RF3Wrapper->>RF3Model: validate trainer binding

Loading

Possibly related PRs

• diff-use/sampleworks#222: Modifies overlapping RF3Wrapper device handling and related construction tests.

Suggested reviewers: dorismai, k-chrispens

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: reusing an initialized RF3 model runtime.
Linked Issues check	✅ Passed	The PR implements #291 by adding RF3Wrapper model reuse and forwarding preloaded RF3 models through get_model_and_device.
Out of Scope Changes check	✅ Passed	The changes stay focused on RF3 model reuse and its tests, with no clear unrelated additions.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[Copilot]

Pull request overview

This PR adds support for reusing an already-initialized RosettaFold3 (RF3) runtime context when constructing RF3Wrapper, and wires that through get_model_and_device so callers can pass pre-loaded RF3 models without re-loading checkpoints. This aligns RF3 with existing model-wrapper reuse patterns (Boltz/Protenix) while preserving RF3’s coupled preprocessing + inference engine initialization.

Changes:

• Extend RF3Wrapper to accept a pre-loaded model and reuse its originating RF3InferenceEngine/trainer/Fabric context instead of creating a new engine.
• Forward the optional model argument through get_model_and_device() for RF3 construction.
• Add tests covering RF3 model reuse behavior and validation (checkpoint/device mismatch, unsafe reuse).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
src/sampleworks/models/rf3/wrapper.py	Accepts pre-loaded RF3 models and reuses the originating inference engine/runtime context with validation.
src/sampleworks/utils/guidance_script_utils.py	Forwards model= into the RF3 wrapper path in get_model_and_device().
tests/models/test_rf3_device_selection.py	Adds unit tests verifying RF3 engine/model reuse and rejection of unsafe/mismatched reuse.
tests/utils/test_guidance_script_utils.py	Adds a test ensuring get_model_and_device() forwards a pre-loaded model into RF3 wrapper construction.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[marcuscollins]

One small nit, just to add a comment on where to find how a particular attribute is set. If you can add that, go ahead and merge. Thanks.

[coderabbitai]

🧹 Nitpick comments (2)

tests/models/test_rf3_device_selection.py (1)

59-71: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add NumPy-style docstrings to the new stub constructors.

• tests/models/test_rf3_device_selection.py#L59-L71: document StubTrainer.__init__ and StubInferenceEngine.__init__.
• tests/utils/test_guidance_script_utils.py#L45-L46: document StubRF3Wrapper.__init__.

As per coding guidelines, every function and class must include a NumPy-style docstring.

🤖 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/models/test_rf3_device_selection.py` around lines 59 - 71, Add
NumPy-style docstrings to the constructors StubTrainer.__init__ and
StubInferenceEngine.__init__ in tests/models/test_rf3_device_selection.py (lines
59-71), documenting their initialization behavior and parameters; also document
StubRF3Wrapper.__init__ in tests/utils/test_guidance_script_utils.py (lines
45-46).

Source: Coding guidelines

src/sampleworks/models/rf3/wrapper.py (1)

274-279: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add regression coverage for checkpoint mismatches.

This required rejection branch is not exercised by the supplied reuse tests, leaving path normalization or validation regressions undetected. Add a test that reuses original.model with a different checkpoint and expects this ValueError.

The PR objectives explicitly require checkpoint mismatches to be rejected.

🤖 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 `@src/sampleworks/models/rf3/wrapper.py` around lines 274 - 279, Add regression
coverage in the reuse tests for the checkpoint validation branch in the RF3
wrapper: reuse an instance loaded from original.model while requesting a
different checkpoint, and assert that initialization raises ValueError. Ensure
the test exercises the normalized engine_checkpoint versus checkpoint_path
comparison and verifies mismatched checkpoints are rejected.

🤖 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.

Nitpick comments:
In `@src/sampleworks/models/rf3/wrapper.py`:
- Around line 274-279: Add regression coverage in the reuse tests for the
checkpoint validation branch in the RF3 wrapper: reuse an instance loaded from
original.model while requesting a different checkpoint, and assert that
initialization raises ValueError. Ensure the test exercises the normalized
engine_checkpoint versus checkpoint_path comparison and verifies mismatched
checkpoints are rejected.

In `@tests/models/test_rf3_device_selection.py`:
- Around line 59-71: Add NumPy-style docstrings to the constructors
StubTrainer.__init__ and StubInferenceEngine.__init__ in
tests/models/test_rf3_device_selection.py (lines 59-71), documenting their
initialization behavior and parameters; also document StubRF3Wrapper.__init__ in
tests/utils/test_guidance_script_utils.py (lines 45-46).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f9e5e410-0172-49a7-a391-9d6c37ab0f38

📥 Commits

Reviewing files that changed from the base of the PR and between 063f519 and cbe1fd0.

📒 Files selected for processing (4)

• src/sampleworks/models/rf3/wrapper.py
• src/sampleworks/utils/guidance_script_utils.py
• tests/models/test_rf3_device_selection.py
• tests/utils/test_guidance_script_utils.py

[manzuoni-astera]

Addressed the remaining nit in 1ab1c4c by documenting exactly where _INFERENCE_ENGINE_ATTR is attached and how it is consumed. The later ownership-validation update also ensures a supplied model is exactly the trainer-owned model for that engine context.