fix(agent): Ensure shared aiodocker client is closed on shutdown and boot failure

Addresses review feedback on PR #11226:
- Wrap shutdown() so self.docker.close() runs in a finally block,
  preventing an aiohttp.ClientSession leak if super().shutdown() raises.
- Release the shared Docker client if __ainit__ raises after the
  session is constructed.
- Add regression tests asserting the session is closed both on
  successful shutdown and when inner shutdown raises.

Refs #11218
Refs #11226

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

Author: rapsealk

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

@@ -1481,106 +1481,115 @@ def __init__(
 
     async def __ainit__(self) -> None:
         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:
-            gwbridge = await 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
-        (ipc_base_path / "container").mkdir(parents=True, exist_ok=True)
-        self.agent_sockpath = ipc_base_path / "container" / f"agent.{self.id}.sock"
-        # Workaround for Docker Desktop for Mac's UNIX socket mount failure with virtiofs
-        if sys.platform != "darwin":
-            socket_relay_name = f"backendai-socket-relay.{self.id}"
-            socket_relay_container = PersistentServiceContainer(
-                "backendai-socket-relay:latest",
-                {
-                    "Cmd": [
-                        f"UNIX-LISTEN:/ipc/{self.agent_sockpath.name},unlink-early,fork,mode=777",
-                        f"TCP-CONNECT:127.0.0.1:{self.local_config.agent.agent_sock_port}",
-                    ],
-                    "HostConfig": {
-                        "Mounts": [
-                            {
-                                "Type": "bind",
-                                "Source": str(ipc_base_path / "container"),
-                                "Target": "/ipc",
-                            },
+            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:
+                gwbridge = await 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
+            (ipc_base_path / "container").mkdir(parents=True, exist_ok=True)
+            self.agent_sockpath = ipc_base_path / "container" / f"agent.{self.id}.sock"
+            # Workaround for Docker Desktop for Mac's UNIX socket mount failure with virtiofs
+            if sys.platform != "darwin":
+                socket_relay_name = f"backendai-socket-relay.{self.id}"
+                socket_relay_container = PersistentServiceContainer(
+                    "backendai-socket-relay:latest",
+                    {
+                        "Cmd": [
+                            f"UNIX-LISTEN:/ipc/{self.agent_sockpath.name},unlink-early,fork,mode=777",
+                            f"TCP-CONNECT:127.0.0.1:{self.local_config.agent.agent_sock_port}",
                         ],
-                        "NetworkMode": "host",
+                        "HostConfig": {
+                            "Mounts": [
+                                {
+                                    "Type": "bind",
+                                    "Source": str(ipc_base_path / "container"),
+                                    "Target": "/ipc",
+                                },
+                            ],
+                            "NetworkMode": "host",
+                        },
                     },
-                },
-                name=socket_relay_name,
-            )
-            await socket_relay_container.ensure_running_latest()
-        self.agent_sock_task = asyncio.create_task(self.handle_agent_socket())
-        self.monitor_docker_task = asyncio.create_task(self.monitor_docker_events())
-        self.docker_ptask_group = aiotools.PersistentTaskGroup()
+                    name=socket_relay_name,
+                )
+                await socket_relay_container.ensure_running_latest()
+            self.agent_sock_task = asyncio.create_task(self.handle_agent_socket())
+            self.monitor_docker_task = asyncio.create_task(self.monitor_docker_events())
+            self.docker_ptask_group = aiotools.PersistentTaskGroup()
 
-        self.network_plugin_ctx = NetworkPluginContext(
-            self.etcd, self.local_config.model_dump(by_alias=True)
-        )
-        await self.network_plugin_ctx.init(
-            context=self,
-            allowlist=self.local_config.agent.allow_network_plugins,
-            blocklist=self.local_config.agent.block_network_plugins,
-        )
+            self.network_plugin_ctx = NetworkPluginContext(
+                self.etcd, self.local_config.model_dump(by_alias=True)
+            )
+            await self.network_plugin_ctx.init(
+                context=self,
+                allowlist=self.local_config.agent.allow_network_plugins,
+                blocklist=self.local_config.agent.block_network_plugins,
+            )
+        except BaseException:
+            # Release the shared aiodocker client if boot fails after its construction
+            # so the underlying aiohttp.ClientSession does not leak.
+            await self.docker.close()
+            raise
 
     async def shutdown(self, stop_signal: signal.Signals) -> None:
-        # Stop handling agent sock.
-        if self.agent_sock_task is not None:
-            self.agent_sock_task.cancel()
-            await self.agent_sock_task
-        if self.docker_ptask_group is not None:
-            await self.docker_ptask_group.shutdown()
-
         try:
-            await super().shutdown(stop_signal)
-        finally:
-            # Stop docker event monitoring.
-            if self.monitor_docker_task is not None:
-                self.monitor_docker_task.cancel()
-                await self.monitor_docker_task
+            # Stop handling agent sock.
+            if self.agent_sock_task is not None:
+                self.agent_sock_task.cancel()
+                await self.agent_sock_task
+            if self.docker_ptask_group is not None:
+                await self.docker_ptask_group.shutdown()
 
-        if self.docker:
-            await self.docker.close()
+            try:
+                await super().shutdown(stop_signal)
+            finally:
+                # Stop docker event monitoring.
+                if self.monitor_docker_task is not None:
+                    self.monitor_docker_task.cancel()
+                    await self.monitor_docker_task
+        finally:
+            # Always release the shared aiodocker client so its aiohttp.ClientSession
+            # does not leak even when inner shutdown steps raise.
+            if self.docker is not None:
+                await self.docker.close()
 
     @override
     async def _load_kernel_registry_from_recovery(self) -> MutableMapping[KernelId, AbstractKernel]:

tests/component/agent/docker/test_agent.py

@@ -183,3 +183,68 @@ async def test_save_last_registry_exception(agent: DockerAgent, mocker: Any) ->
     )
     await agent.save_last_registry()
     assert not registry_state_path.exists()
+
+
+@pytest.fixture
+async def unmanaged_agent(
+    local_config: Any, test_id: str, mocker: Any, socket_relay_image: Any
+) -> Any:
+    """
+    Like the ``agent`` fixture, but leaves shutdown entirely to the test so it
+    can assert against the lifecycle of the shared aiodocker client.
+    """
+    dummy_etcd = DummyEtcd()
+    mocked_etcd_get_prefix = AsyncMock(return_value={})
+    mocker.patch.object(dummy_etcd, "get_prefix", new=mocked_etcd_get_prefix)
+    test_case_id = secrets.token_hex(8)
+    kernel_registry = KernelRegistry()
+    agent = await DockerAgent.new(
+        dummy_etcd,
+        local_config,
+        stats_monitor=None,
+        error_monitor=None,
+        skip_initial_scan=True,
+        agent_public_key=None,
+        kernel_registry=kernel_registry,
+        computers={},
+        slots={},
+        agent_class=AgentClass.PRIMARY,
+    )
+    agent.local_instance_id = test_case_id
+    yield agent
+    # Best-effort cleanup: close the shared client if the test did not already
+    # trigger a full shutdown.  ``Docker.close`` -> ``ClientSession.close`` is
+    # safe to call on an already-closed session.
+    if not agent.docker.session.closed:
+        await agent.docker.close()
+
+
+async def test_shared_docker_client_open_after_ainit(unmanaged_agent: DockerAgent) -> None:
+    assert unmanaged_agent.docker.session.closed is False
+
+
+async def test_shared_docker_client_closed_after_shutdown(
+    unmanaged_agent: DockerAgent,
+) -> None:
+    assert unmanaged_agent.docker.session.closed is False
+    await unmanaged_agent.shutdown(signal.SIGTERM)
+    assert unmanaged_agent.docker.session.closed is True
+
+
+async def test_shared_docker_client_closed_when_super_shutdown_raises(
+    unmanaged_agent: DockerAgent, mocker: Any
+) -> None:
+    # Simulate the base Agent.shutdown raising mid-shutdown — the shared
+    # aiodocker client must still be closed so its aiohttp.ClientSession does
+    # not leak.
+    class _SimulatedShutdownError(Exception):
+        pass
+
+    mocker.patch(
+        "ai.backend.agent.agent.AbstractAgent.shutdown",
+        new=AsyncMock(side_effect=_SimulatedShutdownError("simulated")),
+    )
+    assert unmanaged_agent.docker.session.closed is False
+    with pytest.raises(_SimulatedShutdownError):
+        await unmanaged_agent.shutdown(signal.SIGTERM)
+    assert unmanaged_agent.docker.session.closed is True