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

sky/jobs/client/sdk.py

@@ -7,10 +7,10 @@ import webbrowser
 import click
 
 from sky import sky_logging
-from sky.adaptors import common as adaptors_common
 from sky.client import common as client_common
 from sky.client import sdk
 from sky.server import common as server_common
+from sky.server import rest
 from sky.server.requests import payloads
 from sky.skylet import constants
 from sky.usage import usage_lib
@@ -22,11 +22,7 @@ from sky.utils import dag_utils
 if typing.TYPE_CHECKING:
     import io
 
-    import requests
-
     import sky
-else:
-    requests = adaptors_common.LazyImport('requests')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -86,7 +82,7 @@ def launch(
             task=dag_str,
             name=name,
         )
-        response = requests.post(
+        response = rest.post(
             f'{server_common.get_server_url()}/jobs/launch',
             json=json.loads(body.model_dump_json()),
             timeout=(5, None),
@@ -146,7 +142,7 @@ def queue(refresh: bool,
         all_users=all_users,
         job_ids=job_ids,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/jobs/queue',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -186,7 +182,7 @@ def cancel(
         all=all,
         all_users=all_users,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/jobs/cancel',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),
@@ -197,6 +193,7 @@ def cancel(
 
 @usage_lib.entrypoint
 @server_common.check_server_healthy_or_start
+@rest.retry_on_server_unavailable()
 def tail_logs(name: Optional[str] = None,
               job_id: Optional[int] = None,
               follow: bool = True,
@@ -236,7 +233,7 @@ def tail_logs(name: Optional[str] = None,
         refresh=refresh,
         tail=tail,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/jobs/logs',
         json=json.loads(body.model_dump_json()),
         stream=True,
@@ -244,7 +241,12 @@ def tail_logs(name: Optional[str] = None,
         cookies=server_common.get_api_cookie_jar(),
     )
     request_id = server_common.get_request_id(response)
-    return sdk.stream_response(request_id, response, output_stream)
+    # Log request is idempotent when tail is 0, thus can resume previous
+    # streaming point on retry.
+    return sdk.stream_response(request_id=request_id,
+                               response=response,
+                               output_stream=output_stream,
+                               resumable=(tail == 0))
 
 
 @usage_lib.entrypoint
@@ -281,7 +283,7 @@ def download_logs(
         controller=controller,
         local_dir=local_dir,
     )
-    response = requests.post(
+    response = rest.post(
         f'{server_common.get_server_url()}/jobs/download_logs',
         json=json.loads(body.model_dump_json()),
         timeout=(5, None),

sky/jobs/server/core.py

@@ -145,6 +145,7 @@ def launch(
     entrypoint = task
     dag_uuid = str(uuid.uuid4().hex[:4])
     dag = dag_utils.convert_entrypoint_to_dag(entrypoint)
+    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.
@@ -154,6 +155,9 @@ def launch(
             raise ValueError('Only single-task or chain DAG is '
                              f'allowed for job_launch. Dag: {dag}')
     dag.validate()
+    # TODO(aylei): use consolidated job controller instead of performing
+    # pre-mount operations when submitting jobs.
+    dag.pre_mount_volumes()
 
     user_dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
 

sky/models.py

@@ -6,6 +6,8 @@ import getpass
 import os
 from typing import Any, Dict, Optional
 
+import pydantic
+
 from sky.skylet import constants
 from sky.utils import common_utils
 
@@ -48,6 +50,8 @@ class KubernetesNodeInfo:
     # Resources available on the node. E.g., {'nvidia.com/gpu': '2'}
     total: Dict[str, int]
     free: Dict[str, int]
+    # IP address of the node (external IP preferred, fallback to internal IP)
+    ip_address: Optional[str] = None
 
 
 @dataclasses.dataclass
@@ -76,3 +80,15 @@ class KubernetesNodesInfo:
             },
             hint=data['hint'],
         )
+
+
+class VolumeConfig(pydantic.BaseModel):
+    """Configuration for creating a volume."""
+    name: str
+    type: str
+    cloud: str
+    region: Optional[str]
+    zone: Optional[str]
+    name_on_cloud: str
+    size: Optional[str]
+    config: Dict[str, Any] = {}

sky/provision/__init__.py

@@ -8,6 +8,7 @@ import inspect
 import typing
 from typing import Any, Dict, List, Optional, Type
 
+from sky import models
 from sky import sky_logging
 # These provision.<cloud> modules should never fail even if underlying cloud SDK
 # dependencies are not installed. This is ensured by using sky.adaptors inside
@@ -103,6 +104,31 @@ def bootstrap_instances(
     raise NotImplementedError
 
 
+@_route_to_cloud_impl
+def apply_volume(provider_name: str,
+                 config: models.VolumeConfig) -> models.VolumeConfig:
+    """Create or register a volume.
+
+    This function creates or registers a volume with the provided configuration,
+    and returns a VolumeConfig object with updated configuration.
+    """
+    raise NotImplementedError
+
+
+@_route_to_cloud_impl
+def delete_volume(provider_name: str,
+                  config: models.VolumeConfig) -> models.VolumeConfig:
+    """Delete a volume."""
+    raise NotImplementedError
+
+
+@_route_to_cloud_impl
+def get_volume_usedby(provider_name: str,
+                      config: models.VolumeConfig) -> List[str]:
+    """Get the usedby of a volume."""
+    raise NotImplementedError
+
+
 @_route_to_cloud_impl
 def run_instances(provider_name: str, region: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:

sky/provision/kubernetes/__init__.py

@@ -11,3 +11,6 @@ from sky.provision.kubernetes.instance import wait_instances
 from sky.provision.kubernetes.network import cleanup_ports
 from sky.provision.kubernetes.network import open_ports
 from sky.provision.kubernetes.network import query_ports
+from sky.provision.kubernetes.volume import apply_volume
+from sky.provision.kubernetes.volume import delete_volume
+from sky.provision.kubernetes.volume import get_volume_usedby

sky/provision/kubernetes/instance.py

@@ -3,7 +3,6 @@ import copy
 import json
 import time
 from typing import Any, Callable, Dict, List, Optional, Union
-import uuid
 
 from sky import exceptions
 from sky import sky_logging
@@ -15,6 +14,7 @@ from sky.provision import docker_utils
 from sky.provision.kubernetes import config as config_lib
 from sky.provision.kubernetes import network_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.provision.kubernetes import volume
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import config_utils
@@ -240,7 +240,7 @@ def _raise_pod_scheduling_errors(namespace, context, new_nodes):
                                                        extra_msg,
                                                        details=event_message))
             raise config_lib.KubernetesError(f'{timeout_err_msg} '
-                                             f'Pod status: {pod_status}'
+                                             f'Pod status: {pod_status} '
                                              f'Details: \'{event_message}\' ')
     raise config_lib.KubernetesError(f'{timeout_err_msg}')
 
@@ -673,21 +673,6 @@ def _create_namespaced_pod_with_retries(namespace: str, pod_spec: dict,
             raise e
 
 
-def _create_persistent_volume_claim(namespace: str, context: Optional[str],
-                                    pvc_spec: Dict[str, Any]) -> None:
-    """Creates a persistent volume claim for SkyServe controller."""
-    try:
-        kubernetes.core_api(context).read_namespaced_persistent_volume_claim(
-            name=pvc_spec['metadata']['name'], namespace=namespace)
-        return
-    except kubernetes.api_exception() as e:
-        if e.status != 404:  # Not found
-            raise
-
-    kubernetes.core_api(context).create_namespaced_persistent_volume_claim(
-        namespace=namespace, body=pvc_spec)
-
-
 @timeline.event
 def _wait_for_deployment_pod(context,
                              namespace,
@@ -832,9 +817,12 @@ def _create_pods(region: str, cluster_name_on_cloud: str,
             # Worker pods
             pod_spec_copy['metadata']['labels'].update(
                 constants.WORKER_NODE_TAGS)
-            pod_uuid = str(uuid.uuid4())[:6]
-            pod_name = f'{cluster_name_on_cloud}-{pod_uuid}'
-            pod_spec_copy['metadata']['name'] = f'{pod_name}-worker'
+            pod_name = f'{cluster_name_on_cloud}-worker{i}'
+            if pod_name in running_pods:
+                # If the pod is already running, we skip creating it.
+                return
+            pod_spec_copy['metadata']['name'] = pod_name
+            pod_spec_copy['metadata']['labels']['component'] = pod_name
             # For multi-node support, we put a soft-constraint to schedule
             # worker pods on different nodes than the head pod.
             # This is not set as a hard constraint because if different nodes
@@ -888,7 +876,7 @@ def _create_pods(region: str, cluster_name_on_cloud: str,
             ]
 
         if to_create_deployment:
-            _create_persistent_volume_claim(namespace, context, pvc_spec)
+            volume.create_persistent_volume_claim(namespace, context, pvc_spec)
 
             # It's safe to directly modify the template spec in the deployment spec
             # because controller pod is singleton, i in [0].
@@ -910,6 +898,10 @@ def _create_pods(region: str, cluster_name_on_cloud: str,
                 print('Deployment failed', e)
                 raise e
 
+        # Check if any PVCs with access mode ReadWriteOnce or ReadWriteOncePod
+        # is used by any pod in the namespace.
+        volume.check_pvc_usage_for_pod(context, namespace, pod_spec_copy)
+
         return _create_namespaced_pod_with_retries(namespace, pod_spec_copy,
                                                    context)
 
@@ -1012,40 +1004,6 @@ def stop_instances(
     raise NotImplementedError()
 
 
-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
-
-
 def _delete_services(name_prefix: str, namespace: str,
                      context: Optional[str]) -> None:
     """Delete services with the given name prefix.
@@ -1061,13 +1019,14 @@ def _delete_services(name_prefix: str, namespace: str,
         # TODO(andyl): Wait for
         # https://github.com/pylint-dev/pylint/issues/5263.
         # pylint: disable=cell-var-from-loop
-        _delete_k8s_resource_with_retry(delete_func=lambda: kubernetes.core_api(
-            context).delete_namespaced_service(name=service_name,
-                                               namespace=namespace,
-                                               _request_timeout=config_lib.
-                                               DELETION_TIMEOUT),
-                                        resource_type='service',
-                                        resource_name=service_name)
+        kubernetes_utils.delete_k8s_resource_with_retry(
+            delete_func=lambda: kubernetes.core_api(
+                context).delete_namespaced_service(name=service_name,
+                                                   namespace=namespace,
+                                                   _request_timeout=config_lib.
+                                                   DELETION_TIMEOUT),
+            resource_type='service',
+            resource_name=service_name)
 
 
 def _terminate_node(namespace: str,
@@ -1087,7 +1046,7 @@ def _terminate_node(namespace: str,
     # from within the pod, e.g., for autodown.
     # Note - some misbehaving pods may not terminate gracefully if they have
     # open file descriptors. We force delete pods to avoid this.
-    _delete_k8s_resource_with_retry(
+    kubernetes_utils.delete_k8s_resource_with_retry(
         delete_func=lambda: kubernetes.core_api(context).delete_namespaced_pod(
             name=pod_name,
             namespace=namespace,
@@ -1105,26 +1064,28 @@ def _terminate_deployment(cluster_name: str, namespace: str,
 
     # Delete deployment
     deployment_name = _get_deployment_name(cluster_name)
-    _delete_k8s_resource_with_retry(delete_func=lambda: kubernetes.apps_api(
-        context).delete_namespaced_deployment(name=deployment_name,
-                                              namespace=namespace,
-                                              _request_timeout=config_lib.
-                                              DELETION_TIMEOUT),
-                                    resource_type='deployment',
-                                    resource_name=deployment_name)
+    kubernetes_utils.delete_k8s_resource_with_retry(
+        delete_func=lambda: kubernetes.apps_api(
+            context).delete_namespaced_deployment(name=deployment_name,
+                                                  namespace=namespace,
+                                                  _request_timeout=config_lib.
+                                                  DELETION_TIMEOUT),
+        resource_type='deployment',
+        resource_name=deployment_name)
 
     # Delete PVCs
     pvc_name = _get_pvc_name(
         cluster_name,
         kubernetes_utils.HIGH_AVAILABILITY_DEPLOYMENT_VOLUME_MOUNT_NAME)
     # pylint: disable=cell-var-from-loop
-    _delete_k8s_resource_with_retry(delete_func=lambda: 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)
+    kubernetes_utils.delete_k8s_resource_with_retry(
+        delete_func=lambda: 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)
 
 
 def terminate_instances(

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
@@ -1625,9 +1625,15 @@ def _get_kubeconfig_text_for_context(context: Optional[str] = None) -> str:
     command = 'kubectl config view --minify'
     if context is not None:
         command += f' --context={context}'
+
+    # Ensure subprocess inherits the current environment properly
+    # This fixes the issue where kubectl can't find ~/.kube/config in API server context
+    env = os.environ.copy()
+
     proc = subprocess.run(command,
                           shell=True,
                           check=False,
+                          env=env,
                           stdout=subprocess.PIPE,
                           stderr=subprocess.PIPE)
     if proc.returncode != 0:
@@ -2734,6 +2740,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 +2786,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 '
@@ -3279,5 +3301,49 @@ def format_kubeconfig_exec_auth_with_cache(kubeconfig_path: str) -> str:
     if os.path.isfile(path):
         return path
 
-    format_kubeconfig_exec_auth(config, path)
-    return path
+    try:
+        format_kubeconfig_exec_auth(config, path)
+        return path
+    except Exception as e:  # pylint: disable=broad-except
+        # There may be problems with kubeconfig, but the user is not actually
+        # using Kubernetes (or SSH Node Pools)
+        logger.warning(
+            f'Failed to format kubeconfig at {kubeconfig_path}. '
+            'Please check if the kubeconfig is valid. This may cause '
+            'problems when Kubernetes infra is used. '
+            f'Reason: {common_utils.format_exception(e)}')
+        return kubeconfig_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