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

@@ -1357,6 +1357,7 @@ def __init__(
         error_monitor: ErrorPluginContext,
         skip_initial_scan: bool = False,
         agent_public_key: Optional[PublicKey],
+        metadata_server: MetadataServer,
         kernel_registry: KernelRegistry,
     ) -> None:
         super().__init__(
@@ -1369,6 +1370,7 @@ def __init__(
             kernel_registry=kernel_registry,
         )
         self.checked_invalid_images = set()
+        self.metadata_server = metadata_server
 
     async def __ainit__(self) -> None:
         async with closing_async(Docker()) as docker:
@@ -1414,10 +1416,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 +1445,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 +1473,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,45 @@
+import asyncio
 import signal
 from typing import Optional
 
 from ai.backend.agent.agent import AbstractAgent
 from ai.backend.agent.config.unified import AgentUnifiedConfig
+from ai.backend.agent.docker.metadata.server import MetadataServer
 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
-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
+
+
+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,27 +51,58 @@ 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
 
+        self._default_agent_id = AgentId(self.local_config.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)
+        tasks = []
+        async with asyncio.TaskGroup() as tg:
+            for agent_config in self.local_config.agent_configs:
+                agent_id = AgentId(agent_config.agent.id)
+                etcd_view = AgentEtcdClientView(self.etcd, agent_config)
 
-    async def __aexit__(self, *exc_info) -> None:
-        await self.agent.shutdown(self._stop_signal)
-
-    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
@@ -64,12 +116,15 @@ async def _create_agent(
             case AgentBackend.DOCKER:
                 from .docker.agent import DockerAgent
 
+                await self._initialize_metadata_server()
+
                 return await DockerAgent.new(
                     etcd_view,
                     agent_config,
                     stats_monitor=self.stats_monitor,
                     error_monitor=self.error_monitor,
                     agent_public_key=self.agent_public_key,
+                    metadata_server=self.metadata_server,
                 )
             case AgentBackend.KUBERNETES:
                 from .kubernetes.agent import KubernetesAgent
@@ -91,3 +146,13 @@ async def _create_agent(
                     error_monitor=self.error_monitor,
                     agent_public_key=self.agent_public_key,
                 )
+
+    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()