feat(BA-3023): Add unified multi agent etcd client handling

This change introduces EtcdClientRegistry, which is a class that
can contain etcd clients for multiple agents, each with its own prefix
information at scaling group and individual node level.

Author: hhoikoo

changes/6721.feature.md

@@ -0,0 +1 @@
+Add EtcdClientRegistry for clean handling of etcd clients for multi agents

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/server.py

@@ -49,6 +49,7 @@
 from setproctitle import setproctitle
 from zmq.auth.certs import load_certificate
 
+from ai.backend.agent.etcd import EtcdClientRegistry
 from ai.backend.agent.metrics.metric import RPCMetricObserver
 from ai.backend.agent.resources import scan_gpu_alloc_map
 from ai.backend.common import config, identity, msgpack, utils
@@ -96,6 +97,7 @@
     ServiceMetadata,
 )
 from ai.backend.common.types import (
+    AgentId,
     ClusterInfo,
     CommitStatus,
     ContainerId,
@@ -309,6 +311,7 @@ class AgentRPCServer(aobject):
 
     loop: asyncio.AbstractEventLoop
     agent: AbstractAgent
+    etcd_client_registry: EtcdClientRegistry
     rpc_server: Peer
     rpc_addr: str
     agent_addr: str
@@ -318,17 +321,19 @@ class AgentRPCServer(aobject):
 
     def __init__(
         self,
-        etcd: AsyncEtcd,
+        etcd_client_registry: EtcdClientRegistry,
         local_config: AgentUnifiedConfig,
         *,
         skip_detect_manager: bool = False,
     ) -> None:
         self.loop = current_loop()
-        self.etcd = etcd
+        self.etcd_client_registry = etcd_client_registry
         self.local_config = local_config
         self.skip_detect_manager = skip_detect_manager
         self._stop_signal = signal.SIGTERM
 
+        self.etcd_client_registry.prefill_clients(self.local_config.agent_configs)
+
     async def __ainit__(self) -> None:
         # Start serving requests.
         await self.update_status("starting")
@@ -339,11 +344,12 @@ async def __ainit__(self) -> None:
         await self.read_agent_config()
         await self.read_agent_config_container()
 
+        global_etcd = self.etcd_client_registry.global_etcd
         self.stats_monitor = AgentStatsPluginContext(
-            self.etcd, self.local_config.model_dump(by_alias=True)
+            global_etcd, self.local_config.model_dump(by_alias=True)
         )
         self.error_monitor = AgentErrorPluginContext(
-            self.etcd, self.local_config.model_dump(by_alias=True)
+            global_etcd, self.local_config.model_dump(by_alias=True)
         )
         await self.stats_monitor.init()
         await self.error_monitor.init()
@@ -380,7 +386,7 @@ async def __ainit__(self) -> None:
         backend = self.local_config.agent_common.backend
         agent_mod = importlib.import_module(f"ai.backend.agent.{backend.value}")
         self.agent = await agent_mod.get_agent_cls().new(  # type: ignore
-            self.etcd,
+            self.etcd_client_registry.get_client(AgentId(self.local_config.agent.id)),
             self.local_config,
             stats_monitor=self.stats_monitor,
             error_monitor=self.error_monitor,
@@ -422,13 +428,13 @@ async def _debug_server_task():
 
         self.debug_server_task = asyncio.create_task(_debug_server_task())
 
-        await self.etcd.put("ip", rpc_addr.host, scope=ConfigScopes.NODE)
+        await self.agent.etcd.put("ip", rpc_addr.host, scope=ConfigScopes.NODE)
 
         watcher_port = utils.nmget(
             self.local_config.model_dump(), "watcher.service-addr.port", None
         )
         if watcher_port is not None:
-            await self.etcd.put("watcher_port", watcher_port, scope=ConfigScopes.NODE)
+            await self.agent.etcd.put("watcher_port", watcher_port, scope=ConfigScopes.NODE)
 
         await self.update_status("running")
 
@@ -472,10 +478,11 @@ def _ensure_serializable(o) -> Any:
 
     async def detect_manager(self):
         log.info("detecting the manager...")
-        manager_instances = await self.etcd.get_prefix("nodes/manager")
+        global_etcd = self.etcd_client_registry.global_etcd
+        manager_instances = await global_etcd.get_prefix("nodes/manager")
         if not manager_instances:
             log.warning("watching etcd to wait for the manager being available")
-            async with aclosing(self.etcd.watch_prefix("nodes/manager")) as agen:
+            async with aclosing(global_etcd.watch_prefix("nodes/manager")) as agen:  # type: ignore
                 async for ev in agen:
                     match ev:
                         case QueueSentinel.CLOSED | QueueSentinel.TIMEOUT:
@@ -486,9 +493,11 @@ async def detect_manager(self):
         log.info("detected at least one manager running")
 
     async def read_agent_config(self):
+        global_etcd = self.etcd_client_registry.global_etcd
+
         # Fill up Redis configs from etcd and store as separate attributes
         self._redis_config = config.redis_config_iv.check(
-            await self.etcd.get_prefix("config/redis"),
+            await global_etcd.get_prefix("config/redis"),
         )
         log.info("configured redis: {0}", self._redis_config)
 
@@ -506,7 +515,7 @@ async def read_agent_config(self):
         # Fill up vfolder configs from etcd and store as separate attributes
         # TODO: Integrate vfolder_config into local_config
         self._vfolder_config = config.vfolder_config_iv.check(
-            await self.etcd.get_prefix("volumes"),
+            await global_etcd.get_prefix("volumes"),
         )
         if self._vfolder_config["mount"] is None:
             log.info(
@@ -517,7 +526,7 @@ async def read_agent_config(self):
             log.info("configured vfolder fs prefix: {0}", self._vfolder_config["fsprefix"])
 
         # Fill up shared agent configurations from etcd.
-        agent_etcd_config_raw = await self.etcd.get_prefix("config/agent")
+        agent_etcd_config_raw = await global_etcd.get_prefix("config/agent")
         if agent_etcd_config_raw:
             try:
                 # Parse specific etcd configs and update the unified config
@@ -541,7 +550,9 @@ async def read_agent_config(self):
     async def read_agent_config_container(self):
         # Fill up global container configurations from etcd.
         try:
-            container_etcd_config_raw = await self.etcd.get_prefix("config/container")
+            container_etcd_config_raw = await self.etcd_client_registry.global_etcd.get_prefix(
+                "config/container"
+            )
             if container_etcd_config_raw:
                 # Update config by creating a new instance with modified values
                 container_updates = {}
@@ -586,7 +597,7 @@ async def __aexit__(self, *exc_info) -> None:
 
     @collect_error
     async def update_status(self, status):
-        await self.etcd.put("", status, scope=ConfigScopes.NODE)
+        await self.agent.etcd.put("", status, scope=ConfigScopes.NODE)
 
     @rpc_function
     @collect_error
@@ -1221,29 +1232,12 @@ async def aiomonitor_ctx(
 
 
 @asynccontextmanager
-async def etcd_ctx(local_config: AgentUnifiedConfig) -> AsyncGenerator[AsyncEtcd]:
-    etcd_credentials = None
-    if local_config.etcd.user and local_config.etcd.password:
-        etcd_credentials = {
-            "user": local_config.etcd.user,
-            "password": local_config.etcd.password,
-        }
-    scope_prefix_map = {
-        ConfigScopes.GLOBAL: "",
-        ConfigScopes.SGROUP: f"sgroup/{local_config.agent.scaling_group}",
-        ConfigScopes.NODE: f"nodes/agents/{local_config.agent.id}",
-    }
-    etcd_config_data = local_config.etcd.to_dataclass()
-    etcd = AsyncEtcd(
-        [addr.to_legacy() for addr in etcd_config_data.addrs],
-        local_config.etcd.namespace,
-        scope_prefix_map,
-        credentials=etcd_credentials,
-    )
+async def etcd_ctx(local_config: AgentUnifiedConfig) -> AsyncGenerator[EtcdClientRegistry]:
+    etcd_client_registry = EtcdClientRegistry(local_config.etcd.to_dataclass())
     try:
-        yield etcd
+        yield etcd_client_registry
     finally:
-        await etcd.close()
+        await etcd_client_registry.close()
 
 
 async def prepare_krunner_volumes(local_config: AgentUnifiedConfig) -> AgentUnifiedConfig:
@@ -1253,7 +1247,7 @@ async def prepare_krunner_volumes(local_config: AgentUnifiedConfig) -> AgentUnif
     )
     krunner_volumes: Mapping[str, str] = await kernel_mod.prepare_krunner_env(
         local_config.model_dump(by_alias=True)
-    )  # type: ignore
+    )
     # TODO: merge k8s branch: nfs_mount_path = local_config['baistatic']['mounted-at']
     log.info("Kernel runner environments: {}", [*krunner_volumes.keys()])
     return local_config.with_updates(container_update={"krunner_volumes": krunner_volumes})
@@ -1317,10 +1311,10 @@ async def auto_detect_agent_network(
 
 @asynccontextmanager
 async def agent_server_ctx(
-    local_config: AgentUnifiedConfig, etcd: AsyncEtcd
+    local_config: AgentUnifiedConfig, etcd_client_registry: EtcdClientRegistry
 ) -> AsyncGenerator[AgentRPCServer]:
     agent_server = await AgentRPCServer.new(
-        etcd,
+        etcd_client_registry,
         local_config,
         skip_detect_manager=local_config.agent_common.skip_manager_detection,
     )
@@ -1419,18 +1413,19 @@ async def server_main(
         local_config = await auto_detect_agent_identity(local_config)
 
         # etcd's scope-prefix map depends on the auto-detected identity info.
-        etcd = await agent_init_stack.enter_async_context(etcd_ctx(local_config))
-        local_config = await auto_detect_agent_network(local_config, etcd)
-        plugins = await etcd.get_prefix_dict("config/plugins/accelerator")
+        etcd_client_registry = await agent_init_stack.enter_async_context(etcd_ctx(local_config))
+        global_etcd = etcd_client_registry.global_etcd
+        local_config = await auto_detect_agent_network(local_config, global_etcd)
+        plugins = await global_etcd.get_prefix_dict("config/plugins/accelerator")
         local_config = local_config.with_changes(plugins=plugins)
 
         # Start RPC server.
         agent_server = await agent_init_stack.enter_async_context(
-            agent_server_ctx(local_config, etcd)
+            agent_server_ctx(local_config, etcd_client_registry)
         )
         monitor.console_locals["agent_server"] = agent_server
 
-        await agent_init_stack.enter_async_context(service_discovery_ctx(etcd, agent_server))
+        await agent_init_stack.enter_async_context(service_discovery_ctx(global_etcd, agent_server))
         log.info("Started the agent service.")
     except Exception:
         log.exception("Server initialization failure; triggering shutdown...")

tests/agent/test_agent.py

@@ -19,7 +19,9 @@
 
 
 class Dummy:
-    pass
+    def prefill_clients(self, prefill_data):
+        """Mock implementation of prefill_clients for testing"""
+        pass
 
 
 kgid = "kernel-gid"
@@ -29,8 +31,9 @@ class Dummy:
 
 @pytest.fixture
 async def arpcs_no_ainit(test_id, redis_container):
-    etcd = Dummy()
-    etcd.get_prefix = None
+    etcd_client_registry = Dummy()
+    etcd_client_registry.global_etcd = Dummy()
+    etcd_client_registry.global_etcd.get_prefix = None
 
     # Create a minimal pydantic config for testing
     config = AgentUnifiedConfig(
@@ -40,14 +43,18 @@ async def arpcs_no_ainit(test_id, redis_container):
         etcd=EtcdConfig(namespace="test", addr=HostPortPair(host="127.0.0.1", port=2379)),
     )
 
-    ars = AgentRPCServer(etcd=etcd, local_config=config, skip_detect_manager=True)
+    ars = AgentRPCServer(
+        etcd_client_registry=etcd_client_registry, local_config=config, skip_detect_manager=True
+    )
     yield ars
 
 
 @pytest.mark.asyncio
 async def test_read_agent_config_container_invalid01(arpcs_no_ainit, mocker):
     inspect_mock = AsyncMock(return_value={"a": 1, "b": 2})
-    mocker.patch.object(arpcs_no_ainit.etcd, "get_prefix", new=inspect_mock)
+    mocker.patch.object(
+        arpcs_no_ainit.etcd_client_registry.global_etcd, "get_prefix", new=inspect_mock
+    )
     await arpcs_no_ainit.read_agent_config_container()
     # Check that kernel-gid and kernel-uid are still at their default values (converted from -1)
     assert (
@@ -61,7 +68,9 @@ async def test_read_agent_config_container_invalid01(arpcs_no_ainit, mocker):
 @pytest.mark.asyncio
 async def test_read_agent_config_container_invalid02(arpcs_no_ainit, mocker):
     inspect_mock = AsyncMock(return_value={})
-    mocker.patch.object(arpcs_no_ainit.etcd, "get_prefix", new=inspect_mock)
+    mocker.patch.object(
+        arpcs_no_ainit.etcd_client_registry.global_etcd, "get_prefix", new=inspect_mock
+    )
     await arpcs_no_ainit.read_agent_config_container()
     # Check that kernel-gid and kernel-uid are still at their default values (converted from -1)
     assert (
@@ -75,7 +84,9 @@ async def test_read_agent_config_container_invalid02(arpcs_no_ainit, mocker):
 @pytest.mark.asyncio
 async def test_read_agent_config_container_1valid(arpcs_no_ainit, mocker):
     inspect_mock = AsyncMock(return_value={kgid: 10})
-    mocker.patch.object(arpcs_no_ainit.etcd, "get_prefix", new=inspect_mock)
+    mocker.patch.object(
+        arpcs_no_ainit.etcd_client_registry.global_etcd, "get_prefix", new=inspect_mock
+    )
     await arpcs_no_ainit.read_agent_config_container()
 
     assert arpcs_no_ainit.local_config.container.kernel_gid.real == 10
@@ -87,7 +98,9 @@ async def test_read_agent_config_container_1valid(arpcs_no_ainit, mocker):
 @pytest.mark.asyncio
 async def test_read_agent_config_container_2valid(arpcs_no_ainit, mocker):
     inspect_mock = AsyncMock(return_value={kgid: 10, kuid: 20})
-    mocker.patch.object(arpcs_no_ainit.etcd, "get_prefix", new=inspect_mock)
+    mocker.patch.object(
+        arpcs_no_ainit.etcd_client_registry.global_etcd, "get_prefix", new=inspect_mock
+    )
     await arpcs_no_ainit.read_agent_config_container()
 
     assert arpcs_no_ainit.local_config.container.kernel_gid.real == 10