Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion PanTS-Demo/src/components/segmentation/RefineFlyout.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ export default function RefineFlyout({ classLabel, refine, onApplied }: Props) {
return (
<div style={{ display: "flex", flexDirection: "column", gap: 8 }}>
<p className="atb-flyout-note">
The model redraws this class's outline from its current voxels, with no clicks. It only grows into unlabeled voxels, and Undo reverses it.
The model redraws this class's outline from its current voxels, with no clicks. It can add unlabeled voxels and trim this class's own, never touches another structure, and Undo reverses it.
</p>
{error && <p className="atb-flyout-note atb-flyout-note--error" role="alert">{error}</p>}
<MenuDivider />
Expand Down
5 changes: 3 additions & 2 deletions PanTS-Demo/src/test/refineFlyout.test.tsx
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
/**
* The old standalone RefineFlyout no longer exists. Refinement now shares the
* annotation ribbon's AI slot; this contract tests the current copy/state.
* Refinement shares the annotation ribbon's AI slot (the label and when it shows
* as "AI refine" are tested here); the flyout it opens is RefineFlyout.tsx, whose
* failure messages are tested in refineFlyoutErrors.test.tsx.
*/
import { describe, expect, it } from "vitest";
import { aiSlotCopy, deriveToolbarState, toolbarLayout } from "../components/viewer/annotationToolbarState";
Expand Down
54 changes: 54 additions & 0 deletions docs/operations-reference.md
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,60 @@ reachable server, a plain click or box falls back to the older region-grow tool,
prompts only the model can answer (corrections, scribble, lasso) say the model
is unavailable.

#### Session slots

The model server holds one "slot" (lease) per live prompt session and refuses a request
when all slots are taken, which the viewer shows as "busy". `--max-sessions` sets the
number of slots (default 3). One person can need two at once: the click, box, scribble
and lasso tools keep one session per structure, and "Refine with model" opens a second,
short-lived one. So:

- Set `NNINTERACTIVE_MAX_SESSIONS` in `flask-server/.env` to the same number as the
server's `--max-sessions`, never higher. If the server has fewer, it refuses leases the
backend believes it has, and the backend logs a warning saying so.
- Leave `NNINTERACTIVE_MAX_SESSIONS_PER_USER=2` (prompt session plus Refine).
- With a single slot, Refine fails whenever the same person used a click tool in the last
ten minutes, and everyone else is refused while one person holds it.
- A slot is freed when the viewer leaves the tool or the page closes, or after the idle
timeout (`NNINTERACTIVE_SESSION_TTL`, 600 s).

GPU work is serialised, so with several people predicting at once the later ones wait.
Measured on the lab host: a click or Refine on a real full-resolution scan answered in
about 1 to 4 s alone, and three predictions started together on a synthetic volume
finished about 4 s apart. Every slot also holds its CT in the model server's memory (a
full-resolution scan is about 140 MB as float32), so check memory before raising the
count a lot.

#### Running the model server as a service

On the lab host, run it as a systemd user service so it survives logouts and reboots
(`loginctl enable-linger <user>` once). Save as `~/.config/systemd/user/nninteractive-bodymaps.service`:
```
[Unit]
Description=nnInteractive model server for the BodyMaps website
After=network.target

[Service]
ExecStart=/path/to/nninteractive-env/bin/nninteractive-server --model nnInteractive_v1.0 --device cuda --host 127.0.0.1 --port 1530 --no-torch-compile --max-sessions 3
Restart=on-failure
RestartSec=10
StandardOutput=append:/path/to/nninteractive.log
StandardError=append:/path/to/nninteractive.log

[Install]
WantedBy=default.target
```
Then `systemctl --user daemon-reload && systemctl --user enable --now nninteractive-bodymaps`,
and point `NNINTERACTIVE_URL` at it. To change the slot count, edit `--max-sessions`,
change `NNINTERACTIVE_MAX_SESSIONS` to match, `systemctl --user restart` the service, then
restart the backend (`deploy/start_backend.sh`, at 0 running jobs). Restart when the log
shows `0/N active`: anyone mid-edit loses their model session once, and the viewer
restarts it from their current mask.

To see what is happening: `GET /api/health` shows `interactive.leases` against
`interactive.max` (this backend's own leases), and the model server's log prints
`claimed session ... (k/N active)` and `released session ...` for every slot change.

#### Create Conda Environment
```
conda create -n PanTS_backend python=3.11
Expand Down
5 changes: 3 additions & 2 deletions flask-server/.env.example
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,9 @@ PANTS_LABEL_SCHEMA=ipantsmini-v1
# Preserve this fork's case/class/resolution sessions and baseline-injection
# protocol with nninteractive-client 2.6.0. The model runs in a separate server.
# NNINTERACTIVE_URL=http://127.0.0.1:1527
# Maximum simultaneous leases held by this Flask process (default 2); the
# model server permits 3 by default, leaving one for other processes.
# Maximum simultaneous leases held by this Flask process (default 2). Keep it
# no larger than the model server's --max-sessions (default 3), or the server
# refuses leases this process thinks it can have and everyone sees "busy".
# NNINTERACTIVE_MAX_SESSIONS=2
# Maximum leases one signed-in reader may hold (default 2: prompt + Refine).
# NNINTERACTIVE_MAX_SESSIONS_PER_USER=2
Expand Down
6 changes: 6 additions & 0 deletions flask-server/api/api_blueprint.py
Original file line number Diff line number Diff line change
Expand Up @@ -3985,10 +3985,16 @@ def health():
with _inference_jobs_lock:
statuses = [job.get("status") for job in inference_jobs.values()]
running, queued = statuses.count("running"), statuses.count("queued")
try:
from services import nninteractive_predictor as _nn
interactive = _nn.slot_stats()
except Exception:
interactive = None # health must answer even if the model client is not installed
resp = jsonify({
"status": "ok" if mount["ok"] else "degraded",
"mount": mount,
"inference": {"running": running, "queued": queued, "active": running + queued},
"interactive": interactive,
"threads": threading.active_count(),
})
resp.status_code = 200 if mount["ok"] else 503
Expand Down
39 changes: 38 additions & 1 deletion flask-server/services/nninteractive_predictor.py
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,39 @@ def _evict_oldest_idle(predicate) -> bool:
return False


_LAST_LIMIT_WARNING = 0.0


def _warn_if_model_server_refused_below_local_limit(exc: Exception) -> None:
"""Say so when the model server is full but this backend thinks it has room.

This backend caps its own leases (NNINTERACTIVE_MAX_SESSIONS) and expects the
model server's --max-sessions to be at least as large. If the server refuses a
lease while this process holds fewer than its own limit, the two settings have
drifted apart (or another client is using the server), and every user sees
"busy". Logged at most once a minute. Caller holds _STORE_LOCK.
"""
global _LAST_LIMIT_WARNING
if not _is_capacity_error(exc) or len(_SESSIONS) >= _MAX_SESSIONS:
return
now = time.time()
if now - _LAST_LIMIT_WARNING < 60:
return
_LAST_LIMIT_WARNING = now
print(
f"[nninteractive_predictor] WARNING: the model server refused a session while this backend holds "
f"only {len(_SESSIONS)} of its {_MAX_SESSIONS} (NNINTERACTIVE_MAX_SESSIONS). The model server's "
f"--max-sessions is lower than that, or another client is using it. Raise --max-sessions or lower "
f"NNINTERACTIVE_MAX_SESSIONS so they match."
)


def slot_stats() -> dict:
"""This backend's own model leases, for /api/health. Never raises or blocks on the model server."""
with _STORE_LOCK:
return {"leases": len(_SESSIONS), "max": _MAX_SESSIONS, "per_user": _MAX_SESSIONS_PER_USER}


def session_is_active(
case_id: str,
segment_label: int | None,
Expand Down Expand Up @@ -470,7 +503,11 @@ def get_or_create_session(
raise PromptCapacityError("Interactive segmentation is at capacity; existing prompt sessions were preserved.")

print(f"[nninteractive_predictor] creating session {key} (ct={ct.shape}, baseline={baseline_mask is not None})")
sess = _NNSession(ct, baseline_mask=baseline_mask)
try:
sess = _NNSession(ct, baseline_mask=baseline_mask)
except Exception as exc:
_warn_if_model_server_refused_below_local_limit(exc)
raise
sess._bodymaps_owner = owner_key
sess._bodymaps_token = session_token
sess._bodymaps_case_id = case_id
Expand Down
29 changes: 29 additions & 0 deletions flask-server/tests/unit/test_health_interactive.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
"""/api/health reports this backend's model leases so "busy" can be told apart from "broken"."""
from flask import Flask

from services import nninteractive_predictor as predictor


def _health_json():
from api.api_blueprint import health

app = Flask(__name__)
with app.test_request_context("/api/health"):
response = health()
return response.get_json()


def test_health_includes_the_model_lease_counts():
body = _health_json()
assert body["interactive"]["max"] == predictor._MAX_SESSIONS
assert body["interactive"]["leases"] == len(predictor._SESSIONS)


def test_health_still_answers_when_the_lease_stats_fail(monkeypatch):
def boom():
raise RuntimeError("model client missing")

monkeypatch.setattr(predictor, "slot_stats", boom)
body = _health_json()
assert body["interactive"] is None
assert body["status"] in {"ok", "degraded"}
52 changes: 52 additions & 0 deletions flask-server/tests/unit/test_interactive_session_capacity.py
Original file line number Diff line number Diff line change
Expand Up @@ -211,3 +211,55 @@ def close(self): self.closed = True
assert FakeRemote.instances[1].calls[0] == ("seed", 1)
assert result[0, 1, 1] == 1
assert result[1, 1, 1] == 1


def test_slot_stats_reports_this_backends_leases_and_limits(monkeypatch):
monkeypatch.setattr(predictor, "_NNSession", FakeSession)
ct = np.zeros((2, 2, 2), dtype=np.float32)
assert predictor.slot_stats() == {"leases": 0, "max": 2, "per_user": 2}
predictor.get_or_create_session("case-a", 1, "low", ct, session_token="a", owner="alice")
assert predictor.slot_stats()["leases"] == 1


class RefusingSession:
def __init__(self, *_args, **_kwargs):
raise RuntimeError("server is at capacity")


def test_warns_when_the_model_server_refuses_below_this_backends_own_limit(monkeypatch, capsys):
monkeypatch.setattr(predictor, "_NNSession", RefusingSession)
monkeypatch.setattr(predictor, "_LAST_LIMIT_WARNING", 0.0)
ct = np.zeros((2, 2, 2), dtype=np.float32)

with pytest.raises(RuntimeError, match="at capacity"):
predictor.get_or_create_session("case-a", 1, "low", ct, session_token="a", owner="alice")

out = capsys.readouterr().out
assert "--max-sessions" in out and "NNINTERACTIVE_MAX_SESSIONS" in out
assert "holds only 0 of its 2" in out


def test_the_limit_warning_is_rate_limited(monkeypatch, capsys):
monkeypatch.setattr(predictor, "_NNSession", RefusingSession)
monkeypatch.setattr(predictor, "_LAST_LIMIT_WARNING", 0.0)
ct = np.zeros((2, 2, 2), dtype=np.float32)

for token in ("a", "b", "c"):
with pytest.raises(RuntimeError):
predictor.get_or_create_session("case-a", 1, "low", ct, session_token=token, owner="alice")

assert capsys.readouterr().out.count("WARNING: the model server refused") == 1


def test_no_limit_warning_for_other_failures_or_when_the_local_pool_is_full(monkeypatch, capsys):
monkeypatch.setattr(predictor, "_LAST_LIMIT_WARNING", 0.0)
predictor._warn_if_model_server_refused_below_local_limit(RuntimeError("connection refused"))
assert "WARNING" not in capsys.readouterr().out

monkeypatch.setattr(predictor, "_NNSession", FakeSession)
ct = np.zeros((2, 2, 2), dtype=np.float32)
predictor.get_or_create_session("case-a", 1, "low", ct, session_token="a", owner="alice")
predictor.get_or_create_session("case-b", 1, "low", ct, session_token="b", owner="bob")
predictor._warn_if_model_server_refused_below_local_limit(RuntimeError("server is at capacity"))
assert "WARNING" not in capsys.readouterr().out

Loading