refactor(agent): Reuse the shared aiodocker client in DockerKernel methods

Follow-up to #11226: route DockerKernel's get_logs/commit/download_file/
download_single through the agent-owned Docker client instead of opening
a throwaway connection per call. commit() previously used a bare Docker()
with no closing_async - that leak is also fixed.

Closes #11227
Refs #11216

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: rapsealk

changes/11227.enhance.md

@@ -0,0 +1 @@
+Route DockerKernel's get_logs/commit/download_file/download_single through the agent-owned shared aiodocker client instead of opening a throwaway connection per call.

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

@@ -316,6 +316,7 @@ def __init__(
         agent_sockpath: Path,
         resource_lock: asyncio.Lock,
         network_plugin_ctx: NetworkPluginContext,
+        docker: Docker,
         restarting: bool = False,
         cluster_ssh_port_mapping: ClusterSSHPortMapping | None = None,
         gwbridge_subnet: str | None = None,
@@ -351,6 +352,7 @@ def __init__(
         self.gwbridge_subnet = gwbridge_subnet
 
         self.network_plugin_ctx = network_plugin_ctx
+        self._docker = docker
 
     def _kernel_resource_spec_read(self, filename: Path | str) -> KernelResourceSpec:
         filepath = Path(filename)
@@ -1010,6 +1012,7 @@ def chown_idfile(uid: int | None, gid: int | None) -> None:
             resource_spec=resource_spec,
             environ=environ,
             data={},
+            docker=self._docker,
         )
 
     @property
@@ -1482,47 +1485,49 @@ def __init__(
         )
 
     async def __ainit__(self) -> None:
-        async with closing_async(Docker()) as docker:
-            docker_host = ""
-            match docker.connector:
-                case aiohttp.TCPConnector():
-                    if docker.docker_host is None:
-                        raise InvalidArgumentError("docker_host is not set for TCP connector")
-                    docker_host = docker.docker_host
-                case aiohttp.NamedPipeConnector() | aiohttp.UnixConnector() as connector:
-                    docker_host = connector.path
-                case _:
-                    docker_host = "(unknown)"
-            log.info("accessing the local Docker daemon via {}", docker_host)
-            docker_version = await docker.version()
-            log.info(
-                "running with Docker {0} with API {1}",
-                docker_version["Version"],
-                docker_version["ApiVersion"],
-            )
-            kernel_version = docker_version["KernelVersion"]
-            if "linuxkit" in kernel_version:
-                self.local_config.agent.docker_mode = "linuxkit"
-            else:
-                self.local_config.agent.docker_mode = "native"
-            docker_info = await docker.system.info()
-            docker_info = dict(docker_info)
-            # Assume cgroup v1 if CgroupVersion key is absent
-            if "CgroupVersion" not in docker_info:
-                docker_info["CgroupVersion"] = "1"
-            log.info(
-                "Cgroup Driver: {0}, Cgroup Version: {1}",
-                docker_info["CgroupDriver"],
-                docker_info["CgroupVersion"],
-            )
-            self.docker_info = docker_info
+        # Long-lived shared aiodocker client; must be available before kernel recovery
+        # loads any DockerKernel instances via `super().__ainit__()`.
+        self.docker = Docker()
+        docker = self.docker
+        docker_host = ""
+        match docker.connector:
+            case aiohttp.TCPConnector():
+                if docker.docker_host is None:
+                    raise InvalidArgumentError("docker_host is not set for TCP connector")
+                docker_host = docker.docker_host
+            case aiohttp.NamedPipeConnector() | aiohttp.UnixConnector() as connector:
+                docker_host = connector.path
+            case _:
+                docker_host = "(unknown)"
+        log.info("accessing the local Docker daemon via {}", docker_host)
+        docker_version = await docker.version()
+        log.info(
+            "running with Docker {0} with API {1}",
+            docker_version["Version"],
+            docker_version["ApiVersion"],
+        )
+        kernel_version = docker_version["KernelVersion"]
+        if "linuxkit" in kernel_version:
+            self.local_config.agent.docker_mode = "linuxkit"
+        else:
+            self.local_config.agent.docker_mode = "native"
+        docker_info = await docker.system.info()
+        docker_info = dict(docker_info)
+        # Assume cgroup v1 if CgroupVersion key is absent
+        if "CgroupVersion" not in docker_info:
+            docker_info["CgroupVersion"] = "1"
+        log.info(
+            "Cgroup Driver: {0}, Cgroup Version: {1}",
+            docker_info["CgroupDriver"],
+            docker_info["CgroupVersion"],
+        )
+        self.docker_info = docker_info
         await self._kernel_recovery_adapter.adapt_recovery_data()
         await super().__ainit__()
         try:
-            async with Docker() as docker:
-                gwbridge = await docker.networks.get("docker_gwbridge")
-                gwbridge_info = await gwbridge.show()
-                self.gwbridge_subnet = gwbridge_info["IPAM"]["Config"][0]["Subnet"]
+            gwbridge = await self.docker.networks.get("docker_gwbridge")
+            gwbridge_info = await gwbridge.show()
+            self.gwbridge_subnet = gwbridge_info["IPAM"]["Config"][0]["Subnet"]
         except (DockerError, KeyError, IndexError):
             self.gwbridge_subnet = None
         ipc_base_path = self.local_config.agent.ipc_base_path
@@ -1556,9 +1561,6 @@ async def __ainit__(self) -> None:
         self.monitor_docker_task = asyncio.create_task(self.monitor_docker_events())
         self.docker_ptask_group = aiotools.PersistentTaskGroup()
 
-        # For legacy accelerator plugins
-        self.docker = Docker()
-
         self.network_plugin_ctx = NetworkPluginContext(
             self.etcd, self.local_config.model_dump(by_alias=True)
         )
@@ -1589,7 +1591,11 @@ async def shutdown(self, stop_signal: signal.Signals) -> None:
 
     @override
     async def _load_kernel_registry_from_recovery(self) -> MutableMapping[KernelId, AbstractKernel]:
-        return await self._kernel_recovery.load_kernel_registry()
+        registry = await self._kernel_recovery.load_kernel_registry()
+        for kernel_obj in registry.values():
+            if isinstance(kernel_obj, DockerKernel):
+                kernel_obj.attach_docker(self.docker)
+        return registry
 
     @override
     async def _write_kernel_registry_to_recovery(
@@ -1999,6 +2005,7 @@ async def init_kernel_context(
             self.agent_sockpath,
             self.resource_lock,
             self.network_plugin_ctx,
+            self.docker,
             restarting=restarting,
             cluster_ssh_port_mapping=cluster_ssh_port_mapping,
             gwbridge_subnet=self.gwbridge_subnet,