new network plugin architecture for agent

src/ai/backend/agent/BUILD

@@ -62,6 +62,9 @@ python_distribution(
             "ag": "ai.backend.agent.cli:main",
             "ag.start-server": "ai.backend.agent.server:main",
         },
+        "backendai_network_agent_v1": {
+            "overlay": "ai.backend.agent.docker.intrinsic:OverlayNetworkPlugin",
+        }
     },
     generate_setup=True,
     tags=["wheel"],

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

@@ -13,14 +13,18 @@
 from aiodocker.docker import Docker, DockerContainer
 from aiodocker.exceptions import DockerError
 
+from ai.backend.agent.docker.kernel import DockerKernel
+from ai.backend.agent.plugin.network import AbstractNetworkAgentPlugin
 from ai.backend.agent.types import MountInfo
 from ai.backend.common.logging import BraceStyleAdapter
 from ai.backend.common.netns import nsenter
 from ai.backend.common.types import (
     AcceleratorMetadata,
+    ClusterInfo,
     DeviceId,
     DeviceModelInfo,
     DeviceName,
+    KernelCreationConfig,
     MetricKey,
     SlotName,
     SlotTypes,
@@ -912,3 +916,29 @@ def get_metadata(self) -> AcceleratorMetadata:
             "number_format": {"binary": True, "round_length": 0},
             "display_icon": "ram",
         }
+
+
+class OverlayNetworkPlugin(AbstractNetworkAgentPlugin):
+    async def join_network(
+        self,
+        kernel_config: KernelCreationConfig,
+        cluster_info: ClusterInfo,
+        *,
+        network_name: str,
+        **options,
+    ) -> dict[str, Any]:
+        return {
+            "HostConfig": {
+                "NetworkMode": network_name,
+            },
+            "NetworkingConfig": {
+                "EndpointsConfig": {
+                    network_name: {
+                        "Aliases": [kernel_config["cluster_hostname"]],
+                    },
+                },
+            },
+        }
+
+    async def leave_network(self, kernel: DockerKernel) -> None:
+        pass

src/ai/backend/agent/plugin/BUILD

@@ -0,0 +1 @@
+python_sources(name="src")

src/ai/backend/agent/plugin/network.py

@@ -0,0 +1,34 @@
+from abc import ABCMeta, abstractmethod
+from typing import Any
+
+from ai.backend.agent.kernel import AbstractKernel
+from ai.backend.common.plugin import AbstractPlugin, BasePluginContext
+from ai.backend.common.types import ClusterInfo, KernelCreationConfig
+
+
+class AbstractNetworkAgentPlugin(AbstractPlugin, metaclass=ABCMeta):
+    @abstractmethod
+    async def join_network(
+        self,
+        kernel_config: KernelCreationConfig,
+        cluster_info: ClusterInfo,
+        **kwargs,
+    ) -> dict[str, Any]:
+        """
+        Returns required container config to attach container to network.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def leave_network(
+        self,
+        kernel: AbstractKernel,
+    ) -> None:
+        """
+        Performs extra step to make container leave from the network.
+        """
+        raise NotImplementedError
+
+
+class NetworkPluginContext(BasePluginContext[AbstractNetworkAgentPlugin]):
+    plugin_group = "backendai_network_client_v1"