From 76387b0294e3f091129ec0a95c038030ea12e2c9 Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Sat, 7 Nov 2026 17:27:48 -0800 Subject: [PATCH] Give every kept scan a version from each model the site offers The idle-time second-opinion process ran each kept scan through ePAI, Atlas-Net and LesionSegmenter only. Add SuPreM, R-Super-Pancreas-Kidney and R-Super-Merlin, so a scan ends up with a version from every model in the picker. All three write combined_labels.nii.gz in the viewer's label scheme, so the existing saver works unchanged. Their weights are single files, not nnU-Net folders, so provenance gets a single-file rule (checkpoint and class-list hashes, fixed settings). The R-Super releases have no label map, so label_support says known: false instead of guessing. A test now fails if the default list and the model picker drift apart. --- flask-server/.env.example | 2 +- .../deploy/USER_DATASET_SECOND_OPINIONS.md | 11 +++- .../scripts/user_dataset_second_opinions.py | 2 +- flask-server/services/model_provenance.py | 49 ++++++++++++++ flask-server/services/user_dataset_models.py | 5 +- .../tests/unit/test_model_provenance.py | 66 ++++++++++++++++++- .../tests/unit/test_user_dataset_models.py | 48 +++++++++++--- 7 files changed, 166 insertions(+), 17 deletions(-) diff --git a/flask-server/.env.example b/flask-server/.env.example index d85c87f8..c1964276 100644 --- a/flask-server/.env.example +++ b/flask-server/.env.example @@ -258,6 +258,6 @@ BODYMAPS_LOCAL_AUTO_SIGNIN=false # scripts/user_dataset_second_opinions.py runs every kept scan through the other models, only while the site is # idle and a GPU worker is free, and never on this host's GPU. All settings below are optional; see # deploy/USER_DATASET_SECOND_OPINIONS.md. -# SECOND_OPINIONS_MODELS=ePAI,Atlas-Net,LesionSegmenter +# SECOND_OPINIONS_MODELS=ePAI,Atlas-Net,LesionSegmenter,SuPreM,R-Super-Pancreas-Kidney,R-Super-Merlin # SECOND_OPINIONS_MAX_RUN_SECONDS=1800 # SECOND_OPINIONS_DAILY_CAP=100 diff --git a/flask-server/deploy/USER_DATASET_SECOND_OPINIONS.md b/flask-server/deploy/USER_DATASET_SECOND_OPINIONS.md index bbfb29ca..620d9afa 100644 --- a/flask-server/deploy/USER_DATASET_SECOND_OPINIONS.md +++ b/flask-server/deploy/USER_DATASET_SECOND_OPINIONS.md @@ -2,8 +2,8 @@ Every scan kept in `UserData` carries the mask of the model the user ran. This process runs the same stored CT through the **other** models and saves each -result beside the original, so one scan has an ePAI version, an Atlas-Net -version and a LesionSegmenter version. +result beside the original, so one scan has a version from every model the site offers +(ePAI, Atlas-Net, LesionSegmenter, SuPreM, R-Super-Pancreas-Kidney, R-Super-Merlin). It is **low priority**: it only works while nobody is using the site and a GPU worker is idle, it never uses bdmap1's GPU, and it stops the moment a user job @@ -112,7 +112,7 @@ failure it gives up (listed by `--status`); delete its entry in | Variable | Default | Meaning | | --- | --- | --- | -| `SECOND_OPINIONS_MODELS` | `ePAI,Atlas-Net,LesionSegmenter` | models every scan should have | +| `SECOND_OPINIONS_MODELS` | `ePAI,Atlas-Net,LesionSegmenter,SuPreM,R-Super-Pancreas-Kidney,R-Super-Merlin` | models every scan should have | | `SECOND_OPINIONS_BASE_URL` | `http://127.0.0.1:8000` | the website, for `/api/health` | | `SECOND_OPINIONS_IDLE_CHECKS` / `_CHECK_INTERVAL` | 3 / 20 s | consecutive idle samples needed | | `SECOND_OPINIONS_POLL_SECONDS` | 3 | how often to look for users while a model runs | @@ -124,6 +124,11 @@ failure it gives up (listed by `--status`); delete its entry in ## Notes +* **Adding a model to the list** needs it to run on the workers, to write `combined_labels.nii.gz` in the viewer's label + scheme (all six do), and a checkpoint rule in `services/model_provenance.py`. The R-Super releases have no label map + in `MODEL_TO_VIEWER`, so their `label_support` says `known: false` rather than guess which labels they can produce. +* A bigger model list means more GPU runs per scan; `SECOND_OPINIONS_DAILY_CAP` (100 per rolling 24 hours) is what paces + the backlog of scans already kept, so raise it only when the workers are known to be idle. * The stored scans are **not de-identified** (each case keeps `user_id` and `session_id`), so these versions are internal until that is decided. * The workers must hold the same model environments and weights as the website diff --git a/flask-server/scripts/user_dataset_second_opinions.py b/flask-server/scripts/user_dataset_second_opinions.py index 8ac4a1af..a3376caa 100644 --- a/flask-server/scripts/user_dataset_second_opinions.py +++ b/flask-server/scripts/user_dataset_second_opinions.py @@ -46,7 +46,7 @@ def main(argv=None) -> int: mode.add_argument("--once", action="store_true") mode.add_argument("--loop", action="store_true") parser.add_argument("--limit", type=int, default=None, help="stop after this many finished tasks") - parser.add_argument("--models", default=None, help="comma-separated models (default: ePAI,Atlas-Net,LesionSegmenter)") + parser.add_argument("--models", default=None, help="comma-separated models (default: every model the site offers)") parser.add_argument("--root", default=None, help="UserData folder (default: USER_DATASET_PATH)") parser.add_argument("--json", action="store_true", help="print --status/--dry-run as JSON") args = parser.parse_args(argv) diff --git a/flask-server/services/model_provenance.py b/flask-server/services/model_provenance.py index 5f1ac65e..c167078f 100644 --- a/flask-server/services/model_provenance.py +++ b/flask-server/services/model_provenance.py @@ -37,6 +37,22 @@ "LesionSegmenter": ("LESIONSEG_CKPT_PATH", "/home/visitor/lesionsegmenter/model/nnUNet_results/Dataset_LesionSegmenter/nnUNetTrainer__nnUNetPlans__3d_fullres"), } +# Models whose weights are one file, not an nnU-Net folder: model -> (checkpoint env var, default path, +# class-list env var or None, default class list or None). Same names and defaults as the runners in +# services/auto_segmentor.py (SuPreM, RSUPER_VARIANTS); a test keeps them equal. +MODEL_CHECKPOINT_FILE = { + "SuPreM": ("SUPREM_CHECKPOINT_PATH", + "/home/visitor/suprem_native/workspace/SuPreM/pretrained_checkpoints/supervised_suprem_unet_2100.pth", + None, None), + "R-Super-Pancreas-Kidney": ("RSUPER_PANKID_CHECKPOINT_PATH", + "/home/visitor/rsuper/R-SuperPancreasKidney/R_Super_Kidney_Pancreas_UCSF_Atlas3/fold_0_latest.pth", + "RSUPER_PANKID_CLASS_LIST", + "/home/visitor/rsuper/R-SuperPancreasKidney/labels_rsuper_mask.yaml"), + "R-Super-Merlin": ("RSUPER_MERLIN_CHECKPOINT_PATH", + "/home/visitor/rsuper/Merlin-Cancer-Super/checkpoint/fold_0_latest.pth", + "RSUPER_MERLIN_CLASS_LIST", + "/home/visitor/rsuper/Merlin-Cancer-Super/checkpoint/label_names.yaml"), +} DEFAULT_CHECKPOINT_NAME = "checkpoint_final.pth" # Files of an nnU-Net model folder that define the model (weights + preprocessing/architecture plan). _PLAN_FILES = ("plans.json", "dataset.json") @@ -83,10 +99,36 @@ def checkpoint_name(model: str, env=None) -> str: return DEFAULT_CHECKPOINT_NAME +def _file_checkpoint_fingerprint(model: str, env) -> dict: + """A model whose weights are one file: hash it, and its class list when it has one (that list + decides which label each output channel gets, so it is part of what made the labels).""" + ckpt_var, ckpt_default, list_var, list_default = MODEL_CHECKPOINT_FILE[model] + path = env.get(ckpt_var, ckpt_default) + out: dict = {"path": path, "hashed": False} + if not os.path.isfile(path): + out["reason"] = "checkpoint file not found on this host" + return out + try: + files = {os.path.basename(path): _cached_sha256(path)} + if list_var: + class_list = env.get(list_var, list_default) + if os.path.isfile(class_list): + files[os.path.basename(class_list)] = _cached_sha256(class_list) + else: + out["class_list_missing"] = class_list + except OSError as e: + out["reason"] = f"unreadable:{type(e).__name__}" + return out + out.update(hashed=True, files=files) + return out + + def checkpoint_fingerprint(model: str, env=None) -> dict: """Identify the weights this process would load for `model`. Always returns a dict; `hashed` says whether the files could be read here.""" env = os.environ if env is None else env + if model in MODEL_CHECKPOINT_FILE: + return _file_checkpoint_fingerprint(model, env) rule = MODEL_CHECKPOINT.get(model) if rule is None: return {"hashed": False, "reason": "no checkpoint rule for this model"} @@ -148,6 +190,13 @@ def inference_settings(model: str, env=None) -> Optional[dict]: # the recall-oriented operating point applied when the labels are exported "lesion_threshold": float(env.get("LESIONSEG_LESION_THRESHOLD", "0.05") or "0.05"), "lesion_channel": int(env.get("LESIONSEG_LESION_CHANNEL", "40") or "40")} + if model == "SuPreM": + # Fixed on the runner's command line: the U-Net backbone with the SuPreM weights. + return {"backbone": "unet", "suprem": True, "checkpoint_name": os.path.basename(MODEL_CHECKPOINT_FILE[model][1])} + if model in ("R-Super-Pancreas-Kidney", "R-Super-Merlin"): + # RSUPER_VARIANTS: the Pancreas-Kidney release masks lesions with its own organ masks, Merlin does not. + return {"checkpoint_name": os.path.basename(MODEL_CHECKPOINT_FILE[model][1]), + "organ_mask_on_lesion": model == "R-Super-Pancreas-Kidney"} return None diff --git a/flask-server/services/user_dataset_models.py b/flask-server/services/user_dataset_models.py index d1888925..38376413 100644 --- a/flask-server/services/user_dataset_models.py +++ b/flask-server/services/user_dataset_models.py @@ -3,7 +3,7 @@ A scan is kept in UserData with the mask of the model the user ran. This module (driven by scripts/user_dataset_second_opinions.py, a separate process from the website) runs the same stored CT through the remaining models and saves each result beside the original, so one scan carries an ePAI version, an -Atlas-Net version and a LesionSegmenter version: +Atlas-Net version, a LesionSegmenter version and so on for every model the site offers: mask_only/USER_00000017/ combined_labels.nii.gz, segmentations/, metadata.json the model the user ran (unchanged) @@ -45,7 +45,8 @@ from services.atomic_write import atomic_destination SCHEMA = 1 -DEFAULT_MODELS = ("ePAI", "Atlas-Net", "LesionSegmenter") +# Every model the site offers, so each kept scan ends up with one version from each. +DEFAULT_MODELS = ("ePAI", "Atlas-Net", "LesionSegmenter", "SuPreM", "R-Super-Pancreas-Kidney", "R-Super-Merlin") BY_MODEL = "by_model" STATE_FILE = "second_opinions_state.json" AUDIT_FILE = "second_opinions.jsonl" diff --git a/flask-server/tests/unit/test_model_provenance.py b/flask-server/tests/unit/test_model_provenance.py index 3e348a05..c9b9aee7 100644 --- a/flask-server/tests/unit/test_model_provenance.py +++ b/flask-server/tests/unit/test_model_provenance.py @@ -76,7 +76,7 @@ def test_settings_follow_the_environment(): "checkpoint_name": "checkpoint_best.pth", "tensorrt_plan": "/p/plan.trt"} s = mp.inference_settings("LesionSegmenter", env={"LESIONSEG_LESION_THRESHOLD": "0.2", "LESIONSEG_LESION_CHANNEL": "41"}) assert s["lesion_threshold"] == 0.2 and s["lesion_channel"] == 41 - assert mp.inference_settings("SuPreM", env={}) is None + assert mp.inference_settings("OpenVAE", env={}) is None # no rule for a model that is not a segmenter here # ----------------------------- checkpoint identity ----------------------------- @@ -133,7 +133,7 @@ def test_checkpoint_that_cannot_be_read_says_why(tmp_path): assert fp["hashed"] is False and "checkpoint_final.pth not found" in fp["reason"] fp = mp.checkpoint_fingerprint("ePAI", env={"EPAI_CKPT_PATH": str(empty), "EPAI_CHECKPOINT_NAME": "best.pth"}) assert "best.pth not found" in fp["reason"] # honours the configured file name - assert mp.checkpoint_fingerprint("SuPreM", env={}) == {"hashed": False, "reason": "no checkpoint rule for this model"} + assert mp.checkpoint_fingerprint("OpenVAE", env={}) == {"hashed": False, "reason": "no checkpoint rule for this model"} def test_epai_tensorrt_engine_is_fingerprinted_when_configured(tmp_path): @@ -220,3 +220,65 @@ def boom(*a, **k): out = mp.describe("LesionSegmenter", real_names(), (22,)) assert out["provenance"] == {"error": "RuntimeError: disk on fire"} assert out["label_support"]["known"] is True + + +# ----------------------------- models whose weights are a single file ----------------------------- +FILE_MODELS = ["SuPreM", "R-Super-Pancreas-Kidney", "R-Super-Merlin"] + + +def test_single_file_rules_match_the_runners(): + """Provenance repeats the runners' environment variables and defaults; a stale copy would hash the wrong file.""" + from services.auto_segmentor import RSUPER_VARIANTS + for model in ("R-Super-Pancreas-Kidney", "R-Super-Merlin"): + v = RSUPER_VARIANTS[model] + ckpt_var, ckpt_default, list_var, list_default = mp.MODEL_CHECKPOINT_FILE[model] + assert (ckpt_var, ckpt_default) == (v["checkpoint_env"], v["checkpoint"]) + assert (list_var, list_default) == (v["class_list_env"], v["class_list"]) + assert mp.inference_settings(model, env={})["organ_mask_on_lesion"] is v["organ_mask_on_lesion"] + var, default, _, _ = mp.MODEL_CHECKPOINT_FILE["SuPreM"] + body = _function_source(SEGMENTOR, "_run_suprem_inference") + assert re.search(rf'"{var}",\s*"{re.escape(default)}"', body), "_run_suprem_inference no longer reads this checkpoint" + assert '"--backbone", "unet"' in body and '"--suprem"' in body # the fixed settings recorded for SuPreM + + +@pytest.mark.parametrize("model", FILE_MODELS) +def test_a_single_file_checkpoint_is_hashed_with_its_class_list(model, tmp_path): + ckpt_var, _, list_var, _ = mp.MODEL_CHECKPOINT_FILE[model] + weights = tmp_path / "fold_0_latest.pth"; weights.write_bytes(b"weights-" + model.encode()) + env = {ckpt_var: str(weights)} + if list_var: + classes = tmp_path / "labels.yaml"; classes.write_bytes(b"liver: 1\n") + env[list_var] = str(classes) + fp = mp.checkpoint_fingerprint(model, env=env) + assert fp["hashed"] is True and fp["path"] == str(weights) + assert fp["files"]["fold_0_latest.pth"]["sha256"] == hashlib.sha256(b"weights-" + model.encode()).hexdigest() + assert fp["files"]["fold_0_latest.pth"]["size_bytes"] == len(b"weights-" + model.encode()) + if list_var: + assert fp["files"]["labels.yaml"]["sha256"] == hashlib.sha256(b"liver: 1\n").hexdigest() + + +@pytest.mark.parametrize("model", FILE_MODELS) +def test_a_missing_single_file_checkpoint_says_why(model, tmp_path): + ckpt_var = mp.MODEL_CHECKPOINT_FILE[model][0] + fp = mp.checkpoint_fingerprint(model, env={ckpt_var: str(tmp_path / "absent.pth")}) + assert fp["hashed"] is False and fp["reason"] == "checkpoint file not found on this host" + + +def test_a_missing_class_list_is_noted_but_the_weights_are_still_hashed(tmp_path): + weights = tmp_path / "w.pth"; weights.write_bytes(b"w") + fp = mp.checkpoint_fingerprint("R-Super-Merlin", env={"RSUPER_MERLIN_CHECKPOINT_PATH": str(weights), + "RSUPER_MERLIN_CLASS_LIST": str(tmp_path / "gone.yaml")}) + assert fp["hashed"] is True and list(fp["files"]) == ["w.pth"] and fp["class_list_missing"].endswith("gone.yaml") + + +def test_settings_for_the_single_file_models_are_fixed_by_the_runners(): + assert mp.inference_settings("SuPreM", env={}) == {"backbone": "unet", "suprem": True, + "checkpoint_name": "supervised_suprem_unet_2100.pth"} + assert mp.inference_settings("R-Super-Pancreas-Kidney", env={})["organ_mask_on_lesion"] is True + assert mp.inference_settings("R-Super-Merlin", env={})["organ_mask_on_lesion"] is False + + +def test_label_support_for_the_r_super_releases_says_it_is_unknown_rather_than_guess(): + for model in ("R-Super-Pancreas-Kidney", "R-Super-Merlin"): + assert mp.label_support(model, real_names())["known"] is False + assert mp.label_support("SuPreM", real_names())["known"] is True diff --git a/flask-server/tests/unit/test_user_dataset_models.py b/flask-server/tests/unit/test_user_dataset_models.py index 0ee3b06b..c567b688 100644 --- a/flask-server/tests/unit/test_user_dataset_models.py +++ b/flask-server/tests/unit/test_user_dataset_models.py @@ -8,6 +8,7 @@ import http.server import importlib import json +import re import os import sys import threading @@ -112,6 +113,9 @@ def deps(self): local_hostname=lambda: "bdmap1", now=self.now, sleep=self.sleep) +THREE = ("ePAI", "Atlas-Net", "LesionSegmenter") # the loop tests use three models; DEFAULT_MODELS has its own tests + + @pytest.fixture() def world(tmp_path): root = tmp_path / "UserData" @@ -119,7 +123,7 @@ def world(tmp_path): sessions.mkdir() return types.SimpleNamespace(root=str(root), sessions=str(sessions), rec=Recorder(str(root)), cfg=M.Config(idle_checks=2, check_interval=0, poll_seconds=0.02, max_run_seconds=30, - loop_sleep=0, idle_sleep=0)) + loop_sleep=0, idle_sleep=0, models=THREE)) def run(world, **kw): @@ -135,17 +139,17 @@ def test_pending_tasks_skip_the_users_own_model_finished_versions_and_backoff(wo make_case(world.root, "USER_00000001", "ePAI") make_case(world.root, "USER_00000002", "LesionSegmenter") state = M.load_state(world.root) - tasks = M.pending_tasks(world.root, M.DEFAULT_MODELS, state, 1000.0) + tasks = M.pending_tasks(world.root, THREE, state, 1000.0) assert [(t.case_id, t.model) for t in tasks] == [ ("USER_00000001", "Atlas-Net"), ("USER_00000001", "LesionSegmenter"), ("USER_00000002", "ePAI"), ("USER_00000002", "Atlas-Net")] assert all(t.model != t.origin for t in tasks) M.note_result(state, tasks[0], "failed", 1000.0, "boom") - assert tasks[0].key not in {t.key for t in M.pending_tasks(world.root, M.DEFAULT_MODELS, state, 1001.0)} - assert tasks[0].key in {t.key for t in M.pending_tasks(world.root, M.DEFAULT_MODELS, state, 1000.0 + 3601)} + assert tasks[0].key not in {t.key for t in M.pending_tasks(world.root, THREE, state, 1001.0)} + assert tasks[0].key in {t.key for t in M.pending_tasks(world.root, THREE, state, 1000.0 + 3601)} # a model the user ran that is not in the configured set still gets every configured model make_case(world.root, "USER_00000003", "SuPreM") - assert len([t for t in M.pending_tasks(world.root, M.DEFAULT_MODELS, M.load_state(world.root), 0) if t.case_id == "USER_00000003"]) == 3 + assert len([t for t in M.pending_tasks(world.root, THREE, M.load_state(world.root), 0) if t.case_id == "USER_00000003"]) == 3 def test_failures_back_off_then_give_up_and_preemption_is_not_a_failure(): @@ -371,7 +375,7 @@ def test_the_whole_backlog_is_worked_through_once_and_never_repeated(world): assert counts["ok"] == 6 and counts["failed"] == 0 and len(world.rec.calls) == 6 for i, model in enumerate(["ePAI", "LesionSegmenter", "Atlas-Net"], 1): meta = json.loads(Path(world.root, "mask_only", f"USER_0000000{i}", "metadata.json").read_text()) - assert sorted(v["model"] for v in meta["model_versions"]) == sorted(M.DEFAULT_MODELS) + assert sorted(v["model"] for v in meta["model_versions"]) == sorted(THREE) assert meta["model_versions"][0] == {"model": model, "role": "user_run", "mask": ".", "pipeline": {"segmentation_input": "original"}} assert os.listdir(world.sessions) == [] again = run(world, once=True) @@ -519,6 +523,34 @@ def test_the_safety_relies_on_these_behaviours_of_the_runner_and_the_worker_modu # a worker run happens only for a session folder under .../tmp/ with an id-shaped name: ours qualifies session = "/home/visitor/PanTS-Viewer/tmp/so-USER_00000017-LesionSegmenter-3fa9c2d1" assert gpu_workers._is_session_path(session) - assert M._SAFE_MODEL.match("Atlas-Net") and all(gpu_workers._SESSION_ID.match(f"so-USER_00000017-{m}-3fa9c2d1") for m in M.DEFAULT_MODELS) + assert all(M._SAFE_MODEL.match(m) for m in M.DEFAULT_MODELS) and all(gpu_workers._SESSION_ID.match(f"so-USER_00000017-{m}-3fa9c2d1") for m in M.DEFAULT_MODELS) # every model run here must be one the runner knows, and the user dataset must record the same names - assert all(f"model == '{m}'" in seg for m in M.DEFAULT_MODELS) + # (the R-Super releases are matched as a group: `model in RSUPER_VARIANTS`) + assert all(f"model == '{m}'" in seg or m in auto_segmentor.RSUPER_VARIANTS for m in M.DEFAULT_MODELS) + + +# ------------------------------------------------------------------ the default model list +def test_the_default_models_are_every_model_the_site_offers(): + """A model added to the picker must also be added here, or scans kept afterwards would never get its version.""" + page = (Path(__file__).resolve().parents[3] / "PanTS-Demo" / "src" / "routes" / "UploadPage.tsx").read_text(encoding="utf-8") + block = page[page.index("const MODEL_OPTIONS"): page.index("const MODEL_OPTIONS") + 12000] + offered = [m for m in re.findall(r'^ id: "([^"]+)"', block, re.M) if m != "None"] + assert offered and set(offered) == set(M.DEFAULT_MODELS), (offered, M.DEFAULT_MODELS) + assert len(set(M.DEFAULT_MODELS)) == len(M.DEFAULT_MODELS) + assert M.Config.from_env({}).models == M.DEFAULT_MODELS + + +def test_every_default_model_has_a_checkpoint_rule_and_settings(): + from services import model_provenance as mp + for model in M.DEFAULT_MODELS: + assert model in mp.MODEL_CHECKPOINT or model in mp.MODEL_CHECKPOINT_FILE, model + assert mp.inference_settings(model, env={}) is not None, model + + +def test_a_kept_scan_is_owed_a_version_from_every_other_model(world): + make_case(world.root, "USER_00000001", "SuPreM") + make_case(world.root, "USER_00000002", "R-Super-Merlin") + tasks = M.pending_tasks(world.root, M.DEFAULT_MODELS, M.load_state(world.root), 0) + by_case = {c: sorted(t.model for t in tasks if t.case_id == c) for c in ("USER_00000001", "USER_00000002")} + assert by_case["USER_00000001"] == sorted(set(M.DEFAULT_MODELS) - {"SuPreM"}) + assert by_case["USER_00000002"] == sorted(set(M.DEFAULT_MODELS) - {"R-Super-Merlin"})