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

@@ -138,7 +138,6 @@
     update_nested_dict,
 )
 from .kernel import DockerKernel
-from .metadata.server import MetadataServer
 from .resources import load_resources, scan_available_resources
 from .utils import PersistentServiceContainer
 
@@ -1341,7 +1340,6 @@ class DockerAgent(AbstractAgent[DockerKernel, DockerKernelCreationContext]):
     monitor_docker_task: asyncio.Task
     agent_sockpath: Path
     agent_sock_task: asyncio.Task
-    metadata_server: MetadataServer
     docker_ptask_group: aiotools.PersistentTaskGroup
     gwbridge_subnet: Optional[str]
     checked_invalid_images: Set[str]
@@ -1414,10 +1412,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",
                 {
@@ -1443,12 +1441,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()
 
@@ -1477,7 +1469,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/kubernetes/agent.py

@@ -844,7 +844,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()

src/ai/backend/agent/runtime.py

@@ -1,24 +1,48 @@
+import asyncio
 import importlib
 import signal
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
 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
 from ai.backend.common.etcd import AsyncEtcd, ConfigScopes
-from ai.backend.common.types import aobject
+from ai.backend.common.exception import (
+    BackendAIError,
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+from ai.backend.common.types import AgentId, aobject
+
+if TYPE_CHECKING:
+    from .docker.metadata.server import MetadataServer
+
+
+class AgentIdNotFoundError(BackendAIError):
+    @classmethod
+    def error_code(cls) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.AGENT,
+            operation=ErrorOperation.ACCESS,
+            error_detail=ErrorDetail.NOT_FOUND,
+        )
 
 
 class AgentRuntime(aobject):
     local_config: AgentUnifiedConfig
-    agent: AbstractAgent
+    agents: dict[AgentId, AbstractAgent]
     kernel_registry: KernelRegistry
     etcd: AsyncEtcd
-    etcd_view: AgentEtcdClientView
+    etcd_views: dict[AgentId, AgentEtcdClientView]
+    metadata_server: MetadataServer | None
 
+    _default_agent_id: AgentId
     _stop_signal: signal.Signals
 
     def __init__(
@@ -30,33 +54,68 @@ def __init__(
         agent_public_key: Optional[PublicKey],
     ) -> None:
         self.local_config = local_config
+        self.agents = {}
         self.kernel_registry = KernelRegistry()
         self.etcd = etcd
-        self.etcd_view = AgentEtcdClientView(etcd, self.local_config)
+        self.etcd_views = {}
+        self.metadata_server = None
 
+        agent_configs = self.local_config.get_agent_configs()
+        self._default_agent_id = AgentId(agent_configs[0].agent.id)
         self._stop_signal = signal.SIGTERM
 
         self.stats_monitor = stats_monitor
         self.error_monitor = error_monitor
         self.agent_public_key = agent_public_key
 
     async def __ainit__(self) -> None:
-        self.agent = await self._create_agent(self.etcd_view, self.local_config)
+        if self.local_config.agent_common.backend == AgentBackend.DOCKER:
+            await self._initialize_metadata_server()
 
-    async def __aexit__(self, *exc_info) -> None:
-        await self.agent.shutdown(self._stop_signal)
+        tasks = []
+        async with asyncio.TaskGroup() as tg:
+            for agent_config in self.local_config.get_agent_configs():
+                agent_id = AgentId(agent_config.agent.id)
+                etcd_view = AgentEtcdClientView(self.etcd, agent_config)
 
-    def get_agent(self) -> AbstractAgent:
-        return self.agent
+                self.etcd_views[agent_id] = etcd_view
+                tasks.append(tg.create_task(self._create_agent(etcd_view, agent_config)))
+        self.agents = {(agent := task.result()).id: agent for task in tasks}
 
-    def get_etcd(self) -> AgentEtcdClientView:
-        return self.etcd_view
+    async def __aexit__(self, *exc_info) -> None:
+        for agent in self.agents.values():
+            await agent.shutdown(self._stop_signal)
+        if self.metadata_server is not None:
+            await self.metadata_server.cleanup()
+
+    def get_agents(self) -> list[AbstractAgent]:
+        return list(self.agents.values())
+
+    def get_agent(self, agent_id: Optional[AgentId]) -> AbstractAgent:
+        if agent_id is None:
+            agent_id = self._default_agent_id
+        if agent_id not in self.agents:
+            raise AgentIdNotFoundError(
+                f"Agent '{agent_id}' not found in this runtime. "
+                f"Available agents: {', '.join(self.agents.keys())}"
+            )
+        return self.agents[agent_id]
+
+    def get_etcd(self, agent_id: Optional[AgentId]) -> AgentEtcdClientView:
+        if agent_id is None:
+            agent_id = self._default_agent_id
+        if agent_id not in self.agents:
+            raise AgentIdNotFoundError(
+                f"Agent '{agent_id}' not found in this runtime. "
+                f"Available agents: {', '.join(self.agents.keys())}"
+            )
+        return self.etcd_views[agent_id]
 
     def mark_stop_signal(self, stop_signal: signal.Signals) -> None:
         self._stop_signal = stop_signal
 
-    async def update_status(self, status) -> None:
-        etcd = self.get_etcd()
+    async def update_status(self, status, agent_id: AgentId) -> None:
+        etcd = self.get_etcd(agent_id)
         await etcd.put("", status, scope=ConfigScopes.NODE)
 
     async def _create_agent(
@@ -75,3 +134,13 @@ async def _create_agent(
         agent_cls = agent_mod.get_agent_cls()
 
         return agent_cls.new(etcd_view, agent_config, **agent_kwargs)
+
+    async def _initialize_metadata_server(self) -> None:
+        from .docker.metadata.server import MetadataServer
+
+        self.metadata_server = await MetadataServer.new(
+            self.local_config,
+            self.etcd,
+            kernel_registry=self.kernel_registry.global_view(),
+        )
+        await self.metadata_server.start_server()