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

sky/usage/usage_lib.py

@@ -3,7 +3,6 @@
 import contextlib
 import datetime
 import enum
-import inspect
 import json
 import os
 import time
@@ -11,20 +10,27 @@ import traceback
 import typing
 from typing import Any, Callable, Dict, List, Optional, Union
 
-import click
-import requests
-
 import sky
 from sky import sky_logging
+from sky.adaptors import common as adaptors_common
 from sky.usage import constants
 from sky.utils import common_utils
 from sky.utils import env_options
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    import inspect
+
+    import requests
+
     from sky import resources as resources_lib
-    from sky import status_lib
     from sky import task as task_lib
+    from sky.utils import status_lib
+else:
+    # requests and inspect cost ~100ms to load, which can be postponed to
+    # collection phase or skipped if user specifies no collection
+    requests = adaptors_common.LazyImport('requests')
+    inspect = adaptors_common.LazyImport('inspect')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -36,6 +42,7 @@ def _get_current_timestamp_ns() -> int:
 class MessageType(enum.Enum):
     """Types for messages to be sent to Loki."""
     USAGE = 'usage'
+    HEARTBEAT = 'heartbeat'
     # TODO(zhwu): Add more types, e.g., cluster_lifecycle.
 
 
@@ -59,8 +66,9 @@ class MessageToReport:
         properties = self.__dict__.copy()
         return {k: v for k, v in properties.items() if not k.startswith('_')}
 
-    def __repr__(self):
-        raise NotImplementedError
+    def __repr__(self) -> str:
+        d = self.get_properties()
+        return json.dumps(d)
 
 
 class UsageMessageToReport(MessageToReport):
@@ -75,7 +83,12 @@ class UsageMessageToReport(MessageToReport):
         self.sky_commit: str = sky.__commit__
 
         # Entry
-        self.cmd: str = common_utils.get_pretty_entry_point()
+        self.cmd: Optional[str] = common_utils.get_current_command()
+        # The entrypoint on the client side.
+        self.client_entrypoint: Optional[str] = None
+        # The entrypoint on the server side, where each request has a entrypoint
+        # and a single client_entrypoint can have multiple server-side
+        # entrypoints.
         self.entrypoint: Optional[str] = None  # entrypoint_context
         #: Whether entrypoint is called by sky internal code.
         self.internal: bool = False  # set_internal
@@ -140,6 +153,7 @@ class UsageMessageToReport(MessageToReport):
         #: Requested number of nodes
         self.task_num_nodes: Optional[int] = None  # update_actual_task
         # YAMLs converted to JSON.
+        # TODO: include the skypilot config used in task yaml.
         self.user_task_yaml: Optional[List[Dict[
             str, Any]]] = None  # update_user_task_yaml
         self.actual_task: Optional[List[Dict[str,
@@ -151,11 +165,14 @@ class UsageMessageToReport(MessageToReport):
         self.exception: Optional[str] = None  # entrypoint_context
         self.stacktrace: Optional[str] = None  # entrypoint_context
 
-    def __repr__(self) -> str:
-        d = self.get_properties()
-        return json.dumps(d)
+        # Whether API server is deployed remotely.
+        self.using_remote_api_server: bool = (
+            common_utils.get_using_remote_api_server())
 
     def update_entrypoint(self, msg: str):
+        if self.client_entrypoint is None:
+            self.client_entrypoint = common_utils.get_current_client_entrypoint(
+                msg)
         self.entrypoint = msg
 
     def set_internal(self):
@@ -200,9 +217,11 @@ class UsageMessageToReport(MessageToReport):
     def update_ray_yaml(self, yaml_config_or_path: Union[Dict, str]):
         if self.ray_yamls is None:
             self.ray_yamls = []
-        self.ray_yamls.extend(
-            prepare_json_from_yaml_config(yaml_config_or_path))
-        self.num_tried_regions = len(self.ray_yamls)
+        if self.num_tried_regions is None:
+            self.num_tried_regions = 0
+        # Only keep the latest ray yaml to reduce the size of the message.
+        self.ray_yamls = prepare_json_from_yaml_config(yaml_config_or_path)
+        self.num_tried_regions += 1
 
     def update_cluster_name(self, cluster_name: Union[List[str], str]):
         if isinstance(cluster_name, str):
@@ -266,16 +285,43 @@ class UsageMessageToReport(MessageToReport):
                                            name_or_fn)
 
 
+class HeartbeatMessageToReport(MessageToReport):
+    """Message to be reported to Grafana Loki for heartbeat on a cluster."""
+
+    def __init__(self, interval_seconds: int = 600):
+        super().__init__(constants.USAGE_MESSAGE_SCHEMA_VERSION)
+        # This interval_seconds is mainly for recording the heartbeat interval
+        # in the heartbeat message, so that the collector can use it.
+        self.interval_seconds = interval_seconds
+
+    def get_properties(self) -> Dict[str, Any]:
+        properties = super().get_properties()
+        # The run id is set by the skylet, which will always be the same for
+        # the entire lifetime of the run.
+        with open(os.path.expanduser(constants.USAGE_RUN_ID_FILE),
+                  'r',
+                  encoding='utf-8') as f:
+            properties['run_id'] = f.read().strip()
+        return properties
+
+
 class MessageCollection:
     """A collection of messages."""
 
     def __init__(self):
-        self._messages = {MessageType.USAGE: UsageMessageToReport()}
+        self._messages = {
+            MessageType.USAGE: UsageMessageToReport(),
+            MessageType.HEARTBEAT: HeartbeatMessageToReport()
+        }
 
     @property
-    def usage(self):
+    def usage(self) -> UsageMessageToReport:
         return self._messages[MessageType.USAGE]
 
+    @property
+    def heartbeat(self) -> HeartbeatMessageToReport:
+        return self._messages[MessageType.HEARTBEAT]
+
     def reset(self, message_type: MessageType):
         self._messages[message_type] = self._messages[message_type].__class__()
 
@@ -299,13 +345,25 @@ def _send_to_loki(message_type: MessageType):
 
     message = messages[message_type]
 
+    # In case the message has no start time, set it to the current time.
+    message.start()
     message.send_time = _get_current_timestamp_ns()
-    log_timestamp = message.start_time
+    # Use send time instead of start time to avoid the message being dropped
+    # by Loki, due to the timestamp being too old. We still have the start time
+    # in the message for dashboard.
+    log_timestamp = message.send_time
 
     environment = 'prod'
     if env_options.Options.IS_DEVELOPER.get():
         environment = 'dev'
-    prom_labels = {'type': message_type.value, 'environment': environment}
+    prom_labels = {
+        'type': message_type.value,
+        'environment': environment,
+        'schema_version': message.schema_version,
+    }
+    if message_type == MessageType.USAGE:
+        prom_labels['new_cluster'] = (message.original_cluster_status != 'UP'
+                                      and message.final_cluster_status == 'UP')
 
     headers = {'Content-type': 'application/json'}
     payload = {
@@ -383,7 +441,7 @@ def prepare_json_from_yaml_config(
 def _send_local_messages():
     """Send all messages not been uploaded to Loki."""
     for msg_type, message in messages.items():
-        if not message.message_sent:
+        if not message.message_sent and msg_type != MessageType.HEARTBEAT:
             # Avoid the fallback entrypoint to send the message again
             # in normal case.
             try:
@@ -393,17 +451,26 @@ def _send_local_messages():
                              f'exception caught: {type(e)}({e})')
 
 
-@contextlib.contextmanager
-def entrypoint_context(name: str, fallback: bool = False):
-    """Context manager for entrypoint.
+def store_exception(e: Union[Exception, SystemExit, KeyboardInterrupt]) -> None:
+    with ux_utils.enable_traceback():
+        if hasattr(e, 'stacktrace') and e.stacktrace is not None:
+            messages.usage.stacktrace = e.stacktrace
+        else:
+            trace = traceback.format_exc()
+            messages.usage.stacktrace = trace
+        if hasattr(e, 'detailed_reason') and e.detailed_reason is not None:
+            messages.usage.stacktrace += '\nDetails: ' + e.detailed_reason
+        messages.usage.exception = common_utils.remove_color(
+            common_utils.format_exception(e))
 
-    The context manager will send the usage message to Loki when exiting.
-    The message will only be sent at the outermost level of the context.
 
-    When the outermost context does not cover all the codepaths, an
-    additional entrypoint_context with fallback=True can be used to wrap
-    the global entrypoint to catch any exceptions that are not caught.
-    """
+def send_heartbeat(interval_seconds: int = 600):
+    messages.heartbeat.interval_seconds = interval_seconds
+    _send_to_loki(MessageType.HEARTBEAT)
+
+
+def maybe_show_privacy_policy():
+    """Show the privacy policy if it is not already shown."""
     # Show the policy message only when the entrypoint is used.
     # An indicator for PRIVACY_POLICY has already been shown.
     privacy_policy_indicator = os.path.expanduser(constants.PRIVACY_POLICY_PATH)
@@ -411,10 +478,22 @@ def entrypoint_context(name: str, fallback: bool = False):
         os.makedirs(os.path.dirname(privacy_policy_indicator), exist_ok=True)
         try:
             with open(privacy_policy_indicator, 'x', encoding='utf-8'):
-                click.secho(constants.USAGE_POLICY_MESSAGE, fg='yellow')
+                logger.info(constants.USAGE_POLICY_MESSAGE)
         except FileExistsError:
             pass
 
+
+@contextlib.contextmanager
+def entrypoint_context(name: str, fallback: bool = False):
+    """Context manager for entrypoint.
+
+    The context manager will send the usage message to Loki when exiting.
+    The message will only be sent at the outermost level of the context.
+
+    When the outermost context does not cover all the codepaths, an
+    additional entrypoint_context with fallback=True can be used to wrap
+    the global entrypoint to catch any exceptions that are not caught.
+    """
     is_entry = messages.usage.entrypoint is None
     if is_entry and not fallback:
         for message in messages.values():
@@ -428,13 +507,7 @@ def entrypoint_context(name: str, fallback: bool = False):
     try:
         yield
     except (Exception, SystemExit, KeyboardInterrupt) as e:
-        with ux_utils.enable_traceback():
-            trace = traceback.format_exc()
-            messages.usage.stacktrace = trace
-            if hasattr(e, 'detailed_reason') and e.detailed_reason is not None:
-                messages.usage.stacktrace += '\nDetails: ' + e.detailed_reason
-            messages.usage.exception = common_utils.remove_color(
-                common_utils.format_exception(e))
+        store_exception(e)
         raise
     finally:
         if fallback:
@@ -442,7 +515,27 @@ def entrypoint_context(name: str, fallback: bool = False):
         _send_local_messages()
 
 
-def entrypoint(name_or_fn: Union[str, Callable], fallback: bool = False):
+T = typing.TypeVar('T')
+
+
+@typing.overload
+def entrypoint(
+        name_or_fn: str,
+        fallback: bool = False
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    ...
+
+
+@typing.overload
+def entrypoint(name_or_fn: Callable[..., T],
+               fallback: bool = False) -> Callable[..., T]:
+    ...
+
+
+def entrypoint(
+    name_or_fn: Union[str, Callable[..., T]],
+    fallback: bool = False
+) -> Union[Callable[..., T], Callable[[Callable[..., T]], Callable[..., T]]]:
     return common_utils.make_decorator(entrypoint_context,
                                        name_or_fn,
                                        fallback=fallback)

sky/utils/accelerator_registry.py

@@ -3,12 +3,13 @@ import typing
 from typing import Optional
 
 from sky.clouds import service_catalog
+from sky.utils import rich_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky import clouds
 
-# Canonicalized names of all accelerators (except TPUs) supported by SkyPilot.
+# Canonical names of all accelerators (except TPUs) supported by SkyPilot.
 # NOTE: Must include accelerators supported for local clusters.
 #
 # 1. What if a name is in this list, but not in any catalog?
@@ -30,30 +31,10 @@ if typing.TYPE_CHECKING:
 #
 # Append its case-sensitive canonical name to this list. The name must match
 # `AcceleratorName` in the service catalog.
-_ACCELERATORS = [
-    'A100',
-    'A10G',
-    'Gaudi HL-205',
-    'Inferentia',
-    'Trainium',
-    'K520',
-    'K80',
-    'M60',
-    'Radeon Pro V520',
-    'T4',
-    'T4g',
-    'V100',
-    'V100-32GB',
-    'Virtex UltraScale (VU9P)',
-    'A10',
-    'A100-80GB',
-    'P100',
-    'P40',
-    'Radeon MI25',
-    'P4',
-    'L4',
-    'H100',
-]
+
+# Use a cached version of accelerators to cloud mapping, so that we don't have
+# to download and read the catalog file for every cloud locally.
+_accelerator_df = service_catalog.common.read_catalog('common/accelerators.csv')
 
 # List of non-GPU accelerators that are supported by our backend for job queue
 # scheduling.
@@ -77,42 +58,45 @@ def canonicalize_accelerator_name(accelerator: str,
     """Returns the canonical accelerator name."""
     cloud_str = None
     if cloud is not None:
-        cloud_str = str(cloud).lower()
+        cloud_str = str(cloud)
 
     # TPU names are always lowercase.
     if accelerator.lower().startswith('tpu-'):
         return accelerator.lower()
 
     # Common case: do not read the catalog files.
-    mapping = {name.lower(): name for name in _ACCELERATORS}
-    if accelerator.lower() in mapping:
-        return mapping[accelerator.lower()]
-
-    # _ACCELERATORS may not be comprehensive.
-    # Users may manually add new accelerators to the catalogs, or download new
-    # catalogs (that have new accelerators) without upgrading SkyPilot.
-    # To cover such cases, we should search the accelerator name
-    # in the service catalog.
-    searched = service_catalog.list_accelerators(name_filter=accelerator,
-                                                 case_sensitive=False,
-                                                 clouds=cloud_str)
-    names = list(searched.keys())
-
-    # Exact match.
-    if accelerator in names:
+    df = _accelerator_df[_accelerator_df['AcceleratorName'].str.contains(
+        accelerator, case=False, regex=True)]
+    names = []
+    for name, clouds in df[['AcceleratorName', 'Clouds']].values:
+        if accelerator.lower() == name.lower():
+            return name
+        if cloud_str is None or cloud_str in clouds:
+            names.append(name)
+
+    # Look for Kubernetes accelerators online if the accelerator is not found
+    # in the public cloud catalog. This is to make sure custom accelerators
+    # on Kubernetes can be correctly canonicalized.
+    if not names and cloud_str in ['kubernetes', None]:
+        with rich_utils.safe_status(
+                ux_utils.spinner_message('Listing accelerators on Kubernetes')):
+            searched = service_catalog.list_accelerators(
+                name_filter=accelerator,
+                case_sensitive=False,
+                clouds=cloud_str,
+            )
+        names = list(searched.keys())
+        if accelerator in names:
+            return accelerator
+
+    if not names:
+        # If no match is found, it is fine to return the original name, as
+        # the custom accelerator might be on kubernetes cluster.
         return accelerator
 
     if len(names) == 1:
         return names[0]
-
-    # Do not print an error message here. Optimizer will handle it.
-    if len(names) == 0:
-        return accelerator
-
-    # Currently unreachable.
-    # This can happen if catalogs have the same accelerator with
-    # different names (e.g., A10g and A10G).
-    assert len(names) > 1
+    assert len(names) > 1, names
     with ux_utils.print_exception_no_traceback():
         raise ValueError(f'Accelerator name {accelerator!r} is ambiguous. '
                          f'Please choose one of {names}.')

sky/utils/admin_policy_utils.py

@@ -0,0 +1,147 @@
+"""Admin policy utils."""
+import copy
+import importlib
+import os
+import tempfile
+from typing import Optional, Tuple, Union
+
+import colorama
+
+from sky import admin_policy
+from sky import dag as dag_lib
+from sky import exceptions
+from sky import sky_logging
+from sky import skypilot_config
+from sky import task as task_lib
+from sky.utils import common_utils
+from sky.utils import config_utils
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _get_policy_cls(
+        policy: Optional[str]) -> Optional[admin_policy.AdminPolicy]:
+    """Gets admin-defined policy."""
+    if policy is None:
+        return None
+    try:
+        module_path, class_name = policy.rsplit('.', 1)
+        module = importlib.import_module(module_path)
+    except ImportError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise ImportError(
+                f'Failed to import policy module: {policy}. '
+                'Please check if the module is installed in your Python '
+                'environment.') from e
+
+    try:
+        policy_cls = getattr(module, class_name)
+    except AttributeError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise AttributeError(
+                f'Could not find {class_name} class in module {module_path}. '
+                'Please check with your policy admin for details.') from e
+
+    # Check if the module implements the AdminPolicy interface.
+    if not issubclass(policy_cls, admin_policy.AdminPolicy):
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'Policy class {policy!r} does not implement the AdminPolicy '
+                'interface. Please check with your policy admin for details.')
+    return policy_cls
+
+
+def apply(
+    entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    use_mutated_config_in_current_request: bool = True,
+    request_options: Optional[admin_policy.RequestOptions] = None,
+) -> Tuple['dag_lib.Dag', config_utils.Config]:
+    """Applies an admin policy (if registered) to a DAG or a task.
+
+    It mutates a Dag by applying any registered admin policy and also
+    potentially updates (controlled by `use_mutated_config_in_current_request`)
+    the global SkyPilot config if there is any changes made by the policy.
+
+    Args:
+        dag: The dag to be mutated by the policy.
+        use_mutated_config_in_current_request: Whether to use the mutated
+            config in the current request.
+        request_options: Additional options user passed for the current request.
+
+    Returns:
+        - The new copy of dag after applying the policy
+        - The new copy of skypilot config after applying the policy.
+    """
+    if isinstance(entrypoint, task_lib.Task):
+        dag = dag_lib.Dag()
+        dag.add(entrypoint)
+    else:
+        dag = entrypoint
+
+    policy = skypilot_config.get_nested(('admin_policy',), None)
+    policy_cls = _get_policy_cls(policy)
+    if policy_cls is None:
+        return dag, skypilot_config.to_dict()
+
+    logger.info(f'Applying policy: {policy}')
+    original_config = skypilot_config.to_dict()
+    config = copy.deepcopy(original_config)
+    mutated_dag = dag_lib.Dag()
+    mutated_dag.name = dag.name
+
+    mutated_config = None
+    for task in dag.tasks:
+        user_request = admin_policy.UserRequest(task, config, request_options)
+        try:
+            mutated_user_request = policy_cls.validate_and_mutate(user_request)
+        except Exception as e:  # pylint: disable=broad-except
+            with ux_utils.print_exception_no_traceback():
+                raise exceptions.UserRequestRejectedByPolicy(
+                    f'{colorama.Fore.RED}User request rejected by policy '
+                    f'{policy!r}{colorama.Fore.RESET}: '
+                    f'{common_utils.format_exception(e, use_bracket=True)}'
+                ) from e
+        if mutated_config is None:
+            mutated_config = mutated_user_request.skypilot_config
+        else:
+            if mutated_config != mutated_user_request.skypilot_config:
+                # In the case of a pipeline of tasks, the mutated config
+                # generated should remain the same for all tasks for now for
+                # simplicity.
+                # TODO(zhwu): We should support per-task mutated config or
+                # allowing overriding required global config in task YAML.
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.UserRequestRejectedByPolicy(
+                        'All tasks must have the same SkyPilot config after '
+                        'applying the policy. Please check with your policy '
+                        'admin for details.')
+        mutated_dag.add(mutated_user_request.task)
+    assert mutated_config is not None, dag
+
+    # Update the new_dag's graph with the old dag's graph
+    for u, v in dag.graph.edges:
+        u_idx = dag.tasks.index(u)
+        v_idx = dag.tasks.index(v)
+        mutated_dag.graph.add_edge(mutated_dag.tasks[u_idx],
+                                   mutated_dag.tasks[v_idx])
+
+    if (use_mutated_config_in_current_request and
+            original_config != mutated_config):
+        with tempfile.NamedTemporaryFile(
+                delete=False,
+                mode='w',
+                prefix='policy-mutated-skypilot-config-',
+                suffix='.yaml') as temp_file:
+
+            common_utils.dump_yaml(temp_file.name, dict(**mutated_config))
+            os.environ[skypilot_config.ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
+            logger.debug(f'Updated SkyPilot config: {temp_file.name}')
+            # TODO(zhwu): This is not a clean way to update the SkyPilot config,
+            # because we are resetting the global context for a single DAG,
+            # which is conceptually weird.
+            importlib.reload(skypilot_config)
+
+    logger.debug(f'Mutated user request: {mutated_user_request}')
+    mutated_dag.policy_applied = True
+    return mutated_dag, mutated_config

sky/utils/annotations.py

@@ -0,0 +1,51 @@
+"""Annotations for public APIs."""
+
+import functools
+from typing import Callable, Literal
+
+# Whether the current process is a SkyPilot API server process.
+is_on_api_server = True
+FUNCTIONS_NEED_RELOAD_CACHE = []
+
+
+def client_api(func):
+    """Mark a function as a client-side API.
+
+    Code invoked by server-side functions will find annotations.is_on_api_server
+    to be True, so they can have some server-side handling.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global is_on_api_server
+        is_on_api_server = False
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
+              **lru_cache_kwargs) -> Callable:
+    """LRU cache decorator for functions.
+
+    This decorator allows us to track which functions need to be reloaded for a
+    new request using the scope argument.
+
+    Args:
+        scope: Whether the cache is global or request-specific, i.e. needs to be
+            reloaded for a new request.
+        lru_cache_args: Arguments for functools.lru_cache.
+        lru_cache_kwargs: Keyword arguments for functools.lru_cache.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        if scope == 'global':
+            return functools.lru_cache(*lru_cache_args,
+                                       **lru_cache_kwargs)(func)
+        else:
+            cached_func = functools.lru_cache(*lru_cache_args,
+                                              **lru_cache_kwargs)(func)
+            FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
+            return cached_func
+
+    return decorator