Recover corrupted checkpoint downloads

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)

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)

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":

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"