Preliminary Vast AI support

docs/source/getting-started/installation.rst

@@ -297,6 +297,16 @@ Paperspace
   mkdir -p ~/.paperspace
   echo "{'api_key' : <your_api_key_here>}" > ~/.paperspace/config.json
 
+Vast
+~~~~~~~~~~
+
+`Vast <https://vast.ai/>`__ is a cloud provider that offers low-cost GPUs. To configure Vast access, go to the `Account <https://cloud.vast.ai/account/>`_ page on your Vast console to get your **API key**. Then, run:
+
+.. code-block:: shell
+
+  pip install "vastai-sdk>=0.1.3"
+  echo "<your_api_key_here>" > ~/.vast_api_key
+
 RunPod
 ~~~~~~~~~~
 

sky/__init__.py

@@ -133,6 +133,7 @@ def set_proxy_env_var(proxy_var: str, urllib_var: Optional[str]):
 OCI = clouds.OCI
 Paperspace = clouds.Paperspace
 RunPod = clouds.RunPod
+Vast = clouds.Vast
 Vsphere = clouds.Vsphere
 Fluidstack = clouds.Fluidstack
 optimize = Optimizer.optimize
@@ -150,6 +151,7 @@ def set_proxy_env_var(proxy_var: str, urllib_var: Optional[str]):
     'OCI',
     'Paperspace',
     'RunPod',
+    'Vast',
     'SCP',
     'Vsphere',
     'Fluidstack',

sky/adaptors/vast.py

@@ -0,0 +1,29 @@
+"""Vast cloud adaptor."""
+
+import functools
+
+_vast_sdk = None
+
+
+def import_package(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global _vast_sdk
+
+        if _vast_sdk is None:
+            try:
+                import vastai_sdk as _vast  # pylint: disable=import-outside-toplevel
+                _vast_sdk = _vast.VastAI()
+            except ImportError:
+                raise ImportError('Fail to import dependencies for vast.'
+                                  'Try pip install "skypilot[vast]"') from None
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+@import_package
+def vast():
+    """Return the vast package."""
+    return _vast_sdk

sky/authentication.py

@@ -43,6 +43,7 @@
 from sky.adaptors import ibm
 from sky.adaptors import kubernetes
 from sky.adaptors import runpod
+from sky.adaptors import vast
 from sky.provision.fluidstack import fluidstack_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.provision.lambda_cloud import lambda_utils
@@ -485,6 +486,23 @@ def setup_runpod_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     return configure_ssh_info(config)
 
 
+def setup_vast_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Sets up SSH authentication for Vast.
+    - Generates a new SSH key pair if one does not exist.
+    - Adds the public SSH key to the user's Vast account.
+    """
+    _, public_key_path = get_or_generate_keys()
+    with open(public_key_path, 'r', encoding='UTF-8') as pub_key_file:
+        public_key = pub_key_file.read().strip()
+        current_key_list = vast.vast().show_ssh_keys()  # pylint: disable=assignment-from-no-return
+        # Only add an ssh key if it hasn't already been added
+        if not any(x['public_key'] == public_key for x in current_key_list):
+            vast.vast().create_ssh_key(ssh_key=public_key)
+
+    config['auth']['ssh_public_key'] = PUBLIC_SSH_KEY_PATH
+    return configure_ssh_info(config)
+
+
 def setup_fluidstack_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
 
     get_or_generate_keys()

sky/backends/backend_utils.py

@@ -1056,6 +1056,8 @@ def _add_auth_to_cluster_config(cloud: clouds.Cloud, cluster_config_file: str):
         config = auth.setup_ibm_authentication(config)
     elif isinstance(cloud, clouds.RunPod):
         config = auth.setup_runpod_authentication(config)
+    elif isinstance(cloud, clouds.Vast):
+        config = auth.setup_vast_authentication(config)
     elif isinstance(cloud, clouds.Fluidstack):
         config = auth.setup_fluidstack_authentication(config)
     else:
@@ -2135,7 +2137,8 @@ def run_ray_status_to_check_ray_cluster_healthy() -> bool:
                 except exceptions.CommandError as e:
                     success = False
                     if e.returncode == 255:
-                        logger.debug(f'The cluster is likely {noun}ed.')
+                        word = 'autostopped' if noun == 'autostop' else 'autodowned'
+                        logger.debug(f'The cluster is likely {word}.')
                         reset_local_autostop = False
                 except (Exception, SystemExit) as e:  # pylint: disable=broad-except
                     success = False

sky/backends/cloud_vm_ray_backend.py

@@ -182,6 +182,7 @@ def _get_cluster_config_template(cloud):
         clouds.RunPod: 'runpod-ray.yml.j2',
         clouds.Kubernetes: 'kubernetes-ray.yml.j2',
         clouds.Vsphere: 'vsphere-ray.yml.j2',
+        clouds.Vast: 'vast-ray.yml.j2',
         clouds.Fluidstack: 'fluidstack-ray.yml.j2'
     }
     return cloud_to_template[type(cloud)]

sky/clouds/__init__.py

@@ -25,6 +25,7 @@
 from sky.clouds.paperspace import Paperspace
 from sky.clouds.runpod import RunPod
 from sky.clouds.scp import SCP
+from sky.clouds.vast import Vast
 from sky.clouds.vsphere import Vsphere
 
 __all__ = [
@@ -39,6 +40,7 @@
     'Paperspace',
     'SCP',
     'RunPod',
+    'Vast',
     'OCI',
     'Vsphere',
     'Kubernetes',

sky/clouds/service_catalog/constants.py

@@ -3,5 +3,5 @@
 CATALOG_SCHEMA_VERSION = 'v6'
 CATALOG_DIR = '~/.sky/catalogs'
 ALL_CLOUDS = ('aws', 'azure', 'gcp', 'ibm', 'lambda', 'scp', 'oci',
-              'kubernetes', 'runpod', 'vsphere', 'cudo', 'fluidstack',
+              'kubernetes', 'runpod', 'vast', 'vsphere', 'cudo', 'fluidstack',
               'paperspace', 'do')

sky/clouds/service_catalog/data_fetchers/fetch_vast.py

@@ -0,0 +1,108 @@
+"""A script that generates the Vast Cloud catalog. """
+
+#
+# Due to the design of the sdk, pylint has a false
+# positive for the fnctions.
+#
+# pylint: disable=assignment-from-no-return
+import collections
+import csv
+import json
+import math
+import re
+import sys
+from typing import Any, Dict, List
+
+from sky.adaptors import vast
+
+_map = {
+    'TeslaV100': 'V100',
+    'TeslaT4': 'T4',
+    'TeslaP100': 'P100',
+    'QRTX6000': 'RTX6000',
+    'QRTX8000': 'RTX8000'
+}
+
+
+def create_instance_type(obj: Dict[str, Any]) -> str:
+    stubify = lambda x: re.sub(r'\s', '_', x)
+    return '{}x-{}-{}-{}'.format(obj['num_gpus'], stubify(obj['gpu_name']),
+                                 obj['cpu_cores'], obj['cpu_ram'])
+
+
+def dot_get(d: dict, key: str) -> Any:
+    for k in key.split('.'):
+        d = d[k]
+    return d
+
+
+if __name__ == '__main__':
+    # InstanceType and gpuInfo are basically just stubs
+    # so that the dictwriter is happy without weird
+    # code.
+    mapped_keys = (('gpu_name', 'InstanceType'), ('gpu_name',
+                                                  'AcceleratorName'),
+                   ('num_gpus', 'AcceleratorCount'), ('cpu_cores', 'vCPUs'),
+                   ('cpu_ram', 'MemoryGiB'), ('gpu_name', 'GpuInfo'),
+                   ('search.totalHour', 'Price'), ('min_bid', 'SpotPrice'),
+                   ('geolocation', 'Region'))
+    writer = csv.DictWriter(sys.stdout, fieldnames=[x[1] for x in mapped_keys])
+    writer.writeheader()
+
+    offerList = vast.vast().search_offers(limit=10000)
+    priceMap: Dict[str, List] = collections.defaultdict(list)
+    for offer in offerList:
+        entry = {}
+        for ours, theirs in mapped_keys:
+            field = dot_get(offer, ours)
+            entry[theirs] = field
+
+        instance_type = create_instance_type(offer)
+        entry['InstanceType'] = instance_type
+
+        # the documentation says
+        # "{'gpus': [{
+        #   'name': 'v100',
+        #   'manufacturer': 'nvidia',
+        #   'count': 8.0,
+        #   'memoryinfo': {'sizeinmib': 16384}
+        #   }],
+        #   'totalgpumemoryinmib': 16384}",
+        # we can do that.
+        entry['MemoryGiB'] /= 1024
+
+        gpu = re.sub('Ada', '-Ada', re.sub(r'\s', '', offer['gpu_name']))
+        gpu = re.sub(r'(Ti|PCIE|SXM4|SXM|NVL)$', '', gpu)
+        gpu = re.sub(r'(RTX\d0\d0)(S|D)$', r'\1', gpu)
+
+        if gpu in _map:
+            gpu = _map[gpu]
+
+        entry['AcceleratorName'] = gpu
+        entry['GpuInfo'] = json.dumps({
+            'Gpus': [{
+                'Name': gpu,
+                'Count': offer['num_gpus'],
+                'MemoryInfo': {
+                    'SizeInMiB': offer['gpu_total_ram']
+                }
+            }],
+            'TotalGpuMemoryInMiB': offer['gpu_total_ram']
+        }).replace('"', '\'')
+
+        priceMap[instance_type].append(entry)
+
+    for instanceList in priceMap.values():
+        priceList = sorted([x['Price'] for x in instanceList])
+        index = math.ceil(0.8 * len(priceList)) - 1
+        priceTarget = priceList[index]
+        toList: List = []
+        for instance in instanceList:
+            if instance['Price'] <= priceTarget:
+                instance['Price'] = '{:.2f}'.format(priceTarget)
+                toList.append(instance)
+
+        maxBid = max([x.get('SpotPrice') for x in toList])
+        for instance in toList:
+            instance['SpotPrice'] = '{:.2f}'.format(maxBid)
+            writer.writerow(instance)

sky/clouds/service_catalog/vast_catalog.py

@@ -0,0 +1,104 @@
+""" Vast | Catalog
+
+This module loads the service catalog file and can be used to
+query instance types and pricing information for Vast.ai.
+"""
+
+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('vast/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('Vast does not support zones.')
+    return common.validate_region_zone_impl('vast', _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('Vast 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]:
+    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)
+
+
+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('Vast 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)
+
+
+# TODO: this differs from the fluffy catalog version
+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 Vast offering GPUs."""
+    del require_price  # Unused.
+    return common.list_accelerators_impl('Vast', _df, gpus_only, name_filter,
+                                         region_filter, quantity_filter,
+                                         case_sensitive, all_regions)