From 8d1d9df321f8f1f927c992fd4f30b8857b23885b Mon Sep 17 00:00:00 2001 From: xraymemory Date: Mon, 13 Jul 2026 16:05:08 -0400 Subject: [PATCH 1/3] perf(rf3): reuse initialized model runtime --- src/sampleworks/models/rf3/wrapper.py | 52 ++++++++++--- .../utils/guidance_script_utils.py | 1 + tests/models/test_rf3_device_selection.py | 78 +++++++++++++++++++ tests/utils/test_guidance_script_utils.py | 31 ++++++++ 4 files changed, 153 insertions(+), 9 deletions(-) diff --git a/src/sampleworks/models/rf3/wrapper.py b/src/sampleworks/models/rf3/wrapper.py index fcf4c411..1cf6ab49 100644 --- a/src/sampleworks/models/rf3/wrapper.py +++ b/src/sampleworks/models/rf3/wrapper.py @@ -25,6 +25,9 @@ from sampleworks.utils.msa import MSAManager +_INFERENCE_ENGINE_ATTR = "_sampleworks_rf3_inference_engine" + + @dataclass(frozen=True, slots=True) class RF3Conditioning: """Conditioning tensors from RF3 trunk forward pass. @@ -218,6 +221,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 +236,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 +251,47 @@ 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"] self._device = self.inference_engine.trainer.fabric.device + if device is not None and _cuda_index(device) != _cuda_index(self._device): + raise ValueError( + f"Pre-loaded RF3 model 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 diff --git a/src/sampleworks/utils/guidance_script_utils.py b/src/sampleworks/utils/guidance_script_utils.py index f477ab73..cdff3efc 100644 --- a/src/sampleworks/utils/guidance_script_utils.py +++ b/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}") diff --git a/tests/models/test_rf3_device_selection.py b/tests/models/test_rf3_device_selection.py index 11f94979..3deb7775 100644 --- a/tests/models/test_rf3_device_selection.py +++ b/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,81 @@ 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, + ) + + @pytest.mark.gpu @pytest.mark.slow class TestRF3WrapperDeviceBinding: diff --git a/tests/utils/test_guidance_script_utils.py b/tests/utils/test_guidance_script_utils.py index a91a0fa5..c69641ee 100644 --- a/tests/utils/test_guidance_script_utils.py +++ b/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)} From cbe1fd09ef0a0864cf9e92fce5c719d993b037dc Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 16 Jul 2026 15:20:06 -0400 Subject: [PATCH 2/3] fix(rf3): validate reused model ownership --- src/sampleworks/models/rf3/wrapper.py | 8 +++++--- tests/models/test_rf3_device_selection.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/src/sampleworks/models/rf3/wrapper.py b/src/sampleworks/models/rf3/wrapper.py index 1cf6ab49..6d6d7aa0 100644 --- a/src/sampleworks/models/rf3/wrapper.py +++ b/src/sampleworks/models/rf3/wrapper.py @@ -281,11 +281,13 @@ def __init__( self.inference_engine.trainer = cast( RF3TrainerWithConfidence, self.inference_engine.trainer ) - self._device = self.inference_engine.trainer.fabric.device - if device is not None and _cuda_index(device) != _cuda_index(self._device): + if model is not None and self.inference_engine.trainer.state["model"] is not model: raise ValueError( - f"Pre-loaded RF3 model is on {self._device}, not requested device {device}" + "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"] diff --git a/tests/models/test_rf3_device_selection.py b/tests/models/test_rf3_device_selection.py index 3deb7775..1a5fd9e6 100644 --- a/tests/models/test_rf3_device_selection.py +++ b/tests/models/test_rf3_device_selection.py @@ -115,6 +115,24 @@ def test_rejects_reuse_on_different_device(self, tmp_path, stub_engine): 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 From 1ab1c4ccdd411992607a8204bc55f6d66d560621 Mon Sep 17 00:00:00 2001 From: xraymemory Date: Thu, 16 Jul 2026 15:36:19 -0400 Subject: [PATCH 3/3] docs(rf3): explain inference engine attachment --- src/sampleworks/models/rf3/wrapper.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/sampleworks/models/rf3/wrapper.py b/src/sampleworks/models/rf3/wrapper.py index 6d6d7aa0..f2ecbc5b 100644 --- a/src/sampleworks/models/rf3/wrapper.py +++ b/src/sampleworks/models/rf3/wrapper.py @@ -25,6 +25,8 @@ 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"