kubetorch 0.2.5__py3-none-any.whl

kubetorch/serving/autoscaling.py

@@ -0,0 +1,153 @@
+from dataclasses import asdict, dataclass
+
+from kubetorch.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class AutoScalingError(Exception):
+    pass
+
+
+@dataclass
+class AutoscalingConfig:
+    # The concurrent requests or requests per second threshold that triggers scaling
+    target: int = None
+
+    # The time window used to calculate the average of metrics for scaling decisions
+    window: str = None
+
+    # The metric type to base scaling decisions on:
+    # - concurrency: number of simultaneous requests
+    # - rps: requests per second
+    # - cpu: CPU utilization (requires HPA class)
+    # - memory: Memory utilization (requires HPA class)
+    metric: str = None
+
+    # The percentage of the target value at which to start scaling.
+    # E.g., if target=100 and target_utilization=70, scaling occurs at 70 requests
+    target_utilization: int = None
+
+    # Minimum number of replicas. 0 allows scaling to zero when idle
+    min_scale: int = None
+
+    # Maximum number of replicas the service can scale up to
+    max_scale: int = None
+
+    # Initial number of pods launched by the service
+    initial_scale: int = None
+
+    # Maximum concurrent requests per pod (containerConcurrency).
+    # If not set, pods accept unlimited concurrent requests.
+    concurrency: int = None
+
+    # Time to keep the last pod before scaling to zero (e.g., "30s", "1m5s")
+    scale_to_zero_pod_retention_period: str = None
+
+    # Delay before scaling down (e.g., "15m"). Only for KPA autoscaler.
+    scale_down_delay: str = None
+
+    # Autoscaler class: "kpa.autoscaling.knative.dev" or "hpa.autoscaling.knative.dev"
+    autoscaler_class: str = None
+
+    # Progress deadline for deployment (e.g., "10m"). Time to wait for deployment to be ready.
+    progress_deadline: str = None
+
+    def __init__(self, **kwargs):
+        """Support additional kwargs for autoscaling annotations"""
+        for field in self.__annotations__:
+            setattr(self, field, kwargs.pop(field, getattr(self, field, None)))
+
+        # set additional kwargs as annotations
+        self.extra_annotations = {f"autoscaling.knative.dev/{k}": str(v) for k, v in kwargs.items()}
+
+        self._validate()
+
+    def _validate(self):
+        """Validation logic moved to separate method"""
+        if self.min_scale is not None and self.max_scale is not None and self.min_scale > self.max_scale:
+            raise AutoScalingError("min_scale cannot be greater than max_scale")
+        if self.window is not None and not self.window.endswith(("s", "m", "h")):
+            raise AutoScalingError("window must end with s, m, or h")
+        if self.target_utilization is not None and (self.target_utilization <= 0 or self.target_utilization > 100):
+            raise AutoScalingError("target_utilization must be between 1 and 100")
+        if self.scale_to_zero_pod_retention_period is not None:
+            # Validate time format (e.g., "30s", "1m5s", "2h")
+            import re
+
+            if not re.match(r"^\d+[smh](\d+[smh])*$", self.scale_to_zero_pod_retention_period):
+                raise AutoScalingError(
+                    "scale_to_zero_pod_retention_period must be a valid duration (e.g., '30s', '1m5s')"
+                )
+        if self.scale_down_delay is not None:
+            # Validate time format
+            import re
+
+            if not re.match(r"^\d+[smh](\d+[smh])*$", self.scale_down_delay):
+                raise AutoScalingError("scale_down_delay must be a valid duration (e.g., '15m', '1h')")
+        if self.autoscaler_class is not None and self.autoscaler_class not in [
+            "kpa.autoscaling.knative.dev",
+            "hpa.autoscaling.knative.dev",
+        ]:
+            raise AutoScalingError(
+                "autoscaler_class must be 'kpa.autoscaling.knative.dev' or 'hpa.autoscaling.knative.dev'"
+            )
+        if self.progress_deadline is not None:
+            # Validate time format
+            import re
+
+            if not re.match(r"^\d+[smh](\d+[smh])*$", self.progress_deadline):
+                raise AutoScalingError("progress_deadline must be a valid duration (e.g., '10m', '600s')")
+
+    def __post_init__(self):
+        """Call the same validation for dataclass initialization"""
+        self._validate()
+
+    def dict(self):
+        return asdict(self)
+
+    def convert_to_annotations(self) -> dict:
+        """Convert config to a dictionary of annotations for Knative"""
+        annotations = {}
+
+        # Set autoscaler class if specified, otherwise use default KPA
+        if self.autoscaler_class is not None:
+            annotations["autoscaling.knative.dev/class"] = self.autoscaler_class
+        else:
+            annotations["autoscaling.knative.dev/class"] = "kpa.autoscaling.knative.dev"
+
+        # Only set annotations for values that were explicitly provided
+        if self.target is not None:
+            annotations["autoscaling.knative.dev/target"] = str(self.target)
+
+        if self.min_scale is not None:
+            annotations["autoscaling.knative.dev/min-scale"] = str(self.min_scale)
+
+        if self.max_scale is not None:
+            annotations["autoscaling.knative.dev/max-scale"] = str(self.max_scale)
+
+        if self.window is not None:
+            annotations["autoscaling.knative.dev/window"] = self.window
+
+        if self.metric is not None:
+            annotations["autoscaling.knative.dev/metric"] = self.metric
+
+        if self.target_utilization is not None:
+            annotations["autoscaling.knative.dev/target-utilization-percentage"] = str(self.target_utilization)
+
+        if self.initial_scale is not None:
+            annotations["autoscaling.knative.dev/initial-scale"] = str(self.initial_scale)
+
+        if self.scale_to_zero_pod_retention_period is not None:
+            annotations[
+                "autoscaling.knative.dev/scale-to-zero-pod-retention-period"
+            ] = self.scale_to_zero_pod_retention_period
+
+        if self.scale_down_delay is not None:
+            annotations["autoscaling.knative.dev/scale-down-delay"] = self.scale_down_delay
+
+        # Add any extra annotations from the config
+        if hasattr(self, "extra_annotations"):
+            annotations.update(self.extra_annotations)
+
+        return annotations

kubetorch/serving/base_service_manager.py

@@ -0,0 +1,344 @@
+import importlib
+
+from abc import abstractmethod
+from typing import Dict, List, Optional
+
+import yaml
+from jinja2 import Template
+
+from kubernetes import client, utils
+from kubernetes.client import AppsV1Api, CoreV1Api, CustomObjectsApi
+
+import kubetorch.serving.constants as serving_constants
+from kubetorch import globals
+
+from kubetorch.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class BaseServiceManager:
+    """Base service manager with common functionality for all service types."""
+
+    def __init__(
+        self,
+        objects_api: CustomObjectsApi,
+        core_api: CoreV1Api,
+        apps_v1_api: AppsV1Api,
+        namespace: str,
+    ):
+        self.objects_api = objects_api
+        self.core_api = core_api
+        self.apps_v1_api = apps_v1_api
+
+        # Load config
+        self.global_config = globals.config
+        self.namespace = namespace or self.global_config.namespace
+
+    @property
+    def username(self):
+        return self.global_config.username
+
+    @property
+    def base_labels(self):
+        """Base labels for all resources created by the service manager."""
+        from kubetorch import __version__
+
+        labels = {
+            serving_constants.KT_VERSION_LABEL: __version__,
+        }
+        if self.username:
+            labels[serving_constants.KT_USERNAME_LABEL] = self.username
+
+        return labels
+
+    def _apply_yaml_template(self, yaml_file, replace_existing=False, **kwargs):
+        with importlib.resources.files("kubetorch.serving.templates").joinpath(yaml_file).open("r") as f:
+            template = Template(f.read())
+
+        yaml_content = template.render(**kwargs)
+        yaml_objects = list(yaml.safe_load_all(yaml_content))
+        k8s_client = client.ApiClient()
+
+        for obj in yaml_objects:
+            logger.info(f"Applying {obj.get('kind')}/{obj.get('metadata', {}).get('name')}")
+            try:
+                if replace_existing:
+                    # Try to delete existing resource first
+                    try:
+                        utils.delete_from_dict(k8s_client, obj)
+                    except client.exceptions.ApiException as e:
+                        if e.status != 404:  # Ignore if resource doesn't exist
+                            raise
+
+                utils.create_from_dict(k8s_client, obj)
+                logger.info(f"Successfully applied {obj.get('kind')}/{obj.get('metadata', {}).get('name')}")
+            except utils.FailToCreateError as e:
+                if "already exists" in str(e):
+                    logger.info(f"Resource already exists: {obj.get('kind')}/{obj.get('metadata', {}).get('name')}")
+                else:
+                    raise
+
+    @abstractmethod
+    def get_deployment_timestamp_annotation(self, service_name: str) -> Optional[str]:
+        """Get deployment timestamp annotation for this service type."""
+        pass
+
+    @abstractmethod
+    def update_deployment_timestamp_annotation(self, service_name: str, new_timestamp: str) -> str:
+        """Update deployment timestamp annotation for this service type."""
+        pass
+
+    def fetch_kubetorch_config(self) -> dict:
+        """Fetch the kubetorch configmap from the namespace."""
+        try:
+            kubetorch_config = self.core_api.read_namespaced_config_map(
+                name="kubetorch-config", namespace=globals.config.install_namespace
+            )
+            return kubetorch_config.data
+        except client.exceptions.ApiException as e:
+            if e.status != 404:
+                logger.error(f"Error fetching kubetorch config: {e}")
+            return {}
+
+    @staticmethod
+    def discover_services_static(
+        namespace: str, objects_api=None, apps_v1_api=None, name_filter: str = None
+    ) -> List[Dict]:
+        """Static method to discover Kubetorch services without ServiceManager instance.
+
+        Uses parallel API calls for faster discovery across service types.
+
+        Args:
+            namespace: Kubernetes namespace
+            objects_api: Optional CustomObjectsApi instance (created if None)
+            apps_v1_api: Optional AppsV1Api instance (created if None)
+            name_filter: Optional name filter for services
+
+        Returns:
+            List of service dictionaries with structure:
+            {
+                'name': str,
+                'template_type': str,  # 'ksvc', 'deployment', 'raycluster'
+                'resource': dict,
+                'namespace': str,
+                'creation_timestamp': str  # ISO format
+            }
+        """
+        import concurrent.futures
+        import threading
+
+        if objects_api is None:
+            objects_api = client.CustomObjectsApi()
+        if apps_v1_api is None:
+            apps_v1_api = client.AppsV1Api()
+
+        services = []
+        services_lock = threading.Lock()
+
+        def fetch_knative_services():
+            """Fetch Knative services in parallel."""
+            try:
+                label_selector = f"{serving_constants.KT_TEMPLATE_LABEL}=ksvc"
+                knative_services = objects_api.list_namespaced_custom_object(
+                    group="serving.knative.dev",
+                    version="v1",
+                    namespace=namespace,
+                    plural="services",
+                    label_selector=label_selector,
+                )["items"]
+
+                local_services = []
+                for svc in knative_services:
+                    svc_name = svc["metadata"]["name"]
+                    if name_filter and name_filter not in svc_name:
+                        continue
+
+                    local_services.append(
+                        {
+                            "name": svc_name,
+                            "template_type": "ksvc",
+                            "resource": svc,  # Already a dict
+                            "namespace": namespace,
+                            "creation_timestamp": svc["metadata"]["creationTimestamp"],
+                        }
+                    )
+
+                with services_lock:
+                    services.extend(local_services)
+
+            except client.exceptions.ApiException as e:
+                if e.status != 404:  # Ignore if Knative is not installed
+                    logger.warning(f"Failed to list Knative services: {e}")
+
+        def fetch_deployments():
+            """Fetch Deployments in parallel."""
+            try:
+                label_selector = f"{serving_constants.KT_TEMPLATE_LABEL}=deployment"
+                deployments = apps_v1_api.list_namespaced_deployment(
+                    namespace=namespace,
+                    label_selector=label_selector,
+                )
+
+                local_services = []
+                for deployment in deployments.items:
+                    deploy_name = deployment.metadata.name
+                    if name_filter and name_filter not in deploy_name:
+                        continue
+
+                    # Convert V1Deployment object to dictionary for consistency
+                    deployment_dict = client.ApiClient().sanitize_for_serialization(deployment)
+
+                    # Add kind and apiVersion (not included in V1Deployment object)
+                    deployment_dict["kind"] = "Deployment"
+                    deployment_dict["apiVersion"] = "apps/v1"
+
+                    local_services.append(
+                        {
+                            "name": deploy_name,
+                            "template_type": "deployment",
+                            "resource": deployment_dict,  # Now consistently a dict
+                            "namespace": namespace,
+                            "creation_timestamp": deployment.metadata.creation_timestamp.isoformat() + "Z",
+                        }
+                    )
+
+                with services_lock:
+                    services.extend(local_services)
+
+            except client.exceptions.ApiException as e:
+                logger.warning(f"Failed to list Deployments: {e}")
+
+        def fetch_rayclusters():
+            """Fetch RayClusters in parallel."""
+            try:
+                label_selector = f"{serving_constants.KT_TEMPLATE_LABEL}=raycluster"
+                rayclusters = objects_api.list_namespaced_custom_object(
+                    group="ray.io",
+                    version="v1",
+                    namespace=namespace,
+                    plural="rayclusters",
+                    label_selector=label_selector,
+                )["items"]
+
+                local_services = []
+                for cluster in rayclusters:
+                    cluster_name = cluster["metadata"]["name"]
+                    if name_filter and name_filter not in cluster_name:
+                        continue
+
+                    local_services.append(
+                        {
+                            "name": cluster_name,
+                            "template_type": "raycluster",
+                            "resource": cluster,  # Already a dict
+                            "namespace": namespace,
+                            "creation_timestamp": cluster["metadata"]["creationTimestamp"],
+                        }
+                    )
+
+                with services_lock:
+                    services.extend(local_services)
+
+            except client.exceptions.ApiException as e:
+                if e.status != 404:
+                    logger.warning(f"Failed to list RayClusters: {e}")
+
+        # Execute all API calls in parallel
+        with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+            futures = [
+                executor.submit(fetch_knative_services),
+                executor.submit(fetch_deployments),
+                executor.submit(fetch_rayclusters),
+            ]
+
+            # Wait for all to complete
+            concurrent.futures.wait(futures)
+
+        return services
+
+    @staticmethod
+    def get_pods_for_service_static(
+        service_name: str,
+        namespace: str,
+        core_api=None,
+    ) -> List:
+        """Static method to get pods for a service across different service types.
+
+        Args:
+            service_name: Name of the service
+            namespace: Kubernetes namespace
+            core_api: Optional CoreV1Api instance (created if None)
+
+        Returns:
+            List of pod objects
+        """
+        if core_api is None:
+            core_api = client.CoreV1Api()
+
+        # Build label selector
+        label_selector = f"{serving_constants.KT_SERVICE_LABEL}={service_name}"
+        try:
+            pods = core_api.list_namespaced_pod(namespace=namespace, label_selector=label_selector)
+            return pods.items
+        except client.exceptions.ApiException as e:
+            logger.warning(f"Failed to list pods for service {service_name}: {e}")
+            return []
+
+    def discover_all_services(self, namespace: str = None) -> List[Dict]:
+        """Discover all Kubetorch services across different resource types.
+
+        Returns a list of service dictionaries with normalized structure:
+        {
+            'name': str,
+            'template_type': str,  # 'ksvc', 'deployment', 'raycluster'
+            'resource': object,    # The actual Kubernetes resource object
+            'namespace': str
+        }
+        """
+        return self.discover_services_static(
+            namespace=namespace or self.namespace,
+            objects_api=self.objects_api,
+            apps_v1_api=self.apps_v1_api,
+        )
+
+    # Abstract methods to be implemented by subclasses
+    def create_or_update_service(self, *args, **kwargs):
+        raise NotImplementedError("Subclasses must implement create_or_update_service")
+
+    def get_endpoint(self, service_name: str) -> str:
+        raise NotImplementedError("Subclasses must implement get_endpoint")
+
+    def get_pods_for_service(self, service_name: str, **kwargs) -> List[client.V1Pod]:
+        raise NotImplementedError("Subclasses must implement get_pods_for_service")
+
+    def check_service_ready(self, service_name: str, launch_timeout: int, **kwargs) -> bool:
+        """Check if service is ready to serve requests.
+
+        This method should be implemented by subclasses to provide service-type-specific
+        readiness checking logic.
+
+        Args:
+            service_name: Name of the service to check
+            launch_timeout: Timeout in seconds to wait for service to be ready
+            **kwargs: Additional arguments for readiness checking
+
+        Returns:
+            True if service is ready, raises exception if timeout or error
+        """
+        raise NotImplementedError("Subclasses must implement check_service_ready")
+
+    def teardown_service(self, service_name: str, console=None) -> bool:
+        """Teardown/delete service and associated resources.
+
+        This method should be implemented by subclasses to provide service-type-specific
+        teardown logic.
+
+        Args:
+            service_name: Name of the service to teardown
+            console: Optional Rich console for output
+
+        Returns:
+            True if teardown was successful, False otherwise
+        """
+        raise NotImplementedError("Subclasses must implement teardown_service")

kubetorch/serving/constants.py

@@ -0,0 +1,77 @@
+# K8s Configuration
+KUBECTL_PORT = 6443
+KT_LAUNCH_TIMEOUT = 900  # 15 minutes
+
+# Ports
+DEFAULT_NGINX_PORT = 8080
+LOCAL_NGINX_PORT = 38080
+DEFAULT_KT_SERVER_PORT = 32300  # Standard port of Knative services
+DEFAULT_DEBUG_PORT = 5678
+
+# Namespaces
+KUBETORCH_NAMESPACE = "kubetorch"
+RUNHOUSE_NAMESPACE = "runhouse"
+DEFAULT_NAMESPACE = "default"
+
+# Images
+SERVER_IMAGE_MINIMAL = "ghcr.io/run-house/server:v3"
+SERVER_IMAGE_WITH_OTEL = "ghcr.io/run-house/server-otel:v3"
+
+UBUNTU_IMAGE_MINIMAL = "ghcr.io/run-house/ubuntu:v1"
+UBUNTU_IMAGE_WITH_OTEL = "ghcr.io/run-house/ubuntu:v1"
+
+DEFAULT_PROXY_IMAGE = "ghcr.io/run-house/proxy:v2"
+KUBETORCH_IMAGE_TRAPDOOR = "kubetorch"
+
+# Service Accounts
+DEFAULT_SERVICE_ACCOUNT_NAME = "kubetorch-service-account"
+
+# Annotations
+INACTIVITY_TTL_ANNOTATION = "kubetorch.com/inactivity-ttl"
+KUBECONFIG_PATH_ANNOTATION = "kubetorch.com/kubeconfig-path"
+
+# Labels
+KT_SERVICE_LABEL = "kubetorch.com/service"
+KT_VERSION_LABEL = "kubetorch.com/version"
+KT_MODULE_LABEL = "kubetorch.com/module"
+KT_USER_IDENTIFIER_LABEL = "kubetorch.com/user-identifier"
+KT_USERNAME_LABEL = "kubetorch.com/username"
+KT_POD_TYPE_LABEL = "kubetorch.com/pod-type"
+KT_TEMPLATE_LABEL = "kubetorch.com/template"
+KT_SECRET_NAME_LABEL = "kubetorch.com/secret-name"
+
+# Templates
+TTL_CONTROLLER_CONFIGMAP_NAME = "kubetorch-ttl-controller-config"
+KNATIVE_SERVICE_TEMPLATE_FILE = "knative_service_template.yaml"
+POD_TEMPLATE_FILE = "pod_template.yaml"
+KT_SETUP_TEMPLATE_FILE = "kt_setup_template.sh.j2"
+DEPLOYMENT_TEMPLATE_FILE = "deployment_template.yaml"
+DEPLOYMENT_SERVICE_TEMPLATE_FILE = "service_template.yaml"
+RAYCLUSTER_TEMPLATE_FILE = "raycluster_template.yaml"
+RAYCLUSTER_SERVICE_TEMPLATE_FILE = "raycluster_service_template.yaml"
+
+# Loki
+LOKI_GATEWAY_SERVICE_NAME = "loki-gateway"
+
+# Prometheus
+PROMETHEUS_SERVICE_NAME = "kubetorch-metrics"
+
+# Grafana
+GRAFANA_HEALTH_ENDPOINT = "/api/health"
+PROMETHEUS_HEALTH_ENDPOINT = "/metrics"
+
+# KAI
+KAI_SCHEDULER_NAME = "kai-scheduler"
+KAI_SCHEDULER_LABEL = "kai.scheduler/queue"
+
+# HTTP Client
+KT_TERMINATION_REASONS = ["OOMKilled", "Evicted", "Not Found"]
+
+# NGINX
+NGINX_GATEWAY_PROXY = "kubetorch-proxy"
+DEFAULT_NGINX_HEALTH_ENDPOINT = "/health"
+
+# Rsync
+RSYNC_LOCAL_PORT = 3873
+REMOTE_RSYNC_PORT = 873
+RSYNC_SERVICE_NAME = "kubetorch-rsync"