fix: make the first-use weight download resolve every id and fail loudly

spineps/utils/auto_download.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import shutil
+import tempfile
 import urllib.request
 import zipfile
 from pathlib import Path
@@ -38,8 +39,12 @@
     "vibe": link + current_highest_version + "/vibe.zip",
     "ct": link + current_highest_ct_version + "/ct.zip",
 }
+# Semantic models are also accepted with an explicit "_semantic" suffix: "t2w_semantic" names the same
+# weights as "t2w". The alias has to resolve to the same download name and release version as the
+# canonical id, which is what canonical_model_key() below is for.
+SEMANTIC_ALIAS_SUFFIX = "_semantic"
 for i, j in semantic.copy().items():
-    semantic[i + "_semantic"] = j
+    semantic[i + SEMANTIC_ALIAS_SUFFIX] = j
 
 labeling: dict[str, Path | str] = {
     "t2w_labeling": link + current_labeling_highest_version + "/labeling.zip",
@@ -58,6 +63,26 @@
 }
 
 
+def canonical_model_key(key: str) -> str:
+    """Maps a model id alias onto the id its weights are registered under.
+
+    ``"t2w_semantic"`` and ``"t2w"`` are the same model, but only the canonical id has entries in
+    ``download_names`` and ``phase_to_version``. Without this, asking for an ``*_semantic`` id raised
+    ``KeyError: 't2w_semantic'`` on first use -- even though those ids are listed as available models.
+
+    Args:
+        key (str): Model id as the user (or the model registry) spelled it.
+
+    Returns:
+        str: The canonical model id.
+    """
+    if key.endswith(SEMANTIC_ALIAS_SUFFIX):
+        canonical = key[: -len(SEMANTIC_ALIAS_SUFFIX)]
+        if canonical in download_names:
+            return canonical
+    return key
+
+
 def download_if_missing(key: str, url: Path | str, phase: SpinepsPhase) -> Path:
     """Return the local model folder for a model, downloading and extracting its weights if absent.
 
@@ -71,7 +96,14 @@ def download_if_missing(key: str, url: Path | str, phase: SpinepsPhase) -> Path:
 
     Returns:
         Path: Path to the local model folder containing the (possibly just downloaded) weights.
+
+    Raises:
+        KeyError: If the model id is unknown to the download registry.
+        RuntimeError: If the weights could not be downloaded or extracted.
     """
+    key = canonical_model_key(key)
+    if key not in download_names:
+        raise KeyError(f"no download is registered for model '{key}', known models are {sorted(download_names)}")
     version = phase_to_version.get(f"{phase.name}_{key}", phase_to_version[phase.name])
     out_path = Path(get_mri_segmentor_models_dir(), download_names[key] + "_" + version)
     if not out_path.exists():
@@ -81,50 +113,91 @@ def download_if_missing(key: str, url: Path | str, phase: SpinepsPhase) -> Path:
 
 
 def download_weights(weights_url: Path | str, out_path: Path | str) -> None:
-    """Download a weights zip archive, extract it into ``out_path`` and remove the archive.
+    """Download a weights zip archive and extract it into ``out_path``.
 
-    Shows a progress bar during download. If the extracted archive nests its contents in an extra subfolder
-    (no ``inference_config.json`` at the top level), the inner contents are moved up one level. Returns early
-    without raising if the initial size request fails.
+    Shows a progress bar during download. Everything happens inside a temporary folder next to
+    ``out_path``, which is moved into place only once the archive extracted and actually contains an
+    ``inference_config.json``; on any failure the temporary folder is removed and the error is raised.
+    That matters because the caller treats "``out_path`` exists" as "the model is installed": a download
+    that failed halfway used to leave a partial folder behind that was never retried and never worked,
+    and a failure before the download even started was swallowed entirely, surfacing much later as a
+    confusing "config not found" from the model loader.
 
     Args:
         weights_url: URL of the weights zip archive to download.
-        out_path: Destination folder for the extracted weights (the archive is downloaded next to it as ``.zip``).
+        out_path: Destination folder for the extracted weights.
 
     Raises:
-        AssertionError: If the nested archive layout is detected but the extra entry is not a directory.
+        RuntimeError: If the archive could not be downloaded or extracted, or if what was extracted does
+            not look like a model folder.
     """
     out_path = Path(out_path)
     logger = Print_Logger()
-    try:
-        # Retrieve file size
-        with urllib.request.urlopen(str(weights_url)) as response:
-            file_size = int(response.info().get("Content-Length", -1))
-    except Exception:
-        logger.on_fail("Download attempt failed:", weights_url)
+    if out_path.exists():
         return
-    logger.print("Downloading pretrained weights...")
-
-    with tqdm(total=file_size, unit="B", unit_scale=True, unit_divisor=1024, desc=Path(weights_url).name) as pbar:
-
-        def update_progress(block_num: int, block_size: int, total_size: int) -> None:
-            if pbar.total != total_size:
-                pbar.total = total_size
-            pbar.update(block_num * block_size - pbar.n)
-
-        zip_path = Path(str(out_path) + ".zip")
-        # Download the file
-        urllib.request.urlretrieve(str(weights_url), zip_path, reporthook=update_progress)
-
-    logger.print("Extracting pretrained weights...")
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    # A unique staging folder, so that several SPINEPS processes (e.g. a job array) downloading the same
+    # model at the same time cannot overwrite each other's partial files.
+    staging = Path(tempfile.mkdtemp(prefix=out_path.name + ".download.", dir=out_path.parent))
+    try:
+        zip_path = staging.joinpath("weights.zip")
+        extracted = staging.joinpath("extracted")
+        try:
+            with urllib.request.urlopen(str(weights_url)) as response:
+                file_size = int(response.info().get("Content-Length", -1))
+            logger.print("Downloading pretrained weights...")
+            with tqdm(total=file_size, unit="B", unit_scale=True, unit_divisor=1024, desc=Path(weights_url).name) as pbar:
+
+                def update_progress(block_num: int, block_size: int, total_size: int) -> None:
+                    if pbar.total != total_size:
+                        pbar.total = total_size
+                    pbar.update(block_num * block_size - pbar.n)
+
+                urllib.request.urlretrieve(str(weights_url), zip_path, reporthook=update_progress)
+        except Exception as e:
+            raise RuntimeError(
+                f"could not download the model weights from {weights_url} ({e}). Check your internet "
+                "connection, or download the weights manually from "
+                "https://github.com/Hendrik-code/spineps/releases and extract them into "
+                f"{out_path.parent} (see the README)."
+            ) from e
+
+        logger.print("Extracting pretrained weights...")
+        try:
+            with zipfile.ZipFile(zip_path, "r") as zip_ref:
+                zip_ref.extractall(extracted)
+        except zipfile.BadZipFile as e:
+            raise RuntimeError(f"the downloaded weights archive from {weights_url} is corrupt ({e}); try again.") from e
+
+        model_folder = _model_folder_in(extracted)
+        if model_folder is None:
+            contents = sorted(p.name for p in extracted.rglob("*"))[:10]
+            raise RuntimeError(f"the archive from {weights_url} contains no inference_config.json, found {contents}")
+        # Another process may have finished the same download while this one was busy; theirs is as good.
+        if not out_path.exists():
+            shutil.move(str(model_folder), str(out_path))
+    finally:
+        shutil.rmtree(staging, ignore_errors=True)
+
+
+def _model_folder_in(extracted: Path) -> Path | None:
+    """Picks the folder to install as the model folder out of a freshly extracted archive.
+
+    Archives either hold the model at the top level or wrap it in one extra folder. The config itself may
+    sit deeper still (nnU-Net plan subfolders), which the model loader resolves by globbing, so only its
+    presence somewhere below is verified here.
 
-    with zipfile.ZipFile(zip_path, "r") as zip_ref:
-        zip_ref.extractall(out_path)
-    # Test if there is an additional folder and move the content on up.
-    if not Path(out_path, "inference_config.json").exists():
-        source = next(out_path.iterdir())
-        assert source.is_dir()
-        for i in source.iterdir():
-            shutil.move(i, out_path)
+    Args:
+        extracted (Path): Folder the archive was extracted into.
 
-    zip_path.unlink()
+    Returns:
+        Path | None: The folder to move into place, or None if this is not a SPINEPS model archive.
+    """
+    if next(extracted.rglob("inference_config.json"), None) is None:
+        return None
+    if extracted.joinpath("inference_config.json").is_file():
+        return extracted
+    children = sorted(extracted.iterdir())
+    if len(children) == 1 and children[0].is_dir():
+        return children[0]  # the archive wrapped the model in one extra folder
+    return extracted

unit_tests/test_auto_download.py

@@ -0,0 +1,190 @@
+# Call 'python -m unittest' on this folder
+"""First-use weight download: resolve every advertised model id, and never leave a broken folder behind.
+
+Three defects made the automatic download unreliable:
+
+* the ``*_semantic`` ids (``t2w_semantic``, ``ct_semantic``, ...) are offered in the model registry and
+  even listed in the "Available semantic models" error message, but had no ``download_names`` entry, so
+  asking for one raised ``KeyError: 't2w_semantic'`` on a machine without the weights,
+* ``download_weights`` logged "Download attempt failed" and **returned normally** when the URL could not
+  be opened, so the caller happily returned a path that does not exist and the failure surfaced much
+  later as a confusing config-not-found from the model loader, and
+* the archive was extracted straight into the final folder, which the caller reads as "the model is
+  installed". A download or extraction that broke halfway therefore left a folder that was never retried
+  and never worked -- the user had to know to delete it by hand.
+"""
+
+from __future__ import annotations
+
+import io
+import json
+import tempfile
+import unittest
+import zipfile
+from pathlib import Path
+from unittest import mock
+
+from typing_extensions import Self
+
+from spineps.seg_enums import SpinepsPhase
+from spineps.utils import auto_download
+from spineps.utils.auto_download import (
+    canonical_model_key,
+    download_if_missing,
+    download_names,
+    download_weights,
+    instances,
+    labeling,
+    semantic,
+)
+
+
+class _FakeResponse:
+    """Minimal stand-in for ``urllib.request.urlopen``'s context manager (only ``info()`` is used)."""
+
+    def __init__(self, size: int = 1024) -> None:
+        self._size = size
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *exc) -> None:
+        return None
+
+    def info(self) -> dict:
+        return {"Content-Length": str(self._size)}
+
+
+def _model_zip(path, nested: bool = False, with_config: bool = True) -> None:
+    """Writes a minimal model archive (an inference_config.json plus one weights file) to a path or buffer."""
+    prefix = "wrapper/" if nested else ""
+    with zipfile.ZipFile(path, "w") as zf:
+        if with_config:
+            zf.writestr(prefix + "inference_config.json", json.dumps({"log_name": "test"}))
+        zf.writestr(prefix + "weights.ckpt", "not really a checkpoint")
+
+
+class Test_Every_Advertised_Model_Id_Resolves(unittest.TestCase):
+    def test_all_registry_ids_have_a_download_name(self):
+        for key in (*semantic, *instances, *labeling):
+            with self.subTest(model=key):
+                self.assertIn(canonical_model_key(key), download_names)
+
+    def test_semantic_aliases_map_to_the_canonical_id(self):
+        self.assertEqual(canonical_model_key("t2w_semantic"), "t2w")
+        self.assertEqual(canonical_model_key("ct_semantic"), "ct")
+        # the canonical ids and unrelated ids are untouched
+        self.assertEqual(canonical_model_key("t2w"), "t2w")
+        self.assertEqual(canonical_model_key("t2w_labeling"), "t2w_labeling")
+        self.assertEqual(canonical_model_key("something_semantic"), "something_semantic")
+
+    def test_alias_and_canonical_id_resolve_to_the_same_folder(self):
+        with tempfile.TemporaryDirectory() as td:
+            with (
+                mock.patch.object(auto_download, "get_mri_segmentor_models_dir", return_value=Path(td)),
+                mock.patch.object(auto_download, "download_weights") as download,
+            ):
+                canonical = download_if_missing("ct", semantic["ct"], SpinepsPhase.SEMANTIC)
+                alias = download_if_missing("ct_semantic", semantic["ct_semantic"], SpinepsPhase.SEMANTIC)
+            self.assertEqual(canonical, alias)
+            # ... and at the CT release version, not the generic semantic one
+            self.assertTrue(canonical.name.endswith(auto_download.current_highest_ct_version), canonical.name)
+            self.assertEqual(download.call_count, 2)
+
+    def test_unknown_id_raises_a_readable_error(self):
+        with tempfile.TemporaryDirectory() as td:
+            patched_dir = mock.patch.object(auto_download, "get_mri_segmentor_models_dir", return_value=Path(td))
+            with patched_dir, self.assertRaises(KeyError) as cm:
+                download_if_missing("not_a_model", "http://example.invalid/x.zip", SpinepsPhase.SEMANTIC)
+        self.assertIn("not_a_model", str(cm.exception))
+
+
+class Test_Download_Failure_Is_Loud(unittest.TestCase):
+    def test_unreachable_url_raises(self):
+        with tempfile.TemporaryDirectory() as td:
+            out = Path(td) / "T2w_semantic_v1.0.9"
+            with self.assertRaises(RuntimeError) as cm:
+                download_weights("http://localhost:1/does-not-exist.zip", out)
+            self.assertIn("download the weights manually", str(cm.exception))
+            self.assertFalse(out.exists(), "a failed download must not leave a model folder behind")
+            self.assertEqual([p.name for p in Path(td).iterdir()], [], "staging files were left behind")
+
+    def test_truncated_archive_leaves_nothing(self):
+        with tempfile.TemporaryDirectory() as td:
+            out = Path(td) / "T2w_semantic_v1.0.9"
+
+            def _retrieve(_url, filename, reporthook=None):  # noqa: ARG001
+                # a download that stopped halfway: the file exists but is not a readable zip
+                buffer = io.BytesIO()
+                _model_zip(buffer)
+                Path(filename).write_bytes(buffer.getvalue()[: len(buffer.getvalue()) // 2])
+                return filename, None
+
+            with (
+                mock.patch("urllib.request.urlopen", return_value=_FakeResponse()),
+                mock.patch("urllib.request.urlretrieve", side_effect=_retrieve),
+                self.assertRaises(RuntimeError) as cm,
+            ):
+                download_weights("http://example.invalid/t2w.zip", out)
+            self.assertIn("corrupt", str(cm.exception))
+            self.assertFalse(out.exists())
+            self.assertEqual([p.name for p in Path(td).iterdir()], [])
+
+    def test_archive_without_a_config_is_rejected(self):
+        with tempfile.TemporaryDirectory() as td:
+            out = Path(td) / "T2w_semantic_v1.0.9"
+
+            def _retrieve(_url, filename, reporthook=None):  # noqa: ARG001
+                _model_zip(Path(filename), with_config=False)
+                return filename, None
+
+            with (
+                mock.patch("urllib.request.urlopen", return_value=_FakeResponse()),
+                mock.patch("urllib.request.urlretrieve", side_effect=_retrieve),
+                self.assertRaises(RuntimeError) as cm,
+            ):
+                download_weights("http://example.invalid/t2w.zip", out)
+            self.assertIn("inference_config.json", str(cm.exception))
+            self.assertFalse(out.exists())
+
+
+class Test_Successful_Download_Layout(unittest.TestCase):
+    def _download(self, nested: bool) -> Path:
+        td = Path(tempfile.mkdtemp())
+        out = td / "T2w_semantic_v1.0.9"
+
+        def _retrieve(_url, filename, reporthook=None):  # noqa: ARG001
+            _model_zip(Path(filename), nested=nested)
+            return filename, None
+
+        with (
+            mock.patch("urllib.request.urlopen", return_value=_FakeResponse()),
+            mock.patch("urllib.request.urlretrieve", side_effect=_retrieve),
+        ):
+            download_weights("http://example.invalid/t2w.zip", out)
+        return out
+
+    def test_flat_archive(self):
+        out = self._download(nested=False)
+        self.assertTrue(out.joinpath("inference_config.json").is_file())
+        self.assertTrue(out.joinpath("weights.ckpt").is_file())
+        self.assertEqual([p.name for p in out.parent.iterdir()], [out.name], "staging files were left behind")
+
+    def test_nested_archive_is_unwrapped(self):
+        # The release archives wrap the model in one extra folder; it must not end up as out/wrapper/...
+        out = self._download(nested=True)
+        self.assertTrue(out.joinpath("inference_config.json").is_file())
+        self.assertTrue(out.joinpath("weights.ckpt").is_file())
+
+    def test_existing_folder_is_not_touched(self):
+        with tempfile.TemporaryDirectory() as td:
+            out = Path(td) / "T2w_semantic_v1.0.9"
+            out.mkdir()
+            out.joinpath("inference_config.json").write_text("{}")
+            with mock.patch("urllib.request.urlopen", side_effect=AssertionError("must not download")):
+                download_weights("http://example.invalid/t2w.zip", out)
+            self.assertEqual(out.joinpath("inference_config.json").read_text(), "{}")
+
+
+if __name__ == "__main__":
+    unittest.main()