feat(server): refactor k8s client and add limiter (alibaba#429)

Author: Generalwin

server/src/config.py

@@ -216,6 +216,38 @@ class KubernetesRuntimeConfig(BaseModel):
             "[Beta] Watch timeout (seconds) before restarting the informer stream."
         ),
     )
+    read_qps: float = Field(
+        default=0.0,
+        ge=0,
+        description=(
+            "Maximum read requests per second to the Kubernetes API (get/list). "
+            "0 means unlimited (no rate limiting)."
+        ),
+    )
+    read_burst: int = Field(
+        default=0,
+        ge=0,
+        description=(
+            "Burst size for the read rate limiter. "
+            "0 means use read_qps as burst (minimum 1)."
+        ),
+    )
+    write_qps: float = Field(
+        default=0.0,
+        ge=0,
+        description=(
+            "Maximum write requests per second to the Kubernetes API (create/delete/patch). "
+            "0 means unlimited (no rate limiting)."
+        ),
+    )
+    write_burst: int = Field(
+        default=0,
+        ge=0,
+        description=(
+            "Burst size for the write rate limiter. "
+            "0 means use write_qps as burst (minimum 1)."
+        ),
+    )
     namespace: Optional[str] = Field(
         default=None,
         description="Namespace used for sandbox workloads.",

server/src/services/k8s/agent_sandbox_provider.py

@@ -20,18 +20,16 @@
 import logging
 import re
 from datetime import datetime
-from typing import Dict, List, Any, Optional, Callable
-from threading import Lock
+from typing import Dict, List, Any, Optional
 
 from kubernetes.client import (
     V1Container,
     V1EnvVar,
     V1ResourceRequirements,
     V1VolumeMount,
-    ApiException,
 )
 
-from src.config import AppConfig, IngressConfig, ExecdInitResources
+from src.config import AppConfig
 from src.services.helpers import format_ingress_endpoint
 from src.api.schema import Endpoint, ImageSpec, NetworkPolicy, Volume
 from src.services.k8s.agent_sandbox_template import AgentSandboxTemplateManager
@@ -42,7 +40,6 @@
     build_security_context_from_dict,
     serialize_security_context_to_dict,
 )
-from src.services.k8s.informer import WorkloadInformer
 from src.services.k8s.volume_helper import apply_volumes_to_pod_spec
 from src.services.k8s.workload_provider import WorkloadProvider
 from src.services.runtime_resolver import SecureRuntimeResolver
@@ -83,44 +80,24 @@ class AgentSandboxProvider(WorkloadProvider):
     def __init__(
         self,
         k8s_client: K8sClient,
-        template_file_path: Optional[str] = None,
-        shutdown_policy: str = "Delete",
-        service_account: Optional[str] = None,
-        ingress_config: Optional[IngressConfig] = None,
-        enable_informer: bool = True,
-        informer_factory: Optional[Callable[[str], WorkloadInformer]] = None,
-        informer_resync_seconds: int = 300,
-        informer_watch_timeout_seconds: int = 60,
         app_config: Optional[AppConfig] = None,
-        execd_init_resources: Optional[ExecdInitResources] = None,
     ):
         self.k8s_client = k8s_client
-        self.custom_api = k8s_client.get_custom_objects_api()
-        self.core_api = k8s_client.get_core_v1_api()
 
         self.group = "agents.x-k8s.io"
         self.version = "v1alpha1"
         self.plural = "sandboxes"
 
-        self.shutdown_policy = shutdown_policy
-        self.service_account = service_account
-        self.template_manager = AgentSandboxTemplateManager(template_file_path)
-        self.ingress_config = ingress_config
-        self.execd_init_resources = execd_init_resources
-        self._enable_informer = enable_informer
-        self._informer_factory = informer_factory or (
-            lambda ns: WorkloadInformer(
-                custom_api=self.custom_api,
-                group=self.group,
-                version=self.version,
-                plural=self.plural,
-                namespace=ns,
-                resync_period_seconds=informer_resync_seconds,
-                watch_timeout_seconds=informer_watch_timeout_seconds,
-            )
+        k8s_config = app_config.kubernetes if app_config else None
+        agent_config = app_config.agent_sandbox if app_config else None
+
+        self.shutdown_policy = agent_config.shutdown_policy if agent_config else "Delete"
+        self.service_account = k8s_config.service_account if k8s_config else None
+        self.template_manager = AgentSandboxTemplateManager(
+            agent_config.template_file if agent_config else None
         )
-        self._informers: Dict[str, WorkloadInformer] = {}
-        self._informers_lock = Lock()
+        self.ingress_config = app_config.ingress if app_config else None
+        self.execd_init_resources = k8s_config.execd_init_resources if k8s_config else None
 
         # Initialize secure runtime resolver
         self.resolver = SecureRuntimeResolver(app_config) if app_config else None
@@ -158,14 +135,7 @@ def create_workload(
         egress_image: Optional[str] = None,
         volumes: Optional[List[Volume]] = None,
     ) -> Dict[str, Any]:
-        # Pool mode does not support volumes
-        if extensions and extensions.get("poolRef"):
-            if volumes:
-                raise ValueError(
-                    "Pool mode does not support volumes. "
-                    "Remove 'volumes' from request or use template mode."
-                )
-
+        """Create an agent-sandbox Sandbox CRD workload."""
         if self.runtime_class:
             logger.info(
                 "Using Kubernetes RuntimeClass '%s' for sandbox %s",
@@ -215,21 +185,14 @@ def create_workload(
 
         sandbox = self.template_manager.merge_with_runtime_values(runtime_manifest)
 
-        created = self.custom_api.create_namespaced_custom_object(
+        created = self.k8s_client.create_custom_object(
             group=self.group,
             version=self.version,
             namespace=namespace,
             plural=self.plural,
             body=sandbox,
         )
 
-        informer = self._get_informer(namespace)
-        if informer:
-            try:
-                informer.update_cache(created)
-            except Exception as exc:  # pragma: no cover - defensive
-                logger.warning("Failed to update informer cache for %s: %s", sandbox_id, exc)
-
         return {
             "name": created["metadata"]["name"],
             "uid": created["metadata"]["uid"],
@@ -245,6 +208,7 @@ def _build_pod_spec(
         network_policy: Optional[NetworkPolicy] = None,
         egress_image: Optional[str] = None,
     ) -> Dict[str, Any]:
+        """Build pod spec dict for the Sandbox CRD."""
         init_container = self._build_execd_init_container(execd_image)
         main_container = self._build_main_container(
             image_spec=image_spec,
@@ -284,6 +248,7 @@ def _build_pod_spec(
         return pod_spec
 
     def _build_execd_init_container(self, execd_image: str) -> V1Container:
+        """Build init container that copies execd binary to the shared volume."""
         script = (
             "cp ./execd /opt/opensandbox/bin/execd && "
             "cp ./bootstrap.sh /opt/opensandbox/bin/bootstrap.sh && "
@@ -359,6 +324,7 @@ def _build_main_container(
         )
 
     def _container_to_dict(self, container: V1Container) -> Dict[str, Any]:
+        """Convert a V1Container object to a plain dict for CRD body."""
         result: Dict[str, Any] = {
             "name": container.name,
             "image": container.image,
@@ -388,70 +354,30 @@ def _container_to_dict(self, container: V1Container) -> Dict[str, Any]:
 
         return result
 
-    def _get_informer(self, namespace: str) -> Optional[WorkloadInformer]:
-        if not self._enable_informer:
-            return None
-
-        with self._informers_lock:
-            informer = self._informers.get(namespace)
-            if informer is None:
-                informer = self._informer_factory(namespace)
-                self._informers[namespace] = informer
-                try:
-                    informer.start()
-                except Exception as exc:  # pragma: no cover - defensive
-                    logger.warning(
-                        "Failed to start informer for namespace %s: %s", namespace, exc
-                    )
-                    self._informers.pop(namespace, None)
-                    return None
-        return informer
-
     def get_workload(self, sandbox_id: str, namespace: str) -> Optional[Dict[str, Any]]:
-        informer = self._get_informer(namespace)
-        cache_ready = informer.has_synced if informer else False
-
+        """Get Sandbox CRD by sandbox ID, trying all candidate resource names."""
         candidates = self._resource_name_candidates(sandbox_id)
 
-        if informer and cache_ready:
-            for name in candidates:
-                cached = informer.get(name)
-                if cached:
-                    return cached
-
-        if informer and not cache_ready:
-            logger.warning(
-                f"Informer cache not synced for namespace {namespace}; falling back to direct API get."
-            )
-
         for name in candidates:
-            try:
-                workload = self.custom_api.get_namespaced_custom_object(
-                    group=self.group,
-                    version=self.version,
-                    namespace=namespace,
-                    plural=self.plural,
-                    name=name,
-                )
-                if informer and workload:
-                    informer.update_cache(workload)
+            workload = self.k8s_client.get_custom_object(
+                group=self.group,
+                version=self.version,
+                namespace=namespace,
+                plural=self.plural,
+                name=name,
+            )
+            if workload:
                 return workload
-            except ApiException as e:
-                if e.status != 404:
-                    logger.error(f"Unexpected error getting Sandbox for {sandbox_id}: {e}")
-                    raise
-            except Exception as e:
-                logger.error(f"Unexpected error getting Sandbox for {sandbox_id}: {e}")
-                raise
 
         return None
 
     def delete_workload(self, sandbox_id: str, namespace: str) -> None:
+        """Delete the Sandbox CRD for the given sandbox ID."""
         sandbox = self.get_workload(sandbox_id, namespace)
         if not sandbox:
             raise Exception(f"Sandbox for sandbox {sandbox_id} not found")
 
-        self.custom_api.delete_namespaced_custom_object(
+        self.k8s_client.delete_custom_object(
             group=self.group,
             version=self.version,
             namespace=namespace,
@@ -461,24 +387,17 @@ def delete_workload(self, sandbox_id: str, namespace: str) -> None:
         )
 
     def list_workloads(self, namespace: str, label_selector: str) -> List[Dict[str, Any]]:
-        try:
-            sandbox_list = self.custom_api.list_namespaced_custom_object(
-                group=self.group,
-                version=self.version,
-                namespace=namespace,
-                plural=self.plural,
-                label_selector=label_selector,
-            )
-            return sandbox_list.get("items", [])
-        except ApiException as e:
-            if e.status == 404:
-                return []
-            raise
-        except Exception as e:
-            logger.error(f"Unexpected error listing Sandboxes: {e}")
-            raise
+        """List Sandbox CRDs matching the given label selector."""
+        return self.k8s_client.list_custom_objects(
+            group=self.group,
+            version=self.version,
+            namespace=namespace,
+            plural=self.plural,
+            label_selector=label_selector,
+        )
 
     def update_expiration(self, sandbox_id: str, namespace: str, expires_at: datetime) -> None:
+        """Patch the Sandbox CRD shutdownTime field."""
         sandbox = self.get_workload(sandbox_id, namespace)
         if not sandbox:
             raise Exception(f"Sandbox for sandbox {sandbox_id} not found")
@@ -489,7 +408,7 @@ def update_expiration(self, sandbox_id: str, namespace: str, expires_at: datetim
             }
         }
 
-        self.custom_api.patch_namespaced_custom_object(
+        self.k8s_client.patch_custom_object(
             group=self.group,
             version=self.version,
             namespace=namespace,
@@ -499,6 +418,7 @@ def update_expiration(self, sandbox_id: str, namespace: str, expires_at: datetim
         )
 
     def get_expiration(self, workload: Dict[str, Any]) -> Optional[datetime]:
+        """Parse shutdownTime from Sandbox CRD spec."""
         spec = workload.get("spec", {})
         shutdown_time_str = spec.get("shutdownTime")
 
@@ -508,10 +428,11 @@ def get_expiration(self, workload: Dict[str, Any]) -> Optional[datetime]:
         try:
             return datetime.fromisoformat(shutdown_time_str.replace("Z", "+00:00"))
         except (ValueError, TypeError) as e:
-            logger.warning(f"Invalid shutdownTime format: {shutdown_time_str}, error: {e}")
+            logger.warning("Invalid shutdownTime format: %s, error: %s", shutdown_time_str, e)
             return None
 
     def get_status(self, workload: Dict[str, Any]) -> Dict[str, Any]:
+        """Derive sandbox state from the Sandbox CRD status conditions."""
         status = workload.get("status", {})
         conditions = status.get("conditions", [])
 
@@ -577,10 +498,10 @@ def _pod_state_from_selector(self, workload: Dict[str, Any]) -> Optional[tuple[s
             return None
 
         try:
-            pods = self.core_api.list_namespaced_pod(
+            pods = self.k8s_client.list_pods(
                 namespace=namespace,
                 label_selector=selector,
-            ).items
+            )
         except Exception:
             return None
 
@@ -620,15 +541,15 @@ def get_endpoint_info(self, workload: Dict[str, Any], port: int, sandbox_id: str
         namespace = workload.get("metadata", {}).get("namespace")
         if selector and namespace:
             try:
-                pods = self.core_api.list_namespaced_pod(
+                pods = self.k8s_client.list_pods(
                     namespace=namespace,
                     label_selector=selector,
-                ).items
+                )
                 for pod in pods:
                     if pod.status and pod.status.pod_ip and pod.status.phase == "Running":
                         return Endpoint(endpoint=f"{pod.status.pod_ip}:{port}")
             except Exception as e:
-                logger.warning(f"Failed to resolve pod endpoint: {e}")
+                logger.warning("Failed to resolve pod endpoint: %s", e)
 
         service_fqdn = status.get("serviceFQDN")
         if service_fqdn: