Fix first-run checkpoint download race

CorpusCallosum/localization/inference.py

@@ -22,10 +22,12 @@
 
 from CorpusCallosum.transforms.localization import CropAroundACPCFixedSize
 from CorpusCallosum.utils.types import Points2dType
-from FastSurferCNN.download_checkpoints import load_checkpoint_config_defaults
-from FastSurferCNN.download_checkpoints import main as download_checkpoints
 from FastSurferCNN.utils import Image3d, Vector2d, Vector3d
-from FastSurferCNN.utils.checkpoint import get_config_file
+from FastSurferCNN.utils.checkpoint import (
+    get_checkpoints,
+    get_config_file,
+    load_checkpoint_config_defaults,
+)
 from FastSurferCNN.utils.parser_defaults import FASTSURFER_ROOT
 
 PATCH_SIZE = (64, 64)
@@ -59,10 +61,11 @@ def load_model(device: torch.device) -> DenseNet:
         dropout_prob=0.2
     )
 
-    download_checkpoints(cc=True)
     config_file = get_config_file("CorpusCallosum")
     cc_config = load_checkpoint_config_defaults("checkpoint", filename=config_file)
+    urls = load_checkpoint_config_defaults("url", filename=config_file)
     checkpoint_path = FASTSURFER_ROOT / cc_config['localization']
+    get_checkpoints(checkpoint_path, urls=urls)
 
     # Load state dict
     if isinstance(checkpoint_path, str) or isinstance(checkpoint_path, Path):

CorpusCallosum/segmentation/inference.py

@@ -23,11 +23,13 @@
 
 from CorpusCallosum.data import constants
 from CorpusCallosum.transforms.segmentation import CropAroundACPC
-from FastSurferCNN.download_checkpoints import load_checkpoint_config_defaults
-from FastSurferCNN.download_checkpoints import main as download_checkpoints
 from FastSurferCNN.models.networks import FastSurferVINN
 from FastSurferCNN.utils import Image3d, Image4d, Shape2d, Shape3d, Shape4d, Vector2d, nibabelImage
-from FastSurferCNN.utils.checkpoint import get_config_file
+from FastSurferCNN.utils.checkpoint import (
+    get_checkpoints,
+    get_config_file,
+    load_checkpoint_config_defaults,
+)
 from FastSurferCNN.utils.parallel import thread_executor
 
 
@@ -69,10 +71,11 @@ def load_model(device: torch.device | None = None) -> FastSurferVINN:
     }
     model = FastSurferVINN(params)
 
-    download_checkpoints(cc=True)
     config_file = get_config_file("CorpusCallosum")
     cc_config: dict[str, Path] = load_checkpoint_config_defaults("checkpoint", filename=config_file)
+    urls = load_checkpoint_config_defaults("url", filename=config_file)
     checkpoint_path = constants.FASTSURFER_ROOT / cc_config['segmentation']
+    get_checkpoints(checkpoint_path, urls=urls)
 
     weights = torch.load(checkpoint_path, weights_only=True, map_location=device)
     model.load_state_dict(weights)

FastSurferCNN/utils/checkpoint.py

@@ -19,6 +19,7 @@
 from functools import lru_cache
 from pathlib import Path
 from typing import TYPE_CHECKING, Literal, TypedDict, cast, overload
+from uuid import uuid4
 
 import requests
 import torch
@@ -400,8 +401,20 @@ def download_checkpoint(
             raise RuntimeError(message, responses)
     else:
         response = next(r for r in responses if r.ok)
-        with open(checkpoint_path, "wb") as f:
-            f.write(response.content)
+        checkpoint_path = Path(checkpoint_path)
+        temporary_path = checkpoint_path.with_name(
+            f".{checkpoint_path.name}.{uuid4().hex}.tmp"
+        )
+        try:
+            # Opening a unique file with ``xb`` preserves the permissions dictated by
+            # the process umask and prevents concurrent downloads from sharing a
+            # temporary file. The same-directory replace atomically publishes the
+            # checkpoint only after it has been written and closed completely.
+            with open(temporary_path, "xb") as f:
+                f.write(response.content)
+            os.replace(temporary_path, checkpoint_path)
+        finally:
+            temporary_path.unlink(missing_ok=True)
 
 
 def check_and_download_ckpts(checkpoint_path: Path | str, urls: list[str]) -> None: