From 430f4c5b5eb6b72e90c55de725ce3e43eb3081e4 Mon Sep 17 00:00:00 2001
From: Aditya Sanjeev <adityasanjeev28@gmail.com>
Date: Fri, 14 Aug 2026 22:45:11 -0700
Subject: [PATCH 1/3] Fix epai_predict.py: never actually produced output (2
 bugs)

Discovered while smoke-testing before enabling EPAI_SCRIPT_PATH in prod --
the fast GPU-export ePAI path has been merged but was never turned on, and
for good reason: it was silently broken.

Bug 1 - save_probabilities=False produced zero output, no error. The
fork's export_prediction_from_logits gates its ENTIRE write path (segmentation
write, tumor-stats extraction, CSV update) behind 'if save_probabilities:'.
The unconditional segmentation write that used to run regardless is commented
out at the bottom of that function (a leftover from whatever refactor added
the CSV pipeline). save_probabilities=False looked like the obviously correct
choice (we don't want a probabilities file) but actually means 'do nothing at
all'. Confirmed: rc=0, 'done <case>' printed, zero bytes written, empty CSV.

Bug 2 - wrong reader class. This model's plans specify
NibabelIOWithReorient as the image_reader_writer_class, not SimpleITKIO
(which the script hardcoded for reading the input CT). Reading with the wrong
class produces a properties dict missing 'nibabel_stuff', which the nibabel
writer needs for the reoriented affine -> KeyError deep inside write_seg the
moment bug 1 is fixed and the write path actually runs.

Verified after both fixes: 3 cases (small/medium/a 1060-slice large volume),
all produced real segmentation files and fully-populated findings CSV rows
(shape/spacing/pancreas/duct/PDAC/cyst/PNET stats with sensible confidence
values) end to end.
---
 flask-server/scripts/epai_predict.py | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/flask-server/scripts/epai_predict.py b/flask-server/scripts/epai_predict.py
index f42c0e0b..d21de951 100644
--- a/flask-server/scripts/epai_predict.py
+++ b/flask-server/scripts/epai_predict.py
@@ -27,7 +27,6 @@
 import numpy as np
 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 export_prediction_from_logits
 from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet
@@ -80,9 +79,15 @@ def main():
         checkpoint_name=os.getenv("EPAI_CHECKPOINT_NAME", "checkpoint_final.pth"),
     )
 
-    rw = SimpleITKIO()
     pp = DefaultPreprocessor(verbose=False)
     pm, cm = predictor.plans_manager, predictor.configuration_manager
+    # Must match the writer export_prediction_from_logits uses internally
+    # (also pm.image_reader_writer_class() - this model's plans specify
+    # NibabelIOWithReorient, not SimpleITKIO). Reading with the wrong class
+    # produces a properties dict missing 'nibabel_stuff', which the nibabel
+    # writer needs for the reoriented affine -> KeyError deep inside write_seg.
+    # Confirmed by direct smoke test before this was wired into EPAI_SCRIPT_PATH.
+    rw = pm.image_reader_writer_class()
 
     ct = os.path.join(input_dir, f"{case_id}_0000.nii.gz")
     if not os.path.exists(ct):
@@ -96,10 +101,21 @@ def main():
     # The fork's own export: converts logits -> native-res segmentation (GPU resample,
     # thanks to the patch) AND writes the tumor-findings row into output_csv. Writes the
     # segmentation to <save_dir>/<case_id>.nii.gz.
+    #
+    # save_probabilities MUST be True here despite the name: in this fork it does not
+    # mean "write a probabilities.npz" (that branch is commented out -- "we don't save
+    # probabilities as pickles anymore"). It gates the ENTIRE write path -- segmentation
+    # write, tumor-stats extraction, and the output-CSV update all live inside
+    # `if save_probabilities:`; the unconditional segmentation write that used to run
+    # regardless is commented out at the bottom of the function (a leftover from
+    # whoever refactored this in). Passing False (the "obviously correct" choice,
+    # since we don't want a probabilities file) silently produces ZERO output --
+    # no exception, exit 0, nothing written. Confirmed by direct smoke test before
+    # this was wired into EPAI_SCRIPT_PATH.
     out_trunc = os.path.join(save_dir, case_id)
     export_prediction_from_logits(
         logits, pprops, cm, pm, predictor.dataset_json, out_trunc,
-        save_probabilities=False, output_csv_path=output_csv,
+        save_probabilities=True, output_csv_path=output_csv,
     )
     print(f"done {case_id}", flush=True)
 

From e78c2458b7f817ff855dc5b34399c1e2b6666e05 Mon Sep 17 00:00:00 2001
From: Aditya Sanjeev <adityasanjeev28@gmail.com>
Date: Fri, 14 Aug 2026 23:05:21 -0700
Subject: [PATCH 2/3] Add ePAI warm/persistent predictor

Same proven pattern as scripts/lesionseg_warm_server.py: load the model ONCE
and keep it GPU-resident, eliminating the cold-start subprocess reload
(interpreter + torch import + checkpoint load + cuDNN autotune) that
_run_epai_inference currently pays on every single request via
epai_predict.py.

- epai_warm_server.py (new): persistent HTTP predictor on 127.0.0.1:8766.
  Applies the same two fork fixes as the just-fixed cold path (anisotropy
  axis coercion, correct NibabelIOWithReorient reader class,
  save_probabilities=True), plus the validated GPU export-resample patch.
  Refuses (409) a request whose step_size/disable_tta don't match what it
  was started with, rather than silently serving a different configuration -
  an invisible accuracy change in a cancer-detection tool. Same path-
  containment pattern as lesionseg_warm_server.py (CodeQL-verified there):
  requests may only name a location as a validated relative path under an
  allowed root, never an absolute path.

- run_epai_warm.sh (new): launcher, mirrors run_lesionseg_warm.sh -
  setsid-detached so an SSH drop can't kill it mid-start, waits for a real
  health check rather than a guessed sleep.

- epai_predict.py (modified): now tries EPAI_WARM_URL first if set, falling
  back to the existing (just-fixed) cold in-process path on any failure -
  server down, busy, or a configuration mismatch. Heavy imports (torch,
  nnunetv2) moved inside _cold_predict so the warm path never pays their
  cost either.

Verified end-to-end on bdmap1 against a scratch sessions root (not the real
one): two back-to-back requests for different cases through a running warm
server, both producing real segmentation files + correct findings CSV rows
matching the cold-path smoke test's output (tiny float-noise differences
only, consistent with known GPU nondeterminism). Second request took 8.87s
total with 0.024s of client-side CPU, confirming the model stayed resident
across requests rather than reloading.

NOT enabled by default - EPAI_WARM_URL is unset until the server is started
via run_epai_warm.sh and the env var is set in .env, matching how the
LesionSegmenter warm predictor was rolled out.
---
 flask-server/scripts/epai_predict.py     | 255 ++++++++++++++------
 flask-server/scripts/epai_warm_server.py | 295 +++++++++++++++++++++++
 flask-server/scripts/run_epai_warm.sh    |  64 +++++
 3 files changed, 540 insertions(+), 74 deletions(-)
 create mode 100644 flask-server/scripts/epai_warm_server.py
 create mode 100644 flask-server/scripts/run_epai_warm.sh

diff --git a/flask-server/scripts/epai_predict.py b/flask-server/scripts/epai_predict.py
index d21de951..450f080a 100644
--- a/flask-server/scripts/epai_predict.py
+++ b/flask-server/scripts/epai_predict.py
@@ -1,64 +1,164 @@
 #!/usr/bin/env python3
-"""Optimized ePAI inference: GPU-resident export + the fork's own findings CSV.
+"""Optimized ePAI inference: warm persistent predictor + GPU-resident export.
 
-Drop-in for the bare `nnUNetv2_predict_from_modelfolder` call in
 services/auto_segmentor.py, wired in via the EPAI_SCRIPT_PATH hook (which invokes:
 `bash <script> session_dir case_id input_dir save_dir input_csv output_csv ckpt_path`).
-Runs inside the `epai` conda env, where `import nnunetv2` resolves to qchen76's fork.
-
-Why it exists: profiling showed ePAI's export/convert stage is ~158s of CPU
-resampling of the 26-channel logits back to native resolution -- as expensive as
-inference. Moving that resample to the GPU cuts it to ~0.3s (measured ~50-60x,
-99.99% voxel agreement -- boundary jitter only, no lesion-level change). It reuses
-the fork's own `export_prediction_from_logits`, which does BOTH the segmentation
-export AND the tumor-findings CSV (PDAC/cyst/PNET stats), so the findings report is
-produced identically -- only the resample runs on GPU.
-
-Two fork-specific fixes vs the standard-nnU-Net LesionSegmenter wrapper:
-  1. determine_do_sep_z_and_axis in the fork returns the anisotropy axis as a bare
-     scalar on the separate-z path, so resample_torch's `assert len(axis)==1` crashes
-     (standard nnU-Net returns a 1-element list). Coerce scalar -> [int(axis)].
-  2. Patch ConfigurationManager.resampling_fn_probabilities -> the GPU resampler.
+
+Two execution paths, same pattern as scripts/lesionseg_predict.py:
+
+1. WARM (fast). If EPAI_WARM_URL is set and that service is healthy, the work is
+   handed to the persistent predictor (scripts/epai_warm_server.py), which holds
+   the model in GPU memory across requests -- skipping the cold-start subprocess
+   reload (interpreter + torch import + checkpoint load + cuDNN autotune) this
+   script otherwise pays on EVERY request. Everything above the fallback below is
+   stdlib-only; the heavy imports live inside _cold_predict so the warm path
+   never pays their cost either.
+
+2. COLD (fallback). Load the model in-process and run, exactly as before. Used
+   when no warm URL is configured, the service is down, or it is running a
+   different configuration than the caller asked for. Deliberate: the website
+   must keep working if the warm service dies, just more slowly.
+
+Why this wrapper exists at all instead of the bare CLI: the CLI has no flag for
+GPU-accelerated export resampling. Swapping ConfigurationManager.resampling_fn_
+probabilities for resample_torch_fornnunet cuts that stage from ~158s of CPU
+resampling to ~0.3s (measured ~50-55x at 500-case scale, voxel agreement
+0.99996+) with no accuracy cost.
+
+Three fork-specific fixes vs the standard-nnU-Net LesionSegmenter wrapper (this
+runs inside the `epai` conda env, where `import nnunetv2` resolves to qchen76's
+fork, not vanilla nnU-Net):
+  1. determine_do_sep_z_and_axis in the fork returns the anisotropy axis as a
+     bare scalar on the separate-z path, so resample_torch's `assert
+     len(axis)==1` crashes (standard nnU-Net returns a 1-element list). Coerce
+     scalar -> [int(axis)].
+  2. This model's plans specify NibabelIOWithReorient as the reader/writer
+     class, not SimpleITKIO. Reading with the wrong class produces a properties
+     dict missing 'nibabel_stuff', which the nibabel writer needs for the
+     reoriented affine -> KeyError deep inside write_seg. Must read with
+     plans_manager.image_reader_writer_class(), not a hardcoded reader.
+  3. export_prediction_from_logits's entire write path (segmentation file,
+     tumor-stats extraction, output-CSV update) lives inside `if
+     save_probabilities:` in this fork -- the unconditional write that used to
+     run regardless is commented out (a leftover from whatever refactor added
+     the CSV pipeline). save_probabilities=False looks like the obviously
+     correct choice (we don't want a probabilities.npz -- that branch is itself
+     commented out, "we don't save probabilities as pickles anymore") but
+     actually means "write nothing at all," silently: rc=0, no exception, zero
+     bytes written. MUST be True.
+All three found and fixed via direct smoke test (3 cases: small/medium/a
+1060-slice large volume, all producing real segmentation files and fully
+populated findings CSV rows) before this was ever wired into EPAI_SCRIPT_PATH.
 """
+import json
 import os
 import sys
-from functools import partial
-
-import numpy as np
-import torch
-from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
-from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor
-from nnunetv2.inference.export_prediction import export_prediction_from_logits
-from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet
-import nnunetv2.preprocessing.resampling.resample_torch as _rt
-import nnunetv2.utilities.plans_handling.plans_handler as _ph
-
-
-def _apply_fork_fixes():
-    # (1) coerce the anisotropy axis back to a list so `assert len(axis)==1` holds
-    _orig_determine = _rt.determine_do_sep_z_and_axis
-
-    def _fixed_determine(*a, **k):
-        do_sep, axis = _orig_determine(*a, **k)
-        if do_sep and np.isscalar(axis):
-            axis = [int(axis)]
-        return do_sep, axis
-
-    _rt.determine_do_sep_z_and_axis = _fixed_determine
-
-    # (2) GPU-resident probabilities resampling (the ~158s -> ~0.3s export lever)
-    _ph.ConfigurationManager.resampling_fn_probabilities = property(
-        lambda self: partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False)
-    )
-
-
-def main():
-    # Positional contract from services/auto_segmentor.py's EPAI_SCRIPT_PATH hook.
-    if len(sys.argv) < 8:
-        print("usage: epai_predict.py session_dir case_id input_dir save_dir "
-              "input_csv output_csv ckpt_path", file=sys.stderr)
-        sys.exit(2)
-    _session_dir, case_id, input_dir, save_dir, _input_csv, output_csv, ckpt_path = sys.argv[1:8]
+import urllib.error
+import urllib.request
+
+
+def _warm_predict(case_id, input_dir, output_dir, output_csv, step_size, disable_tta, url, timeout):
+    """Try the persistent predictor. Returns True if it produced the output.
+
+    Mirrors lesionseg_predict.py's _warm_predict: 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
+
+    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
+
+    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)
+        rel_csv = os.path.relpath(os.path.realpath(output_csv), root)
+    except ValueError as e:
+        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, rel_csv)):
+        print(f"[warm] paths are outside the server root {root}; using cold path", flush=True)
+        return False
+
+    body = json.dumps({
+        "case_id": case_id,
+        "input_rel": rel_in.replace(os.sep, "/"),
+        "output_rel": rel_out.replace(os.sep, "/"),
+        "output_csv_rel": rel_csv.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 {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
+    return True
+
+
+def _cold_predict(case_id, input_dir, save_dir, output_csv, ckpt_path):
+    """In-process inference. Heavy imports are local so the warm path never pays them."""
+    from functools import partial
+
+    import numpy as np
+    import torch
+    from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+    from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor
+    from nnunetv2.inference.export_prediction import export_prediction_from_logits
+    from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet
+    import nnunetv2.preprocessing.resampling.resample_torch as _rt
+    import nnunetv2.utilities.plans_handling.plans_handler as _ph
+
+    def _apply_fork_fixes():
+        # (1) coerce the anisotropy axis back to a list -- see module docstring, fix 1.
+        _orig_determine = _rt.determine_do_sep_z_and_axis
+
+        def _fixed_determine(*a, **k):
+            do_sep, axis = _orig_determine(*a, **k)
+            if do_sep and np.isscalar(axis):
+                axis = [int(axis)]
+            return do_sep, axis
+
+        _rt.determine_do_sep_z_and_axis = _fixed_determine
+
+        # (2) GPU-resident probabilities resampling (the ~158s -> ~0.3s export lever)
+        _ph.ConfigurationManager.resampling_fn_probabilities = property(
+            lambda self: partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False)
+        )
 
     if not torch.cuda.is_available():
         print("ERROR: CUDA not available", file=sys.stderr)
@@ -81,12 +181,9 @@ def main():
 
     pp = DefaultPreprocessor(verbose=False)
     pm, cm = predictor.plans_manager, predictor.configuration_manager
-    # Must match the writer export_prediction_from_logits uses internally
-    # (also pm.image_reader_writer_class() - this model's plans specify
-    # NibabelIOWithReorient, not SimpleITKIO). Reading with the wrong class
-    # produces a properties dict missing 'nibabel_stuff', which the nibabel
-    # writer needs for the reoriented affine -> KeyError deep inside write_seg.
-    # Confirmed by direct smoke test before this was wired into EPAI_SCRIPT_PATH.
+    # Must match the writer export_prediction_from_logits uses internally -- see
+    # module docstring, fix 2. Reading with SimpleITKIO produces a properties
+    # dict missing 'nibabel_stuff', which the nibabel writer needs.
     rw = pm.image_reader_writer_class()
 
     ct = os.path.join(input_dir, f"{case_id}_0000.nii.gz")
@@ -98,20 +195,7 @@ def main():
     dpp, _, pprops = pp.run_case_npy(data, None, props, pm, cm, predictor.dataset_json)
     logits = predictor.predict_logits_from_preprocessed_data(torch.from_numpy(dpp))
 
-    # The fork's own export: converts logits -> native-res segmentation (GPU resample,
-    # thanks to the patch) AND writes the tumor-findings row into output_csv. Writes the
-    # segmentation to <save_dir>/<case_id>.nii.gz.
-    #
-    # save_probabilities MUST be True here despite the name: in this fork it does not
-    # mean "write a probabilities.npz" (that branch is commented out -- "we don't save
-    # probabilities as pickles anymore"). It gates the ENTIRE write path -- segmentation
-    # write, tumor-stats extraction, and the output-CSV update all live inside
-    # `if save_probabilities:`; the unconditional segmentation write that used to run
-    # regardless is commented out at the bottom of the function (a leftover from
-    # whoever refactored this in). Passing False (the "obviously correct" choice,
-    # since we don't want a probabilities file) silently produces ZERO output --
-    # no exception, exit 0, nothing written. Confirmed by direct smoke test before
-    # this was wired into EPAI_SCRIPT_PATH.
+    # save_probabilities MUST be True -- see module docstring, fix 3.
     out_trunc = os.path.join(save_dir, case_id)
     export_prediction_from_logits(
         logits, pprops, cm, pm, predictor.dataset_json, out_trunc,
@@ -120,5 +204,28 @@ def main():
     print(f"done {case_id}", flush=True)
 
 
+def main():
+    # Positional contract from services/auto_segmentor.py's EPAI_SCRIPT_PATH hook.
+    if len(sys.argv) < 8:
+        print("usage: epai_predict.py session_dir case_id input_dir save_dir "
+              "input_csv output_csv ckpt_path", file=sys.stderr)
+        sys.exit(2)
+    _session_dir, case_id, input_dir, save_dir, _input_csv, output_csv, ckpt_path = sys.argv[1:8]
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    step_size = float(os.getenv("EPAI_STEP_SIZE", "0.5"))
+    disable_tta = os.getenv("EPAI_DISABLE_TTA", "0").strip().lower() in {"1", "true", "yes", "on"}
+
+    warm_url = os.getenv("EPAI_WARM_URL", "").strip()
+    if warm_url:
+        timeout = int(os.getenv("EPAI_WARM_TIMEOUT", "3600"))
+        if _warm_predict(case_id, input_dir, save_dir, output_csv, step_size, disable_tta,
+                         warm_url, timeout):
+            return
+
+    _cold_predict(case_id, input_dir, save_dir, output_csv, ckpt_path)
+
+
 if __name__ == "__main__":
     main()
diff --git a/flask-server/scripts/epai_warm_server.py b/flask-server/scripts/epai_warm_server.py
new file mode 100644
index 00000000..381192b4
--- /dev/null
+++ b/flask-server/scripts/epai_warm_server.py
@@ -0,0 +1,295 @@
+#!/usr/bin/env python3
+"""
+Warm/persistent ePAI (qchen76 fork) inference service.
+
+Same pattern as scripts/lesionseg_warm_server.py -- loads the model ONCE and keeps
+it resident, eliminating the cold-start subprocess reload (interpreter + torch
+import + checkpoint load + cuDNN autotune) that services/auto_segmentor.py's
+EPAI_SCRIPT_PATH hook currently pays on EVERY request via epai_predict.py.
+
+Deliberately stdlib-only for the HTTP layer, binds 127.0.0.1 only -- see
+lesionseg_warm_server.py's docstring for the full rationale, unchanged here.
+
+Inference settings mirror the validated configuration:
+  - EPAI_DISABLE_TTA -- validated on 366 cases across two independent case
+    populations (lesion-focused + general pool): dropping TTA does not cost
+    detection (recall and detection rate were HIGHER without it in both
+    samples), at ~2.1x the speed. See epai_step_tta_val.py results.
+  - GPU export-resident resampling -- validated at 500-case scale, ~51-55x on
+    the export stage, voxel agreement 0.99996+ (see epai_export_bench.py).
+  - tile_step_size is NOT overridden from the nnU-Net default (0.5) here --
+    unlike LesionSegmenter's banked 0.7, ePAI's own step size was never swept
+    against real ground truth. Do not change EPAI_STEP_SIZE's default without
+    that validation first (same rigor as the TTA frontier: real GT, report
+    detection rate and whole-lesion misses, not just mean dice).
+
+A client requesting settings this server was not started with is REFUSED
+(HTTP 409), matching lesionseg_warm_server.py's reasoning: silently serving a
+different configuration would be an invisible accuracy change in a
+cancer-detection tool.
+
+Two fork-specific fixes vs the standard-nnU-Net LesionSegmenter wrapper (see
+epai_predict.py's docstring for the full story of how each was found):
+  1. determine_do_sep_z_and_axis in the fork returns the anisotropy axis as a
+     bare scalar on the separate-z path; resample_torch's `assert len(axis)==1`
+     needs a 1-element list. Coerce scalar -> [int(axis)].
+  2. This model's plans specify NibabelIOWithReorient, not SimpleITKIO --
+     reading with the wrong reader class produces a properties dict missing
+     'nibabel_stuff', which the nibabel writer needs -> KeyError in write_seg.
+  3. export_prediction_from_logits's write path (segmentation file, tumor-stats
+     extraction, output-CSV update) is entirely gated behind
+     `if save_probabilities:` in this fork -- despite the name, this does NOT
+     mean "write a probabilities.npz" (that branch is commented out). It MUST
+     be True or nothing is written at all, silently, rc=0.
+
+Usage:
+    EPAI_CKPT_PATH=/home/visitor/ePAI/model/... PORT=8766 python epai_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 numpy as np
+import torch
+from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor
+from nnunetv2.preprocessing.preprocessors.default_preprocessor import DefaultPreprocessor
+from nnunetv2.inference.export_prediction import export_prediction_from_logits
+from nnunetv2.preprocessing.resampling.resample_torch import resample_torch_fornnunet
+import nnunetv2.preprocessing.resampling.resample_torch as _rt
+import nnunetv2.utilities.plans_handling.plans_handler as _ph
+
+MODEL = os.environ.get(
+    "EPAI_CKPT_PATH",
+    "/home/visitor/ePAI/model/qchen76_2025_0421/nnUNetTrainer__nnUNetPlans__3d_fullres",
+)
+PORT = int(os.environ.get("PORT", "8766"))
+CHECKPOINT = os.environ.get("EPAI_CHECKPOINT_NAME", "checkpoint_final.pth")
+STEP_SIZE = float(os.environ.get("EPAI_STEP_SIZE", "0.5"))
+DISABLE_TTA = os.environ.get("EPAI_DISABLE_TTA", "0").strip().lower() in {"1", "true", "yes", "on"}
+LOCK_TIMEOUT = int(os.environ.get("LOCK_TIMEOUT", "1800"))
+
+
+def _default_root():
+    """Directory tree the service is willing to read from and write to.
+
+    Mirrors lesionseg_warm_server.py's _default_root(): the Flask app's
+    sessions directory, where auto_segmentor.py builds the per-session
+    input/output dirs this service is asked to process.
+    """
+    env = os.environ.get("EPAI_WARM_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()
+
+# Same containment pattern as lesionseg_warm_server.py (CodeQL-verified there):
+# a request may only name a location as a validated RELATIVE path under
+# ALLOWED_ROOT. No absolute paths, no "..". The absolute path is always built
+# server-side; the request's path components are never trusted directly.
+_SAFE_COMPONENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+def resolve_under_root(rel_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 _apply_fork_fixes():
+    _orig_determine = _rt.determine_do_sep_z_and_axis
+
+    def _fixed_determine(*a, **k):
+        do_sep, axis = _orig_determine(*a, **k)
+        if do_sep and np.isscalar(axis):
+            axis = [int(axis)]
+        return do_sep, axis
+
+    _rt.determine_do_sep_z_and_axis = _fixed_determine
+
+    _ph.ConfigurationManager.resampling_fn_probabilities = property(
+        lambda self: partial(resample_torch_fornnunet, device=torch.device("cuda"), is_seg=False)
+    )
+
+
+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()
+_apply_fork_fixes()
+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
+)
+preprocessor = DefaultPreprocessor(verbose=False)
+# Must match the writer export_prediction_from_logits uses internally (see
+# module docstring, fix 2) -- this model's plans specify NibabelIOWithReorient.
+reader_writer = predictor.plans_manager.image_reader_writer_class()
+
+_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)
+
+_lock = threading.Lock()
+_stats = {"requests": 0, "failures": 0, "last_seconds": None}
+
+
+def run_inference(case_id, input_dir, output_dir, output_csv_path):
+    os.makedirs(output_dir, exist_ok=True)
+    ct = os.path.join(input_dir, f"{case_id}_0000.nii.gz")
+    if not os.path.exists(ct):
+        raise RuntimeError(f"input not found: {ct}")
+
+    t0 = time.time()
+    data, props = reader_writer.read_images([ct])
+    pm, cm = predictor.plans_manager, predictor.configuration_manager
+    preprocessed, _, preprocessed_props = preprocessor.run_case_npy(
+        data, None, props, pm, cm, 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()
+    out_trunc = os.path.join(output_dir, case_id)
+    # save_probabilities=True is required -- see module docstring, fix 3.
+    export_prediction_from_logits(
+        logits, preprocessed_props, cm, pm, predictor.dataset_json, out_trunc,
+        save_probabilities=True, output_csv_path=output_csv_path,
+    )
+    export_s = round(time.time() - t_exp, 2)
+
+    del logits, preprocessed
+    torch.cuda.empty_cache()
+
+    seconds = round(time.time() - t0, 2)
+    out_path = out_trunc + predictor.dataset_json.get("file_ending", ".nii.gz")
+    print(f"[warm] {case_id}: done in {seconds}s (infer {infer_s}s, export {export_s}s)", flush=True)
+    return {
+        "case_id": case_id, "output_path": out_path, "seconds": seconds,
+        "infer_seconds": infer_s, "export_seconds": export_s,
+    }
+
+
+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))
+            case_id = req["case_id"]
+            rel_input = req["input_rel"]
+            rel_output = req["output_rel"]
+            rel_output_csv = req["output_csv_rel"]
+        except Exception as e:
+            self._json(400, {"error": f"bad request: {e}"})
+            return
+
+        try:
+            input_dir = resolve_under_root(rel_input)
+            output_dir = resolve_under_root(rel_output)
+            output_csv_path = resolve_under_root(rel_output_csv)
+        except ValueError as e:
+            self._json(400, {"error": str(e)})
+            return
+
+        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
+        if not os.path.isfile(output_csv_path):
+            # auto_segmentor.py's _ensure_output_csv_template writes this before
+            # invoking the script -- the fork's export function updates this
+            # existing row rather than appending a fresh one (see epai_predict.py).
+            self._json(400, {"error": f"output_csv_path does not exist: {output_csv_path}"})
+            return
+
+        t0 = time.time()
+        if not _lock.acquire(timeout=LOCK_TIMEOUT):
+            self._json(503, {"error": "predictor busy, timed out waiting for GPU"})
+            return
+        try:
+            result = run_inference(case_id, input_dir, output_dir, output_csv_path)
+            _stats["requests"] += 1
+            _stats["last_seconds"] = round(time.time() - t0, 2)
+            self._json(200, {
+                "status": "ok",
+                "results": [result],
+                "total_seconds": _stats["last_seconds"],
+            })
+        except Exception as e:
+            _stats["failures"] += 1
+            traceback.print_exc()
+            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_epai_warm.sh b/flask-server/scripts/run_epai_warm.sh
new file mode 100644
index 00000000..462133e8
--- /dev/null
+++ b/flask-server/scripts/run_epai_warm.sh
@@ -0,0 +1,64 @@
+#!/bin/bash
+# Start (or restart) the persistent ePAI predictor.
+#
+# Run from a script FILE, never as a nested `bash -lc "..."` over SSH -- see
+# run_lesionseg_warm.sh's comment on why (a real production outage already
+# happened from that pattern).
+#
+# 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.
+set -u
+
+PORT="${PORT:-8766}"
+MODEL="${EPAI_CKPT_PATH:-/home/visitor/ePAI/model/qchen76_2025_0421/nnUNetTrainer__nnUNetPlans__3d_fullres}"
+ENV_NAME="${CONDA_ENV_EPAI:-epai}"
+PYBIN="${PYBIN:-$HOME/.conda/envs/$ENV_NAME/bin/python}"
+LOG="${LOG:-/tmp/epai_warm.log}"
+SCRIPT="$(cd "$(dirname "$0")" && pwd)/epai_warm_server.py"
+
+# Settings must match what auto_segmentor.py requests, or the server answers
+# 409 and every request falls back to the slow cold path.
+STEP_SIZE="${EPAI_STEP_SIZE:-0.5}"
+DISABLE_TTA="${EPAI_DISABLE_TTA:-1}"
+
+FLASK_DIR="$(cd "$(dirname "$0")/.." && pwd)"
+EPAI_WARM_ROOT="${EPAI_WARM_ROOT:-${SESSIONS_DIR_PATH:-$FLASK_DIR/sessions}}"
+
+if [ ! -x "$PYBIN" ]; then
+  echo "ERROR: python not found at $PYBIN (set PYBIN or CONDA_ENV_EPAI)" >&2
+  exit 1
+fi
+if [ ! -f "$SCRIPT" ]; then
+  echo "ERROR: $SCRIPT not found" >&2
+  exit 1
+fi
+
+echo "[run_warm] stopping any existing ePAI warm predictor on port $PORT"
+pkill -f "epai_warm_server.py" 2>/dev/null
+sleep 3
+if pgrep -f "epai_warm_server.py" >/dev/null; then
+  echo "[run_warm] still alive after SIGTERM, sending SIGKILL"
+  pkill -9 -f "epai_warm_server.py" 2>/dev/null
+  sleep 2
+fi
+
+echo "[run_warm] starting: model=$MODEL port=$PORT step=$STEP_SIZE disable_tta=$DISABLE_TTA"
+echo "[run_warm] permitted path root: $EPAI_WARM_ROOT"
+mkdir -p "$EPAI_WARM_ROOT"
+setsid env \
+  EPAI_CKPT_PATH="$MODEL" PORT="$PORT" EPAI_STEP_SIZE="$STEP_SIZE" \
+  EPAI_DISABLE_TTA="$DISABLE_TTA" EPAI_WARM_ROOT="$EPAI_WARM_ROOT" \
+  "$PYBIN" -u "$SCRIPT" > "$LOG" 2>&1 < /dev/null &
+
+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

From 793f6199d9e0a1d61925be23b8431dc869685e79 Mon Sep 17 00:00:00 2001
From: Aditya Sanjeev <adityasanjeev28@gmail.com>
Date: Sat, 15 Aug 2026 00:09:06 -0700
Subject: [PATCH 3/3] Fix path-injection in epai_warm_server.py (CodeQL
 py/path-injection, high)

case_id came straight from the request JSON and was interpolated into
os.path.join(input_dir, f"{case_id}_0000.nii.gz") and
os.path.join(output_dir, case_id) with no validation -- unlike
input_rel/output_rel/output_csv_rel, which all go through
resolve_under_root(). A case_id like "/etc/passwd" discards input_dir
entirely (os.path.join treats an absolute-path component as a full
reset), giving arbitrary read via read_images and arbitrary write via
the export path.

Fix mirrors the already CodeQL-clean lesionseg_warm_server.py: never
take case_id from the client at all. _discover_case() derives it from
whatever *_0000.nii.gz file is actually present in the (already
validated) input_dir, so the untrusted string never reaches a path
operation. Verified end-to-end on bdmap1: real inference request
through the fixed server produces the correct segmentation file and
CSV row, keyed by the same case_id as before.
---
 flask-server/scripts/epai_warm_server.py | 31 +++++++++++++++++++-----
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/flask-server/scripts/epai_warm_server.py b/flask-server/scripts/epai_warm_server.py
index 381192b4..cbb440af 100644
--- a/flask-server/scripts/epai_warm_server.py
+++ b/flask-server/scripts/epai_warm_server.py
@@ -157,11 +157,28 @@ def _fixed_determine(*a, **k):
 _stats = {"requests": 0, "failures": 0, "last_seconds": None}
 
 
-def run_inference(case_id, input_dir, output_dir, output_csv_path):
+def _discover_case(input_dir):
+    """Derive case_id from a file actually present in the validated input_dir,
+    the same pattern lesionseg_warm_server.py uses (CodeQL-verified there): the
+    client's case_id string must never reach a path operation, even after
+    input_dir itself is validated -- a containment check on input_dir alone
+    does not make case_id safe to interpolate into a filename (CodeQL flagged
+    exactly this: os.path.join(dir, f"{case_id}...") lets an absolute-path
+    case_id discard dir entirely). Removing the tainted flow, not auditing it.
+    """
+    candidates = [f for f in os.listdir(input_dir) if f.endswith("_0000.nii.gz")]
+    if len(candidates) != 1:
+        raise RuntimeError(
+            f"expected exactly one *_0000.nii.gz case in {input_dir}, found {len(candidates)}"
+        )
+    fname = candidates[0]
+    return fname, fname[: -len("_0000.nii.gz")]
+
+
+def run_inference(input_dir, output_dir, output_csv_path):
     os.makedirs(output_dir, exist_ok=True)
-    ct = os.path.join(input_dir, f"{case_id}_0000.nii.gz")
-    if not os.path.exists(ct):
-        raise RuntimeError(f"input not found: {ct}")
+    fname, case_id = _discover_case(input_dir)
+    ct = os.path.join(input_dir, fname)
 
     t0 = time.time()
     data, props = reader_writer.read_images([ct])
@@ -231,7 +248,9 @@ def do_POST(self):
         try:
             length = int(self.headers.get("Content-Length", 0))
             req = json.loads(self.rfile.read(length))
-            case_id = req["case_id"]
+            # case_id is accepted for schema compatibility with the client but is
+            # NEVER used to build a path -- see _discover_case's docstring.
+            _ = req["case_id"]
             rel_input = req["input_rel"]
             rel_output = req["output_rel"]
             rel_output_csv = req["output_csv_rel"]
@@ -273,7 +292,7 @@ def do_POST(self):
             self._json(503, {"error": "predictor busy, timed out waiting for GPU"})
             return
         try:
-            result = run_inference(case_id, input_dir, output_dir, output_csv_path)
+            result = run_inference(input_dir, output_dir, output_csv_path)
             _stats["requests"] += 1
             _stats["last_seconds"] = round(time.time() - t0, 2)
             self._json(200, {