nnInteractive 2.1.0__tar.gz → 2.2.0__tar.gz

{nninteractive-2.1.0 → nninteractive-2.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nnInteractive
-Version: 2.1.0
+Version: 2.2.0
 Summary: Inference code for nnInteractive
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>
@@ -226,10 +226,14 @@ Requires-Dist: batchgenerators>=0.25.1
 Requires-Dist: fastapi>=0.110
 Requires-Dist: uvicorn[standard]>=0.27
 Requires-Dist: httpx>=0.27
+Requires-Dist: blosc2
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
+Provides-Extra: client
+Requires-Dist: httpx>=0.27; extra == "client"
+Requires-Dist: blosc2; extra == "client"
 Dynamic: license-file
 
 <img src="imgs/nnInteractive_header_white.png">

{nninteractive-2.1.0 → nninteractive-2.2.0}/nnInteractive/inference/inference_session.py

@@ -50,14 +50,16 @@ class nnInteractiveInferenceSession:
     ):
         """
         Only intended to work with nnInteractiveTrainerV2 and its derivatives
+
+        ``use_torch_compile``: compile the network with ``torch.compile``. The
+        first prediction after enabling this is slow (compilation happens lazily
+        on the first forward pass), but every subsequent prediction is faster.
+        This is recommended for the persistent inference server, where the
+        process is long-lived so the one-time compile cost is paid only once and
+        amortized across the whole session lifetime.
         """
         print("session initialized")
 
-        # set as part of initialization
-        assert use_torch_compile is False, (
-            "torch.compile is not supported. The blosc2-backed interaction tensor "
-            "requires numpy↔torch round-trips that break compile tracing."
-        )
         self.network = None
         self.label_manager = None
         self.dataset_json = None
@@ -883,6 +885,42 @@ class nnInteractiveInferenceSession:
         else:
             del initial_seg
 
+    @torch.inference_mode()
+    def warmup(self) -> bool:
+        """Run a single dummy forward pass to trigger lazy ``torch.compile`` compilation up front.
+
+        With ``torch.compile`` enabled the network is compiled lazily on its first
+        forward pass, which would otherwise make the user's *first* real prediction
+        slow. Every prediction path — the initial coarse pass, the zoom-out
+        iterations, and the refinement patches — feeds the network an input of
+        identical shape ``[1, num_input_channels + num_interaction_channels,
+        *patch_size]`` (``_build_network_input`` always resizes the crop to
+        ``patch_size``, and refinement crops at exactly ``patch_size``). So a single
+        dummy pass at that shape populates the compile cache and every subsequent
+        real prediction is fast.
+
+        Returns ``True`` if a warmup pass was run, ``False`` if it was a no-op
+        (network not compiled — there is nothing to pre-compile, so a dummy pass
+        would not save the user any time). Mirrors ``_predict``'s autocast/
+        inference-mode context and the float32 input dtype that ``torch.cat``
+        produces when concatenating the float32 image with the fp16 interactions.
+        """
+        if self.network is None or self.configuration_manager is None:
+            raise RuntimeError("warmup() requires an initialized network; call initialize_* first")
+        if not isinstance(self.network, OptimizedModule):
+            return False
+        num_input_channels = (
+            determine_num_input_channels(self.plans_manager, self.configuration_manager, self.dataset_json)
+            + self.num_interaction_channels
+        )
+        patch_size = self.configuration_manager.patch_size
+        dummy = torch.zeros((1, num_input_channels, *patch_size), dtype=torch.float32, device=self.device)
+        with torch.autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context():
+            self.network(dummy)
+        del dummy
+        empty_cache(self.device)
+        return True
+
     @torch.inference_mode()
     def _predict(self, force_full_refine: bool = False):
         """
@@ -1307,9 +1345,16 @@ class nnInteractiveInferenceSession:
 
         Returns an artifact dict that can be applied to this or any other freshly
         constructed session via :meth:`initialize_from_loaded_artifacts`. The
-        returned references (network, plans_manager, configuration_manager, ...)
-        are safe to share by reference across sessions — they are treated as
-        read-only after construction.
+        returned values are the actual objects (the ``nn.Module`` with its
+        weights and buffers, the plans/configuration managers, the dataset
+        json, the label manager) — not copies. Multiple sessions calling
+        :meth:`initialize_from_loaded_artifacts` with the same dict will all
+        end up with ``self.network`` pointing at the same module instance and
+        the same weight tensors on the GPU. This is safe as long as callers
+        treat these objects as read-only after construction; in the multi-
+        session server that is enforced by running inference under
+        ``@torch.inference_mode()`` and serializing predict calls with a
+        global GPU lock.
 
         Note: this also mutates ``self`` (applies capability, sets pad/decay/
         thickness) because ``num_interaction_channels`` is required to build the
@@ -1402,7 +1447,10 @@ class nnInteractiveInferenceSession:
 
         ``artifacts`` is the dict returned by :meth:`_load_model_artifacts_from_disk`.
         Useful for spawning multiple sessions that share one loaded model (e.g.
-        the multi-session inference server).
+        the multi-session inference server). All artifact entries — including
+        ``self.network`` — are stored by reference; passing the same dict to
+        multiple sessions does not duplicate the network or its weights in
+        memory.
         """
         self.preferred_scribble_thickness = artifacts["preferred_scribble_thickness"]
         self.interaction_decay = artifacts["interaction_decay"]

{nninteractive-2.1.0 → nninteractive-2.2.0}/nnInteractive/inference/remote/remote_session.py

@@ -10,6 +10,7 @@ from __future__ import annotations
 
 import json
 import os
+import threading
 import warnings
 from typing import List, Optional, Tuple, Union
 
@@ -165,8 +166,16 @@ class nnInteractiveRemoteInferenceSession:
         claim_info = claim_resp.json()
         self._lease_token = claim_info["lease_token"]
         self.idle_timeout_seconds: float = float(claim_info.get("idle_timeout_seconds", 0.0))
+        self.liveness_timeout_seconds: float = float(claim_info.get("liveness_timeout_seconds", 0.0))
         self._http.headers[LEASE_HEADER] = self._lease_token
 
+        # Background liveness heartbeat bookkeeping. Defined before any code that
+        # might raise so close()/__del__ can always reference them safely. The
+        # thread itself is started at the end of __init__, once construction has
+        # fully succeeded.
+        self._stop_heartbeat = threading.Event()
+        self._heartbeat_thread: Optional[threading.Thread] = None
+
         caps = self._get_json(PATH_CAPABILITIES)
 
         # Attributes that mirror the local session so the GUI can introspect them
@@ -187,6 +196,17 @@ class nnInteractiveRemoteInferenceSession:
         self.target_buffer: Union[np.ndarray, torch.Tensor, None] = None
         self.do_autozoom: bool = bool(caps.get("do_autozoom", True))
 
+        # Construction succeeded — start auto-heartbeating to keep the server
+        # from reaping us as a dead client. Beat at half the liveness timeout so
+        # one dropped request still leaves margin. Daemon thread: it never blocks
+        # interpreter exit, and close() joins it cleanly.
+        if self.liveness_timeout_seconds > 0:
+            self._heartbeat_interval = max(5.0, self.liveness_timeout_seconds / 2.0)
+            self._heartbeat_thread = threading.Thread(
+                target=self._heartbeat_loop, name="nnInteractive-heartbeat", daemon=True
+            )
+            self._heartbeat_thread.start()
+
     # ------------------------------------------------------------------ #
     #                       HTTP helpers (private)                       #
     # ------------------------------------------------------------------ #
@@ -272,7 +292,10 @@ class nnInteractiveRemoteInferenceSession:
         self.target_buffer = target_buffer
         self._post_json(
             PATH_SET_TARGET_BUFFER,
-            {"shape": list(target_buffer.shape), "dtype": _buffer_dtype_str(target_buffer)},
+            {
+                "shape": list(target_buffer.shape),
+                "dtype": _buffer_dtype_str(target_buffer),
+            },
         )
 
     def set_do_autozoom(self, do_autozoom: bool) -> None:
@@ -420,17 +443,44 @@ class nnInteractiveRemoteInferenceSession:
             return False
 
     def heartbeat(self) -> float:
-        """Extend this session's idle timeout. Returns remaining seconds.
+        """Tell the server this client is still alive. Returns remaining seconds
+        until the *idle* timeout.
+
+        This proves liveness only: it stops the server from reaping the session
+        as a crashed/dead client, but it does NOT postpone the idle timeout —
+        that is refreshed solely by real user actions (``set_image``,
+        ``add_*_interaction``, …). A session left untouched will therefore still
+        be reaped at the idle timeout even while heartbeats keep arriving.
 
-        For GUIs that keep the app open across long idle stretches: drive
-        this from a timer (e.g. every 60 s) to avoid SessionExpiredError
-        after the configured idle timeout.
+        You normally never call this yourself: the session auto-heartbeats from
+        a background thread for the lifetime of the object.
         """
         resp = self._http.post(PATH_HEARTBEAT)
         _raise_for_lease_errors(resp)
         resp.raise_for_status()
         return float(resp.json().get("remaining_seconds", 0.0))
 
+    def _heartbeat_loop(self) -> None:
+        """Background daemon: prove liveness every ``_heartbeat_interval`` seconds.
+
+        Stops when the session is closed or once the lease is gone. Transient
+        network errors are swallowed so a brief blip doesn't kill the heartbeat;
+        the server's liveness timeout tolerates a few missed beats. Lease expiry
+        (idle reap or server restart) is surfaced to the user on their next real
+        call, not from this thread.
+        """
+        while not self._stop_heartbeat.wait(self._heartbeat_interval):
+            try:
+                self.heartbeat()
+            except SessionExpiredError:
+                break
+            except httpx.HTTPError:
+                continue
+            except Exception:
+                # Never let the daemon thread die noisily (e.g. client closing
+                # concurrently). Bail out quietly.
+                break
+
     def lease_status(self) -> dict:
         """Read-only probe: how much time is left before this session is reaped?
 
@@ -444,6 +494,14 @@ class nnInteractiveRemoteInferenceSession:
         return resp.json()
 
     def close(self) -> None:
+        # Stop the heartbeat thread first so it can't use self._http after we
+        # close it. join() with a short timeout: the thread spends almost all
+        # its time in Event.wait(), which the set() interrupts immediately.
+        self._stop_heartbeat.set()
+        if self._heartbeat_thread is not None:
+            self._heartbeat_thread.join(timeout=5.0)
+            self._heartbeat_thread = None
+
         # Best-effort release so the server can free our slot for other users
         # without waiting for the idle reaper. Swallow errors: the server may
         # already be gone, our lease may already be expired, etc. close()

{nninteractive-2.1.0 → nninteractive-2.2.0}/nnInteractive/inference/server/app.py

@@ -2,15 +2,24 @@
 
 The server hosts up to ``max_sessions`` concurrent
 :class:`nnInteractiveInferenceSession` instances, one per connected client. The
-model weights are loaded once at startup and shared by reference across all
-sessions — per-session state (image, target buffer, interactions tensor) is
-isolated.
+model artifacts (``nn.Module`` network with its weights and buffers, plans/
+configuration managers, dataset json, label manager) are loaded once at startup;
+each session's ``self.network`` is a plain Python reference to that single
+module — there is exactly one network and one copy of the weights on the GPU
+regardless of session count. Per-session state (image, target buffer,
+interactions tensor) is isolated. Safety of sharing relies on (a) inference
+running under ``@torch.inference_mode()`` and (b) a global ``gpu_lock``
+serializing predict-capable endpoints, so no two sessions ever touch the
+module concurrently and nothing mutates it after construction.
 
 Each client identifies itself via a lease token issued by ``POST /claim``. The
 token rides along on every subsequent request in the ``X-Lease-Token`` header.
-Sessions idle for longer than ``idle_timeout_seconds`` are reaped automatically;
-subsequent requests bearing a reaped lease receive HTTP 410 Gone so the client
-can surface a "session expired" message.
+Sessions are reaped automatically for either of two reasons: the user went idle
+(no real interaction for longer than ``idle_timeout_seconds``) or the client
+went dead (no request of any kind — not even a heartbeat — for longer than the
+much shorter ``liveness_timeout_seconds``, which reclaims slots held by crashed
+clients quickly). Subsequent requests bearing a reaped lease receive HTTP 410
+Gone so the client can surface a "session expired" message.
 
 Concurrency model:
   - Each session has its own ``threading.Lock`` that serializes the per-session
@@ -72,13 +81,28 @@ class SessionEntry:
         self.session = session
         self.lock = threading.Lock()
         self.created_at = time.monotonic()
+        # Two independent clocks. ``last_active_at`` tracks real user actions
+        # (drives the idle timeout); ``last_seen_at`` tracks any sign of life
+        # including heartbeats (drives the much shorter liveness timeout used to
+        # reclaim slots held by crashed/disconnected clients).
         self.last_active_at = self.created_at
+        self.last_seen_at = self.created_at
 
-    def touch(self) -> None:
-        self.last_active_at = time.monotonic()
+    def mark_seen(self) -> None:
+        """Record that the client is still alive (liveness only)."""
+        self.last_seen_at = time.monotonic()
+
+    def mark_active(self) -> None:
+        """Record real user activity. Activity implies liveness, so this bumps
+        both clocks."""
+        now = time.monotonic()
+        self.last_active_at = now
+        self.last_seen_at = now
 
     def close(self) -> None:
-        """Free the session's per-instance state. Shared model artifacts are NOT freed.
+        """Free the session's per-instance state. The shared network module and
+        other model artifacts are NOT freed — they live in the registry and are
+        reused by future sessions.
 
         Best-effort: any exception here is logged but not re-raised; cleanup must
         not block reaping or shutdown.
@@ -100,8 +124,10 @@ class SessionFull(Exception):
 class SessionRegistry:
     """Threadsafe lease-keyed dict of :class:`SessionEntry`.
 
-    The model artifacts loaded at server startup are stashed here and reused
-    whenever a new session is created.
+    The model artifacts loaded at server startup are stashed here and handed to
+    every newly created session by reference (the ``nn.Module`` instance, its
+    weights, and the plans/configuration/label-manager objects are not copied
+    or re-instantiated per session).
     """
 
     def __init__(
@@ -109,17 +135,21 @@ class SessionRegistry:
         artifacts: dict,
         max_sessions: int,
         idle_timeout_seconds: float,
+        liveness_timeout_seconds: float,
         device: torch.device,
         torch_n_threads: int,
         do_autozoom: bool,
+        use_torch_compile: bool,
         verbose: bool,
     ) -> None:
         self._artifacts = artifacts
         self._max_sessions = int(max_sessions)
         self._idle_timeout_seconds = float(idle_timeout_seconds)
+        self._liveness_timeout_seconds = float(liveness_timeout_seconds)
         self._device = device
         self._torch_n_threads = torch_n_threads
         self._do_autozoom = do_autozoom
+        self._use_torch_compile = use_torch_compile
         self._verbose = verbose
         self._entries: dict[str, SessionEntry] = {}
         self._mu = threading.Lock()
@@ -132,6 +162,10 @@ class SessionRegistry:
     def idle_timeout_seconds(self) -> float:
         return self._idle_timeout_seconds
 
+    @property
+    def liveness_timeout_seconds(self) -> float:
+        return self._liveness_timeout_seconds
+
     def claim(self) -> str:
         """Create a new session and return its lease token. Raises SessionFull if at cap."""
         with self._mu:
@@ -140,7 +174,7 @@ class SessionRegistry:
             token = uuid.uuid4().hex
             session = nnInteractiveInferenceSession(
                 device=self._device,
-                use_torch_compile=False,
+                use_torch_compile=self._use_torch_compile,
                 verbose=self._verbose,
                 torch_n_threads=self._torch_n_threads,
                 do_autozoom=self._do_autozoom,
@@ -148,18 +182,28 @@ class SessionRegistry:
             session.initialize_from_loaded_artifacts(self._artifacts)
             entry = SessionEntry(session)
             self._entries[token] = entry
-            logger.info("claimed session %s (%d/%d active)", token, len(self._entries), self._max_sessions)
+            logger.info(
+                "claimed session %s (%d/%d active)",
+                token,
+                len(self._entries),
+                self._max_sessions,
+            )
             return token
 
     def get(self, token: Optional[str]) -> SessionEntry:
-        """Look up a session by lease token, touching last_active_at on success."""
+        """Look up a session by lease token, marking it seen (liveness) on success.
+
+        Note this only refreshes the liveness clock, not activity: a bare
+        ``/heartbeat`` keeps the session from being reaped as dead but does not
+        postpone the idle timeout. Endpoints that represent real user actions
+        call ``entry.mark_active()`` explicitly (see the lock helpers)."""
         if not token:
             raise HTTPException(status.HTTP_410_GONE, detail="lease token missing")
         with self._mu:
             entry = self._entries.get(token)
             if entry is None:
                 raise HTTPException(status.HTTP_410_GONE, detail="lease expired or unknown")
-        entry.touch()
+        entry.mark_seen()
         return entry
 
     def peek(self, token: Optional[str]) -> SessionEntry:
@@ -181,21 +225,37 @@ class SessionRegistry:
         if entry is None:
             return False
         entry.close()
-        logger.info("released session %s (%d/%d active)", token, len(self._entries), self._max_sessions)
+        logger.info(
+            "released session %s (%d/%d active)",
+            token,
+            len(self._entries),
+            self._max_sessions,
+        )
         return True
 
     def sweep(self) -> int:
-        """Drop sessions idle for more than idle_timeout_seconds. Returns the number reaped."""
+        """Drop sessions that have either gone idle or stopped showing signs of life.
+
+        A session is reaped if it has seen no real user activity for longer than
+        ``idle_timeout_seconds`` (the user walked away) OR if it has not been
+        seen at all for longer than ``liveness_timeout_seconds`` (the client
+        process crashed/disconnected and stopped heartbeating). Returns the
+        number reaped."""
         now = time.monotonic()
-        reaped: list[tuple[str, SessionEntry]] = []
+        reaped: list[tuple[str, SessionEntry, str]] = []
         with self._mu:
             for token, entry in list(self._entries.items()):
-                if (now - entry.last_active_at) > self._idle_timeout_seconds:
-                    self._entries.pop(token, None)
-                    reaped.append((token, entry))
-        for token, entry in reaped:
+                if (now - entry.last_seen_at) > self._liveness_timeout_seconds:
+                    reason = "dead"
+                elif (now - entry.last_active_at) > self._idle_timeout_seconds:
+                    reason = "idle"
+                else:
+                    continue
+                self._entries.pop(token, None)
+                reaped.append((token, entry, reason))
+        for token, entry, reason in reaped:
             entry.close()
-            logger.info("reaped idle session %s", token)
+            logger.info("reaped %s session %s", reason, token)
         return len(reaped)
 
     def close_all(self) -> None:
@@ -215,19 +275,23 @@ def make_app(
     device: torch.device,
     max_sessions: int = 1,
     idle_timeout_seconds: float = 600.0,
+    liveness_timeout_seconds: float = 60.0,
     torch_n_threads: int = 8,
     do_autozoom: bool = True,
+    use_torch_compile: bool = False,
     verbose: bool = False,
     api_key: Optional[str] = None,
-    sweep_interval_seconds: float = 30.0,
+    sweep_interval_seconds: float = 15.0,
 ) -> FastAPI:
     registry = SessionRegistry(
         artifacts=artifacts,
         max_sessions=max_sessions,
         idle_timeout_seconds=idle_timeout_seconds,
+        liveness_timeout_seconds=liveness_timeout_seconds,
         device=device,
         torch_n_threads=torch_n_threads,
         do_autozoom=do_autozoom,
+        use_torch_compile=use_torch_compile,
         verbose=verbose,
     )
     gpu_lock = threading.Lock()
@@ -266,9 +330,12 @@ def make_app(
 
     app = FastAPI(title="nnInteractive Inference Server", lifespan=lifespan)
 
-    # Capability snapshot is computed once and never changes (model is loaded
-    # once at startup and shared by all sessions). We build it from a fresh
-    # session initialized off the artifacts.
+    # Capability snapshot is computed once and never changes (the network
+    # module is loaded once at startup and the same instance is referenced by
+    # every session). We build it from a fresh session initialized off the
+    # artifacts.
+    # This throwaway session only reads static capability metadata and never runs
+    # a forward pass, so there is no point compiling its network reference.
     _capability_session = nnInteractiveInferenceSession(
         device=device,
         use_torch_compile=False,
@@ -370,7 +437,11 @@ def make_app(
         )
 
     def _under_session_lock(entry: SessionEntry, fn):
-        """Run ``fn(session)`` under the session's lock, converting known errors to HTTP 400."""
+        """Run ``fn(session)`` under the session's lock, converting known errors to HTTP 400.
+
+        Every endpoint routed through here represents a real user action, so we
+        also refresh the activity clock (postponing the idle timeout)."""
+        entry.mark_active()
         with entry.lock:
             try:
                 return fn(entry.session)
@@ -379,7 +450,10 @@ def make_app(
 
     def _under_session_and_gpu_lock(entry: SessionEntry, fn):
         """Run ``fn(session)`` under session lock + global GPU lock. Acquisition order
-        is always session-then-gpu to avoid deadlocks."""
+        is always session-then-gpu to avoid deadlocks.
+
+        Like ``_under_session_lock``, this marks real user activity."""
+        entry.mark_active()
         with entry.lock:
             with gpu_lock:
                 try:
@@ -410,6 +484,7 @@ def make_app(
         body = {
             "lease_token": token,
             "idle_timeout_seconds": registry.idle_timeout_seconds,
+            "liveness_timeout_seconds": registry.liveness_timeout_seconds,
             "max_sessions": registry.max_sessions,
         }
         return Response(
@@ -580,7 +655,7 @@ def _build_capability_snapshot(session: nnInteractiveInferenceSession) -> dict:
         "supports_initial_label": bool(session.supports_initial_label),
         "supports_zero_shot_label_refinement": bool(session.supports_zero_shot_label_refinement),
         "preferred_scribble_thickness": session.preferred_scribble_thickness,
-        "interaction_decay": float(session.interaction_decay) if session.interaction_decay is not None else None,
+        "interaction_decay": (float(session.interaction_decay) if session.interaction_decay is not None else None),
         "patch_size": list(cfg.patch_size) if cfg is not None else None,
         "do_autozoom": bool(session.do_autozoom),
         "inference_session_version": session.INFERENCE_SESSION_VERSION,