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

sky/provision/lambda_cloud/__init__.py

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

sky/provision/lambda_cloud/config.py

@@ -0,0 +1,10 @@
+"""Lambda Cloud configuration bootstrapping"""
+
+from sky.provision import common
+
+
+def bootstrap_instances(
+        region: str, cluster_name: str,
+        config: common.ProvisionConfig) -> common.ProvisionConfig:
+    del region, cluster_name  # unused
+    return config

sky/provision/lambda_cloud/instance.py

@@ -0,0 +1,265 @@
+"""Lambda instance provisioning."""
+
+import time
+from typing import Any, Dict, List, Optional
+
+from sky import sky_logging
+from sky.provision import common
+import sky.provision.lambda_cloud.lambda_utils as lambda_utils
+from sky.utils import common_utils
+from sky.utils import status_lib
+from sky.utils import ux_utils
+
+POLL_INTERVAL = 1
+
+logger = sky_logging.init_logger(__name__)
+_lambda_client = None
+
+
+def _get_lambda_client():
+    global _lambda_client
+    if _lambda_client is None:
+        _lambda_client = lambda_utils.LambdaCloudClient()
+    return _lambda_client
+
+
+def _filter_instances(
+        cluster_name_on_cloud: str,
+        status_filters: Optional[List[str]]) -> Dict[str, Dict[str, Any]]:
+    lambda_client = _get_lambda_client()
+    instances = lambda_client.list_instances()
+    possible_names = [
+        f'{cluster_name_on_cloud}-head',
+        f'{cluster_name_on_cloud}-worker',
+    ]
+
+    filtered_instances = {}
+    for instance in instances:
+        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]:
+    head_instance_id = None
+    for instance_id, instance in instances.items():
+        if instance['name'].endswith('-head'):
+            head_instance_id = instance_id
+            break
+    return head_instance_id
+
+
+def _get_private_ip(instance_info: Dict[str, Any], single_node: bool) -> str:
+    private_ip = instance_info.get('private_ip')
+    if private_ip is None:
+        if single_node:
+            # The Lambda cloud API may return an instance info without
+            # private IP. It does not align with their docs, but we still
+            # allow single-node cluster to proceed with provisioning, by using
+            # 127.0.0.1, as private IP is not critical for single-node case.
+            return '127.0.0.1'
+        msg = f'Failed to retrieve private IP for instance {instance_info}.'
+        logger.error(msg)
+        raise RuntimeError(msg)
+    return private_ip
+
+
+def run_instances(region: str, cluster_name_on_cloud: str,
+                  config: common.ProvisionConfig) -> common.ProvisionRecord:
+    """Runs instances for the given cluster"""
+    lambda_client = _get_lambda_client()
+    pending_status = ['booting']
+    while True:
+        instances = _filter_instances(cluster_name_on_cloud, pending_status)
+        if not instances:
+            break
+        logger.info(f'Waiting for {len(instances)} instances to be ready.')
+        time.sleep(POLL_INTERVAL)
+    exist_instances = _filter_instances(cluster_name_on_cloud, ['active'])
+    head_instance_id = _get_head_instance_id(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_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(exist_instances)} nodes, no need to start more.')
+        return common.ProvisionRecord(
+            provider_name='lambda',
+            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 = []
+    remote_ssh_key_name = config.authentication_config['remote_key_name']
+
+    def launch_nodes(node_type: str, quantity: int) -> List[str]:
+        try:
+            instance_ids = lambda_client.create_instances(
+                instance_type=config.node_config['InstanceType'],
+                region=region,
+                name=f'{cluster_name_on_cloud}-{node_type}',
+                quantity=quantity,
+                ssh_key_name=remote_ssh_key_name,
+            )
+            logger.info(f'Launched {len(instance_ids)} {node_type} node(s), '
+                        f'instance_ids: {instance_ids}')
+            return instance_ids
+        except Exception as e:
+            logger.warning(f'run_instances error: {e}')
+            raise
+
+    if head_instance_id is None:
+        instance_ids = launch_nodes('head', 1)
+        assert len(instance_ids) == 1
+        created_instance_ids.append(instance_ids[0])
+        head_instance_id = instance_ids[0]
+
+    assert head_instance_id is not None, 'head_instance_id should not be None'
+
+    worker_node_count = to_start_count - 1
+    if worker_node_count > 0:
+        instance_ids = launch_nodes('worker', worker_node_count)
+        created_instance_ids.extend(instance_ids)
+
+    while True:
+        instances = _filter_instances(cluster_name_on_cloud, ['active'])
+        if len(instances) == config.count:
+            break
+
+        time.sleep(POLL_INTERVAL)
+
+    return common.ProvisionRecord(
+        provider_name='lambda',
+        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  # Unused.
+
+
+def stop_instances(
+    cluster_name_on_cloud: str,
+    provider_config: Optional[Dict[str, Any]] = None,
+    worker_only: bool = False,
+) -> None:
+    raise NotImplementedError(
+        'stop_instances is not supported for Lambda Cloud')
+
+
+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
+    lambda_client = _get_lambda_client()
+    instances = _filter_instances(cluster_name_on_cloud, None)
+
+    instance_ids_to_terminate = []
+    for instance_id, instance in instances.items():
+        if worker_only and not instance['name'].endswith('-worker'):
+            continue
+        instance_ids_to_terminate.append(instance_id)
+
+    try:
+        logger.debug(
+            f'Terminating instances {", ".join(instance_ids_to_terminate)}')
+        lambda_client.remove_instances(instance_ids_to_terminate)
+    except Exception as e:  # pylint: disable=broad-except
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError(
+                f'Failed to terminate instances {instance_ids_to_terminate}: '
+                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, ['active'])
+    single_node = len(running_instances) == 1
+    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=_get_private_ip(instance_info, single_node),
+                external_ip=instance_info['ip'],
+                ssh_port=22,
+                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='lambda',
+        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"""
+    assert provider_config is not None, (cluster_name_on_cloud, provider_config)
+    instances = _filter_instances(cluster_name_on_cloud, None)
+
+    status_map = {
+        'booting': status_lib.ClusterStatus.INIT,
+        'active': status_lib.ClusterStatus.UP,
+        'unhealthy': status_lib.ClusterStatus.INIT,
+        'terminating': None,
+    }
+    statuses: Dict[str, Optional[status_lib.ClusterStatus]] = {}
+    for instance_id, instance in instances.items():
+        status = status_map.get(instance['status'])
+        if non_terminated_only and status is None:
+            continue
+        statuses[instance_id] = status
+    return statuses
+
+
+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 Lambda Cloud')
+
+
+def cleanup_ports(
+    cluster_name_on_cloud: str,
+    ports: List[str],
+    provider_config: Optional[Dict[str, Any]] = None,
+) -> None:
+    """See sky/provision/__init__.py"""
+    del cluster_name_on_cloud, ports, provider_config  # Unused.

sky/{clouds/utils → provision/lambda_cloud}/lambda_utils.py

@@ -1,4 +1,5 @@
 """Lambda Cloud helper functions."""
+
 import json
 import os
 import time
@@ -49,7 +50,7 @@ class Metadata:
         if value is None:
             if instance_id in metadata:
                 metadata.pop(instance_id)  # del entry
-            if len(metadata) == 0:
+            if not metadata:
                 if os.path.exists(self.path):
                     os.remove(self.path)
                 return
@@ -68,7 +69,7 @@ class Metadata:
         for instance_id in list(metadata.keys()):
             if instance_id not in instance_ids:
                 del metadata[instance_id]
-        if len(metadata) == 0:
+        if not metadata:
             os.remove(self.path)
             return
         with open(self.path, 'w', encoding='utf-8') as f:
@@ -76,12 +77,12 @@ class Metadata:
 
 
 def raise_lambda_error(response: requests.Response) -> None:
-    """Raise LambdaCloudError if appropriate. """
+    """Raise LambdaCloudError if appropriate."""
     status_code = response.status_code
     if status_code == 200:
         return
     if status_code == 429:
-        # https://docs.lambdalabs.com/cloud/rate-limiting/
+        # https://docs.lambdalabs.com/public-cloud/cloud-api/
         raise LambdaCloudError('Your API requests are being rate limited.')
     try:
         resp_json = response.json()
@@ -131,23 +132,25 @@ class LambdaCloudClient:
         self.api_key = self._credentials['api_key']
         self.headers = {'Authorization': f'Bearer {self.api_key}'}
 
-    def create_instances(self,
-                         instance_type: str = 'gpu_1x_a100_sxm4',
-                         region: str = 'us-east-1',
-                         quantity: int = 1,
-                         name: str = '',
-                         ssh_key_name: str = '') -> List[str]:
+    def create_instances(
+        self,
+        instance_type: str = 'gpu_1x_a100_sxm4',
+        region: str = 'us-east-1',
+        quantity: int = 1,
+        name: str = '',
+        ssh_key_name: str = '',
+    ) -> List[str]:
         """Launch new instances."""
         # Optimization:
         # Most API requests are rate limited at ~1 request every second but
         # launch requests are rate limited at ~1 request every 10 seconds.
         # So don't use launch requests to check availability.
-        # See https://docs.lambdalabs.com/cloud/rate-limiting/ for more.
-        available_regions = self.list_catalog()[instance_type]\
-                ['regions_with_capacity_available']
+        # See https://docs.lambdalabs.com/public-cloud/cloud-api/ for more.
+        available_regions = (self.list_catalog()[instance_type]
+                             ['regions_with_capacity_available'])
         available_regions = [reg['name'] for reg in available_regions]
         if region not in available_regions:
-            if len(available_regions) > 0:
+            if available_regions:
                 aval_reg = ' '.join(available_regions)
             else:
                 aval_reg = 'None'
@@ -163,27 +166,25 @@ class LambdaCloudClient:
             'instance_type_name': instance_type,
             'ssh_key_names': [ssh_key_name],
             'quantity': quantity,
-            'name': name
+            'name': name,
         })
         response = _try_request_with_backoff(
             'post',
             f'{API_ENDPOINT}/instance-operations/launch',
             data=data,
-            headers=self.headers)
+            headers=self.headers,
+        )
         return response.json().get('data', []).get('instance_ids', [])
 
-    def remove_instances(self, *instance_ids: str) -> Dict[str, Any]:
+    def remove_instances(self, instance_ids: List[str]) -> Dict[str, Any]:
         """Terminate instances."""
-        data = json.dumps({
-            'instance_ids': [
-                instance_ids[0]  # TODO(ewzeng) don't hardcode
-            ]
-        })
+        data = json.dumps({'instance_ids': instance_ids})
         response = _try_request_with_backoff(
             'post',
             f'{API_ENDPOINT}/instance-operations/terminate',
             data=data,
-            headers=self.headers)
+            headers=self.headers,
+        )
         return response.json().get('data', []).get('terminated_instances', [])
 
     def list_instances(self) -> List[Dict[str, Any]]:

sky/provision/logging.py

@@ -41,7 +41,7 @@ def setup_provision_logging(log_dir: str):
         # Disable propagation to avoid streaming logs to the console, which is
         # set up for sky root logger.
         provision_logger.propagate = False
-        stream_handler = sky_logging.RichSafeStreamHandler(sys.stdout)
+        stream_handler = logging.StreamHandler(sys.stdout)
         stream_handler.flush = sys.stdout.flush  # type: ignore
         stream_handler.setFormatter(sky_logging.DIM_FORMATTER)
         stream_handler.setLevel(logging.WARNING)

sky/provision/nebius/__init__.py

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

sky/provision/nebius/config.py

@@ -0,0 +1,11 @@
+"""Nebius 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