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

sky/provision/do/instance.py

@@ -0,0 +1,287 @@
+"""DigitalOcean instance provisioning."""
+
+import time
+from typing import Any, Dict, List, Optional
+import uuid
+
+from sky import sky_logging
+from sky.provision import common
+from sky.provision.do import constants
+from sky.provision.do import utils
+from sky.utils import status_lib
+
+# The maximum number of times to poll for the status of an operation
+MAX_POLLS = 60 // constants.POLL_INTERVAL
+# Stopping instances can take several minutes, so we increase the timeout
+MAX_POLLS_FOR_UP_OR_STOP = MAX_POLLS * 8
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _get_head_instance(
+        instances: Dict[str, Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    for instance_name, instance_meta in instances.items():
+        if instance_name.endswith('-head'):
+            return instance_meta
+    return None
+
+
+def run_instances(region: str, cluster_name_on_cloud: str,
+                  config: common.ProvisionConfig) -> common.ProvisionRecord:
+    """Runs instances for the given cluster."""
+
+    pending_status = ['new']
+    newly_started_instances = utils.filter_instances(cluster_name_on_cloud,
+                                                     pending_status + ['off'])
+    while True:
+        instances = utils.filter_instances(cluster_name_on_cloud,
+                                           pending_status)
+        if not instances:
+            break
+        instance_statuses = [
+            instance['status'] for instance in instances.values()
+        ]
+        logger.info(f'Waiting for {len(instances)} instances to be ready: '
+                    f'{instance_statuses}')
+        time.sleep(constants.POLL_INTERVAL)
+
+    exist_instances = utils.filter_instances(cluster_name_on_cloud,
+                                             status_filters=pending_status +
+                                             ['active', 'off'])
+    if len(exist_instances) > config.count:
+        raise RuntimeError(
+            f'Cluster {cluster_name_on_cloud} already has '
+            f'{len(exist_instances)} nodes, but {config.count} are required.')
+
+    stopped_instances = utils.filter_instances(cluster_name_on_cloud,
+                                               status_filters=['off'])
+    for instance in stopped_instances.values():
+        utils.start_instance(instance)
+    for _ in range(MAX_POLLS_FOR_UP_OR_STOP):
+        instances = utils.filter_instances(cluster_name_on_cloud, ['off'])
+        if len(instances) == 0:
+            break
+        num_stopped_instances = len(stopped_instances)
+        num_restarted_instances = num_stopped_instances - len(instances)
+        logger.info(
+            f'Waiting for {num_restarted_instances}/{num_stopped_instances} '
+            'stopped instances to be restarted.')
+        time.sleep(constants.POLL_INTERVAL)
+    else:
+        msg = ('run_instances: Failed to restart all'
+               'instances possibly due to to capacity issue.')
+        logger.warning(msg)
+        raise RuntimeError(msg)
+
+    exist_instances = utils.filter_instances(cluster_name_on_cloud,
+                                             status_filters=['active'])
+    head_instance = _get_head_instance(exist_instances)
+    to_start_count = config.count - len(exist_instances)
+    if to_start_count < 0:
+        raise RuntimeError(
+            f'Cluster {cluster_name_on_cloud} already has '
+            f'{len(exist_instances)} nodes, but {config.count} are required.')
+    if to_start_count == 0:
+        if head_instance is None:
+            head_instance = list(exist_instances.values())[0]
+            utils.rename_instance(
+                head_instance,
+                f'{cluster_name_on_cloud}-{uuid.uuid4().hex[:4]}-head')
+        assert head_instance is not None, ('`head_instance` should not be None')
+        logger.info(f'Cluster {cluster_name_on_cloud} already has '
+                    f'{len(exist_instances)} nodes, no need to start more.')
+        return common.ProvisionRecord(
+            provider_name='do',
+            cluster_name=cluster_name_on_cloud,
+            region=region,
+            zone=None,
+            head_instance_id=head_instance['name'],
+            resumed_instance_ids=list(newly_started_instances.keys()),
+            created_instance_ids=[],
+        )
+
+    created_instances: List[Dict[str, Any]] = []
+    for _ in range(to_start_count):
+        instance_type = 'head' if head_instance is None else 'worker'
+        instance = utils.create_instance(
+            region=region,
+            cluster_name_on_cloud=cluster_name_on_cloud,
+            instance_type=instance_type,
+            config=config)
+        logger.info(f'Launched instance {instance["name"]}.')
+        created_instances.append(instance)
+        if head_instance is None:
+            head_instance = instance
+
+    # Wait for instances to be ready.
+    for _ in range(MAX_POLLS_FOR_UP_OR_STOP):
+        instances = utils.filter_instances(cluster_name_on_cloud,
+                                           status_filters=['active'])
+        logger.info('Waiting for instances to be ready: '
+                    f'({len(instances)}/{config.count}).')
+        if len(instances) == config.count:
+            break
+
+        time.sleep(constants.POLL_INTERVAL)
+    else:
+        # Failed to launch config.count of instances after max retries
+        msg = 'run_instances: Failed to create the instances'
+        logger.warning(msg)
+        raise RuntimeError(msg)
+    assert head_instance is not None, 'head_instance should not be None'
+    return common.ProvisionRecord(
+        provider_name='do',
+        cluster_name=cluster_name_on_cloud,
+        region=region,
+        zone=None,
+        head_instance_id=head_instance['name'],
+        resumed_instance_ids=list(stopped_instances.keys()),
+        created_instance_ids=[
+            instance['name'] for instance in created_instances
+        ],
+    )
+
+
+def wait_instances(region: str, cluster_name_on_cloud: str,
+                   state: Optional[status_lib.ClusterStatus]) -> None:
+    del region, cluster_name_on_cloud, state  # unused
+    # We already wait on ready state in `run_instances` no need
+
+
+def stop_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    del provider_config  # unused
+    all_instances = utils.filter_instances(cluster_name_on_cloud,
+                                           status_filters=None)
+    num_instances = len(all_instances)
+
+    # Request a stop on all instances
+    for instance_name, instance_meta in all_instances.items():
+        if worker_only and instance_name.endswith('-head'):
+            num_instances -= 1
+            continue
+        utils.stop_instance(instance_meta)
+
+    # Wait for instances to stop
+    for _ in range(MAX_POLLS_FOR_UP_OR_STOP):
+        all_instances = utils.filter_instances(cluster_name_on_cloud, ['off'])
+        if len(all_instances) >= num_instances:
+            break
+        time.sleep(constants.POLL_INTERVAL)
+    else:
+        raise RuntimeError(f'Maximum number of polls: '
+                           f'{MAX_POLLS_FOR_UP_OR_STOP} reached. '
+                           f'Instance {all_instances} is still not in '
+                           'STOPPED status.')
+
+
+def terminate_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    """See sky/provision/__init__.py"""
+    del provider_config  # unused
+    instances = utils.filter_instances(cluster_name_on_cloud,
+                                       status_filters=None)
+    for instance_name, instance_meta in instances.items():
+        logger.debug(f'Terminating instance {instance_name}')
+        if worker_only and instance_name.endswith('-head'):
+            continue
+        utils.down_instance(instance_meta)
+
+    for _ in range(MAX_POLLS_FOR_UP_OR_STOP):
+        instances = utils.filter_instances(cluster_name_on_cloud,
+                                           status_filters=None)
+        if len(instances) == 0 or len(instances) <= 1 and worker_only:
+            break
+        time.sleep(constants.POLL_INTERVAL)
+    else:
+        msg = ('Failed to delete all instances')
+        logger.warning(msg)
+        raise RuntimeError(msg)
+
+
+def get_cluster_info(
+    region: str,
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> common.ClusterInfo:
+    del region  # unused
+    running_instances = utils.filter_instances(cluster_name_on_cloud,
+                                               ['active'])
+    instances: Dict[str, List[common.InstanceInfo]] = {}
+    head_instance: Optional[str] = None
+    for instance_name, instance_meta in running_instances.items():
+        if instance_name.endswith('-head'):
+            head_instance = instance_name
+        for net in instance_meta['networks']['v4']:
+            if net['type'] == 'public':
+                instance_ip = net['ip_address']
+                break
+        instances[instance_name] = [
+            common.InstanceInfo(
+                instance_id=instance_meta['name'],
+                internal_ip=instance_ip,
+                external_ip=instance_ip,
+                ssh_port=22,
+                tags={},
+            )
+        ]
+
+    assert head_instance is not None, 'no head instance found'
+    return common.ClusterInfo(
+        instances=instances,
+        head_instance_id=head_instance,
+        provider_name='do',
+        provider_config=provider_config,
+    )
+
+
+def query_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    non_terminated_only: bool = True,
+) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+    """See sky/provision/__init__.py"""
+    # terminated instances are not retrieved by the
+    # API making `non_terminated_only` argument moot.
+    del non_terminated_only
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+    instances = utils.filter_instances(cluster_name_on_cloud,
+                                       status_filters=None)
+
+    status_map = {
+        'new': status_lib.ClusterStatus.INIT,
+        'archive': status_lib.ClusterStatus.INIT,
+        'active': status_lib.ClusterStatus.UP,
+        'off': status_lib.ClusterStatus.STOPPED,
+    }
+    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    for instance_meta in instances.values():
+        status = status_map[instance_meta['status']]
+        statuses[instance_meta['name']] = status
+    return statuses
+
+
+def open_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    """See sky/provision/__init__.py"""
+    logger.debug(
+        f'Skip opening ports {ports} for DigitalOcean instances, as all '
+        'ports are open by default.')
+    del cluster_name_on_cloud, provider_config, ports
+
+
+def cleanup_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    del cluster_name_on_cloud, provider_config, ports

sky/provision/do/utils.py

@@ -0,0 +1,301 @@
+"""DigitalOcean API client wrapper for SkyPilot.
+
+Example usage of `pydo` client library was mostly taken from here:
+https://github.com/digitalocean/pydo/blob/main/examples/poc_droplets_volumes_sshkeys.py
+"""
+
+import copy
+import os
+from typing import Any, Dict, List, Optional
+import urllib
+import uuid
+
+from sky import sky_logging
+from sky.adaptors import do
+from sky.provision import common
+from sky.provision import constants as provision_constants
+from sky.provision.do import constants
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+POSSIBLE_CREDENTIALS_PATHS = [
+    os.path.expanduser(
+        '~/Library/Application Support/doctl/config.yaml'),  # OS X
+    os.path.expanduser(
+        os.path.join(os.getenv('XDG_CONFIG_HOME', '~/.config/'),
+                     'doctl/config.yaml')),  # Linux
+]
+INITIAL_BACKOFF_SECONDS = 10
+MAX_BACKOFF_FACTOR = 10
+MAX_ATTEMPTS = 6
+SSH_KEY_NAME_ON_DO = f'sky-key-{common_utils.get_user_hash()}'
+
+CREDENTIALS_PATH = '~/.config/doctl/config.yaml'
+_client = None
+_ssh_key_id = None
+
+
+class DigitalOceanError(Exception):
+    pass
+
+
+def _init_client():
+    global _client, CREDENTIALS_PATH
+    assert _client is None
+    CREDENTIALS_PATH = None
+    credentials_found = 0
+    for path in POSSIBLE_CREDENTIALS_PATHS:
+        if os.path.exists(path):
+            CREDENTIALS_PATH = path
+            credentials_found += 1
+            logger.debug(f'Digital Ocean credential path found at {path}')
+    if not credentials_found > 1:
+        logger.debug('more than 1 credential file found')
+    if CREDENTIALS_PATH is None:
+        raise DigitalOceanError(
+            'no credentials file found from '
+            f'the following paths {POSSIBLE_CREDENTIALS_PATHS}')
+
+    # attempt default context
+    credentials = common_utils.read_yaml(CREDENTIALS_PATH)
+    default_token = credentials.get('access-token', None)
+    if default_token is not None:
+        try:
+            test_client = do.pydo.Client(token=default_token)
+            test_client.droplets.list()
+            logger.debug('trying `default` context')
+            _client = test_client
+            return _client
+        except do.exceptions().HttpResponseError:
+            pass
+
+    auth_contexts = credentials.get('auth-contexts', None)
+    if auth_contexts is not None:
+        for context, api_token in auth_contexts.items():
+            try:
+                test_client = do.pydo.Client(token=api_token)
+                test_client.droplets.list()
+                logger.debug(f'using {context} context')
+                _client = test_client
+                break
+            except do.exceptions().HttpResponseError:
+                continue
+        else:
+            raise DigitalOceanError(
+                'no valid api tokens found try '
+                'setting a new API token with `doctl auth init`')
+    return _client
+
+
+def client():
+    global _client
+    if _client is None:
+        _client = _init_client()
+    return _client
+
+
+def ssh_key_id(public_key: str):
+    global _ssh_key_id
+    if _ssh_key_id is None:
+        page = 1
+        paginated = True
+        while paginated:
+            try:
+                resp = client().ssh_keys.list(per_page=50, page=page)
+                for ssh_key in resp['ssh_keys']:
+                    if ssh_key['public_key'] == public_key:
+                        _ssh_key_id = ssh_key
+                        return _ssh_key_id
+            except do.exceptions().HttpResponseError as err:
+                raise DigitalOceanError(
+                    f'Error: {err.status_code} {err.reason}: '
+                    f'{err.error.message}') from err
+
+            pages = resp['links']
+            if 'pages' in pages and 'next' in pages['pages']:
+                pages = pages['pages']
+                parsed_url = urllib.parse.urlparse(pages['next'])
+                page = int(urllib.parse.parse_qs(parsed_url.query)['page'][0])
+            else:
+                paginated = False
+
+        request = {
+            'public_key': public_key,
+            'name': SSH_KEY_NAME_ON_DO,
+        }
+        _ssh_key_id = client().ssh_keys.create(body=request)['ssh_key']
+    return _ssh_key_id
+
+
+def _create_volume(request: Dict[str, Any]) -> Dict[str, Any]:
+    try:
+        resp = client().volumes.create(body=request)
+        volume = resp['volume']
+    except do.exceptions().HttpResponseError as err:
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+    else:
+        return volume
+
+
+def _create_droplet(request: Dict[str, Any]) -> Dict[str, Any]:
+    try:
+        resp = client().droplets.create(body=request)
+        droplet_id = resp['droplet']['id']
+
+        get_resp = client().droplets.get(droplet_id)
+        droplet = get_resp['droplet']
+    except do.exceptions().HttpResponseError as err:
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+    return droplet
+
+
+def create_instance(region: str, cluster_name_on_cloud: str, instance_type: str,
+                    config: common.ProvisionConfig) -> Dict[str, Any]:
+    """Creates a instance and mounts the requested block storage
+
+    Args:
+        region (str): instance region
+        instance_name (str): name of instance
+        config (common.ProvisionConfig): provisioner configuration
+
+    Returns:
+        Dict[str, Any]: instance metadata
+    """
+    # sort tags by key to support deterministic unit test stubbing
+    tags = dict(sorted(copy.deepcopy(config.tags).items()))
+    tags = {
+        'Name': cluster_name_on_cloud,
+        provision_constants.TAG_RAY_CLUSTER_NAME: cluster_name_on_cloud,
+        provision_constants.TAG_SKYPILOT_CLUSTER_NAME: cluster_name_on_cloud,
+        **tags
+    }
+    tags = [f'{key}:{value}' for key, value in tags.items()]
+    default_image = constants.GPU_IMAGES.get(
+        config.node_config['InstanceType'],
+        'gpu-h100x1-base',
+    )
+    image_id = config.node_config['ImageId']
+    image_id = image_id if image_id is not None else default_image
+    instance_name = (f'{cluster_name_on_cloud}-'
+                     f'{uuid.uuid4().hex[:4]}-{instance_type}')
+    instance_request = {
+        'name': instance_name,
+        'region': region,
+        'size': config.node_config['InstanceType'],
+        'image': image_id,
+        'ssh_keys': [
+            ssh_key_id(
+                config.authentication_config['ssh_public_key'])['fingerprint']
+        ],
+        'tags': tags,
+    }
+    instance = _create_droplet(instance_request)
+
+    volume_request = {
+        'size_gigabytes': config.node_config['DiskSize'],
+        'name': instance_name,
+        'region': region,
+        'filesystem_type': 'ext4',
+        'tags': tags
+    }
+    volume = _create_volume(volume_request)
+
+    attach_request = {'type': 'attach', 'droplet_id': instance['id']}
+    try:
+        client().volume_actions.post_by_id(volume['id'], attach_request)
+    except do.exceptions().HttpResponseError as err:
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+    logger.debug(f'{instance_name} created')
+    return instance
+
+
+def start_instance(instance: Dict[str, Any]):
+    try:
+        client().droplet_actions.post(droplet_id=instance['id'],
+                                      body={'type': 'power_on'})
+    except do.exceptions().HttpResponseError as err:
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+
+
+def stop_instance(instance: Dict[str, Any]):
+    try:
+        client().droplet_actions.post(
+            droplet_id=instance['id'],
+            body={'type': 'shutdown'},
+        )
+    except do.exceptions().HttpResponseError as err:
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+
+
+def down_instance(instance: Dict[str, Any]):
+    # We use dangerous destroy to atomically delete
+    # block storage and instance for autodown
+    try:
+        client().droplets.destroy_with_associated_resources_dangerous(
+            droplet_id=instance['id'], x_dangerous=True)
+    except do.exceptions().HttpResponseError as err:
+        if 'a destroy is already in progress' in err.error.message:
+            return
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+
+
+def rename_instance(instance: Dict[str, Any], new_name: str):
+    try:
+        client().droplet_actions.rename(droplet=instance['id'],
+                                        body={
+                                            'type': 'rename',
+                                            'name': new_name
+                                        })
+    except do.exceptions().HttpResponseError as err:
+        raise DigitalOceanError(
+            f'Error: {err.status_code} {err.reason}: {err.error.message}'
+        ) from err
+
+
+def filter_instances(
+        cluster_name_on_cloud: str,
+        status_filters: Optional[List[str]] = None) -> Dict[str, Any]:
+    """Returns Dict mapping instance name
+    to instance metadata filtered by status
+    """
+
+    filtered_instances: Dict[str, Any] = {}
+    page = 1
+    paginated = True
+    while paginated:
+        try:
+            resp = client().droplets.list(
+                tag_name=f'{provision_constants.TAG_SKYPILOT_CLUSTER_NAME}:'
+                f'{cluster_name_on_cloud}',
+                per_page=50,
+                page=page)
+            for instance in resp['droplets']:
+                if status_filters is None or instance[
+                        'status'] in status_filters:
+                    filtered_instances[instance['name']] = instance
+        except do.exceptions().HttpResponseError as err:
+            raise DigitalOceanError(
+                f'Error: {err.status_code} {err.reason}: {err.error.message}'
+            ) from err
+
+        pages = resp['links']
+        if 'pages' in pages and 'next' in pages['pages']:
+            pages = pages['pages']
+            parsed_url = urllib.parse.urlparse(pages['next'])
+            page = int(urllib.parse.parse_qs(parsed_url.query)['page'][0])
+        else:
+            paginated = False
+    return filtered_instances