skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/utils/resources_utils.py

@@ -181,57 +181,81 @@ def simplify_ports(ports: List[str]) -> List[str]:
 
 
 def format_resource(resource: 'resources_lib.Resources',
-                    simplify: bool = False) -> str:
+                    simplified_only: bool = False) -> Tuple[str, Optional[str]]:
     resource = resource.assert_launchable()
-    vcpu, mem = resource.cloud.get_vcpus_mem_from_instance_type(
-        resource.instance_type)
+    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)
 
-    components = []
+    elements_simple = []
+    elements_full = []
 
     if resource.accelerators is not None:
         acc, count = list(resource.accelerators.items())[0]
-        components.append(f'gpus={acc}:{count}')
+        elements_simple.append(f'gpus={acc}:{count}')
+        elements_full.append(f'gpus={acc}:{count}')
 
-    is_k8s = str(resource.cloud).lower() == 'kubernetes'
-    if (resource.accelerators is None or is_k8s or not simplify):
+    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:
-            components.append(f'cpus={int(vcpu)}')
+            elements_full.append(f'cpus={int(vcpu)}')
         if mem is not None:
-            components.append(f'mem={int(mem)}')
+            elements_full.append(f'mem={int(mem)}')
 
-    instance_type = resource.instance_type
-    if simplify:
-        instance_type = common_utils.truncate_long_string(instance_type, 15)
     if not is_k8s:
-        components.append(instance_type)
-    if simplify:
-        components.append('...')
-    else:
+        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:
-                components.append(f'image_id={image_id[None]}')
+                elements_full.append(f'image_id={image_id[None]}')
             else:
-                components.append(f'image_id={image_id}')
-        components.append(f'disk={resource.disk_size}')
+                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:
-            components.append(f'disk_tier={disk_tier.value}')
+            elements_full.append(f'disk_tier={disk_tier.value}')
         ports = resource.ports
         if ports is not None:
-            components.append(f'ports={ports}')
+            elements_full.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:
+    resources_str_simple = (
+        f'{spot}({"" if not elements_simple else ", ".join(elements_simple)})')
+    if simplified_only:
+        return resources_str_simple, None
+    else:
+        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(

sky/utils/rich_utils.py

@@ -193,7 +193,8 @@ class _RevertibleStatus:
                     self.get_status_fn().__exit__(exc_type, exc_val, exc_tb)
                     self.set_status_fn(None)
             else:
-                self.get_status_fn().update(self.previous_message)
+                if self.previous_message is not None:
+                    self.get_status_fn().update(self.previous_message)
 
     def update(self, *args, **kwargs):
         self.get_status_fn().update(*args, **kwargs)
@@ -263,11 +264,12 @@ def safe_logger():
         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()
 
 
@@ -421,7 +423,7 @@ async def decode_rich_status_async(
         undecoded_buffer = b''
 
         # Iterate over the response content in chunks
-        async for chunk in response.content.iter_chunked(8192):
+        async for chunk, _ in response.content.iter_chunks():
             if chunk is None:
                 return
 
@@ -481,6 +483,8 @@ async def decode_rich_status_async(
                     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)

sky/utils/schemas.py

@@ -460,8 +460,8 @@ def get_volume_schema():
                 'type': 'string',
                 'pattern': constants.MEMORY_SIZE_PATTERN,
             },
-            'resource_name': {
-                'type': 'string',
+            'use_existing': {
+                'type': 'boolean',
             },
             'config': {
                 'type': 'object',
@@ -574,6 +574,9 @@ def get_volume_mount_schema():
             'volume_name': {
                 'type': 'string',
             },
+            'is_ephemeral': {
+                'type': 'boolean',
+            },
             'volume_config': {
                 'type': 'object',
                 'required': [],
@@ -791,23 +794,6 @@ def _filter_schema(schema: dict, keys_to_keep: List[Tuple[str, ...]]) -> dict:
     return new_schema
 
 
-def _experimental_task_schema() -> dict:
-    # TODO: experimental.config_overrides has been deprecated in favor of the
-    # top-level `config` field. Remove in v0.11.0.
-    config_override_schema = _filter_schema(
-        get_config_schema(), constants.OVERRIDEABLE_CONFIG_KEYS_IN_TASK)
-    return {
-        'experimental': {
-            'type': 'object',
-            'required': [],
-            'additionalProperties': False,
-            'properties': {
-                'config_overrides': config_override_schema,
-            }
-        }
-    }
-
-
 def get_task_schema():
     return {
         '$schema': 'https://json-schema.org/draft/2020-12/schema',
@@ -918,7 +904,6 @@ def get_task_schema():
             '_metadata': {
                 'type': 'object',
             },
-            **_experimental_task_schema(),
         }
     }
 
@@ -1043,11 +1028,21 @@ class RemoteIdentityOptions(enum.Enum):
 
 def get_default_remote_identity(cloud: str) -> str:
     """Get the default remote identity for the specified cloud."""
-    if cloud == 'kubernetes':
+    if cloud in ('kubernetes', 'ssh'):
         return RemoteIdentityOptions.SERVICE_ACCOUNT.value
     return RemoteIdentityOptions.LOCAL_CREDENTIALS.value
 
 
+_CAPABILITIES_SCHEMA = {
+    'capabilities': {
+        'type': 'array',
+        'items': {
+            'type': 'string',
+            'case_insensitive_enum': ['compute', 'storage']
+        },
+    }
+}
+
 _REMOTE_IDENTITY_SCHEMA = {
     'remote_identity': {
         'type': 'string',
@@ -1070,25 +1065,16 @@ _REMOTE_IDENTITY_SCHEMA_KUBERNETES = {
     },
 }
 
-_CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
-    'networking': {
-        'type': 'string',
-        'case_insensitive_enum': [
-            type.value for type in kubernetes_enums.KubernetesNetworkingMode
-        ],
-    },
-    'ports': {
-        'type': 'string',
-        'case_insensitive_enum': [
-            type.value for type in kubernetes_enums.KubernetesPortMode
-        ],
-    },
+_CONTEXT_CONFIG_SCHEMA_MINIMAL = {
     'pod_config': {
         'type': 'object',
         'required': [],
         # Allow arbitrary keys since validating pod spec is hard
         'additionalProperties': True,
     },
+    'provision_timeout': {
+        'type': 'integer',
+    },
     'custom_metadata': {
         'type': 'object',
         'required': [],
@@ -1103,9 +1089,23 @@ _CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
             }]
         },
     },
-    'provision_timeout': {
-        'type': 'integer',
+}
+
+_CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
+    # TODO(kevin): Remove 'networking' in v0.13.0.
+    'networking': {
+        'type': 'string',
+        'case_insensitive_enum': [
+            type.value for type in kubernetes_enums.KubernetesNetworkingMode
+        ],
+    },
+    'ports': {
+        'type': 'string',
+        'case_insensitive_enum': [
+            type.value for type in kubernetes_enums.KubernetesPortMode
+        ],
     },
+    **_CONTEXT_CONFIG_SCHEMA_MINIMAL,
     'autoscaler': {
         'type': 'string',
         'case_insensitive_enum': [
@@ -1153,6 +1153,12 @@ _CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
     },
     'remote_identity': {
         'type': 'string',
+    },
+    'post_provision_runcmd': {
+        'type': 'array',
+        'items': {
+            'type': 'string'
+        },
     }
 }
 
@@ -1189,7 +1195,13 @@ def get_config_schema():
                         'consolidation_mode': {
                             'type': 'boolean',
                             'default': False,
-                        }
+                        },
+                        'controller_logs_gc_retention_hours': {
+                            'type': 'integer',
+                        },
+                        'task_logs_gc_retention_hours': {
+                            'type': 'integer',
+                        },
                     },
                 },
                 'bucket': {
@@ -1234,6 +1246,9 @@ def get_config_schema():
                         'type': 'null',
                     }],
                 },
+                'use_ssm': {
+                    'type': 'boolean',
+                },
                 'post_provision_runcmd': {
                     'type': 'array',
                     'items': {
@@ -1247,6 +1262,7 @@ def get_config_schema():
                         }]
                     },
                 },
+                **_CAPABILITIES_SCHEMA,
                 **_LABELS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,
             },
@@ -1304,6 +1320,7 @@ def get_config_schema():
                         }
                     ],
                 },
+                **_CAPABILITIES_SCHEMA,
                 **_LABELS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,
             },
@@ -1328,10 +1345,15 @@ def get_config_schema():
             'additionalProperties': False,
             'properties': {
                 'allowed_contexts': {
-                    'type': 'array',
-                    'items': {
+                    'oneOf': [{
+                        'type': 'array',
+                        'items': {
+                            'type': 'string',
+                        },
+                    }, {
                         'type': 'string',
-                    },
+                        'pattern': '^all$'
+                    }]
                 },
                 'context_configs': {
                     'type': 'object',
@@ -1361,11 +1383,42 @@ def get_config_schema():
                         'type': 'string',
                     },
                 },
-                'pod_config': {
+                'context_configs': {
                     'type': 'object',
                     'required': [],
-                    # Allow arbitrary keys since validating pod spec is hard
-                    'additionalProperties': True,
+                    'properties': {},
+                    # Properties are ssh cluster names, which are the
+                    # kubernetes context names without `ssh-` prefix.
+                    'additionalProperties': {
+                        'type': 'object',
+                        'required': [],
+                        'additionalProperties': False,
+                        'properties': {
+                            **_CONTEXT_CONFIG_SCHEMA_MINIMAL,
+                        },
+                    },
+                },
+                **_CONTEXT_CONFIG_SCHEMA_MINIMAL,
+            }
+        },
+        'slurm': {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'allowed_clusters': {
+                    'oneOf': [{
+                        'type': 'array',
+                        'items': {
+                            'type': 'string',
+                        },
+                    }, {
+                        'type': 'string',
+                        'pattern': '^all$'
+                    }]
+                },
+                'provision_timeout': {
+                    'type': 'integer',
                 },
             }
         },
@@ -1403,11 +1456,24 @@ def get_config_schema():
                 }
             },
         },
+        'vast': {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'secure_only': {
+                    'type': 'boolean',
+                },
+            }
+        },
         'nebius': {
             'type': 'object',
             'required': [],
             'properties': {
-                **_NETWORK_CONFIG_SCHEMA, 'tenant_id': {
+                **_NETWORK_CONFIG_SCHEMA, 'use_static_ip_address': {
+                    'type': 'boolean',
+                },
+                'tenant_id': {
                     'type': 'string',
                 },
                 'domain': {
@@ -1520,7 +1586,7 @@ def get_config_schema():
         }
     }
 
-    daemon_schema = {
+    daemon_schema: Dict[str, Any] = {
         'type': 'object',
         'required': [],
         'additionalProperties': False,
@@ -1580,10 +1646,10 @@ def get_config_schema():
 
     allowed_workspace_cloud_names = list(constants.ALL_CLOUDS) + ['cloudflare']
     # Create pattern for not supported clouds, i.e.
-    # all clouds except gcp, kubernetes, ssh
+    # all clouds except aws, gcp, kubernetes, ssh, nebius
     not_supported_clouds = [
         cloud for cloud in allowed_workspace_cloud_names
-        if cloud.lower() not in ['gcp', 'kubernetes', 'ssh', 'nebius']
+        if cloud.lower() not in ['aws', 'gcp', 'kubernetes', 'ssh', 'nebius']
     ]
     not_supported_cloud_regex = '|'.join(not_supported_clouds)
     workspaces_schema = {
@@ -1594,7 +1660,8 @@ def get_config_schema():
             'type': 'object',
             'additionalProperties': False,
             'patternProperties': {
-                # Pattern for non-GCP clouds - only allows 'disabled' property
+                # Pattern for clouds with no workspace-specific config -
+                # only allow 'disabled' property.
                 f'^({not_supported_cloud_regex})$': {
                     'type': 'object',
                     'additionalProperties': False,
@@ -1625,7 +1692,21 @@ def get_config_schema():
                         },
                         'disabled': {
                             'type': 'boolean'
-                        }
+                        },
+                        **_CAPABILITIES_SCHEMA,
+                    },
+                    'additionalProperties': False,
+                },
+                'aws': {
+                    'type': 'object',
+                    'properties': {
+                        'profile': {
+                            'type': 'string'
+                        },
+                        'disabled': {
+                            'type': 'boolean'
+                        },
+                        **_CAPABILITIES_SCHEMA,
                     },
                     'additionalProperties': False,
                 },
@@ -1650,10 +1731,15 @@ def get_config_schema():
                     'required': [],
                     'properties': {
                         'allowed_contexts': {
-                            'type': 'array',
-                            'items': {
+                            'oneOf': [{
+                                'type': 'array',
+                                'items': {
+                                    'type': 'string',
+                                },
+                            }, {
                                 'type': 'string',
-                            },
+                                'pattern': '^all$'
+                            }]
                         },
                         'disabled': {
                             'type': 'boolean'

sky/utils/subprocess_utils.py

@@ -10,7 +10,8 @@ import sys
 import threading
 import time
 import typing
-from typing import Any, Callable, Dict, List, Optional, Protocol, Tuple, Union
+from typing import (Any, Callable, Dict, List, Optional, Protocol, Set, Tuple,
+                    Union)
 
 import colorama
 
@@ -18,6 +19,7 @@ from sky import exceptions
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
 from sky.skylet import log_lib
+from sky.skylet import subprocess_daemon
 from sky.utils import common_utils
 from sky.utils import timeline
 from sky.utils import ux_utils
@@ -107,7 +109,7 @@ def get_parallel_threads(cloud_str: Optional[str] = None) -> int:
 
 
 def run_in_parallel(func: Callable,
-                    args: List[Any],
+                    args: Union[List[Any], Set[Any]],
                     num_threads: Optional[int] = None) -> List[Any]:
     """Run a function in parallel on a list of arguments.
 
@@ -128,7 +130,7 @@ def run_in_parallel(func: Callable,
     if len(args) == 0:
         return []
     if len(args) == 1:
-        return [func(args[0])]
+        return [func(list(args)[0])]
 
     processes = (num_threads
                  if num_threads is not None else get_parallel_threads())
@@ -305,11 +307,17 @@ def run_with_retries(
     return returncode, stdout, stderr
 
 
-def kill_process_daemon(process_pid: int) -> None:
+def kill_process_daemon(process_pid: int, use_kill_pg: bool = False) -> None:
     """Start a daemon as a safety net to kill the process.
 
     Args:
         process_pid: The PID of the process to kill.
+        use_kill_pg: Whether to use kill process group to kill the process. If
+            True, the process will use os.killpg() to kill the target process
+            group on UNIX system, which is more efficient than using the daemon
+            to refresh the process tree in the daemon. Note that both
+            implementations have corner cases where subprocesses might not be
+            killed. Refer to subprocess_daemon.py for more details.
     """
     # Get initial children list
     try:
@@ -336,6 +344,10 @@ def kill_process_daemon(process_pid: int) -> None:
         ','.join(map(str, initial_children)),
     ]
 
+    env = os.environ.copy()
+    if use_kill_pg:
+        env[subprocess_daemon.USE_KILL_PG_ENV_VAR] = '1'
+
     # We do not need to set `start_new_session=True` here, as the
     # daemon script will detach itself from the parent process with
     # fork to avoid being killed by parent process. See the reason we
@@ -347,6 +359,7 @@ def kill_process_daemon(process_pid: int) -> None:
         stderr=subprocess.DEVNULL,
         # Disable input
         stdin=subprocess.DEVNULL,
+        env=env,
     )
 
 

sky/utils/thread_utils.py

@@ -0,0 +1,91 @@
+"""Utility functions for threads."""
+
+import threading
+from typing import Any, Dict, Generic, Optional, overload, TypeVar
+
+from sky.utils import common_utils
+
+
+class SafeThread(threading.Thread):
+    """A thread that can catch exceptions."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._exc = None
+
+    def run(self):
+        try:
+            super().run()
+        except BaseException as e:  # pylint: disable=broad-except
+            self._exc = e
+
+    @property
+    def format_exc(self) -> Optional[str]:
+        if self._exc is None:
+            return None
+        return common_utils.format_exception(self._exc)
+
+
+# pylint: disable=invalid-name
+KeyType = TypeVar('KeyType')
+ValueType = TypeVar('ValueType')
+
+
+# Google style guide: Do not rely on the atomicity of built-in types.
+# Our launch and down process pool will be used by multiple threads,
+# therefore we need to use a thread-safe dict.
+# see https://google.github.io/styleguide/pyguide.html#218-threading
+class ThreadSafeDict(Generic[KeyType, ValueType]):
+    """A thread-safe dict."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._dict: Dict[KeyType, ValueType] = dict(*args, **kwargs)
+        self._lock = threading.Lock()
+
+    def __getitem__(self, key: KeyType) -> ValueType:
+        with self._lock:
+            return self._dict.__getitem__(key)
+
+    def __setitem__(self, key: KeyType, value: ValueType) -> None:
+        with self._lock:
+            return self._dict.__setitem__(key, value)
+
+    def __delitem__(self, key: KeyType) -> None:
+        with self._lock:
+            return self._dict.__delitem__(key)
+
+    def __len__(self) -> int:
+        with self._lock:
+            return self._dict.__len__()
+
+    def __contains__(self, key: KeyType) -> bool:
+        with self._lock:
+            return self._dict.__contains__(key)
+
+    def items(self):
+        with self._lock:
+            return self._dict.items()
+
+    def values(self):
+        with self._lock:
+            return self._dict.values()
+
+    @overload
+    def get(self, key: KeyType, default: ValueType) -> ValueType:
+        ...
+
+    @overload
+    def get(self,
+            key: KeyType,
+            default: Optional[ValueType] = None) -> Optional[ValueType]:
+        ...
+
+    def get(self,
+            key: KeyType,
+            default: Optional[ValueType] = None) -> Optional[ValueType]:
+        with self._lock:
+            return self._dict.get(key, default)
+
+    def pop(self, key: KeyType) -> Optional[ValueType]:
+        with self._lock:
+            return self._dict.pop(key, None)

sky/utils/timeline.py

@@ -58,7 +58,8 @@ class Event:
         })
         event_begin['args'] = {'stack': '\n'.join(traceback.format_stack())}
         if self._message is not None:
-            event_begin['args']['message'] = self._message
+            event_begin['args'][
+                'message'] = self._message  # type: ignore[index]
         _events.append(event_begin)
 
     def end(self):