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

This change introduces AgentEtcdClientView, which is a class that acts
as an adaptor layer for ensuring that the config scope of etcd is always
in sync with the specific agent's scaling group and agent ID.

Author: hhoikoo

changes/6721.feature.md

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

src/ai/backend/agent/agent.py

@@ -64,6 +64,7 @@
 )
 from trafaret import DataError
 
+from ai.backend.agent.etcd import AgentEtcdClientView
 from ai.backend.agent.metrics.metric import (
     StatScope,
     StatTaskObserver,
@@ -762,7 +763,7 @@ class AbstractAgent(
     id: AgentId
     loop: asyncio.AbstractEventLoop
     local_config: AgentUnifiedConfig
-    etcd: AsyncEtcd
+    etcd: AgentEtcdClientView
     local_instance_id: str
     kernel_registry: MutableMapping[KernelId, AbstractKernel]
     computers: MutableMapping[DeviceName, ComputerContext]
@@ -839,7 +840,7 @@ def __init__(
     ) -> None:
         self._skip_initial_scan = skip_initial_scan
         self.loop = current_loop()
-        self.etcd = etcd
+        self.etcd = AgentEtcdClientView(etcd, config_container=self)
         self.local_config = local_config
         self.id = AgentId(local_config.agent.id or f"agent-{uuid4()}")
         self.local_instance_id = generate_local_instance_id(__file__)

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

@@ -16,6 +16,7 @@
     Any,
     Mapping,
     Optional,
+    Protocol,
     Self,
     Sequence,
     TypeVar,
@@ -1309,3 +1310,7 @@ def _validate_agent_configs(self) -> Self:
             config.validate_agent_specific_config()
 
         return self
+
+
+class AgentConfigContainer(Protocol):
+    local_config: AgentUnifiedConfig

src/ai/backend/agent/etcd.py

@@ -0,0 +1,231 @@
+from collections.abc import AsyncGenerator, Iterable, Mapping
+from typing import Optional, override
+
+from etcd_client import CondVar
+
+from ai.backend.agent.config.unified import AgentConfigContainer
+from ai.backend.common.etcd import (
+    AbstractKVStore,
+    AsyncEtcd,
+    ConfigScopes,
+    Event,
+    GetPrefixValue,
+    NestedStrKeyedMapping,
+)
+from ai.backend.common.exception import (
+    BackendAIError,
+    ErrorCode,
+    ErrorDetail,
+    ErrorDomain,
+    ErrorOperation,
+)
+from ai.backend.common.types import QueueSentinel
+
+
+class AgentEtcdError(BackendAIError):
+    @classmethod
+    def error_code(cls) -> ErrorCode:
+        return ErrorCode(
+            domain=ErrorDomain.AGENT,
+            operation=ErrorOperation.SETUP,
+            error_detail=ErrorDetail.NOT_FOUND,
+        )
+
+
+# AgentEtcdClientView inherits from AsyncEtcd, but really it's just composing an AsyncEtcd instance
+# and acting as an adaptor. Inheritance is made only to make the type checker happy, and the places
+# that use AsyncEtcd really should not take the concrete implementation type AsyncEtcd, but rather
+# the interface type AbstractKVStore. In the current codebase, manually modifying places that
+# currently take in an AsyncEtcd instance to instead take in AbstractKVStore would be too invasive.
+class AgentEtcdClientView(AsyncEtcd, AbstractKVStore):
+    _etcd: AsyncEtcd
+    # We take in a config container instead of the unified config object itself, because config
+    # objects are never mutated, but rather replaced with a copy.
+    _config_container: AgentConfigContainer
+
+    def __init__(
+        self,
+        etcd: AsyncEtcd,
+        config_container: AgentConfigContainer,
+    ) -> None:
+        self._etcd = etcd
+        self._config_container = config_container
+
+    def _augment_scope_prefix_map(
+        self,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]],
+    ) -> Mapping[ConfigScopes, str]:
+        if scope_prefix_map is None:
+            scope_prefix_map = {}
+
+        agent_config = self._config_container.local_config.agent
+        return {
+            **scope_prefix_map,
+            ConfigScopes.SGROUP: f"sgroup/{agent_config.scaling_group}",
+            ConfigScopes.NODE: f"nodes/agents/{agent_config.id}",
+        }
+
+    @override
+    async def put(
+        self,
+        key: str,
+        val: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ):
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        await self._etcd.put(key, val, scope=scope, scope_prefix_map=scope_prefix_map)
+
+    @override
+    async def put_prefix(
+        self,
+        key: str,
+        dict_obj: NestedStrKeyedMapping,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ):
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        await self._etcd.put_prefix(key, dict_obj, scope=scope, scope_prefix_map=scope_prefix_map)
+
+    @override
+    async def put_dict(
+        self,
+        flattened_dict_obj: Mapping[str, str],
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ):
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        await self._etcd.put_dict(
+            flattened_dict_obj, scope=scope, scope_prefix_map=scope_prefix_map
+        )
+
+    @override
+    async def get(
+        self,
+        key: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.MERGED,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ) -> Optional[str]:
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        return await self._etcd.get(key, scope=scope, scope_prefix_map=scope_prefix_map)
+
+    @override
+    async def get_prefix(
+        self,
+        key_prefix: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.MERGED,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ) -> GetPrefixValue:
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        return await self._etcd.get_prefix(
+            key_prefix,
+            scope=scope,
+            scope_prefix_map=scope_prefix_map,
+        )
+
+    @override
+    async def replace(
+        self,
+        key: str,
+        initial_val: str,
+        new_val: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ) -> bool:
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        return await self._etcd.replace(
+            key,
+            initial_val,
+            new_val,
+            scope=scope,
+            scope_prefix_map=scope_prefix_map,
+        )
+
+    @override
+    async def delete(
+        self,
+        key: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ):
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        await self._etcd.delete(key, scope=scope, scope_prefix_map=scope_prefix_map)
+
+    @override
+    async def delete_multi(
+        self,
+        keys: Iterable[str],
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ):
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        await self._etcd.delete_multi(keys, scope=scope, scope_prefix_map=scope_prefix_map)
+
+    @override
+    async def delete_prefix(
+        self,
+        key_prefix: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+    ):
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        await self._etcd.delete_prefix(key_prefix, scope=scope, scope_prefix_map=scope_prefix_map)
+
+    @override
+    async def watch(
+        self,
+        key: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+        once: bool = False,
+        ready_event: Optional[CondVar] = None,
+        cleanup_event: Optional[CondVar] = None,
+        wait_timeout: Optional[float] = None,
+    ) -> AsyncGenerator[QueueSentinel | Event, None]:
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        watch_result = self._etcd.watch(
+            key,
+            scope=scope,
+            scope_prefix_map=scope_prefix_map,
+            once=once,
+            ready_event=ready_event,
+            cleanup_event=cleanup_event,
+            wait_timeout=wait_timeout,
+        )
+        async for item in watch_result:
+            yield item
+
+    @override
+    async def watch_prefix(
+        self,
+        key_prefix: str,
+        *,
+        scope: ConfigScopes = ConfigScopes.GLOBAL,
+        scope_prefix_map: Optional[Mapping[ConfigScopes, str]] = None,
+        once: bool = False,
+        ready_event: Optional[CondVar] = None,
+        cleanup_event: Optional[CondVar] = None,
+        wait_timeout: Optional[float] = None,
+    ) -> AsyncGenerator[QueueSentinel | Event, None]:
+        scope_prefix_map = self._augment_scope_prefix_map(scope_prefix_map)
+        watch_prefix_result = self._etcd.watch_prefix(
+            key_prefix,
+            scope=scope,
+            scope_prefix_map=scope_prefix_map,
+            once=once,
+            ready_event=ready_event,
+            cleanup_event=cleanup_event,
+            wait_timeout=wait_timeout,
+        )
+        async for item in watch_prefix_result:
+            yield item

src/ai/backend/agent/server.py

@@ -309,6 +309,7 @@ class AgentRPCServer(aobject):
 
     loop: asyncio.AbstractEventLoop
     agent: AbstractAgent
+    etcd: AsyncEtcd
     rpc_server: Peer
     rpc_addr: str
     agent_addr: str
@@ -1253,7 +1254,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})