kubetorch 0.2.0__py3-none-any.whl

kubetorch/serving/utils.py

@@ -0,0 +1,377 @@
+import os
+import socket
+import time
+import warnings
+
+from dataclasses import dataclass
+from typing import Literal, Optional, Union
+
+import httpx
+from kubernetes.client import ApiException, CoreV1Api, V1Pod
+from kubernetes.utils import parse_quantity
+
+from kubetorch import globals
+from kubetorch.logger import get_logger
+from kubetorch.servers.http.utils import is_running_in_kubernetes
+from kubetorch.serving.constants import (
+    KUBETORCH_MONITORING_NAMESPACE,
+    LOKI_GATEWAY_SERVICE_NAME,
+    PROMETHEUS_SERVICE_NAME,
+    PROMETHEUS_URL,
+)
+from kubetorch.utils import load_kubeconfig
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class GPUConfig:
+    count: Optional[int] = None
+    memory: Optional[str] = None
+    sharing_type: Optional[Literal["memory", "fraction"]] = None
+    gpu_memory: Optional[str] = None
+    gpu_fraction: Optional[str] = None
+    gpu_type: Optional[str] = None
+
+    def __post_init__(self):
+        self.validate()
+
+    def validate(self) -> bool:
+        if self.count and not isinstance(self.count, int):
+            raise ValueError("GPU count must an int")
+
+        if self.sharing_type == "memory":
+            if not self.gpu_memory:
+                raise ValueError(
+                    "GPU memory must be specified when using memory sharing"
+                )
+        elif self.sharing_type == "fraction":
+            if not self.gpu_fraction:
+                raise ValueError(
+                    "GPU fraction must be specified when using fraction sharing"
+                )
+            try:
+                fraction = float(self.gpu_fraction)
+                if not 0 < fraction <= 1:
+                    raise ValueError("GPU fraction must be between 0 and 1")
+            except ValueError:
+                raise ValueError("GPU fraction must be a valid float between 0 and 1")
+
+        return True
+
+    def to_dict(self) -> dict:
+        base_dict = {
+            "sharing_type": self.sharing_type,
+            "count": self.count,
+        }
+
+        if self.memory is not None:
+            base_dict["memory"] = self.memory
+
+        if self.sharing_type == "memory" and self.gpu_memory:
+            base_dict["gpu_memory"] = self.gpu_memory
+        if self.sharing_type == "fraction" and self.gpu_fraction:
+            # Convert to millicores format
+            fraction = float(self.gpu_fraction)
+            base_dict["gpu_fraction"] = f"{int(fraction * 1000)}m"
+        if self.gpu_type is not None:
+            base_dict["gpu_type"] = self.gpu_type
+
+        return base_dict
+
+
+class RequestedPodResources:
+    """Resources requested in a Kubetorch cluster/compute object. Note these are the values we receive
+    from launcher the cluster via a Sky dryrun."""
+
+    # Default overhead percentages to account for filesystem overhead, OS files, logs, container runtime, etc.
+    MEMORY_OVERHEAD = 0.20
+    CPU_OVERHEAD = 0.10
+    DISK_OVERHEAD = 0.15
+    GPU_OVERHEAD = 0.0
+
+    MIN_MEMORY_GB = 0.1  # 100Mi minimum
+    MIN_CPU_CORES = 0.1  # 100m minimum
+
+    CPU_STEPS = [1, 2, 4, 8, 16, 32, 48, 64, 96, 128, 192]
+    MEMORY_STEPS = [0.5, 1, 2, 4, 8, 16, 32, 48, 64, 96, 128, 192, 256, 384, 512, 768]
+
+    def __init__(
+        self,
+        memory: Optional[Union[str, float]] = None,
+        cpus: Optional[Union[int, float]] = None,
+        disk_size: Optional[int] = None,
+        num_gpus: Optional[Union[int, dict]] = None,
+    ):
+
+        self.memory = (
+            max(float(memory), self.MIN_MEMORY_GB) if memory is not None else None
+        )
+        self.cpus = (
+            max(self.normalize_cpu_value(cpus), self.MIN_CPU_CORES)
+            if cpus is not None
+            else None
+        )
+        self.disk_size = disk_size
+        self.num_gpus = num_gpus
+
+    def __str__(self):
+        # Example: RequestedPodResources(memory=16.0, cpus=4.0, disk=NoneGB, gpus={'A10G': 1})"
+        disk_str = f"{self.disk_size}GB" if self.disk_size is not None else "None"
+        memory = f"{self.memory}GB" if self.memory is not None else "None"
+
+        return (
+            f"RequestedPodResources(memory={memory}, cpus={self.cpus}, disk_size={disk_str}, "
+            f"num_gpus={self.num_gpus})"
+        )
+
+    def __repr__(self):
+        return (
+            f"RequestedPodResources(memory={self.memory}, cpus={self.cpus}, "
+            f"disk_size={self.disk_size}, num_gpus={self.num_gpus})"
+        )
+
+    @classmethod
+    def cpu_for_resource_request(cls, cpu_val: int = None):
+        if cpu_val is None:
+            return None
+
+        # Ensure minimum CPU value
+        cpu_val = max(float(cpu_val), cls.MIN_CPU_CORES)
+
+        # Convert to millicores (ex: '4.0' -> 4000m)
+        return f"{int(float(cpu_val) * 1000)}m"
+
+    @classmethod
+    def memory_for_resource_request(cls, memory_val: Union[str, float, int] = None):
+        if memory_val is None:
+            return None
+
+        # If it's a number, treat as GB
+        if isinstance(memory_val, (int, float)):
+            gb_val = max(float(memory_val), cls.MIN_MEMORY_GB)
+            memory_val = f"{gb_val}Gi"
+
+        # Validate the string - if invalid will throw a ValueError
+        parse_quantity(str(memory_val))
+
+        return str(memory_val)
+
+    @classmethod
+    def normalize_cpu_value(
+        cls, cpu_value: Optional[Union[int, str, float]]
+    ) -> Optional[float]:
+        """Convert CPU value to float, handling string values with '+' allowed by Sky and Kubetorch."""
+        if cpu_value is None:
+            return None
+
+        if isinstance(cpu_value, str):
+            # Strip the '+' if present and convert to float
+            return float(cpu_value.rstrip("+"))
+
+        return float(cpu_value)
+
+
+def check_kubetorch_versions(response):
+    from kubetorch import __version__ as python_client_version, VersionMismatchError
+
+    try:
+        data = response.json()
+    except ValueError:
+        # older nginx proxy versions won't return a JSON
+        return
+
+    helm_installed_version = data.get("version")
+    if not helm_installed_version:
+        logger.debug("No 'version' found in health check response")
+        return
+
+    if python_client_version != helm_installed_version:
+        msg = (
+            f"client={python_client_version}, cluster={helm_installed_version}. "
+            "To suppress this error, set the environment variable "
+            "`KUBETORCH_IGNORE_VERSION_MISMATCH=1`."
+        )
+        if not os.getenv("KUBETORCH_IGNORE_VERSION_MISMATCH"):
+            raise VersionMismatchError(msg)
+
+        warnings.warn(f"Kubetorch version mismatch: {msg}")
+
+
+def extract_config_from_nginx_health_check(response):
+    """Extract the config from the nginx health check response."""
+    try:
+        data = response.json()
+    except ValueError:
+        return
+    config = data.get("config", {})
+    return config
+
+
+def wait_for_port_forward(
+    process,
+    local_port,
+    timeout=30,
+    health_endpoint: str = None,
+    validate_kubetorch_versions: bool = True,
+):
+    from kubetorch import VersionMismatchError
+
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        if process.poll() is not None:
+            stderr = process.stderr.read().decode()
+            raise Exception(f"Port forward failed: {stderr}")
+
+        try:
+            # Check if socket is open
+            with socket.create_connection(("localhost", local_port), timeout=1):
+                if not health_endpoint:
+                    # If we are not checking HTTP (ex: rsync)
+                    return True
+        except OSError:
+            time.sleep(0.2)
+            continue
+
+        if health_endpoint:
+            url = f"http://localhost:{local_port}" + health_endpoint
+            try:
+                # Check if HTTP endpoint is ready
+                resp = httpx.get(url, timeout=2)
+                if resp.status_code == 200:
+                    if validate_kubetorch_versions:
+                        check_kubetorch_versions(resp)
+                    # Extract config to set outside of function scope
+                    config = extract_config_from_nginx_health_check(resp)
+                    return config
+            except VersionMismatchError as e:
+                raise e
+            except Exception as e:
+                logger.debug(f"Waiting for HTTP endpoint to be ready: {e}")
+
+        time.sleep(0.2)
+
+    raise TimeoutError("Timeout waiting for port forward to be ready")
+
+
+def pod_is_running(pod: V1Pod):
+    return pod.status.phase == "Running" and pod.metadata.deletion_timestamp is None
+
+
+def check_loki_enabled(core_api: CoreV1Api = None) -> bool:
+    """Check if loki is enabled"""
+    if core_api is None:
+        load_kubeconfig()
+        core_api = CoreV1Api()
+
+    kt_namespace = globals.config.install_namespace
+
+    try:
+        # Check if loki-gateway service exists in the namespace
+        core_api.read_namespaced_service(
+            name=LOKI_GATEWAY_SERVICE_NAME, namespace=kt_namespace
+        )
+        logger.debug(f"Loki gateway service found in namespace {kt_namespace}")
+    except ApiException as e:
+        if e.status == 404:
+            logger.debug(f"Loki gateway service not found in namespace {kt_namespace}")
+            return False
+
+        # Additional permission-proof check: try to ping the internal Loki gateway URL
+        # Needed if running in kubernetes without full kubeconfig permissions
+        if is_running_in_kubernetes():
+            try:
+                loki_url = f"http://loki-gateway.{kt_namespace}.svc.cluster.local/loki/api/v1/labels"
+                response = httpx.get(loki_url, timeout=2)
+                if response.status_code == 200:
+                    logger.debug("Loki gateway is reachable")
+                else:
+                    logger.debug(f"Loki gateway returned status {response.status_code}")
+                    return False
+            except Exception as e:
+                logger.debug(f"Loki gateway is not reachable: {e}")
+                return False
+
+    return True
+
+
+def check_prometheus_enabled(
+    prometheus_url: str, namespace: str, core_api: CoreV1Api = None
+) -> bool:
+    """Check if Prometheus is enabled and reachable."""
+    if prometheus_url and prometheus_url != PROMETHEUS_URL:
+        return True
+
+    if core_api is None:
+        load_kubeconfig()
+        core_api = CoreV1Api()
+
+    is_in_kubernetes = is_running_in_kubernetes()
+
+    # Check namespace exists
+    try:
+        core_api.read_namespace(name=namespace)
+    except ApiException as e:
+        if e.status == 404:
+            logger.debug(f"Prometheus namespace not found: {namespace}")
+            return False
+
+    # Check Prometheus service exists
+    try:
+        core_api.read_namespaced_service(
+            name=PROMETHEUS_SERVICE_NAME, namespace=namespace
+        )
+        logger.debug(f"Prometheus service found in namespace {namespace}")
+    except ApiException as e:
+        if e.status == 404:
+            logger.debug(f"Prometheus service not found: {PROMETHEUS_SERVICE_NAME}")
+            return False
+
+        # If running inside the cluster, try hitting the service directly
+        if is_in_kubernetes:
+            try:
+                response = httpx.get(PROMETHEUS_URL, timeout=2)
+                if response.status_code == 200:
+                    logger.debug("Prometheus is reachable and healthy")
+                else:
+                    logger.debug(f"Prometheus returned status {response.status_code}")
+                    return False
+            except Exception as e:
+                logger.debug(f"Prometheus is not reachable: {e}")
+                return False
+
+    return True
+
+
+def check_tempo_enabled(core_api: CoreV1Api = None) -> bool:
+    if core_api is None:
+        load_kubeconfig()
+        core_api = CoreV1Api()
+
+    try:
+        otel = core_api.read_namespaced_service(
+            name="kubetorch-otel-opentelemetry-collector",
+            namespace=KUBETORCH_MONITORING_NAMESPACE,
+        )
+        tempo = core_api.read_namespaced_service(
+            name="kubetorch-otel-tempo-distributor",
+            namespace=KUBETORCH_MONITORING_NAMESPACE,
+        )
+        return otel is not None and tempo is not None
+
+    except ApiException as e:
+        if e.status == 404:
+            return False
+        raise
+
+
+def nested_override(original_dict, override_dict):
+    for key, value in override_dict.items():
+        if key in original_dict:
+            if isinstance(original_dict[key], dict) and isinstance(value, dict):
+                # Recursively merge nested dictionaries
+                nested_override(original_dict[key], value)
+            else:
+                original_dict[key] = value  # Custom wins
+        else:
+            original_dict[key] = value

kubetorch/utils.py

@@ -0,0 +1,284 @@
+import enum
+import importlib.util
+import json
+import os
+import re
+import subprocess
+import sys
+from datetime import datetime, timezone
+from io import StringIO
+from pathlib import Path
+from urllib.parse import urlparse
+
+from kubernetes import client, config
+
+from kubetorch.constants import DEFAULT_KUBECONFIG_PATH, MAX_USERNAME_LENGTH
+from kubetorch.globals import config as kt_config
+from kubetorch.logger import get_logger
+from kubetorch.resources.callables.utils import get_local_install_path
+
+logger = get_logger(__name__)
+
+
+def extract_host_port(url: str):
+    """Extract host and port when needed separately from a URL."""
+    p = urlparse(url)
+    return p.hostname, (p.port or (443 if p.scheme == "https" else 80))
+
+
+def http_to_ws(url: str) -> str:
+    """Convert HTTP/HTTPS URLs to WebSocket URLs, or return as-is if already WS."""
+    if url.startswith("https://"):
+        return "wss://" + url[len("https://") :]
+    if url.startswith("http://"):
+        return "ws://" + url[len("http://") :]
+    if url.startswith(("ws://", "wss://")):
+        return url  # already WebSocket URL
+    # Default to ws:// for bare hostnames
+    return "ws://" + url
+
+
+def validate_username(username):
+    if username is None:  # will be used in case we run kt config user username
+        return username
+    # Kubernetes requires service names to follow DNS-1035 label standards
+    original_username = username  # if an exception is raised because the username is invalid, we want to print the original provided name
+    username = username.lower().replace("_", "-").replace("/", "-")
+    # Make sure the first character is a letter
+    if not re.match(r"^[a-z]", username):
+        # Strip out all the characters before the first letter with a regex
+        username = re.sub(r"^[^a-z]*", "", username)
+    username = username[:MAX_USERNAME_LENGTH]
+    # Make sure username doesn't end or start with a hyphen
+    if username.startswith("-") or username.endswith("-"):
+        username = username.strip("-")
+    reserved = ["kt", "kubetorch", "knative"]
+    if username in reserved:
+        raise ValueError(
+            f"{original_username} is one of the reserved names: {', '.join(reserved)}"
+        )
+    if not re.match(r"^[a-z]([-a-z0-9]*[a-z0-9])?$", username):
+        raise ValueError(f"{original_username} must be a valid k8s name")
+    return username
+
+
+def load_kubeconfig():
+    try:
+        config.load_incluster_config()
+    except config.config_exception.ConfigException:
+        kubeconfig_path = os.getenv("KUBECONFIG") or DEFAULT_KUBECONFIG_PATH
+        abs_path = Path(kubeconfig_path).expanduser()
+        if not abs_path.exists():
+            raise FileNotFoundError(f"Kubeconfig file not found in path: {abs_path}")
+        config.load_kube_config(str(abs_path))
+
+
+def current_git_branch():
+    try:
+        # For CI env
+        branch = (
+            os.environ.get("GITHUB_HEAD_REF")  # PRs: source branch name
+            or os.environ.get("GITHUB_REF_NAME")  # Pushes: branch name
+            or os.environ.get("CI_COMMIT_REF_NAME")  # GitLab
+            or os.environ.get("CIRCLE_BRANCH")  # CircleCI
+        )
+        if not branch:
+            branch = (
+                subprocess.check_output(
+                    ["git", "rev-parse", "--abbrev-ref", "HEAD"],
+                    stderr=subprocess.DEVNULL,
+                )
+                .decode("utf-8")
+                .strip()
+            )
+        return branch
+    except Exception as e:
+        logger.debug(f"Failed to load current git branch: {e}")
+        return None
+
+
+def iso_timestamp_to_nanoseconds(timestamp):
+    if timestamp is None:
+        dt = datetime.now(timezone.utc)
+    elif isinstance(timestamp, datetime):
+        dt = timestamp
+        if dt.tzinfo is None:
+            dt = dt.replace(tzinfo=timezone.utc)
+    elif isinstance(timestamp, str):
+        dt = datetime.fromisoformat(timestamp)
+        if dt.tzinfo is None:
+            dt = dt.replace(tzinfo=timezone.utc)
+    else:
+        raise ValueError(f"Unsupported timestamp type: {type(timestamp)}")
+    return int(dt.timestamp() * 1e9)
+
+
+def get_kt_install_url(freeze: bool = False):
+    # Returns:
+    #   str: kubetorch install url
+    #   bool: whether to install in editable mode
+    if kt_config.install_url or freeze:
+        return kt_config.install_url, False
+    local_kt_path = get_local_install_path("kubetorch")
+    if local_kt_path and (Path(local_kt_path) / "pyproject.toml").exists():
+        return local_kt_path, True
+    else:
+        # If the user is using uv, sometimes the pip freeze command won't return the full "kubetorch @ ..." install
+        # URL. However, running `uv pip freeze` will return the full URL, even when the user isn't actually in a uv
+        # venv. So we default to using `uv pip freeze` if uv is installed, otherwise we use the regular `pip freeze`.
+        uv_installed = importlib.util.find_spec("uv") is not None
+        # check if installed from presigned url
+        freeze_cmd = (
+            "pip freeze | grep kubetorch"
+            if not uv_installed
+            else "uv pip freeze | grep kubetorch"
+        )
+        output = subprocess.run(
+            freeze_cmd,
+            shell=True,
+            capture_output=True,
+            text=True,
+        ).stdout
+        if not output.startswith("kubetorch @ "):
+            raise Exception(
+                "Could not find kubetorch version to install on the pod. You must either set the "
+                "``install_url`` in the config or ``KT_INSTALL_URL`` env var, have kubetorch installed "
+                "locally, or set ``compute.freeze`` to ``True``."
+            )
+        install_url = output[len("kubetorch @ ") :].strip()
+        return install_url, False
+
+
+class LogVerbosity(str, enum.Enum):
+    DEBUG = "debug"
+    INFO = "info"
+    CRITICAL = "critical"
+
+
+####################################################################################################
+# Logging redirection
+####################################################################################################
+class StreamTee(object):
+    def __init__(self, instream, outstreams):
+        self.instream = instream
+        self.outstreams = outstreams
+
+    def write(self, message):
+        self.instream.write(message)
+        for stream in self.outstreams:
+            if message:
+                stream.write(message)
+                # We flush here to ensure that the logs are written to the file immediately
+                # see https://github.com/run-house/runhouse/pull/724
+                stream.flush()
+
+    def writelines(self, lines):
+        self.instream.writelines(lines)
+        for stream in self.outstreams:
+            stream.writelines(lines)
+            stream.flush()
+
+    def flush(self):
+        self.instream.flush()
+        for stream in self.outstreams:
+            stream.flush()
+
+    def __getattr__(self, item):
+        # Needed in case someone calls a method on instream, such as Ray calling sys.stdout.istty()
+        return getattr(self.instream, item)
+
+
+class capture_stdout:
+    """Context manager for capturing stdout to a file, list, or stream, while still printing to stdout."""
+
+    def __init__(self, output=None):
+        self.output = output
+        self._stream = None
+
+    def __enter__(self):
+        if self.output is None:
+            self.output = StringIO()
+
+        if isinstance(self.output, str):
+            self._stream = open(self.output, "w")
+        else:
+            self._stream = self.output
+        sys.stdout = StreamTee(sys.stdout, [self])
+        sys.stderr = StreamTee(sys.stderr, [self])
+        return self
+
+    def write(self, message):
+        self._stream.write(message)
+
+    def flush(self):
+        self._stream.flush()
+
+    @property
+    def stream(self):
+        if isinstance(self.output, str):
+            return open(self.output, "r")
+        return self._stream
+
+    def list(self):
+        if isinstance(self.output, str):
+            return self.stream.readlines()
+        return (self.stream.getvalue() or "").splitlines()
+
+    def __str__(self):
+        return self.stream.getvalue()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if hasattr(sys.stdout, "instream"):
+            sys.stdout = sys.stdout.instream
+        if hasattr(sys.stderr, "instream"):
+            sys.stderr = sys.stderr.instream
+        self._stream.close()
+        return False
+
+
+####################################################################################################
+# Logging formatting
+####################################################################################################
+class ColoredFormatter:
+    COLORS = {
+        "black": "\u001b[30m",
+        "red": "\u001b[31m",
+        "green": "\u001b[32m",
+        "yellow": "\u001b[33m",
+        "blue": "\u001b[34m",
+        "magenta": "\u001b[35m",
+        "cyan": "\u001b[36m",
+        "white": "\u001b[37m",
+        "bold": "\u001b[1m",
+        "italic": "\u001b[3m",
+        "reset": "\u001b[0m",
+    }
+
+    @classmethod
+    def get_color(cls, color: str):
+        return cls.COLORS.get(color, "")
+
+
+class ServerLogsFormatter:
+    def __init__(self, name: str = None):
+        self.name = name
+        self.start_color = ColoredFormatter.get_color("cyan")
+        self.reset_color = ColoredFormatter.get_color("reset")
+
+
+def initialize_k8s_clients():
+    """Initialize Kubernetes API clients."""
+    load_kubeconfig()
+    return (
+        client.CoreV1Api(),
+        client.CustomObjectsApi(),
+        client.AppsV1Api(),
+    )
+
+
+def string_to_dict(value):
+    try:
+        result = json.loads(value or "{}")
+        return result if isinstance(result, dict) else {}
+    except (json.JSONDecodeError, TypeError):
+        return {}