ePAI: fix the write-path bug, then ship a warm predictor - #156
Merged
Conversation
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.
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.
aperson30
changed the title
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Two commits, same underlying investigation.
1. Fix epai_predict.py: silently produced zero output (2 bugs)
Found while smoke-testing before finally enabling
EPAI_SCRIPT_PATHin prod (the GPU-resident export path has been sitting merged-but-off — this is why: it never actually worked).Bug 1 —
save_probabilities=Falseproduced zero output, no error. The fork'sexport_prediction_from_logitsgates its entire write path (segmentation file, tumor-stats extraction, CSV update) behindif save_probabilities:. The unconditional write that used to run regardless is commented out at the bottom of that function — a leftover from whichever refactor added the CSV pipeline.Falsereads as the obviously-correct choice (we don't want a probabilities file) but in this fork it means "do nothing at all." Confirmed:rc=0,done <case>printed, zero bytes written, CSV untouched.Bug 2 — wrong reader class. This model's plans specify
NibabelIOWithReorient, not the hardcodedSimpleITKIO. Reading with the wrong class produces a properties dict missingnibabel_stuff, which the nibabel writer needs for the reoriented affine —KeyErrorinsidewrite_segthe moment bug 1's fix lets the write path actually run.Verified: 3 cases (small, medium, a 1060-slice large volume) — all produced real segmentation files and fully-populated findings CSV rows end to end.
2. Add a warm/persistent predictor
Same proven pattern as
lesionseg_warm_server.py— load the model once, keep it GPU-resident, skip the cold-start subprocess reload on every request.epai_warm_server.py(new) — persistent HTTP predictor, applies both fixes above plus the validated GPU export-resample patch. Refuses (409) a request whose settings don't match what it was started with rather than silently serving a different configuration. Same path-containment pattern as the LesionSegmenter warm server (CodeQL-verified there).run_epai_warm.sh(new) — launcher, setsid-detached, waits for a real health check.epai_predict.py(modified) — triesEPAI_WARM_URLfirst if set, falls back to the (now-fixed) cold path on any failure. Heavy imports moved inside the cold path so the warm path never pays their cost.Verified end-to-end on a scratch sessions root: two back-to-back requests for different cases through one running warm server, both producing correct output matching the cold-path test (tiny float-noise only, known GPU nondeterminism). Second request: 8.87s total, 0.024s of client-side CPU — confirms the model stays resident, no reload between requests.
Not enabled by default —
EPAI_WARM_URLstays unset until the server is started and the var is set in.env, same rollout pattern as the LesionSegmenter warm predictor.This unblocks the full ePAI speed stack once
.envis configured: GPU export-resample (~51-55x, validated at 500-case scale) + no-TTA (~2.1x, validated at 366-case scale, no detection cost) + warm predictor (kills the cold-start reload on every request).