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

sky/backends/wheel_utils.py

@@ -100,6 +100,13 @@ def _build_sky_wheel() -> pathlib.Path:
                 # modify the commit hash in the file later.
                 # Symlink other files/folders.
                 target.symlink_to(item, target_is_directory=item.is_dir())
+
+        # Symlink sky_templates directory from repo root
+        sky_templates_src = SKY_PACKAGE_PATH.parent / 'sky_templates'
+        if sky_templates_src.exists():
+            sky_templates_target = tmp_dir / 'sky_templates'
+            sky_templates_target.symlink_to(sky_templates_src,
+                                            target_is_directory=True)
         setup_files_dir = SKY_PACKAGE_PATH / 'setup_files'
 
         setup_content = (setup_files_dir / 'setup.py').read_text()
@@ -244,6 +251,17 @@ def build_sky_wheel() -> Tuple[pathlib.Path, str]:
         # protocol. "compare, update and clone" has to be atomic to avoid
         # race conditions.
         last_modification_time = _get_latest_modification_time(SKY_PACKAGE_PATH)
+        # Also check sky_templates directory modification time
+        sky_templates_path = SKY_PACKAGE_PATH.parent / 'sky_templates'
+        if sky_templates_path.exists():
+            sky_templates_mtime = _get_latest_modification_time(
+                sky_templates_path)
+            if (last_modification_time is not None and
+                    sky_templates_mtime is not None):
+                last_modification_time = max(last_modification_time,
+                                             sky_templates_mtime)
+            elif last_modification_time is None:
+                last_modification_time = sky_templates_mtime
         last_wheel_modification_time = _get_latest_modification_time(WHEEL_DIR)
 
         # Only build wheels if the wheel is outdated, wheel does not exist

sky/catalog/__init__.py

@@ -247,6 +247,13 @@ def get_accelerators_from_instance_type(
                                instance_type)
 
 
+def get_arch_from_instance_type(instance_type: str,
+                                clouds: CloudFilter = None) -> Optional[str]:
+    """Returns the arch from a instance type."""
+    return _map_clouds_catalog(clouds, 'get_arch_from_instance_type',
+                               instance_type)
+
+
 def get_instance_type_for_accelerator(
     acc_name: str,
     acc_count: Union[int, float],
@@ -326,6 +333,7 @@ def get_common_gpus() -> List[str]:
         'A10G',
         'A100',
         'A100-80GB',
+        'B200',
         'H100',
         'H200',
         'L4',

sky/catalog/aws_catalog.py

@@ -271,6 +271,10 @@ def get_accelerators_from_instance_type(
         _get_df(), instance_type)
 
 
+def get_arch_from_instance_type(instance_type: str) -> Optional[str]:
+    return common.get_arch_from_instance_type_impl(_get_df(), instance_type)
+
+
 def get_instance_type_for_accelerator(
     acc_name: str,
     acc_count: int,

sky/catalog/common.py

@@ -385,7 +385,7 @@ def get_hourly_cost_impl(
                              f'{instance_type!r}.')
     cheapest_idx = df[price_str].idxmin()
     cheapest = df.loc[cheapest_idx]
-    return cheapest[price_str]
+    return float(cheapest[price_str])
 
 
 def _get_value(value):
@@ -527,6 +527,24 @@ def get_accelerators_from_instance_type_impl(
     return {acc_name: _convert(acc_count)}
 
 
+def get_arch_from_instance_type_impl(
+    df: 'pd.DataFrame',
+    instance_type: str,
+) -> Optional[str]:
+    df = _get_instance_type(df, instance_type, None)
+    if df.empty:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'No instance type {instance_type} found.')
+    row = df.iloc[0]
+    if 'Arch' not in row:
+        return None
+    arch = row['Arch']
+    if pd.isnull(arch):
+        return None
+
+    return arch
+
+
 def get_instance_type_for_accelerator_impl(
     df: 'pd.DataFrame',
     acc_name: str,

sky/catalog/data_fetchers/fetch_aws.py

@@ -60,6 +60,7 @@ ALL_REGIONS = [
     'ap-northeast-2',
     'ap-southeast-1',
     'ap-southeast-2',
+    'ap-southeast-4',
     'ap-northeast-1',
 ]
 US_REGIONS = ['us-east-1', 'us-east-2', 'us-west-1', 'us-west-2']
@@ -67,17 +68,13 @@ US_REGIONS = ['us-east-1', 'us-east-2', 'us-west-1', 'us-west-2']
 # The following columns will be included in the final catalog.
 USEFUL_COLUMNS = [
     'InstanceType', 'AcceleratorName', 'AcceleratorCount', 'vCPUs', 'MemoryGiB',
-    'GpuInfo', 'Price', 'SpotPrice', 'Region', 'AvailabilityZone'
+    'GpuInfo', 'Price', 'SpotPrice', 'Region', 'AvailabilityZone', 'Arch'
 ]
 
 # NOTE: the hard-coded us-east-1 URL is not a typo. AWS pricing endpoint is
 # only available in this region, but it serves pricing information for all
 # regions.
 PRICING_TABLE_URL_FMT = 'https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonEC2/current/{region}/index.csv'  # pylint: disable=line-too-long
-# Hardcode the regions that offer p4de.24xlarge as our credential does not have
-# the permission to query the offerings of the instance.
-# Ref: https://aws.amazon.com/ec2/instance-types/p4/
-P4DE_REGIONS = ['us-east-1', 'us-west-2']
 # g6f instances have fractional GPUs, but the API returns Count: 1 under
 # GpuInfo. However, the GPU memory is properly scaled. Taking the instance GPU
 # divided by the total memory of an L4 will give us the fraction of the GPU.
@@ -214,35 +211,6 @@ def _get_spot_pricing_table(region: str) -> 'pd.DataFrame':
     return df
 
 
-def _patch_p4de(region: str, df: 'pd.DataFrame',
-                pricing_df: 'pd.DataFrame') -> 'pd.DataFrame':
-    # Hardcoded patch for p4de.24xlarge, as our credentials doesn't have access
-    # to the instance type.
-    # Columns:
-    # InstanceType,AcceleratorName,AcceleratorCount,vCPUs,MemoryGiB,GpuInfo,
-    # Price,SpotPrice,Region,AvailabilityZone
-    records = []
-    for zone in df[df['Region'] == region]['AvailabilityZone'].unique():
-        records.append({
-            'InstanceType': 'p4de.24xlarge',
-            'AcceleratorName': 'A100-80GB',
-            'AcceleratorCount': 8,
-            'vCPUs': 96,
-            'MemoryGiB': 1152,
-            'GpuInfo':
-                ('{\'Gpus\': [{\'Name\': \'A100-80GB\', \'Manufacturer\': '
-                 '\'NVIDIA\', \'Count\': 8, \'MemoryInfo\': {\'SizeInMiB\': '
-                 '81920}}], \'TotalGpuMemoryInMiB\': 655360}'),
-            'AvailabilityZone': zone,
-            'Region': region,
-            'Price': pricing_df[pricing_df['InstanceType'] == 'p4de.24xlarge']
-                     ['Price'].values[0],
-            'SpotPrice': np.nan,
-        })
-    df = pd.concat([df, pd.DataFrame.from_records(records)])
-    return df
-
-
 def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
     try:
         # Fetch the zone info first to make sure the account has access to the
@@ -266,7 +234,7 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
         def get_acc_info(row) -> Tuple[Optional[str], float]:
             accelerator = None
             for col, info_key in [('GpuInfo', 'Gpus'),
-                                  ('InferenceAcceleratorInfo', 'Accelerators'),
+                                  ('NeuronInfo', 'NeuronDevices'),
                                   ('FpgaInfo', 'Fpgas')]:
                 info = row.get(col)
                 if isinstance(info, dict):
@@ -275,6 +243,17 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
                 return None, np.nan
             return accelerator['Name'], accelerator['Count']
 
+        def get_arch(row) -> Optional[str]:
+            if 'ProcessorInfo' in row:
+                processor = row['ProcessorInfo']
+                if 'SupportedArchitectures' in processor:
+                    archs = processor['SupportedArchitectures']
+                    if isinstance(archs, list):
+                        return archs[0]
+                    elif isinstance(archs, str):
+                        return archs
+            return None
+
         def get_vcpus(row) -> float:
             if not np.isnan(row['vCPU']):
                 return float(row['vCPU'])
@@ -299,18 +278,6 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
             if row['InstanceType'] == 'p4de.24xlarge':
                 acc_name = 'A100-80GB'
                 acc_count = 8
-            if row['InstanceType'].startswith('trn1'):
-                # Trainium instances does not have a field for information of
-                # the accelerators. We need to infer the accelerator info from
-                # the instance type name.
-                # aws ec2 describe-instance-types --region us-east-1
-                # https://aws.amazon.com/ec2/instance-types/trn1/
-                acc_name = 'Trainium'
-                find_num_in_name = re.search(r'(\d+)xlarge',
-                                             row['InstanceType'])
-                assert find_num_in_name is not None, row['InstanceType']
-                num_in_name = find_num_in_name.group(1)
-                acc_count = int(num_in_name) // 2
             if row['InstanceType'] == 'p5en.48xlarge':
                 # TODO(andyl): Check if this workaround still needed after
                 # v0.10.0 released. Currently, the acc_name returned by the
@@ -320,10 +287,15 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
             if (row['InstanceType'].startswith('g6f') or
                     row['InstanceType'].startswith('gr6f')):
                 # These instance actually have only fractional GPUs, but the API
-                # returns Count: 1 under GpuInfo. We need to check the GPU
-                # memory to get the actual fraction of the GPU.
+                # returns Count: 1 or Count: 0 under GpuInfo. We need to
+                # directly check the GPU memory to get the actual fraction of
+                # the GPU. Note that TotalGpuMemoryInMiB seems unreliable here -
+                # sometimes it is unexpectedly 0.
                 # See also Standard_NV{vcpu}ads_A10_v5 support on Azure.
-                fraction = row['GpuInfo']['TotalGpuMemoryInMiB'] / L4_GPU_MEMORY
+                assert len(row['GpuInfo']['Gpus']) == 1
+                assert row['GpuInfo']['Gpus'][0]['Name'] == 'L4'
+                fraction = row['GpuInfo']['Gpus'][0]['MemoryInfo'][
+                    'SizeInMiB'] / L4_GPU_MEMORY
                 acc_count = round(fraction, 3)
             if row['InstanceType'] == 'p5.4xlarge':
                 acc_count = 1
@@ -332,6 +304,7 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
                 'AcceleratorCount': acc_count,
                 'vCPUs': get_vcpus(row),
                 'MemoryGiB': get_memory_gib(row),
+                'Arch': get_arch(row),
             })
 
         # The AWS API may not have all the instance types in the pricing table,
@@ -355,11 +328,21 @@ def _get_instance_types_df(region: str) -> Union[str, 'pd.DataFrame']:
         df = pd.concat(
             [df, df.apply(get_additional_columns, axis='columns')],
             axis='columns')
-        # patch the df for p4de.24xlarge
-        if region in P4DE_REGIONS:
-            df = _patch_p4de(region, df, pricing_df)
         if 'GpuInfo' not in df.columns:
             df['GpuInfo'] = np.nan
+        if 'NeuronInfo' in df.columns:
+            # The AWS Neuron API uses 'NeuronDevices' instead of 'Gpus'
+            # in its dict; for consistency with GPU handling, rename key.
+            def map_neuroninfo(neuroninfo):
+                if isinstance(neuroninfo,
+                              dict) and 'NeuronDevices' in neuroninfo:
+                    # Rename 'NeuronDevices' to 'Gpus'
+                    neuroninfo = neuroninfo.copy()
+                    neuroninfo['Gpus'] = neuroninfo.pop('NeuronDevices')
+                return neuroninfo
+
+            df['NeuronInfo'] = df['NeuronInfo'].apply(map_neuroninfo)
+            df['GpuInfo'] = df['GpuInfo'].fillna(df['NeuronInfo'])
         df = df[USEFUL_COLUMNS]
     except Exception as e:  # pylint: disable=broad-except
         print(traceback.format_exc())
@@ -407,44 +390,70 @@ def get_all_regions_instance_types_df(regions: Set[str]) -> 'pd.DataFrame':
 # TODO(tian): find out the driver version.
 #   Neuron driver:
 _GPU_DESC_UBUNTU_DATE = [
-    ('gpu', 'AMI GPU PyTorch 2.1.0', '20.04', '20231103'),
-    ('gpu', 'AMI GPU PyTorch 1.10.0', '18.04', '20221114'),
-    ('k80', 'AMI GPU PyTorch 1.10.0', '20.04', '20211208'),
-    ('k80', 'AMI GPU PyTorch 1.10.0', '18.04', '20211208'),
-    ('neuron', 'Base Neuron AMI', '22.04', '20240923'),
+    ('neuron', '/aws/service/neuron/dlami/multi-framework', '22.04'),
 ]
 
 
-def _fetch_image_id(region: str, description: str, ubuntu_version: str,
-                    creation_date: str) -> Optional[str]:
+def _fetch_image_creation_date(region: str,
+                               image_id: Optional[str]) -> Optional[str]:
+    if image_id is None:
+        return None
     try:
         image = subprocess.check_output(f"""\
-            aws ec2 describe-images --region {region} --owners amazon \\
-                --filters 'Name=name,Values="Deep Learning {description} (Ubuntu {ubuntu_version}) {creation_date}"' \\
-                    'Name=state,Values=available' --query 'Images[:1].ImageId' --output text
+            aws ec2 describe-images --region {region} --image-ids {image_id} \\
+                --query 'Images[0].Name' --output text
             """,
                                         shell=True)
     except subprocess.CalledProcessError as e:
-        print(f'Failed {region}, {description}, {ubuntu_version}, '
-              f'{creation_date}. Trying next date.')
+        print(f'Failed to fetch image creation date for {region}, {image_id}')
         print(f'{type(e)}: {e}')
         image_id = None
+    else:
+        assert image is not None
+        image_name = image.decode('utf-8').strip()
+        match = re.search(r'(\d+)$', image_name)
+        if match:
+            return match.group(1)
+    return None
+
+
+def _fetch_image_id_from_ssm_param(
+        region: str,
+        ssm_prefix: str,
+        ubuntu_version: str = '22.04') -> Optional[str]:
+    try:
+        image = subprocess.check_output(f"""\
+            aws ssm get-parameter --region {region} --name "{ssm_prefix}/ubuntu-{ubuntu_version}/latest/image_id" \\
+                --query 'Parameter.Value' --output text
+            """,
+                                        shell=True)
+    except subprocess.CalledProcessError as e:
+        print(
+            f'Failed to fetch image ID from SSM parameter for {region}, {ssm_prefix}, {ubuntu_version}'
+        )
+        print(f'{type(e)}: {e}')
+        return None
     else:
         assert image is not None
         image_id = image.decode('utf-8').strip()
     return image_id
 
 
-def _get_image_row(region: str, gpu: str, description: str, ubuntu_version: str,
-                   date: str) -> Tuple[str, str, str, str, Optional[str], str]:
-    print(f'Getting image for {region}, {description}, {ubuntu_version}, {gpu}')
-    image_id = _fetch_image_id(region, description, ubuntu_version, date)
-    if image_id is None:
-        # not found
-        print(f'Failed to find image for {region}, {description}, '
-              f'{ubuntu_version}, {gpu}')
+def _get_image_row(
+    region: str,
+    gpu: str,
+    ssm_prefix: str,
+    ubuntu_version: str = '22.04'
+) -> Tuple[str, str, str, str, Optional[str], Optional[str]]:
+    print(f'Getting image for {region}, {ssm_prefix}, {ubuntu_version}, {gpu}')
+    image_id = _fetch_image_id_from_ssm_param(region, ssm_prefix,
+                                              ubuntu_version)
+    if image_id is not None:
+        creation_date = _fetch_image_creation_date(region, image_id)
+    else:
+        creation_date = None
     tag = f'skypilot:{gpu}-ubuntu-{ubuntu_version.replace(".", "")}'
-    return tag, region, 'ubuntu', ubuntu_version, image_id, date
+    return tag, region, 'ubuntu', ubuntu_version, image_id, creation_date
 
 
 def get_all_regions_images_df(regions: Set[str]) -> 'pd.DataFrame':
@@ -559,13 +568,26 @@ if __name__ == '__main__':
     instance_df.to_csv('aws/vms.csv', index=False)
     print('AWS Service Catalog saved to aws/vms.csv')
 
-    # Disable refreshing images.csv as we are using skypilot custom AMIs
+    # Disable refreshing images.csv for skypilot custom AMIs
+    # refresh only the neuron based images
     # See sky/clouds/catalog/images/README.md for more details.
-    # image_df = get_all_regions_images_df(user_regions)
-    # _check_regions_integrity(image_df, 'images')
-
-    # image_df.to_csv('aws/images.csv', index=False)
-    # print('AWS Images saved to aws/images.csv')
+    image_df = get_all_regions_images_df(user_regions)
+    _check_regions_integrity(image_df, 'images')
+    # filter out rows where ImageId is None
+    image_df = image_df[image_df['ImageId'].notna()]
+
+    # check if aws/images.csv exists
+    if os.path.exists('aws/images.csv'):
+        # load the data from aws/images.csv
+        existing_image_df = pd.read_csv('aws/images.csv')
+        # filter out the neuron based images
+        existing_image_df = existing_image_df[~existing_image_df['Tag'].
+                                              eq('skypilot:neuron-ubuntu-2204')]
+        # concat the new neuron based images with the existing images
+        image_df = pd.concat([existing_image_df, image_df])
+
+    image_df.to_csv('aws/images.csv', index=False)
+    print('AWS Images saved to aws/images.csv')
 
     if args.az_mappings:
         az_mappings_df = fetch_availability_zone_mappings()

sky/catalog/data_fetchers/fetch_gcp.py

@@ -182,8 +182,9 @@ TPU_V4_HOST_DF = pd.read_csv(
 SERIES_TO_DESCRIPTION = {
     'a2': 'A2 Instance',
     'a3': 'A3 Instance',
-    # TODO(zhwu): GCP does not have A4 instance in SKUs API yet. We keep it here
-    # for completeness.
+    # NOTE: GCP does not provide separate CPU/RAM pricing for A4 instances.
+    # The B200 GPU pricing includes the full VM cost. See special handling in
+    # get_vm_price() which sets A4 VM price to 0.
     'a4': 'A4 Instance',
     'c2': 'Compute optimized',
     'c2d': 'C2D AMD Instance',
@@ -394,6 +395,15 @@ def get_vm_df(skus: List[Dict[str, Any]], region_prefix: str) -> 'pd.DataFrame':
         if series in ['f1', 'g1']:
             memory_price = 0.0
 
+        # Special case for A4 instances.
+        # GCP does not provide separate CPU/RAM pricing for A4 instances in the
+        # SKUs API. The GPU pricing (B200) includes the full VM cost.
+        # We set the VM price to 0 so the entry is not dropped, and the GPU
+        # pricing will provide the total cost.
+        if series == 'a4':
+            cpu_price = 0.0
+            memory_price = 0.0
+
         # TODO(tian): (2024/11/10) Some SKUs are missing in the SKUs API. We
         # skip them in the catalog for now. We should investigate why they are
         # missing and add them back.
@@ -525,7 +535,24 @@ def get_gpu_df(skus: List[Dict[str, Any]],
             row_gpu_name = row['AcceleratorName']
             if row['Region'] not in sku['serviceRegions']:
                 continue
-            if sku['category']['usageType'] != ondemand_or_spot:
+
+            # Check usageType matches, with special handling for B200 spot.
+            # GCP has a bug where some B200 spot SKUs have usageType='OnDemand'
+            # but the description contains 'Spot Preemptible'.
+            usage_type = sku['category']['usageType']
+            description = sku['description']
+            is_spot_description = 'spot preemptible' in description.lower()
+
+            if usage_type != ondemand_or_spot:
+                # For B200 spot pricing, also accept SKUs where description
+                # says "Spot Preemptible" even if usageType is wrong.
+                if not (spot and row_gpu_name == 'B200' and
+                        is_spot_description):
+                    continue
+
+            # For B200 on-demand, skip SKUs that are actually spot (description
+            # says "Spot Preemptible" but usageType is incorrectly 'OnDemand').
+            if not spot and row_gpu_name == 'B200' and is_spot_description:
                 continue
 
             gpu_names = [f'{row_gpu_name} GPU']

sky/catalog/data_fetchers/fetch_nebius.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 import decimal
 import json
 import logging
+import os
 import re
 from typing import Any, Dict, List, Optional
 
@@ -22,8 +23,6 @@ TIMEOUT = 10
 PARENT_ID_TEMPLATE = 'project-{}public-images'
 ACCELERATOR_MANUFACTURER = 'NVIDIA'
 
-VRAM = {'L40S': 49152, 'H100': 81920, 'H200': 144384, 'B200': 184320}
-
 
 @dataclass
 class PresetInfo:
@@ -38,6 +37,7 @@ class PresetInfo:
         platform_name (str): The name of the platform the preset belongs to.
         gpu (int): The number of GPUs in the preset.
         vcpu (int): The number of virtual CPUs in the preset.
+        gpu_memory_gibibytes (int): size of gpu memory in GiB.
         memory_gib (int): The amount of memory in GiB in the preset.
         accelerator_manufacturer (str | None): The manufacturer of the
             accelerator (e.g., "NVIDIA"), or None if no accelerator.
@@ -54,6 +54,7 @@ class PresetInfo:
     platform_name: str
     gpu: int
     vcpu: int
+    gpu_memory_gibibytes: int
     memory_gib: int
     accelerator_manufacturer: Optional[str]
     accelerator_name: Optional[str]
@@ -157,6 +158,7 @@ def _estimate_platforms(platforms: List[Any], parent_id: str,
                 platform_name=platform_name,
                 gpu=preset.resources.gpu_count or 0,
                 vcpu=preset.resources.vcpu_count,
+                gpu_memory_gibibytes=platform.spec.gpu_memory_gibibytes,
                 memory_gib=preset.resources.memory_gibibytes,
                 accelerator_manufacturer=ACCELERATOR_MANUFACTURER
                 if platform_name.startswith('gpu-') else '',
@@ -178,6 +180,7 @@ def _write_preset_prices(presets: List[PresetInfo], output_file: str) -> None:
         presets (List[PresetInfo]): A list of PresetInfo objects to write.
         output_file (str): The path to the output CSV file.
     """
+    os.makedirs(os.path.dirname(output_file))
     # Set up the CSV writer to output to stdout
     with open(output_file, 'w', encoding='utf-8') as out:
         header = [
@@ -193,23 +196,23 @@ def _write_preset_prices(presets: List[PresetInfo], output_file: str) -> None:
         ]
         writer = csv.DictWriter(out, fieldnames=header)
         writer.writeheader()
-
+        # logger.info(presets)
         for preset in sorted(presets,
                              key=lambda x:
                              (bool(x.gpu), x.region, x.platform_name, x.vcpu)):
             gpu_info = ''
             if preset.gpu > 0 and preset.accelerator_name:
+                vram = preset.gpu_memory_gibibytes * 1024
                 gpu_info_dict = {
                     'Gpus': [{
                         'Name': preset.accelerator_name,
                         'Manufacturer': preset.accelerator_manufacturer,
                         'Count': preset.gpu,
                         'MemoryInfo': {
-                            'SizeInMiB': VRAM.get(preset.accelerator_name, 0)
+                            'SizeInMiB': vram
                         },
                     }],
-                    'TotalGpuMemoryInMiB': VRAM.get(preset.accelerator_name, 0)
-                                           * preset.gpu,
+                    'TotalGpuMemoryInMiB': vram * preset.gpu,
                 }
                 gpu_info = json.dumps(gpu_info_dict).replace('"', '\'')