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

sky/utils/resources_utils.py

@@ -1,14 +1,20 @@
 """Utility functions for resources."""
+import dataclasses
 import enum
 import itertools
+import json
+import math
 import re
 import typing
-from typing import List, Optional, Set
+from typing import Dict, List, Optional, Set, Union
 
+from sky import skypilot_config
+from sky.utils import registry
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky import backends
+    from sky import resources as resources_lib
 
 _PORT_RANGE_HINT_MSG = ('Invalid port range {}. Please use the format '
                         '"from-to", in which from <= to. e.g. "1-3".')
@@ -22,6 +28,7 @@ class DiskTier(enum.Enum):
     LOW = 'low'
     MEDIUM = 'medium'
     HIGH = 'high'
+    ULTRA = 'ultra'
     BEST = 'best'
 
     @classmethod
@@ -43,6 +50,18 @@ class DiskTier(enum.Enum):
         return types.index(self) <= types.index(other)
 
 
+@dataclasses.dataclass
+class ClusterName:
+    display_name: str
+    name_on_cloud: str
+
+    def __repr__(self) -> str:
+        return repr(self.display_name)
+
+    def __str__(self) -> str:
+        return self.display_name
+
+
 def check_port_str(port: str) -> None:
     if not port.isdigit():
         with ux_utils.print_exception_no_traceback():
@@ -118,29 +137,82 @@ def simplify_ports(ports: List[str]) -> List[str]:
     return port_set_to_ranges(port_ranges_to_set(ports))
 
 
+def format_resource(resource: 'resources_lib.Resources',
+                    simplify: bool = False) -> str:
+    if simplify:
+        cloud = resource.cloud
+        if resource.accelerators is None:
+            vcpu, _ = cloud.get_vcpus_mem_from_instance_type(
+                resource.instance_type)
+            hardware = f'vCPU={int(vcpu)}'
+        else:
+            hardware = f'{resource.accelerators}'
+        spot = '[Spot]' if resource.use_spot else ''
+        return f'{cloud}({spot}{hardware})'
+    else:
+        # accelerator_args is way too long.
+        # Convert from:
+        #  GCP(n1-highmem-8, {'tpu-v2-8': 1}, accelerator_args={'runtime_version': '2.12.0'}  # pylint: disable=line-too-long
+        # to:
+        #  GCP(n1-highmem-8, {'tpu-v2-8': 1}...)
+        pattern = ', accelerator_args={.*}'
+        launched_resource_str = re.sub(pattern, '...', str(resource))
+        return launched_resource_str
+
+
 def get_readable_resources_repr(handle: 'backends.CloudVmRayResourceHandle',
                                 simplify: bool = False) -> str:
     if (handle.launched_nodes is not None and
             handle.launched_resources is not None):
-        if simplify:
-            cloud = handle.launched_resources.cloud
-            if handle.launched_resources.accelerators is None:
-                vcpu, _ = cloud.get_vcpus_mem_from_instance_type(
-                    handle.launched_resources.instance_type)
-                hardware = f'vCPU={int(vcpu)}'
-            else:
-                hardware = f'{handle.launched_resources.accelerators}'
-            spot = '[Spot]' if handle.launched_resources.use_spot else ''
-            return f'{handle.launched_nodes}x {cloud}({spot}{hardware})'
-        else:
-            launched_resource_str = str(handle.launched_resources)
-            # accelerator_args is way too long.
-            # Convert from:
-            #  GCP(n1-highmem-8, {'tpu-v2-8': 1}, accelerator_args={'runtime_version': '2.12.0'}  # pylint: disable=line-too-long
-            # to:
-            #  GCP(n1-highmem-8, {'tpu-v2-8': 1}...)
-            pattern = ', accelerator_args={.*}'
-            launched_resource_str = re.sub(pattern, '...',
-                                           launched_resource_str)
-            return f'{handle.launched_nodes}x {launched_resource_str}'
+        return (f'{handle.launched_nodes}x '
+                f'{format_resource(handle.launched_resources, simplify)}')
     return _DEFAULT_MESSAGE_HANDLE_INITIALIZING
+
+
+def make_ray_custom_resources_str(
+        resource_dict: Optional[Dict[str, Union[int, float]]]) -> Optional[str]:
+    """Convert resources to Ray custom resources format."""
+    if resource_dict is None:
+        return None
+    # Ray does not allow fractional resources, so we need to ceil the values.
+    ceiled_dict = {k: math.ceil(v) for k, v in resource_dict.items()}
+    return json.dumps(ceiled_dict, separators=(',', ':'))
+
+
+@dataclasses.dataclass
+class FeasibleResources:
+    """Feasible resources returned by cloud.
+
+    Used to represent a collection of feasible resources returned by cloud,
+    any fuzzy candidates, and optionally a string hint if no feasible resources
+    are found.
+
+    Fuzzy candidates example: when the requested GPU is A100:1 but is not
+    available in a cloud/region, the fuzzy candidates are results of a fuzzy
+    search in the catalog that are offered in the location. E.g.,
+    ['A100-80GB:1', 'A100-80GB:2', 'A100-80GB:4', 'A100:8']
+    """
+    resources_list: List['resources_lib.Resources']
+    fuzzy_candidate_list: List[str]
+    hint: Optional[str]
+
+
+def need_to_query_reservations() -> bool:
+    """Checks if we need to query reservations from cloud APIs.
+
+    We need to query reservations if:
+    - The cloud has specific reservations.
+    - The cloud prioritizes reservations over on-demand instances.
+
+    This is useful to skip the potentially expensive reservation query for
+    clouds that do not use reservations.
+    """
+    for cloud_str in registry.CLOUD_REGISTRY.keys():
+        cloud_specific_reservations = skypilot_config.get_nested(
+            (cloud_str, 'specific_reservations'), None)
+        cloud_prioritize_reservations = skypilot_config.get_nested(
+            (cloud_str, 'prioritize_reservations'), False)
+        if (cloud_specific_reservations is not None or
+                cloud_prioritize_reservations):
+            return True
+    return False

sky/utils/rich_utils.py

@@ -1,16 +1,107 @@
 """Rich status spinner utils."""
 import contextlib
+import enum
+import logging
 import threading
-from typing import Union
+import typing
+from typing import Dict, Iterator, Optional, Tuple, Union
 
 import rich.console as rich_console
 
-console = rich_console.Console()
-_status = None
+from sky.utils import annotations
+from sky.utils import message_utils
+
+if typing.TYPE_CHECKING:
+    import requests
+
+console = rich_console.Console(soft_wrap=True)
+_statuses: Dict[str, Optional[Union['EncodedStatus',
+                                    'rich_console.Status']]] = {
+                                        'server': None,
+                                        'client': None,
+                                    }
+_status_nesting_level = 0
 
 _logging_lock = threading.RLock()
 
 
+class Control(enum.Enum):
+    """Control codes for the status spinner."""
+    INIT = 'rich_init'
+    START = 'rich_start'
+    STOP = 'rich_stop'
+    EXIT = 'rich_exit'
+    UPDATE = 'rich_update'
+
+    def encode(self, msg: str) -> str:
+        return f'<{self.value}>{msg}</{self.value}>'
+
+    @classmethod
+    def decode(cls, encoded_msg: str) -> Tuple[Optional['Control'], str]:
+        # Find the control code
+        control_str = None
+        for control in cls:
+            if f'<{control.value}>' in encoded_msg:
+                control_str = control.value
+                encoded_msg = encoded_msg.replace(f'<{control.value}>', '')
+                encoded_msg = encoded_msg.replace(f'</{control.value}>', '')
+                break
+        else:
+            return None, encoded_msg
+        return cls(control_str), encoded_msg
+
+
+class EncodedStatusMessage:
+    """A class to encode status messages."""
+
+    def __init__(self, msg: str):
+        self.msg = msg
+
+    def init(self) -> str:
+        return message_utils.encode_payload(Control.INIT.encode(self.msg))
+
+    def enter(self) -> str:
+        return message_utils.encode_payload(Control.START.encode(self.msg))
+
+    def exit(self) -> str:
+        return message_utils.encode_payload(Control.EXIT.encode(''))
+
+    def update(self, msg: str) -> str:
+        return message_utils.encode_payload(Control.UPDATE.encode(msg))
+
+    def stop(self) -> str:
+        return message_utils.encode_payload(Control.STOP.encode(''))
+
+    def start(self) -> str:
+        return message_utils.encode_payload(Control.START.encode(self.msg))
+
+
+class EncodedStatus:
+    """A class to encode status messages."""
+
+    def __init__(self, msg: str):
+        self.status = msg
+        self.encoded_msg = EncodedStatusMessage(msg)
+        print(self.encoded_msg.init(), end='', flush=True)
+
+    def __enter__(self):
+        print(self.encoded_msg.enter(), end='', flush=True)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        print(self.encoded_msg.exit(), end='', flush=True)
+
+    def update(self, msg: str):
+        self.status = msg
+        print(self.encoded_msg.update(msg), end='', flush=True)
+
+    def stop(self):
+        print(self.encoded_msg.stop(), end='', flush=True)
+
+    def start(self):
+        print(self.encoded_msg.start(), end='', flush=True)
+
+
 class _NoOpConsoleStatus:
     """An empty class for multi-threaded console.status."""
 
@@ -30,34 +121,172 @@ class _NoOpConsoleStatus:
         pass
 
 
+# TODO(SKY-1216): we need a wrapper for the rich.progress in our code as well.
+class _RevertibleStatus:
+    """A wrapper for status that can revert to previous message after exit."""
+
+    def __init__(self, message: str, status_type: str):
+        self.previous_message = None
+        self.status_type = status_type
+        status = _statuses[status_type]
+        if status is not None:
+            self.previous_message = status.status
+        self.message = message
+
+    def __enter__(self):
+        global _status_nesting_level
+        _statuses[self.status_type].update(self.message)
+        _status_nesting_level += 1
+        _statuses[self.status_type].__enter__()
+        return _statuses[self.status_type]
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        global _status_nesting_level
+        _status_nesting_level -= 1
+        if _status_nesting_level <= 0:
+            _status_nesting_level = 0
+            if _statuses[self.status_type] is not None:
+                _statuses[self.status_type].__exit__(exc_type, exc_val, exc_tb)
+                _statuses[self.status_type] = None
+        else:
+            _statuses[self.status_type].update(self.previous_message)
+
+    def update(self, *args, **kwargs):
+        _statuses[self.status_type].update(*args, **kwargs)
+
+    def stop(self):
+        _statuses[self.status_type].stop()
+
+    def start(self):
+        _statuses[self.status_type].start()
+
+
 def safe_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
     """A wrapper for multi-threaded console.status."""
     from sky import sky_logging  # pylint: disable=import-outside-toplevel
-    if (threading.current_thread() is threading.main_thread() and
+    if (annotations.is_on_api_server and
+            threading.current_thread() is threading.main_thread() and
             not sky_logging.is_silent()):
-        global _status
-        if _status is None:
-            _status = console.status(msg)
-        _status.update(msg)
-        return _status
+        if _statuses['server'] is None:
+            _statuses['server'] = EncodedStatus(msg)
+        return _RevertibleStatus(msg, 'server')
     return _NoOpConsoleStatus()
 
 
+def stop_safe_status():
+    """Stops all nested statuses.
+
+    This is useful when we need to stop all statuses, e.g., when we are going to
+    stream logs from user program and do not want it to interfere with the
+    spinner display.
+    """
+    if (threading.current_thread() is threading.main_thread() and
+            _statuses['server'] is not None):
+        _statuses['server'].stop()
+
+
 def force_update_status(msg: str):
     """Update the status message even if sky_logging.is_silent() is true."""
     if (threading.current_thread() is threading.main_thread() and
-            _status is not None):
-        _status.update(msg)
+            _statuses['server'] is not None):
+        _statuses['server'].update(msg)
 
 
 @contextlib.contextmanager
 def safe_logger():
-    logged = False
     with _logging_lock:
-        if _status is not None and _status._live.is_started:  # pylint: disable=protected-access
-            _status.stop()
-            yield
-            logged = True
-            _status.start()
-    if not logged:
+        client_status_obj = _statuses['client']
+
+        client_status_live = (client_status_obj is not None and
+                              client_status_obj._live.is_started)  # pylint: disable=protected-access
+        if client_status_live:
+            client_status_obj.stop()
         yield
+        if client_status_live:
+            client_status_obj.start()
+
+
+class RichSafeStreamHandler(logging.StreamHandler):
+
+    def emit(self, record: logging.LogRecord) -> None:
+        with safe_logger():
+            return super().emit(record)
+
+
+def client_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
+    """A wrapper for multi-threaded console.status."""
+    from sky import sky_logging  # pylint: disable=import-outside-toplevel
+    if (threading.current_thread() is threading.main_thread() and
+            not sky_logging.is_silent()):
+        if _statuses['client'] is None:
+            _statuses['client'] = console.status(msg)
+        return _RevertibleStatus(msg, 'client')
+    return _NoOpConsoleStatus()
+
+
+def decode_rich_status(
+        response: 'requests.Response') -> Iterator[Optional[str]]:
+    """Decode the rich status message from the response."""
+    decoding_status = None
+    try:
+        last_line = ''
+        # Iterate over the response content in chunks. We do not use iter_lines
+        # because it will strip the trailing newline characters, causing the
+        # progress bar ending with `\r` becomes a pyramid.
+        for encoded_msg in response.iter_content(chunk_size=None):
+            if encoded_msg is None:
+                return
+            encoded_msg = encoded_msg.decode('utf-8')
+            lines = encoded_msg.splitlines(keepends=True)
+
+            lines[0] = last_line + lines[0]
+            last_line = lines[-1]
+            # If the last line is not ended with `\r` or `\n` (with ending
+            # spaces stripped), it means the last line is not a complete line.
+            # We keep the last line in the buffer and continue.
+            if (not last_line.strip(' ').endswith('\r') and
+                    not last_line.strip(' ').endswith('\n')):
+                lines = lines[:-1]
+            else:
+                # Reset the buffer for the next line, as the last line is a
+                # complete line.
+                last_line = ''
+
+            for line in lines:
+                if line.endswith('\r\n'):
+                    # Replace `\r\n` with `\n`, as printing a line ends with
+                    # `\r\n` in linux will cause the line to be empty.
+                    line = line[:-2] + '\n'
+                is_payload, line = message_utils.decode_payload(
+                    line, raise_for_mismatch=False)
+                control = None
+                if is_payload:
+                    control, encoded_status = Control.decode(line)
+                if control is None:
+                    yield line
+                    continue
+
+                # control is not None, i.e. it is a rich status control message.
+                if threading.current_thread() is not threading.main_thread():
+                    yield None
+                    continue
+                if control == Control.INIT:
+                    decoding_status = client_status(encoded_status)
+                else:
+                    if decoding_status is None:
+                        # status may not be initialized if a user use --tail for
+                        # sky api logs.
+                        continue
+                    assert decoding_status is not None, (
+                        f'Rich status not initialized: {line}')
+                    if control == Control.UPDATE:
+                        decoding_status.update(encoded_status)
+                    elif control == Control.STOP:
+                        decoding_status.stop()
+                    elif control == Control.EXIT:
+                        decoding_status.__exit__(None, None, None)
+                    elif control == Control.START:
+                        decoding_status.start()
+    finally:
+        if decoding_status is not None:
+            decoding_status.__exit__(None, None, None)