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

@@ -15,6 +15,34 @@
 # If agents field is populated, this field indicates the default values for all
 # agents.
 [agent]
+  # Agent ID
+  ## 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
+  scaling-group-type = "compute"
+  # Allowed compute plugins
+  ## allow-compute-plugins = [ "ai.backend.activator.agent", "ai.backend.accelerator.cuda_open",]
+  # Blocked compute plugins
+  ## block-compute-plugins = [ "ai.backend.accelerator.mock",]
+  # Allowed network plugins
+  ## 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
+  kernel-creation-concurrency = 4
   # Backend type for the agent.
   # This determines how the agent interacts with the underlying infrastructure.
   # Available options are:
@@ -40,14 +68,6 @@
   ## rpc-auth-manager-public-key = "/path/to/public.key"
   # Path to RPC auth agent keypair
   ## rpc-auth-agent-keypair = "/path/to/keypair.key"
-  # Agent socket port
-  agent-sock-port = 6007
-  # Agent ID
-  ## id = "agent-001"
-  # 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
@@ -58,10 +78,6 @@
   ## region = "us-east-1"
   # Instance type
   ## instance-type = "m5.large"
-  # Scaling group name
-  scaling-group = "default"
-  # Scaling group type
-  scaling-group-type = "compute"
   # Path to PID file
   pid-file = "/dev/null"
   # Event loop type
@@ -76,22 +92,6 @@
   metadata-server-bind-host = "0.0.0.0"
   # Metadata server port
   metadata-server-port = 40128
-  # Allowed compute plugins
-  ## allow-compute-plugins = [ "ai.backend.activator.agent", "ai.backend.accelerator.cuda_open",]
-  # Blocked compute plugins
-  ## block-compute-plugins = [ "ai.backend.accelerator.mock",]
-  # Allowed network plugins
-  ## 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
-  kernel-creation-concurrency = 4
   # Whether to use experimental Redis event dispatcher
   use-experimental-redis-event-dispatcher = false
   # Docker mode detected based on kernel version (linuxkit/native)
@@ -114,10 +114,6 @@
   kernel-uid = -1
   # Kernel group ID
   kernel-gid = -1
-  # Bind host for containers
-  bind-host = ""
-  # Advertised host for containers
-  ## advertised-host = "192.168.1.100"
   # Port range for containers.
   # If multiple agents are used, user must ensure that the port ranges
   # do not overlap between the agent, else it may cause subtle issues
@@ -147,6 +143,10 @@
   # networks, and services.
   # This field is only used when backend is set to 'docker'.
   swarm-enabled = false
+  # Bind host for containers
+  bind-host = ""
+  # Advertised host for containers
+  ## advertised-host = "192.168.1.100"
 
   # KRunner volumes configuration, mapping container names to host paths.
   # This is used to specify volumes that should be mounted into containers
@@ -368,13 +368,17 @@
   # Only override fields if necessary.
   [agents.agent]
     # Agent ID
-    id = "agent-001"
+    ## id = "agent-001"
     # Agent socket port
-    ## agent-sock-port = 6007  # min=1024 max=65535
-    # Mount path for containers
-    ## mount-path = "/mnt/backend.ai"
-    # Whether to enable cohabiting storage proxy
-    ## cohabiting-storage-proxy = true
+    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
+    scaling-group-type = "compute"
     # Allowed compute plugins
     ## allow-compute-plugins = [ "ai.backend.activator.agent", "ai.backend.accelerator.cuda_open",]
     # Blocked compute plugins
@@ -383,14 +387,14 @@
     ## 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 = true
+    force-terminate-abusing-containers = false
     # Kernel creation concurrency
-    ## kernel-creation-concurrency = 4  # min=1 max=32
-    # Docker mode detected based on kernel version (linuxkit/native)
-    ## docker-mode = "linuxkit"
-    # Owner uid:gid of the mount directory
-    ## mount-path-uid-gid = "root:root"
+    kernel-creation-concurrency = 4
 
     # Container lifecycle synchronization config
     [agents.agent.sync-container-lifecycles]
@@ -402,26 +406,26 @@
   # Container config overrides for the individual agent
   [agents.container]
     # Kernel user ID
-    ## kernel-uid = 1000
+    kernel-uid = -1
     # Kernel group ID
-    ## kernel-gid = 1000
+    kernel-gid = -1
     # Port range for containers.
     # If multiple agents are used, user must ensure that the port ranges
     # do not overlap between the agent, else it may cause subtle issues
     # late into the agent's runtime.
-    ## port-range = [ 30000, 31000,]
+    port-range = [ 30000, 31000,]
     # Statistics type
     ## stats-type = "cgroup"
     # Sandbox type
-    ## sandbox-type = "docker"
+    sandbox-type = "docker"
     # Jail arguments
-    ## jail-args = [ "--mount", "/tmp",]
+    jail-args = [ "--mount", "/tmp",]
     # Scratch type
-    ## scratch-type = "hostdir"
+    scratch-type = "hostdir"
     # Scratch root directory
-    ## scratch-root = "./scratches"
+    scratch-root = "scratches"
     # Scratch size
-    ## scratch-size = "1G"
+    scratch-size = 0
     # Scratch NFS address
     ## scratch-nfs-address = "192.168.1.100:/export"
     # Scratch NFS options
@@ -433,7 +437,7 @@
     # When enabled, the agent will use Docker Swarm APIs to manage containers,
     # networks, and services.
     # This field is only used when backend is set to 'docker'.
-    ## swarm-enabled = true
+    swarm-enabled = false
 
   # Resource config overrides for the individual agent
   [agents.resource]

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)