Skip to content

ePAI: fix the write-path bug, then ship a warm predictor #156

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aperson30 merged 3 commits into from
Aug 15, 2026
Merged

ePAI: fix the write-path bug, then ship a warm predictor #156

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
430f4c5
Fix epai_predict.py: never actually produced output (2 bugs)
aperson30 Aug 15, 2026
e78c245
Add ePAI warm/persistent predictor
aperson30 Aug 15, 2026
793f619
Fix path-injection in epai_warm_server.py (CodeQL py/path-injection, …
aperson30 Aug 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
243 changes: 183 additions & 60 deletions flask-server/scripts/epai_predict.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,65 +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.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
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)
Expand All @@ -80,9 +179,12 @@ 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 -- 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")
if not os.path.exists(ct):
Expand All @@ -93,16 +195,37 @@ 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 -- 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,
save_probabilities=False, output_csv_path=output_csv,
save_probabilities=True, output_csv_path=output_csv,
)
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()
Loading
Loading