spikeforge-server 0.1.0__py3-none-any.whl

server/__init__.py

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

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()

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 []

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"}

server/backend_handlers.py

@@ -0,0 +1,59 @@
+"""Handler for the ``deploy_run`` backend-execution WebSocket action.
+
+Mirrors :mod:`server.target_handlers`: it stays outside :mod:`server.handlers`
+so that module keeps within the line limit. The capability view is unchanged
+(it is still the ``deployment_report`` action); ``deploy_run`` adds the
+*executed* view, compiling the target-ready graph and running it on the
+backend. An unknown target is reported on the existing ``error`` channel and a
+missing sample degrades to an error-shaped backend payload with a named
+reason, never a raise.
+"""
+
+from typing import Any, Dict
+
+from fastapi import WebSocket
+
+from server.backend_payloads import backend_run_payload
+from server.messages import send_backend_run, send_locked
+from server.schemas import ClientMessage
+from server.session import Session
+from server.target_handlers import DEFAULT_TARGET, shaped_sample, target_pair
+from spikeforge_targets.registry import target_names
+
+
+async def _resolve(
+    ws: WebSocket, session: Session, message: ClientMessage
+) -> Any:
+    """Return the active ``(spec, module)`` pair or report an error."""
+    try:
+        return target_pair(session, message.train)
+    except ValueError as exc:
+        await send_locked(ws, session, {"type": "error", "payload": str(exc)})
+        return None
+
+
+async def handle_deploy_run(
+    ws: WebSocket, session: Session, message: ClientMessage
+) -> None:
+    """Compile and run the active topology on the named backend."""
+    target = message.name or DEFAULT_TARGET
+    if target not in target_names():
+        await send_locked(ws, session, {
+            "type": "error", "payload": f"unknown target: {target!r}",
+        })
+        return
+    pair = await _resolve(ws, session, message)
+    if pair is None:
+        return
+    spec, module = pair
+    payload: Dict[str, Any] = backend_run_payload(
+        spec, module, target, shaped_sample(session, spec)
+    )
+    await send_backend_run(ws, session, payload)
+
+
+async def dispatch_backend(
+    ws: WebSocket, session: Session, message: ClientMessage
+) -> None:
+    """Route the backend-execution action."""
+    await handle_deploy_run(ws, session, message)

server/backend_payloads.py

@@ -0,0 +1,53 @@
+"""JSON-able payloads for the backend-run WebSocket action.
+
+The payload is the :meth:`BackendResult.to_dict` of an executed (or refused)
+backend run. A missing sample or an export failure is reported as an honest
+``error`` result with a named reason rather than a raised exception, matching
+the deployment-report handler's degrade-don't-raise contract.
+"""
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from spikeforge.nir_bridge.exporter import to_nir
+from spikeforge.topology.spec import TopologySpec
+from spikeforge_targets.backends import STATUS_ERROR, compile_run
+from spikeforge_targets.backends.result import BackendResult
+
+#: Reason reported when no shaped sample is available to run.
+NO_SAMPLE = "no sample is loaded; select a sample before running a backend"
+
+
+def _device_input(module: Any, spikes: torch.Tensor) -> torch.Tensor:
+    """Move ``spikes`` onto the module's device for the backend run."""
+    param = next(module.parameters(), None)
+    return spikes if param is None else spikes.to(param.device)
+
+
+def _error(target: str, reason: str) -> Dict[str, Any]:
+    """Return an error-shaped backend payload with a named reason."""
+    return BackendResult(
+        target=target, status=STATUS_ERROR, notes=(reason,)
+    ).to_dict()
+
+
+def backend_run_payload(
+    spec: TopologySpec,
+    module: Any,
+    target: str,
+    spikes: Optional[torch.Tensor] = None,
+) -> Dict[str, Any]:
+    """Return the backend run payload for a spec, module, and shaped spikes.
+
+    ``spikes`` must already be shaped for ``spec``; without them the payload
+    is an honest error rather than a fabricated run.
+    """
+    if spikes is None:
+        return _error(target, NO_SAMPLE)
+    try:
+        graph = to_nir(spec, module)
+        result = compile_run(target, graph, _device_input(module, spikes))
+    except Exception as exc:
+        return _error(target, f"{type(exc).__name__}: {exc}")
+    return result.to_dict()

server/download_errors.py

@@ -0,0 +1,5 @@
+"""Errors raised by the dataset download service."""
+
+
+class DownloadCancelledError(Exception):
+    """Raised when a dataset download is cancelled by the user."""

server/downloads.py

@@ -0,0 +1,129 @@
+"""Dataset downloads: isolated worker process, progress, and cancellation."""
+
+import asyncio
+import os
+import subprocess
+import sys
+from typing import Any, Awaitable, Callable, Dict, Optional, Set
+
+from server.download_errors import DownloadCancelledError
+from spikeforge.config import DATA_DIR
+from spikeforge.data.datasets import build_dataset
+
+_POLL_SECONDS = 0.4
+
+Emit = Callable[[Dict[str, Any]], Awaitable[None]]
+
+
+def _spawn(dataset: str, train: bool) -> subprocess.Popen:
+    """Start the isolated downloader child process."""
+    return subprocess.Popen([
+        sys.executable,
+        "-m",
+        "spikeforge.data.download_cli",
+        dataset,
+        "1" if train else "0",
+    ])
+
+
+def _dir_bytes(path: str) -> int:
+    """Sum the sizes of every file under ``path`` (best effort)."""
+    total = 0
+    for root, _dirs, files in os.walk(path):
+        for name in files:
+            try:
+                total += os.path.getsize(os.path.join(root, name))
+            except OSError:
+                continue
+    return total
+
+
+class DownloadManager:
+    """Prepare datasets off the event loop and stream progress to a client."""
+
+    def __init__(self) -> None:
+        """Start idle, with an empty set of already-prepared datasets."""
+        self._prepared: Set[str] = set()
+        self._process: Optional[subprocess.Popen] = None
+        self._dataset = ""
+        self._status = "idle"
+        self._bytes = 0
+        self._baseline = 0
+
+    def snapshot(self) -> Dict[str, Any]:
+        """Return the current download state for the client."""
+        return {
+            "dataset": self._dataset,
+            "status": self._status,
+            "bytes": self._bytes,
+        }
+
+    def cancel(self) -> None:
+        """Terminate the worker process if a download is in flight."""
+        if self._process is not None and self._process.poll() is None:
+            self._process.terminate()
+            self._status = "cancelled"
+
+    async def ensure(self, dataset: str, train: bool, emit: Emit) -> None:
+        """Ensure ``dataset`` is on disk, emitting progress until it is."""
+        if dataset in self._prepared:
+            return
+        loop = asyncio.get_running_loop()
+        if await loop.run_in_executor(None, self._loads, dataset, train):
+            self._prepared.add(dataset)
+            return
+        await self._download(loop, dataset, train, emit)
+        self._prepared.add(dataset)
+
+    @staticmethod
+    def _loads(dataset: str, train: bool) -> bool:
+        """Return True when the dataset is already available locally."""
+        try:
+            build_dataset(dataset, train=train, download=False)
+            return True
+        except (RuntimeError, OSError, ValueError, KeyError):
+            return False
+
+    async def _download(
+        self,
+        loop: asyncio.AbstractEventLoop,
+        dataset: str,
+        train: bool,
+        emit: Emit,
+    ) -> None:
+        """Run the worker process, polling progress until it exits."""
+        self._dataset = dataset
+        self._status = "downloading"
+        self._baseline = _dir_bytes(DATA_DIR)
+        self._bytes = 0
+        await emit(self.snapshot())
+        self._process = _spawn(dataset, train)
+        await self._poll(loop, emit)
+        await self._finish(emit)
+
+    async def _poll(
+        self, loop: asyncio.AbstractEventLoop, emit: Emit
+    ) -> None:
+        """Emit progress snapshots until the worker process exits."""
+        while self._process is not None and self._process.poll() is None:
+            current = await loop.run_in_executor(None, _dir_bytes, DATA_DIR)
+            self._bytes = max(0, current - self._baseline)
+            await emit(self.snapshot())
+            await asyncio.sleep(_POLL_SECONDS)
+
+    async def _finish(self, emit: Emit) -> None:
+        """Resolve the terminal state and report it to the client."""
+        process = self._process
+        self._process = None
+        if self._status == "cancelled":
+            await emit(self.snapshot())
+            raise DownloadCancelledError()
+        if process is None or process.returncode != 0:
+            self._status = "error"
+            await emit(self.snapshot())
+            raise RuntimeError("dataset download failed")
+        self._status = "done"
+        await emit(self.snapshot())
+
+
+manager = DownloadManager()

server/encoder.py

@@ -0,0 +1,106 @@
+"""Render raw samples and encoded spike volumes from a client config."""
+
+from typing import Any, Dict, Optional
+
+import torch
+
+from server.schemas import EncodeConfig
+from spikeforge.data.sample_source import SampleSource
+from spikeforge.encoding.spike_encoder import SpikeEncoder
+
+
+def to_list(tensor: torch.Tensor) -> Any:
+    """Convert a tensor to a nested python list of floats."""
+    return tensor.detach().cpu().float().tolist()
+
+
+class EncoderEngine:
+    """Build a dataset-backed sample and its encoded spike volume."""
+
+    def __init__(self, config: EncodeConfig) -> None:
+        """Encode the selected sample under the client's config."""
+        self._config = config
+        self._source = SampleSource(config.dataset, size=config.input_size)
+        self._index = self._source.clamp(config.sample_index)
+        self._encoder = SpikeEncoder.from_encode_config(config)
+        self._image = self._source.image(self._index)
+        self._spikes = self._encoder.encode_image(self._image)
+
+    # --- payload builders -------------------------------------------------
+
+    def sample_image(self) -> Any:
+        """Return the raw input image grid as a nested list."""
+        return to_list(self._image[0])
+
+    def sample_tensor(self) -> torch.Tensor:
+        """Return the raw input image as a ``[C, H, W]`` tensor."""
+        return self._image
+
+    def reconstruction(self) -> Optional[Dict[str, Any]]:
+        """Return averaged spike reconstructions for the rate coding."""
+        if self._config.coding != "rate":
+            return None
+        height, width = self._source.size
+        gain1 = self._spikes.mean(dim=0)[0].reshape(height, width)
+        low = gain1 * self._config.gain
+        return {
+            "gain1": to_list(gain1),
+            "low": to_list(low),
+            "size": [height, width],
+        }
+
+    def spike_frame(self, step: int) -> Any:
+        """Return one 2-D spike frame of the selected sample."""
+        height, width = self._source.size
+        index = min(max(int(step), 0), self._spikes.size(0) - 1)
+        return to_list(self._spikes[index, 0].reshape(height, width))
+
+    def spike_tensor(self) -> torch.Tensor:
+        """Return the full [T,1,784] spike volume."""
+        return self._spikes
+
+    def raster(self, max_neurons: int = 784) -> Dict[str, Any]:
+        """Return (time, neuron) spike coordinates for the sample."""
+        matrix = self._spikes[:, 0, :]
+        time_idx, neuron_idx = torch.where(matrix > 0)
+        return {
+            "time": to_list(time_idx),
+            "neurons": to_list(neuron_idx),
+            "num_steps": int(matrix.size(0)),
+            "num_neurons": int(matrix.size(1)),
+        }
+
+    def num_steps(self) -> int:
+        """Return the number of encoded time steps."""
+        return int(self._spikes.size(0))
+
+    def target_label(self) -> int:
+        """Return the label of the selected sample."""
+        return self.sample_label()
+
+    def sample_label(self) -> int:
+        """Return the label of the selected sample."""
+        return self._source.label(self._index)
+
+    def sample_index(self) -> int:
+        """Return the clamped index of the selected sample."""
+        return self._index
+
+    def spike_input(self) -> torch.Tensor:
+        """Return the [T,1,784] spikes for training/inference."""
+        return self._spikes
+
+    @property
+    def encoder(self) -> SpikeEncoder:
+        """Return the encoder that produced the sample's spikes."""
+        return self._encoder
+
+    @property
+    def dataset(self) -> str:
+        """Return the registry key of the loaded sample source."""
+        return self._source.dataset
+
+    @property
+    def modality(self) -> str:
+        """Return the sample modality, always ``image``."""
+        return "image"