feat(BA-2753): Spawn multiple agents and route RPC appropriately

This change adds support for actually spawning multiple agents within
the same agent server and adding agent_id field for all appropriate RPC
calls in the agent server, then ensuring that the manager sends that
info such that the agent server can correctly route the RPC calls to the
correct agent.

Author: hhoikoo

changes/6320.feature.md

@@ -0,0 +1 @@
+Update Agent server RPC functions to include agent ID for agent runtime with multiple agents

src/ai/backend/agent/agent.py

@@ -2276,7 +2276,7 @@ async def scan_running_kernels(self) -> None:
         """
         ipc_base_path = self.local_config.agent.ipc_base_path
         var_base_path = self.local_config.agent.var_base_path
-        last_registry_file = f"last_registry.{self.local_instance_id}.dat"
+        last_registry_file = f"last_registry.{self.id}.dat"
         if os.path.isfile(ipc_base_path / last_registry_file):
             shutil.move(ipc_base_path / last_registry_file, var_base_path / last_registry_file)
         try:
@@ -3745,7 +3745,7 @@ async def save_last_registry(self, force=False) -> None:
         if (not force) and (now <= self.last_registry_written_time + 60):
             return  # don't save too frequently
         var_base_path = self.local_config.agent.var_base_path
-        last_registry_file = f"last_registry.{self.local_instance_id}.dat"
+        last_registry_file = f"last_registry.{self.id}.dat"
         try:
             with open(var_base_path / last_registry_file, "wb") as f:
                 pickle.dump(self.kernel_registry, f)

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

@@ -1314,6 +1314,7 @@ def __init__(
         error_monitor: ErrorPluginContext,
         skip_initial_scan: bool = False,
         agent_public_key: Optional[PublicKey],
+        metadata_server: MetadataServer,
     ) -> None:
         super().__init__(
             etcd,
@@ -1324,6 +1325,8 @@ def __init__(
             agent_public_key=agent_public_key,
         )
         self.checked_invalid_images = set()
+        self.metadata_server = metadata_server
+        self.metadata_server.register_agent(self)
 
     async def __ainit__(self) -> None:
         async with closing_async(Docker()) as docker:
@@ -1369,10 +1372,10 @@ async def __ainit__(self) -> None:
             self.gwbridge_subnet = None
         ipc_base_path = self.local_config.agent.ipc_base_path
         (ipc_base_path / "container").mkdir(parents=True, exist_ok=True)
-        self.agent_sockpath = ipc_base_path / "container" / f"agent.{self.local_instance_id}.sock"
+        self.agent_sockpath = ipc_base_path / "container" / f"agent.{self.id}.sock"
         # Workaround for Docker Desktop for Mac's UNIX socket mount failure with virtiofs
         if sys.platform != "darwin":
-            socket_relay_name = f"backendai-socket-relay.{self.local_instance_id}"
+            socket_relay_name = f"backendai-socket-relay.{self.id}"
             socket_relay_container = PersistentServiceContainer(
                 "backendai-socket-relay:latest",
                 {
@@ -1398,12 +1401,6 @@ async def __ainit__(self) -> None:
         self.monitor_docker_task = asyncio.create_task(self.monitor_docker_events())
         self.docker_ptask_group = aiotools.PersistentTaskGroup()
 
-        self.metadata_server = await MetadataServer.new(
-            self.local_config,
-            self.etcd,
-            self.kernel_registry,
-        )
-        await self.metadata_server.start_server()
         # For legacy accelerator plugins
         self.docker = Docker()
 
@@ -1432,7 +1429,6 @@ async def shutdown(self, stop_signal: signal.Signals):
                 self.monitor_docker_task.cancel()
                 await self.monitor_docker_task
 
-        await self.metadata_server.cleanup()
         if self.docker:
             await self.docker.close()
 

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

@@ -1,14 +1,15 @@
 import logging
 from collections.abc import Sequence
 from http import HTTPStatus
-from typing import Any, List, Mapping, MutableMapping, cast
+from typing import Any, Iterator, List, Mapping, MutableMapping, cast, override
 from uuid import UUID
 
 import attr
 from aiodocker.docker import Docker
 from aiohttp import web
 from aiohttp.typedefs import Handler, Middleware
 
+from ai.backend.agent.agent import AbstractAgent
 from ai.backend.agent.config.unified import AgentUnifiedConfig
 from ai.backend.agent.docker.kernel import prepare_kernel_metadata_uri_handling
 from ai.backend.agent.kernel import AbstractKernel
@@ -17,7 +18,7 @@
 from ai.backend.common.etcd import AsyncEtcd
 from ai.backend.common.json import dump_json_str
 from ai.backend.common.plugin import BasePluginContext
-from ai.backend.common.types import KernelId, aobject
+from ai.backend.common.types import AgentId, KernelId, aobject
 from ai.backend.logging import BraceStyleAdapter
 
 from .plugin import MetadataPlugin
@@ -87,17 +88,50 @@ async def list_versions(request: web.Request) -> web.Response:
     return web.Response(body="latest/")
 
 
+class AggregateKernelRegistry(Mapping[KernelId, AbstractKernel]):
+    _agents: dict[AgentId, AbstractAgent]
+
+    def __init__(self) -> None:
+        self._agents = {}
+
+    def register_agent(self, agent: AbstractAgent) -> None:
+        self._agents[agent.id] = agent
+
+    @override
+    def __getitem__(self, kernel_id: KernelId) -> AbstractKernel:
+        for agent in self._agents.values():
+            if kernel_id in agent.kernel_registry:
+                return agent.kernel_registry[kernel_id]
+        raise KeyError(kernel_id)
+
+    @override
+    def __iter__(self) -> Iterator[KernelId]:
+        for agent in self._agents.values():
+            yield from agent.kernel_registry.keys()
+
+    @override
+    def __len__(self) -> int:
+        return sum(len(agent.kernel_registry) for agent in self._agents.values())
+
+    @override
+    def __contains__(self, x: object, /) -> bool:
+        if not isinstance(x, str):
+            return False
+        return any(agent.__contains__(x) for agent in self._agents)
+
+
 class MetadataServer(aobject):
     app: web.Application
     runner: web.AppRunner
     route_structure: MutableMapping[str, Any]
     loaded_apps: List[str]
+    kernel_registry: AggregateKernelRegistry
 
     def __init__(
         self,
         local_config: AgentUnifiedConfig,
         etcd: AsyncEtcd,
-        kernel_registry: Mapping[KernelId, AbstractKernel],
+        kernel_registry: AggregateKernelRegistry,
     ) -> None:
         app = web.Application(
             middlewares=[
@@ -112,6 +146,7 @@ def __init__(
         self.app = app
         self.loaded_apps = []
         self.route_structure = {"latest": {"extension": {}}}
+        self.kernel_registry = kernel_registry
 
     async def __ainit__(self):
         local_config = cast(AgentUnifiedConfig, self.app["_root.context"].local_config)
@@ -132,6 +167,9 @@ async def __ainit__(self):
         self.app.router.add_route("GET", "/", list_versions)
         self.app.router.add_route("GET", "/{version}", self.list_available_apps)
 
+    def register_agent(self, agent: AbstractAgent) -> None:
+        self.kernel_registry.register_agent(agent)
+
     async def list_available_apps(self, request: web.Request) -> web.Response:
         return web.Response(body="\n".join([x + "/" for x in self.loaded_apps]))
 

src/ai/backend/agent/etcd.py

@@ -0,0 +1,58 @@
+from typing import Collection
+
+from ai.backend.agent.config.unified import AgentUnifiedConfig
+from ai.backend.common.data.config.types import EtcdConfigData
+from ai.backend.common.etcd import AsyncEtcd, ConfigScopes
+from ai.backend.common.types import AgentId
+
+
+class EtcdClientRegistry:
+    _etcd_config: EtcdConfigData
+    _etcd_clients: dict[AgentId, AsyncEtcd]
+    _global_etcd: AsyncEtcd
+
+    @property
+    def global_etcd(self) -> AsyncEtcd:
+        return self._global_etcd
+
+    def __init__(self, etcd_config: EtcdConfigData) -> None:
+        self._etcd_config = etcd_config
+        self._etcd_clients = {}
+        self._global_etcd = self._create_client(agent_id=None, scaling_group=None)
+
+    async def close(self) -> None:
+        for etcd in self._etcd_clients.values():
+            await etcd.close()
+        await self._global_etcd.close()
+
+    def get_client(self, agent_id: AgentId) -> AsyncEtcd:
+        return self._etcd_clients[agent_id]
+
+    def prefill_clients(self, prefill_data: Collection[AgentUnifiedConfig]) -> None:
+        for agent_config in prefill_data:
+            agent_id = AgentId(agent_config.agent.id)
+            self._etcd_clients[agent_id] = self._create_client(
+                agent_id, agent_config.agent.scaling_group
+            )
+
+    def _create_client(self, agent_id: AgentId | None, scaling_group: str | None) -> AsyncEtcd:
+        scope_prefix_map = {ConfigScopes.GLOBAL: ""}
+        if scaling_group is not None:
+            scope_prefix_map[ConfigScopes.SGROUP] = f"sgroup/{scaling_group}"
+        if agent_id is not None:
+            scope_prefix_map[ConfigScopes.NODE] = f"nodes/agents/{agent_id}"
+
+        if self._etcd_config.user and self._etcd_config.password:
+            etcd_credentials = {
+                "user": self._etcd_config.user,
+                "password": self._etcd_config.password,
+            }
+        else:
+            etcd_credentials = None
+
+        return AsyncEtcd(
+            [addr.to_legacy() for addr in self._etcd_config.addrs],
+            self._etcd_config.namespace,
+            scope_prefix_map,
+            credentials=etcd_credentials,
+        )

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

@@ -842,7 +842,7 @@ def __init__(
     async def __ainit__(self) -> None:
         await super().__ainit__()
         ipc_base_path = self.local_config.agent.ipc_base_path
-        self.agent_sockpath = ipc_base_path / "container" / f"agent.{self.local_instance_id}.sock"
+        self.agent_sockpath = ipc_base_path / "container" / f"agent.{self.id}.sock"
 
         await self.check_krunner_pv_status()
         await self.fetch_workers()