segresnet-mlx 0.2.1__py3-none-any.whl

segresnet_mlx/network.py

@@ -0,0 +1,195 @@
+"""MONAI 1.6 SegResNet inference graph in MLX, using NCDHW tensors."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from medmlx_core.compensated import CompensatedFloat32Operators, CompensatedFloat48Operators
+
+    from segresnet_mlx._types import MLX, Array
+
+from medmlx_core.layout import require_ncdhw
+from medmlx_core.ops import conv_transpose3d_ncdhw, upsample_trilinear_ncdhw
+from medmlx_core.precision import Float32Operators
+
+from segresnet_mlx.topology import ArithmeticProfile, SegResNetSpec
+
+
+def group_mean(values: Array, mx: MLX) -> Array:
+    """Mean over the last axis of ``(N, G, S)`` values, summed in 256-element tiles.
+
+    Tiling bounds the float32 accumulation depth of a group that spans a whole
+    volume. Returns ``(N, G, 1)``. ``SegResNet`` does not call this (its GroupNorm
+    comes from the ``medmlx_core`` operators); it is exercised by
+    ``tests/test_equivalence.py``.
+    """
+    size = int(values.shape[-1])
+    padded = mx.pad(values, [(0, 0), (0, 0), (0, (-size) % 256)])
+    tiles = mx.reshape(padded, (*values.shape[:2], -1, 256))
+    return mx.sum(mx.sum(tiles, axis=-1), axis=-1, keepdims=True) / size
+
+
+class SegResNet:
+    """Raw-logit forward pass for MONAI's 3-D, group-normalized SegResNet.
+
+    Tensors are channels-first ``(N, C, D, H, W)`` float32 at every layer boundary.
+    ``weights`` uses MONAI's state-dict keys and Torch parameter layouts
+    (convolution weights are ``(out, in, kD, kH, kW)``); ``required_shapes`` in
+    ``segresnet_mlx.weights`` lists them. Inference only: dropout is omitted and
+    no activation is applied to the output.
+
+    ``arithmetic_profile`` selects the accumulation scheme (see the README):
+    ``"torch-order-fp32"`` (default), ``"compensated-fp32"``, or
+    ``"compensated-fp48"``, which needs ``upsample_mode="deconv"`` and raises
+    ``ValueError`` otherwise.
+    """
+
+    def __init__(
+        self,
+        spec: SegResNetSpec,
+        weights: Mapping[str, Array],
+        *,
+        mx: MLX,
+        arithmetic_profile: ArithmeticProfile = "torch-order-fp32",
+    ) -> None:
+        self.spec = spec
+        self._weights = weights
+        self._mx = mx
+        self.arithmetic_profile = arithmetic_profile
+        # Both convolution and normalization accumulation order affect sparse
+        # released-checkpoint logits; correcting only one leaves amplified drift.
+        if arithmetic_profile == "torch-order-fp32":
+            self._operators = Float32Operators(mx, spec.num_groups, spec.norm_eps)
+        elif arithmetic_profile == "compensated-fp32":
+            from medmlx_core.compensated import CompensatedFloat32Operators as AccuracyOperators
+
+            self._operators = AccuracyOperators(mx, spec.num_groups, spec.norm_eps)
+        elif arithmetic_profile == "compensated-fp48":
+            if spec.upsample_mode != "deconv":
+                raise ValueError("Paired precision requires transpose-convolution upsampling")
+            from medmlx_core.compensated import CompensatedFloat48Operators as PairedOperators
+
+            self._operators = PairedOperators(mx, spec.num_groups, spec.norm_eps)
+        else:
+            raise ValueError(f"Unknown SegResNet arithmetic profile {arithmetic_profile!r}")
+
+    def __call__(self, values: Array) -> Array:
+        """Return float32 logits ``(N, out_channels, D, H, W)`` for an ``(N, C, D, H, W)`` input.
+
+        Raises:
+            ValueError: If the input is not rank 5, has the wrong channel count, or a
+                spatial size is not divisible by ``2 ** len(spec.blocks_up)``.
+        """
+        if self.arithmetic_profile == "compensated-fp48":
+            operators = cast("CompensatedFloat48Operators", self._operators)
+            return operators.round(self.paired_forward(values))
+        return self._forward(values)
+
+    def paired_forward(self, values: Array) -> Array:
+        """Run the paired-precision graph without the final rounding.
+
+        Returns ``(2N, out_channels, D, H, W)``: the high words of the ``N`` logit
+        volumes followed by their low words. Their float64 sum is the logit; the
+        caller (``predict_wide_logits``) does that sum on the host.
+        """
+        if self.arithmetic_profile != "compensated-fp48":
+            raise ValueError("Paired logits require the compensated-fp48 profile")
+        operators = cast("CompensatedFloat48Operators", self._operators)
+        return self._forward(operators.pack(values))
+
+    def _forward(self, values: Array) -> Array:
+        require_ncdhw(values, name="SegResNet input")
+        if int(values.shape[1]) != self.spec.in_channels:
+            raise ValueError(
+                f"SegResNet input has {values.shape[1]} channels; expected {self.spec.in_channels}"
+            )
+
+        divisor = 2 ** len(self.spec.blocks_up)
+        if any(int(size) % divisor for size in values.shape[2:]):
+            raise ValueError(f"SegResNet spatial sizes must be divisible by {divisor}")
+        hidden = self.conv(values, "convInit.conv", padding=1)
+        skips: list[Array] = []
+        for stage, block_count in enumerate(self.spec.blocks_down):
+            if stage > 0:
+                hidden = self.conv(hidden, f"down_layers.{stage}.0.conv", padding=1, stride=2)
+            block_offset = 1
+            for block in range(block_count):
+                hidden = self.residual(hidden, f"down_layers.{stage}.{block + block_offset}")
+            skips.append(hidden)
+
+        for stage, block_count in enumerate(self.spec.blocks_up):
+            hidden = self.conv(hidden, f"up_samples.{stage}.0.conv", padding=0)
+            if self.spec.upsample_mode == "deconv":
+                prefix = f"up_samples.{stage}.1.deconv"
+                if self.arithmetic_profile != "torch-order-fp32":
+                    hidden = cast("CompensatedFloat32Operators", self._operators).conv_transpose3d(
+                        hidden, self._weights[f"{prefix}.weight"], self._weights[f"{prefix}.bias"]
+                    )
+                else:
+                    hidden = conv_transpose3d_ncdhw(
+                        hidden,
+                        self._weights[f"{prefix}.weight"],
+                        self._weights[f"{prefix}.bias"],
+                        padding=0,
+                        output_padding=0,
+                        stride=2,
+                        mx=self._mx,
+                    )
+                hidden = self._add(hidden, skips[-(stage + 2)])
+            else:
+                hidden = upsample_trilinear_ncdhw(hidden, mx=self._mx) + skips[-(stage + 2)]
+            for block in range(block_count):
+                hidden = self.residual(hidden, f"up_layers.{stage}.{block}")
+
+        hidden = self._norm(hidden, "conv_final.0")
+        hidden = self._relu(hidden)
+        return self.conv(hidden, "conv_final.2.conv", padding=0, bias=True)
+
+    def residual(self, values: Array, prefix: str) -> Array:
+        """Pre-activation residual block: (norm, ReLU, 3x3x3 conv) twice, plus the input."""
+        hidden = self._norm(values, f"{prefix}.norm1")
+        hidden = self._relu(hidden)
+        hidden = self.conv(hidden, f"{prefix}.conv1.conv", padding=1)
+        hidden = self._norm(hidden, f"{prefix}.norm2")
+        hidden = self._relu(hidden)
+        hidden = self.conv(hidden, f"{prefix}.conv2.conv", padding=1)
+        return self._add(hidden, values)
+
+    def _relu(self, values: Array) -> Array:
+        return (
+            cast("CompensatedFloat48Operators", self._operators).relu(values)
+            if self.arithmetic_profile == "compensated-fp48"
+            else self._mx.maximum(values, 0.0)
+        )
+
+    def _add(self, values: Array, other: Array) -> Array:
+        return (
+            cast("CompensatedFloat48Operators", self._operators).add(values, other)
+            if self.arithmetic_profile == "compensated-fp48"
+            else values + other
+        )
+
+    def _norm(self, values: Array, prefix: str) -> Array:
+        return self._operators.norm(
+            values, self._weights[f"{prefix}.weight"], self._weights[f"{prefix}.bias"]
+        )
+
+    def conv(
+        self,
+        values: Array,
+        prefix: str,
+        *,
+        padding: int,
+        stride: int = 1,
+        bias: bool = False,
+    ) -> Array:
+        """Convolve with ``{prefix}.weight`` (plus ``{prefix}.bias`` if ``bias``)."""
+        return self._operators.conv(
+            values,
+            self._weights[f"{prefix}.weight"],
+            self._weights.get(f"{prefix}.bias") if bias else None,
+            padding=padding,
+            stride=stride,
+        )

segresnet_mlx/predictor.py

@@ -0,0 +1,303 @@
+"""Measured MLX window execution and NumPy full-volume raw-logit inference."""
+
+from __future__ import annotations
+
+import time
+from collections.abc import Callable
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal, cast
+
+from segresnet_mlx._records import BundleConfig, InfererConfig
+from segresnet_mlx._runtime import import_mlx
+
+if TYPE_CHECKING:
+    from ignite.engine import Engine, Events
+
+    from segresnet_mlx._types import MLX, Array, CoreSlidingWindow
+    from segresnet_mlx.network import SegResNet
+
+import numpy as np
+from medmlx_core.errors import InvalidInputError
+from medmlx_core.layout import require_ncdhw
+from medmlx_core.typing import ArrayFactory
+
+from segresnet_mlx._arrays import Float32Array, HostArray
+from segresnet_mlx.medmlx import BUNDLES
+from segresnet_mlx.telemetry import SCHEMA_VERSION, WindowTelemetry
+
+
+@dataclass(frozen=True, slots=True)
+class SlidingWindowConfig:
+    """The bundle's ``SlidingWindowInferer`` settings; other options use MONAI defaults.
+
+    Attributes:
+        roi_size: Window size in voxels, ``(D, H, W)`` of the preprocessed volume.
+        sw_batch_size: Windows sent to the network per call.
+        overlap: Fraction of each window shared with its neighbour, in ``[0, 1)``.
+        mode: Window blending, ``"constant"`` or ``"gaussian"`` (sigma scale 0.125).
+        padding_mode: How a volume smaller than ``roi_size`` is padded: ``"constant"``
+            (zeros), ``"reflect"``, ``"replicate"`` or ``"circular"``.
+    """
+
+    roi_size: tuple[int, int, int]
+    sw_batch_size: int = 1
+    overlap: float = 0.25
+    mode: str = "constant"
+    padding_mode: str = "constant"
+
+    @classmethod
+    def from_bundle_config(cls, config: InfererConfig | BundleConfig) -> SlidingWindowConfig:
+        """Read the staged JSON's inferer section without evaluating bundle expressions.
+
+        Missing keys fall back to MONAI's defaults. Values are taken literally, so a
+        reference such as ``@displayable_configs#sw_overlap`` is not resolved.
+        """
+        inferer = cast(InfererConfig, config.get("inferer", config))
+        return cls(
+            cast(tuple[int, int, int], tuple(inferer["roi_size"])),
+            inferer.get("sw_batch_size", 1),
+            inferer.get("overlap", 0.25),
+            inferer.get("mode", "constant"),
+            inferer.get("padding_mode", "constant"),
+        )
+
+
+# Window settings of each bundle in ``segresnet_mlx.medmlx.BUNDLES``.
+BUNDLE_CONFIGS = {
+    name: SlidingWindowConfig(
+        bundle.roi_size, bundle.sw_batch_size, bundle.overlap, bundle.mode, bundle.padding_mode
+    )
+    for name, bundle in BUNDLES.items()
+}
+
+_PAD_MODES: dict[str, Literal["constant", "reflect", "edge", "wrap"]] = {
+    "constant": "constant",
+    "reflect": "reflect",
+    "replicate": "edge",
+    "circular": "wrap",
+}
+
+
+class BundlePredictor:
+    """Synchronized MLX window callback with an injected network.
+
+    ``forward`` retains the MLX-array callback API; ``predict_logits`` passes
+    synchronized NumPy logits directly to the host stitcher. Telemetry
+    lifecycle remains caller-owned here, allowing an engine to finish only after
+    postprocessing. The module-level ``predict_logits`` owns a direct-call session.
+    """
+
+    schema_version = SCHEMA_VERSION
+    network_label = "SegResNet"
+
+    def __init__(
+        self,
+        network: Callable[[Array], Array],
+        *,
+        telemetry: WindowTelemetry | None = None,
+        mx: MLX | None = None,
+    ) -> None:
+        self.network = network
+        self.telemetry = telemetry
+        try:
+            self.mx = import_mlx() if mx is None else mx
+            self.mx.reset_peak_memory()
+            if telemetry is not None:
+                telemetry.set_metadata(
+                    {
+                        "schema_version": self.schema_version,
+                        "dtype": "float32",
+                        "network_device": str(self.mx.default_device()),
+                        "transforms_and_stitching_device": "cpu",
+                        "compile": False,
+                        "synchronization": "mx.eval and mx.synchronize before host export",
+                        "window_timing_scope": (
+                            "synchronized network, MLX input creation and host export; "
+                            "excludes validation and telemetry"
+                        ),
+                        "peak_memory_scope": "predictor lifetime, including warmup",
+                    }
+                )
+        except BaseException as error:
+            if telemetry is not None:
+                telemetry.fail(error)
+            raise
+
+    def attach(self, engine: Engine, *, events: type[Events]) -> None:
+        """Attach the telemetry (if any) to an Ignite ``engine`` using the caller's ``events``."""
+        if self.telemetry is not None:
+            self.telemetry.attach(engine, events=events)
+
+    def __call__(self, values: HostArray | Array) -> Array:
+        return self.forward(values)
+
+    def forward(self, values: HostArray | Array) -> Array:
+        """Run one finite float32 window and return its logits as an MLX array.
+
+        The logits are exported to host NumPy for the finite check and timing, then
+        wrapped again as an MLX array for callers that expect one.
+        """
+        try:
+            return cast(ArrayFactory, self.mx.array)(self._forward_window(values))
+        except BaseException as error:
+            if self.telemetry is not None:
+                self.telemetry.fail(error)
+            raise
+
+    def _forward_window(self, values: HostArray | Array) -> Float32Array:
+        host = self._host_input(values)
+        self.mx.synchronize()
+        started = time.perf_counter()
+        result = self.network(cast(ArrayFactory, self.mx.array)(host))
+        self.mx.eval(result)
+        self.mx.synchronize()
+        output = cast(Float32Array, np.array(result))
+        elapsed = time.perf_counter() - started
+        if not np.isfinite(output).all():
+            raise RuntimeError(f"{self.network_label} produced non-finite logits")
+        if self.telemetry is not None:
+            self.telemetry.observe(elapsed, self.mx.get_peak_memory())
+        # The allocator cache is kept across windows: every window has the same
+        # shapes, so later windows reuse the buffers instead of reallocating
+        # them. ``predict_logits`` releases the cache once the volume is done.
+        del result
+        return output
+
+    def _host_input(self, values: HostArray | Array) -> Float32Array:
+        host = np.asarray(values)
+        if host.dtype != np.float32 or not np.isfinite(host).all():
+            raise InvalidInputError(f"{self.network_label} input must be finite float32")
+        return cast(Float32Array, host)
+
+    def predict_logits(
+        self, values: HostArray | Array, configuration: SlidingWindowConfig
+    ) -> Float32Array:
+        """Stitch preprocessed ``(N, C, D, H, W)`` float32 data into full-volume logits.
+
+        Returns float32 NumPy ``(N, out_channels, D, H, W)`` at the input spatial
+        size. Volumes smaller than ``roi_size`` are padded and cropped back. Window
+        placement and blending follow MONAI's ``SlidingWindowInferer`` via the
+        shared core. Callers own telemetry completion.
+
+        Raises:
+            InvalidInputError: If ``values`` is not finite float32.
+            ValueError: If ``values`` is not rank 5.
+        """
+        import medmlx_core.sliding_window as core_sliding_window
+
+        sliding_window_inference = cast(
+            "CoreSlidingWindow", core_sliding_window
+        ).sliding_window_inference
+
+        try:
+            host = self._host_input(values)
+            require_ncdhw(host, name="SegResNet input")
+            if self.telemetry is not None:
+                self.telemetry.set_metadata({"sw_batch_size": configuration.sw_batch_size})
+            callback: (
+                Callable[[HostArray | Array], Float32Array]
+                | Callable[[HostArray | Array, list[list[slice]]], Float32Array]
+            ) = self._forward_window
+            use_coordinates = configuration.sw_batch_size == 1
+            if use_coordinates:
+                # MONAI preserves the sliced input's strides for batch size 1;
+                # the shared core concatenates even a single patch. MLX selects
+                # different reduction kernels for strided/contiguous inputs.
+                # Open gap: a core change that skips the single-patch
+                # concatenation would make this workaround unnecessary.
+                padding = [(0, 0), (0, 0)]
+                for size, roi in zip(host.shape[2:], configuration.roi_size, strict=True):
+                    extra = max(roi - size, 0)
+                    padding.append((extra // 2, extra - extra // 2))
+                # Torch padding materializes a contiguous array; unpadded data
+                # retains its original strides. Match the bundle's padding_mode
+                # so this fast path agrees with the core's own padding.
+                pad_kind = _PAD_MODES[configuration.padding_mode]
+                padded = (
+                    np.ascontiguousarray(np.pad(host, padding, mode=pad_kind))
+                    if any(before or after for before, after in padding)
+                    else host
+                )
+
+                def coordinate_callback(
+                    values: HostArray | Array, coordinates: list[list[slice]]
+                ) -> Float32Array:
+                    return self._forward_window(padded[tuple(coordinates[0])])
+
+                callback = coordinate_callback
+
+            return sliding_window_inference(
+                host,
+                configuration.roi_size,
+                configuration.sw_batch_size,
+                callback,
+                overlap=configuration.overlap,
+                mode=configuration.mode,
+                padding_mode=configuration.padding_mode,
+                with_coord=use_coordinates,
+                predictor_input="numpy",
+            )
+        except BaseException as error:
+            if self.telemetry is not None:
+                self.telemetry.fail(error)
+            raise
+        finally:
+            self.mx.clear_cache()
+
+
+def predict_logits(
+    values: HostArray | Array,
+    configuration: SlidingWindowConfig,
+    *,
+    network: Callable[[Array], Array],
+    telemetry: WindowTelemetry | None = None,
+) -> Float32Array:
+    """Return float32 NumPy NCDHW logits, completing optional window telemetry.
+
+    ``values`` is preprocessed, finite float32 ``(N, C, D, H, W)``; the result has the
+    same spatial size. Input preparation, activation, label mapping, image I/O and QC
+    are external.
+    """
+    with telemetry.session() if telemetry is not None else nullcontext():
+        predictor = BundlePredictor(network, telemetry=telemetry)
+        return predictor.predict_logits(values, configuration)
+
+
+def predict_wide_logits(
+    network: SegResNet, values: Float32Array, configuration: SlidingWindowConfig, *, mx: MLX
+) -> Float32Array:
+    """Carry paired Metal logits into FP64 host blending, rounding only at export.
+
+    Preprocessed values and checkpoint parameters remain the same FP32 inputs.
+    The original MONAI float64 network/inferer is the independent oracle for this
+    explicit profile. The learned graph runs on Metal; host work is stitching.
+    """
+    from medmlx_core.sliding_window import sliding_window_inference
+
+    def window(
+        patch: np.ndarray[tuple[int, ...], np.dtype[np.float64]],
+    ) -> np.ndarray[tuple[int, ...], np.dtype[np.float64]]:
+        result = network.paired_forward(cast(ArrayFactory, mx.array)(patch.astype(np.float32)))
+        mx.eval(result)
+        mx.synchronize()
+        packed = np.array(result, dtype=np.float32)
+        batch = packed.shape[0] // 2
+        output = packed[:batch].astype(np.float64) + packed[batch:].astype(np.float64)
+        if not np.isfinite(output).all():
+            raise RuntimeError("SegResNet produced non-finite paired logits")
+        del result
+        mx.clear_cache()
+        return output
+
+    logits = sliding_window_inference(
+        values.astype(np.float64),
+        configuration.roi_size,
+        configuration.sw_batch_size,
+        window,
+        overlap=configuration.overlap,
+        mode=configuration.mode,
+        padding_mode=configuration.padding_mode,
+        predictor_input="numpy",
+    )
+    return logits.astype(np.float32)

segresnet_mlx/telemetry.py

@@ -0,0 +1,166 @@
+"""Bounded progress records and explicit SegResNet inference lifecycle."""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+import tempfile
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING, Literal, Protocol, cast
+
+from segresnet_mlx._records import Json
+
+if TYPE_CHECKING:
+    from ignite.engine import Events
+
+import numpy as np
+
+
+class EventEngine(Protocol):
+    def add_event_handler(self, event_name: Events, handler: Callable[..., None]) -> object: ...
+
+
+SCHEMA_VERSION = "medmlx.segresnet.mlx.v1"
+
+
+class WindowTelemetry:
+    """Atomic JSON progress record for sliding-window inference, times in seconds.
+
+    ``record`` is updated in place and written to ``path`` (parent directories are
+    created). While running it holds window counts and totals, refreshed after the
+    first window and then every 32nd. The raw ``window_seconds`` samples and their
+    p50/p95 are written only when the session finishes: samples stay in memory so the
+    growing list is not rewritten per window.
+    """
+
+    def __init__(self, path: Path, record: dict[str, Json]) -> None:
+        self.path = path
+        self.record = record
+        self._metadata = {"schema_version": SCHEMA_VERSION, **record}
+        self.timings: list[float] = []
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.begin()
+
+    def set_metadata(self, values: dict[str, Json]) -> None:
+        """Merge run-level metadata that is kept across ``begin`` calls."""
+        self._metadata.update(values)
+        self.record.update(values)
+        self._write()
+
+    def begin(self) -> None:
+        """Discard samples and start a new running record (``__init__`` calls this)."""
+        self.timings = []
+        self.record.clear()
+        self.record.update(self._metadata)
+        self.record.update(
+            status="running",
+            window_count=0,
+            window_total_seconds=0.0,
+            window_max_seconds=None,
+            peak_mlx_bytes=None,
+            window_samples_complete=False,
+        )
+        self._write()
+
+    def observe(self, seconds: float, peak_bytes: int | None) -> None:
+        """Add one window's duration (seconds) and the MLX peak memory so far (bytes).
+
+        Raises:
+            RuntimeError: If the session is closed.
+            ValueError: If ``seconds`` is negative or not finite.
+        """
+        if self.record["status"] != "running":
+            raise RuntimeError("SegResNet telemetry is closed; begin a new session")
+        if not math.isfinite(seconds) or seconds < 0:
+            raise ValueError("Window duration must be finite and nonnegative")
+        self.timings.append(seconds)
+        self.record["window_count"] = len(self.timings)
+        self.record["window_total_seconds"] = (
+            cast(float, self.record["window_total_seconds"]) + seconds
+        )
+        self.record["window_max_seconds"] = max(
+            cast(float | None, self.record["window_max_seconds"]) or 0, seconds
+        )
+        if peak_bytes is not None:
+            self.record["peak_mlx_bytes"] = max(
+                cast(int | None, self.record["peak_mlx_bytes"]) or 0, peak_bytes
+            )
+        # Constant-size snapshots. Quantiles and full raw samples are final-only.
+        if len(self.timings) == 1 or len(self.timings) % 32 == 0:
+            self._write()
+
+    def finish(self, status: Literal["completed", "failed", "interrupted"] = "completed") -> None:
+        """Close the record with ``status`` and write the raw samples and quantiles.
+
+        Does nothing if the record already has that status or has failed, or if a
+        completed finish follows an interruption.
+        """
+        previous = self.record["status"]
+        if previous in {status, "failed"} or (previous == "interrupted" and status == "completed"):
+            return
+        p50, p95 = np.percentile(self.timings, [50, 95]) if self.timings else (None, None)
+        self.record.update(
+            status=status,
+            window_seconds=list(self.timings),
+            window_samples_complete=True,
+            window_p50_seconds=None if p50 is None else float(p50),
+            window_p95_seconds=None if p95 is None else float(p95),
+        )
+        self._write()
+
+    def fail(self, error: BaseException) -> None:
+        """Record ``error``'s type and finish as failed, without masking ``error``."""
+        self.record["error_type"] = type(error).__name__
+        try:
+            self.finish("failed")
+        except Exception as record_error:
+            # An unavailable output directory must not replace the inference error.
+            error.add_note(f"Could not finalize SegResNet telemetry: {record_error!r}")
+
+    @contextmanager
+    def session(self) -> Generator[None]:
+        """Use around direct inferer calls that are not owned by a MONAI engine."""
+        self.begin()
+        try:
+            yield
+            self.finish()
+        except BaseException as error:
+            self.fail(error)
+            raise
+
+    def attach(self, engine: object, *, events: type[Events]) -> None:
+        """Attach caller-provided engine events before handlers that re-raise failures."""
+        attached = cast(EventEngine, engine)
+        attached.add_event_handler(events.STARTED, self._started)
+        attached.add_event_handler(events.COMPLETED, self._completed)
+        attached.add_event_handler(events.EXCEPTION_RAISED, self._failed)
+        for event in (events.TERMINATE, events.TERMINATE_SINGLE_EPOCH, events.INTERRUPT):
+            attached.add_event_handler(event, self._interrupted)
+
+    def _started(self, engine: object) -> None:
+        self.begin()
+
+    def _completed(self, engine: object) -> None:
+        self.finish()
+
+    def _failed(self, engine: object, error: BaseException) -> None:
+        self.fail(error)
+        # Registering EXCEPTION_RAISED otherwise makes Ignite swallow the error.
+        raise error
+
+    def _interrupted(self, engine: object, **kwargs: object) -> None:
+        self.finish("interrupted")
+
+    def _write(self) -> None:
+        # Atomic visibility for progress readers, not a crash-durability promise.
+        descriptor, name = tempfile.mkstemp(prefix=f".{self.path.name}.", dir=self.path.parent)
+        try:
+            with os.fdopen(descriptor, "w", encoding="utf-8") as handle:
+                json.dump(self.record, handle, indent=2, allow_nan=False)
+                handle.write("\n")
+            os.replace(name, self.path)
+        finally:
+            Path(name).unlink(missing_ok=True)