diff --git a/clients/python/agentic-sandbox-client/k8s_agent_sandbox/constants.py b/clients/python/agentic-sandbox-client/k8s_agent_sandbox/constants.py index 6c52ac695..d6e13dab1 100644 --- 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 new file mode 100644 index 000000000..ac0c5f9ad --- /dev/null +++ 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 new file mode 100644 index 000000000..909085497 --- /dev/null +++ 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. \ No newline at end of file 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 new file mode 100644 index 000000000..59725b69c --- /dev/null +++ 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 new file mode 100644 index 000000000..1ae581867 --- /dev/null +++ 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 new file mode 100644 index 000000000..2185acbc7 --- /dev/null +++ b/clients/python/agentic-sandbox-client/python-counter-template.yaml @@ -0,0 +1,29 @@ +apiVersion: extensions.agents.x-k8s.io/v1alpha1 +kind: SandboxTemplate +metadata: + name: python-counter-template + namespace: sandbox-test + labels: + language: python +spec: + #enableDisruptionControl: true + #shutdownTime: "2025-12-31T23:59:59Z" + podTemplate: + metadata: + labels: + app: agent-sandbox-workload + spec: + serviceAccountName: sandbox-test + runtimeClassName: gvisor + containers: + - name: my-container1 + image: python:3.10-slim + command: ["python3", "-c"] + args: + - | + import time + i = 0 + while True: + print(f"Count: {i}", flush=True) + i += 1 + time.sleep(1) \ No newline at end of file diff --git a/clients/python/agentic-sandbox-client/test_podsnapshot_extension.py b/clients/python/agentic-sandbox-client/test_podsnapshot_extension.py new file mode 100644 index 000000000..383e12d33 --- /dev/null +++ b/clients/python/agentic-sandbox-client/test_podsnapshot_extension.py @@ -0,0 +1,92 @@ +# 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 argparse +from kubernetes import config +from k8s_agent_sandbox.gke_extensions import PodSnapshotSandboxClient + + +def main( + template_name: str, + api_url: str | None, + namespace: str, + server_port: int, +): + """ + Tests the Sandbox client by creating a sandbox, running a command, + and then cleaning up. + """ + + print( + f"--- Starting Sandbox Client Test (Namespace: {namespace}, Port: {server_port}) ---" + ) + + # Load kube config + try: + config.load_incluster_config() + except config.ConfigException: + config.load_kube_config() + + try: + print("\n***** Phase 1: Starting Counter *****") + + with PodSnapshotSandboxClient( + template_name=template_name, + namespace=namespace, + api_url=api_url, + server_port=server_port, + ) as sandbox: + print("\n======= Testing Pod Snapshot Extension =======") + assert ( + sandbox.snapshot_crd_installed == True + ), "Pod Snapshot CRD is not installed." + + except Exception as e: + print(f"\n--- An error occurred during the test: {e} ---") + # The __exit__ method of the Sandbox class will handle cleanup. + finally: + print("\n--- Sandbox Client Test Finished ---") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Test the Sandbox client.") + parser.add_argument( + "--template-name", + default="python-sandbox-template", + help="The name of the sandbox template to use for the test.", + ) + + parser.add_argument( + "--api-url", + help="Direct URL to router (e.g. http://localhost:8080)", + default=None, + ) + parser.add_argument( + "--namespace", default="default", help="Namespace to create sandbox in" + ) + parser.add_argument( + "--server-port", + type=int, + default=8888, + help="Port the sandbox container listens on", + ) + + args = parser.parse_args() + + main( + template_name=args.template_name, + api_url=args.api_url, + namespace=args.namespace, + server_port=args.server_port, + )