Extract OffloadingService and test the inference hot path (#45)

src/server/communication/offloading_service.py

@@ -0,0 +1,184 @@
+"""The offloading decision: picking the split layer and recording why.
+
+`RequestHandler` used to build the algorithm, run it, unpick its candidate list
+and publish the decision telemetry inline. This module owns that, so the handler
+is left orchestrating rather than deciding.
+"""
+
+from dataclasses import dataclass, field
+
+from server.commons import EvaluationFiles
+from server.offloading_algo.factory import (
+    OffloadingContext,
+    configured_algorithm_class_name,
+    configured_algorithm_name,
+    create_offloading_algorithm,
+)
+from server.telemetry.offloading_decisions import (
+    append_offloading_decision,
+    append_offloading_decision_jsonl,
+)
+
+DEVICE_ONLY_LAYER = -1
+
+REASON_FORCED_LOCAL = "forced_local_inference"
+REASON_MISSING_METADATA = "fallback_missing_model_metadata"
+REASON_DEFAULT = "lowest_estimated_cost"
+
+
+@dataclass
+class OffloadingDecision:
+    """The chosen split point, plus everything telemetry needs to explain it."""
+
+    layer: int
+    reason: str
+    strategy: str
+    algorithm_class: str
+    candidates: list = field(default_factory=list)
+    estimated_total_cost_ms: float = 0.0
+
+    @property
+    def selected_candidate(self) -> dict | None:
+        """The candidate the algorithm settled on, if it is still in the list."""
+        return next(
+            (
+                candidate
+                for candidate in self.candidates
+                if candidate.get("offloading_layer_index") == self.layer
+                and candidate.get("considered_for_selection", True)
+            ),
+            None,
+        )
+
+    @property
+    def switch_penalty(self) -> float:
+        return float((self.selected_candidate or {}).get("switch_penalty", 0.0) or 0.0)
+
+    @property
+    def was_computed(self) -> bool:
+        """True when the algorithm actually ran and produced a layer.
+
+        False for a forced-local decision or a metadata fallback, neither of
+        which reflects a real cost comparison.
+        """
+        return self.reason not in (REASON_FORCED_LOCAL, REASON_MISSING_METADATA)
+
+
+class OffloadingService:
+    """Runs the configured offloading algorithm and publishes its decisions."""
+
+    def __init__(self, config: dict, device_state, firestore_publisher=None):
+        self.config = config or {}
+        self.device_state = device_state
+        self.firestore_publisher = firestore_publisher
+
+    def _create_algorithm(
+        self,
+        *,
+        device_id: str,
+        avg_speed: float,
+        device_inference_times: list,
+        edge_inference_times: list,
+        layers_sizes: list,
+        device_cpu_percent: float,
+    ):
+        state = self.device_state
+        return create_offloading_algorithm(
+            self.config,
+            OffloadingContext(
+                avg_speed=avg_speed,
+                num_layers=len(layers_sizes),
+                layers_sizes=list(layers_sizes),
+                inference_time_device=list(device_inference_times),
+                inference_time_edge=list(edge_inference_times),
+                model_dir=state.model_dir(device_id),
+                device_cpu_percent=device_cpu_percent,
+                adaptive_state=state.adaptive_state(device_id),
+                variance_stats=state.variance_stats(device_id),
+            ),
+        )
+
+    def forced_local_decision(self) -> OffloadingDecision:
+        """The decision used when local-only inference is being forced."""
+        return OffloadingDecision(
+            layer=DEVICE_ONLY_LAYER,
+            reason=REASON_FORCED_LOCAL,
+            strategy=configured_algorithm_name(self.config),
+            algorithm_class=configured_algorithm_class_name(self.config),
+        )
+
+    def decide(
+        self,
+        *,
+        device_id: str,
+        avg_speed: float,
+        device_inference_times: list,
+        edge_inference_times: list,
+        layers_sizes: list,
+        device_cpu_percent: float,
+    ) -> OffloadingDecision:
+        """Pick the next offloading layer for a device."""
+        algorithm = self._create_algorithm(
+            device_id=device_id,
+            avg_speed=avg_speed,
+            device_inference_times=device_inference_times,
+            edge_inference_times=edge_inference_times,
+            layers_sizes=layers_sizes,
+            device_cpu_percent=device_cpu_percent,
+        )
+        strategy = algorithm.strategy
+        algorithm_class = algorithm.__class__.__name__
+
+        try:
+            layer = algorithm.select_offloading_layer()
+        except IndexError:
+            # Model metadata (layer sizes) is missing: fall back to device-only
+            # rather than guessing a split point.
+            return OffloadingDecision(
+                layer=DEVICE_ONLY_LAYER,
+                reason=REASON_MISSING_METADATA,
+                strategy=strategy,
+                algorithm_class=algorithm_class,
+                candidates=algorithm.candidate_evaluations,
+            )
+
+        return OffloadingDecision(
+            layer=layer,
+            reason=getattr(algorithm, "selection_reason", REASON_DEFAULT),
+            strategy=strategy,
+            algorithm_class=algorithm_class,
+            candidates=algorithm.candidate_evaluations,
+            estimated_total_cost_ms=(
+                algorithm.lowest_evaluation * 1000
+                if hasattr(algorithm, "lowest_evaluation")
+                else 0.0
+            ),
+        )
+
+    def append_and_publish_decision(self, event: dict, split_config: dict) -> None:
+        """Persist decision telemetry locally and optionally to Firestore."""
+        append_offloading_decision(
+            EvaluationFiles.offloading_decisions_file_path,
+            event,
+            max_rows=split_config["max_rows"],
+            max_interval_seconds=split_config["max_interval_seconds"],
+        )
+        jsonl_path = append_offloading_decision_jsonl(
+            EvaluationFiles.offloading_decisions_jsonl_base_path,
+            event,
+            max_rows=split_config["max_rows"],
+            max_interval_seconds=split_config["max_interval_seconds"],
+        )
+        if self.firestore_publisher is None:
+            return
+        self.firestore_publisher.publish_file(
+            "offloading_decisions",
+            jsonl_path,
+            metadata={
+                "offloading_decisions_jsonl_file": jsonl_path.name,
+                "offloading_decisions_jsonl_path": str(jsonl_path),
+                "last_event_timestamp": event.get("timestamp"),
+            },
+            document_id=jsonl_path.name,
+            delete_after_publish=True,
+        )