diff --git a/src/simplefold/inference.py b/src/simplefold/inference.py index bccaf97..cd59b8b 100644 --- a/src/simplefold/inference.py +++ b/src/simplefold/inference.py @@ -22,6 +22,7 @@ from utils.datamodule_utils import process_one_inference_structure from utils.esm_utils import _af2_to_esm, esm_registry from utils.boltz_utils import process_structure, save_structure +from utils.checkpoint_utils import load_torch_checkpoint from utils.fasta_utils import process_fastas, download_fasta_utilities, check_fasta_inputs from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer @@ -82,14 +83,9 @@ def initialize_folding_model(args): ckpt_dir = Path(args.ckpt_dir) ckpt_path = os.path.join(ckpt_dir, f"{simplefold_model}.ckpt") - # create folding model - ckpt_path = os.path.join(ckpt_dir, f"{simplefold_model}.ckpt") - if not os.path.exists(ckpt_path): - os.makedirs(ckpt_dir, exist_ok=True) - os.system(f"curl -L {ckpt_url_dict[simplefold_model]} -o {ckpt_path}") cfg_path = get_config_path(f"configs/model/architecture/foldingdit_{simplefold_model[11:]}.yaml") - checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) + checkpoint = load_torch_checkpoint(ckpt_path, ckpt_url_dict[simplefold_model]) # load model checkpoint if args.backend == 'torch': @@ -122,12 +118,9 @@ def initialize_plddt_module(args, device): # load pLDDT module if specified plddt_ckpt_path = os.path.join(args.ckpt_dir, "plddt.ckpt") - if not os.path.exists(plddt_ckpt_path): - os.makedirs(args.ckpt_dir, exist_ok=True) - os.system(f"curl -L {plddt_ckpt_url} -o {plddt_ckpt_path}") plddt_module_path = get_config_path("configs/model/architecture/plddt_module.yaml") - plddt_checkpoint = torch.load(plddt_ckpt_path, map_location="cpu", weights_only=False) + plddt_checkpoint = load_torch_checkpoint(plddt_ckpt_path, plddt_ckpt_url) if args.backend == "torch": plddt_config = omegaconf.OmegaConf.load(plddt_module_path) @@ -150,12 +143,12 @@ def initialize_plddt_module(args, device): print(f"pLDDT output module loaded with {args.backend} backend.") plddt_latent_ckpt_path = os.path.join(args.ckpt_dir, "simplefold_1.6B.ckpt") - if not os.path.exists(plddt_latent_ckpt_path): - os.makedirs(args.ckpt_dir, exist_ok=True) - os.system(f"curl -L {ckpt_url_dict['simplefold_1.6B']} -o {plddt_latent_ckpt_path}") plddt_latent_config_path = get_config_path("configs/model/architecture/foldingdit_1.6B.yaml") - plddt_latent_checkpoint = torch.load(plddt_latent_ckpt_path, map_location="cpu", weights_only=False) + plddt_latent_checkpoint = load_torch_checkpoint( + plddt_latent_ckpt_path, + ckpt_url_dict["simplefold_1.6B"], + ) if args.backend == "torch": plddt_latent_config = omegaconf.OmegaConf.load(plddt_latent_config_path) diff --git a/src/simplefold/utils/checkpoint_utils.py b/src/simplefold/utils/checkpoint_utils.py new file mode 100644 index 0000000..57a04c7 --- /dev/null +++ b/src/simplefold/utils/checkpoint_utils.py @@ -0,0 +1,60 @@ +# +# For licensing see accompanying LICENSE file. +# Copyright (c) 2025 Apple Inc. Licensed under MIT License. +# + +import os +import pickle +import zipfile +from pathlib import Path +from urllib import request + +import torch + + +_CHECKPOINT_READ_ERROR_SNIPPETS = ( + "PytorchStreamReader failed reading zip archive", + "failed finding central directory", +) + + +def _download_checkpoint(url: str, checkpoint_path: Path) -> None: + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = checkpoint_path.with_name(f".{checkpoint_path.name}.download") + + if tmp_path.exists(): + tmp_path.unlink() + + try: + request.urlretrieve(url, tmp_path) + os.replace(tmp_path, checkpoint_path) + finally: + if tmp_path.exists(): + tmp_path.unlink() + + +def _is_checkpoint_read_error(error: BaseException) -> bool: + if isinstance(error, (EOFError, pickle.UnpicklingError, zipfile.BadZipFile)): + return True + + if isinstance(error, RuntimeError): + message = str(error) + return any(snippet in message for snippet in _CHECKPOINT_READ_ERROR_SNIPPETS) + + return False + + +def load_torch_checkpoint(path: os.PathLike[str] | str, url: str): + checkpoint_path = Path(path) + if not checkpoint_path.exists(): + _download_checkpoint(url, checkpoint_path) + + try: + return torch.load(checkpoint_path, map_location="cpu", weights_only=False) + except Exception as error: + if not _is_checkpoint_read_error(error): + raise + + checkpoint_path.unlink(missing_ok=True) + _download_checkpoint(url, checkpoint_path) + return torch.load(checkpoint_path, map_location="cpu", weights_only=False) diff --git a/src/simplefold/wrapper.py b/src/simplefold/wrapper.py index add3734..0713452 100644 --- a/src/simplefold/wrapper.py +++ b/src/simplefold/wrapper.py @@ -17,6 +17,7 @@ from utils.datamodule_utils import process_one_inference_structure from utils.esm_utils import _af2_to_esm, esm_registry from utils.boltz_utils import process_structure, save_structure +from utils.checkpoint_utils import load_torch_checkpoint from utils.fasta_utils import process_fastas, download_fasta_utilities from boltz_data_pipeline.feature.featurizer import BoltzFeaturizer from boltz_data_pipeline.tokenize.boltz_protein import BoltzTokenizer @@ -81,10 +82,8 @@ def from_pretrained_folding_model(self): # create folding model ckpt_path = os.path.join(self.ckpt_dir, f"{simplefold_model}.ckpt") - if not os.path.exists(ckpt_path): - os.system(f"curl -L -o {ckpt_path} {ckpt_url_dict[simplefold_model]}") - checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False) + checkpoint = load_torch_checkpoint(ckpt_path, ckpt_url_dict[simplefold_model]) # load model checkpoint cfg_path = os.path.join( @@ -123,13 +122,9 @@ def from_pretrained_plddt_model(self): # load pLDDT module if specified plddt_ckpt_path = os.path.join(self.ckpt_dir, "plddt.ckpt") - if not os.path.exists(plddt_ckpt_path): - os.system(f"curl -L -o {plddt_ckpt_path} {plddt_ckpt_url}") plddt_module_path = "configs/model/architecture/plddt_module.yaml" - plddt_checkpoint = torch.load( - plddt_ckpt_path, map_location="cpu", weights_only=False - ) + plddt_checkpoint = load_torch_checkpoint(plddt_ckpt_path, plddt_ckpt_url) if self.backend == "torch": plddt_config = omegaconf.OmegaConf.load(plddt_module_path) @@ -156,15 +151,11 @@ def from_pretrained_plddt_model(self): print(f"pLDDT output module loaded with {self.backend} backend.") plddt_latent_ckpt_path = os.path.join(self.ckpt_dir, "simplefold_1.6B.ckpt") - if not os.path.exists(plddt_latent_ckpt_path): - os.makedirs(self.ckpt_dir, exist_ok=True) - os.system( - f"curl -L -o {plddt_latent_ckpt_path} {ckpt_url_dict['simplefold_1.6B']}" - ) plddt_latent_config_path = "configs/model/architecture/foldingdit_1.6B.yaml" - plddt_latent_checkpoint = torch.load( - plddt_latent_ckpt_path, map_location="cpu", weights_only=False + plddt_latent_checkpoint = load_torch_checkpoint( + plddt_latent_ckpt_path, + ckpt_url_dict["simplefold_1.6B"], ) if self.backend == "torch": diff --git a/tests/test_checkpoint_utils.py b/tests/test_checkpoint_utils.py new file mode 100644 index 0000000..219af79 --- /dev/null +++ b/tests/test_checkpoint_utils.py @@ -0,0 +1,71 @@ +# +# For licensing see accompanying LICENSE file. +# Copyright (c) 2025 Apple Inc. Licensed under MIT License. +# + +from pathlib import Path + +from utils import checkpoint_utils + + +def test_missing_checkpoint_downloads_to_temp_file_before_loading(monkeypatch, tmp_path): + checkpoint_path = tmp_path / "model.ckpt" + seen_temp_paths = [] + + def fake_urlretrieve(url, filename): + temp_path = Path(filename) + seen_temp_paths.append(temp_path) + assert url == "https://example.com/model.ckpt" + assert temp_path != checkpoint_path + assert not checkpoint_path.exists() + temp_path.write_bytes(b"complete checkpoint") + + def fake_torch_load(path, map_location, weights_only): + assert Path(path) == checkpoint_path + assert map_location == "cpu" + assert weights_only is False + return Path(path).read_bytes() + + monkeypatch.setattr(checkpoint_utils.request, "urlretrieve", fake_urlretrieve) + monkeypatch.setattr(checkpoint_utils.torch, "load", fake_torch_load) + + result = checkpoint_utils.load_torch_checkpoint( + checkpoint_path, + "https://example.com/model.ckpt", + ) + + assert result == b"complete checkpoint" + assert checkpoint_path.read_bytes() == b"complete checkpoint" + assert seen_temp_paths == [tmp_path / ".model.ckpt.download"] + assert not seen_temp_paths[0].exists() + + +def test_corrupt_checkpoint_is_replaced_and_loaded(monkeypatch, tmp_path): + checkpoint_path = tmp_path / "model.ckpt" + checkpoint_path.write_bytes(b"truncated") + load_attempts = [] + + def fake_urlretrieve(url, filename): + assert url == "https://example.com/model.ckpt" + Path(filename).write_bytes(b"fresh checkpoint") + + def fake_torch_load(path, map_location, weights_only): + load_attempts.append(Path(path).read_bytes()) + if len(load_attempts) == 1: + raise RuntimeError( + "PytorchStreamReader failed reading zip archive: " + "failed finding central directory" + ) + return {"loaded": True} + + monkeypatch.setattr(checkpoint_utils.request, "urlretrieve", fake_urlretrieve) + monkeypatch.setattr(checkpoint_utils.torch, "load", fake_torch_load) + + result = checkpoint_utils.load_torch_checkpoint( + checkpoint_path, + "https://example.com/model.ckpt", + ) + + assert result == {"loaded": True} + assert load_attempts == [b"truncated", b"fresh checkpoint"] + assert checkpoint_path.read_bytes() == b"fresh checkpoint"