skypilot-nightly 1.0.0.dev20250624__py3-none-any.whl → 1.0.0.dev20250625__py3-none-any.whl

sky/provision/kubernetes/utils.py

@@ -10,7 +10,7 @@ import shutil
 import subprocess
 import time
 import typing
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 from urllib.parse import urlparse
 
 import sky
@@ -2734,6 +2734,21 @@ def get_kubernetes_node_info(
                     node.metadata.labels.get(label_key))
                 break
 
+        # Extract IP address from node addresses (prefer external, fallback to internal)
+        node_ip = None
+        if node.status.addresses:
+            # First try to find external IP
+            for address in node.status.addresses:
+                if address.type == 'ExternalIP':
+                    node_ip = address.address
+                    break
+            # If no external IP, try to find internal IP
+            if node_ip is None:
+                for address in node.status.addresses:
+                    if address.type == 'InternalIP':
+                        node_ip = address.address
+                        break
+
         allocated_qty = 0
         accelerator_count = get_node_accelerator_count(node.status.allocatable)
 
@@ -2765,7 +2780,8 @@ def get_kubernetes_node_info(
             name=node.metadata.name,
             accelerator_type=accelerator_name,
             total={'accelerator_count': int(accelerator_count)},
-            free={'accelerators_available': int(accelerators_available)})
+            free={'accelerators_available': int(accelerators_available)},
+            ip_address=node_ip)
     hint = ''
     if has_multi_host_tpu:
         hint = ('(Note: Multi-host TPUs are detected and excluded from the '
@@ -3281,3 +3297,37 @@ def format_kubeconfig_exec_auth_with_cache(kubeconfig_path: str) -> str:
 
     format_kubeconfig_exec_auth(config, path)
     return path
+
+
+def delete_k8s_resource_with_retry(delete_func: Callable, resource_type: str,
+                                   resource_name: str) -> None:
+    """Helper to delete Kubernetes resources with 404 handling and retries.
+
+    Args:
+        delete_func: Function to call to delete the resource
+        resource_type: Type of resource being deleted (e.g. 'service'),
+            used in logging
+        resource_name: Name of the resource being deleted, used in logging
+    """
+    max_retries = 3
+    retry_delay = 5  # seconds
+
+    for attempt in range(max_retries):
+        try:
+            delete_func()
+            return
+        except kubernetes.api_exception() as e:
+            if e.status == 404:
+                logger.warning(
+                    f'terminate_instances: Tried to delete {resource_type} '
+                    f'{resource_name}, but the {resource_type} was not '
+                    'found (404).')
+                return
+            elif attempt < max_retries - 1:
+                logger.warning(f'terminate_instances: Failed to delete '
+                               f'{resource_type} {resource_name} (attempt '
+                               f'{attempt + 1}/{max_retries}). Error: {e}. '
+                               f'Retrying in {retry_delay} seconds...')
+                time.sleep(retry_delay)
+            else:
+                raise

sky/provision/kubernetes/volume.py

@@ -0,0 +1,147 @@
+"""Kubernetes pvc provisioning."""
+from typing import Any, Dict, List, Optional, Tuple
+
+from sky import models
+from sky import sky_logging
+from sky.adaptors import kubernetes
+from sky.provision.kubernetes import config as config_lib
+from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.volumes import volume as volume_lib
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _get_context_namespace(config: models.VolumeConfig) -> Tuple[str, str]:
+    """Gets the context and namespace of a volume."""
+    if config.region is None:
+        context = kubernetes_utils.get_current_kube_config_context_name()
+        config.region = context
+    else:
+        context = config.region
+    namespace = config.config.get('namespace')
+    if namespace is None:
+        namespace = kubernetes_utils.get_kube_config_context_namespace(context)
+        config.config['namespace'] = namespace
+    return context, namespace
+
+
+def check_pvc_usage_for_pod(context: Optional[str], namespace: str,
+                            pod_spec: Dict[str, Any]) -> None:
+    """Checks if the PVC is used by any pod in the namespace."""
+    volumes = pod_spec.get('spec', {}).get('volumes', [])
+    if not volumes:
+        return
+    once_modes = [
+        volume_lib.VolumeAccessMode.READ_WRITE_ONCE.value,
+        volume_lib.VolumeAccessMode.READ_WRITE_ONCE_POD.value
+    ]
+    for volume in volumes:
+        pvc_name = volume.get('persistentVolumeClaim', {}).get('claimName')
+        if not pvc_name:
+            continue
+        pvc = kubernetes.core_api(
+            context).read_namespaced_persistent_volume_claim(
+                name=pvc_name, namespace=namespace)
+        access_mode = pvc.spec.access_modes[0]
+        if access_mode not in once_modes:
+            continue
+        usedby = _get_volume_usedby(context, namespace, pvc_name)
+        if usedby:
+            raise config_lib.KubernetesError(f'Volume {pvc_name} with access '
+                                             f'mode {access_mode} is already '
+                                             f'in use by {usedby}.')
+
+
+def apply_volume(config: models.VolumeConfig) -> models.VolumeConfig:
+    """Creates or registers a volume."""
+    context, namespace = _get_context_namespace(config)
+    pvc_spec = _get_pvc_spec(namespace, config)
+    create_persistent_volume_claim(namespace, context, pvc_spec)
+    return config
+
+
+def delete_volume(config: models.VolumeConfig) -> models.VolumeConfig:
+    """Deletes a volume."""
+    context, namespace = _get_context_namespace(config)
+    pvc_name = config.name_on_cloud
+    logger.info(f'Deleting PVC {pvc_name}')
+    kubernetes_utils.delete_k8s_resource_with_retry(
+        delete_func=lambda pvc_name=pvc_name: kubernetes.core_api(
+            context).delete_namespaced_persistent_volume_claim(
+                name=pvc_name,
+                namespace=namespace,
+                _request_timeout=config_lib.DELETION_TIMEOUT),
+        resource_type='pvc',
+        resource_name=pvc_name)
+    return config
+
+
+def _get_volume_usedby(context: Optional[str], namespace: str,
+                       pvc_name: str) -> List[str]:
+    """Gets the usedby resources of a volume."""
+    usedby = []
+    # Get all pods in the namespace
+    pods = kubernetes.core_api(context).list_namespaced_pod(namespace=namespace)
+    for pod in pods.items:
+        if pod.spec.volumes is not None:
+            for volume in pod.spec.volumes:
+                if volume.persistent_volume_claim is not None:
+                    if volume.persistent_volume_claim.claim_name == pvc_name:
+                        usedby.append(pod.metadata.name)
+    return usedby
+
+
+def get_volume_usedby(config: models.VolumeConfig) -> List[str]:
+    """Gets the usedby resources of a volume."""
+    context, namespace = _get_context_namespace(config)
+    pvc_name = config.name_on_cloud
+    return _get_volume_usedby(context, namespace, pvc_name)
+
+
+def create_persistent_volume_claim(namespace: str, context: Optional[str],
+                                   pvc_spec: Dict[str, Any]) -> None:
+    """Creates a persistent volume claim for SkyServe controller."""
+    pvc_name = pvc_spec['metadata']['name']
+    try:
+        kubernetes.core_api(context).read_namespaced_persistent_volume_claim(
+            name=pvc_name, namespace=namespace)
+        logger.debug(f'PVC {pvc_name} already exists')
+        return
+    except kubernetes.api_exception() as e:
+        if e.status != 404:  # Not found
+            raise
+    logger.info(f'Creating PVC {pvc_name}')
+    kubernetes.core_api(context).create_namespaced_persistent_volume_claim(
+        namespace=namespace, body=pvc_spec)
+
+
+def _get_pvc_spec(namespace: str,
+                  config: models.VolumeConfig) -> Dict[str, Any]:
+    """Gets the PVC spec for the given storage config."""
+    access_mode = config.config.get('access_mode')
+    size = config.size
+    # The previous code assumes that the access_mode and size are always set.
+    assert access_mode is not None
+    assert size is not None
+    pvc_spec: Dict[str, Any] = {
+        'metadata': {
+            'name': config.name_on_cloud,
+            'namespace': namespace,
+            'labels': {
+                'parent': 'skypilot',
+                'skypilot-name': config.name,
+            }
+        },
+        'spec': {
+            'accessModes': [access_mode],
+            'resources': {
+                'requests': {
+                    'storage': f'{size}Gi'
+                }
+            },
+        }
+    }
+    storage_class = config.config.get('storage_class_name')
+    if storage_class is not None:
+        pvc_spec['spec']['storageClassName'] = storage_class
+    return pvc_spec

sky/resources.py

@@ -30,6 +30,9 @@ from sky.utils import resources_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
 
+if typing.TYPE_CHECKING:
+    from sky.volumes import volume as volume_lib
+
 logger = sky_logging.init_logger(__name__)
 
 _DEFAULT_DISK_SIZE_GB = 256
@@ -289,7 +292,8 @@ class Resources:
                 self._job_recovery = job_recovery
 
         if disk_size is not None:
-            self._disk_size = int(parse_memory_resource(disk_size, 'disk_size'))
+            self._disk_size = int(
+                resources_utils.parse_memory_resource(disk_size, 'disk_size'))
         else:
             self._disk_size = _DEFAULT_DISK_SIZE_GB
 
@@ -707,11 +711,11 @@ class Resources:
             self._memory = None
             return
 
-        memory = parse_memory_resource(str(memory),
-                                       'memory',
-                                       ret_type=float,
-                                       allow_plus=True,
-                                       allow_x=True)
+        memory = resources_utils.parse_memory_resource(str(memory),
+                                                       'memory',
+                                                       ret_type=float,
+                                                       allow_plus=True,
+                                                       allow_x=True)
         self._memory = memory
         if memory.endswith(('+', 'x')):
             # 'x' is used internally for make sure our resources used by
@@ -1465,11 +1469,15 @@ class Resources:
     def get_spot_str(self) -> str:
         return '[Spot]' if self.use_spot else ''
 
-    def make_deploy_variables(self, cluster_name: resources_utils.ClusterName,
-                              region: clouds.Region,
-                              zones: Optional[List[clouds.Zone]],
-                              num_nodes: int,
-                              dryrun: bool) -> Dict[str, Optional[str]]:
+    def make_deploy_variables(
+        self,
+        cluster_name: resources_utils.ClusterName,
+        region: clouds.Region,
+        zones: Optional[List[clouds.Zone]],
+        num_nodes: int,
+        dryrun: bool,
+        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
+    ) -> Dict[str, Optional[str]]:
         """Converts planned sky.Resources to resource variables.
 
         These variables are divided into two categories: cloud-specific and
@@ -1491,7 +1499,7 @@ class Resources:
         # Cloud specific variables
         assert self.cloud is not None, 'Cloud must be specified'
         cloud_specific_variables = self.cloud.make_deploy_resources_variables(
-            self, cluster_name, region, zones, num_nodes, dryrun)
+            self, cluster_name, region, zones, num_nodes, dryrun, volume_mounts)
 
         # TODO(andyl): Should we print some warnings if users' envs share
         # same names with the cloud specific variables, but not enabled
@@ -2291,67 +2299,3 @@ def parse_time_minutes(time: str) -> int:
                 continue
 
     raise ValueError(f'Invalid time format: {time}')
-
-
-def parse_memory_resource(resource_qty_str: Union[str, int, float],
-                          field_name: str,
-                          ret_type: type = int,
-                          unit: str = 'gb',
-                          allow_plus: bool = False,
-                          allow_x: bool = False,
-                          allow_rounding: bool = False) -> str:
-    """Returns memory size in chosen units given a resource quantity string.
-
-    Args:
-        resource_qty_str: Resource quantity string
-        unit: Unit to convert to
-        allow_plus: Whether to allow '+' prefix
-        allow_x: Whether to allow 'x' suffix
-    """
-    assert unit in constants.MEMORY_SIZE_UNITS, f'Invalid unit: {unit}'
-
-    error_msg = f'"{field_name}" field should be a <int><b|k|m|g|t|p><+?>,'\
-                f' got {resource_qty_str}'
-
-    resource_str = str(resource_qty_str)
-
-    # Handle plus and x suffixes, x is only used internally for jobs controller
-    plus = ''
-    if resource_str.endswith('+'):
-        if allow_plus:
-            resource_str = resource_str[:-1]
-            plus = '+'
-        else:
-            raise ValueError(error_msg)
-
-    x = ''
-    if resource_str.endswith('x'):
-        if allow_x:
-            resource_str = resource_str[:-1]
-            x = 'x'
-        else:
-            raise ValueError(error_msg)
-
-    try:
-        # We assume it is already in the wanted units to maintain backwards
-        # compatibility
-        ret_type(resource_str)
-        return f'{resource_str}{plus}{x}'
-    except ValueError:
-        pass
-
-    resource_str = resource_str.lower()
-    for mem_unit, multiplier in constants.MEMORY_SIZE_UNITS.items():
-        if resource_str.endswith(mem_unit):
-            try:
-                value = ret_type(resource_str[:-len(mem_unit)])
-                converted = (value * multiplier /
-                             constants.MEMORY_SIZE_UNITS[unit])
-                if not allow_rounding and ret_type(converted) != converted:
-                    raise ValueError(error_msg)
-                converted = ret_type(converted)
-                return f'{converted}{plus}{x}'
-            except ValueError:
-                continue
-
-    raise ValueError(error_msg)

sky/serve/client/sdk.py

@@ -5,9 +5,9 @@ from typing import List, Optional, Union
 
 import click
 
-from sky.adaptors import common as adaptors_common
 from sky.client import common as client_common
 from sky.server import common as server_common
+from sky.server import rest
 from sky.server.requests import payloads
 from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
@@ -17,12 +17,8 @@ from sky.utils import dag_utils
 if typing.TYPE_CHECKING:
     import io
 
-    import requests
-
     import sky
     from sky.serve import serve_utils
-else:
-    requests = adaptors_common.LazyImport('requests')
 
 
 @context.contextual
@@ -78,7 +74,7 @@ def up(
             task=dag_str,
             service_name=service_name,
         )
-        response = requests.post(
+        response = rest.post(
             f'{server_common.get_server_url()}/serve/up',
             json=json.loads(body.model_dump_json()),
             timeout=(5, None),
@@ -140,7 +136,7 @@ def update(
             mode=mode,
         )
 
-        response = requests.post(
+        response = rest.post(
             f'{server_common.get_server_url()}/serve/update',
             json=json.loads(body.model_dump_json()),
             timeout=(5, None),
@@ -182,7 +178,7 @@ def down(
         all=all,
         purge=purge,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/down',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -217,7 +213,7 @@ def terminate_replica(service_name: str, replica_id: int,
         replica_id=replica_id,
         purge=purge,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/terminate-replica',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -290,7 +286,7 @@ def status(
         exceptions.ClusterNotUpError: if the sky serve controller is not up.
     """
     body = payloads.ServeStatusBody(service_names=service_names,)
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/status',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -301,6 +297,7 @@ def status(
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
+@rest.retry_on_server_unavailable()
 def tail_logs(service_name: str,
               target: Union[str, 'serve_utils.ServiceComponent'],
               replica_id: Optional[int] = None,
@@ -376,7 +373,7 @@ def tail_logs(service_name: str,
         replica_id=replica_id,
         follow=follow,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/logs',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -384,7 +381,10 @@ def tail_logs(service_name: str,
         cookies=server_common.get_api_cookie_jar(),
     )
     request_id = server_common.get_request_id(response)
-    sdk.stream_response(request_id, response, output_stream)
+    return sdk.stream_response(request_id=request_id,
+                               response=response,
+                               output_stream=output_stream,
+                               resumable=True)
 
 
 @usage_lib.entrypoint
@@ -436,7 +436,7 @@ def sync_down_logs(service_name: str,
         targets=targets,
         replica_ids=replica_ids,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/serve/sync-down-logs',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),

sky/serve/server/core.py

@@ -28,6 +28,7 @@ from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
 from sky.utils import controller_utils
+from sky.utils import dag_utils
 from sky.utils import rich_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
@@ -139,10 +140,13 @@ def up(
                              f'{constants.CLUSTER_NAME_VALID_REGEX}')
 
     serve_utils.validate_service_task(task)
+    dag = dag_utils.convert_entrypoint_to_dag(task)
+    dag.resolve_and_validate_volumes()
     # Always apply the policy again here, even though it might have been applied
     # in the CLI. This is to ensure that we apply the policy to the final DAG
     # and get the mutated config.
-    dag, mutated_user_config = admin_policy_utils.apply(task)
+    dag, mutated_user_config = admin_policy_utils.apply(dag)
+    dag.pre_mount_volumes()
     task = dag.tasks[0]
 
     with rich_utils.safe_status(

sky/server/common.py

@@ -9,11 +9,13 @@ import json
 import os
 import pathlib
 import re
+import shutil
 import subprocess
 import sys
+import tempfile
 import time
 import typing
-from typing import Any, Dict, Literal, Optional, Tuple
+from typing import Any, Dict, Literal, Optional, Tuple, Union
 from urllib import parse
 import uuid
 
@@ -27,6 +29,7 @@ from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.data import data_utils
 from sky.server import constants as server_constants
+from sky.server import rest
 from sky.skylet import constants
 from sky.usage import usage_lib
 from sky.utils import annotations
@@ -240,9 +243,9 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
     server_url = endpoint if endpoint is not None else get_server_url()
     while time_out_try_count <= RETRY_COUNT_ON_TIMEOUT:
         try:
-            response = requests.get(f'{server_url}/api/health',
-                                    timeout=2.5,
-                                    cookies=get_api_cookie_jar())
+            response = rest.get(f'{server_url}/api/health',
+                                timeout=2.5,
+                                cookies=get_api_cookie_jar())
         except requests.exceptions.Timeout:
             if time_out_try_count == RETRY_COUNT_ON_TIMEOUT:
                 return ApiServerInfo(status=ApiServerStatus.UNHEALTHY)
@@ -327,6 +330,8 @@ def get_request_id(response: 'requests.Response') -> RequestId:
 def _start_api_server(deploy: bool = False,
                       host: str = '127.0.0.1',
                       foreground: bool = False,
+                      metrics: bool = False,
+                      metrics_port: Optional[int] = None,
                       enable_basic_auth: bool = False):
     """Starts a SkyPilot API server locally."""
     server_url = get_server_url(host)
@@ -357,10 +362,13 @@ def _start_api_server(deploy: bool = False,
             args += ['--deploy']
         if host is not None:
             args += [f'--host={host}']
+        if metrics_port is not None:
+            args += [f'--metrics-port={metrics_port}']
 
         if foreground:
             # Replaces the current process with the API server
             os.environ[constants.ENV_VAR_IS_SKYPILOT_SERVER] = 'true'
+            _set_metrics_env_var(os.environ, metrics, deploy)
             if enable_basic_auth:
                 os.environ[constants.ENV_VAR_ENABLE_BASIC_AUTH] = 'true'
             os.execvp(args[0], args)
@@ -368,6 +376,10 @@ def _start_api_server(deploy: bool = False,
         log_path = os.path.expanduser(constants.API_SERVER_LOGS)
         os.makedirs(os.path.dirname(log_path), exist_ok=True)
 
+        # For spawn mode, copy the environ to avoid polluting the SDK process.
+        server_env = os.environ.copy()
+        server_env[constants.ENV_VAR_IS_SKYPILOT_SERVER] = 'true'
+        _set_metrics_env_var(server_env, metrics, deploy)
         # Start the API server process in the background and don't wait for it.
         # If this is called from a CLI invocation, we need
         # start_new_session=True so that SIGINT on the CLI will not also kill
@@ -437,6 +449,26 @@ def _start_api_server(deploy: bool = False,
                 f'SkyPilot API server started. {dashboard_msg}'))
 
 
+def _set_metrics_env_var(env: Union[Dict[str, str], os._Environ], metrics: bool,
+                         deploy: bool):
+    """Sets the metrics environment variables.
+
+    Args:
+        env: The environment variables to set.
+        metrics: Whether to enable metrics.
+        deploy: Whether the server is running in deploy mode, which means
+            multiple processes might be running.
+    """
+    if metrics:
+        env[constants.ENV_VAR_SERVER_METRICS_ENABLED] = 'true'
+        if deploy:
+            metrics_dir = os.path.join(tempfile.gettempdir(), 'metrics')
+            shutil.rmtree(metrics_dir, ignore_errors=True)
+            os.makedirs(metrics_dir, exist_ok=True)
+            # Refer to https://prometheus.github.io/client_python/multiprocess/
+            env['PROMETHEUS_MULTIPROC_DIR'] = metrics_dir
+
+
 def check_server_healthy(
     endpoint: Optional[str] = None
 ) -> Tuple[Literal[
@@ -571,6 +603,8 @@ def get_skypilot_version_on_disk() -> str:
 def check_server_healthy_or_start_fn(deploy: bool = False,
                                      host: str = '127.0.0.1',
                                      foreground: bool = False,
+                                     metrics: bool = False,
+                                     metrics_port: Optional[int] = None,
                                      enable_basic_auth: bool = False):
     api_server_status = None
     try:
@@ -592,7 +626,8 @@ def check_server_healthy_or_start_fn(deploy: bool = False,
             # have started the server while we were waiting for the lock.
             api_server_info = get_api_server_status(endpoint)
             if api_server_info.status == ApiServerStatus.UNHEALTHY:
-                _start_api_server(deploy, host, foreground, enable_basic_auth)
+                _start_api_server(deploy, host, foreground, metrics,
+                                  metrics_port, enable_basic_auth)
 
 
 def check_server_healthy_or_start(func):

sky/server/constants.py

@@ -7,7 +7,7 @@ from sky.skylet import constants
 # API server version, whenever there is a change in API server that requires a
 # restart of the local API server or error out when the client does not match
 # the server version.
-API_VERSION = '9'
+API_VERSION = '10'
 
 # Prefix for API request names.
 REQUEST_NAME_PREFIX = 'sky.'
@@ -22,6 +22,10 @@ API_SERVER_REQUEST_DB_PATH = '~/.sky/api_server/requests.db'
 # background.
 CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS = 60
 
+# The interval (seconds) for the volume status to be refreshed in the
+# background.
+VOLUME_REFRESH_DAEMON_INTERVAL_SECONDS = 60
+
 # Environment variable for a file path to the API cookie file.
 # Keep in sync with websocket_proxy.py
 API_COOKIE_FILE_ENV_VAR = f'{constants.SKYPILOT_ENV_VAR_PREFIX}API_COOKIE_FILE'