feat(BA-3026): Extract agent common resources to AgentRuntime

This change introduces AgentRuntime, which contains common resources
shared by all agents. Currently the class is minimal but future multi
agent changes will make more use of the class.

Author: hhoikoo

changes/6728.feature.md

@@ -0,0 +1 @@
+Extract agent common resources to AgentRuntime

src/ai/backend/agent/docker/__init__.py

@@ -1,8 +0,0 @@
-from typing import Type
-
-from ..agent import AbstractAgent
-from .agent import DockerAgent
-
-
-def get_agent_cls() -> Type[AbstractAgent]:
-    return DockerAgent

src/ai/backend/agent/docker/agent.py

@@ -119,14 +119,14 @@
     known_slot_types,
 )
 from ..scratch import create_loop_filesystem, destroy_loop_filesystem
-from ..server import get_extra_volumes
 from ..types import (
     AgentEventData,
     Container,
     KernelOwnershipData,
     LifecycleEvent,
     MountInfo,
     Port,
+    VolumeInfo,
 )
 from ..utils import (
     closing_async,
@@ -161,6 +161,49 @@
     2.39: "ubuntu24.04",
 }
 
+deeplearning_image_keys = {
+    "tensorflow",
+    "caffe",
+    "keras",
+    "torch",
+    "mxnet",
+    "theano",
+}
+
+deeplearning_sample_volume = VolumeInfo(
+    "deeplearning-samples",
+    "/home/work/samples",
+    "ro",
+)
+
+
+async def get_extra_volumes(docker, lang):
+    avail_volumes = (await docker.volumes.list())["Volumes"]
+    if not avail_volumes:
+        return []
+    avail_volume_names = set(v["Name"] for v in avail_volumes)
+
+    # deeplearning specialization
+    # TODO: extract as config
+    volume_list = []
+    for k in deeplearning_image_keys:
+        if k in lang:
+            volume_list.append(deeplearning_sample_volume)
+            break
+
+    # Mount only actually existing volumes
+    mount_list = []
+    for vol in volume_list:
+        if vol.name in avail_volume_names:
+            mount_list.append(vol)
+        else:
+            log.info(
+                "skipped attaching extra volume {0} to a kernel based on image {1}",
+                vol.name,
+                lang,
+            )
+    return mount_list
+
 
 def container_from_docker_container(src: DockerContainer) -> Container:
     ports = []

src/ai/backend/agent/dummy/__init__.py

@@ -1,8 +0,0 @@
-from typing import Type
-
-from ..agent import AbstractAgent
-from .agent import DummyAgent
-
-
-def get_agent_cls() -> Type[AbstractAgent]:
-    return DummyAgent

src/ai/backend/agent/kubernetes/__init__.py

@@ -1,8 +0,0 @@
-from typing import Type
-
-from ..agent import AbstractAgent
-from .agent import KubernetesAgent
-
-
-def get_agent_cls() -> Type[AbstractAgent]:
-    return KubernetesAgent

src/ai/backend/agent/runtime.py

@@ -0,0 +1,75 @@
+import signal
+from typing import Optional
+
+from ai.backend.agent.agent import AbstractAgent
+from ai.backend.agent.config.unified import AgentUnifiedConfig
+from ai.backend.agent.monitor import AgentErrorPluginContext, AgentStatsPluginContext
+from ai.backend.agent.types import AgentBackend
+from ai.backend.common.auth import PublicKey
+from ai.backend.common.etcd import AsyncEtcd
+from ai.backend.common.types import aobject
+
+
+class AgentRuntime(aobject):
+    local_config: AgentUnifiedConfig
+    agent: AbstractAgent
+    etcd: AsyncEtcd
+
+    _stop_signal: signal.Signals
+
+    def __init__(
+        self,
+        local_config: AgentUnifiedConfig,
+        etcd: AsyncEtcd,
+        stats_monitor: AgentStatsPluginContext,
+        error_monitor: AgentErrorPluginContext,
+        agent_public_key: Optional[PublicKey],
+    ) -> None:
+        self.local_config = local_config
+        self.etcd = etcd
+
+        self._stop_signal = signal.SIGTERM
+
+        self.stats_monitor = stats_monitor
+        self.error_monitor = error_monitor
+        self.agent_public_key = agent_public_key
+
+    async def __ainit__(self) -> None:
+        self.agent = await self._create_agent(self.etcd, self.local_config)
+
+    async def __aexit__(self, *exc_info) -> None:
+        await self.agent.shutdown(self._stop_signal)
+
+    def get_agent(self) -> AbstractAgent:
+        return self.agent
+
+    def get_etcd(self) -> AsyncEtcd:
+        return self.etcd
+
+    def mark_stop_signal(self, stop_signal: signal.Signals) -> None:
+        self._stop_signal = stop_signal
+
+    async def _create_agent(
+        self,
+        etcd: AsyncEtcd,
+        agent_config: AgentUnifiedConfig,
+    ) -> AbstractAgent:
+        agent_kwargs = {
+            "stats_monitor": self.stats_monitor,
+            "error_monitor": self.error_monitor,
+            "agent_public_key": self.agent_public_key,
+        }
+
+        match self.local_config.agent_common.backend:
+            case AgentBackend.DOCKER:
+                from .docker.agent import DockerAgent
+
+                return await DockerAgent.new(etcd, agent_config, **agent_kwargs)
+            case AgentBackend.KUBERNETES:
+                from .kubernetes.agent import KubernetesAgent
+
+                return await KubernetesAgent.new(etcd, agent_config, **agent_kwargs)
+            case AgentBackend.DUMMY:
+                from ai.backend.agent.dummy.agent import DummyAgent
+
+                return await DummyAgent.new(etcd, agent_config, **agent_kwargs)