refactor(BA-3028): Move kernel registry ownership to AgentRuntime

This change moves where the agent's kernel registry is stored from
within the agent to outside the agent at a global level in AgentRuntime.
This is to make future change of pulling out Metadata server outside of
DockerAgent easier, which needs to take in a global view of all kernels
across all agents.

Author: hhoikoo

changes/6730.enhance.md

@@ -0,0 +1 @@
+Move kernel registry ownership to agent runtime

src/ai/backend/agent/agent.py

@@ -229,6 +229,7 @@
 from .kernel import (
     RUN_ID_FOR_BATCH_JOB,
     AbstractKernel,
+    KernelRegistry,
     match_distro_data,
 )
 from .observer.heartbeat import HeartbeatObserver
@@ -836,6 +837,7 @@ def __init__(
         error_monitor: ErrorPluginContext,
         skip_initial_scan: bool = False,
         agent_public_key: Optional[PublicKey],
+        kernel_registry: KernelRegistry,
     ) -> None:
         self._skip_initial_scan = skip_initial_scan
         self.loop = current_loop()
@@ -844,7 +846,7 @@ def __init__(
         self.id = AgentId(local_config.agent.id or f"agent-{uuid4()}")
         self.local_instance_id = generate_local_instance_id(__file__)
         self.agent_public_key = agent_public_key
-        self.kernel_registry = {}
+        self.kernel_registry = kernel_registry.agent_view(self.id)
         self.computers = {}
         self.images = {}
         self.restarting_kernels = {}

src/ai/backend/agent/docker/agent.py

@@ -109,7 +109,7 @@
 from ..config.unified import AgentUnifiedConfig, ContainerSandboxType, ScratchType
 from ..exception import ContainerCreationError, UnsupportedResource
 from ..fs import create_scratch_filesystem, destroy_scratch_filesystem
-from ..kernel import AbstractKernel
+from ..kernel import AbstractKernel, KernelRegistry
 from ..plugin.network import ContainerNetworkCapability, ContainerNetworkInfo, NetworkPluginContext
 from ..proxy import DomainSocketProxy, proxy_connection
 from ..resources import (
@@ -1357,6 +1357,7 @@ def __init__(
         error_monitor: ErrorPluginContext,
         skip_initial_scan: bool = False,
         agent_public_key: Optional[PublicKey],
+        kernel_registry: KernelRegistry,
     ) -> None:
         super().__init__(
             etcd,
@@ -1365,6 +1366,7 @@ def __init__(
             error_monitor=error_monitor,
             skip_initial_scan=skip_initial_scan,
             agent_public_key=agent_public_key,
+            kernel_registry=kernel_registry,
         )
         self.checked_invalid_images = set()
 

src/ai/backend/agent/kernel.py

@@ -21,15 +21,18 @@
     Any,
     Dict,
     FrozenSet,
+    Iterator,
     List,
     Literal,
+    MutableMapping,
     NotRequired,
     Optional,
     Set,
     Tuple,
     TypedDict,
     Union,
     cast,
+    overload,
 )
 
 import zmq
@@ -452,6 +455,102 @@ async def execute(
             raise
 
 
+@dataclass(frozen=True)
+class AgentKernelRegistryKey:
+    agent_id: AgentId
+    kernel_id: KernelId
+
+
+class KernelRegistry(MutableMapping[AgentKernelRegistryKey, AbstractKernel]):
+    _registry: MutableMapping[AgentKernelRegistryKey, AbstractKernel]
+    _global_registry: MutableMapping[KernelId, AbstractKernel]
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._registry = {}
+        self._global_registry = {}
+
+    def agent_view(self, agent_id: AgentId) -> "KernelRegistryAgentView":
+        return KernelRegistryAgentView(self, agent_id)
+
+    def global_view(self) -> "KernelRegistryGlobalView":
+        return KernelRegistryGlobalView(self)
+
+    @overload
+    def __getitem__(self, key: KernelId) -> AbstractKernel: ...
+
+    @overload
+    def __getitem__(self, key: AgentKernelRegistryKey) -> AbstractKernel: ...
+
+    def __getitem__(self, key: KernelId | AgentKernelRegistryKey) -> AbstractKernel:
+        if isinstance(key, AgentKernelRegistryKey):
+            return self._registry[key]
+        else:
+            return self._global_registry[key]
+
+    def __setitem__(self, key: AgentKernelRegistryKey, value: AbstractKernel) -> None:
+        self._registry[key] = value
+        self._global_registry[key.kernel_id] = value
+
+    def __delitem__(self, key: AgentKernelRegistryKey) -> None:
+        del self._registry[key]
+        del self._global_registry[key.kernel_id]
+
+    def __iter__(self) -> Iterator[AgentKernelRegistryKey]:
+        return iter(self._registry)
+
+    def __len__(self) -> int:
+        return len(self._registry)
+
+
+class KernelRegistryAgentView(MutableMapping[KernelId, AbstractKernel]):
+    _registry: KernelRegistry
+    _agent_id: AgentId
+
+    def __init__(self, kernel_registry: KernelRegistry, agent_id: AgentId) -> None:
+        super().__init__()
+
+        self._registry = kernel_registry
+        self._agent_id = agent_id
+
+    def __getitem__(self, key: KernelId) -> AbstractKernel:
+        return self._registry[AgentKernelRegistryKey(self._agent_id, key)]
+
+    def __setitem__(self, key: KernelId, value: AbstractKernel) -> None:
+        self._registry[AgentKernelRegistryKey(self._agent_id, key)] = value
+
+    def __delitem__(self, key: KernelId) -> None:
+        del self._registry[AgentKernelRegistryKey(self._agent_id, key)]
+
+    def __iter__(self) -> Iterator[KernelId]:
+        for registry_key in self._registry:
+            if registry_key.agent_id == self._agent_id:
+                yield registry_key.kernel_id
+
+    def __len__(self) -> int:
+        return sum(1 for key in self._registry if key.agent_id == self._agent_id)
+
+
+class KernelRegistryGlobalView(Mapping[KernelId, AbstractKernel]):
+    _registry: KernelRegistry
+
+    def __init__(self, kernel_registry: KernelRegistry) -> None:
+        super().__init__()
+
+        self._registry = kernel_registry
+
+    def __getitem__(self, key: KernelId) -> AbstractKernel:
+        return self._registry[key]
+
+    def __iter__(self) -> Iterator[KernelId]:
+        for registry_key in self._registry:
+            yield registry_key.kernel_id
+
+    def __len__(self) -> int:
+        return len(self._registry)
+
+
 _zctx = None
 
 

src/ai/backend/agent/kubernetes/agent.py

@@ -68,7 +68,7 @@
 )
 from ..config.unified import AgentUnifiedConfig, ScratchType
 from ..exception import K8sError, UnsupportedResource
-from ..kernel import AbstractKernel
+from ..kernel import AbstractKernel, KernelRegistry
 from ..resources import (
     AbstractComputePlugin,
     ComputerContext,
@@ -829,6 +829,7 @@ def __init__(
         error_monitor: ErrorPluginContext,
         skip_initial_scan: bool = False,
         agent_public_key: Optional[PublicKey],
+        kernel_registry: KernelRegistry,
     ) -> None:
         super().__init__(
             etcd,
@@ -837,6 +838,7 @@ def __init__(
             error_monitor=error_monitor,
             skip_initial_scan=skip_initial_scan,
             agent_public_key=agent_public_key,
+            kernel_registry=kernel_registry,
         )
 
     async def __ainit__(self) -> None:

src/ai/backend/agent/runtime.py

@@ -4,6 +4,7 @@
 from ai.backend.agent.agent import AbstractAgent
 from ai.backend.agent.config.unified import AgentUnifiedConfig
 from ai.backend.agent.etcd import AgentEtcdClientView
+from ai.backend.agent.kernel import KernelRegistry
 from ai.backend.agent.monitor import AgentErrorPluginContext, AgentStatsPluginContext
 from ai.backend.agent.types import AgentBackend
 from ai.backend.common.auth import PublicKey
@@ -14,6 +15,7 @@
 class AgentRuntime(aobject):
     local_config: AgentUnifiedConfig
     agent: AbstractAgent
+    kernel_registry: KernelRegistry
     etcd: AsyncEtcd
     etcd_view: AgentEtcdClientView
 
@@ -28,6 +30,7 @@ def __init__(
         agent_public_key: Optional[PublicKey],
     ) -> None:
         self.local_config = local_config
+        self.kernel_registry = KernelRegistry()
         self.etcd = etcd
         self.etcd_view = AgentEtcdClientView(etcd, self.local_config)
 

tests/agent/docker/test_agent.py

@@ -11,6 +11,7 @@
 from aiodocker.exceptions import DockerError
 
 from ai.backend.agent.docker.agent import DockerAgent
+from ai.backend.agent.kernel import KernelRegistry
 from ai.backend.common.arch import DEFAULT_IMAGE_ARCH
 from ai.backend.common.docker import ImageRef
 from ai.backend.common.exception import ImageNotAvailable
@@ -28,13 +29,15 @@ async def agent(local_config, test_id, mocker, socket_relay_image):
     mocked_etcd_get_prefix = AsyncMock(return_value={})
     mocker.patch.object(dummy_etcd, "get_prefix", new=mocked_etcd_get_prefix)
     test_case_id = secrets.token_hex(8)
+    kernel_registry = KernelRegistry()
     agent = await DockerAgent.new(
         dummy_etcd,
         local_config,
         stats_monitor=None,
         error_monitor=None,
         skip_initial_scan=True,
         agent_public_key=None,
+        kernel_registry=kernel_registry,
     )  # for faster test iteration
     agent.local_instance_id = test_case_id  # use per-test private registry file
     try: