Extract per-device state out of RequestHandler (#11)

.gitignore

@@ -116,3 +116,4 @@ data
 client-sciot.json
 .python-version
 uv.lock
+sciot-*.json

src/server/communication/device_state.py

@@ -0,0 +1,217 @@
+"""Per-device runtime state for the inference/offloading path.
+
+``RequestHandler`` used to keep every piece of runtime state in class-level
+dictionaries mutated from both the ASGI thread and the background I/O thread.
+This module owns that state instead, behind a single reentrant lock, and scopes
+the variance detector and the adaptive-risk state to one device each.
+
+The legacy ``device_profiles`` mapping stays the canonical storage for a
+device's EMA tables: ``http_server`` and the dashboard read it as a plain dict,
+so the manager hands out the live dict rather than a copy.
+"""
+
+import json
+import threading
+
+from server.commons import OffloadingDataFiles
+from server.core.variance_detector import VarianceDetector
+from server.offloading_algo.adaptive_risk import AdaptiveRiskState
+
+DEFAULT_MODEL_DIR = "test_model_96x96"
+DEFAULT_MODEL_KEY = "fomo_96x96"
+
+VARIANCE_WINDOW_SIZE = 10
+VARIANCE_THRESHOLD = 0.15
+
+
+def build_profile(
+    device_id: str,
+    *,
+    model_dir: str = DEFAULT_MODEL_DIR,
+    model_key: str = DEFAULT_MODEL_KEY,
+    num_layers: int = 0,
+    last_offloading_layer: int | None = None,
+) -> dict:
+    """Build a device profile in the legacy ``device_profiles`` shape."""
+    profile = {
+        "id": device_id,
+        "model_dir": model_dir,
+        "model_key": model_key,
+        "device_inference_times": {f"layer_{i}": 1 for i in range(num_layers)},
+        "edge_inference_times": {f"layer_{i}": 0.1 for i in range(num_layers)},
+    }
+    if last_offloading_layer is not None:
+        profile["ultimo_layer"] = last_offloading_layer
+    return profile
+
+
+class DeviceStateManager:
+    """Thread-safe registry of per-device runtime state.
+
+    Holds the profile (EMA tables), the variance detector, and the adaptive-risk
+    state for each device. Every mutation goes through ``_lock``: the EMA updates
+    run on the ASGI thread while the background I/O worker snapshots the same
+    dicts.
+    """
+
+    def __init__(self):
+        self._lock = threading.RLock()
+        # device_id -> legacy profile dict; exposed live via `profiles`.
+        self._profiles: dict[str, dict] = {}
+        self._variance_detectors: dict[str, VarianceDetector] = {}
+        self._adaptive_states: dict[str, AdaptiveRiskState] = {}
+        self._layer_sizes_cache: dict[str, list] = {}
+
+    @property
+    def profiles(self) -> dict:
+        """The live ``device_id -> profile`` mapping.
+
+        Returned by reference for backwards compatibility with direct readers.
+        """
+        return self._profiles
+
+    # ── device lifecycle ────────────────────────────────────────────────
+    def register(
+        self,
+        device_id: str,
+        *,
+        model_dir: str = DEFAULT_MODEL_DIR,
+        model_key: str = DEFAULT_MODEL_KEY,
+        num_layers: int = 0,
+        last_offloading_layer: int | None = None,
+    ) -> dict:
+        """Return the device's profile, creating it on first sight."""
+        with self._lock:
+            profile = self._profiles.get(device_id)
+            if profile is None:
+                profile = build_profile(
+                    device_id,
+                    model_dir=model_dir,
+                    model_key=model_key,
+                    num_layers=num_layers,
+                    last_offloading_layer=last_offloading_layer,
+                )
+                self._profiles[device_id] = profile
+            return profile
+
+    def __contains__(self, device_id: str) -> bool:
+        with self._lock:
+            return device_id in self._profiles
+
+    def model_dir(self, device_id: str, default: str = DEFAULT_MODEL_DIR) -> str:
+        with self._lock:
+            return self._profiles.get(device_id, {}).get("model_dir", default)
+
+    def model_key(self, device_id: str, default: str = DEFAULT_MODEL_KEY) -> str:
+        with self._lock:
+            return self._profiles.get(device_id, {}).get("model_key", default)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._profiles.clear()
+            self._variance_detectors.clear()
+            self._adaptive_states.clear()
+            self._layer_sizes_cache.clear()
+
+    # ── per-device collaborators ────────────────────────────────────────
+    def variance_detector(self, device_id: str) -> VarianceDetector:
+        """The device's own detector.
+
+        Scoped per device so heterogeneous hardware reporting the same layer
+        indices cannot pollute each other's variance statistics.
+        """
+        with self._lock:
+            detector = self._variance_detectors.get(device_id)
+            if detector is None:
+                detector = VarianceDetector(
+                    window_size=VARIANCE_WINDOW_SIZE,
+                    variance_threshold=VARIANCE_THRESHOLD,
+                )
+                self._variance_detectors[device_id] = detector
+            return detector
+
+    def adaptive_state(self, device_id: str) -> AdaptiveRiskState:
+        with self._lock:
+            state = self._adaptive_states.get(device_id)
+            if state is None:
+                state = AdaptiveRiskState()
+                self._adaptive_states[device_id] = state
+            return state
+
+    def variance_stats(self, device_id: str) -> dict:
+        """Variance statistics for one device, for the offloading context."""
+        return self.variance_detector(device_id).get_all_stats()
+
+    # ── EMA updates ─────────────────────────────────────────────────────
+    def update_device_times(self, device_id: str, layer_times, *, alpha: float) -> None:
+        """Apply the EMA update for device-measured layer times."""
+        profile = self.register(device_id)
+        detector = self.variance_detector(device_id)
+        with self._lock:
+            times = profile["device_inference_times"]
+            for layer_id, measured in enumerate(layer_times):
+                self._apply_ema(times, layer_id, float(measured), alpha)
+                detector.add_device_measurement(layer_id, measured)
+
+    def update_edge_times(
+        self, device_id: str, layer_times, *, start_layer: int, alpha: float
+    ) -> None:
+        """Apply the EMA update for edge-measured layer times."""
+        profile = self.register(device_id)
+        detector = self.variance_detector(device_id)
+        with self._lock:
+            times = profile["edge_inference_times"]
+            for offset, measured in enumerate(layer_times):
+                layer_id = start_layer + offset
+                self._apply_ema(times, layer_id, float(measured), alpha)
+                detector.add_edge_measurement(layer_id, measured)
+
+    @staticmethod
+    def _apply_ema(times: dict, layer_id: int, measured: float, alpha: float) -> None:
+        layer_key = f"layer_{layer_id}"
+        if layer_key in times:
+            times[layer_key] = alpha * measured + (1 - alpha) * times[layer_key]
+        else:
+            times[layer_key] = measured
+
+    # ── reads for the offloading decision ───────────────────────────────
+    def snapshot_times(self, device_id: str) -> tuple[list, list]:
+        """Copy the EMA tables so callers never iterate a dict being mutated."""
+        with self._lock:
+            profile = self._profiles.get(device_id)
+            if profile is None:
+                raise ValueError(f"Unknown device_id '{device_id}'")
+            return (
+                list(profile["device_inference_times"].values()),
+                list(profile["edge_inference_times"].values()),
+            )
+
+    def snapshot_debug_times(self, device_id: str) -> tuple[dict, dict]:
+        """Copy the EMA tables for the background debug-JSON writer."""
+        with self._lock:
+            profile = self._profiles.get(device_id)
+            if profile is None:
+                return {}, {}
+            return (
+                dict(profile["device_inference_times"]),
+                dict(profile["edge_inference_times"]),
+            )
+
+    def layer_sizes(self, model_dir: str) -> list:
+        """Return per-layer output sizes, reading each model's file once."""
+        with self._lock:
+            cached = self._layer_sizes_cache.get(model_dir)
+        if cached is not None:
+            return cached
+
+        # Read outside the lock: this touches the disk, and a lookup for another
+        # model must not block behind it.
+        with open(OffloadingDataFiles.get_sizes_path(model_dir), "r") as file:
+            sizes = list(json.load(file).values())
+        with self._lock:
+            return self._layer_sizes_cache.setdefault(model_dir, sizes)
+
+    def load_stats(self, device_id: str) -> tuple[list, list, list]:
+        """Return (device EMA times, edge EMA times, layer sizes) for a device."""
+        device_times, edge_times = self.snapshot_times(device_id)
+        return device_times, edge_times, self.layer_sizes(self.model_dir(device_id))