fix(BA-2753): Disaggregate components not supposed to be shared

This change pull out many connections/components like etcd connections
etc. that are not supposed to be shared between the multi agent
instances out such that they are distinct from each other.

Author: hhoikoo

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)

src/ai/backend/agent/config/unified.py

@@ -687,6 +687,46 @@ class OverridableAgentConfig(BaseConfigSchema):
         extra="allow",
     )
 
+    def with_derived_fields(
+        self,
+        default_agent: AgentConfig,
+        agent_idx: int,
+    ) -> OverridableAgentConfig:
+        if self.id is None and default_agent.id is None:
+            return self
+
+        agent_id = self.id if self.id is not None else f"{default_agent.id}-{agent_idx}"
+
+        if "ipc_base_path" in self.model_fields_set:
+            ipc_base_path = self.ipc_base_path
+        else:
+            ipc_base_path = default_agent.ipc_base_path / agent_id
+
+        if "var_base_path" in self.model_fields_set:
+            var_base_path = self.var_base_path
+        else:
+            var_base_path = default_agent.var_base_path / agent_id
+
+        if "image_commit_path" in self.model_fields_set:
+            image_commit_path = self.image_commit_path
+        else:
+            image_commit_path = default_agent.image_commit_path / agent_id
+
+        if default_agent.abuse_report_path is None or "abuse_report_path" in self.model_fields_set:
+            abuse_report_path = self.abuse_report_path
+        else:
+            abuse_report_path = default_agent.abuse_report_path / agent_id
+
+        return self.model_copy(
+            update={
+                "id": agent_id,
+                "ipc_base_path": ipc_base_path,
+                "var_base_path": var_base_path,
+                "image_commit_path": image_commit_path,
+                "abuse_report_path": abuse_report_path,
+            }
+        )
+
 
 class AgentConfig(CommonAgentConfig, OverridableAgentConfig):
     pass
@@ -1130,10 +1170,16 @@ 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)
+            agent = self.agent.with_derived_fields(default.agent, agent_idx)
+            agent_override_fields = agent.model_dump(include=agent.model_fields_set)
             agent_updates["agent"] = default.agent.model_copy(update=agent_override_fields)
         if self.container is not None:
             container_override_fields = self.container.model_dump(
@@ -1169,9 +1215,20 @@ 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
 
     def with_updates(
         self,
@@ -1257,7 +1314,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/etcd.py

@@ -0,0 +1,52 @@
+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)
+
+    def get_client(self, agent_id: AgentId) -> AsyncEtcd:
+        return self._etcd_clients[agent_id]
+
+    def prefill_clients(self, prefill_data: Collection[tuple[AgentId, AgentUnifiedConfig]]) -> None:
+        for agent_id, agent_config in prefill_data:
+            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 is not None and self._etcd_config.password is not None:
+            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,
+        )