refactor: remove dead code and deduplicate scheduler setup

astrai/inference/cache/buffer.py

@@ -72,16 +72,6 @@ class KVStorage:
             (n_layers, size, n_kv_heads, head_dim), device=device, dtype=dtype
         )
 
-    def get_key_buffer(self, layer_id: int) -> Tensor:
-        return self.k_buffer[layer_id]
-
-    def get_value_buffer(self, layer_id: int) -> Tensor:
-        return self.v_buffer[layer_id]
-
-    def set_kv_buffer(self, layer_id: int, loc: Tensor, k: Tensor, v: Tensor) -> None:
-        self.k_buffer[layer_id, loc] = k
-        self.v_buffer[layer_id, loc] = v
-
 
 @dataclass
 class KVCache:

astrai/inference/metrics.py

@@ -148,17 +148,6 @@ class MetricsCollector:
         self._completed.append(timing)
         self._accumulate(timing)
 
-    def clear(self):
-        """Reset all state (e.g. on engine shutdown)."""
-        self._timings.clear()
-        self._completed.clear()
-        self._ttft_ms_sum = 0.0
-        self._ttft_ms_count = 0
-        self._decode_tps_sum = 0.0
-        self._decode_tps_count = 0
-        self._e2e_ms_sum = 0.0
-        self._e2e_ms_count = 0
-
     # timing scopes
 
     @contextmanager
@@ -180,17 +169,6 @@ class MetricsCollector:
                 t._decode_steps += 1
                 t._decode_total_s += dt
 
-    # access
-
-    def get_timing(self, task_id: str) -> Optional[TaskTiming]:
-        """Return the timing record for *task_id* (active or completed)."""
-        if task_id in self._timings:
-            return self._timings[task_id]
-        for t in self._completed:
-            if t.task_id == task_id:
-                return t
-        return None
-
     # aggregate stats
 
     def get_stats(self) -> Dict[str, Any]:

astrai/inference/scheduler.py

@@ -79,29 +79,21 @@ class InferenceScheduler:
 
         if backend is None:
             self._backend = None
-            default_backend = get_backend()
-            self._backend_name = type(default_backend).__name__
-            with attn_backend(default_backend):
-                self._executor = Executor(
-                    model=model,
-                    kv_cache=self._cache,
-                    task_cache=self._task_cache,
-                    device=self.device,
-                    dtype=self.dtype,
-                    enable_cuda_graph=enable_cuda_graph,
-                )
+            active_backend = get_backend()
         else:
-            with attn_backend(backend):
+            active_backend = backend
+        with attn_backend(active_backend):
+            if backend is not None:
                 self._backend = get_backend()
-                self._backend_name = type(self._backend).__name__
-                self._executor = Executor(
-                    model=model,
-                    kv_cache=self._cache,
-                    task_cache=self._task_cache,
-                    device=self.device,
-                    dtype=self.dtype,
-                    enable_cuda_graph=enable_cuda_graph,
-                )
+            self._backend_name = type(get_backend()).__name__
+            self._executor = Executor(
+                model=model,
+                kv_cache=self._cache,
+                task_cache=self._task_cache,
+                device=self.device,
+                dtype=self.dtype,
+                enable_cuda_graph=enable_cuda_graph,
+            )
 
         self._stop_event = threading.Event()
         self._loop_thread: Optional[threading.Thread] = None
@@ -283,12 +275,7 @@ class InferenceScheduler:
         except Exception as e:
             self._stop_event.set()
             logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
-            for task in self._task_mgr.get_active_tasks():
-                self._task_mgr.invoke_callback(task.task_id, STOP)
-                self._task_cache.task_free(task.task_id)
-            for task in self._task_mgr.get_waiting_tasks():
-                self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._task_mgr.clear_queues()
+            self._abort_and_clear(free_waiting=False)
 
     def start(self):
         if self._loop_thread is not None and self._loop_thread.is_alive():
@@ -304,15 +291,20 @@ class InferenceScheduler:
         if self._loop_thread is not None:
             self._loop_thread.join(timeout=2.0)
         self._loop_thread = None
+        self._abort_and_clear(free_waiting=True)
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    def _abort_and_clear(self, free_waiting: bool):
+        """Invoke STOP callbacks, release cache slots, and clear task queues."""
         for task in self._task_mgr.get_active_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
             self._task_cache.task_free(task.task_id)
         for task in self._task_mgr.get_waiting_tasks():
             self._task_mgr.invoke_callback(task.task_id, STOP)
-            self._task_cache.task_free(task.task_id)
+            if free_waiting:
+                self._task_cache.task_free(task.task_id)
         self._task_mgr.clear_queues()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
 
     def run_batch(
         self,

astrai/inference/workspace.py

@@ -124,13 +124,6 @@ class InferenceWorkspace:
             device=device,
         )
 
-    def decode_buffers(self, batch: int, q_heads: int):
-        """Return ``(o_part, ml_part)`` view sliced to live dimensions."""
-        return (
-            self.decode_o_part[:batch, :q_heads],
-            self.decode_ml_part[:batch, :q_heads],
-        )
-
     def fill_input_ids(self, ids: "list[int]") -> Tensor:
         """Write ``ids`` into the device buffer and return ``[B]``.