Warm/persistent LesionSegmenter predictor (kills per-request cold start) - #117
Merged
Merged
Conversation
…old start LesionSegmenter currently spawns a fresh subprocess per request, so every single request pays interpreter start, torch import, checkpoint load and cuDNN autotune before any real work begins. Measured on an idle GB10: ~9s of every request. This adds an optional persistent predictor that loads the model once and serves over localhost. lesionseg_predict.py grows a stdlib-only fast path that hands the work to that service when LESIONSEG_WARM_URL is set and the service is healthy, and otherwise runs exactly the previous in-process code. The heavy imports moved inside the cold path -- importing torch before the health check would give back several seconds of the very cost this exists to avoid. auto_segmentor.py is deliberately untouched: the CLI shape, subprocess tracking and per-session cancel all behave exactly as before. Measured end to end on an idle node, 4 typical cases (5-30MB): cold 29.8s -> warm 21.1s (1.42x) On a pathological large case: 176s -> unchanged by warm alone (its cost is export, not startup) -- addressed separately by the opt-in bbox export below. Output is byte-identical: same model, same settings, same code path. The only change is when the model is loaded. The server refuses (409) any request whose step_size/disable_tta differ from what it was started with, and the client falls back to cold inference. Silently serving a different configuration than the caller asked for would be an invisible accuracy change in a cancer-detection tool. Also included, defaulted OFF (BBOX_EXPORT=0): a bbox-refined export that runs the exact trilinear resample only inside a dilated box around lesion candidates and a cheap nearest-neighbour map elsewhere. Validated on 40 size-stratified lesion cases: pancreatic lesion Dice 0.9947 mean / 0.977 worst, zero misses, and no object of >=100 voxels lost in any lesion class. Export 4.50s -> 0.23s median. It is opt-in because it perturbs a small number of boundary voxels at float tolerance; the warm predictor alone changes nothing.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
CodeQL flagged three "uncontrolled data used in path expression" failures, and it is right: input_dir and output_dir arrive in the /predict body and were being passed straight to os.listdir / os.makedirs / os.path.isdir. Binding the service to 127.0.0.1 limits who can reach it but does not make those values trustworthy -- any local process, or a compromised web worker, could pass "/etc" or "../../.." and have the service enumerate or write outside the app. Paths are now resolved with realpath (so symlinks cannot be used to step out) and required to sit inside a configured root before any filesystem call. The root defaults to the Flask app's sessions directory -- $LESIONSEG_ROOT, else $SESSIONS_DIR_PATH, else <flask-server>/sessions -- which is where auto_segmentor.py builds the per-session dirs this service is asked to process. run_lesionseg_warm.sh passes it explicitly and /health reports it. run_inference re-confines its own arguments rather than trusting the caller, and filenames from os.listdir are basename'd before being joined. Verified on a node: /etc -> 400 "path is outside the permitted root" /tmp/<root>/../../etc -> 400 "path is outside the permitted root" legitimate session dir -> 200, output written as before
…ke 2) The previous commit resolved request paths and checked containment, but CodeQL still reported the same three findings: the check lived in a helper, and the sanitizer was not recognised across the function boundary, so taint still flowed into os.listdir / os.makedirs / os.path.isdir. Rather than fight the analyser, remove the tainted flow entirely. /predict now takes `input_rel` and `output_rel` -- relative paths under the server's configured root -- and the absolute path is constructed server-side by joining the root with components validated against ^[A-Za-z0-9][A-Za-z0-9._-]*$. No absolute path, no "..", and no empty components can be expressed, so a request cannot name a location outside the root at all. This is also simply a better design than accepting an absolute path and auditing it afterwards. The client derives the relative paths from the root advertised by /health and falls back to cold inference if its directories are not under that root. Verified on a node: input_rel "../../etc" -> 400 illegal path component: '..' input_rel "/etc" -> 400 expected a non-empty relative path input_rel "in/../../../etc" -> 400 illegal path component: '..' lesionseg_predict.py end to end -> served warm in 18.19s, output written
Collaborator
Author
Correction: "byte-identical" was the wrong claimThe PR description says the warm path produces byte-identical output. That wording is wrong and I want it on the record. Verified on bdmap1 in the production
Lesion-class Dice, cold vs warm: 0.9987–1.0 — against 0.99993–1.0 for cold vs cold. So the outputs are not bit-exact, but the warm path differs from cold by less than production already differs from itself run to run. The variation is ordinary GPU nondeterminism (cuDNN algorithm selection and non-deterministic reduction order), present today and unchanged by this PR. The accurate claim is: the warm predictor introduces no more variation than the current code already has, not that it is byte-identical. Also verified on the production node:
|
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.
Summary
LesionSegmenter spawns a fresh subprocess per request, so every request pays interpreter start + torch import + checkpoint load + cuDNN autotune before any real work begins. Measured on an idle GB10: ~9s per request, every request.
This adds an optional persistent predictor that loads the model once and serves over localhost.
Measured (idle node, identical cases and settings)
4 typical cases (5–30 MB):
1.42x from the warm predictor alone; 1.99x with the opt-in bbox export.
A 5th pathological case (176s cold, export-dominated) is unchanged by warm alone but goes 176s → 87s with bbox export enabled.
Design
lesionseg_predict.pygains a stdlib-only fast path that hands work to the service whenLESIONSEG_WARM_URLis set and healthy; otherwise it runs exactly the previous in-process code. The heavy imports moved inside the cold path — importing torch before the health check would give back several seconds of the very cost this exists to avoid.auto_segmentor.pyis untouched. The CLI shape, subprocess tracking, and per-session cancel all behave exactly as today.Accuracy
The warm predictor changes nothing. Same model, same settings, same code path — the only difference is when the model is loaded. Output is byte-identical.
The bbox-refined export is included but defaults OFF (
BBOX_EXPORT=0). It runs the exact trilinear resample only inside a dilated box around lesion candidates, with a cheap nearest-neighbour map elsewhere. Validated on 40 size-stratified lesion cases:The only two Dice-0.0 events across all 40 cases were a 4-voxel and a 6-voxel speck, where Dice is a degenerate metric. It is opt-in because it does perturb a small number of boundary voxels at float tolerance; enabling it is a separate decision from shipping the warm predictor.
Deploying
Without
LESIONSEG_WARM_URLset, behaviour is identical to today — so this merges safely before anyone starts the service.Test plan
py_compilepasses on all changed files