spikeforge-server 0.1.0__tar.gz

spikeforge_server-0.1.0/LICENSE

@@ -0,0 +1,30 @@
+BSD 3-Clause License
+
+Copyright (c) 2018 Capsize LLC
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+* Neither the name of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+

spikeforge_server-0.1.0/PKG-INFO

@@ -0,0 +1,130 @@
+Metadata-Version: 2.4
+Name: spikeforge-server
+Version: 0.1.0
+Summary: FastAPI + WebSocket server adapter for the spikeforge dashboard.
+Author-email: Capsize LLC <contact@capsizegames.com>
+Maintainer-email: Capsize LLC <contact@capsizegames.com>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/capsize-games/spikeforge
+Project-URL: Repository, https://github.com/capsize-games/spikeforge
+Keywords: spiking neural networks,fastapi,websockets,dashboard
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: fastapi>=0.110
+Requires-Dist: uvicorn[standard]>=0.27
+Requires-Dist: websockets>=12.0
+Requires-Dist: pydantic>=2.5
+Requires-Dist: spikeforge~=0.3.0
+Requires-Dist: spikeforge-targets~=0.1.0
+Requires-Dist: spikeforge-hub~=0.1.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0; extra == "dev"
+Requires-Dist: ruff>=0.0.280; extra == "dev"
+Dynamic: license-file
+
+# spikeforge
+
+[![CI](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml/badge.svg)](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml)
+[![Status: pre-1.0](https://img.shields.io/badge/status-pre--1.0-orange.svg)](OPEN_SOURCE_CHECKLIST.md)
+[![License: BSD-3-Clause](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](LICENSE)
+[![Python 3.10–3.13](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-3776AB.svg)](https://www.python.org/downloads/)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg)](https://github.com/capsize-games/spikeforge/blob/main/CONTRIBUTING.md)
+[![Docs](https://img.shields.io/badge/docs-long--form%20reference-blue.svg)](documentation/README.md)
+
+A spiking-neural-network (SNN) toolkit built on
+[snnTorch](https://snntorch.readthedocs.io/) and PyTorch. Loads MNIST-style
+and neuromorphic event datasets, encodes them into rate, latency, delta, and
+random spikes, and trains, validates, exports, and deploys LIF networks —
+with a live browser dashboard served over WebSockets.
+
+> **Pre-1.0 and unpublished.** `install.sh` is the supported path today;
+> PyPI ships once the distributions are published. Before trusting any
+> number this produces, read
+> [Implications and boundaries](documentation/implications-and-boundaries.md).
+
+## Quickstart
+
+```bash
+git clone https://github.com/capsize-games/spikeforge.git
+cd spikeforge
+docker compose up --build
+```
+
+Open <http://localhost:8877> — the dashboard connects to the WebSocket on
+the same host and port. No separate backend or proxy to run.
+
+Prefer a local install, a headless example, or the CLI tools instead? See
+[Usage](documentation/usage.md) and [Quickstart](documentation/quickstart.md)
+for every path (`./install.sh`, local dev with Vite, `examples/`, and the
+eight `spikeforge-*` console scripts).
+
+## Features
+
+- **Encoding** — rate, latency, delta, and random spike coders.
+- **Training** — fully-connected and convolutional LIF networks with
+  surrogate-gradient cross-entropy, checkpointing, and opt-in AMP / gradient
+  checkpointing / truncated BPTT / multi-GPU.
+- **Topologies** — `fc_legacy`, `fc_small`, `conv_net`, `recurrent_net`, plus
+  the sequence presets `sequence_mlp` and `sequence_attn`.
+- **Datasets** — MNIST, Fashion-MNIST, KMNIST, QMNIST, USPS, EMNIST,
+  CIFAR-10, and (via the `events` extra) N-MNIST, DVS128 Gesture,
+  CIFAR10-DVS, and Spiking Speech Commands.
+- **Interpreter spine** — NIR export, an independent NIR interpreter, and
+  numerical drift validation.
+- **Introspection** — educational-mode `U[t]`/`I[t]`/`S[t]` traces,
+  trajectory metrics, and surrogate-derivative curves.
+- **Deployment** — a capability matrix, weight quantization, energy
+  accounting, and executable `reference`, `norse`, and `lava_loihi2`
+  backends.
+- **Model hub** — a curated, offline-first catalog plus optional live
+  Hugging Face search.
+- **Dashboard** — a React + TypeScript UI with training, introspection,
+  analysis, targets, energy, and hub panels, and seven guided walkthroughs.
+
+## Packages
+
+This repository is a single workspace that publishes four distributions:
+
+| Distribution | Import root | Purpose |
+|---|---|---|
+| `spikeforge` | `spikeforge` | Core package: encoders, topologies, training, simulator, NIR bridge, tracking |
+| `spikeforge-targets` | `spikeforge_targets` | Deployment targets, quantization, energy accounting, sparse event runtime |
+| `spikeforge-hub` | `spikeforge_hub` | Curated model hub and optional Hugging Face access |
+| `spikeforge-server` | `server` | FastAPI + WebSocket server and dashboard hosting |
+
+## Documentation
+
+This README stays short on purpose.
+[`documentation/`](documentation/README.md) is the full reference — install
+paths, the CLI tools, architecture, module layout, and the dev workflow —
+written for contributors and coding agents alike. Also see
+[COOKBOOK.md](COOKBOOK.md) for copy-pasteable recipes,
+[examples/](examples/) for runnable end-to-end scripts, and
+[plans/](plans/) for design documents and the roadmap.
+
+See [CONTRIBUTING.md](https://github.com/capsize-games/spikeforge/blob/main/CONTRIBUTING.md)
+and [rules.md](rules.md) before opening a pull request.
+
+## Citing
+
+If spikeforge is useful in your research, please cite it — see
+[CITATION.cff](CITATION.cff) (GitHub renders a "Cite this repository"
+button from it automatically).
+
+## License
+
+Released under the **BSD 3-Clause License** — see [LICENSE](LICENSE) and
+[AUTHORS](AUTHORS).

spikeforge_server-0.1.0/README.md

@@ -0,0 +1,94 @@
+# spikeforge
+
+[![CI](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml/badge.svg)](https://github.com/capsize-games/spikeforge/actions/workflows/ci.yml)
+[![Status: pre-1.0](https://img.shields.io/badge/status-pre--1.0-orange.svg)](OPEN_SOURCE_CHECKLIST.md)
+[![License: BSD-3-Clause](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](LICENSE)
+[![Python 3.10–3.13](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12%20%7C%203.13-3776AB.svg)](https://www.python.org/downloads/)
+[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
+[![PRs Welcome](https://img.shields.io/badge/PRs-welcome-brightgreen.svg)](https://github.com/capsize-games/spikeforge/blob/main/CONTRIBUTING.md)
+[![Docs](https://img.shields.io/badge/docs-long--form%20reference-blue.svg)](documentation/README.md)
+
+A spiking-neural-network (SNN) toolkit built on
+[snnTorch](https://snntorch.readthedocs.io/) and PyTorch. Loads MNIST-style
+and neuromorphic event datasets, encodes them into rate, latency, delta, and
+random spikes, and trains, validates, exports, and deploys LIF networks —
+with a live browser dashboard served over WebSockets.
+
+> **Pre-1.0 and unpublished.** `install.sh` is the supported path today;
+> PyPI ships once the distributions are published. Before trusting any
+> number this produces, read
+> [Implications and boundaries](documentation/implications-and-boundaries.md).
+
+## Quickstart
+
+```bash
+git clone https://github.com/capsize-games/spikeforge.git
+cd spikeforge
+docker compose up --build
+```
+
+Open <http://localhost:8877> — the dashboard connects to the WebSocket on
+the same host and port. No separate backend or proxy to run.
+
+Prefer a local install, a headless example, or the CLI tools instead? See
+[Usage](documentation/usage.md) and [Quickstart](documentation/quickstart.md)
+for every path (`./install.sh`, local dev with Vite, `examples/`, and the
+eight `spikeforge-*` console scripts).
+
+## Features
+
+- **Encoding** — rate, latency, delta, and random spike coders.
+- **Training** — fully-connected and convolutional LIF networks with
+  surrogate-gradient cross-entropy, checkpointing, and opt-in AMP / gradient
+  checkpointing / truncated BPTT / multi-GPU.
+- **Topologies** — `fc_legacy`, `fc_small`, `conv_net`, `recurrent_net`, plus
+  the sequence presets `sequence_mlp` and `sequence_attn`.
+- **Datasets** — MNIST, Fashion-MNIST, KMNIST, QMNIST, USPS, EMNIST,
+  CIFAR-10, and (via the `events` extra) N-MNIST, DVS128 Gesture,
+  CIFAR10-DVS, and Spiking Speech Commands.
+- **Interpreter spine** — NIR export, an independent NIR interpreter, and
+  numerical drift validation.
+- **Introspection** — educational-mode `U[t]`/`I[t]`/`S[t]` traces,
+  trajectory metrics, and surrogate-derivative curves.
+- **Deployment** — a capability matrix, weight quantization, energy
+  accounting, and executable `reference`, `norse`, and `lava_loihi2`
+  backends.
+- **Model hub** — a curated, offline-first catalog plus optional live
+  Hugging Face search.
+- **Dashboard** — a React + TypeScript UI with training, introspection,
+  analysis, targets, energy, and hub panels, and seven guided walkthroughs.
+
+## Packages
+
+This repository is a single workspace that publishes four distributions:
+
+| Distribution | Import root | Purpose |
+|---|---|---|
+| `spikeforge` | `spikeforge` | Core package: encoders, topologies, training, simulator, NIR bridge, tracking |
+| `spikeforge-targets` | `spikeforge_targets` | Deployment targets, quantization, energy accounting, sparse event runtime |
+| `spikeforge-hub` | `spikeforge_hub` | Curated model hub and optional Hugging Face access |
+| `spikeforge-server` | `server` | FastAPI + WebSocket server and dashboard hosting |
+
+## Documentation
+
+This README stays short on purpose.
+[`documentation/`](documentation/README.md) is the full reference — install
+paths, the CLI tools, architecture, module layout, and the dev workflow —
+written for contributors and coding agents alike. Also see
+[COOKBOOK.md](COOKBOOK.md) for copy-pasteable recipes,
+[examples/](examples/) for runnable end-to-end scripts, and
+[plans/](plans/) for design documents and the roadmap.
+
+See [CONTRIBUTING.md](https://github.com/capsize-games/spikeforge/blob/main/CONTRIBUTING.md)
+and [rules.md](rules.md) before opening a pull request.
+
+## Citing
+
+If spikeforge is useful in your research, please cite it — see
+[CITATION.cff](CITATION.cff) (GitHub renders a "Cite this repository"
+button from it automatically).
+
+## License
+
+Released under the **BSD 3-Clause License** — see [LICENSE](LICENSE) and
+[AUTHORS](AUTHORS).

spikeforge_server-0.1.0/pyproject.toml

@@ -0,0 +1,86 @@
+# PEP 621 packaging authority for the server distribution (ARCH-0001 Phase
+# 1b).
+#
+# The server distribution owns the top-level ``server`` package. It promotes
+# the dependencies that used to be the core ``web`` extra to its own base
+# dependencies, so core stays headless and the server is the only thing that
+# pulls FastAPI/uvicorn/websockets/pydantic.
+#
+# ``server/`` stays at its current repository path and is exposed to the build
+# via the ``server`` symlink next to this file (see the ``packages.find`` table
+# below). The package discovery here depends only on the top-level ``server``
+# package, never on core.
+
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "spikeforge-server"
+version = "0.1.0"
+description = "FastAPI + WebSocket server adapter for the spikeforge dashboard."
+# ``README.md`` and ``LICENSE`` are committed symlinks to the repository-root
+# files, mirroring the core distribution's metadata approach.
+readme = "README.md"
+requires-python = ">=3.10"
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+authors = [
+    { name = "Capsize LLC", email = "contact@capsizegames.com" },
+]
+maintainers = [
+    { name = "Capsize LLC", email = "contact@capsizegames.com" },
+]
+keywords = [
+    "spiking neural networks",
+    "fastapi",
+    "websockets",
+    "dashboard",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "fastapi>=0.110",
+    "uvicorn[standard]>=0.27",
+    "websockets>=12.0",
+    "pydantic>=2.5",
+    "spikeforge~=0.3.0",
+    "spikeforge-targets~=0.1.0",
+    "spikeforge-hub~=0.1.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "pytest-cov>=4.0",
+    "ruff>=0.0.280",
+]
+
+[project.scripts]
+spikeforge-server = "server.__main__:main"
+
+[project.urls]
+Homepage = "https://github.com/capsize-games/spikeforge"
+Repository = "https://github.com/capsize-games/spikeforge"
+
+[tool.setuptools]
+zip-safe = false
+
+[tool.setuptools.packages.find]
+# Mirror the core distribution: the `server` import root is exposed next to
+# this file as a symlink to `../../server` so the default via-sdist
+# `python -m build` produces a complete wheel. Discovery depends only on this
+# top-level package, never on core.
+where = ["."]
+include = ["server*"]
+exclude = ["spikeforge*", "spikeforge_targets*", "spikeforge_hub*", "tests*"]

spikeforge_server-0.1.0/server/__init__.py

@@ -0,0 +1 @@
+"""FastAPI server exposing SNN encoding experiments over WebSocket."""

spikeforge_server-0.1.0/server/__main__.py

@@ -0,0 +1,55 @@
+"""Run the FastAPI server with uvicorn: ``python -m server``."""
+
+import argparse
+from typing import Optional, Sequence
+
+import uvicorn
+
+#: Default bind address and port (the single-port dashboard contract on :8877).
+DEFAULT_HOST = "127.0.0.1"
+DEFAULT_PORT = 8877
+
+
+def _parse_args(argv: Optional[Sequence[str]]) -> argparse.Namespace:
+    """Parse the optional bind arguments, keeping the historical defaults."""
+    parser = argparse.ArgumentParser(
+        prog="spikeforge-server",
+        description="Run the spikeforge dashboard/WebSocket server.",
+    )
+    parser.add_argument(
+        "--host",
+        default=DEFAULT_HOST,
+        help=f"bind address (default: {DEFAULT_HOST})",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=DEFAULT_PORT,
+        help=f"bind port (default: {DEFAULT_PORT})",
+    )
+    parser.add_argument(
+        "--no-reload",
+        action="store_true",
+        help="disable uvicorn's auto-reload watcher",
+    )
+    return parser.parse_args(argv)
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    """Start the dashboard server with uvicorn.
+
+    ``python -m server`` and the ``spikeforge-server`` console script both
+    land here; the defaults are the historical ``127.0.0.1:8877`` with
+    auto-reload, matching the project's single-port contract.
+    """
+    args = _parse_args(argv)
+    uvicorn.run(
+        "server.app:app",
+        host=args.host,
+        port=args.port,
+        reload=not args.no_reload,
+    )
+
+
+if __name__ == "__main__":
+    main()

spikeforge_server-0.1.0/server/animation.py

@@ -0,0 +1,76 @@
+"""Per-step hidden-layer animation frames and their availability message.
+
+The frames come from the loaded model's hidden stage. Without a model there
+is nothing honest to animate, so availability and its named reason are
+reported to the client rather than a fabricated frame; with no model the
+input-frame stream stays exactly as it was.
+"""
+
+from typing import Any, List, Optional, Tuple
+
+import torch
+from fastapi import WebSocket
+
+from server.messages import send_locked
+from server.session import Session
+from spikeforge.network.hidden_frames import hidden_frame_series
+
+#: Named reasons the hidden animation cannot run.
+NO_MODEL = "no model loaded; train or load one"
+NO_SAMPLE = "no sample configured"
+
+
+def _on_device(net: Any, spikes: torch.Tensor) -> torch.Tensor:
+    """Move ``spikes`` onto the network's parameter device."""
+    param = next(net.parameters(), None)
+    return spikes if param is None else spikes.to(param.device)
+
+
+def hidden_series(session: Session) -> Tuple[Optional[List[Any]], str]:
+    """Return the per-step hidden frames and a reason when unavailable."""
+    engine = session.training.engine
+    if engine is None:
+        return None, NO_MODEL
+    sample = session.engine
+    if sample is None:
+        return None, NO_SAMPLE
+    try:
+        spikes = _on_device(engine.net, sample.spike_input())
+        return hidden_frame_series(engine.net, spikes), ""
+    except Exception as exc:
+        return None, f"{type(exc).__name__}: {exc}"
+
+
+async def send_animation_state(
+    ws: WebSocket, session: Session, available: bool, reason: str
+) -> None:
+    """Report whether the hidden animation is available, and why not."""
+    await send_locked(ws, session, {
+        "type": "animation_state",
+        "payload": {"available": available, "reason": reason,
+                    "source": "hidden"},
+    })
+
+
+async def emit_hidden_frame(
+    ws: WebSocket, session: Session, frames: List[Any], step: int
+) -> None:
+    """Send one step's hidden-layer frame as a ``spike_frame``."""
+    index = min(max(int(step), 0), len(frames) - 1)
+    await send_locked(ws, session, {
+        "type": "spike_frame",
+        "payload": frames[index],
+        "step": index,
+        "source": "hidden",
+    })
+
+
+async def prepare(
+    ws: WebSocket, session: Session, animate: bool
+) -> List[Any]:
+    """Return the hidden frame series, reporting availability once."""
+    if not animate:
+        return []
+    frames, reason = hidden_series(session)
+    await send_animation_state(ws, session, frames is not None, reason)
+    return frames or []

spikeforge_server-0.1.0/server/app.py

@@ -0,0 +1,154 @@
+"""FastAPI application exposing the encoding engine over WebSocket."""
+
+import asyncio
+from typing import Any, Dict, Optional
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+
+from server.downloads import manager
+from server.handlers import dispatch
+from server.hub_downloads import manager as hub_manager
+from server.messages import send_locked
+from server.protocol_version import PROTOCOL_MAJOR, PROTOCOL_VERSION
+from server.schemas import ClientMessage
+from server.session import Session
+from server.web import mount_client
+from spikeforge.runtime import device as device_mod
+
+app = FastAPI(title="spikeforge server")
+
+# Warm the CUDA context once at startup so training/benchmarks don't stall.
+device_mod.prime()
+
+# Serve the built React app (no-op when client/dist is absent).
+mount_client(app)
+
+# One session per connected client.
+_sessions: Dict[int, Session] = {}
+
+#: Payload code returned when an inbound version is missing or mismatched.
+VERSION_MISMATCH = "protocol_version_mismatch"
+
+
+def _version_error(raw: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    """Return a mismatch payload, or None when the inbound version fits.
+
+    From Phase 2 onward ``protocol_version`` is required: a message that
+    omits it -- or carries one whose MAJOR component differs from ours -- is
+    rejected with ``payload.code = "protocol_version_mismatch"``.
+    """
+    incoming = raw.get("protocol_version")
+    if incoming is None:
+        return {
+            "code": VERSION_MISMATCH,
+            "message": (
+                "client message is missing protocol_version; server "
+                f"protocol is {PROTOCOL_VERSION}"
+            ),
+            "client": None,
+            "server": PROTOCOL_VERSION,
+        }
+    major = str(incoming).split(".", 1)[0]
+    if major == PROTOCOL_MAJOR:
+        return None
+    return {
+        "code": VERSION_MISMATCH,
+        "message": (
+            f"client protocol {incoming!r} is incompatible with server "
+            f"protocol {PROTOCOL_VERSION}"
+        ),
+        "client": incoming,
+        "server": PROTOCOL_VERSION,
+    }
+
+
+async def _drain_training(ws: WebSocket, session: Session) -> None:
+    """Forward worker-produced training messages to the socket."""
+    while True:
+        message = await session.training.queue.get()
+        await send_locked(ws, session, message)
+
+
+async def _dispatch_inbox(ws: WebSocket, session: Session) -> None:
+    """Process queued client messages one at a time."""
+    while True:
+        message = await session.inbox.get()
+        try:
+            await dispatch(ws, session, message)
+        except Exception as exc:  # keep the connection alive on bad input
+            await send_locked(ws, session, {
+                "type": "error", "payload": str(exc),
+            })
+
+
+async def _cleanup(
+    session: Session,
+    drain: asyncio.Task,
+    worker: asyncio.Task,
+    session_id: int,
+) -> None:
+    """Stop background work and forget the session."""
+    session.training.stop()
+    drain.cancel()
+    worker.cancel()
+    await session.cancel()
+    _sessions.pop(session_id, None)
+
+
+async def _serve(ws: WebSocket, session: Session) -> None:
+    """Read client messages, giving downloads their own cancel path."""
+    while True:
+        raw = await ws.receive_json()
+        version_error = _version_error(raw)
+        if version_error is not None:
+            await send_locked(ws, session, {
+                "type": "error", "payload": version_error,
+            })
+            continue
+        try:
+            message = ClientMessage.model_validate(raw)
+        except Exception as exc:  # keep the connection alive on bad input
+            await send_locked(ws, session, {
+                "type": "error", "payload": str(exc),
+            })
+            continue
+        # Cancels must bypass the queue so they land while a download blocks
+        # the dispatcher on its progress stream.
+        if message.type == "cancel_download":
+            manager.cancel()
+            continue
+        if message.type == "hub_cancel":
+            hub_manager.cancel()
+            await send_locked(ws, session, {
+                "type": "hub_download_state",
+                "payload": hub_manager.snapshot(),
+            })
+            continue
+        await session.inbox.put(message)
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(ws: WebSocket) -> None:
+    """Accept a client, run its session, and always clean up."""
+    await ws.accept()
+    session_id = id(ws)
+    session = Session(asyncio.get_running_loop())
+    _sessions[session_id] = session
+    drain = asyncio.create_task(_drain_training(ws, session))
+    worker = asyncio.create_task(_dispatch_inbox(ws, session))
+    try:
+        await _serve(ws, session)
+    except WebSocketDisconnect:
+        pass
+    except Exception as exc:  # surface unexpected errors to the client
+        await send_locked(ws, session, {
+            "type": "error", "payload": str(exc),
+        })
+    finally:
+        await _cleanup(session, drain, worker, session_id)
+
+
+@app.get("/health")
+async def health() -> Dict[str, str]:
+    """Liveness probe for container/orchestrator health checks."""
+    return {"status": "ok"}