skypilot-nightly 1.0.0.dev20251203__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/catalog/shadeform_catalog.py

@@ -7,12 +7,15 @@ and can be used to query instance types and pricing information for Shadeform.
 import typing
 from typing import Dict, List, Optional, Tuple, Union
 
-import pandas as pd
-
+from sky.adaptors import common as adaptors_common
 from sky.catalog import common
 
 if typing.TYPE_CHECKING:
+    import pandas as pd
+
     from sky.clouds import cloud
+else:
+    pd = adaptors_common.LazyImport('pandas')
 
 # We'll use dynamic fetching, so no static CSV file to load
 _df = None

sky/catalog/slurm_catalog.py

@@ -0,0 +1,236 @@
+"""Slurm Catalog."""
+
+import collections
+import re
+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.catalog import common
+from sky.clouds import cloud
+from sky.provision.slurm import utils as slurm_utils
+from sky.utils import resources_utils
+
+logger = sky_logging.init_logger(__name__)
+
+_DEFAULT_NUM_VCPUS = 2
+_DEFAULT_MEMORY_CPU_RATIO = 1
+
+
+def instance_type_exists(instance_type: str) -> bool:
+    """Check if the given instance type is valid for Slurm."""
+    return slurm_utils.SlurmInstanceType.is_valid_instance_type(instance_type)
+
+
+def get_default_instance_type(cpus: Optional[str] = None,
+                              memory: Optional[str] = None,
+                              disk_tier: Optional[
+                                  resources_utils.DiskTier] = None,
+                              region: Optional[str] = None,
+                              zone: Optional[str] = None) -> Optional[str]:
+    # Delete unused parameters.
+    del disk_tier, region, zone
+
+    # Slurm provisions resources via --cpus-per-task and --mem.
+    instance_cpus = float(
+        cpus.strip('+')) if cpus is not None else _DEFAULT_NUM_VCPUS
+    if memory is not None:
+        if memory.endswith('+'):
+            instance_mem = float(memory[:-1])
+        elif memory.endswith('x'):
+            instance_mem = float(memory[:-1]) * instance_cpus
+        else:
+            instance_mem = float(memory)
+    else:
+        instance_mem = instance_cpus * _DEFAULT_MEMORY_CPU_RATIO
+    virtual_instance_type = slurm_utils.SlurmInstanceType(
+        instance_cpus, instance_mem).name
+    return virtual_instance_type
+
+
+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]]:
+    """List accelerators in Slurm clusters.
+
+    Returns a dictionary mapping GPU type to a list of InstanceTypeInfo objects.
+    """
+    return list_accelerators_realtime(gpus_only, name_filter, region_filter,
+                                      quantity_filter, case_sensitive,
+                                      all_regions, require_price)[0]
+
+
+def list_accelerators_realtime(
+    gpus_only: bool = True,
+    name_filter: Optional[str] = None,
+    region_filter: Optional[str] = None,
+    quantity_filter: Optional[int] = None,
+    case_sensitive: bool = True,
+    all_regions: bool = False,
+    require_price: bool = False,
+) -> Tuple[Dict[str, List[common.InstanceTypeInfo]], Dict[str, int], Dict[str,
+                                                                          int]]:
+    """Fetches real-time accelerator information from the Slurm cluster.
+
+    Uses the `get_slurm_node_info_list` helper function.
+
+    Args:
+        gpus_only: If True, only return GPU accelerators.
+        name_filter: Regex filter for accelerator names (e.g., 'V100', 'gpu').
+        region_filter: Optional filter for Slurm partitions.
+        quantity_filter: Minimum number of accelerators required per node.
+        case_sensitive: Whether name_filter is case-sensitive.
+        all_regions: Unused in Slurm context.
+        require_price: Unused in Slurm context.
+
+    Returns:
+        A tuple of three dictionaries:
+        - qtys_map: Maps GPU type to set of InstanceTypeInfo objects for unique
+          counts found per node.
+        - total_capacity: Maps GPU type to total count across all nodes.
+        - total_available: Maps GPU type to total free count across all nodes.
+    """
+    del gpus_only, all_regions, require_price
+
+    enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh(
+        cloud.CloudCapability.COMPUTE)
+    if not sky_clouds.cloud_in_iterable(sky_clouds.Slurm(), enabled_clouds):
+        return {}, {}, {}
+
+    if region_filter is None:
+        # Get the first available cluster as default
+        all_clusters = slurm_utils.get_all_slurm_cluster_names()
+        if not all_clusters:
+            return {}, {}, {}
+        slurm_cluster = all_clusters[0]
+    else:
+        slurm_cluster = region_filter
+
+    slurm_nodes_info = slurm_utils.slurm_node_info(
+        slurm_cluster_name=slurm_cluster)
+
+    if not slurm_nodes_info:
+        # Customize error message based on filters
+        err_msg = 'No matching GPU nodes found in the Slurm cluster'
+        filters_applied = []
+        if name_filter:
+            filters_applied.append(f'gpu_name={name_filter!r}')
+        if quantity_filter:
+            filters_applied.append(f'quantity>={quantity_filter}')
+        if filters_applied:
+            err_msg += f' with filters ({", ".join(filters_applied)})'
+        err_msg += '.'
+        logger.error(
+            err_msg)  # Log as error as it indicates no usable resources found
+        raise ValueError(err_msg)
+
+    # Aggregate results into the required format
+    qtys_map: Dict[str,
+                   Set[common.InstanceTypeInfo]] = collections.defaultdict(set)
+    total_capacity: Dict[str, int] = collections.defaultdict(int)
+    total_available: Dict[str, int] = collections.defaultdict(int)
+
+    for node_info in slurm_nodes_info:
+        gpu_type = node_info['gpu_type']
+        node_total_gpus = node_info['total_gpus']
+        node_free_gpus = node_info['free_gpus']
+        partition = node_info['partition']
+
+        # Apply name filter to the determined GPU type
+        regex_flags = 0 if case_sensitive else re.IGNORECASE
+        if name_filter and not re.match(
+                name_filter, gpu_type, flags=regex_flags):
+            continue
+
+        # Apply quantity filter (total GPUs on node must meet this)
+        if quantity_filter and node_total_gpus < quantity_filter:
+            continue
+
+        # Apply partition filter if specified
+        # TODO(zhwu): when a node is in multiple partitions, the partition
+        # mapping from node to partition does not work.
+        # if partition_filter and partition != partition_filter:
+        #     continue
+
+        # Create InstanceTypeInfo objects for various GPU counts
+        # Similar to Kubernetes, generate powers of 2 up to node_total_gpus
+        if node_total_gpus > 0:
+            count = 1
+            while count <= node_total_gpus:
+                instance_info = common.InstanceTypeInfo(
+                    instance_type=None,  # Slurm doesn't have instance types
+                    accelerator_name=gpu_type,
+                    accelerator_count=count,
+                    cpu_count=node_info['vcpu_count'],
+                    memory=node_info['memory_gb'],
+                    price=0.0,  # Slurm doesn't have price info
+                    region=partition,  # Use partition as region
+                    cloud='slurm',  # Specify cloud as 'slurm'
+                    device_memory=0.0,  # No GPU memory info from Slurm
+                    spot_price=0.0,  # Slurm doesn't have spot pricing
+                )
+                qtys_map[gpu_type].add(instance_info)
+                count *= 2
+
+            # Add the actual total if it's not already included
+            # (e.g., if node has 12 GPUs, include counts 1, 2, 4, 8, 12)
+            if count // 2 != node_total_gpus:
+                instance_info = common.InstanceTypeInfo(
+                    instance_type=None,
+                    accelerator_name=gpu_type,
+                    accelerator_count=node_total_gpus,
+                    cpu_count=node_info['vcpu_count'],
+                    memory=node_info['memory_gb'],
+                    price=0.0,
+                    region=partition,
+                    cloud='slurm',
+                    device_memory=0.0,
+                    spot_price=0.0,
+                )
+                qtys_map[gpu_type].add(instance_info)
+
+        # Map of GPU type -> total count across all matched nodes
+        total_capacity[gpu_type] += node_total_gpus
+
+        # Map of GPU type -> total *free* count across all matched nodes
+        total_available[gpu_type] += node_free_gpus
+
+    # Check if any GPUs were found after applying filters
+    if not total_capacity:
+        err_msg = 'No matching GPU nodes found in the Slurm cluster'
+        filters_applied = []
+        if name_filter:
+            filters_applied.append(f'gpu_name={name_filter!r}')
+        if quantity_filter:
+            filters_applied.append(f'quantity>={quantity_filter}')
+        if filters_applied:
+            err_msg += f' with filters ({", ".join(filters_applied)})'
+        err_msg += '.'
+        logger.error(err_msg)
+        raise ValueError(err_msg)
+
+    # Convert sets of InstanceTypeInfo to sorted lists
+    final_qtys_map = {
+        gpu: sorted(list(instances), key=lambda x: x.accelerator_count)
+        for gpu, instances in qtys_map.items()
+    }
+
+    logger.debug(f'Aggregated Slurm GPU Info: '
+                 f'qtys={final_qtys_map}, '
+                 f'capacity={dict(total_capacity)}, '
+                 f'available={dict(total_available)}')
+
+    return final_qtys_map, dict(total_capacity), dict(total_available)
+
+
+def validate_region_zone(
+    region_name: Optional[str],
+    zone_name: Optional[str],
+) -> Tuple[Optional[str], Optional[str]]:
+    return (region_name, zone_name)

sky/catalog/vast_catalog.py

@@ -7,7 +7,10 @@ query instance types and pricing information for Vast.ai.
 import typing
 from typing import Dict, List, Optional, Tuple, Union
 
+import pandas as pd
+
 from sky.catalog import common
+from sky.utils import resources_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -16,6 +19,17 @@ if typing.TYPE_CHECKING:
 _df = common.read_catalog('vast/vms.csv')
 
 
+def _apply_datacenter_filter(df: pd.DataFrame,
+                             datacenter_only: bool) -> pd.DataFrame:
+    """Filter dataframe by hosting_type if datacenter_only is True.
+
+    hosting_type: 0 = Consumer hosted, 1 = Datacenter hosted
+    """
+    if not datacenter_only or 'HostingType' not in df.columns:
+        return df
+    return df[df['HostingType'] >= 1]
+
+
 def instance_type_exists(instance_type: str) -> bool:
     return common.instance_type_exists_impl(_df, instance_type)
 
@@ -48,13 +62,16 @@ def get_vcpus_mem_from_instance_type(
 
 def get_default_instance_type(cpus: Optional[str] = None,
                               memory: Optional[str] = None,
-                              disk_tier: Optional[str] = None,
+                              disk_tier: Optional[
+                                  resources_utils.DiskTier] = None,
                               region: Optional[str] = None,
-                              zone: Optional[str] = None) -> Optional[str]:
+                              zone: Optional[str] = None,
+                              datacenter_only: bool = False) -> Optional[str]:
     del disk_tier
     # NOTE: After expanding catalog to multiple entries, you may
     # want to specify a default instance type or family.
-    return common.get_instance_type_for_cpus_mem_impl(_df, cpus, memory, region,
+    df = _apply_datacenter_filter(_df, datacenter_only)
+    return common.get_instance_type_for_cpus_mem_impl(df, cpus, memory, region,
                                                       zone)
 
 
@@ -70,12 +87,19 @@ def get_instance_type_for_accelerator(
         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."""
+        zone: Optional[str] = None,
+        datacenter_only: bool = False) -> Tuple[Optional[List[str]], List[str]]:
+    """Returns a list of instance types that have the given accelerator.
+
+    Args:
+        datacenter_only: If True, only return instances hosted in datacenters
+            (hosting_type >= 1).
+    """
     if zone is not None:
         with ux_utils.print_exception_no_traceback():
             raise ValueError('Vast does not support zones.')
-    return common.get_instance_type_for_accelerator_impl(df=_df,
+    df = _apply_datacenter_filter(_df, datacenter_only)
+    return common.get_instance_type_for_accelerator_impl(df=df,
                                                          acc_name=acc_name,
                                                          acc_count=acc_count,
                                                          cpus=cpus,

sky/check.py

@@ -528,8 +528,9 @@ def _print_checked_cloud(
         # `dict` reasons for K8s and SSH will be printed in detail in
         # _format_enabled_cloud. Skip here unless the cloud is disabled.
         if not isinstance(reason, str):
-            if not ok and isinstance(cloud_tuple[1],
-                                     (sky_clouds.SSH, sky_clouds.Kubernetes)):
+            if not ok and isinstance(
+                    cloud_tuple[1],
+                (sky_clouds.SSH, sky_clouds.Kubernetes, sky_clouds.Slurm)):
                 if reason is not None:
                     reason_str = _format_context_details(cloud_tuple[1],
                                                          show_details=True,
@@ -555,7 +556,9 @@ def _print_checked_cloud(
         capability_string = f'[{", ".join(enabled_capabilities)}]'
         if verbose and cloud is not cloudflare and cloud is not coreweave:
             activated_account = cloud.get_active_user_identity_str()
-        if isinstance(cloud_tuple[1], (sky_clouds.SSH, sky_clouds.Kubernetes)):
+        if isinstance(
+                cloud_tuple[1],
+            (sky_clouds.SSH, sky_clouds.Kubernetes, sky_clouds.Slurm)):
             detail_string = _format_context_details(cloud_tuple[1],
                                                     show_details=True,
                                                     ctx2text=ctx2text)
@@ -586,6 +589,9 @@ def _format_context_details(cloud: Union[str, sky_clouds.Cloud],
     if isinstance(cloud_type, sky_clouds.SSH):
         # Get the cluster names by reading from the node pools file
         contexts = sky_clouds.SSH.get_ssh_node_pool_contexts()
+    elif isinstance(cloud_type, sky_clouds.Slurm):
+        # Get the cluster names from SLURM config
+        contexts = sky_clouds.Slurm.existing_allowed_clusters()
     else:
         assert isinstance(cloud_type, sky_clouds.Kubernetes)
         contexts = sky_clouds.Kubernetes.existing_allowed_contexts()
@@ -650,15 +656,19 @@ def _format_context_details(cloud: Union[str, sky_clouds.Cloud],
                                               'configuration.'))
                 else:
                     # Default case - not set up
-                    text_suffix = (': ' + _red_color('disabled. ') +
-                                   _dim_color('Reason: Not set up. Use '
-                                              '`sky ssh up --infra '
-                                              f'{context.lstrip("ssh-")}` '
-                                              'to set up.'))
+                    text_suffix = (': ' + _red_color('disabled. ') + _dim_color(
+                        'Reason: Not set up. Use '
+                        '`sky ssh up --infra '
+                        f'{common_utils.removeprefix(context, "ssh-")}` '
+                        'to set up.'))
         contexts_formatted.append(
             f'\n    {symbol}{cleaned_context}{text_suffix}')
-    identity_str = ('SSH Node Pools' if isinstance(cloud_type, sky_clouds.SSH)
-                    else 'Allowed contexts')
+    if isinstance(cloud_type, sky_clouds.SSH):
+        identity_str = 'SSH Node Pools'
+    elif isinstance(cloud_type, sky_clouds.Slurm):
+        identity_str = 'Allowed clusters'
+    else:
+        identity_str = 'Allowed contexts'
     return f'\n    {identity_str}:{"".join(contexts_formatted)}'
 
 
@@ -677,7 +687,11 @@ def _format_enabled_cloud(cloud_name: str,
     cloud_and_capabilities = f'{cloud_name} [{", ".join(capabilities)}]'
     title = _green_color(cloud_and_capabilities)
 
-    if cloud_name in [repr(sky_clouds.Kubernetes()), repr(sky_clouds.SSH())]:
+    if cloud_name in [
+            repr(sky_clouds.Kubernetes()),
+            repr(sky_clouds.SSH()),
+            repr(sky_clouds.Slurm())
+    ]:
         return (f'{title}' + _format_context_details(
             cloud_name, show_details=False, ctx2text=ctx2text))
     return _green_color(cloud_and_capabilities)