From 2a216e9e7dc5d884ddf9a2b145a05b3f6db02e76 Mon Sep 17 00:00:00 2001 From: iback Date: Mon, 5 Oct 2026 08:21:01 +0000 Subject: [PATCH] fix: make the first-use weight download resolve every id and fail loudly Three defects in the automatic download: * The `*_semantic` ids are registered as available models -- they even show up in the "Available semantic models: [...]" error message -- but had no `download_names` entry, so `--model-semantic t2w_semantic` on a machine without the weights died with `KeyError: 't2w_semantic'`. `ct_semantic` additionally missed the CT version override and resolved to the v1.0.9 folder name while downloading the v2.0.0 archive. A single `canonical_model_key()` now maps the alias onto the id the weights are registered under. * `download_weights` caught every exception, logged "Download attempt failed" and returned normally, so `download_if_missing` handed back a path that does not exist and the real failure surfaced later as a config-not-found from the model loader. It now raises a RuntimeError naming the URL, the cause, and the manual-download fallback. * The archive was extracted straight into the final model folder, whose existence is the only "is this model installed?" check there is. An extraction interrupted halfway (Ctrl-C, full disk) left a partial folder that was never retried and never worked, and a truncated download left a stray .zip plus a raw BadZipFile. Download and extraction now happen in a unique staging folder that is moved into place only after the contents are verified to hold an inference_config.json, and is removed on any failure. The unique staging name also stops concurrent SPINEPS processes (e.g. a job array) from overwriting each other's partial download. Co-Authored-By: Claude Opus 5 --- spineps/utils/auto_download.py | 145 +++++++++++++++++------ unit_tests/test_auto_download.py | 190 +++++++++++++++++++++++++++++++ 2 files changed, 299 insertions(+), 36 deletions(-) create mode 100644 unit_tests/test_auto_download.py diff --git a/spineps/utils/auto_download.py b/spineps/utils/auto_download.py index 7beaa50..0d52bc2 100644 --- a/spineps/utils/auto_download.py +++ b/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 diff --git a/unit_tests/test_auto_download.py b/unit_tests/test_auto_download.py new file mode 100644 index 0000000..af63a81 --- /dev/null +++ b/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()