skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/provision/kubernetes/volume.py

@@ -0,0 +1,282 @@
+"""Kubernetes pvc provisioning."""
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+from sky import global_user_state
+from sky import models
+from sky import sky_logging
+from sky.adaptors import kubernetes
+from sky.provision import constants
+from sky.provision.kubernetes import config as config_lib
+from sky.provision.kubernetes import constants as k8s_constants
+from sky.provision.kubernetes import utils as kubernetes_utils
+from sky.utils 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_pods, _ = _get_volume_usedby(context, namespace, pvc_name)
+        if usedby_pods:
+            raise config_lib.KubernetesError(f'Volume {pvc_name} with access '
+                                             f'mode {access_mode} is already '
+                                             f'in use by Pods {usedby_pods}.')
+
+
+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)
+    # Check if the storage class exists
+    storage_class_name = pvc_spec['spec'].get('storageClassName')
+    if storage_class_name is not None:
+        try:
+            kubernetes.storage_api(context).read_storage_class(
+                name=storage_class_name)
+        except kubernetes.api_exception() as e:
+            raise config_lib.KubernetesError(
+                f'Check storage class {storage_class_name} error: {e}')
+    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
+    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)
+    logger.info(f'Deleted PVC {pvc_name} in namespace {namespace}')
+    return config
+
+
+def _get_volume_usedby(
+    context: Optional[str],
+    namespace: str,
+    pvc_name: str,
+) -> Tuple[List[str], List[str]]:
+    """Gets the usedby resources of a volume.
+
+    This function returns the pods and clusters that are using the volume.
+    The usedby_pods is accurate, which also includes the Pods that are not
+    managed by SkyPilot.
+
+    Args:
+        context: Kubernetes context
+        namespace: Kubernetes namespace
+        pvc_name: PVC name
+
+    Returns:
+        usedby_pods: List of pods using the volume. These may include pods
+                     not created by SkyPilot.
+        usedby_clusters: List of clusters using the volume.
+    """
+    usedby_pods = []
+    usedby_clusters = []
+    field_selector = ','.join([
+        f'status.phase!={phase}'
+        for phase in k8s_constants.PVC_NOT_HOLD_POD_PHASES
+    ])
+    cloud_to_name_map = _get_cluster_name_on_cloud_to_cluster_name_map()
+    # Get all pods in the namespace
+    pods = kubernetes.core_api(context).list_namespaced_pod(
+        namespace=namespace, field_selector=field_selector)
+    for pod in pods.items:
+        if pod.spec.volumes is None:
+            continue
+        for volume in pod.spec.volumes:
+            if volume.persistent_volume_claim is None:
+                continue
+            if volume.persistent_volume_claim.claim_name == pvc_name:
+                usedby_pods.append(pod.metadata.name)
+                # Get the real cluster name
+                cluster_name_on_cloud = pod.metadata.labels.get(
+                    constants.TAG_SKYPILOT_CLUSTER_NAME)
+                if cluster_name_on_cloud is None:
+                    continue
+                cluster_name = cloud_to_name_map.get(cluster_name_on_cloud)
+                if cluster_name is not None:
+                    usedby_clusters.append(cluster_name)
+    if usedby_pods:
+        logger.debug(f'Volume {pvc_name} is used by Pods {usedby_pods}'
+                     f' and clusters {usedby_clusters}')
+    return usedby_pods, usedby_clusters
+
+
+def _get_cluster_name_on_cloud_to_cluster_name_map() -> Dict[str, str]:
+    """Gets the map from cluster name on cloud to cluster name."""
+    clusters = global_user_state.get_clusters()
+    cloud_to_name_map = {}
+    for cluster in clusters:
+        handle = cluster['handle']
+        if handle is None:
+            continue
+        cloud_to_name_map[handle.cluster_name_on_cloud] = cluster['name']
+    return cloud_to_name_map
+
+
+def get_volume_usedby(
+    config: models.VolumeConfig,) -> Tuple[List[str], 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 get_all_volumes_usedby(
+    configs: List[models.VolumeConfig],
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Gets the usedby resources of all volumes."""
+    field_selector = ','.join([
+        f'status.phase!={phase}'
+        for phase in k8s_constants.PVC_NOT_HOLD_POD_PHASES
+    ])
+    label_selector = 'parent=skypilot'
+    context_to_namespaces: Dict[str, Set[str]] = {}
+    pvc_names = set()
+    for config in configs:
+        context, namespace = _get_context_namespace(config)
+        if context not in context_to_namespaces:
+            context_to_namespaces[context] = set()
+        context_to_namespaces[context].add(namespace)
+        pvc_names.add(config.name_on_cloud)
+    cloud_to_name_map = _get_cluster_name_on_cloud_to_cluster_name_map()
+    # Get all pods in the namespace
+    used_by_pods: Dict[str, Dict[str, Dict[str, List[str]]]] = {}
+    used_by_clusters: Dict[str, Dict[str, Dict[str, List[str]]]] = {}
+    for context, namespaces in context_to_namespaces.items():
+        used_by_pods[context] = {}
+        used_by_clusters[context] = {}
+        for namespace in namespaces:
+            used_by_pods[context][namespace] = {}
+            used_by_clusters[context][namespace] = {}
+            pods = kubernetes.core_api(context).list_namespaced_pod(
+                namespace=namespace,
+                field_selector=field_selector,
+                label_selector=label_selector)
+            for pod in pods.items:
+                if pod.spec.volumes is None:
+                    continue
+                for volume in pod.spec.volumes:
+                    if volume.persistent_volume_claim is None:
+                        continue
+                    volume_name = volume.persistent_volume_claim.claim_name
+                    if volume_name not in pvc_names:
+                        continue
+                    if volume_name not in used_by_pods[context][namespace]:
+                        used_by_pods[context][namespace][volume_name] = []
+                    used_by_pods[context][namespace][volume_name].append(
+                        pod.metadata.name)
+                    cluster_name_on_cloud = pod.metadata.labels.get(
+                        constants.TAG_SKYPILOT_CLUSTER_NAME)
+                    if cluster_name_on_cloud is None:
+                        continue
+                    cluster_name = cloud_to_name_map.get(cluster_name_on_cloud)
+                    if cluster_name is None:
+                        continue
+                    if cluster_name not in used_by_clusters[context][namespace]:
+                        used_by_clusters[context][namespace][cluster_name] = []
+                    used_by_clusters[context][namespace][cluster_name].append(
+                        cluster_name)
+    return used_by_pods, used_by_clusters
+
+
+def map_all_volumes_usedby(
+        used_by_pods: Dict[str, Any], used_by_clusters: Dict[str, Any],
+        config: models.VolumeConfig) -> Tuple[List[str], List[str]]:
+    """Maps the usedby resources of a volume."""
+    context, namespace = _get_context_namespace(config)
+    pvc_name = config.name_on_cloud
+
+    return (used_by_pods.get(context, {}).get(namespace, {}).get(pvc_name, []),
+            used_by_clusters.get(context, {}).get(namespace,
+                                                  {}).get(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
+    kubernetes.core_api(context).create_namespaced_persistent_volume_claim(
+        namespace=namespace, body=pvc_spec)
+    logger.info(f'Created PVC {pvc_name} in namespace {namespace}')
+
+
+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'
+                }
+            },
+        }
+    }
+    if config.labels:
+        pvc_spec['metadata']['labels'].update(config.labels)
+    storage_class = config.config.get('storage_class_name')
+    if storage_class is not None:
+        pvc_spec['spec']['storageClassName'] = storage_class
+    return pvc_spec

sky/provision/lambda_cloud/instance.py

@@ -1,7 +1,7 @@
 """Lambda Cloud instance provisioning."""
 
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from sky import sky_logging
 from sky.provision import common
@@ -68,9 +68,10 @@ def _get_private_ip(instance_info: Dict[str, Any], single_node: bool) -> str:
     return private_ip
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster"""
+    del cluster_name  # unused
     lambda_client = _get_lambda_client()
     pending_status = ['booting']
     while True:
@@ -106,34 +107,35 @@ def run_instances(region: str, cluster_name_on_cloud: str,
     created_instance_ids = []
     remote_ssh_key_name = config.authentication_config['remote_key_name']
 
-    def launch_nodes(node_type: str, quantity: int) -> List[str]:
+    def launch_node(node_type: str) -> str:
         try:
             instance_ids = lambda_client.create_instances(
                 instance_type=config.node_config['InstanceType'],
                 region=region,
                 name=f'{cluster_name_on_cloud}-{node_type}',
-                quantity=quantity,
+                # Quantity cannot actually be greater than 1; see:
+                # https://github.com/skypilot-org/skypilot/issues/7084
+                quantity=1,
                 ssh_key_name=remote_ssh_key_name,
             )
-            logger.info(f'Launched {len(instance_ids)} {node_type} node(s), '
-                        f'instance_ids: {instance_ids}')
-            return instance_ids
+            logger.info(f'Launched {node_type} node, '
+                        f'instance_id: {instance_ids[0]}')
+            return instance_ids[0]
         except Exception as e:
             logger.warning(f'run_instances error: {e}')
             raise
 
     if head_instance_id is None:
-        instance_ids = launch_nodes('head', 1)
-        assert len(instance_ids) == 1
-        created_instance_ids.append(instance_ids[0])
-        head_instance_id = instance_ids[0]
+        head_instance_id = launch_node('head')
+        created_instance_ids.append(head_instance_id)
 
     assert head_instance_id is not None, 'head_instance_id should not be None'
 
     worker_node_count = to_start_count - 1
     if worker_node_count > 0:
-        instance_ids = launch_nodes('worker', worker_node_count)
-        created_instance_ids.extend(instance_ids)
+        for _ in range(worker_node_count):
+            worker_instance_id = launch_node('worker')
+            created_instance_ids.append(worker_instance_id)
 
     while True:
         instances = _filter_instances(cluster_name_on_cloud, ['active'])
@@ -226,11 +228,14 @@ def get_cluster_info(
 
 
 def query_instances(
+    cluster_name: str,
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
-) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    retry_if_missing: bool = False,
+) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(cluster_name_on_cloud, None)
 
@@ -240,12 +245,13 @@ def query_instances(
         'unhealthy': status_lib.ClusterStatus.INIT,
         'terminating': None,
     }
-    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    statuses: Dict[str, Tuple[Optional['status_lib.ClusterStatus'],
+                              Optional[str]]] = {}
     for instance_id, instance in instances.items():
         status = status_map.get(instance['status'])
         if non_terminated_only and status is None:
             continue
-        statuses[instance_id] = status
+        statuses[instance_id] = (status, None)
     return statuses
 
 

sky/provision/nebius/constants.py

@@ -0,0 +1,50 @@
+"""Constants used by the Nebius provisioner."""
+
+VERSION = 'v1'
+
+# InfiniBand-capable instance platforms
+INFINIBAND_INSTANCE_PLATFORMS = [
+    'gpu-h100-sxm',
+    'gpu-h200-sxm',
+]
+
+# InfiniBand environment variables for NCCL and UCX
+INFINIBAND_ENV_VARS = {
+    'NCCL_IB_HCA': 'mlx5',
+    'UCX_NET_DEVICES': ('mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,'
+                        'mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1')
+}
+
+# pylint: disable=line-too-long
+INFINIBAND_IMAGE_ID = 'docker:cr.eu-north1.nebius.cloud/nebius-benchmarks/nccl-tests:2.23.4-ubu22.04-cu12.4'
+
+# Docker run options for InfiniBand support
+INFINIBAND_DOCKER_OPTIONS = ['--device=/dev/infiniband', '--cap-add=IPC_LOCK']
+
+# InfiniBand fabric mapping by platform and region
+# Based on Nebius documentation
+INFINIBAND_FABRIC_MAPPING = {
+    # H100 platforms
+    ('gpu-h100-sxm', 'eu-north1'): [
+        'fabric-2', 'fabric-3', 'fabric-4', 'fabric-6'
+    ],
+
+    # H200 platforms
+    ('gpu-h200-sxm', 'eu-north1'): ['fabric-7'],
+    ('gpu-h200-sxm', 'eu-west1'): ['fabric-5'],
+    ('gpu-h200-sxm', 'us-central1'): ['us-central1-a'],
+}
+
+
+def get_default_fabric(platform: str, region: str) -> str:
+    """Get the default (first) fabric for a given platform and region."""
+    fabrics = INFINIBAND_FABRIC_MAPPING.get((platform, region), [])
+    if not fabrics:
+        # Select north europe region as default
+        fabrics = INFINIBAND_FABRIC_MAPPING.get(('gpu-h100-sxm', 'eu-north1'),
+                                                [])
+        if not fabrics:
+            raise ValueError(
+                f'No InfiniBand fabric available for platform {platform} '
+                f'in region {region}')
+    return fabrics[0]

sky/provision/nebius/instance.py

@@ -1,6 +1,6 @@
 """Nebius instance provisioning."""
 import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from sky import sky_logging
 from sky.provision import common
@@ -65,9 +65,10 @@ def _wait_until_no_pending(region: str, cluster_name_on_cloud: str) -> None:
                            f' to be ready.')
 
 
-def run_instances(region: str, cluster_name_on_cloud: str,
+def run_instances(region: str, cluster_name: str, cluster_name_on_cloud: str,
                   config: common.ProvisionConfig) -> common.ProvisionRecord:
     """Runs instances for the given cluster."""
+    del cluster_name  # unused
     _wait_until_no_pending(region, cluster_name_on_cloud)
     running_instances = _filter_instances(region, cluster_name_on_cloud,
                                           ['RUNNING'])
@@ -124,6 +125,7 @@ def run_instances(region: str, cluster_name_on_cloud: str,
         node_type = 'head' if head_instance_id is None else 'worker'
         try:
             platform, preset = config.node_config['InstanceType'].split('_')
+
             instance_id = utils.launch(
                 cluster_name_on_cloud=cluster_name_on_cloud,
                 node_type=node_type,
@@ -132,7 +134,14 @@ def run_instances(region: str, cluster_name_on_cloud: str,
                 region=region,
                 image_family=config.node_config['ImageId'],
                 disk_size=config.node_config['DiskSize'],
-                user_data=config.node_config['UserData'])
+                user_data=config.node_config['UserData'],
+                use_spot=config.node_config['use_spot'],
+                associate_public_ip_address=(
+                    not config.provider_config['use_internal_ips']),
+                use_static_ip_address=config.provider_config.get(
+                    'use_static_ip_address', False),
+                filesystems=config.node_config.get('filesystems', []),
+                network_tier=config.node_config.get('network_tier'))
         except Exception as e:  # pylint: disable=broad-except
             logger.warning(f'run_instances error: {e}')
             raise
@@ -241,11 +250,14 @@ def get_cluster_info(
 
 
 def query_instances(
+    cluster_name: str,
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
-) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    retry_if_missing: bool = False,
+) -> Dict[str, Tuple[Optional['status_lib.ClusterStatus'], Optional[str]]]:
     """See sky/provision/__init__.py"""
+    del cluster_name, retry_if_missing  # unused
     assert provider_config is not None, (cluster_name_on_cloud, provider_config)
     instances = _filter_instances(provider_config['region'],
                                   cluster_name_on_cloud, None)
@@ -257,12 +269,13 @@ def query_instances(
         'STOPPING': status_lib.ClusterStatus.STOPPED,
         'DELETING': status_lib.ClusterStatus.STOPPED,
     }
-    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    statuses: Dict[str, Tuple[Optional['status_lib.ClusterStatus'],
+                              Optional[str]]] = {}
     for inst_id, inst in instances.items():
         status = status_map[inst['status']]
         if non_terminated_only and status is None:
             continue
-        statuses[inst_id] = status
+        statuses[inst_id] = (status, None)
     return statuses