skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/utils/resources_utils.py

@@ -4,11 +4,12 @@ import enum
 import itertools
 import json
 import math
-import re
 import typing
-from typing import Dict, List, Optional, Set, Union
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 from sky import skypilot_config
+from sky.skylet import constants
+from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import ux_utils
 
@@ -50,6 +51,48 @@ class DiskTier(enum.Enum):
         return types.index(self) <= types.index(other)
 
 
+class NetworkTier(enum.Enum):
+    """All network tiers supported by SkyPilot."""
+    STANDARD = 'standard'
+    BEST = 'best'
+
+    @classmethod
+    def supported_tiers(cls) -> List[str]:
+        return [tier.value for tier in cls]
+
+    @classmethod
+    def cli_help_message(cls) -> str:
+        return (
+            f'Network tier. Could be one of {", ".join(cls.supported_tiers())}'
+            f'. If {cls.BEST.value} is specified, use the best network tier '
+            'available on the specified instance. '
+            f'Default: {cls.STANDARD.value}')
+
+    @classmethod
+    def from_str(cls, tier: str) -> 'NetworkTier':
+        if tier not in cls.supported_tiers():
+            raise ValueError(f'Invalid network tier: {tier}')
+        return cls(tier)
+
+    def __le__(self, other: 'NetworkTier') -> bool:
+        types = list(NetworkTier)
+        return types.index(self) <= types.index(other)
+
+
+class StorageType(enum.Enum):
+    """Storage type."""
+    # Durable network storage, e.g. GCP persistent disks
+    NETWORK = 'network'
+    # Local instance storage, e.g. GCP local SSDs
+    INSTANCE = 'instance'
+
+
+class DiskAttachMode(enum.Enum):
+    """Disk attach mode."""
+    READ_ONLY = 'read_only'
+    READ_WRITE = 'read_write'
+
+
 @dataclasses.dataclass
 class ClusterName:
     display_name: str
@@ -138,35 +181,81 @@ def simplify_ports(ports: List[str]) -> List[str]:
 
 
 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})'
+                    simplified_only: bool = False) -> Tuple[str, Optional[str]]:
+    resource = resource.assert_launchable()
+    is_k8s = str(resource.cloud).lower() == 'kubernetes'
+    if resource.accelerators is None or is_k8s or not simplified_only:
+        vcpu, mem = resource.cloud.get_vcpus_mem_from_instance_type(
+            resource.instance_type)
+
+    elements_simple = []
+    elements_full = []
+
+    if resource.accelerators is not None:
+        acc, count = list(resource.accelerators.items())[0]
+        elements_simple.append(f'gpus={acc}:{count}')
+        elements_full.append(f'gpus={acc}:{count}')
+
+    if (resource.accelerators is None or is_k8s):
+        if vcpu is not None:
+            elements_simple.append(f'cpus={int(vcpu)}')
+            elements_full.append(f'cpus={int(vcpu)}')
+        if mem is not None:
+            elements_simple.append(f'mem={int(mem)}')
+            elements_full.append(f'mem={int(mem)}')
+    elif not simplified_only:
+        if vcpu is not None:
+            elements_full.append(f'cpus={int(vcpu)}')
+        if mem is not None:
+            elements_full.append(f'mem={int(mem)}')
+
+    if not is_k8s:
+        instance_type_full = resource.instance_type
+        instance_type_simple = common_utils.truncate_long_string(
+            instance_type_full, 15)
+        elements_simple.append(instance_type_simple)
+        elements_full.append(instance_type_full)
+    elements_simple.append('...')
+    if not simplified_only:
+        image_id = resource.image_id
+        if image_id is not None:
+            if None in image_id:
+                elements_full.append(f'image_id={image_id[None]}')
+            else:
+                elements_full.append(f'image_id={image_id}')
+        elements_full.append(f'disk={resource.disk_size}')
+        disk_tier = resource.disk_tier
+        if disk_tier is not None:
+            elements_full.append(f'disk_tier={disk_tier.value}')
+        ports = resource.ports
+        if ports is not None:
+            elements_full.append(f'ports={ports}')
+
+    spot = '[spot]' if resource.use_spot else ''
+    resources_str_simple = (
+        f'{spot}({"" if not elements_simple else ", ".join(elements_simple)})')
+    if simplified_only:
+        return resources_str_simple, None
     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:
+        resources_str_full = (
+            f'{spot}({"" if not elements_full else ", ".join(elements_full)})')
+        return resources_str_simple, resources_str_full
+
+
+def get_readable_resources_repr(
+        handle: 'backends.CloudVmRayResourceHandle',
+        simplified_only: bool = False) -> Tuple[str, Optional[str]]:
+    resource_str_simple, resource_str_full = format_resource(
+        handle.launched_resources, simplified_only)
+    if not simplified_only:
+        assert resource_str_full is not None
     if (handle.launched_nodes is not None and
             handle.launched_resources is not None):
-        return (f'{handle.launched_nodes}x '
-                f'{format_resource(handle.launched_resources, simplify)}')
-    return _DEFAULT_MESSAGE_HANDLE_INITIALIZING
+        return (f'{handle.launched_nodes}x{resource_str_simple}',
+                None if simplified_only else
+                f'{handle.launched_nodes}x{resource_str_full}')
+    return (_DEFAULT_MESSAGE_HANDLE_INITIALIZING,
+            _DEFAULT_MESSAGE_HANDLE_INITIALIZING)
 
 
 def make_ray_custom_resources_str(
@@ -208,10 +297,18 @@ def need_to_query_reservations() -> bool:
     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)
+        cloud_specific_reservations = (
+            skypilot_config.get_effective_region_config(
+                cloud=cloud_str,
+                region=None,
+                keys=('specific_reservations',),
+                default_value=None))
+        cloud_prioritize_reservations = (
+            skypilot_config.get_effective_region_config(
+                cloud=cloud_str,
+                region=None,
+                keys=('prioritize_reservations',),
+                default_value=False))
         if (cloud_specific_reservations is not None or
                 cloud_prioritize_reservations):
             return True
@@ -248,6 +345,7 @@ def make_launchables_for_valid_region_zones(
     launchables = []
     regions = launchable_resources.get_valid_regions_for_launchable()
     for region in regions:
+        assert launchable_resources.cloud is not None, 'Cloud must be specified'
         optimize_by_zone = (override_optimize_by_zone or
                             launchable_resources.cloud.optimize_by_zone())
         # It is possible that we force the optimize_by_zone but some clouds
@@ -266,3 +364,122 @@ def make_launchables_for_valid_region_zones(
             # Batch the requests at the granularity of a single region.
             launchables.append(launchable_resources.copy(region=region.name))
     return launchables
+
+
+def parse_memory_resource(resource_qty_str: Union[str, int, float],
+                          field_name: str,
+                          ret_type: type = int,
+                          unit: str = 'gb',
+                          allow_plus: bool = False,
+                          allow_x: bool = False,
+                          allow_rounding: bool = False) -> str:
+    """Returns memory size in chosen units given a resource quantity string.
+
+    Args:
+        resource_qty_str: Resource quantity string
+        unit: Unit to convert to
+        allow_plus: Whether to allow '+' prefix
+        allow_x: Whether to allow 'x' suffix
+    """
+    assert unit in constants.MEMORY_SIZE_UNITS, f'Invalid unit: {unit}'
+
+    error_msg = (f'"{field_name}" field should be a '
+                 f'{constants.MEMORY_SIZE_PATTERN}+?,'
+                 f' got {resource_qty_str}')
+
+    resource_str = str(resource_qty_str)
+
+    # Handle plus and x suffixes, x is only used internally for jobs controller
+    plus = ''
+    if resource_str.endswith('+'):
+        if allow_plus:
+            resource_str = resource_str[:-1]
+            plus = '+'
+        else:
+            raise ValueError(error_msg)
+
+    x = ''
+    if resource_str.endswith('x'):
+        if allow_x:
+            resource_str = resource_str[:-1]
+            x = 'x'
+        else:
+            raise ValueError(error_msg)
+
+    try:
+        # We assume it is already in the wanted units to maintain backwards
+        # compatibility
+        ret_type(resource_str)
+        return f'{resource_str}{plus}{x}'
+    except ValueError:
+        pass
+
+    resource_str = resource_str.lower()
+    for mem_unit, multiplier in constants.MEMORY_SIZE_UNITS.items():
+        if resource_str.endswith(mem_unit):
+            try:
+                value = ret_type(resource_str[:-len(mem_unit)])
+                converted = (value * multiplier /
+                             constants.MEMORY_SIZE_UNITS[unit])
+                if not allow_rounding and ret_type(converted) != converted:
+                    raise ValueError(error_msg)
+                converted = ret_type(converted)
+                return f'{converted}{plus}{x}'
+            except ValueError:
+                continue
+
+    raise ValueError(error_msg)
+
+
+def parse_time_minutes(time: str) -> int:
+    """Convert a time string to minutes.
+
+    Args:
+        time: Time string with optional unit suffix (e.g., '30m', '2h', '1d')
+
+    Returns:
+        Time in minutes as an integer
+    """
+    time_str = str(time)
+
+    if time_str.isdecimal():
+        # We assume it is already in minutes to maintain backwards
+        # compatibility
+        return int(time_str)
+
+    time_str = time_str.lower()
+    for unit, multiplier in constants.TIME_UNITS.items():
+        if time_str.endswith(unit):
+            try:
+                value = float(time_str[:-len(unit)])
+                final_value = math.ceil(value * multiplier)
+                if final_value >= 0:
+                    return final_value
+            except ValueError:
+                continue
+
+    raise ValueError(f'Invalid time format: {time}')
+
+
+def normalize_any_of_resources_config(
+        any_of: List[Dict[str, Any]]) -> Tuple[str, ...]:
+    """Normalize a list of any_of resources config to a canonical form.
+
+    Args:
+        any_of: A list of any_of resources config.
+
+    Returns:
+        A normalized tuple representation that can be compared for equality.
+        Two lists with the same resource configurations in different orders
+        will produce the same normalized result.
+    """
+    if not any_of:
+        return tuple()
+
+    # Convert each config to JSON string with sorted keys, then sort the list
+    normalized_configs = [
+        json.dumps(config, sort_keys=True, separators=(',', ':'))
+        for config in any_of
+    ]
+
+    return tuple(sorted(normalized_configs))

sky/utils/rich_utils.py

@@ -1,28 +1,53 @@
 """Rich status spinner utils."""
 import contextlib
+import contextvars
 import enum
 import logging
 import threading
 import typing
-from typing import Dict, Iterator, Optional, Tuple, Union
+from typing import Callable, Iterator, Optional, Tuple, Union
 
+from sky import exceptions
 from sky.adaptors import common as adaptors_common
 from sky.utils import annotations
+from sky.utils import context
 from sky.utils import message_utils
 from sky.utils import rich_console_utils
 
 if typing.TYPE_CHECKING:
+    import aiohttp
     import requests
     import rich.console as rich_console
 else:
     requests = adaptors_common.LazyImport('requests')
     rich_console = adaptors_common.LazyImport('rich.console')
+    aiohttp = adaptors_common.LazyImport('aiohttp')
+
+GeneralStatus = Union['rich_console.Status', 'EncodedStatus']
+
+_client_status: Optional[GeneralStatus] = None
+_server_status: contextvars.ContextVar[
+    Optional[GeneralStatus]] = contextvars.ContextVar('server_status',
+                                                      default=None)
+
+
+def _get_client_status() -> Optional[GeneralStatus]:
+    return _client_status
+
+
+def _get_server_status() -> Optional[GeneralStatus]:
+    return _server_status.get()
+
+
+def _set_client_status(status: Optional[GeneralStatus]):
+    global _client_status
+    _client_status = status
+
+
+def _set_server_status(status: Optional[GeneralStatus]):
+    _server_status.set(status)
+
 
-_statuses: Dict[str, Optional[Union['EncodedStatus',
-                                    'rich_console.Status']]] = {
-                                        'server': None,
-                                        'client': None,
-                                    }
 _status_nesting_level = 0
 
 _logging_lock = threading.RLock()
@@ -35,6 +60,8 @@ class Control(enum.Enum):
     STOP = 'rich_stop'
     EXIT = 'rich_exit'
     UPDATE = 'rich_update'
+    HEARTBEAT = 'heartbeat'
+    RETRY = 'retry'
 
     def encode(self, msg: str) -> str:
         return f'<{self.value}>{msg}</{self.value}>'
@@ -128,20 +155,22 @@ class _NoOpConsoleStatus:
 class _RevertibleStatus:
     """A wrapper for status that can revert to previous message after exit."""
 
-    def __init__(self, message: str, status_type: str):
+    def __init__(self, message: str, get_status_fn: Callable[[], GeneralStatus],
+                 set_status_fn: Callable[[Optional[GeneralStatus]], None]):
         self.previous_message = None
-        self.status_type = status_type
-        status = _statuses[status_type]
+        self.get_status_fn = get_status_fn
+        self.set_status_fn = set_status_fn
+        status = self.get_status_fn()
         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)
+        self.get_status_fn().update(self.message)
         _status_nesting_level += 1
-        _statuses[self.status_type].__enter__()
-        return _statuses[self.status_type]
+        self.get_status_fn().__enter__()
+        return self.get_status_fn()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         # We use the same lock with the `safe_logger` to avoid the following 2
@@ -160,32 +189,49 @@ class _RevertibleStatus:
             _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
+                if self.get_status_fn() is not None:
+                    self.get_status_fn().__exit__(exc_type, exc_val, exc_tb)
+                    self.set_status_fn(None)
             else:
-                _statuses[self.status_type].update(self.previous_message)
+                if self.previous_message is not None:
+                    self.get_status_fn().update(self.previous_message)
 
     def update(self, *args, **kwargs):
-        _statuses[self.status_type].update(*args, **kwargs)
+        self.get_status_fn().update(*args, **kwargs)
 
     def stop(self):
-        _statuses[self.status_type].stop()
+        self.get_status_fn().stop()
 
     def start(self):
-        _statuses[self.status_type].start()
+        self.get_status_fn().start()
+
+
+def _is_thread_safe() -> bool:
+    """Check if the current status context is thread-safe.
+
+    We are thread-safe if we are on the main thread or the server_status is
+    context-local, i.e. an async context has been initialized.
+    """
+    return (threading.current_thread() is threading.main_thread() or
+            context.get() is not None)
 
 
 def safe_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
-    """A wrapper for multi-threaded console.status."""
+    """A wrapper for multi-threaded server-side console.status.
+
+    This function will encode rich status with control codes and output the
+    encoded string to stdout. Client-side decode control codes from server
+    output and update the rich status. This function is safe to be called in
+    async/multi-threaded context.
+
+    See also: :func:`client_status`, :class:`EncodedStatus`.
+    """
     from sky import sky_logging  # pylint: disable=import-outside-toplevel
-    if (annotations.is_on_api_server and
-            threading.current_thread() is threading.main_thread() and
+    if (annotations.is_on_api_server and _is_thread_safe() and
             not sky_logging.is_silent()):
-        if _statuses['server'] is None:
-            _statuses['server'] = EncodedStatus(msg)
-        return _RevertibleStatus(msg, 'server')
+        if _get_server_status() is None:
+            _set_server_status(EncodedStatus(msg))
+        return _RevertibleStatus(msg, _get_server_status, _set_server_status)
     return _NoOpConsoleStatus()
 
 
@@ -196,29 +242,34 @@ def stop_safe_status():
     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()
+    if _is_thread_safe():
+        return
+    server_status = _get_server_status()
+    if server_status is not None:
+        server_status.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
-            _statuses['server'] is not None):
-        _statuses['server'].update(msg)
+    if not _is_thread_safe():
+        return
+    server_status = _get_server_status()
+    if server_status is not None:
+        server_status.update(msg)
 
 
 @contextlib.contextmanager
 def safe_logger():
     with _logging_lock:
-        client_status_obj = _statuses['client']
+        client_status_obj = _get_client_status()
 
         client_status_live = (client_status_obj is not None and
+                              hasattr(client_status_obj, '_live') and
                               client_status_obj._live.is_started)  # pylint: disable=protected-access
-        if client_status_live:
+        if client_status_live and client_status_obj is not None:
             client_status_obj.stop()
         yield
-        if client_status_live:
+        if client_status_live and client_status_obj is not None:
             client_status_obj.start()
 
 
@@ -230,13 +281,13 @@ class RichSafeStreamHandler(logging.StreamHandler):
 
 
 def client_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
-    """A wrapper for multi-threaded console.status."""
+    """A wrapper for multi-threaded client-side 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'] = rich_console_utils.get_console().status(msg)
-        return _RevertibleStatus(msg, 'client')
+        if _get_client_status() is None:
+            _set_client_status(rich_console_utils.get_console().status(msg))
+        return _RevertibleStatus(msg, _get_client_status, _set_client_status)
     return _NoOpConsoleStatus()
 
 
@@ -320,6 +371,9 @@ def decode_rich_status(
                     yield line
                     continue
 
+                if control == Control.RETRY:
+                    raise exceptions.RequestInterruptedError(
+                        'Streaming interrupted. Please retry.')
                 # control is not None, i.e. it is a rich status control message.
                 if threading.current_thread() is not threading.main_thread():
                     yield None
@@ -341,6 +395,130 @@ def decode_rich_status(
                         decoding_status.__exit__(None, None, None)
                     elif control == Control.START:
                         decoding_status.start()
+                    elif control == Control.HEARTBEAT:
+                        # Heartbeat is not displayed to the user, so we do not
+                        # need to update the status.
+                        pass
+    finally:
+        if decoding_status is not None:
+            decoding_status.__exit__(None, None, None)
+
+
+async def decode_rich_status_async(
+        response: 'aiohttp.ClientResponse'
+) -> typing.AsyncIterator[Optional[str]]:
+    """Async version of rich_utils.decode_rich_status that decodes rich status
+    messages from an aiohttp response.
+
+    Args:
+        response: The aiohttp response.
+
+    Yields:
+        Optional[str]: Decoded lines or None for control messages.
+    """
+    decoding_status = None
+    try:
+        last_line = ''
+        # Buffer to store incomplete UTF-8 bytes between chunks
+        undecoded_buffer = b''
+
+        # Iterate over the response content in chunks
+        async for chunk, _ in response.content.iter_chunks():
+            if chunk is None:
+                return
+
+            # Append the new chunk to any leftover bytes from previous iteration
+            current_bytes = undecoded_buffer + chunk
+            undecoded_buffer = b''
+
+            # Try to decode the combined bytes
+            try:
+                encoded_msg = current_bytes.decode('utf-8')
+            except UnicodeDecodeError as e:
+                # Check if this is potentially an incomplete sequence at the end
+                if e.start > 0:
+                    # Decode the valid part
+                    encoded_msg = current_bytes[:e.start].decode('utf-8')
+
+                    # Check if the remaining bytes are likely a partial char
+                    # or actually invalid UTF-8
+                    remaining_bytes = current_bytes[e.start:]
+                    if len(remaining_bytes) < 4:  # Max UTF-8 char is 4 bytes
+                        # Likely incomplete - save for next chunk
+                        undecoded_buffer = remaining_bytes
+                    else:
+                        # Likely invalid - replace with replacement character
+                        encoded_msg += remaining_bytes.decode('utf-8',
+                                                              errors='replace')
+                        undecoded_buffer = b''
+                else:
+                    # Error at the very beginning of the buffer - invalid UTF-8
+                    encoded_msg = current_bytes.decode('utf-8',
+                                                       errors='replace')
+                    undecoded_buffer = b''
+
+            lines = encoded_msg.splitlines(keepends=True)
+
+            # Skip processing if lines is empty to avoid IndexError
+            if not lines:
+                continue
+
+            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)
+                if line is None:
+                    continue
+                control = None
+                if is_payload:
+                    control, encoded_status = Control.decode(line)
+                if control is None:
+                    yield line
+                    continue
+
+                if control == Control.RETRY:
+                    raise exceptions.RequestInterruptedError(
+                        'Streaming interrupted. Please retry.')
+                # control is not None, i.e. it is a rich status control message.
+                # In async context, we'll handle rich status controls normally
+                # since async typically runs in main thread
+                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()
+                    elif control == Control.HEARTBEAT:
+                        # Heartbeat is not displayed to the user, so we do not
+                        # need to update the status.
+                        pass
     finally:
         if decoding_status is not None:
             decoding_status.__exit__(None, None, None)