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

@@ -888,6 +888,8 @@ async def __ainit__(self) -> None:
                 "AbstractAgent.__ainit__": "Redis runtime configuration is not set."
             })
 
+        self.local_config.agent.image_commit_path.mkdir(parents=True, exist_ok=True)
+
         redis_profile_target = self.local_config.redis.to_redis_profile_target()
         stream_redis_target = redis_profile_target.profile_target(RedisRole.STREAM)
         mq = await self._make_message_queue(stream_redis_target)
@@ -2276,7 +2278,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 +3747,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/config/unified.py

@@ -1140,7 +1140,12 @@ class AgentOverrideConfig(BaseConfigSchema):
         description="Resource config overrides for the individual agent",
     )
 
-    def construct_unified_config(self, *, default: AgentUnifiedConfig) -> AgentUnifiedConfig:
+    def construct_unified_config(
+        self,
+        *,
+        default: AgentUnifiedConfig,
+        agent_idx: int,
+    ) -> AgentUnifiedConfig:
         agent_updates: dict[str, Any] = {}
         if self.agent is not None:
             agent_override_fields = self.agent.model_dump(include=self.agent.model_fields_set)
@@ -1282,7 +1287,10 @@ def validate(config: AgentSpecificConfig) -> None:
         return self
 
     def _for_each_agent(self, func: Callable[[AgentUnifiedConfig], R]) -> list[R]:
-        agents = [agent.construct_unified_config(default=self) for agent in self.agents]
+        agents = [
+            agent.construct_unified_config(default=self, agent_idx=i)
+            for i, agent in enumerate(self.agents)
+        ]
         if not agents:
             agents.append(self)
 

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

@@ -1298,7 +1298,7 @@ class DockerAgent(AbstractAgent[DockerKernel, DockerKernelCreationContext]):
     monitor_docker_task: asyncio.Task
     agent_sockpath: Path
     agent_sock_task: asyncio.Task
-    metadata_server: MetadataServer
+    metadata_server: Optional[MetadataServer]
     docker_ptask_group: aiotools.PersistentTaskGroup
     gwbridge_subnet: Optional[str]
     checked_invalid_images: Set[str]
@@ -1324,6 +1324,9 @@ def __init__(
             agent_public_key=agent_public_key,
         )
         self.checked_invalid_images = set()
+        # MetadataServer must be shared across all instances of DockerAgent.
+        # metadata_server is initialized by AgentRPCServer and assigned via assign_metadata_server()
+        self.metadata_server = None
 
     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()
 
@@ -1416,6 +1413,9 @@ async def __ainit__(self) -> None:
             blocklist=self.local_config.agent.block_network_plugins,
         )
 
+    def assign_metadata_server(self, metadata_server: MetadataServer) -> None:
+        self.metadata_server = metadata_server
+
     async def shutdown(self, stop_signal: signal.Signals):
         # Stop handling agent sock.
         if self.agent_sock_task is not None:
@@ -1432,7 +1432,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/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()