skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/clouds/service_catalog/do_catalog.py

@@ -0,0 +1,111 @@
+"""Digital ocean service catalog.
+
+This module loads the service catalog file and can be used to
+query instance types and pricing information for digital ocean.
+"""
+
+import typing
+from typing import Dict, List, Optional, Tuple, Union
+
+from sky.clouds.service_catalog import common
+from sky.utils import ux_utils
+
+if typing.TYPE_CHECKING:
+    from sky.clouds import cloud
+
+_df = common.read_catalog('do/vms.csv')
+
+
+def instance_type_exists(instance_type: str) -> bool:
+    return common.instance_type_exists_impl(_df, instance_type)
+
+
+def validate_region_zone(
+        region: Optional[str],
+        zone: Optional[str]) -> Tuple[Optional[str], Optional[str]]:
+    if zone is not None:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError('DO does not support zones.')
+    return common.validate_region_zone_impl('DO', _df, region, zone)
+
+
+def get_hourly_cost(
+    instance_type: str,
+    use_spot: bool = False,
+    region: Optional[str] = None,
+    zone: Optional[str] = None,
+) -> float:
+    """Returns the cost, or the cheapest cost among all zones for spot."""
+    if zone is not None:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError('DO does not support zones.')
+    return common.get_hourly_cost_impl(_df, instance_type, use_spot, region,
+                                       zone)
+
+
+def get_vcpus_mem_from_instance_type(
+    instance_type: str,) -> Tuple[Optional[float], Optional[float]]:
+    return common.get_vcpus_mem_from_instance_type_impl(_df, instance_type)
+
+
+def get_default_instance_type(
+    cpus: Optional[str] = None,
+    memory: Optional[str] = None,
+    disk_tier: Optional[str] = None,
+) -> Optional[str]:
+    # NOTE: After expanding catalog to multiple entries, you may
+    # want to specify a default instance type or family.
+    del disk_tier  # unused
+    return common.get_instance_type_for_cpus_mem_impl(_df, cpus, memory)
+
+
+def get_accelerators_from_instance_type(
+        instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
+    return common.get_accelerators_from_instance_type_impl(_df, instance_type)
+
+
+def get_instance_type_for_accelerator(
+    acc_name: str,
+    acc_count: int,
+    cpus: Optional[str] = None,
+    memory: Optional[str] = None,
+    use_spot: bool = False,
+    region: Optional[str] = None,
+    zone: Optional[str] = None,
+) -> Tuple[Optional[List[str]], List[str]]:
+    """Returns a list of instance types that have the given accelerator."""
+    if zone is not None:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError('DO does not support zones.')
+    return common.get_instance_type_for_accelerator_impl(
+        df=_df,
+        acc_name=acc_name,
+        acc_count=acc_count,
+        cpus=cpus,
+        memory=memory,
+        use_spot=use_spot,
+        region=region,
+        zone=zone,
+    )
+
+
+def get_region_zones_for_instance_type(instance_type: str,
+                                       use_spot: bool) -> List['cloud.Region']:
+    df = _df[_df['InstanceType'] == instance_type]
+    return common.get_region_zones(df, use_spot)
+
+
+def list_accelerators(
+    gpus_only: bool,
+    name_filter: Optional[str],
+    region_filter: Optional[str],
+    quantity_filter: Optional[int],
+    case_sensitive: bool = True,
+    all_regions: bool = False,
+    require_price: bool = True,
+) -> Dict[str, List[common.InstanceTypeInfo]]:
+    """Returns all instance types in DO offering GPUs."""
+    del require_price  # unused
+    return common.list_accelerators_impl('DO', _df, gpus_only, name_filter,
+                                         region_filter, quantity_filter,
+                                         case_sensitive, all_regions)

sky/clouds/service_catalog/fluidstack_catalog.py

@@ -4,7 +4,7 @@ This module loads the service catalog file and can be used to query
 instance types and pricing information for FluidStack.
 """
 import typing
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from sky.clouds.service_catalog import common
 from sky.utils import ux_utils
@@ -65,7 +65,7 @@ def get_default_instance_type(cpus: Optional[str] = None,
 
 
 def get_accelerators_from_instance_type(
-        instance_type: str) -> Optional[Dict[str, int]]:
+        instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
     return common.get_accelerators_from_instance_type_impl(_df, instance_type)
 
 

sky/clouds/service_catalog/gcp_catalog.py

@@ -9,6 +9,7 @@ from typing import Dict, List, Optional, Tuple
 from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
+from sky.clouds import GCP
 from sky.clouds.service_catalog import common
 from sky.utils import resources_utils
 from sky.utils import ux_utils
@@ -96,7 +97,13 @@ _ACC_INSTANCE_TYPE_DICTS = {
         8: ['g2-standard-96'],
     },
     'H100': {
+        1: ['a3-highgpu-1g'],
+        2: ['a3-highgpu-2g'],
+        4: ['a3-highgpu-4g'],
         8: ['a3-highgpu-8g'],
+    },
+    'H100-MEGA': {
+        8: ['a3-megagpu-8g'],
     }
 }
 
@@ -243,7 +250,6 @@ def get_default_instance_type(
         cpus: Optional[str] = None,
         memory: Optional[str] = None,
         disk_tier: Optional[resources_utils.DiskTier] = None) -> Optional[str]:
-    del disk_tier  # unused
     if cpus is None and memory is None:
         cpus = f'{_DEFAULT_NUM_VCPUS}+'
     if memory is None:
@@ -254,6 +260,12 @@ def get_default_instance_type(
         f'{family}-' for family in _DEFAULT_INSTANCE_FAMILY)
     df = _df[_df['InstanceType'].notna()]
     df = df[df['InstanceType'].str.startswith(instance_type_prefix)]
+
+    def _filter_disk_type(instance_type: str) -> bool:
+        valid, _ = GCP.check_disk_tier(instance_type, disk_tier)
+        return valid
+
+    df = df.loc[df['InstanceType'].apply(_filter_disk_type)]
     return common.get_instance_type_for_cpus_mem_impl(df, cpus,
                                                       memory_gb_or_ratio)
 
@@ -280,7 +292,9 @@ def get_instance_type_for_accelerator(
 
     if acc_name in _ACC_INSTANCE_TYPE_DICTS:
         df = _df[_df['InstanceType'].notna()]
-        instance_types = _ACC_INSTANCE_TYPE_DICTS[acc_name][acc_count]
+        instance_types = _ACC_INSTANCE_TYPE_DICTS[acc_name].get(acc_count, None)
+        if instance_types is None:
+            return None, []
         df = df[df['InstanceType'].isin(instance_types)]
 
         # Check the cpus and memory specified by the user.

sky/clouds/service_catalog/ibm_catalog.py

@@ -4,7 +4,7 @@ This module loads the service catalog file and can be used to query
 instance types and pricing information for IBM.
 """
 
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from sky import sky_logging
 from sky.adaptors import ibm
@@ -43,7 +43,7 @@ def get_vcpus_mem_from_instance_type(
 
 
 def get_accelerators_from_instance_type(
-        instance_type: str) -> Optional[Dict[str, int]]:
+        instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
     return common.get_accelerators_from_instance_type_impl(_df, instance_type)
 
 

sky/clouds/service_catalog/kubernetes_catalog.py

@@ -8,17 +8,23 @@ import typing
 from typing import Dict, List, Optional, Set, Tuple
 
 from sky import check as sky_check
+from sky import clouds as sky_clouds
+from sky import sky_logging
 from sky.adaptors import common as adaptors_common
-from sky.clouds import Kubernetes
+from sky.adaptors import kubernetes
 from sky.clouds.service_catalog import CloudFilter
 from sky.clouds.service_catalog import common
 from sky.provision.kubernetes import utils as kubernetes_utils
 
+logger = sky_logging.init_logger(__name__)
+
 if typing.TYPE_CHECKING:
     import pandas as pd
 else:
     pd = adaptors_common.LazyImport('pandas')
 
+logger = sky_logging.init_logger(__name__)
+
 _PULL_FREQUENCY_HOURS = 7
 
 # We keep pull_frequency_hours so we can remotely update the default image paths
@@ -31,7 +37,16 @@ _image_df = common.read_catalog('kubernetes/images.csv',
 
 def get_image_id_from_tag(tag: str, region: Optional[str]) -> Optional[str]:
     """Returns the image id from the tag."""
-    return common.get_image_id_from_tag_impl(_image_df, tag, region)
+    global _image_df
+    image_id = common.get_image_id_from_tag_impl(_image_df, tag, region)
+    if image_id is None:
+        # Refresh the image catalog and try again, if the image tag is not
+        # found.
+        logger.debug('Refreshing the image catalog and trying again.')
+        _image_df = common.read_catalog('kubernetes/images.csv',
+                                        pull_frequency_hours=0)
+        image_id = common.get_image_id_from_tag_impl(_image_df, tag, region)
+    return image_id
 
 
 def is_image_tag_valid(tag: str, region: Optional[str]) -> bool:
@@ -50,9 +65,14 @@ def list_accelerators(
     # TODO(romilb): We should consider putting a lru_cache() with TTL to
     #   avoid multiple calls to kubernetes API in a short period of time (e.g.,
     #   from the optimizer).
-    return list_accelerators_realtime(gpus_only, name_filter, region_filter,
-                                      quantity_filter, case_sensitive,
-                                      all_regions, require_price)[0]
+    return _list_accelerators(gpus_only,
+                              name_filter,
+                              region_filter,
+                              quantity_filter,
+                              case_sensitive,
+                              all_regions,
+                              require_price,
+                              realtime=False)[0]
 
 
 def list_accelerators_realtime(
@@ -65,27 +85,100 @@ def list_accelerators_realtime(
     require_price: bool = True
 ) -> Tuple[Dict[str, List[common.InstanceTypeInfo]], Dict[str, int], Dict[str,
                                                                           int]]:
+    return _list_accelerators(gpus_only,
+                              name_filter,
+                              region_filter,
+                              quantity_filter,
+                              case_sensitive,
+                              all_regions,
+                              require_price,
+                              realtime=True)
+
+
+def _list_accelerators(
+    gpus_only: bool,
+    name_filter: Optional[str],
+    region_filter: Optional[str],
+    quantity_filter: Optional[int],
+    case_sensitive: bool = True,
+    all_regions: bool = False,
+    require_price: bool = True,
+    realtime: bool = False
+) -> Tuple[Dict[str, List[common.InstanceTypeInfo]], Dict[str, int], Dict[str,
+                                                                          int]]:
+    """List accelerators in the Kubernetes cluster.
+
+    If realtime is True, the function will query the cluster to fetch real-time
+    GPU usage, which is returned in total_accelerators_available. Note that
+    this may require an expensive list_pod_for_all_namespaces call, which
+    requires cluster-wide pod read permissions.
+
+    If the user does not have sufficient permissions to list pods in all
+    namespaces, the function will return free GPUs as -1.
+
+    Returns:
+        A tuple of three dictionaries:
+        - qtys_map: Dict mapping accelerator names to lists of InstanceTypeInfo
+            objects with quantity information.
+        - total_accelerators_capacity: Dict mapping accelerator names to their
+            total capacity in the cluster.
+        - total_accelerators_available: Dict mapping accelerator names to their
+            current availability. Returns -1 for each accelerator if
+            realtime=False or if insufficient permissions.
+    """
+    # TODO(romilb): This should be refactored to use get_kubernetes_node_info()
+    #   function from kubernetes_utils.
     del all_regions, require_price  # Unused.
-    k8s_cloud = Kubernetes()
-    if not any(
-            map(k8s_cloud.is_same_cloud,
-                sky_check.get_cached_enabled_clouds_or_refresh())
-    ) or not kubernetes_utils.check_credentials()[0]:
+
+    # First check if Kubernetes is enabled. This ensures k8s python client is
+    # installed. Do not put any k8s-specific logic before this check.
+    enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh()
+    if not sky_clouds.cloud_in_iterable(sky_clouds.Kubernetes(),
+                                        enabled_clouds):
+        return {}, {}, {}
+
+    # TODO(zhwu): this should return all accelerators in multiple kubernetes
+    # clusters defined by allowed_contexts.
+    if region_filter is None:
+        context = kubernetes_utils.get_current_kube_config_context_name()
+        if context is None and kubernetes_utils.is_incluster_config_available():
+            # If context is None and we are running in a kubernetes pod, use the
+            # in-cluster context as the current context.
+            context = kubernetes.in_cluster_context_name()
+    else:
+        context = region_filter
+    if context is None:
         return {}, {}, {}
 
-    has_gpu = kubernetes_utils.detect_gpu_resource()
+    # Verify that the credentials are still valid.
+    if not kubernetes_utils.check_credentials(context)[0]:
+        return {}, {}, {}
+
+    has_gpu = kubernetes_utils.detect_accelerator_resource(context)
     if not has_gpu:
         return {}, {}, {}
 
-    label_formatter, _ = kubernetes_utils.detect_gpu_label_formatter()
-    if not label_formatter:
+    lf, _ = kubernetes_utils.detect_gpu_label_formatter(context)
+    if not lf:
         return {}, {}, {}
 
     accelerators_qtys: Set[Tuple[str, int]] = set()
-    key = label_formatter.get_label_key()
-    nodes = kubernetes_utils.get_kubernetes_nodes()
-    # Get the pods to get the real-time GPU usage
-    pods = kubernetes_utils.get_kubernetes_pods()
+    keys = lf.get_label_keys()
+    nodes = kubernetes_utils.get_kubernetes_nodes(context)
+    pods = None
+    if realtime:
+        # Get the pods to get the real-time GPU usage
+        try:
+            pods = kubernetes_utils.get_all_pods_in_kubernetes_cluster(context)
+        except kubernetes.api_exception() as e:
+            if e.status == 403:
+                logger.warning(
+                    'Failed to get pods in the Kubernetes cluster '
+                    '(forbidden). Please check if your account has '
+                    'necessary permissions to list pods. Realtime GPU '
+                    'availability information may be incorrect.')
+            else:
+                raise
     # Total number of GPUs in the cluster
     total_accelerators_capacity: Dict[str, int] = {}
     # Total number of GPUs currently available in the cluster
@@ -93,58 +186,88 @@ def list_accelerators_realtime(
     min_quantity_filter = quantity_filter if quantity_filter else 1
 
     for node in nodes:
-        if key in node.metadata.labels:
-            allocated_qty = 0
-            accelerator_name = label_formatter.get_accelerator_from_label_value(
-                node.metadata.labels.get(key))
-
-            # Check if name_filter regex matches the accelerator_name
-            regex_flags = 0 if case_sensitive else re.IGNORECASE
-            if name_filter and not re.match(
-                    name_filter, accelerator_name, flags=regex_flags):
-                continue
-
-            accelerator_count = int(
-                node.status.allocatable.get('nvidia.com/gpu', 0))
-
-            # Generate the GPU quantities for the accelerators
-            if accelerator_name and accelerator_count > 0:
-                for count in range(1, accelerator_count + 1):
-                    accelerators_qtys.add((accelerator_name, count))
-
-            for pod in pods:
-                # Get all the pods running on the node
-                if (pod.spec.node_name == node.metadata.name and
-                        pod.status.phase in ['Running', 'Pending']):
-                    # Iterate over all the containers in the pod and sum the
-                    # GPU requests
-                    for container in pod.spec.containers:
-                        if container.resources.requests:
-                            allocated_qty += int(
-                                container.resources.requests.get(
-                                    'nvidia.com/gpu', 0))
-
-            accelerators_available = accelerator_count - allocated_qty
-
-            if accelerator_count >= min_quantity_filter:
-                quantized_count = (min_quantity_filter *
-                                   (accelerator_count // min_quantity_filter))
-                if accelerator_name not in total_accelerators_capacity:
-                    total_accelerators_capacity[
-                        accelerator_name] = quantized_count
-                else:
-                    total_accelerators_capacity[
-                        accelerator_name] += quantized_count
-
-            if accelerators_available >= min_quantity_filter:
-                quantized_availability = min_quantity_filter * (
-                    accelerators_available // min_quantity_filter)
+        for key in keys:
+            if key in node.metadata.labels:
+                allocated_qty = 0
+                accelerator_name = lf.get_accelerator_from_label_value(
+                    node.metadata.labels.get(key))
+
+                # Exclude multi-host TPUs from being processed.
+                # TODO(Doyoung): Remove the logic when adding support for
+                # multi-host TPUs.
+                if kubernetes_utils.is_multi_host_tpu(node.metadata.labels):
+                    continue
+
+                # Check if name_filter regex matches the accelerator_name
+                regex_flags = 0 if case_sensitive else re.IGNORECASE
+                if name_filter and not re.match(
+                        name_filter, accelerator_name, flags=regex_flags):
+                    continue
+
+                # Generate the accelerator quantities
+                accelerator_count = (
+                    kubernetes_utils.get_node_accelerator_count(
+                        node.status.allocatable))
+
+                if accelerator_name and accelerator_count > 0:
+                    # TPUs are counted in a different way compared to GPUs.
+                    # Multi-node GPUs can be split into smaller units and be
+                    # provisioned, but TPUs are considered as an atomic unit.
+                    if kubernetes_utils.is_tpu_on_gke(accelerator_name):
+                        accelerators_qtys.add(
+                            (accelerator_name, accelerator_count))
+                    else:
+                        count = 1
+                        while count <= accelerator_count:
+                            accelerators_qtys.add((accelerator_name, count))
+                            count *= 2
+                        # Add the accelerator count if it's not already in the
+                        # set (e.g., if there's 12 GPUs, we should have qtys 1,
+                        # 2, 4, 8, 12)
+                        if accelerator_count not in accelerators_qtys:
+                            accelerators_qtys.add(
+                                (accelerator_name, accelerator_count))
+
+                if accelerator_count >= min_quantity_filter:
+                    quantized_count = (
+                        min_quantity_filter *
+                        (accelerator_count // min_quantity_filter))
+                    if accelerator_name not in total_accelerators_capacity:
+                        total_accelerators_capacity[
+                            accelerator_name] = quantized_count
+                    else:
+                        total_accelerators_capacity[
+                            accelerator_name] += quantized_count
+
+                if pods is None:
+                    # If we can't get the pods, we can't get the GPU usage
+                    total_accelerators_available[accelerator_name] = -1
+                    continue
+
+                for pod in pods:
+                    # Get all the pods running on the node
+                    if (pod.spec.node_name == node.metadata.name and
+                            pod.status.phase in ['Running', 'Pending']):
+                        # Iterate over all the containers in the pod and sum
+                        # the GPU requests
+                        for container in pod.spec.containers:
+                            if container.resources.requests:
+                                allocated_qty += (
+                                    kubernetes_utils.get_node_accelerator_count(
+                                        container.resources.requests))
+
+                accelerators_available = accelerator_count - allocated_qty
+
+                # Initialize the entry if it doesn't exist yet
                 if accelerator_name not in total_accelerators_available:
-                    total_accelerators_available[
-                        accelerator_name] = quantized_availability
-                else:
-                    total_accelerators_available[
-                        accelerator_name] += quantized_availability
+                    total_accelerators_available[accelerator_name] = 0
+
+                if accelerators_available >= min_quantity_filter:
+                    quantized_availability = min_quantity_filter * (
+                        accelerators_available // min_quantity_filter)
+                    total_accelerators_available[accelerator_name] = (
+                        total_accelerators_available.get(accelerator_name, 0) +
+                        quantized_availability)
 
     result = []
 
@@ -160,7 +283,7 @@ def list_accelerators_realtime(
                                     memory=None,
                                     price=0.0,
                                     spot_price=0.0,
-                                    region='kubernetes'))
+                                    region=context))
 
     df = pd.DataFrame(result,
                       columns=[
@@ -175,7 +298,6 @@ def list_accelerators_realtime(
     qtys_map = common.list_accelerators_impl('Kubernetes', df, gpus_only,
                                              name_filter, region_filter,
                                              quantity_filter, case_sensitive)
-
     return qtys_map, total_accelerators_capacity, total_accelerators_available
 
 

sky/clouds/service_catalog/lambda_catalog.py

@@ -4,7 +4,7 @@ This module loads the service catalog file and can be used to query
 instance types and pricing information for Lambda.
 """
 import typing
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 from sky.clouds.service_catalog import common
 from sky.utils import resources_utils
@@ -13,7 +13,12 @@ from sky.utils import ux_utils
 if typing.TYPE_CHECKING:
     from sky.clouds import cloud
 
-_df = common.read_catalog('lambda/vms.csv')
+# Keep it synced with the frequency in
+# skypilot-catalog/.github/workflows/update-lambda-catalog.yml
+_PULL_FREQUENCY_HOURS = 7
+
+_df = common.read_catalog('lambda/vms.csv',
+                          pull_frequency_hours=_PULL_FREQUENCY_HOURS)
 
 # Number of vCPUS for gpu_1x_a10
 _DEFAULT_NUM_VCPUS = 30
@@ -67,7 +72,7 @@ def get_default_instance_type(
 
 
 def get_accelerators_from_instance_type(
-        instance_type: str) -> Optional[Dict[str, int]]:
+        instance_type: str) -> Optional[Dict[str, Union[int, float]]]:
     return common.get_accelerators_from_instance_type_impl(_df, instance_type)