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

configs/agent/sample.toml

@@ -16,13 +16,9 @@
 # agents.
 [agent]
   # Agent ID
-  ## id = "agent-001"
+  id = "agent-001"
   # Agent socket port
   agent-sock-port = 6007
-  # Base path for IPC
-  ipc-base-path = "/tmp/backend.ai/ipc"
-  # Base path for variable data
-  var-base-path = "var/lib/backend.ai"
   # Scaling group name
   scaling-group = "default"
   # Scaling group type
@@ -35,10 +31,6 @@
   ## allow-network-plugins = [ "ai.backend.manager.network.overlay",]
   # Blocked network plugins
   ## block-network-plugins = [ "ai.backend.manager.network.overlay",]
-  # Path for image commit
-  image-commit-path = "tmp/backend.ai/commit"
-  # Path for abuse reports
-  ## abuse-report-path = "/var/log/backend.ai/abuse"
   # Whether to force terminate abusing containers
   force-terminate-abusing-containers = false
   # Kernel creation concurrency
@@ -68,6 +60,10 @@
   ## rpc-auth-manager-public-key = "/path/to/public.key"
   # Path to RPC auth agent keypair
   ## rpc-auth-agent-keypair = "/path/to/keypair.key"
+  # Base path for IPC
+  ipc-base-path = "/tmp/backend.ai/ipc"
+  # Base path for variable data
+  var-base-path = "var/lib/backend.ai"
   # Mount path for containers
   ## mount-path = "/mnt/backend.ai"
   # Whether to enable cohabiting storage proxy
@@ -92,6 +88,10 @@
   metadata-server-bind-host = "0.0.0.0"
   # Metadata server port
   metadata-server-port = 40128
+  # Path for image commit
+  image-commit-path = "tmp/backend.ai/commit"
+  # Path for abuse reports
+  ## abuse-report-path = "/var/log/backend.ai/abuse"
   # Whether to use experimental Redis event dispatcher
   use-experimental-redis-event-dispatcher = false
   # Docker mode detected based on kernel version (linuxkit/native)
@@ -368,13 +368,9 @@
   # Only override fields if necessary.
   [agents.agent]
     # Agent ID
-    ## id = "agent-001"
+    id = "agent-001"
     # Agent socket port
     agent-sock-port = 6007
-    # Base path for IPC
-    ipc-base-path = "/tmp/backend.ai/ipc"
-    # Base path for variable data
-    var-base-path = "var/lib/backend.ai"
     # Scaling group name
     scaling-group = "default"
     # Scaling group type
@@ -387,10 +383,6 @@
     ## allow-network-plugins = [ "ai.backend.manager.network.overlay",]
     # Blocked network plugins
     ## block-network-plugins = [ "ai.backend.manager.network.overlay",]
-    # Path for image commit
-    image-commit-path = "tmp/backend.ai/commit"
-    # Path for abuse reports
-    ## abuse-report-path = "/var/log/backend.ai/abuse"
     # Whether to force terminate abusing containers
     force-terminate-abusing-containers = false
     # Kernel creation concurrency

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

@@ -22,6 +22,7 @@
     Sequence,
     TypeVar,
 )
+from uuid import uuid4
 
 from pydantic import (
     AfterValidator,
@@ -48,6 +49,7 @@
     UserID,
 )
 from ai.backend.common.types import (
+    AgentId,
     BinarySize,
     BinarySizeField,
     ResourceGroupType,
@@ -438,6 +440,20 @@ class CommonAgentConfig(BaseConfigSchema):
         validation_alias=AliasChoices("rpc-auth-agent-keypair", "rpc_auth_agent_keypair"),
         serialization_alias="rpc-auth-agent-keypair",
     )
+    ipc_base_path: AutoDirectoryPath = Field(
+        default=AutoDirectoryPath("/tmp/backend.ai/ipc"),
+        description="Base path for IPC",
+        examples=["/tmp/backend.ai/ipc"],
+        validation_alias=AliasChoices("ipc-base-path", "ipc_base_path"),
+        serialization_alias="ipc-base-path",
+    )
+    var_base_path: AutoDirectoryPath = Field(
+        default=AutoDirectoryPath("./var/lib/backend.ai"),
+        description="Base path for variable data",
+        examples=["./var/lib/backend.ai"],
+        validation_alias=AliasChoices("var-base-path", "var_base_path"),
+        serialization_alias="var-base-path",
+    )
     mount_path: Optional[AutoDirectoryPath] = Field(
         default=None,
         description="Mount path for containers",
@@ -528,6 +544,20 @@ class CommonAgentConfig(BaseConfigSchema):
         validation_alias=AliasChoices("metadata-server-port", "metadata_server_port"),
         serialization_alias="metadata-server-port",
     )
+    image_commit_path: AutoDirectoryPath = Field(
+        default=AutoDirectoryPath("./tmp/backend.ai/commit"),
+        description="Path for image commit",
+        examples=["./tmp/backend.ai/commit"],
+        validation_alias=AliasChoices("image-commit-path", "image_commit_path"),
+        serialization_alias="image-commit-path",
+    )
+    abuse_report_path: Optional[Path] = Field(
+        default=None,
+        description="Path for abuse reports",
+        examples=["/var/log/backend.ai/abuse"],
+        validation_alias=AliasChoices("abuse-report-path", "abuse_report_path"),
+        serialization_alias="abuse-report-path",
+    )
     use_experimental_redis_event_dispatcher: bool = Field(
         default=False,
         description="Whether to use experimental Redis event dispatcher",
@@ -574,8 +604,8 @@ def _validate_rpc_listen_addr(cls, rpc_listen_addr: HostPortPair) -> HostPortPai
 
 
 class OverridableAgentConfig(BaseConfigSchema):
-    id: Optional[str] = Field(
-        default=None,
+    id: str = Field(
+        default_factory=lambda: f"agent-{uuid4()}",
         description="Agent ID",
         examples=["agent-001"],
     )
@@ -588,20 +618,6 @@ class OverridableAgentConfig(BaseConfigSchema):
         validation_alias=AliasChoices("agent-sock-port", "agent_sock_port"),
         serialization_alias="agent-sock-port",
     )
-    ipc_base_path: AutoDirectoryPath = Field(
-        default=AutoDirectoryPath("/tmp/backend.ai/ipc"),
-        description="Base path for IPC",
-        examples=["/tmp/backend.ai/ipc"],
-        validation_alias=AliasChoices("ipc-base-path", "ipc_base_path"),
-        serialization_alias="ipc-base-path",
-    )
-    var_base_path: AutoDirectoryPath = Field(
-        default=AutoDirectoryPath("./var/lib/backend.ai"),
-        description="Base path for variable data",
-        examples=["./var/lib/backend.ai"],
-        validation_alias=AliasChoices("var-base-path", "var_base_path"),
-        serialization_alias="var-base-path",
-    )
     scaling_group: str = Field(
         default="default",
         description="Scaling group name",
@@ -644,20 +660,6 @@ class OverridableAgentConfig(BaseConfigSchema):
         validation_alias=AliasChoices("block-network-plugins", "block_network_plugins"),
         serialization_alias="block-network-plugins",
     )
-    image_commit_path: AutoDirectoryPath = Field(
-        default=AutoDirectoryPath("./tmp/backend.ai/commit"),
-        description="Path for image commit",
-        examples=["./tmp/backend.ai/commit"],
-        validation_alias=AliasChoices("image-commit-path", "image_commit_path"),
-        serialization_alias="image-commit-path",
-    )
-    abuse_report_path: Optional[Path] = Field(
-        default=None,
-        description="Path for abuse reports",
-        examples=["/var/log/backend.ai/abuse"],
-        validation_alias=AliasChoices("abuse-report-path", "abuse_report_path"),
-        serialization_alias="abuse-report-path",
-    )
     force_terminate_abusing_containers: bool = Field(
         default=False,
         description="Whether to force terminate abusing containers",
@@ -1138,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)
@@ -1177,9 +1184,24 @@ class AgentUnifiedConfig(AgentGlobalConfig, AgentSpecificConfig):
         extra="allow",
     )
 
+    @property
+    def agent_common(self) -> CommonAgentConfig:
+        return self.agent
+
+    @property
+    def agent_default(self) -> OverridableAgentConfig:
+        return self.agent
+
     @property
     def agent_configs(self) -> Sequence[AgentUnifiedConfig]:
-        return self._for_each_agent(lambda x: x)
+        agent_configs = self._for_each_agent(lambda x: x)
+        if not agent_configs:
+            raise ValueError("There must be at least one agent config")
+        return agent_configs
+
+    @property
+    def agent_ids(self) -> Sequence[AgentId]:
+        return [AgentId(agent_config.agent.id) for agent_config in self.agent_configs]
 
     def with_updates(
         self,
@@ -1265,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()