skypilot-nightly 1.0.0.dev20250520__py3-none-any.whl → 1.0.0.dev20250522__py3-none-any.whl

sky/utils/cli_utils/status_utils.py

@@ -33,17 +33,15 @@ class StatusColumn:
     def __init__(self,
                  name: str,
                  calc_func: Callable,
-                 trunc_length: int = 0,
+                 truncate: bool = True,
                  show_by_default: bool = True):
         self.name = name
         self.calc_func = calc_func
-        self.trunc_length = trunc_length
+        self.truncate: bool = truncate
         self.show_by_default = show_by_default
 
     def calc(self, record):
-        val = self.calc_func(record)
-        if self.trunc_length != 0:
-            val = common_utils.truncate_long_string(str(val), self.trunc_length)
+        val = self.calc_func(record, self.truncate)
         return val
 
 
@@ -68,19 +66,20 @@ def show_status_table(cluster_records: List[_ClusterRecord],
             StatusColumn('USER_ID', _get_user_hash, show_by_default=False))
 
     status_columns += [
-        StatusColumn('LAUNCHED', _get_launched),
-        StatusColumn('RESOURCES',
-                     _get_resources,
-                     trunc_length=70 if not show_all else 0),
-        StatusColumn('REGION', _get_region, show_by_default=False),
-        StatusColumn('ZONE', _get_zone, show_by_default=False),
+        StatusColumn('INFRA', _get_infra, truncate=not show_all),
+        StatusColumn('RESOURCES', _get_resources, truncate=not show_all),
         StatusColumn('STATUS', _get_status_colored),
         StatusColumn('AUTOSTOP', _get_autostop),
-        StatusColumn('HEAD_IP', _get_head_ip, show_by_default=False),
-        StatusColumn('COMMAND',
-                     _get_command,
-                     trunc_length=COMMAND_TRUNC_LENGTH if not show_all else 0),
+        StatusColumn('LAUNCHED', _get_launched),
     ]
+    if show_all:
+        status_columns += [
+            StatusColumn('HEAD_IP', _get_head_ip, show_by_default=False),
+            StatusColumn('COMMAND',
+                         _get_command,
+                         truncate=not show_all,
+                         show_by_default=False),
+        ]
 
     columns = []
     for status_column in status_columns:
@@ -160,10 +159,10 @@ def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
     status_columns = [
         StatusColumn('NAME', _get_name),
         StatusColumn('LAUNCHED', _get_launched),
-        StatusColumn('DURATION', _get_duration, trunc_length=20),
+        StatusColumn('DURATION', _get_duration, truncate=False),
         StatusColumn('RESOURCES',
                      _get_resources_for_cost_report,
-                     trunc_length=70 if not show_all else 0),
+                     truncate=False),
         StatusColumn('STATUS',
                      _get_status_for_cost_report,
                      show_by_default=True),
@@ -221,47 +220,68 @@ def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
 # Some of these lambdas are invoked on both _ClusterRecord and
 # _ClusterCostReportRecord, which is okay as we guarantee the queried fields
 # exist in those cases.
-_get_name = (lambda cluster_record: cluster_record['name'])
-_get_user_hash = (lambda cluster_record: cluster_record['user_hash'])
-_get_user_name = (lambda cluster_record: cluster_record.get('user_name', '-'))
-_get_launched = (lambda cluster_record: log_utils.readable_time_duration(
+_get_name = (lambda cluster_record, _: cluster_record['name'])
+_get_user_hash = (lambda cluster_record, _: cluster_record['user_hash'])
+_get_user_name = (
+    lambda cluster_record, _: cluster_record.get('user_name', '-'))
+_get_launched = (lambda cluster_record, _: log_utils.readable_time_duration(
     cluster_record['launched_at']))
-_get_region = (
-    lambda clusters_status: clusters_status['handle'].launched_resources.region)
-_get_command = (lambda cluster_record: cluster_record['last_use'])
-_get_duration = (lambda cluster_record: log_utils.readable_time_duration(
+_get_duration = (lambda cluster_record, _: log_utils.readable_time_duration(
     0, cluster_record['duration'], absolute=True))
 
 
-def _get_status(cluster_record: _ClusterRecord) -> status_lib.ClusterStatus:
-    return cluster_record['status']
-
+def _get_command(cluster_record: _ClusterRecord, truncate: bool = True) -> str:
+    command = cluster_record.get('last_use', '-')
+    if truncate:
+        return common_utils.truncate_long_string(command, COMMAND_TRUNC_LENGTH)
+    return command
 
-def _get_status_colored(cluster_record: _ClusterRecord) -> str:
-    return _get_status(cluster_record).colored_str()
 
+def _get_status(cluster_record: _ClusterRecord,
+                truncate: bool = True) -> status_lib.ClusterStatus:
+    del truncate
+    return cluster_record['status']
 
-def _get_resources(cluster_record: _ClusterRecord) -> str:
-    if 'resources_str' in cluster_record:
-        return cluster_record['resources_str']
-    handle = cluster_record['handle']
-    if isinstance(handle, backends.LocalDockerResourceHandle):
-        resources_str = 'docker'
-    elif isinstance(handle, backends.CloudVmRayResourceHandle):
-        resources_str = resources_utils.get_readable_resources_repr(handle)
-    else:
-        raise ValueError(f'Unknown handle type {type(handle)} encountered.')
-    return resources_str
 
+def _get_status_colored(cluster_record: _ClusterRecord,
+                        truncate: bool = True) -> str:
+    del truncate
+    return _get_status(cluster_record).colored_str()
 
-def _get_zone(cluster_record: _ClusterRecord) -> str:
-    zone_str = cluster_record['handle'].launched_resources.zone
-    if zone_str is None:
-        zone_str = '-'
-    return zone_str
 
+def _get_resources(cluster_record: _ClusterRecord,
+                   truncate: bool = True) -> str:
+    """Get the resources information for a cluster.
 
-def _get_autostop(cluster_record: _ClusterRecord) -> str:
+    Returns:
+        A string in one of the following formats:
+        - For cloud VMs: "Nx instance_type" (e.g., "1x m6i.2xlarge")
+        - For K8S/SSH: "Nx (...)"
+        - "-" if no resource information is available
+    """
+    handle = cluster_record['handle']
+    if isinstance(handle, backends.CloudVmRayResourceHandle):
+        launched_resources = handle.launched_resources
+        if launched_resources is None:
+            return '-'
+
+        # For cloud VMs, show instance type directly
+        # For K8S/SSH, show (...) as the resource type
+        resources_str = cluster_record.get('resources_str', None)
+        if not truncate:
+            resources_str_full = cluster_record.get('resources_str_full', None)
+            if resources_str_full is not None:
+                resources_str = resources_str_full
+        if resources_str is None:
+            resources_str = resources_utils.get_readable_resources_repr(
+                handle, simplify=truncate)
+
+        return resources_str
+    return '-'
+
+
+def _get_autostop(cluster_record: _ClusterRecord, truncate: bool = True) -> str:
+    del truncate
     autostop_str = ''
     separation = ''
     if cluster_record['autostop'] >= 0:
@@ -276,7 +296,8 @@ def _get_autostop(cluster_record: _ClusterRecord) -> str:
     return autostop_str
 
 
-def _get_head_ip(cluster_record: _ClusterRecord) -> str:
+def _get_head_ip(cluster_record: _ClusterRecord, truncate: bool = True) -> str:
+    del truncate  # Unused
     handle = cluster_record['handle']
     if not isinstance(handle, backends.CloudVmRayResourceHandle):
         return '-'
@@ -291,6 +312,25 @@ def _is_pending_autostop(cluster_record: _ClusterRecord) -> bool:
         cluster_record) != status_lib.ClusterStatus.STOPPED
 
 
+def _get_infra(cluster_record: _ClusterRecord, truncate: bool = True) -> str:
+    """Get the infrastructure information for a cluster.
+
+    Returns:
+        A string in one of the following formats:
+        - AWS/region (e.g., "AWS/us-east-1")
+        - K8S/context (e.g., "K8S/my-ctx")
+        - SSH/hostname (e.g., "SSH/my-tobi-box")
+        - "-" if no infrastructure information is available
+    """
+    handle = cluster_record['handle']
+    if isinstance(handle, backends.CloudVmRayResourceHandle):
+        if handle.launched_resources is None:
+            # If launched_resources is None, try to get infra from the record
+            return cluster_record.get('infra', '-')
+        return handle.launched_resources.infra.formatted_str(truncate)
+    return '-'
+
+
 # ---- 'sky cost-report' helper functions below ----
 
 
@@ -347,14 +387,13 @@ def show_kubernetes_cluster_status_table(
         show_all: bool) -> None:
     """Compute cluster table values and display for Kubernetes clusters."""
     status_columns = [
-        StatusColumn('USER', lambda c: c.user),
-        StatusColumn('NAME', lambda c: c.cluster_name),
-        StatusColumn('LAUNCHED',
-                     lambda c: log_utils.readable_time_duration(c.launched_at)),
-        StatusColumn('RESOURCES',
-                     lambda c: c.resources_str,
-                     trunc_length=70 if not show_all else 0),
-        StatusColumn('STATUS', lambda c: c.status.colored_str()),
+        StatusColumn('USER', lambda c, _: c.user),
+        StatusColumn('NAME', lambda c, _: c.cluster_name),
+        StatusColumn('RESOURCES', lambda c, _: c.resources_str, truncate=False),
+        StatusColumn('STATUS', lambda c, _: c.status.colored_str()),
+        StatusColumn(
+            'LAUNCHED',
+            lambda c, _: log_utils.readable_time_duration(c.launched_at)),
         # TODO(romilb): We should consider adding POD_NAME field here when --all
         #  is passed to help users fetch pod name programmatically.
     ]

sky/utils/common_utils.py

@@ -723,10 +723,43 @@ def deprecated_function(
     return new_func
 
 
-def truncate_long_string(s: str, max_length: int = 35) -> str:
-    """Truncate a string to a maximum length, preserving whole words."""
+def truncate_long_string(s: str,
+                         max_length: int = 35,
+                         truncate_middle: bool = False) -> str:
+    """Truncate a string to a maximum length.
+
+    Args:
+        s: String to truncate.
+        max_length: Maximum length of the truncated string.
+        truncate_middle: Whether to truncate in the middle of the string.
+            If True, the middle part of the string is replaced with '...'.
+            If False, truncation happens at the end preserving whole words.
+
+    Returns:
+        Truncated string.
+    """
     if len(s) <= max_length:
         return s
+
+    if truncate_middle:
+        # Reserve 3 characters for '...'
+        if max_length <= 3:
+            return '...'
+
+        # Calculate how many characters to keep from beginning and end
+        half_length = (max_length - 3) // 2
+        remainder = (max_length - 3) % 2
+
+        # Keep one more character at the beginning if max_length - 3 is odd
+        start_length = half_length + remainder
+        end_length = half_length
+
+        # When end_length is 0, just show the start part and '...'
+        if end_length == 0:
+            return s[:start_length] + '...'
+        return s[:start_length] + '...' + s[-end_length:]
+
+    # Original end-truncation logic
     splits = s.split(' ')
     if len(splits[0]) > max_length:
         return splits[0][:max_length] + '...'  # Use '…'?

sky/utils/context.py

@@ -57,6 +57,7 @@ class Context(object):
         self._log_file = None
         self._log_file_handle = None
         self.env_overrides = {}
+        self.config_context = None
 
     def cancel(self):
         """Cancel the context."""
@@ -159,17 +160,25 @@ class ContextualEnviron(MutableMapping):
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
-                return ctx.env_overrides[key]
+                value = ctx.env_overrides[key]
+                # None is used to indicate that the key is deleted in the
+                # context.
+                if value is None:
+                    raise KeyError(key)
+                return value
         return self._environ[key]
 
     def __iter__(self):
         ctx = get()
+        deleted_keys = set()
         if ctx is not None:
-            for key in ctx.env_overrides:
+            for key, value in ctx.env_overrides.items():
+                if value is None:
+                    deleted_keys.add(key)
                 yield key
             for key in self._environ:
                 # Deduplicate the keys
-                if key not in ctx.env_overrides:
+                if key not in ctx.env_overrides and key not in deleted_keys:
                     yield key
         else:
             return self._environ.__iter__()
@@ -178,10 +187,27 @@ class ContextualEnviron(MutableMapping):
         return len(dict(self))
 
     def __setitem__(self, key, value):
-        return self._environ.__setitem__(key, value)
+        ctx = get()
+        if ctx is not None:
+            ctx.env_overrides[key] = value
+        else:
+            self._environ.__setitem__(key, value)
 
     def __delitem__(self, key):
-        return self._environ.__delitem__(key)
+        ctx = get()
+        if ctx is not None:
+            if key in ctx.env_overrides:
+                del ctx.env_overrides[key]
+            elif key in self._environ:
+                # If the key is not set in the context but set in the environ
+                # of the process, we mark it as deleted in the context by
+                # setting the value to None.
+                ctx.env_overrides[key] = None
+            else:
+                # The key is not set in the context nor the process.
+                raise KeyError(key)
+        else:
+            self._environ.__delitem__(key)
 
     def __repr__(self):
         return self._environ.__repr__()
@@ -190,7 +216,11 @@ class ContextualEnviron(MutableMapping):
         copied = self._environ.copy()
         ctx = get()
         if ctx is not None:
-            copied.update(ctx.env_overrides)
+            for key in ctx.env_overrides:
+                if ctx.env_overrides[key] is None:
+                    copied.pop(key)
+                else:
+                    copied[key] = ctx.env_overrides[key]
         return copied
 
     def setdefault(self, key, default=None):

sky/utils/context_utils.py

@@ -1,5 +1,6 @@
 """Utilities for SkyPilot context."""
 import asyncio
+import contextvars
 import functools
 import io
 import multiprocessing
@@ -170,3 +171,17 @@ def cancellation_guard(func: F) -> F:
         return func(*args, **kwargs)
 
     return typing.cast(F, wrapper)
+
+
+# TODO(aylei): replace this with asyncio.to_thread once we drop support for
+# python 3.8
+def to_thread(func, /, *args, **kwargs):
+    """Asynchronously run function *func* in a separate thread.
+
+    This is same as asyncio.to_thread added in python 3.9
+    """
+    loop = asyncio.get_running_loop()
+    # This is critical to pass the current coroutine context to the new thread
+    pyctx = contextvars.copy_context()
+    func_call = functools.partial(pyctx.run, func, *args, **kwargs)
+    return loop.run_in_executor(None, func_call)

sky/utils/infra_utils.py

@@ -0,0 +1,175 @@
+"""Utility functions for handling infrastructure specifications."""
+import dataclasses
+from typing import Optional
+
+from sky.utils import common_utils
+from sky.utils import ux_utils
+
+_REGION_OR_ZONE_TRUNCATION_LENGTH = 25
+
+
+@dataclasses.dataclass
+class InfraInfo:
+    """Infrastructure information parsed from infra string.
+
+    When a field is None, it means the field is not specified.
+    """
+    cloud: Optional[str] = None
+    region: Optional[str] = None
+    zone: Optional[str] = None
+
+    def __init__(self,
+                 cloud: Optional[str] = None,
+                 region: Optional[str] = None,
+                 zone: Optional[str] = None):
+        assert cloud not in ['none', 'None', 'NONE'], 'cloud must be specified'
+        if not cloud or cloud == '*':
+            cloud = None
+        if not region or region == '*':
+            region = None
+        if not zone or zone == '*':
+            zone = None
+
+        self.cloud = cloud
+        self.region = region
+        self.zone = zone
+
+    @staticmethod
+    def from_str(infra: Optional[str]) -> 'InfraInfo':
+        """Parse the infra string into cloud, region, and zone components.
+
+        The format of the infra string is `cloud`, `cloud/region`, or
+        `cloud/region/zone`. Examples: `aws`, `aws/us-east-1`,
+        `aws/us-east-1/us-east-1a`. For any field, you can use `*` to indicate
+        that any value is acceptable.
+
+        If `*` is used for any field, the InfraInfo will have None for that
+        field.
+
+        Args:
+            infra: A string in the format of `cloud`, `cloud/region`, or
+                `cloud/region/zone`. Examples: `aws`, `aws/us-east-1`,
+                `aws/us-east-1/us-east-1a`.
+
+        Returns:
+            An InfraInfo object containing cloud, region, and zone information.
+
+        Raises:
+            ValueError: If the infra string is malformed.
+        """
+        if infra is None or not infra.strip():
+            return InfraInfo()
+
+        infra = infra.strip().strip('/')
+
+        # Split on / to get cloud, region, zone
+        parts = [p.strip() for p in infra.strip().split('/')]
+
+        if '' in parts:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Invalid infra format: {infra}. Format should not contain '
+                    'empty parts (e.g., double slashes "//").')
+
+        if not parts or not parts[0]:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Invalid infra format: {infra}. Expected format is '
+                    '"cloud", "cloud/region", or "cloud/region/zone".')
+
+        cloud_name: Optional[str] = parts[0].lower()
+
+        # Handle Kubernetes contexts specially, as they can contain slashes
+        if cloud_name in ['k8s', 'kubernetes']:
+            # For Kubernetes, the entire string after "k8s/" is the
+            # context name (region)
+            cloud_name = 'kubernetes'  # Normalize k8s to kubernetes
+            region = '/'.join(parts[1:]) if len(parts) >= 2 else None
+            zone = None
+        else:
+            # For non-Kubernetes clouds, continue with regular parsing
+            # but be careful to only split into max 3 parts
+            region_zone_parts = parts[1:]
+            region = None
+            zone = None
+            if region_zone_parts:
+                region = region_zone_parts[0]
+                if len(region_zone_parts) > 1:
+                    zone = region_zone_parts[1]
+                if len(region_zone_parts) > 2:
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(
+                            f'Invalid infra format: {infra}. Expected format '
+                            'is "cloud", "cloud/region", or '
+                            '"cloud/region/zone".')
+
+        if cloud_name == '*':
+            cloud_name = None
+        if region == '*':
+            region = None
+        if zone == '*':
+            zone = None
+        return InfraInfo(cloud=cloud_name, region=region, zone=zone)
+
+    def to_str(self) -> Optional[str]:
+        """Formats cloud, region, and zone into an infra string.
+
+        Args:
+            cloud: The cloud object
+            region: The region name
+            zone: The zone name
+
+        Returns:
+            A formatted infra string, or None if cloud is None or '*'
+        """
+        cloud = self.cloud
+        region = self.region
+        zone = self.zone
+
+        if cloud is None:
+            cloud = '*'
+        if region is None:
+            region = '*'
+        if zone is None:
+            zone = '*'
+
+        # Build the parts list and filter out trailing wildcards
+        parts = [cloud.lower(), region, zone]
+        while parts and parts[-1] == '*':
+            parts.pop()
+
+        if not parts:
+            return None
+
+        # Join the parts with '/'
+        return '/'.join(parts)
+
+    def formatted_str(self, truncate: bool = True) -> str:
+        """Formats cloud, region, and zone into an infra string.
+
+        Args:
+            truncate: Whether to truncate the region or zone
+
+        Returns:
+            A formatted infra string, or None if cloud is None or '*'
+        """
+        if self.cloud is None or self.cloud == '*':
+            return '-'
+
+        region_or_zone = None
+        if self.zone is not None and self.zone != '*':
+            region_or_zone = self.zone
+        elif self.region is not None and self.region != '*':
+            region_or_zone = self.region
+
+        if region_or_zone is not None and truncate:
+            region_or_zone = common_utils.truncate_long_string(
+                region_or_zone,
+                _REGION_OR_ZONE_TRUNCATION_LENGTH,
+                truncate_middle=True)
+
+        formatted_str = f'{self.cloud}'
+        if region_or_zone is not None:
+            formatted_str += f' ({region_or_zone})'
+
+        return formatted_str

sky/utils/resources_utils.py

@@ -4,11 +4,11 @@ import enum
 import itertools
 import json
 import math
-import re
 import typing
 from typing import Dict, List, Optional, Set, Union
 
 from sky import skypilot_config
+from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import ux_utils
 
@@ -50,6 +50,20 @@ class DiskTier(enum.Enum):
         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
@@ -139,34 +153,54 @@ def simplify_ports(ports: List[str]) -> List[str]:
 
 def format_resource(resource: 'resources_lib.Resources',
                     simplify: bool = False) -> str:
+    resource = resource.assert_launchable()
+    vcpu, mem = resource.cloud.get_vcpus_mem_from_instance_type(
+        resource.instance_type)
+
+    components = []
+
+    if resource.accelerators is not None:
+        acc, count = list(resource.accelerators.items())[0]
+        components.append(f'gpus={acc}:{count}')
+
+    is_k8s = str(resource.cloud).lower() == 'kubernetes'
+    if (resource.accelerators is None or is_k8s or not simplify):
+        if vcpu is not None:
+            components.append(f'cpus={int(vcpu)}')
+        if mem is not None:
+            components.append(f'mem={int(mem)}')
+
+    instance_type = resource.instance_type
     if simplify:
-        resource = resource.assert_launchable()
-        cloud = resource.cloud
-        if resource.accelerators is None:
-            vcpu, _ = cloud.get_vcpus_mem_from_instance_type(
-                resource.instance_type)
-            assert vcpu is not None, 'vCPU must be specified'
-            hardware = f'vCPU={int(vcpu)}'
-        else:
-            hardware = f'{resource.accelerators}'
-        spot = '[Spot]' if resource.use_spot else ''
-        return f'{cloud}({spot}{hardware})'
+        instance_type = common_utils.truncate_long_string(instance_type, 15)
+    if not is_k8s:
+        components.append(instance_type)
+    if simplify:
+        components.append('...')
     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
+        image_id = resource.image_id
+        if image_id is not None:
+            if None in image_id:
+                components.append(f'image_id={image_id[None]}')
+            else:
+                components.append(f'image_id={image_id}')
+        components.append(f'disk={resource.disk_size}')
+        disk_tier = resource.disk_tier
+        if disk_tier is not None:
+            components.append(f'disk_tier={disk_tier.value}')
+        ports = resource.ports
+        if ports is not None:
+            components.append(f'ports={ports}')
+
+    spot = '[spot]' if resource.use_spot else ''
+    return f'{spot}({"" if not components else ", ".join(components)})'
 
 
 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):
-        return (f'{handle.launched_nodes}x '
+        return (f'{handle.launched_nodes}x'
                 f'{format_resource(handle.launched_resources, simplify)}')
     return _DEFAULT_MESSAGE_HANDLE_INITIALIZING