feat: add Pod Snapshot extension for Python SDK (#338)

shrutiyam-glitch · web-flow · commit 15026b786c59 · 2026-03-06T03:30:19.000+05:30
* Add constants

* PodSanpshotSandboxClient Class

* Fix minor nits

* Address nit

* Update comment

* Use label-selector to check if pod snapshot agent is running

* Updated snapshot client check

* Update variable names

* address comments
diff --git a/clients/python/agentic-sandbox-client/k8s_agent_sandbox/constants.py b/clients/python/agentic-sandbox-client/k8s_agent_sandbox/constants.py
@@ -26,3 +26,7 @@
 SANDBOX_PLURAL_NAME = "sandboxes"
 
 POD_NAME_ANNOTATION = "agents.x-k8s.io/pod-name"
+
+PODSNAPSHOT_API_GROUP = "podsnapshot.gke.io"
+PODSNAPSHOT_API_VERSION = "v1alpha1"
+PODSNAPSHOT_API_KIND = "PodSnapshot"
diff --git a/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/__init__.py b/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/__init__.py
@@ -0,0 +1,15 @@
+# Copyright 2026 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .podsnapshot_client import PodSnapshotSandboxClient
diff --git a/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot.md b/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot.md
@@ -0,0 +1,62 @@
+# Agentic Sandbox Pod Snapshot Extension
+
+This directory contains the Python client extension for interacting with the Agentic Sandbox to manage Pod Snapshots. This extension allows you to trigger snapshots of a running sandbox and restore a new sandbox from the recently created snapshot.
+
+## `podsnapshot_client.py`
+
+This file defines the `PodSnapshotSandboxClient` class, which extend the base `SandboxClient` to provide snapshot capabilities.
+
+### `PodSnapshotSandboxClient`
+
+A specialized Sandbox client for interacting with the gke pod snapshot controller.
+
+### Key Features:
+
+*   **`PodSnapshotSandboxClient(template_name: str, ...)`**:
+    *   Initializes the client with optional server port.
+
+*   **`__exit__(self)`**:
+    *   Cleans up the `SandboxClaim` resources.
+
+## `test_podsnapshot_extension.py`
+
+This file, located in the parent directory (`clients/python/agentic-sandbox-client/`), contains an integration test script for the `PodSnapshotSandboxClient` extension. It verifies the snapshot and restore functionality.
+
+### Test Phases:
+
+1.  **Phase 1: Starting Counter Sandbox**:
+    *   Starts a sandbox with a counter application.
+
+### Prerequisites
+
+1.  **Python Virtual Environment**:
+    ```bash
+    python3 -m venv .venv
+    source .venv/bin/activate
+    ```
+
+2.  **Install Dependencies**:
+    ```bash
+    pip install kubernetes
+    pip install -e clients/python/agentic-sandbox-client/
+    ```
+
+3.  **Pod Snapshot Controller**: The Pod Snapshot controller must be installed in a **GKE standard cluster** running with **gVisor**. 
+   * For detailed setup instructions, refer to the [GKE Pod Snapshots public documentation](https://docs.cloud.google.com/kubernetes-engine/docs/how-to/pod-snapshots).
+   * Ensure a GCS bucket is configured to store the pod snapshot states and that the necessary IAM permissions are applied.
+
+4.  **CRDs**: `PodSnapshotStorageConfig`, `PodSnapshotPolicy` CRDs must be applied. `PodSnapshotPolicy` should specify the selector match labels.
+
+5.  **Sandbox Template**: A `SandboxTemplate` (e.g., `python-counter-template`) with runtime gVisor, appropriate KSA and label that matches that selector label in `PodSnapshotPolicy` must be available in the cluster.
+
+### Running Tests:
+
+To run the integration test, execute the script with the appropriate arguments:
+
+```bash
+python3 clients/python/agentic-sandbox-client/test_podsnapshot_extension.py \
+  --template-name python-counter-template \
+  --namespace sandbox-test
+```
+
+Adjust the `--namespace`, `--template-name` as needed for your environment.
diff --git a/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot_client.py b/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/podsnapshot_client.py
@@ -0,0 +1,94 @@
+# Copyright 2026 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from kubernetes import client
+from kubernetes.client import ApiException
+from ..sandbox_client import SandboxClient
+from ..constants import (
+    PODSNAPSHOT_API_GROUP,
+    PODSNAPSHOT_API_VERSION,
+    PODSNAPSHOT_API_KIND,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class PodSnapshotSandboxClient(SandboxClient):
+    """
+    A specialized Sandbox client for interacting with the GKE Pod Snapshot Controller.
+
+    TODO: This class enables users to take a snapshot of their sandbox and restore from the taken snapshot.
+    """
+
+    def __init__(
+        self,
+        template_name: str,
+        **kwargs,
+    ):
+        super().__init__(template_name, **kwargs)
+
+        self.snapshot_crd_installed = False
+        self.core_v1_api = client.CoreV1Api()
+
+    def __enter__(self) -> "PodSnapshotSandboxClient":
+        try:
+            self.snapshot_crd_installed = self._check_snapshot_crd_installed()
+            if not self.snapshot_crd_installed:
+                raise RuntimeError(
+                    "Pod Snapshot Controller is not ready. "
+                    "Ensure the PodSnapshot CRD is installed."
+                )
+            super().__enter__()
+            return self
+        except Exception as e:
+            self.__exit__(None, None, None)
+            raise RuntimeError(
+                f"Failed to initialize PodSnapshotSandboxClient. Ensure that you are connected to a GKE cluster "
+                f"with the Pod Snapshot Controller enabled. Error details: {e}"
+            ) from e
+
+    def _check_snapshot_crd_installed(self) -> bool:
+        """
+        Checks if the PodSnapshot CRD is installed in the cluster.
+        """
+
+        if self.snapshot_crd_installed:
+            return True
+
+        try:
+            # Check if the API resource exists using CustomObjectsApi
+            resource_list = self.custom_objects_api.get_api_resources(
+                group=PODSNAPSHOT_API_GROUP,
+                version=PODSNAPSHOT_API_VERSION,
+            )
+
+            if not resource_list or not resource_list.resources:
+                return False
+
+            for resource in resource_list.resources:
+                if resource.kind == PODSNAPSHOT_API_KIND:
+                    return True
+            return False
+        except ApiException as e:
+            # If discovery fails with 403/404, we assume not ready/accessible
+            if e.status == 403 or e.status == 404:
+                return False
+            raise
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Automatically cleans up the Sandbox.
+        """
+        super().__exit__(exc_type, exc_val, exc_tb)
diff --git a/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/test_podsnapshot_client.py b/clients/python/agentic-sandbox-client/k8s_agent_sandbox/gke_extensions/test_podsnapshot_client.py
@@ -0,0 +1,198 @@
+# Copyright 2026 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+import os
+import logging
+from unittest.mock import MagicMock, patch
+from k8s_agent_sandbox.gke_extensions.podsnapshot_client import (
+    PodSnapshotSandboxClient,
+)
+from k8s_agent_sandbox.constants import (
+    PODSNAPSHOT_API_KIND,
+    PODSNAPSHOT_API_GROUP,
+    PODSNAPSHOT_API_VERSION,
+)
+
+from kubernetes.client import ApiException
+
+from kubernetes import config
+
+logger = logging.getLogger(__name__)
+
+
+class TestPodSnapshotSandboxClient(unittest.TestCase):
+
+    @patch("kubernetes.config")
+    def setUp(self, mock_config):
+        logger.info("Setting up TestPodSnapshotSandboxClient...")
+        # Mock kubernetes config loading
+        mock_config.load_incluster_config.side_effect = config.ConfigException(
+            "Not in cluster"
+        )
+        mock_config.load_kube_config.return_value = None
+
+        # Create client without patching super, as it's tested separately
+        with patch.object(
+            PodSnapshotSandboxClient, "_check_snapshot_crd_installed", return_value=True
+        ):
+            self.client = PodSnapshotSandboxClient("test-template")
+
+        # Mock the kubernetes APIs on the client instance
+        self.client.custom_objects_api = MagicMock()
+        self.client.core_v1_api = MagicMock()
+
+        logger.info("Finished setting up TestPodSnapshotSandboxClient.")
+
+    def test_init(self):
+        """Test initialization of PodSnapshotSandboxClient."""
+        logger.info("Starting test_init...")
+        with patch(
+            "k8s_agent_sandbox.sandbox_client.SandboxClient.__init__", return_value=None
+        ) as mock_super:
+            with patch.object(
+                PodSnapshotSandboxClient,
+                "_check_snapshot_crd_installed",
+                return_value=True,
+            ):
+                client = PodSnapshotSandboxClient("test-template")
+            mock_super.assert_called_once_with("test-template")
+        self.assertFalse(client.snapshot_crd_installed)
+        logger.info("Finished test_init.")
+
+    def test_check_snapshot_crd_installed_success(self):
+        """Test _check_snapshot_crd_installed success scenarios (Check CRD Existence)."""
+        logger.info("TEST: CRD Existence Success")
+        mock_resource_list = MagicMock()
+        mock_resource = MagicMock()
+        mock_resource.kind = PODSNAPSHOT_API_KIND
+        mock_resource_list.resources = [mock_resource]
+        self.client.custom_objects_api.get_api_resources.return_value = (
+            mock_resource_list
+        )
+
+        self.client.snapshot_crd_installed = False
+        self.assertTrue(self.client._check_snapshot_crd_installed())
+        self.client.custom_objects_api.get_api_resources.assert_called_with(
+            group=PODSNAPSHOT_API_GROUP, version=PODSNAPSHOT_API_VERSION
+        )
+
+    def test_check_snapshot_crd_installed_failures(self):
+        """Test _check_snapshot_crd_installed failure scenarios."""
+
+        # 1. No CRDs found
+        self.client.custom_objects_api.get_api_resources.return_value = None
+        self.client.snapshot_crd_installed = False
+        self.assertFalse(self.client._check_snapshot_crd_installed())
+
+        # 2. CRD Kind mismatch
+        mock_resource_list = MagicMock()
+        mock_resource = MagicMock()
+        mock_resource.kind = "SomeOtherKind"
+        mock_resource_list.resources = [mock_resource]
+        self.client.custom_objects_api.get_api_resources.return_value = (
+            mock_resource_list
+        )
+        self.client.snapshot_crd_installed = False
+        self.assertFalse(self.client._check_snapshot_crd_installed())
+
+        # 3. 404 on CRD check
+        self.client.custom_objects_api.get_api_resources.side_effect = ApiException(
+            status=404
+        )
+        self.client.snapshot_crd_installed = False
+        self.assertFalse(self.client._check_snapshot_crd_installed())
+
+        # 4. 403 on CRD check
+        self.client.custom_objects_api.get_api_resources.side_effect = ApiException(
+            status=403
+        )
+        self.client.snapshot_crd_installed = False
+        self.assertFalse(self.client._check_snapshot_crd_installed())
+
+    def test_check_snapshot_crd_installed_exceptions(self):
+        """Test API exceptions during snapshot readiness checks."""
+
+        # 1. 500 on CRD Check
+        self.client.custom_objects_api.get_api_resources.side_effect = ApiException(
+            status=500
+        )
+        self.client.snapshot_crd_installed = False
+        with self.assertRaises(ApiException):
+            self.client._check_snapshot_crd_installed()
+
+    def test_enter_exit(self):
+        """Test context manager __enter__ implementation."""
+        # Success path
+        self.client.snapshot_crd_installed = False
+        with patch.object(
+            self.client, "_check_snapshot_crd_installed", return_value=True
+        ) as mock_ready:
+            with patch(
+                "k8s_agent_sandbox.sandbox_client.SandboxClient.__enter__"
+            ) as mock_super_enter:
+                result = self.client.__enter__()
+                self.assertEqual(result, self.client)
+                mock_ready.assert_called_once()
+                mock_super_enter.assert_called_once()
+                self.assertTrue(self.client.snapshot_crd_installed)
+
+        # Failure path: Controller not ready (return False)
+        self.client.snapshot_crd_installed = False
+        with patch.object(
+            self.client, "_check_snapshot_crd_installed", return_value=False
+        ) as mock_ready:
+            with patch.object(self.client, "__exit__") as mock_exit:
+                with self.assertRaises(RuntimeError) as context:
+                    self.client.__enter__()
+                self.assertIn(
+                    "Pod Snapshot Controller is not ready",
+                    str(context.exception),
+                )
+                mock_exit.assert_called_once_with(None, None, None)
+
+        # Failure path: Exception during check
+        self.client.snapshot_crd_installed = False
+        with patch.object(
+            self.client,
+            "_check_snapshot_crd_installed",
+            side_effect=ValueError("Test error"),
+        ) as mock_ready:
+            with patch.object(self.client, "__exit__") as mock_exit:
+                with self.assertRaises(RuntimeError) as context:
+                    self.client.__enter__()
+                self.assertIn(
+                    "Failed to initialize PodSnapshotSandboxClient",
+                    str(context.exception),
+                )
+                mock_exit.assert_called_once_with(None, None, None)
+
+        # Test Exit
+        with patch(
+            "k8s_agent_sandbox.sandbox_client.SandboxClient.__exit__"
+        ) as mock_super_exit:
+            exc_val = ValueError("test")
+            self.client.__exit__(ValueError, exc_val, None)
+            mock_super_exit.assert_called_once_with(ValueError, exc_val, None)
+
+    def test_check_snapshot_crd_installed_already_ready(self):
+        """Test early return if snapshot controller is already ready."""
+        self.client.snapshot_crd_installed = True
+        result = self.client._check_snapshot_crd_installed()
+        self.assertTrue(result)
+        self.client.custom_objects_api.get_api_resources.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/clients/python/agentic-sandbox-client/python-counter-template.yaml b/clients/python/agentic-sandbox-client/python-counter-template.yaml
diff --git a/clients/python/agentic-sandbox-client/test_podsnapshot_extension.py b/clients/python/agentic-sandbox-client/test_podsnapshot_extension.py
