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

sky/provision/runpod/utils.py

@@ -1,10 +1,13 @@
 """RunPod library wrapper for SkyPilot."""
 
+import base64
 import time
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Tuple
 
 from sky import sky_logging
 from sky.adaptors import runpod
+from sky.provision import docker_utils
+import sky.provision.runpod.api.commands as runpod_commands
 from sky.skylet import constants
 from sky.utils import common_utils
 
@@ -45,6 +48,11 @@ GPU_NAME_MAP = {
 }
 
 
+def _construct_docker_login_template_name(cluster_name: str) -> str:
+    """Constructs the registry auth template name."""
+    return f'{cluster_name}-docker-login-template'
+
+
 def retry(func):
     """Decorator to retry a function."""
 
@@ -64,9 +72,83 @@ def retry(func):
     return wrapper
 
 
+# Adapted from runpod.api.queries.pods.py::QUERY_POD.
+# Adding containerRegistryAuthId to the query.
+_QUERY_POD = """
+query myPods {
+    myself {
+        pods {
+            id
+            containerDiskInGb
+            containerRegistryAuthId
+            costPerHr
+            desiredStatus
+            dockerArgs
+            dockerId
+            env
+            gpuCount
+            imageName
+            lastStatusChange
+            machineId
+            memoryInGb
+            name
+            podType
+            port
+            ports
+            uptimeSeconds
+            vcpuCount
+            volumeInGb
+            volumeMountPath
+            runtime {
+                ports{
+                    ip
+                    isIpPublic
+                    privatePort
+                    publicPort
+                    type
+                }
+            }
+            machine {
+                gpuDisplayName
+            }
+        }
+    }
+}
+"""
+
+
+def _sky_get_pods() -> dict:
+    """List all pods with extra registry auth information.
+
+    Adapted from runpod.get_pods() to include containerRegistryAuthId.
+    """
+    raw_return = runpod.runpod.api.graphql.run_graphql_query(_QUERY_POD)
+    cleaned_return = raw_return['data']['myself']['pods']
+    return cleaned_return
+
+
+_QUERY_POD_TEMPLATE_WITH_REGISTRY_AUTH = """
+query myself {
+    myself {
+        podTemplates {
+            name
+            containerRegistryAuthId
+        }
+    }
+}
+"""
+
+
+def _list_pod_templates_with_container_registry() -> dict:
+    """List all pod templates."""
+    raw_return = runpod.runpod.api.graphql.run_graphql_query(
+        _QUERY_POD_TEMPLATE_WITH_REGISTRY_AUTH)
+    return raw_return['data']['myself']['podTemplates']
+
+
 def list_instances() -> Dict[str, Dict[str, Any]]:
     """Lists instances associated with API key."""
-    instances = runpod.runpod.get_pods()
+    instances = _sky_get_pods()
 
     instance_dict: Dict[str, Dict[str, Any]] = {}
     for instance in instances:
@@ -74,13 +156,23 @@ def list_instances() -> Dict[str, Dict[str, Any]]:
 
         info['status'] = instance['desiredStatus']
         info['name'] = instance['name']
+        info['port2endpoint'] = {}
 
-        if instance['desiredStatus'] == 'RUNNING' and instance.get('runtime'):
+        # Sometimes when the cluster is in the process of being created,
+        # the `port` field in the runtime is None and we need to check for it.
+        if (instance['desiredStatus'] == 'RUNNING' and
+                instance.get('runtime') and
+                instance.get('runtime').get('ports')):
             for port in instance['runtime']['ports']:
-                if port['privatePort'] == 22 and port['isIpPublic']:
-                    info['external_ip'] = port['ip']
-                    info['ssh_port'] = port['publicPort']
-                elif not port['isIpPublic']:
+                if port['isIpPublic']:
+                    if port['privatePort'] == 22:
+                        info['external_ip'] = port['ip']
+                        info['ssh_port'] = port['publicPort']
+                    info['port2endpoint'][port['privatePort']] = {
+                        'host': port['ip'],
+                        'port': port['publicPort']
+                    }
+                else:
                     info['internal_ip'] = port['ip']
 
         instance_dict[instance['id']] = info
@@ -88,37 +180,161 @@ def list_instances() -> Dict[str, Dict[str, Any]]:
     return instance_dict
 
 
-def launch(name: str, instance_type: str, region: str, disk_size: int) -> str:
+def delete_pod_template(template_name: str) -> None:
+    """Deletes a pod template."""
+    try:
+        runpod.runpod.api.graphql.run_graphql_query(
+            f'mutation {{deleteTemplate(templateName: "{template_name}")}}')
+    except runpod.runpod.error.QueryError as e:
+        logger.warning(f'Failed to delete template {template_name}: {e} '
+                       'Please delete it manually.')
+
+
+def delete_register_auth(registry_auth_id: str) -> None:
+    """Deletes a registry auth."""
+    try:
+        runpod.runpod.delete_container_registry_auth(registry_auth_id)
+    except runpod.runpod.error.QueryError as e:
+        logger.warning(
+            f'Failed to delete registry auth {registry_auth_id}: {e} '
+            'Please delete it manually.')
+
+
+def _create_template_for_docker_login(
+    cluster_name: str,
+    image_name: str,
+    docker_login_config: Optional[Dict[str, str]],
+) -> Tuple[str, Optional[str]]:
+    """Creates a template for the given image with the docker login config.
+
+    Returns:
+        formatted_image_name: The formatted image name.
+        template_id: The template ID. None for no docker login config.
+    """
+    if docker_login_config is None:
+        return image_name, None
+    login_config = docker_utils.DockerLoginConfig(**docker_login_config)
+    container_registry_auth_name = f'{cluster_name}-registry-auth'
+    container_template_name = _construct_docker_login_template_name(
+        cluster_name)
+    # The `name` argument is only for display purpose and the registry server
+    # will be splitted from the docker image name (Tested with AWS ECR).
+    # Here we only need the username and password to create the registry auth.
+    # TODO(tian): Now we create a template and a registry auth for each cluster.
+    # Consider create one for each server and reuse them. Challenges including
+    # calculate the reference count and delete them when no longer needed.
+    create_auth_resp = runpod.runpod.create_container_registry_auth(
+        name=container_registry_auth_name,
+        username=login_config.username,
+        password=login_config.password,
+    )
+    registry_auth_id = create_auth_resp['id']
+    create_template_resp = runpod.runpod.create_template(
+        name=container_template_name,
+        image_name=None,
+        registry_auth_id=registry_auth_id,
+    )
+    return login_config.format_image(image_name), create_template_resp['id']
+
+
+def launch(cluster_name: str, node_type: str, instance_type: str, region: str,
+           disk_size: int, image_name: str, ports: Optional[List[int]],
+           public_key: str, preemptible: Optional[bool], bid_per_gpu: float,
+           docker_login_config: Optional[Dict[str, str]]) -> str:
     """Launches an instance with the given parameters.
 
     Converts the instance_type to the RunPod GPU name, finds the specs for the
     GPU, and launches the instance.
+
+    Returns:
+        instance_id: The instance ID.
     """
+    name = f'{cluster_name}-{node_type}'
     gpu_type = GPU_NAME_MAP[instance_type.split('_')[1]]
     gpu_quantity = int(instance_type.split('_')[0].replace('x', ''))
     cloud_type = instance_type.split('_')[2]
 
     gpu_specs = runpod.runpod.get_gpu(gpu_type)
+    # TODO(zhwu): keep this align with setups in
+    # `provision.kuberunetes.instance.py`
+    setup_cmd = (
+        'prefix_cmd() '
+        '{ if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }; '
+        '$(prefix_cmd) apt update;'
+        'export DEBIAN_FRONTEND=noninteractive;'
+        '$(prefix_cmd) apt install openssh-server rsync curl patch -y;'
+        '$(prefix_cmd) mkdir -p /var/run/sshd; '
+        '$(prefix_cmd) '
+        'sed -i "s/PermitRootLogin prohibit-password/PermitRootLogin yes/" '
+        '/etc/ssh/sshd_config; '
+        '$(prefix_cmd) sed '
+        '"s@session\\s*required\\s*pam_loginuid.so@session optional '
+        'pam_loginuid.so@g" -i /etc/pam.d/sshd; '
+        'cd /etc/ssh/ && $(prefix_cmd) ssh-keygen -A; '
+        '$(prefix_cmd) mkdir -p ~/.ssh; '
+        '$(prefix_cmd) chown -R $(whoami) ~/.ssh;'
+        '$(prefix_cmd) chmod 700 ~/.ssh; '
+        f'$(prefix_cmd) echo "{public_key}" >> ~/.ssh/authorized_keys; '
+        '$(prefix_cmd) chmod 644 ~/.ssh/authorized_keys; '
+        '$(prefix_cmd) service ssh restart; '
+        '[ $(id -u) -eq 0 ] && echo alias sudo="" >> ~/.bashrc;sleep infinity')
+    # Use base64 to deal with the tricky quoting issues caused by runpod API.
+    encoded = base64.b64encode(setup_cmd.encode('utf-8')).decode('utf-8')
 
-    new_instance = runpod.runpod.create_pod(
-        name=name,
-        image_name='runpod/base:0.0.2',
-        gpu_type_id=gpu_type,
-        cloud_type=cloud_type,
-        container_disk_in_gb=disk_size,
-        min_vcpu_count=4 * gpu_quantity,
-        min_memory_in_gb=gpu_specs['memoryInGb'] * gpu_quantity,
-        gpu_count=gpu_quantity,
-        country_code=region,
-        ports=(f'22/tcp,'
-               f'{constants.SKY_REMOTE_RAY_DASHBOARD_PORT}/http,'
-               f'{constants.SKY_REMOTE_RAY_PORT}/http'),
-        support_public_ip=True,
-    )
+    docker_args = (f'bash -c \'echo {encoded} | base64 --decode > init.sh; '
+                   f'bash init.sh\'')
+
+    # Port 8081 is occupied for nginx in the base image.
+    custom_ports_str = ''
+    if ports is not None:
+        custom_ports_str = ''.join([f'{p}/tcp,' for p in ports])
+    ports_str = (f'22/tcp,'
+                 f'{custom_ports_str}'
+                 f'{constants.SKY_REMOTE_RAY_DASHBOARD_PORT}/http,'
+                 f'{constants.SKY_REMOTE_RAY_PORT}/http')
+
+    image_name_formatted, template_id = _create_template_for_docker_login(
+        cluster_name, image_name, docker_login_config)
+
+    params = {
+        'name': name,
+        'image_name': image_name_formatted,
+        'gpu_type_id': gpu_type,
+        'cloud_type': cloud_type,
+        'container_disk_in_gb': disk_size,
+        'min_vcpu_count': 4 * gpu_quantity,
+        'min_memory_in_gb': gpu_specs['memoryInGb'] * gpu_quantity,
+        'gpu_count': gpu_quantity,
+        'country_code': region,
+        'ports': ports_str,
+        'support_public_ip': True,
+        'docker_args': docker_args,
+        'template_id': template_id,
+    }
+
+    if preemptible is None or not preemptible:
+        new_instance = runpod.runpod.create_pod(**params)
+    else:
+        new_instance = runpod_commands.create_spot_pod(
+            bid_per_gpu=bid_per_gpu,
+            **params,
+        )
 
     return new_instance['id']
 
 
+def get_registry_auth_resources(
+        cluster_name: str) -> Tuple[Optional[str], Optional[str]]:
+    """Gets the registry auth resources."""
+    container_registry_auth_name = _construct_docker_login_template_name(
+        cluster_name)
+    for template in _list_pod_templates_with_container_registry():
+        if template['name'] == container_registry_auth_name:
+            return container_registry_auth_name, template[
+                'containerRegistryAuthId']
+    return None, None
+
+
 def remove(instance_id: str) -> None:
     """Terminates the given instance."""
     runpod.runpod.terminate_pod(instance_id)

sky/provision/vast/__init__.py

@@ -0,0 +1,10 @@
+"""Vast provisioner for SkyPilot."""
+
+from sky.provision.vast.config import bootstrap_instances
+from sky.provision.vast.instance import cleanup_ports
+from sky.provision.vast.instance import get_cluster_info
+from sky.provision.vast.instance import query_instances
+from sky.provision.vast.instance import run_instances
+from sky.provision.vast.instance import stop_instances
+from sky.provision.vast.instance import terminate_instances
+from sky.provision.vast.instance import wait_instances

sky/provision/vast/config.py

@@ -0,0 +1,11 @@
+"""Vast configuration bootstrapping."""
+
+from sky.provision import common
+
+
+def bootstrap_instances(
+        region: str, cluster_name: str,
+        config: common.ProvisionConfig) -> common.ProvisionConfig:
+    """Bootstraps instances for the given cluster."""
+    del region, cluster_name  # unused
+    return config

sky/provision/vast/instance.py

@@ -0,0 +1,247 @@
+"""Vast instance provisioning."""
+import time
+from typing import Any, Dict, List, Optional
+
+from sky import sky_logging
+from sky.provision import common
+from sky.provision.vast import utils
+from sky.utils import common_utils
+from sky.utils import status_lib
+from sky.utils import ux_utils
+
+POLL_INTERVAL = 10
+
+logger = sky_logging.init_logger(__name__)
+# a much more convenient method
+status_filter = lambda machine_dict, stat_list: {
+    k: v for k, v in machine_dict.items() if v['status'] in stat_list
+}
+
+
+def _filter_instances(cluster_name_on_cloud: str,
+                      status_filters: Optional[List[str]],
+                      head_only: bool = False) -> Dict[str, Any]:
+
+    instances = utils.list_instances()
+    possible_names = [f'{cluster_name_on_cloud}-head']
+    if not head_only:
+        possible_names.append(f'{cluster_name_on_cloud}-worker')
+
+    filtered_instances = {}
+    for instance_id, instance in instances.items():
+        if (status_filters is not None and
+                instance['status'] not in status_filters):
+            continue
+        if instance.get('name') in possible_names:
+            filtered_instances[instance_id] = instance
+    return filtered_instances
+
+
+def _get_head_instance_id(instances: Dict[str, Any]) -> Optional[str]:
+    for inst_id, inst in instances.items():
+        if inst['name'].endswith('-head'):
+            return inst_id
+    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 = ['CREATED', 'RESTARTING']
+
+    created_instance_ids = []
+    instances: Dict[str, Any] = {}
+
+    while True:
+        instances = _filter_instances(cluster_name_on_cloud, None)
+        if not status_filter(instances, pending_status):
+            break
+        logger.info(f'Waiting for {len(instances)} instances to be ready.')
+        time.sleep(POLL_INTERVAL)
+
+    running_instances = status_filter(instances, ['RUNNING'])
+    head_instance_id = _get_head_instance_id(running_instances)
+    stopped_instances = status_filter(instances, ['EXITED', 'STOPPED'])
+
+    if config.resume_stopped_nodes and stopped_instances:
+        for instance in stopped_instances.values():
+            utils.start(instance['id'])
+    else:
+        to_start_count = config.count - (len(running_instances) +
+                                         len(stopped_instances))
+        if to_start_count < 0:
+            raise RuntimeError(f'Cluster {cluster_name_on_cloud} already has '
+                               f'{len(running_instances)} nodes,'
+                               f'but {config.count} are required.')
+        if to_start_count == 0:
+            if head_instance_id is None:
+                raise RuntimeError(
+                    f'Cluster {cluster_name_on_cloud} has no head node.')
+            logger.info(
+                f'Cluster {cluster_name_on_cloud} already has '
+                f'{len(running_instances)} nodes, no need to start more.')
+            return common.ProvisionRecord(provider_name='vast',
+                                          cluster_name=cluster_name_on_cloud,
+                                          region=region,
+                                          zone=None,
+                                          head_instance_id=head_instance_id,
+                                          resumed_instance_ids=[],
+                                          created_instance_ids=[])
+
+        for _ in range(to_start_count):
+            node_type = 'head' if head_instance_id is None else 'worker'
+            try:
+                instance_id = utils.launch(
+                    name=f'{cluster_name_on_cloud}-{node_type}',
+                    instance_type=config.node_config['InstanceType'],
+                    region=region,
+                    disk_size=config.node_config['DiskSize'],
+                    preemptible=config.node_config['Preemptible'],
+                    image_name=config.node_config['ImageId'])
+            except Exception as e:  # pylint: disable=broad-except
+                logger.warning(f'run_instances error: {e}')
+                raise
+            logger.info(f'Launched instance {instance_id}.')
+            created_instance_ids.append(instance_id)
+            if head_instance_id is None:
+                head_instance_id = instance_id
+
+    # Wait for instances to be ready.
+    while True:
+        instances = _filter_instances(cluster_name_on_cloud, ['RUNNING'])
+        ready_instance_cnt = 0
+        for instance_id, instance in instances.items():
+            if instance.get('ssh_port') is not None:
+                ready_instance_cnt += 1
+        logger.info('Waiting for instances to be ready: '
+                    f'({ready_instance_cnt}/{config.count}).')
+        if ready_instance_cnt == config.count:
+            break
+
+        time.sleep(POLL_INTERVAL)
+
+    head_instance_id = _get_head_instance_id(utils.list_instances())
+    assert head_instance_id is not None, 'head_instance_id should not be None'
+    return common.ProvisionRecord(provider_name='vast',
+                                  cluster_name=cluster_name_on_cloud,
+                                  region=region,
+                                  zone=None,
+                                  head_instance_id=head_instance_id,
+                                  resumed_instance_ids=[],
+                                  created_instance_ids=created_instance_ids)
+
+
+def wait_instances(region: str, cluster_name_on_cloud: str,
+                   state: Optional[status_lib.ClusterStatus]) -> None:
+    del region, cluster_name_on_cloud, state
+
+
+def stop_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    return action_instances('stop', cluster_name_on_cloud, provider_config,
+                            worker_only)
+
+
+def terminate_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    return action_instances('remove', cluster_name_on_cloud, provider_config,
+                            worker_only)
+
+
+def action_instances(
+    fn: str,
+    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 = _filter_instances(cluster_name_on_cloud, None)
+    for inst_id, inst in instances.items():
+        logger.debug(f'Instance {fn} {inst_id}: {inst}')
+        if worker_only and inst['name'].endswith('-head'):
+            continue
+        try:
+            getattr(utils, fn)(inst_id)
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError(
+                    f'Failed to {fn} instance {inst_id}: '
+                    f'{common_utils.format_exception(e, use_bracket=False)}'
+                ) from e
+
+
+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 = _filter_instances(cluster_name_on_cloud, ['RUNNING'])
+    instances: Dict[str, List[common.InstanceInfo]] = {}
+    head_instance_id = None
+    for instance_id, instance_info in running_instances.items():
+        instances[instance_id] = [
+            common.InstanceInfo(
+                instance_id=instance_id,
+                internal_ip=instance_info['local_ipaddrs'].strip(),
+                external_ip=instance_info['public_ipaddr'],
+                ssh_port=instance_info['ports']['22/tcp'][0]['HostPort'],
+                tags={},
+            )
+        ]
+        if instance_info['name'].endswith('-head'):
+            head_instance_id = instance_id
+
+    return common.ClusterInfo(
+        instances=instances,
+        head_instance_id=head_instance_id,
+        provider_name='vast',
+        provider_config=provider_config,
+    )
+
+
+def open_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    raise NotImplementedError('open_ports is not supported for Vast')
+
+
+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"""
+
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+    instances = _filter_instances(cluster_name_on_cloud, None)
+    # "running", "frozen", "stopped", "unknown", "loading"
+    status_map = {
+        'LOADING': status_lib.ClusterStatus.INIT,
+        'EXITED': status_lib.ClusterStatus.STOPPED,
+        'STOPPED': status_lib.ClusterStatus.STOPPED,
+        'RUNNING': status_lib.ClusterStatus.UP,
+    }
+    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    for inst_id, inst in instances.items():
+        status = status_map[inst['status']]
+        if non_terminated_only and status is None:
+            continue
+        statuses[inst_id] = status
+    return statuses
+
+
+def cleanup_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  # Unused.