Resource pools: shared engine resources with a staged migration

mstar/engine/base.py

@@ -9,7 +9,7 @@
 from mstar.communication.tensors import NameToTensorList
 from mstar.conductor.request_info import CurrentForwardPassInfo
 from mstar.distributed.communication import WorkerParallelGroups
-from mstar.engine.kv_store import KVCacheConfig, StoreWritePolicy
+from mstar.engine.kv_store import KVCacheConfig
 from mstar.profile.worker import ExecTimings
 
 
@@ -23,14 +23,12 @@ class EngineCapabilities:
     """Static declaration of optional surfaces an engine implements.
 
     The worker and ``EngineManager`` consult these flags instead of
-    ``isinstance`` / ``hasattr`` probes to decide whether to iterate or
-    dispatch into engine-specific code paths (e.g. CPU-offload victim
-    selection, KV-cache LRU tracking, store write-policy push). The
-    default ``EngineCapabilities()`` declares an engine that needs none
-    of the optional surfaces — stateless engines leave it untouched.
+    ``isinstance`` / ``hasattr`` probes. The default declares an engine
+    that needs none of the optional surfaces; stateless engines leave it
+    untouched. Per-node resource questions (offload tiers, LRU tracking,
+    write policy) go through ``node_resources`` instead of flags here.
     """
     requires_kv_cache: bool = False
-    supports_cpu_offload: bool = False
 
 
 @dataclass
@@ -308,58 +306,23 @@ def reset_pre_plan_for_batch(self, batch: NodeBatch) -> None:
         """
         return
 
-    # ── Capabilities + optional surfaces ────────────────────────────────
+    # ── Capabilities + per-node resources ───────────────────────────────
     #
     # ``capabilities`` is a class-level declaration of which optional
-    # surfaces this engine class implements. Worker / EngineManager check
-    # it instead of ``isinstance`` / ``hasattr`` probes. The methods below
-    # are the corresponding surfaces — all safe no-op defaults so engines
-    # that don't opt in can still be called uniformly. KVCacheEngine
-    # overrides both the capability flags and the methods; stateless
-    # engines leave them at default.
+    # surfaces this engine class implements; Worker / EngineManager check
+    # it instead of ``isinstance`` / ``hasattr`` probes. ``node_resources``
+    # is the per-node counterpart: the worker reaches an engine's KV cache
+    # pool (offload, LRU eligibility, write policy) through it rather than
+    # through per-question engine methods.
 
     capabilities = EngineCapabilities()
 
-    def lru_tracked_nodes(self) -> list[str]:
-        """Nodes for which the worker should LRU-track per-request activity
-        (used to pick CPU-offload victims). Default: no nodes — stateless
-        engines have no KV state to age out.
+    def node_resources(self, node_name: str) -> dict[str, Any]:
+        """The named resources this engine built for ``node_name`` (KV
+        cache pool, embedder, cross pools, scratch caches). Engines that
+        build none return an empty dict.
         """
-        return []
-
-    def set_alloc_write_policy(self, policy: StoreWritePolicy) -> None:
-        """Apply a store write policy. Default: no-op — engines without an
-        alloc manager have nothing to set.
-        """
-        return
-
-    def offload_candidates(self, node_name: str) -> list[tuple[str, int]]:
-        """Return ``(request_id, gpu_pages_held)`` for every request with
-        GPU pages on ``node_name``. The worker partitions the result into
-        in-batch vs external candidates and picks an eviction victim.
-        Default: empty list — no offloadable state.
-        """
-        return []
-
-    def offload_request(self, node_name: str, request_id: str) -> int:
-        """Offload ``request_id``'s KV pages on ``node_name`` to CPU.
-        Returns the number of GPU pages freed (0 if nothing was freed or
-        the engine doesn't support offload).
-        """
-        return 0
-
-    def reload_request(self, node_name: str, request_id: str) -> bool:
-        """Reload an offloaded request back to GPU on ``node_name``.
-        Returns True on success; False if the request isn't offloaded, GPU
-        pages are insufficient, or the engine doesn't support offload.
-        """
-        return False
-
-    def is_offloaded(self, node_name: str, request_id: str) -> bool:
-        """Whether ``request_id`` is currently CPU-offloaded on ``node_name``.
-        Default: False.
-        """
-        return False
+        return {}
 
     def execute_with_max_batch_size(self, batch: NodeBatch) -> NodeOutput:
         if self.enable_profile: