test(agent): Lock DockerKernel pickling invariant and harden attach_docker

Addresses review feedback on PR #11228:
- Add a pickle round-trip test asserting __getstate__ excludes _docker,
  locking the RPC-marshalling invariant for future contributors.
- Document on DockerKernel that _docker is local-only and must be
  re-attached by the agent via attach_docker().
- Warn (but still overwrite) when attach_docker is called with a
  connector different from the currently-attached one; same-connector
  reattach remains silent.
- Align the self.docker type annotation on DockerAgent with the one
  introduced in #11226 (no-op if already present).

Refs #11227
Refs #11228

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

Author: rapsealk

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

@@ -1418,6 +1418,7 @@ async def _rollback_container_creation() -> None:
 
 
 class DockerAgent(AbstractAgent[DockerKernel, DockerKernelCreationContext]):
+    docker: Docker
     docker_info: Mapping[str, Any]
     monitor_docker_task: asyncio.Task[Any]
     agent_sockpath: Path

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

@@ -45,6 +45,15 @@
 
 
 class DockerKernel(AbstractKernel):
+    """
+    Docker-backed kernel implementation.
+
+    Note: ``_docker`` is a local, process-bound reference to the agent-owned
+    ``aiodocker.Docker`` client and MUST NOT be serialized (aiohttp sessions
+    are not picklable). It is excluded from ``__getstate__`` / ``__setstate__``
+    and must be re-attached by the agent on recovery via ``attach_docker()``.
+    """
+
     network_driver: str
     _docker: Docker
 
@@ -80,6 +89,12 @@ def __init__(
 
     def attach_docker(self, docker: Docker) -> None:
         """Re-attach the agent-owned aiodocker client after unpickling or recovery."""
+        existing = getattr(self, "_docker", None)
+        if existing is not None and existing is not docker:
+            log.warning(
+                "DockerKernel {} re-attached with a different Docker client; this is likely a bug",
+                self.kernel_id,
+            )
         self._docker = docker
 
     @override

tests/unit/agent/test_container_id_sync.py

@@ -7,15 +7,17 @@
 
 from __future__ import annotations
 
+import pickle
 from unittest.mock import Mock
 from uuid import uuid4
 
 import pytest
 
 from ai.backend.agent.docker.kernel import DockerKernel
+from ai.backend.agent.resources import KernelResourceSpec
 from ai.backend.agent.types import KernelOwnershipData
 from ai.backend.common.docker import ImageRef
-from ai.backend.common.types import AgentId, ContainerId, KernelId, SessionId
+from ai.backend.common.types import AgentId, ContainerId, KernelId, ResourceSlot, SessionId
 
 
 @pytest.fixture
@@ -80,3 +82,89 @@ def test_second_call_updates_both_sources(self, mock_kernel_obj: DockerKernel) -
         mock_kernel_obj.set_container_id(second)
         assert mock_kernel_obj.container_id == second
         assert mock_kernel_obj["container_id"] == second
+
+
+@pytest.fixture
+def picklable_kernel_obj() -> DockerKernel:
+    """
+    Create a DockerKernel whose state (apart from the excluded ``_docker``
+    reference) is fully picklable.
+
+    The shared ``mock_kernel_obj`` fixture uses ``Mock()`` for
+    ``resource_spec``, which cannot be pickled. Pickling tests therefore
+    need a kernel constructed with a real, lightweight ``KernelResourceSpec``
+    so the only non-picklable member is the intentionally-excluded
+    ``_docker`` client.
+    """
+    kernel_id = KernelId(uuid4())
+    session_id = SessionId(uuid4())
+    agent_id = AgentId("test-agent-id")
+    ownership_data = KernelOwnershipData(
+        kernel_id=kernel_id,
+        session_id=session_id,
+        agent_id=agent_id,
+    )
+    image = ImageRef(
+        name="test-image",
+        project="test-project",
+        registry="registry.local",
+        tag="latest",
+        architecture="x86_64",
+        is_local=False,
+    )
+    resource_spec = KernelResourceSpec(
+        slots=ResourceSlot(),
+        allocations={},
+        scratch_disk_size=0,
+    )
+    return DockerKernel(
+        ownership_data=ownership_data,
+        network_id="test-network",
+        image=image,
+        version=1,
+        network_driver="bridge",
+        agent_config={},
+        resource_spec=resource_spec,
+        service_ports=[],
+        environ={},
+        data={},
+        docker=Mock(),
+    )
+
+
+class TestDockerKernelPickling:
+    """
+    Tests locking the DockerKernel pickling invariant.
+
+    The agent marshals DockerKernel instances via pickle for RPC / recovery.
+    The live ``_docker`` client (aiohttp-backed aiodocker.Docker) is NOT
+    picklable, so ``__getstate__`` must drop it and ``__setstate__`` must
+    tolerate its absence. The agent re-attaches the client after unpickling
+    via ``attach_docker()``. These tests guard against future regressions
+    that would accidentally put ``_docker`` back into the pickled payload.
+    """
+
+    def test_pickle_round_trip_excludes_docker(self, picklable_kernel_obj: DockerKernel) -> None:
+        """pickle.dumps/loads must succeed and drop the _docker reference."""
+        # The Mock() Docker client is not picklable; the test confirms that
+        # __getstate__ drops it so pickling succeeds anyway, and that non-
+        # excluded state is preserved across the round trip.
+        cid = ContainerId("pickled-container-xyz")
+        picklable_kernel_obj.set_container_id(cid)
+
+        data = pickle.dumps(picklable_kernel_obj)
+        restored = pickle.loads(data)
+
+        assert isinstance(restored, DockerKernel)
+        # _docker must not be carried across pickle; the agent re-attaches
+        # it via attach_docker() on recovery.
+        assert getattr(restored, "_docker", None) is None
+        # Sanity check: non-excluded state survives the round trip.
+        assert restored.container_id == cid
+        assert restored["container_id"] == cid
+        assert restored.network_driver == picklable_kernel_obj.network_driver
+
+    def test_getstate_does_not_contain_docker_key(self, picklable_kernel_obj: DockerKernel) -> None:
+        """__getstate__ output must not carry the live _docker reference."""
+        state = picklable_kernel_obj.__getstate__()
+        assert "_docker" not in state