perf(rf3): reuse initialized model runtime

src/sampleworks/models/rf3/wrapper.py

@@ -25,6 +25,11 @@
 from sampleworks.utils.msa import MSAManager
 
 
+# Attached to the trainer-owned model after RF3InferenceEngine initialization in
+# RF3Wrapper.__init__; reused wrappers retrieve the complete runtime through it.
+_INFERENCE_ENGINE_ATTR = "_sampleworks_rf3_inference_engine"
+
+
 @dataclass(frozen=True, slots=True)
 class RF3Conditioning:
     """Conditioning tensors from RF3 trunk forward pass.
@@ -218,6 +223,7 @@ def __init__(
         checkpoint_path: str | Path,
         msa_manager: MSAManager | None = None,
         device: torch.device | str | None = None,
+        model: Any | None = None,
     ):
         """
         Parameters
@@ -232,6 +238,11 @@ def __init__(
             parallel jobs that must target distinct GPUs — passing an ``int``
             to Fabric (the default) always resolves to GPU 0, which serialises
             otherwise-parallel workers onto a single device.
+        model: Any | None
+            Model previously obtained from another ``RF3Wrapper``. RF3 model
+            reuse also reuses the originating inference engine because its
+            trainer, preprocessing pipeline, Fabric strategy, and model are a
+            single initialized runtime context.
 
             References: https://lightning.ai/docs/fabric/stable/fundamentals/launch.html
               devices argument to fabric run
@@ -242,22 +253,49 @@ def __init__(
         self.checkpoint_path = Path(checkpoint_path).expanduser().resolve()
         self.msa_manager = msa_manager
         self.msa_pairing_strategy = "greedy"
+        self.inference_engine: RF3InferenceEngine
 
-        engine_kwargs: dict[str, Any] = {
-            "ckpt_path": str(self.checkpoint_path),
-            "diffusion_batch_size": 1,
-        }
-        if device is not None:
-            engine_kwargs["devices_per_node"] = [_cuda_index(device)]
+        if model is None:
+            engine_kwargs: dict[str, Any] = {
+                "ckpt_path": str(self.checkpoint_path),
+                "diffusion_batch_size": 1,
+            }
+            if device is not None:
+                engine_kwargs["devices_per_node"] = [_cuda_index(device)]
 
-        self.inference_engine = RF3InferenceEngine(**engine_kwargs)
-        self.inference_engine.initialize()
+            self.inference_engine = RF3InferenceEngine(**engine_kwargs)
+            self.inference_engine.initialize()
+        else:
+            inference_engine = getattr(model, _INFERENCE_ENGINE_ATTR, None)
+            if inference_engine is None:
+                raise ValueError(
+                    "RF3 model reuse requires a model created by RF3Wrapper so its "
+                    "initialized inference engine and Fabric context can also be reused"
+                )
+            self.inference_engine = cast(RF3InferenceEngine, inference_engine)
+            engine_checkpoint = Path(self.inference_engine.ckpt_path).expanduser().resolve()
+            if engine_checkpoint != self.checkpoint_path:
+                raise ValueError(
+                    f"Pre-loaded RF3 model uses checkpoint {engine_checkpoint}, "
+                    f"not requested checkpoint {self.checkpoint_path}"
+                )
 
         self.inference_engine.trainer = cast(
             RF3TrainerWithConfidence, self.inference_engine.trainer
         )
-        self.model = self.inference_engine.trainer.state["model"]
+        if model is not None and self.inference_engine.trainer.state["model"] is not model:
+            raise ValueError(
+                "Pre-loaded RF3 model does not belong to its attached inference engine"
+            )
         self._device = self.inference_engine.trainer.fabric.device
+        if device is not None and _cuda_index(device) != _cuda_index(self._device):
+            raise ValueError(f"RF3 runtime is on {self._device}, not requested device {device}")
+
+        if model is None:
+            self.model = self.inference_engine.trainer.state["model"]
+            object.__setattr__(self.model, _INFERENCE_ENGINE_ATTR, self.inference_engine)
+        else:
+            self.model = model
 
         # Chiral feature state, set in featurize()
         self._track_chiral_features: bool = False

src/sampleworks/utils/guidance_script_utils.py

@@ -213,6 +213,7 @@ def get_model_and_device(
             checkpoint_path=validated_checkpoint_path,
             msa_manager=MSAManager(),
             device=device,
+            model=model,
         )
     else:
         raise ValueError(f"Unknown model type: {model_type}")

tests/models/test_rf3_device_selection.py

@@ -6,6 +6,8 @@
 that specific GPU index via ``devices_per_node=[idx]``.
 """
 
+from pathlib import Path
+
 import pytest
 import torch
 from sampleworks.utils.imports import require_rf3, RF3_AVAILABLE
@@ -14,6 +16,7 @@
 pytestmark = pytest.mark.skipif(not RF3_AVAILABLE, reason="RF3 dependencies not installed")
 
 if RF3_AVAILABLE:
+    import sampleworks.models.rf3.wrapper as rf3_wrapper_module
     from sampleworks.models.rf3.wrapper import _cuda_index, RF3Wrapper
 
 
@@ -38,6 +41,99 @@ def test_rejects_non_cuda(self, device):
             _cuda_index(device)
 
 
+class TestRF3ModelReuse:
+    """Verify model reuse preserves the complete initialized RF3 runtime."""
+
+    @pytest.fixture
+    def stub_engine(self, monkeypatch):
+        """Replace RF3InferenceEngine with a lightweight initialized runtime."""
+
+        class StubFabric:
+            """Provide the Fabric attributes RF3Wrapper reads."""
+
+            device = torch.device("cuda:0")
+
+        class StubTrainer:
+            """Provide initialized trainer state with a model."""
+
+            def __init__(self):
+                self.fabric = StubFabric()
+                self.state = {"model": torch.nn.Linear(1, 1)}
+
+        class StubInferenceEngine:
+            """Track construction while avoiding checkpoint and GPU work."""
+
+            instances = 0
+
+            def __init__(self, **kwargs):
+                type(self).instances += 1
+                self.ckpt_path = Path(kwargs["ckpt_path"]).resolve()
+                self.trainer = StubTrainer()
+
+            def initialize(self):
+                """Match the real inference engine initialization interface."""
+
+        monkeypatch.setattr(rf3_wrapper_module, "RF3InferenceEngine", StubInferenceEngine)
+        return StubInferenceEngine
+
+    def test_reuses_engine_trainer_and_model(self, tmp_path, stub_engine):
+        """A reused model must not initialize a second RF3 inference engine."""
+        checkpoint = tmp_path / "rf3.ckpt"
+        original = RF3Wrapper(checkpoint_path=checkpoint, device="cuda:0")
+
+        reused = RF3Wrapper(
+            checkpoint_path=checkpoint,
+            device="cuda:0",
+            model=original.model,
+        )
+
+        assert stub_engine.instances == 1
+        assert reused.inference_engine is original.inference_engine
+        assert reused.inference_engine.trainer is original.inference_engine.trainer
+        assert reused.model is original.model
+
+    def test_rejects_model_without_inference_context(self, tmp_path, stub_engine):
+        """A bare RF3 module cannot be safely attached to a new Fabric runtime."""
+        with pytest.raises(ValueError, match="model created by RF3Wrapper"):
+            RF3Wrapper(
+                checkpoint_path=tmp_path / "rf3.ckpt",
+                device="cuda:0",
+                model=torch.nn.Linear(1, 1),
+            )
+
+        assert stub_engine.instances == 0
+
+    def test_rejects_reuse_on_different_device(self, tmp_path, stub_engine):
+        """An initialized RF3 runtime cannot be silently moved between GPUs."""
+        checkpoint = tmp_path / "rf3.ckpt"
+        original = RF3Wrapper(checkpoint_path=checkpoint, device="cuda:0")
+
+        with pytest.raises(ValueError, match="not requested device"):
+            RF3Wrapper(
+                checkpoint_path=checkpoint,
+                device="cuda:1",
+                model=original.model,
+            )
+
+    def test_rejects_model_spoofing_engine_context(self, tmp_path, stub_engine):
+        """The supplied model must be the model owned by the attached trainer."""
+        checkpoint = tmp_path / "rf3.ckpt"
+        original = RF3Wrapper(checkpoint_path=checkpoint, device="cuda:0")
+        other_model = torch.nn.Linear(1, 1)
+        object.__setattr__(
+            other_model,
+            rf3_wrapper_module._INFERENCE_ENGINE_ATTR,
+            original.inference_engine,
+        )
+
+        with pytest.raises(ValueError, match="does not belong"):
+            RF3Wrapper(
+                checkpoint_path=checkpoint,
+                device="cuda:0",
+                model=other_model,
+            )
+
+
 @pytest.mark.gpu
 @pytest.mark.slow
 class TestRF3WrapperDeviceBinding:

tests/utils/test_guidance_script_utils.py

@@ -5,11 +5,13 @@
 from unittest.mock import MagicMock
 
 import pytest
+import sampleworks.utils.guidance_script_utils as guidance_script_utils
 import torch
 from sampleworks.utils.guidance_script_arguments import GuidanceConfig, JobResult
 from sampleworks.utils.guidance_script_utils import (
     _three_state_resolver,
     _write_job_metadata,
+    get_model_and_device,
     get_reward_function_and_structure,
     save_everything,
 )
@@ -33,6 +35,35 @@ def test_resolve_alignment_reverse_diffusion(override, is_boltz, expected):
     assert _three_state_resolver(override, is_boltz) is expected
 
 
+def test_get_model_and_device_forwards_preloaded_model_to_rf3(monkeypatch):
+    """RF3 construction receives the pre-loaded model supplied by callers."""
+    preloaded_model = object()
+
+    class StubRF3Wrapper:
+        """Capture RF3 constructor arguments without loading model dependencies."""
+
+        def __init__(self, **kwargs):
+            self.kwargs = kwargs
+
+    monkeypatch.setattr(guidance_script_utils, "RF3Wrapper", StubRF3Wrapper)
+    monkeypatch.setattr(
+        guidance_script_utils,
+        "validate_model_checkpoint",
+        lambda model_type, checkpoint: checkpoint,
+    )
+    monkeypatch.setattr(guidance_script_utils, "MSAManager", lambda: object())
+
+    device, wrapper = get_model_and_device(
+        "cpu",
+        "rf3.ckpt",
+        "rf3",
+        model=preloaded_model,
+    )
+
+    assert device == torch.device("cpu")
+    assert wrapper.kwargs["model"] is preloaded_model
+
+
 def test_save_everything_uses_model_atom_array_for_mismatch(tmp_path: Path):
     """Mismatch final_state should save with model template when provided."""
     refined_structure = {"asym_unit": build_test_atom_array(n_atoms=3, with_occupancy=True)}