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

sky/models.py

@@ -0,0 +1,27 @@
+"""Data Models for SkyPilot."""
+
+import collections
+import dataclasses
+from typing import Dict, Optional
+
+
+@dataclasses.dataclass
+class User:
+    # User hash
+    id: str
+    # Display name of the user
+    name: Optional[str] = None
+
+
+RealtimeGpuAvailability = collections.namedtuple(
+    'RealtimeGpuAvailability', ['gpu', 'counts', 'capacity', 'available'])
+
+
+@dataclasses.dataclass
+class KubernetesNodeInfo:
+    """Dataclass to store Kubernetes node information."""
+    name: str
+    accelerator_type: Optional[str]
+    # Resources available on the node. E.g., {'nvidia.com/gpu': '2'}
+    total: Dict[str, int]
+    free: Dict[str, int]

sky/optimizer.py

@@ -1,7 +1,6 @@
 """Optimizer: assigns best resources to user tasks."""
 import collections
 import copy
-import enum
 import json
 import typing
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
@@ -17,8 +16,14 @@ from sky import resources as resources_lib
 from sky import sky_logging
 from sky import task as task_lib
 from sky.adaptors import common as adaptors_common
+from sky.usage import usage_lib
+from sky.utils import common
 from sky.utils import env_options
 from sky.utils import log_utils
+from sky.utils import resources_utils
+from sky.utils import rich_utils
+from sky.utils import subprocess_utils
+from sky.utils import timeline
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -41,12 +46,6 @@ _PerCloudCandidates = Dict[clouds.Cloud, List[resources_lib.Resources]]
 _TaskToPerCloudCandidates = Dict[task_lib.Task, _PerCloudCandidates]
 
 
-# Constants: minimize what target?
-class OptimizeTarget(enum.Enum):
-    COST = 0
-    TIME = 1
-
-
 # For logging purposes.
 def _create_table(field_names: List[str]) -> prettytable.PrettyTable:
     table_kwargs = {
@@ -102,11 +101,13 @@ class Optimizer:
         return egress_time
 
     @staticmethod
+    @timeline.event
+    @usage_lib.entrypoint('sky.optimizer.optimize')
     def optimize(dag: 'dag_lib.Dag',
-                 minimize: OptimizeTarget = OptimizeTarget.COST,
+                 minimize: common.OptimizeTarget = common.OptimizeTarget.COST,
                  blocked_resources: Optional[Iterable[
                      resources_lib.Resources]] = None,
-                 quiet: bool = False):
+                 quiet: bool = False) -> 'dag_lib.Dag':
         """Find the best execution plan for the given DAG.
 
         Args:
@@ -120,22 +121,22 @@ class Optimizer:
                 for a task.
             exceptions.NoCloudAccessError: if no public clouds are enabled.
         """
-        _check_specified_clouds(dag)
-
-        # This function is effectful: mutates every node in 'dag' by setting
-        # node.best_resources if it is None.
-        Optimizer._add_dummy_source_sink_nodes(dag)
-        try:
-            unused_best_plan = Optimizer._optimize_dag(
-                dag=dag,
-                minimize_cost=minimize == OptimizeTarget.COST,
-                blocked_resources=blocked_resources,
-                quiet=quiet)
-        finally:
-            # Make sure to remove the dummy source/sink nodes, even if the
-            # optimization fails.
-            Optimizer._remove_dummy_source_sink_nodes(dag)
-        return dag
+        with rich_utils.safe_status(ux_utils.spinner_message('Optimizing')):
+            _check_specified_clouds(dag)
+            # This function is effectful: mutates every node in 'dag' by setting
+            # node.best_resources if it is None.
+            Optimizer._add_dummy_source_sink_nodes(dag)
+            try:
+                unused_best_plan = Optimizer._optimize_dag(
+                    dag=dag,
+                    minimize_cost=minimize == common.OptimizeTarget.COST,
+                    blocked_resources=blocked_resources,
+                    quiet=quiet)
+            finally:
+                # Make sure to remove the dummy source/sink nodes, even if the
+                # optimization fails.
+                Optimizer._remove_dummy_source_sink_nodes(dag)
+            return dag
 
     @staticmethod
     def _add_dummy_source_sink_nodes(dag: 'dag_lib.Dag'):
@@ -182,7 +183,7 @@ class Optimizer:
         """Removes special Source and Sink nodes."""
         source = [t for t in dag.tasks if t.name == _DUMMY_SOURCE_NAME]
         sink = [t for t in dag.tasks if t.name == _DUMMY_SINK_NAME]
-        if len(source) == len(sink) == 0:
+        if not source and not sink:
             return
         assert len(source) == len(sink) == 1, dag.tasks
         dag.remove(source[0])
@@ -252,6 +253,29 @@ class Optimizer:
         # node -> cloud -> list of resources that satisfy user's requirements.
         node_to_candidate_map: _TaskToPerCloudCandidates = {}
 
+        def get_available_reservations(
+            launchable_resources: Dict[resources_lib.Resources,
+                                       List[resources_lib.Resources]]
+        ) -> Dict[resources_lib.Resources, int]:
+            if not resources_utils.need_to_query_reservations():
+                return {}
+
+            num_available_reserved_nodes_per_resource = {}
+
+            def get_reservations_available_resources(
+                    resources: resources_lib.Resources):
+                num_available_reserved_nodes_per_resource[resources] = sum(
+                    resources.get_reservations_available_resources().values())
+
+            launchable_resources_list: List[resources_lib.Resources] = sum(
+                launchable_resources.values(), [])
+            with rich_utils.safe_status(
+                    ux_utils.spinner_message('Checking reserved resources')):
+                subprocess_utils.run_in_parallel(
+                    get_reservations_available_resources,
+                    launchable_resources_list)
+            return num_available_reserved_nodes_per_resource
+
         # Compute the estimated cost/time for each node.
         for node_i, node in enumerate(topo_order):
             if node_i == 0:
@@ -261,8 +285,6 @@ class Optimizer:
 
             # Don't print for the last node, Sink.
             do_print = node_i != len(topo_order) - 1
-            if do_print:
-                logger.debug('#### {} ####'.format(node))
 
             fuzzy_candidates: List[str] = []
             if node_i < len(topo_order) - 1:
@@ -273,13 +295,21 @@ class Optimizer:
                         blocked_resources=blocked_resources,
                         quiet=quiet))
                 node_to_candidate_map[node] = cloud_candidates
+                # Has to call the printing after the launchable resources are
+                # computed, because the missing fields of the resources are
+                # inferred in the _fill_in_launchable_resources function.
+                logger.debug('#### {} ####'.format(node))
             else:
                 # Dummy sink node.
                 launchable_resources = {
                     list(node.resources)[0]: list(node.resources)
                 }
 
+            # Fetch reservations in advance and in parallel to speed up the
+            # reservation info fetching.
             num_resources = len(list(node.resources))
+            num_available_reserved_nodes_per_resource = (
+                get_available_reservations(launchable_resources))
 
             for orig_resources, launchable_list in launchable_resources.items():
                 if num_resources == 1 and node.time_estimator_func is None:
@@ -302,15 +332,16 @@ class Optimizer:
                     else:
                         estimated_runtime = node.estimate_runtime(
                             orig_resources)
+
                 for resources in launchable_list:
                     if do_print:
                         logger.debug(f'resources: {resources}')
 
                     if minimize_cost:
                         cost_per_node = resources.get_cost(estimated_runtime)
-                        num_available_reserved_nodes = sum(
-                            resources.get_reservations_available_resources(
-                            ).values())
+                        num_available_reserved_nodes = (
+                            num_available_reserved_nodes_per_resource.get(
+                                resources, 0))
 
                         # We consider the cost of the unused reservation
                         # resources to be 0 since we are already paying for
@@ -348,10 +379,6 @@ class Optimizer:
                             for orig_resources in node.resources):
                         source_hint = 'kubernetes cluster'
 
-                # TODO(romilb): When `sky show-gpus` supports Kubernetes,
-                #  add a hint to run `sky show-gpus --kubernetes` to list
-                #  available accelerators on Kubernetes.
-
                 bold = colorama.Style.BRIGHT
                 cyan = colorama.Fore.CYAN
                 reset = colorama.Style.RESET_ALL
@@ -360,10 +387,14 @@ class Optimizer:
                     fuzzy_candidates_str = (
                         f'\nTry one of these offered accelerators: {cyan}'
                         f'{fuzzy_candidates}{reset}')
+                node_resources_reprs = ', '.join(f'{node.num_nodes}x ' +
+                                                 r.repr_with_region_zone
+                                                 for r in node.resources)
                 error_msg = (
                     f'{source_hint.capitalize()} does not contain any '
-                    f'instances satisfying the request:\n{node}.'
-                    f'\n\nTo fix: relax or change the '
+                    f'instances satisfying the request: '
+                    f'{node_resources_reprs}.'
+                    f'\nTo fix: relax or change the '
                     f'resource requirements.{fuzzy_candidates_str}\n\n'
                     f'Hint: {bold}sky show-gpus{reset} '
                     'to list available accelerators.\n'
@@ -692,7 +723,6 @@ class Optimizer:
         node_to_cost_map: _TaskToCostMap,
         minimize_cost: bool,
     ):
-        logger.info('== Optimizer ==')
         ordered_node_to_cost_map = collections.OrderedDict()
         ordered_best_plan = collections.OrderedDict()
         for node in topo_order:
@@ -714,15 +744,18 @@ class Optimizer:
                     node.get_inputs() is None and node.get_outputs() is None):
                 print_hourly_cost = True
 
-        if print_hourly_cost:
-            logger.info(f'{colorama.Style.BRIGHT}Estimated cost: '
-                        f'{colorama.Style.RESET_ALL}${total_cost:.1f} / hour\n')
-        else:
-            logger.info(f'{colorama.Style.BRIGHT}Estimated total runtime: '
-                        f'{colorama.Style.RESET_ALL}{total_time / 3600:.1f} '
-                        'hours\n'
-                        f'{colorama.Style.BRIGHT}Estimated total cost: '
-                        f'{colorama.Style.RESET_ALL}${total_cost:.1f}\n')
+        if not env_options.Options.MINIMIZE_LOGGING.get():
+            if print_hourly_cost:
+                logger.info(
+                    f'{colorama.Style.BRIGHT}Estimated cost: '
+                    f'{colorama.Style.RESET_ALL}${total_cost:.1f} / hour\n')
+            else:
+                logger.info(
+                    f'{colorama.Style.BRIGHT}Estimated total runtime: '
+                    f'{colorama.Style.RESET_ALL}{total_time / 3600:.1f} '
+                    'hours\n'
+                    f'{colorama.Style.BRIGHT}Estimated total cost: '
+                    f'{colorama.Style.RESET_ALL}${total_cost:.1f}\n')
 
         def _get_resources_element_list(
                 resources: 'resources_lib.Resources') -> List[str]:
@@ -797,13 +830,17 @@ class Optimizer:
             return row
 
         def _get_resource_group_hash(resources: 'resources_lib.Resources'):
-            return json.dumps(
-                {
-                    'cloud': f'{resources.cloud}',
-                    'accelerators': f'{resources.accelerators}',
-                    'use_spot': resources.use_spot
-                },
-                sort_keys=True)
+            resource_key_dict = {
+                'cloud': f'{resources.cloud}',
+                'accelerators': f'{resources.accelerators}',
+                'use_spot': resources.use_spot
+            }
+            if isinstance(resources.cloud, clouds.Kubernetes):
+                # Region for Kubernetes is the context name, i.e. different
+                # Kubernetes clusters. We add region to the key to show all the
+                # Kubernetes clusters in the optimizer table for better UX.
+                resource_key_dict['region'] = resources.region
+            return json.dumps(resource_key_dict, sort_keys=True)
 
         # Print the list of resouces that the optimizer considered.
         resource_fields = [
@@ -821,7 +858,7 @@ class Optimizer:
             best_plan_table = _create_table(['TASK', '#NODES'] +
                                             resource_fields)
             best_plan_table.add_rows(best_plan_rows)
-            logger.info(f'{best_plan_table}\n')
+            logger.info(f'{best_plan_table}')
 
         # Print the egress plan if any data egress is scheduled.
         Optimizer._print_egress_plan(graph, best_plan, minimize_cost)
@@ -840,10 +877,12 @@ class Optimizer:
             }
             task_str = (f'for task {task.name!r} ' if num_tasks > 1 else '')
             plural = 's' if task.num_nodes > 1 else ''
-            logger.info(
-                f'{colorama.Style.BRIGHT}Considered resources {task_str}'
-                f'({task.num_nodes} node{plural}):'
-                f'{colorama.Style.RESET_ALL}')
+            if num_tasks > 1:
+                # Add a new line for better readability, when there are multiple
+                # tasks.
+                logger.info('')
+            logger.info(f'Considered resources {task_str}'
+                        f'({task.num_nodes} node{plural}):')
 
             # Only print 1 row per cloud.
             # The following code is to generate the table
@@ -910,7 +949,16 @@ class Optimizer:
 
             table = _create_table(field_names)
             table.add_rows(rows)
-            logger.info(f'{table}\n')
+            logger.info(f'{table}')
+
+            # Warning message for using disk_tier=ultra
+            # TODO(yi): Consider price of disks in optimizer and
+            # move this warning there.
+            if chosen_resources.disk_tier == resources_utils.DiskTier.ULTRA:
+                logger.warning(
+                    'Using disk_tier=ultra will utilize more advanced disks '
+                    '(io2 Block Express on AWS and extreme persistent disk on '
+                    'GCP), which can lead to significant higher costs (~$2/h).')
 
     @staticmethod
     def _print_candidates(node_to_candidate_map: _TaskToPerCloudCandidates):
@@ -932,10 +980,10 @@ class Optimizer:
                             f'Multiple {cloud} instances satisfy '
                             f'{acc_name}:{int(acc_count)}. '
                             f'The cheapest {candidate_list[0]!r} is considered '
-                            f'among:\n{instance_list}.\n')
+                            f'among:\n{instance_list}.')
             if is_multi_instances:
                 logger.info(
-                    f'To list more details, run \'sky show-gpus {acc_name}\'.')
+                    f'To list more details, run: sky show-gpus {acc_name}\n')
 
     @staticmethod
     def _optimize_dag(
@@ -1068,8 +1116,7 @@ class Optimizer:
             Optimizer.print_optimized_plan(graph, topo_order, best_plan,
                                            total_time, total_cost,
                                            node_to_cost_map, minimize_cost)
-            if not env_options.Options.MINIMIZE_LOGGING.get():
-                Optimizer._print_candidates(local_node_to_candidate_map)
+            Optimizer._print_candidates(local_node_to_candidate_map)
         return best_plan
 
 
@@ -1120,7 +1167,7 @@ def _make_launchables_for_valid_region_zones(
     regions = launchable_resources.get_valid_regions_for_launchable()
     for region in regions:
         if (launchable_resources.use_spot and region.zones is not None or
-                isinstance(launchable_resources.cloud, clouds.GCP)):
+                launchable_resources.cloud.optimize_by_zone()):
             # Spot instances.
             # Do not batch the per-zone requests.
             for zone in region.zones:
@@ -1231,6 +1278,9 @@ def _fill_in_launchable_resources(
     if blocked_resources is None:
         blocked_resources = []
     for resources in task.resources:
+        # Validate the resources first which may fill in missing fields
+        # automatically for the resources.
+        resources.validate()
         if (resources.cloud is not None and
                 not clouds.cloud_in_iterable(resources.cloud, enabled_clouds)):
             # Skip the resources that are on a cloud that is not enabled. The
@@ -1239,22 +1289,29 @@ def _fill_in_launchable_resources(
             continue
         clouds_list = ([resources.cloud]
                        if resources.cloud is not None else enabled_clouds)
-        for cloud in clouds_list:
-            (feasible_resources,
-             fuzzy_candidate_list) = cloud.get_feasible_launchable_resources(
-                 resources, num_nodes=task.num_nodes)
-            if len(feasible_resources) > 0:
+        # If clouds provide hints, store them for later printing.
+        hints: Dict[clouds.Cloud, str] = {}
+
+        feasible_list = subprocess_utils.run_in_parallel(
+            lambda cloud, r=resources, n=task.num_nodes:
+            (cloud, cloud.get_feasible_launchable_resources(r, n)),
+            clouds_list)
+        for cloud, feasible_resources in feasible_list:
+            if feasible_resources.hint is not None:
+                hints[cloud] = feasible_resources.hint
+            if feasible_resources.resources_list:
                 # Assume feasible_resources is sorted by prices. Guaranteed by
                 # the implementation of get_feasible_launchable_resources and
                 # the underlying service_catalog filtering
-                cheapest = feasible_resources[0]
+                cheapest = feasible_resources.resources_list[0]
                 # Generate region/zone-specified resources.
                 launchable[resources].extend(
                     _make_launchables_for_valid_region_zones(cheapest))
-                cloud_candidates[cloud] = feasible_resources
+                cloud_candidates[cloud] = feasible_resources.resources_list
             else:
-                all_fuzzy_candidates.update(fuzzy_candidate_list)
-        if len(launchable[resources]) == 0:
+                all_fuzzy_candidates.update(
+                    feasible_resources.fuzzy_candidate_list)
+        if not launchable[resources]:
             clouds_str = str(clouds_list) if len(clouds_list) > 1 else str(
                 clouds_list[0])
             num_node_str = ''
@@ -1269,15 +1326,17 @@ def _fill_in_launchable_resources(
                                 f'{colorama.Fore.CYAN}'
                                 f'{sorted(all_fuzzy_candidates)}'
                                 f'{colorama.Style.RESET_ALL}')
-            else:
-                if resources.cpus is not None:
-                    logger.info('Try specifying a different CPU count, '
-                                'or add "+" to the end of the CPU count '
-                                'to allow for larger instances.')
-                if resources.memory is not None:
-                    logger.info('Try specifying a different memory size, '
-                                'or add "+" to the end of the memory size '
-                                'to allow for larger instances.')
+                else:
+                    if resources.cpus is not None:
+                        logger.info('Try specifying a different CPU count, '
+                                    'or add "+" to the end of the CPU count '
+                                    'to allow for larger instances.')
+                    if resources.memory is not None:
+                        logger.info('Try specifying a different memory size, '
+                                    'or add "+" to the end of the memory size '
+                                    'to allow for larger instances.')
+                for cloud, hint in hints.items():
+                    logger.info(f'{repr(cloud)}: {hint}')
 
         launchable[resources] = _filter_out_blocked_launchable_resources(
             launchable[resources], blocked_resources)

sky/provision/__init__.py

@@ -5,10 +5,10 @@ providers supported by SkyPilot need to follow.
 """
 import functools
 import inspect
+import typing
 from typing import Any, Dict, List, Optional, Type
 
 from sky import sky_logging
-from sky import status_lib
 # These provision.<cloud> modules should never fail even if underlying cloud SDK
 # dependencies are not installed. This is ensured by using sky.adaptors inside
 # these modules, for lazy loading of cloud SDKs.
@@ -19,9 +19,17 @@ from sky.provision import cudo
 from sky.provision import fluidstack
 from sky.provision import gcp
 from sky.provision import kubernetes
+from sky.provision import lambda_cloud
+from sky.provision import nebius
+from sky.provision import oci
 from sky.provision import runpod
+from sky.provision import vast
 from sky.provision import vsphere
 from sky.utils import command_runner
+from sky.utils import timeline
+
+if typing.TYPE_CHECKING:
+    from sky.utils import status_lib
 
 logger = sky_logging.init_logger(__name__)
 
@@ -39,6 +47,8 @@ def _route_to_cloud_impl(func):
             provider_name = kwargs.pop('provider_name')
 
         module_name = provider_name.lower()
+        if module_name == 'lambda':
+            module_name = 'lambda_cloud'
         module = globals().get(module_name)
         assert module is not None, f'Unknown provider: {module_name}'
 
@@ -55,13 +65,14 @@ def _route_to_cloud_impl(func):
 # pylint: disable=unused-argument
 
 
+@timeline.event
 @_route_to_cloud_impl
 def query_instances(
     provider_name: str,
     cluster_name_on_cloud: str,
     provider_config: Optional[Dict[str, Any]] = None,
     non_terminated_only: bool = True,
-) -> Dict[str, Optional[status_lib.ClusterStatus]]:
+) -> Dict[str, Optional['status_lib.ClusterStatus']]:
     """Query instances.
 
     Returns a dictionary of instance IDs and status.
@@ -155,6 +166,10 @@ def query_ports(
     return the endpoint without querying the cloud provider. If head_ip is not
     provided, the cloud provider will be queried to get the endpoint info.
 
+    The underlying implementation is responsible for retries and timeout, e.g.
+    kubernetes will wait for the service that expose the ports to be ready
+    before returning the endpoint info.
+
     Returns a dict with port as the key and a list of common.Endpoint.
     """
     del provider_name, provider_config, cluster_name_on_cloud  # unused
@@ -163,7 +178,7 @@ def query_ports(
 
 @_route_to_cloud_impl
 def wait_instances(provider_name: str, region: str, cluster_name_on_cloud: str,
-                   state: Optional[status_lib.ClusterStatus]) -> None:
+                   state: Optional['status_lib.ClusterStatus']) -> None:
     """Wait instances until they ends up in the given state."""
     raise NotImplementedError
 
@@ -182,12 +197,12 @@ def get_cluster_info(
 def get_command_runners(
     provider_name: str,
     cluster_info: common.ClusterInfo,
-    **crednetials: Dict[str, Any],
+    **credentials: Dict[str, Any],
 ) -> List[command_runner.CommandRunner]:
     """Get a command runner for the given cluster."""
     ip_list = cluster_info.get_feasible_ips()
     port_list = cluster_info.get_ssh_ports()
     return command_runner.SSHCommandRunner.make_runner_list(
         node_list=zip(ip_list, port_list),
-        **crednetials,
+        **credentials,
     )