refactor: remove dead code and deduplicate scheduler setup

astrai/dataset/dataset.py

@@ -492,9 +492,6 @@ class DPODataset(BaseDataset):
 
     required_keys = ["chosen", "rejected", "chosen_mask", "rejected_mask"]
 
-    def make_processor(self, tokenizer, max_len: int):
-        return partial(dpo_processor, tokenizer=tokenizer, max_len=max_len)
-
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
         return {
             "chosen": self.store.fetch_record(index, "chosen").to(dtype=torch.long),

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()
+            active_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,
-                )
         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._backend_name = type(get_backend()).__name__
-                self._executor = Executor(
+            self._executor = Executor(
-                    model=model,
+                model=model,
-                    kv_cache=self._cache,
+                kv_cache=self._cache,
-                    task_cache=self._task_cache,
+                task_cache=self._task_cache,
-                    device=self.device,
+                device=self.device,
-                    dtype=self.dtype,
+                dtype=self.dtype,
-                    enable_cuda_graph=enable_cuda_graph,
+                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._abort_and_clear(free_waiting=False)
-                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()
 
     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]``.
 

astrai/model/automodel.py

@@ -4,7 +4,7 @@ AutoModel base class for model loading and saving.
 
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Self, Union
+from typing import Union
 
 import torch.nn as nn
 
@@ -90,7 +90,3 @@ class AutoModel(nn.Module):
             state_dict=self.state_dict(),
             save_directory=str(save_directory),
         )
-
-    def to(self, *args, **kwargs) -> Self:
-        """Move model to device/dtype."""
-        return super().to(*args, **kwargs)

astrai/trainer/schedule.py

@@ -2,7 +2,7 @@
 
 import math
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List
+from typing import List
 
 from torch.optim.lr_scheduler import LRScheduler
 
@@ -20,12 +20,6 @@ class BaseScheduler(LRScheduler, ABC):
         """Calculate the current learning rate."""
         raise NotImplementedError
 
-    def state_dict(self) -> Dict[str, Any]:
-        return super().state_dict()
-
-    def load_state_dict(self, state_dict: Dict[str, Any]):
-        super().load_state_dict(state_dict)
-
 
 class SchedulerFactory(BaseFactory["BaseScheduler"]):
     """Factory class for creating learning rate schedulers.

tests/inference/test_cache.py

@@ -201,23 +201,6 @@ def test_req_to_token_pool_write():
 # ---- KVStorage ----
 
 
-def test_kv_storage_set_and_get():
-    storage = KVStorage(
-        size=16,
-        n_layers=2,
-        n_kv_heads=4,
-        head_dim=8,
-        device=torch.device("cpu"),
-        dtype=torch.float32,
-    )
-    loc = torch.tensor([[0, 1]], dtype=torch.long)
-    k = torch.randn(1, 2, 4, 8)
-    v = torch.randn(1, 2, 4, 8)
-    storage.set_kv_buffer(0, loc, k, v)
-    assert torch.allclose(storage.get_key_buffer(0)[loc], k)
-    assert torch.allclose(storage.get_value_buffer(0)[loc], v)
-
-
 def test_kv_storage_buffer_shape():
     storage = KVStorage(
         size=32,