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

sky/provision/common.py

@@ -1,9 +1,11 @@
 """Common data structures for provisioning"""
 import abc
 import dataclasses
+import functools
 import os
 from typing import Any, Dict, List, Optional, Tuple
 
+from sky import sky_logging
 from sky.utils import resources_utils
 
 # NOTE: we can use pydantic instead of dataclasses or namedtuples, because
@@ -14,6 +16,10 @@ from sky.utils import resources_utils
 # -------------------- input data model -------------------- #
 
 InstanceId = str
+_START_TITLE = '\n' + '-' * 20 + 'Start: {} ' + '-' * 20
+_END_TITLE = '-' * 20 + 'End:   {} ' + '-' * 20 + '\n'
+
+logger = sky_logging.init_logger(__name__)
 
 
 class ProvisionerError(RuntimeError):
@@ -46,6 +52,8 @@ class ProvisionConfig:
     tags: Dict[str, str]
     # Whether or not to resume stopped instances.
     resume_stopped_nodes: bool
+    # Optional ports to open on launch of the cluster.
+    ports_to_open_on_launch: Optional[List[int]]
 
 
 # -------------------- output data model -------------------- #
@@ -123,7 +131,8 @@ class ClusterInfo:
         if self.head_instance_id is None:
             return None
         if self.head_instance_id not in self.instances:
-            raise ValueError('Head instance ID not in the cluster metadata.')
+            raise ValueError('Head instance ID not in the cluster metadata. '
+                             f'ClusterInfo: {self.__dict__}')
         return self.instances[self.head_instance_id][0]
 
     def get_worker_instances(self) -> List[InstanceInfo]:
@@ -197,8 +206,14 @@ class ClusterInfo:
         return ip_list
 
     def get_feasible_ips(self, force_internal_ips: bool = False) -> List[str]:
-        """Get external IPs if they exist, otherwise get internal ones."""
-        return self._get_ips(not self.has_external_ips() or force_internal_ips)
+        """Get internal or external IPs depends on the settings."""
+        if self.provider_config is not None:
+            use_internal_ips = self.provider_config.get('use_internal_ips',
+                                                        False)
+        else:
+            use_internal_ips = False
+        return self._get_ips(use_internal_ips or not self.has_external_ips() or
+                             force_internal_ips)
 
     def get_ssh_ports(self) -> List[int]:
         """Get the SSH port of all the instances."""
@@ -268,3 +283,16 @@ def query_ports_passthrough(
     for port in ports:
         result[port] = [SocketEndpoint(port=port, host=head_ip)]
     return result
+
+
+def log_function_start_end(func):
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        logger.info(_START_TITLE.format(func.__name__))
+        try:
+            return func(*args, **kwargs)
+        finally:
+            logger.info(_END_TITLE.format(func.__name__))
+
+    return wrapper

sky/provision/constants.py

@@ -0,0 +1,25 @@
+"""Constants used in the SkyPilot provisioner."""
+
+# Tag uniquely identifying all nodes of a cluster
+TAG_RAY_CLUSTER_NAME = 'ray-cluster-name'
+TAG_SKYPILOT_CLUSTER_NAME = 'skypilot-cluster-name'
+# Legacy tag for backward compatibility to distinguish head and worker nodes.
+TAG_RAY_NODE_KIND = 'ray-node-type'
+TAG_SKYPILOT_HEAD_NODE = 'skypilot-head-node'
+
+HEAD_NODE_TAGS = {
+    TAG_RAY_NODE_KIND: 'head',
+    TAG_SKYPILOT_HEAD_NODE: '1',
+}
+
+WORKER_NODE_TAGS = {
+    TAG_RAY_NODE_KIND: 'worker',
+    TAG_SKYPILOT_HEAD_NODE: '0',
+}
+
+# Names for Azure Deployments.
+DEPLOYMENT_NAME = 'skypilot-config'
+LEGACY_DEPLOYMENT_NAME = 'ray-config'
+EXTERNAL_RG_BOOTSTRAP_DEPLOYMENT_NAME = (
+    'skypilot-bootstrap-{cluster_name_on_cloud}')
+EXTERNAL_RG_VM_DEPLOYMENT_NAME = 'skypilot-vm-{cluster_name_on_cloud}'

sky/provision/cudo/__init__.py

@@ -3,6 +3,7 @@
 from sky.provision.cudo.config import bootstrap_instances
 from sky.provision.cudo.instance import cleanup_ports
 from sky.provision.cudo.instance import get_cluster_info
+from sky.provision.cudo.instance import open_ports
 from sky.provision.cudo.instance import query_instances
 from sky.provision.cudo.instance import run_instances
 from sky.provision.cudo.instance import stop_instances
@@ -11,4 +12,4 @@ from sky.provision.cudo.instance import wait_instances
 
 __all__ = ('bootstrap_instances', 'run_instances', 'stop_instances',
            'terminate_instances', 'wait_instances', 'get_cluster_info',
-           'cleanup_ports', 'query_instances')
+           'cleanup_ports', 'query_instances', 'open_ports')

sky/provision/cudo/cudo_utils.py

@@ -0,0 +1,112 @@
+"""Cudo catalog helper."""
+
+cudo_gpu_model = {
+    'NVIDIA V100': 'V100',
+    'NVIDIA A40': 'A40',
+    'RTX 3080': 'RTX3080',
+    'RTX A4000': 'RTXA4000',
+    'RTX A4500': 'RTXA4500',
+    'RTX A5000': 'RTXA5000',
+    'RTX A6000': 'RTXA6000',
+}
+
+cudo_gpu_mem = {
+    'RTX3080': 12,
+    'A40': 48,
+    'RTXA4000': 16,
+    'RTXA4500': 20,
+    'RTXA5000': 24,
+    'RTXA6000': 48,
+    'V100': 16,
+}
+
+machine_specs = [
+    # Low
+    {
+        'vcpu': 2,
+        'mem': 4,
+        'gpu': 1,
+    },
+    {
+        'vcpu': 4,
+        'mem': 8,
+        'gpu': 1,
+    },
+    {
+        'vcpu': 8,
+        'mem': 16,
+        'gpu': 2,
+    },
+    {
+        'vcpu': 16,
+        'mem': 32,
+        'gpu': 2,
+    },
+    {
+        'vcpu': 32,
+        'mem': 64,
+        'gpu': 4,
+    },
+    {
+        'vcpu': 64,
+        'mem': 128,
+        'gpu': 8,
+    },
+    # Mid
+    {
+        'vcpu': 96,
+        'mem': 192,
+        'gpu': 8
+    },
+    {
+        'vcpu': 48,
+        'mem': 96,
+        'gpu': 4
+    },
+    {
+        'vcpu': 24,
+        'mem': 48,
+        'gpu': 2
+    },
+    {
+        'vcpu': 12,
+        'mem': 24,
+        'gpu': 1
+    },
+    # Hi
+    {
+        'vcpu': 96,
+        'mem': 192,
+        'gpu': 4
+    },
+    {
+        'vcpu': 48,
+        'mem': 96,
+        'gpu': 2
+    },
+    {
+        'vcpu': 24,
+        'mem': 48,
+        'gpu': 1
+    },
+]
+
+
+def cudo_gpu_to_skypilot_gpu(model):
+    if model in cudo_gpu_model:
+        return cudo_gpu_model[model]
+    else:
+        return model
+
+
+def skypilot_gpu_to_cudo_gpu(model):
+    for key, value in cudo_gpu_model.items():
+        if value == model:
+            return key
+    return model
+
+
+def gpu_exists(model):
+    if model in cudo_gpu_model:
+        return True
+    return False

sky/provision/cudo/cudo_wrapper.py

@@ -4,29 +4,29 @@ from typing import Dict
 
 from sky import sky_logging
 from sky.adaptors import cudo
+import sky.provision.cudo.cudo_utils as utils
 
 logger = sky_logging.init_logger(__name__)
 
 
 def launch(name: str, data_center_id: str, ssh_key: str, machine_type: str,
-           memory_gib: int, vcpu_count: int, gpu_count: int, gpu_model: str,
+           memory_gib: int, vcpu_count: int, gpu_count: int,
            tags: Dict[str, str], disk_size: int):
     """Launches an instance with the given parameters."""
-    disk = cudo.cudo.Disk(storage_class='STORAGE_CLASS_NETWORK',
-                          size_gib=disk_size)
-
-    request = cudo.cudo.CreateVMBody(ssh_key_source='SSH_KEY_SOURCE_NONE',
-                                     custom_ssh_keys=[ssh_key],
-                                     vm_id=name,
-                                     machine_type=machine_type,
-                                     data_center_id=data_center_id,
-                                     boot_disk_image_id='ubuntu-nvidia-docker',
-                                     memory_gib=memory_gib,
-                                     vcpus=vcpu_count,
-                                     gpus=gpu_count,
-                                     gpu_model=gpu_model,
-                                     boot_disk=disk,
-                                     metadata=tags)
+
+    request = cudo.cudo.CreateVMBody(
+        ssh_key_source='SSH_KEY_SOURCE_NONE',
+        custom_ssh_keys=[ssh_key],
+        vm_id=name,
+        machine_type=machine_type,
+        data_center_id=data_center_id,
+        boot_disk_image_id='ubuntu-2204-nvidia-535-docker-v20240214',
+        memory_gib=memory_gib,
+        vcpus=vcpu_count,
+        gpus=gpu_count,
+        boot_disk=cudo.cudo.Disk(storage_class='STORAGE_CLASS_NETWORK',
+                                 size_gib=disk_size),
+        metadata=tags)
 
     try:
         api = cudo.cudo.cudo_api.virtual_machines()
@@ -121,3 +121,24 @@ def list_instances():
         return instances
     except cudo.cudo.rest.ApiException as e:
         raise e
+
+
+def vm_available(to_start_count, gpu_count, gpu_model, data_center_id, mem,
+                 cpus):
+    try:
+        gpu_model = utils.skypilot_gpu_to_cudo_gpu(gpu_model)
+        api = cudo.cudo.cudo_api.virtual_machines()
+        types = api.list_vm_machine_types(mem,
+                                          cpus,
+                                          gpu=gpu_count,
+                                          gpu_model=gpu_model,
+                                          data_center_id=data_center_id)
+        types_dict = types.to_dict()
+        hc = types_dict['host_configs']
+        total_count = sum(item['count_vm_available'] for item in hc)
+        if total_count < to_start_count:
+            raise Exception(
+                'Too many VMs requested, try another gpu type or region')
+        return total_count
+    except cudo.cudo.rest.ApiException as e:
+        raise e

sky/provision/cudo/instance.py

@@ -4,10 +4,10 @@ import time
 from typing import Any, Dict, List, Optional
 
 from sky import sky_logging
-from sky import status_lib
 from sky.provision import common
 from sky.provision.cudo import cudo_machine_type
 from sky.provision.cudo import cudo_wrapper
+from sky.utils import status_lib
 
 POLL_INTERVAL = 10
 
@@ -16,7 +16,6 @@ logger = sky_logging.init_logger(__name__)
 
 def _filter_instances(cluster_name_on_cloud: str,
                       status_filters: Optional[List[str]]) -> Dict[str, Any]:
-
     instances = cudo_wrapper.list_instances()
     possible_names = [
         f'{cluster_name_on_cloud}-head', f'{cluster_name_on_cloud}-worker'
@@ -77,10 +76,19 @@ def run_instances(region: str, cluster_name_on_cloud: str,
 
     created_instance_ids = []
     public_key = config.node_config['AuthorizedKey']
-
+    instance_type = config.node_config['InstanceType']
+    spec = cudo_machine_type.get_spec_from_instance(instance_type, region)
+    gpu_count = int(float(spec['gpu_count']))
+    vcpu_count = int(spec['vcpu_count'])
+    memory_gib = int(spec['mem_gb'])
+    gpu_model = spec['gpu_model']
+    try:
+        cudo_wrapper.vm_available(to_start_count, gpu_count, gpu_model, region,
+                                  memory_gib, vcpu_count)
+    except Exception as e:
+        logger.warning(f'run_instances: {e}')
+        raise
     for _ in range(to_start_count):
-        instance_type = config.node_config['InstanceType']
-        spec = cudo_machine_type.get_spec_from_instance(instance_type, region)
 
         node_type = 'head' if head_instance_id is None else 'worker'
         try:
@@ -89,10 +97,9 @@ def run_instances(region: str, cluster_name_on_cloud: str,
                 ssh_key=public_key,
                 data_center_id=region,
                 machine_type=spec['machine_type'],
-                memory_gib=int(spec['mem_gb']),
-                vcpu_count=int(spec['vcpu_count']),
-                gpu_count=int(float(spec['gpu_count'])),
-                gpu_model=spec['gpu_model'],
+                memory_gib=memory_gib,
+                vcpu_count=vcpu_count,
+                gpu_count=gpu_count,
                 tags={},
                 disk_size=config.node_config['DiskSize'])
         except Exception as e:  # pylint: disable=broad-except
@@ -150,11 +157,10 @@ def terminate_instances(
     del provider_config
     instances = _filter_instances(cluster_name_on_cloud, None)
     for inst_id, inst in instances.items():
-        logger.info(f'Terminating instance {inst_id}.'
-                    f'{inst}')
         if worker_only and inst['name'].endswith('-head'):
             continue
-        logger.info(f'Removing {inst_id}: {inst}')
+        logger.debug(f'Terminating Cudo instance {inst_id}.'
+                     f'{inst}')
         cudo_wrapper.remove(inst_id)
 
 
@@ -213,6 +219,16 @@ def query_instances(
     return statuses
 
 
+def open_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    del cluster_name_on_cloud, ports, provider_config
+    # Cudo has all ports open by default. Nothing to do here.
+    return
+
+
 def cleanup_ports(
     cluster_name_on_cloud: str,
     ports: List[str],

sky/provision/do/__init__.py

@@ -0,0 +1,11 @@
+"""DO provisioner for SkyPilot."""
+
+from sky.provision.do.config import bootstrap_instances
+from sky.provision.do.instance import cleanup_ports
+from sky.provision.do.instance import get_cluster_info
+from sky.provision.do.instance import open_ports
+from sky.provision.do.instance import query_instances
+from sky.provision.do.instance import run_instances
+from sky.provision.do.instance import stop_instances
+from sky.provision.do.instance import terminate_instances
+from sky.provision.do.instance import wait_instances

sky/provision/do/config.py

@@ -0,0 +1,14 @@
+"""Paperspace configuration bootstrapping."""
+
+from sky import sky_logging
+from sky.provision import common
+
+logger = sky_logging.init_logger(__name__)
+
+
+def bootstrap_instances(
+        region: str, cluster_name: str,
+        config: common.ProvisionConfig) -> common.ProvisionConfig:
+    """Bootstraps instances for the given cluster."""
+    del region, cluster_name
+    return config

sky/provision/do/constants.py

@@ -0,0 +1,10 @@
+"""DO cloud constants
+"""
+
+POLL_INTERVAL = 5
+WAIT_DELETE_VOLUMES = 5
+
+GPU_IMAGES = {
+    'gpu-h100x1-80gb': 'gpu-h100x1-base',
+    'gpu-h100x8-640gb': 'gpu-h100x8-base',
+}