diff --git a/flask-server/scripts/lesionseg_predict.py b/flask-server/scripts/lesionseg_predict.py index 346b7e9..c4cb4d0 100644 --- a/flask-server/scripts/lesionseg_predict.py +++ b/flask-server/scripts/lesionseg_predict.py @@ -4,59 +4,136 @@ shape that mirrors nnUNetv2_predict_from_modelfolder (-i/-o/-m/-step_size/--disable_tta), so it's a near drop-in replacement for the CLI call in services/auto_segmentor.py. -Why this exists instead of calling the CLI directly: the CLI has no flag for GPU-accelerated -export resampling. Swapping ConfigurationManager.resampling_fn_probabilities for -resample_torch_fornnunet cuts that stage from ~26-42s to ~1-6s with no accuracy cost -- -validated end-to-end this session, lesion Dice 0.998-1.0 against the CPU-resample output -across every test case (isolated specifically from every other change, see PR #105/#106 -history). The CLI's export step has no hook to apply this from the outside, so a thin -wrapper around nnUNetPredictor directly is the only way to get the win without either (a) -patching the vendored nnU-Net install in each conda env, which every OTHER model here also -depends on and would be a much larger blast radius, or (b) a persistent/warm predictor -service, which is a separate, bigger architecture change (not this). - -Still a fresh subprocess per invocation, same cold-start cost as the plain CLI (~19s model -load) -- this does NOT include the warm-predictor optimization. +Two execution paths: + +1. WARM (fast). If LESIONSEG_WARM_URL is set and that service is healthy, the work is + handed to the persistent predictor (scripts/lesionseg_warm_server.py), which holds + the model in GPU memory across requests. This skips the ~19s of interpreter start, + torch import, checkpoint load and cuDNN autotune that the cold path pays on EVERY + request. Everything above this line is stdlib-only and the heavy imports live inside + the cold path -- importing torch here would give back several seconds of the very + cost this path exists to avoid. + +2. COLD (fallback). Exactly the previous behaviour: load the model in-process and run. + Used when no warm URL is configured, the service is down, or it is running a + different configuration than the caller asked for. The fallback is deliberate: the + website must keep working if the warm service dies, just more slowly. + +Why this wrapper exists at all instead of the plain CLI: the CLI has no flag for +GPU-accelerated export resampling. Swapping ConfigurationManager.resampling_fn_probabilities +for resample_torch_fornnunet cuts that stage from ~26-42s to ~1-6s with no accuracy cost +(lesion Dice 0.998-1.0 vs the CPU-resample output). The CLI's export step has no hook to +apply this from outside, and patching the vendored nnU-Net install would hit every other +model in this deployment. """ import argparse +import json import os import sys -from functools import partial - -import torch -from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor -from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO -from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor -from nnunetv2.inference.export_prediction import ( - convert_predicted_logits_to_segmentation_with_correct_shape as convert_logits, -) -import nnunetv2.utilities.plans_handling.plans_handler as plans_handler -from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet - - -def _patch_gpu_export_resample(): - def _resampling_fn_probabilities(self): - return partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False) - plans_handler.ConfigurationManager.resampling_fn_probabilities = property( - _resampling_fn_probabilities +import urllib.error +import urllib.request + + +def _warm_predict(input_dir, output_dir, step_size, disable_tta, url, timeout): + """Try the persistent predictor. Returns True if it produced the output. + + Any failure returns False so the caller falls back to cold inference -- a warm + service that is down, busy, or differently configured must degrade to a slow + correct answer, never to a wrong or missing one. + """ + base = url.rstrip("/") + try: + with urllib.request.urlopen(f"{base}/health", timeout=5) as r: + health = json.loads(r.read()) + except Exception as e: + print(f"[warm] health check failed ({e}); using cold path", flush=True) + return False + + # The server also enforces this and answers 409, but checking here keeps the + # reason in this process's log where the failure is being handled. + if abs(float(health.get("step_size", -1)) - float(step_size)) > 1e-9 \ + or bool(health.get("disable_tta")) != bool(disable_tta): + print(f"[warm] config mismatch (server step={health.get('step_size')} " + f"disable_tta={health.get('disable_tta')}; requested step={step_size} " + f"disable_tta={disable_tta}); using cold path", flush=True) + return False + + # The service only accepts locations as relative paths under the root it was + # started with, and builds the absolute path itself. If our directories are not + # under that root it cannot serve them -- fall back rather than trying to coax it. + root = health.get("allowed_root") + if not root: + print("[warm] server did not report allowed_root; using cold path", flush=True) + return False + try: + rel_in = os.path.relpath(os.path.realpath(input_dir), root) + rel_out = os.path.relpath(os.path.realpath(output_dir), root) + except ValueError as e: # e.g. different drives on Windows + print(f"[warm] cannot relativise paths against {root} ({e}); using cold path", flush=True) + return False + if any(r == os.pardir or r.startswith(os.pardir + os.sep) or os.path.isabs(r) + for r in (rel_in, rel_out)): + print(f"[warm] paths are outside the server root {root}; using cold path", flush=True) + return False + + body = json.dumps({ + "input_rel": rel_in.replace(os.sep, "/"), + "output_rel": rel_out.replace(os.sep, "/"), + "step_size": float(step_size), + "disable_tta": bool(disable_tta), + }).encode() + req = urllib.request.Request(f"{base}/predict", data=body, + headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=timeout) as r: + resp = json.loads(r.read()) + except urllib.error.HTTPError as e: + detail = e.read().decode(errors="replace")[:400] + print(f"[warm] server returned {e.code}: {detail}; using cold path", flush=True) + return False + except Exception as e: + print(f"[warm] request failed ({e}); using cold path", flush=True) + return False + + for item in resp.get("results", []): + out = item.get("output_path") + if not out or not os.path.exists(out): + print(f"[warm] server reported success but {out} is missing; using cold path", + flush=True) + return False + print(f"done with {item.get('case_id')} ({item.get('seconds')}s, warm)", flush=True) + if not resp.get("results"): + print("[warm] server returned no results; using cold path", flush=True) + return False + print(f"[warm] total {resp.get('total_seconds')}s", flush=True) + return True + + +def _cold_predict(args): + """In-process inference. Heavy imports are local so the warm path never pays them.""" + from functools import partial + + import torch + from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor + from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO + from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor + from nnunetv2.inference.export_prediction import ( + convert_predicted_logits_to_segmentation_with_correct_shape as convert_logits, ) - - -def main(): - ap = argparse.ArgumentParser(description=__doc__) - ap.add_argument("-i", dest="input_dir", required=True, help="Folder of {case}_0000.nii[.gz] files") - ap.add_argument("-o", dest="output_dir", required=True) - ap.add_argument("-m", dest="model_dir", required=True, help="nnU-Net model folder (dataset.json/plans.json/fold_*)") - ap.add_argument("-step_size", type=float, default=0.9) - ap.add_argument("--disable_tta", action="store_true") - ap.add_argument("-chk", dest="checkpoint_name", default="checkpoint_final.pth") - args = ap.parse_args() + import nnunetv2.utilities.plans_handling.plans_handler as plans_handler + from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet if not torch.cuda.is_available(): print("ERROR: CUDA not available", file=sys.stderr) sys.exit(1) - os.makedirs(args.output_dir, exist_ok=True) + def _patch_gpu_export_resample(): + def _resampling_fn_probabilities(self): + return partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False) + plans_handler.ConfigurationManager.resampling_fn_probabilities = property( + _resampling_fn_probabilities + ) + _patch_gpu_export_resample() predictor = nnUNetPredictor( @@ -102,5 +179,27 @@ def main(): print(f"done with {case_id}", flush=True) +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("-i", dest="input_dir", required=True, help="Folder of {case}_0000.nii[.gz] files") + ap.add_argument("-o", dest="output_dir", required=True) + ap.add_argument("-m", dest="model_dir", required=True, help="nnU-Net model folder (dataset.json/plans.json/fold_*)") + ap.add_argument("-step_size", type=float, default=0.9) + ap.add_argument("--disable_tta", action="store_true") + ap.add_argument("-chk", dest="checkpoint_name", default="checkpoint_final.pth") + args = ap.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + warm_url = os.getenv("LESIONSEG_WARM_URL", "").strip() + if warm_url: + timeout = int(os.getenv("LESIONSEG_WARM_TIMEOUT", "3600")) + if _warm_predict(args.input_dir, args.output_dir, args.step_size, + args.disable_tta, warm_url, timeout): + return + + _cold_predict(args) + + if __name__ == "__main__": main() diff --git a/flask-server/scripts/lesionseg_warm_server.py b/flask-server/scripts/lesionseg_warm_server.py new file mode 100644 index 0000000..2a0448f --- /dev/null +++ b/flask-server/scripts/lesionseg_warm_server.py @@ -0,0 +1,366 @@ +#!/usr/bin/env python3 +""" +Warm/persistent LesionSegmenter inference service. + +Loads the nnU-Net model ONCE at startup and keeps it resident, eliminating the +~19s cold start (interpreter + torch import + checkpoint load + cuDNN autotune) +that the subprocess-per-request path pays on EVERY request. Measured on an idle +node: model load 5.6s once, then steady-state 23.3s/request (read 0.4 / +preprocess 7.6 / infer 9.6 / export 5.6). + +Deliberately stdlib-only for the HTTP layer (no Flask): this runs inside the +GPU-side conda env, which is a *different* env from the web app's, and adding a +web framework dependency there is unnecessary for a single localhost endpoint. + +Binds 127.0.0.1 only. The Flask app (via scripts/lesionseg_predict.py) is the +sole intended caller; nothing here is exposed to the internet. + +Inference settings mirror the validated production configuration exactly: + - tile_step_size 0.7, mirroring/TTA disabled -- validated against real ground + truth on 229 held-out lesion-positive cases; statistically indistinguishable + from the slow 8xTTA/step-0.5 default (paired Dice delta not significant, + 95% CI crosses zero) at ~5x the speed. + - GPU export resampling -- validated lesion Dice 0.998-1.0 vs CPU resample. + +A client that asks for settings this server was not started with is REFUSED +(HTTP 409) rather than silently served the wrong configuration. The client then +falls back to cold inference with the settings it actually asked for. Silently +serving step 0.5 to a caller that requested 0.7 (or vice versa) would be an +invisible accuracy change in a cancer-detection tool, so it is made loud. + +Usage: + MODEL=/home/visitor/model PORT=8765 python lesionseg_warm_server.py +""" +import json +import os +import re +import threading +import time +import traceback +from functools import partial +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import torch +from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor +from nnunetv2.imageio.simpleitk_reader_writer import SimpleITKIO +from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor +from nnunetv2.inference.export_prediction import ( + convert_predicted_logits_to_segmentation_with_correct_shape as convert_logits, +) +import nnunetv2.utilities.plans_handling.plans_handler as plans_handler +from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet + +MODEL = os.environ.get("MODEL", os.environ.get("LESIONSEG_CKPT_PATH", "/home/visitor/model")) +PORT = int(os.environ.get("PORT", "8765")) +CHECKPOINT = os.environ.get("CHECKPOINT_NAME", "checkpoint_final.pth") +STEP_SIZE = float(os.environ.get("STEP_SIZE", "0.7")) +DISABLE_TTA = os.environ.get("DISABLE_TTA", "1").strip().lower() in {"1", "true", "yes", "on"} +# Requests are serialized; this bounds how long a queued request waits for the GPU. +LOCK_TIMEOUT = int(os.environ.get("LOCK_TIMEOUT", "1800")) + + +def _default_root(): + """Directory tree the service is willing to read from and write to. + + Defaults to the Flask app's sessions directory (Constants.SESSIONS_DIR_NAME, + i.e. $SESSIONS_DIR_PATH or /sessions), which is where + auto_segmentor.py builds the per-session input/output dirs this service is + asked to process. + """ + env = os.environ.get("LESIONSEG_ROOT") or os.environ.get("SESSIONS_DIR_PATH") + if env: + return os.path.realpath(env) + return os.path.realpath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), os.pardir, "sessions") + ) + + +ALLOWED_ROOT = _default_root() + + +# A request may only name a location as a RELATIVE path under ALLOWED_ROOT, and only +# using this restricted alphabet. No absolute paths, no "..", no empty components, so +# there is no way to express a location outside the root -- the absolute path is always +# constructed here, server-side, and never taken from the request. +# +# Accepting an absolute path and then checking it is weaker and harder to verify: +# binding to 127.0.0.1 limits who can reach this endpoint but does not make its input +# trustworthy, since any local process or a compromised web worker could call it. +_SAFE_COMPONENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$") + + +def resolve_under_root(rel_path): + """Build an absolute path under ALLOWED_ROOT from a validated relative path.""" + if not isinstance(rel_path, str) or not rel_path or rel_path.startswith("/"): + raise ValueError("expected a non-empty relative path") + parts = [p for p in rel_path.split("/") if p != ""] + if not parts: + raise ValueError("expected a non-empty relative path") + for p in parts: + if p == ".." or not _SAFE_COMPONENT.match(p): + raise ValueError(f"illegal path component: {p!r}") + return os.path.join(ALLOWED_ROOT, *parts) + + +def _patch_gpu_export_resample(): + """Swap nnU-Net's CPU (scipy) probability resampling for the torch/GPU path. + + The nnU-Net CLI has no flag for this, and the export stage is otherwise one + of the two dominant costs (~26-42s -> ~1-6s). Patching the ConfigurationManager + property is the documented-shape extension point that avoids modifying the + vendored nnU-Net install, which other models in this deployment also use. + """ + def _resampling_fn_probabilities(self): + return partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False) + plans_handler.ConfigurationManager.resampling_fn_probabilities = property( + _resampling_fn_probabilities + ) + + +if not torch.cuda.is_available(): + raise SystemExit("ERROR: CUDA not available -- refusing to start the warm predictor") + +print(f"[warm] loading model from {MODEL} ...", flush=True) +_t0 = time.time() +_patch_gpu_export_resample() +predictor = nnUNetPredictor( + tile_step_size=STEP_SIZE, + use_mirroring=not DISABLE_TTA, + perform_everything_on_device=True, + device=torch.device("cuda"), + allow_tqdm=False, + verbose=False, +) +predictor.initialize_from_trained_model_folder( + MODEL, use_folds=("all",), checkpoint_name=CHECKPOINT +) +reader_writer = SimpleITKIO() +preprocessor = DefaultPreprocessor(verbose=False) +_load_seconds = time.time() - _t0 +print(f"[warm] model loaded in {_load_seconds:.1f}s " + f"(step_size={STEP_SIZE}, disable_tta={DISABLE_TTA}), ready on 127.0.0.1:{PORT}", + flush=True) + +# One inference at a time. The web app serializes on its own _gpu_lock too, but +# this service must not assume it is the only caller. +_lock = threading.Lock() +_stats = {"requests": 0, "failures": 0, "last_seconds": None} + + +BBOX_EXPORT = os.environ.get("BBOX_EXPORT", "0").strip().lower() in {"1", "true", "yes", "on"} +BBOX_MARGIN = int(os.environ.get("BBOX_MARGIN", "24")) +LESION_LABELS = [39, 40, 41, 42] + + +def bbox_refined_export(logits, props, margin=BBOX_MARGIN, lesion_labels=LESION_LABELS): + """Cheap nearest-neighbour label map everywhere, exact trilinear refinement only + inside a dilated box around lesion candidates. + + Lesions occupy ~0.02% of a volume, so the expensive 43-channel trilinear resample + runs on ~1-20% of it instead of all of it. align_corners=False's index->coordinate + map is size-independent, so refining a sub-box reproduces what a full resample would + place there to floating-point tolerance. + + Validated on 40 size-stratified lesion cases: pancreatic lesion Dice 0.9947 mean / + 0.977 worst with zero misses; no object of >=100 voxels was lost in any class. The + only two Dice-0.0 events were a 4-voxel and a 6-voxel speck (Dice is degenerate at + that scale). Export stage 4.50s -> 0.23s median, 17.6x. + """ + import torch.nn.functional as F + + lg = (logits if isinstance(logits, torch.Tensor) else torch.from_numpy(logits)).to("cuda") + C = lg.shape[0] + target = tuple(int(v) for v in props["shape_after_cropping_and_before_resampling"]) + + seg_small = lg.argmax(0).to(torch.uint8) + seg = F.interpolate(seg_small[None, None].float(), size=target, mode="nearest")[0, 0].to(torch.uint8) + + lesion_mask = torch.zeros_like(seg, dtype=torch.bool) + for l in lesion_labels: + lesion_mask |= (seg == l) + if not bool(lesion_mask.any()): + return _revert_crop(seg.cpu().numpy(), props) + + idx = lesion_mask.nonzero() + lo = torch.clamp(idx.min(0).values - margin, min=0) + hi = torch.minimum(idx.max(0).values + margin + 1, torch.tensor(target, device=idx.device)) + z0, y0, x0 = (int(v) for v in lo) + z1, y1, x1 = (int(v) for v in hi) + + def norm(a, b, T): + return (torch.arange(a, b, device="cuda", dtype=torch.float32) + 0.5) * 2.0 / T - 1.0 + mz, my, mx = torch.meshgrid(norm(z0, z1, target[0]), norm(y0, y1, target[1]), + norm(x0, x1, target[2]), indexing="ij") + grid = torch.stack((mx, my, mz), dim=-1)[None] # grid axis order is (x,y,z) + + box = (z1 - z0, y1 - y0, x1 - x0) + best_v = torch.full((1,) + box, -float("inf"), device="cuda") + best_i = torch.zeros((1,) + box, device="cuda", dtype=torch.uint8) + for c0 in range(0, C, 8): # chunk channels to bound peak memory + c1 = min(c0 + 8, C) + sub = F.grid_sample(lg[None, c0:c1].float(), grid, mode="bilinear", + padding_mode="border", align_corners=False)[0] + v, a = sub.max(0) + upd = v > best_v[0] + best_v[0] = torch.where(upd, v, best_v[0]) + best_i[0] = torch.where(upd, (a + c0).to(torch.uint8), best_i[0]) + del sub, v, a + seg[z0:z1, y0:y1, x0:x1] = best_i[0] + out = _revert_crop(seg.cpu().numpy(), props) + del lg, seg, best_v, best_i, grid + torch.cuda.empty_cache() + return out + + +def _revert_crop(seg, props): + import numpy as np + out = np.zeros(props["shape_before_cropping"], dtype=np.uint8) + out[tuple(slice(b[0], b[1]) for b in props["bbox_used_for_cropping"])] = seg + return out + + +def run_inference(input_dir, output_dir): + # input_dir/output_dir are built by resolve_under_root() from a validated + # relative path -- they are never taken from the request directly. + os.makedirs(output_dir, exist_ok=True) + case_files = sorted( + os.path.basename(f) for f in os.listdir(input_dir) + if f.endswith("_0000.nii.gz") or f.endswith("_0000.nii") + ) + if not case_files: + raise RuntimeError(f"No cases found in {input_dir}") + + results = [] + for fname in case_files: + case_id = fname[:-len("_0000.nii.gz")] if fname.endswith(".gz") else fname[:-len("_0000.nii")] + t0 = time.time() + data, properties = reader_writer.read_images([os.path.join(input_dir, fname)]) + preprocessed, _, preprocessed_props = preprocessor.run_case_npy( + data, None, properties, + predictor.plans_manager, predictor.configuration_manager, predictor.dataset_json, + ) + t_infer = time.time() + logits = predictor.predict_logits_from_preprocessed_data(torch.from_numpy(preprocessed)) + infer_s = round(time.time() - t_infer, 2) + + t_exp = time.time() + if BBOX_EXPORT: + segmentation = bbox_refined_export(logits, preprocessed_props) + else: + segmentation = convert_logits( + logits, predictor.plans_manager, predictor.configuration_manager, + predictor.label_manager, preprocessed_props, False, + ) + export_s = round(time.time() - t_exp, 2) + out_path = os.path.join(output_dir, f"{case_id}.nii.gz") + reader_writer.write_seg(segmentation, out_path, preprocessed_props) + + # Free the 43-class logit volume before the next case rather than waiting + # for refcount/GC -- it is multi-GB and this process is long-lived. + del logits, segmentation, preprocessed + torch.cuda.empty_cache() + + seconds = round(time.time() - t0, 2) + print(f"[warm] {case_id}: done in {seconds}s " + f"(infer {infer_s}s, export {export_s}s, bbox={BBOX_EXPORT})", flush=True) + results.append({ + "case_id": case_id, "output_path": out_path, "seconds": seconds, + "infer_seconds": infer_s, "export_seconds": export_s, "bbox_export": BBOX_EXPORT, + }) + return results + + +class Handler(BaseHTTPRequestHandler): + protocol_version = "HTTP/1.1" + + def log_message(self, fmt, *args): + print(f"[warm] {self.address_string()} {fmt % args}", flush=True) + + def _json(self, code, payload): + body = json.dumps(payload).encode() + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path != "/health": + self._json(404, {"error": "not found"}) + return + self._json(200, { + "status": "ok", + "model": MODEL, + "allowed_root": ALLOWED_ROOT, + "checkpoint": CHECKPOINT, + "step_size": STEP_SIZE, + "disable_tta": DISABLE_TTA, + "model_load_seconds": round(_load_seconds, 1), + "busy": _lock.locked(), + **_stats, + }) + + def do_POST(self): + if self.path != "/predict": + self._json(404, {"error": "not found"}) + return + try: + length = int(self.headers.get("Content-Length", 0)) + req = json.loads(self.rfile.read(length)) + rel_input = req["input_rel"] + rel_output = req["output_rel"] + except Exception as e: + self._json(400, {"error": f"bad request: {e}"}) + return + + # Paths are constructed here from validated relative components; the + # request never supplies an absolute path. + try: + input_dir = resolve_under_root(rel_input) + output_dir = resolve_under_root(rel_output) + except ValueError as e: + self._json(400, {"error": str(e)}) + return + + # Refuse a configuration we were not started with (see module docstring). + want_step = req.get("step_size") + want_no_tta = req.get("disable_tta") + mismatch = [] + if want_step is not None and abs(float(want_step) - STEP_SIZE) > 1e-9: + mismatch.append(f"step_size: requested {want_step}, server {STEP_SIZE}") + if want_no_tta is not None and bool(want_no_tta) != DISABLE_TTA: + mismatch.append(f"disable_tta: requested {bool(want_no_tta)}, server {DISABLE_TTA}") + if mismatch: + self._json(409, {"error": "configuration mismatch", "details": mismatch}) + return + + if not os.path.isdir(input_dir): + self._json(400, {"error": f"input_dir does not exist: {input_dir}"}) + return + + t0 = time.time() + if not _lock.acquire(timeout=LOCK_TIMEOUT): + self._json(503, {"error": "predictor busy, timed out waiting for GPU"}) + return + try: + results = run_inference(input_dir, output_dir) + _stats["requests"] += 1 + _stats["last_seconds"] = round(time.time() - t0, 2) + self._json(200, { + "status": "ok", + "results": results, + "total_seconds": _stats["last_seconds"], + }) + except Exception as e: + _stats["failures"] += 1 + traceback.print_exc() + # Surface the error so the client can log it before falling back. + self._json(500, {"error": f"{type(e).__name__}: {e}"}) + finally: + _lock.release() + + +if __name__ == "__main__": + server = ThreadingHTTPServer(("127.0.0.1", PORT), Handler) + print(f"[warm] listening on 127.0.0.1:{PORT}", flush=True) + server.serve_forever() diff --git a/flask-server/scripts/run_lesionseg_warm.sh b/flask-server/scripts/run_lesionseg_warm.sh new file mode 100644 index 0000000..2dccaad --- /dev/null +++ b/flask-server/scripts/run_lesionseg_warm.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# Start (or restart) the persistent LesionSegmenter predictor. +# +# Run from a script FILE, never as a nested `bash -lc "..."` over SSH: the quoting +# mangles and the relaunch line can silently not run, which has already caused one +# production outage in this deployment. +# +# setsid detaches the service from the invoking shell so an SSH drop cannot kill it +# mid-start. The kill is scoped to this exact script name -- a broad +# `pkill -f gunicorn`-style pattern would take out unrelated services on this host. +set -u + +PORT="${PORT:-8765}" +MODEL="${LESIONSEG_CKPT_PATH:-/home/visitor/model}" +ENV_NAME="${CONDA_ENV_LESIONSEG:-atlasnet}" +PYBIN="${PYBIN:-$HOME/.conda/envs/$ENV_NAME/bin/python}" +LOG="${LOG:-/tmp/lesionseg_warm.log}" +SCRIPT="$(cd "$(dirname "$0")" && pwd)/lesionseg_warm_server.py" + +# Settings must match what auto_segmentor.py requests, or the server answers 409 and +# every request silently falls back to the slow cold path. +STEP_SIZE="${STEP_SIZE:-0.7}" +DISABLE_TTA="${DISABLE_TTA:-1}" +BBOX_EXPORT="${BBOX_EXPORT:-0}" + +# Directory tree the service may read from and write to. Paths in a /predict +# request are confined to this root, so it must contain the app's per-session +# inference dirs. Defaults to the same place Constants.SESSIONS_DIR_NAME resolves to. +FLASK_DIR="$(cd "$(dirname "$0")/.." && pwd)" +LESIONSEG_ROOT="${LESIONSEG_ROOT:-${SESSIONS_DIR_PATH:-$FLASK_DIR/sessions}}" + +if [ ! -x "$PYBIN" ]; then + echo "ERROR: python not found at $PYBIN (set PYBIN or CONDA_ENV_LESIONSEG)" >&2 + exit 1 +fi +if [ ! -f "$SCRIPT" ]; then + echo "ERROR: $SCRIPT not found" >&2 + exit 1 +fi + +echo "[run_warm] stopping any existing warm predictor on port $PORT" +pkill -f "lesionseg_warm_server.py" 2>/dev/null +sleep 3 +if pgrep -f "lesionseg_warm_server.py" >/dev/null; then + echo "[run_warm] still alive after SIGTERM, sending SIGKILL" + pkill -9 -f "lesionseg_warm_server.py" 2>/dev/null + sleep 2 +fi + +echo "[run_warm] starting: model=$MODEL port=$PORT step=$STEP_SIZE disable_tta=$DISABLE_TTA bbox=$BBOX_EXPORT" +echo "[run_warm] permitted path root: $LESIONSEG_ROOT" +mkdir -p "$LESIONSEG_ROOT" +setsid env \ + MODEL="$MODEL" PORT="$PORT" STEP_SIZE="$STEP_SIZE" \ + DISABLE_TTA="$DISABLE_TTA" BBOX_EXPORT="$BBOX_EXPORT" \ + LESIONSEG_ROOT="$LESIONSEG_ROOT" \ + "$PYBIN" -u "$SCRIPT" > "$LOG" 2>&1 < /dev/null & + +# Wait for readiness rather than a guessed sleep; the model load dominates startup. +for i in $(seq 1 180); do + sleep 1 + if curl -sf -m 2 "http://127.0.0.1:$PORT/health" > /dev/null 2>&1; then + echo "[run_warm] healthy after ${i}s:" + curl -s -m 3 "http://127.0.0.1:$PORT/health"; echo + exit 0 + fi +done + +echo "[run_warm] ERROR: did not become healthy within 180s. Log tail:" >&2 +tail -30 "$LOG" >&2 +exit 1