skypilot-nightly 1.0.0.dev20250215__py3-none-any.whl → 1.0.0.dev20250217__py3-none-any.whl

sky/data/storage.py

@@ -18,7 +18,6 @@ from sky import exceptions
 from sky import global_user_state
 from sky import sky_logging
 from sky import skypilot_config
-from sky import status_lib
 from sky.adaptors import aws
 from sky.adaptors import azure
 from sky.adaptors import cloudflare
@@ -34,6 +33,7 @@ from sky.skylet import constants
 from sky.utils import common_utils
 from sky.utils import rich_utils
 from sky.utils import schemas
+from sky.utils import status_lib
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -203,9 +203,8 @@ class StoreType(enum.Enum):
     @classmethod
     def get_fields_from_store_url(
         cls, store_url: str
-    ) -> Tuple['StoreType', Type['AbstractStore'], str, str, Optional[str],
-               Optional[str]]:
-        """Returns the store type, store class, bucket name, and sub path from
+    ) -> Tuple['StoreType', str, str, Optional[str], Optional[str]]:
+        """Returns the store type, bucket name, and sub path from
         a store URL, and the storage account name and region if applicable.
 
         Args:
@@ -221,21 +220,16 @@ class StoreType(enum.Enum):
                 if store_type == StoreType.AZURE:
                     storage_account_name, bucket_name, sub_path = \
                         data_utils.split_az_path(store_url)
-                    store_cls: Type['AbstractStore'] = AzureBlobStore
                 elif store_type == StoreType.IBM:
                     bucket_name, sub_path, region = data_utils.split_cos_path(
                         store_url)
-                    store_cls = IBMCosStore
                 elif store_type == StoreType.R2:
                     bucket_name, sub_path = data_utils.split_r2_path(store_url)
-                    store_cls = R2Store
                 elif store_type == StoreType.GCS:
                     bucket_name, sub_path = data_utils.split_gcs_path(store_url)
-                    store_cls = GcsStore
                 elif store_type == StoreType.S3:
                     bucket_name, sub_path = data_utils.split_s3_path(store_url)
-                    store_cls = S3Store
-                return store_type, store_cls,bucket_name, \
+                return store_type, bucket_name, \
                     sub_path, storage_account_name, region
         raise ValueError(f'Unknown store URL: {store_url}')
 
@@ -546,7 +540,7 @@ class Storage(object):
         self,
         name: Optional[str] = None,
         source: Optional[SourceType] = None,
-        stores: Optional[Dict[StoreType, AbstractStore]] = None,
+        stores: Optional[List[StoreType]] = None,
         persistent: Optional[bool] = True,
         mode: StorageMode = StorageMode.MOUNT,
         sync_on_reconstruction: bool = True,
@@ -605,26 +599,47 @@ class Storage(object):
           _bucket_sub_path: Optional[str]; The subdirectory to use for the
             storage object.
         """
-        self.name: str
+        self.name = name
         self.source = source
         self.persistent = persistent
         self.mode = mode
         assert mode in StorageMode
+        self.stores: Dict[StoreType, Optional[AbstractStore]] = {}
+        if stores is not None:
+            for store in stores:
+                self.stores[store] = None
         self.sync_on_reconstruction = sync_on_reconstruction
         self._is_sky_managed = _is_sky_managed
         self._bucket_sub_path = _bucket_sub_path
 
+        self._constructed = False
         # TODO(romilb, zhwu): This is a workaround to support storage deletion
-        # for spot. Once sky storage supports forced management for external
-        # buckets, this can be deprecated.
+        # for managed jobs. Once sky storage supports forced management for
+        # external buckets, this can be deprecated.
         self.force_delete = False
 
-        # Validate and correct inputs if necessary
-        self._validate_storage_spec(name)
+    def construct(self):
+        """Constructs the storage object.
+
+        The Storage object is lazily initialized, so that when a user
+        initializes a Storage object on client side, it does not trigger the
+        actual storage creation on the client side.
 
-        # Sky optimizer either adds a storage object instance or selects
-        # from existing ones
-        self.stores = {} if stores is None else stores
+        Instead, once the specification of the storage object is uploaded to the
+        SkyPilot API server side, the server side should use this construct()
+        method to actually create the storage object. The construct() method
+        will:
+
+        1. Set the stores field if not specified
+        2. Create the bucket or check the existence of the bucket
+        3. Sync the data from the source to the bucket if necessary
+        """
+        if self._constructed:
+            return
+        self._constructed = True
+
+        # Validate and correct inputs if necessary
+        self._validate_storage_spec(self.name)
 
         # Logic to rebuild Storage if it is in global user state
         handle = global_user_state.get_handle_from_storage_name(self.name)
@@ -635,6 +650,28 @@ class Storage(object):
                          f'loading Storage: {self.name}')
             self._add_store_from_metadata(self.handle.sky_stores)
 
+            # When a storage object is reconstructed from global_user_state,
+            # the user may have specified a new store type in the yaml file that
+            # was not used with the storage object. We should error out in this
+            # case, as we don't support having multiple stores for the same
+            # storage object.
+            if any(s is None for s in self.stores.values()):
+                new_store_type = None
+                previous_store_type = None
+                for store_type, store in self.stores.items():
+                    if store is not None:
+                        previous_store_type = store_type
+                    else:
+                        new_store_type = store_type
+                with ux_utils.print_exception_no_traceback():
+                    raise exceptions.StorageBucketCreateError(
+                        f'Bucket {self.name} was previously created for '
+                        f'{previous_store_type.value.lower()!r}, but a new '
+                        f'store type {new_store_type.value.lower()!r} is '
+                        'requested. This is not supported yet. Please specify '
+                        'the same store type: '
+                        f'{previous_store_type.value.lower()!r}.')
+
             # TODO(romilb): This logic should likely be in add_store to move
             # syncing to file_mount stage..
             if self.sync_on_reconstruction:
@@ -648,15 +685,16 @@ class Storage(object):
 
         else:
             # Storage does not exist in global_user_state, create new stores
-            sky_managed_stores = {
-                t: s.get_metadata()
-                for t, s in self.stores.items()
-                if s.is_sky_managed
-            }
+            # Sky optimizer either adds a storage object instance or selects
+            # from existing ones
+            input_stores = self.stores
+            self.stores = {}
             self.handle = self.StorageMetadata(storage_name=self.name,
                                                source=self.source,
-                                               mode=self.mode,
-                                               sky_stores=sky_managed_stores)
+                                               mode=self.mode)
+
+            for store in input_stores:
+                self.add_store(store)
 
             if self.source is not None:
                 # If source is a pre-existing bucket, connect to the bucket
@@ -987,10 +1025,13 @@ class Storage(object):
           region: str; Region to place the bucket in. Caller must ensure that
             the region is valid for the chosen store_type.
         """
+        assert self._constructed, self
+        assert self.name is not None, self
+
         if isinstance(store_type, str):
             store_type = StoreType(store_type)
 
-        if store_type in self.stores:
+        if self.stores.get(store_type) is not None:
             if store_type == StoreType.AZURE:
                 azure_store_obj = self.stores[store_type]
                 assert isinstance(azure_store_obj, AzureBlobStore)
@@ -999,8 +1040,9 @@ class Storage(object):
                             f'storage account {storage_account_name!r}.')
             else:
                 logger.info(f'Storage type {store_type} already exists.')
-
-            return self.stores[store_type]
+            store = self.stores[store_type]
+            assert store is not None, self
+            return store
 
         store_cls: Type[AbstractStore]
         if store_type == StoreType.S3:
@@ -1068,6 +1110,7 @@ class Storage(object):
         if store.is_sky_managed:
             self.handle.add_store(store)
             if not is_reconstructed:
+                assert self.name is not None, self
                 global_user_state.add_or_update_storage(self.name, self.handle,
                                                         StorageStatus.INIT)
 
@@ -1083,17 +1126,26 @@ class Storage(object):
         """
         if not self.stores:
             logger.info('No backing stores found. Deleting storage.')
+            assert self.name is not None
             global_user_state.remove_storage(self.name)
         if store_type is not None:
+            assert self.name is not None
             store = self.stores[store_type]
+            assert store is not None, self
+            is_sky_managed = store.is_sky_managed
             # We delete a store from the cloud if it's sky managed. Else just
             # remove handle and return
-            if store.is_sky_managed:
+            if is_sky_managed:
                 self.handle.remove_store(store)
                 store.delete()
                 # Check remaining stores - if none is sky managed, remove
                 # the storage from global_user_state.
-                delete = all(not s.is_sky_managed for s in self.stores.values())
+                delete = True
+                for store in self.stores.values():
+                    assert store is not None, self
+                    if store.is_sky_managed:
+                        delete = False
+                        break
                 if delete:
                     global_user_state.remove_storage(self.name)
                 else:
@@ -1104,6 +1156,7 @@ class Storage(object):
             del self.stores[store_type]
         else:
             for _, store in self.stores.items():
+                assert store is not None, self
                 if store.is_sky_managed:
                     self.handle.remove_store(store)
                     store.delete()
@@ -1111,15 +1164,19 @@ class Storage(object):
                     store.delete()
             self.stores = {}
             # Remove storage from global_user_state if present
-            global_user_state.remove_storage(self.name)
+            if self.name is not None:
+                global_user_state.remove_storage(self.name)
 
     def sync_all_stores(self):
         """Syncs the source and destinations of all stores in the Storage"""
         for _, store in self.stores.items():
+            assert store is not None, self
             self._sync_store(store)
 
     def _sync_store(self, store: AbstractStore):
         """Runs the upload routine for the store and handles failures"""
+        assert self._constructed, self
+        assert self.name is not None, self
 
         def warn_for_git_dir(source: str):
             if os.path.isdir(os.path.join(source, '.git')):
@@ -1175,14 +1232,17 @@ class Storage(object):
         assert not config, f'Invalid storage args: {config.keys()}'
 
         # Validation of the config object happens on instantiation.
+        if store is not None:
+            stores = [StoreType(store.upper())]
+        else:
+            stores = None
         storage_obj = cls(name=name,
                           source=source,
                           persistent=persistent,
                           mode=mode,
+                          stores=stores,
                           _is_sky_managed=_is_sky_managed,
                           _bucket_sub_path=_bucket_sub_path)
-        if store is not None:
-            storage_obj.add_store(StoreType(store.upper()))
 
         # Add force deletion flag
         storage_obj.force_delete = force_delete
@@ -1207,7 +1267,9 @@ class Storage(object):
         is_sky_managed = self._is_sky_managed
         if self.stores:
             stores = ','.join([store.value for store in self.stores])
-            is_sky_managed = list(self.stores.values())[0].is_sky_managed
+            store = list(self.stores.values())[0]
+            if store is not None:
+                is_sky_managed = store.is_sky_managed
         add_if_not_none('store', stores)
         add_if_not_none('_is_sky_managed', is_sky_managed)
         add_if_not_none('persistent', self.persistent)

sky/data/storage_utils.py

@@ -1,9 +1,12 @@
 """Utility functions for the storage module."""
 import glob
 import os
+import pathlib
 import shlex
 import subprocess
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, TextIO, Union
+import warnings
+import zipfile
 
 import colorama
 
@@ -213,13 +216,70 @@ def get_excluded_files(src_dir_path: str) -> List[str]:
     skyignore_path = os.path.join(expand_src_dir_path,
                                   constants.SKY_IGNORE_FILE)
     if os.path.exists(skyignore_path):
-        logger.info(f'  {colorama.Style.DIM}'
-                    f'Excluded files to sync to cluster based on '
-                    f'{constants.SKY_IGNORE_FILE}.'
-                    f'{colorama.Style.RESET_ALL}')
+        logger.debug(f'  {colorama.Style.DIM}'
+                     f'Excluded files to sync to cluster based on '
+                     f'{constants.SKY_IGNORE_FILE}.'
+                     f'{colorama.Style.RESET_ALL}')
         return get_excluded_files_from_skyignore(src_dir_path)
-    logger.info(f'  {colorama.Style.DIM}'
-                f'Excluded files to sync to cluster based on '
-                f'{constants.GIT_IGNORE_FILE}.'
-                f'{colorama.Style.RESET_ALL}')
+    logger.debug(f'  {colorama.Style.DIM}'
+                 f'Excluded files to sync to cluster based on '
+                 f'{constants.GIT_IGNORE_FILE}.'
+                 f'{colorama.Style.RESET_ALL}')
     return get_excluded_files_from_gitignore(src_dir_path)
+
+
+def zip_files_and_folders(items: List[str],
+                          output_file: Union[str, pathlib.Path],
+                          log_file: Optional[TextIO] = None):
+
+    def _store_symlink(zipf, path: str, is_dir: bool):
+        # Get the target of the symlink
+        target = os.readlink(path)
+        # Use relative path as absolute path will not be able to resolve on
+        # remote API server.
+        if os.path.isabs(target):
+            target = os.path.relpath(target, os.path.dirname(path))
+        # Create a ZipInfo instance
+        zi = zipfile.ZipInfo(path + '/') if is_dir else zipfile.ZipInfo(path)
+        # Set external attributes to mark as symlink
+        zi.external_attr = 0xA1ED0000
+        # Write symlink target as content
+        zipf.writestr(zi, target)
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore',
+                                category=UserWarning,
+                                message='Duplicate name:')
+        with zipfile.ZipFile(output_file, 'w') as zipf:
+            for item in items:
+                item = os.path.expanduser(item)
+                if not os.path.isfile(item) and not os.path.isdir(item):
+                    raise ValueError(f'{item} does not exist.')
+                excluded_files = set(
+                    [os.path.join(item, f) for f in get_excluded_files(item)])
+                if os.path.isfile(item) and item not in excluded_files:
+                    zipf.write(item)
+                elif os.path.isdir(item):
+                    for root, dirs, files in os.walk(item, followlinks=False):
+                        # Store directory entries (important for empty
+                        # directories)
+                        for dir_name in dirs:
+                            dir_path = os.path.join(root, dir_name)
+                            if dir_path in excluded_files:
+                                continue
+                            # If it's a symlink, store it as a symlink
+                            if os.path.islink(dir_path):
+                                _store_symlink(zipf, dir_path, is_dir=True)
+                            else:
+                                zipf.write(dir_path)
+
+                        for file in files:
+                            file_path = os.path.join(root, file)
+                            if file_path in excluded_files:
+                                continue
+                            if os.path.islink(file_path):
+                                _store_symlink(zipf, file_path, is_dir=False)
+                            else:
+                                zipf.write(file_path)
+                if log_file is not None:
+                    log_file.write(f'Zipped {item}\n')

sky/exceptions.py

@@ -1,13 +1,16 @@
 """Exceptions."""
+import builtins
 import enum
+import traceback
+import types
 import typing
-from typing import List, Optional, Sequence
+from typing import Any, Dict, List, Optional, Sequence
 
 from sky.utils import env_options
 
 if typing.TYPE_CHECKING:
-    from sky import status_lib
     from sky.backends import backend
+    from sky.utils import status_lib
 
 # Return code for keyboard interruption and SIGTSTP
 KEYBOARD_INTERRUPT_CODE = 130
@@ -21,6 +24,107 @@ INSUFFICIENT_PRIVILEGES_CODE = 52
 GIT_FATAL_EXIT_CODE = 128
 
 
+def is_safe_exception(exc: Exception) -> bool:
+    """Returns True if the exception is safe to send to clients.
+
+    Safe exceptions are:
+    1. Built-in exceptions
+    2. SkyPilot's own exceptions
+    """
+    module = type(exc).__module__
+
+    # Builtin exceptions (e.g., ValueError, RuntimeError)
+    if module == 'builtins':
+        return True
+
+    # SkyPilot exceptions
+    if module.startswith('sky.'):
+        return True
+
+    return False
+
+
+def wrap_exception(exc: Exception) -> Exception:
+    """Wraps non-safe exceptions into SkyPilot exceptions
+
+    This is used to wrap exceptions that are not safe to deserialize at clients.
+
+    Examples include exceptions from cloud providers whose packages are not
+    available at clients.
+    """
+    if is_safe_exception(exc):
+        return exc
+
+    return CloudError(message=str(exc),
+                      cloud_provider=type(exc).__module__.split('.')[0],
+                      error_type=type(exc).__name__)
+
+
+def serialize_exception(e: Exception) -> Dict[str, Any]:
+    """Serialize the exception.
+
+    This function also wraps any unsafe exceptions (e.g., cloud exceptions)
+    into SkyPilot's CloudError before serialization to ensure clients can
+    deserialize them without needing cloud provider packages installed.
+    """
+    # Wrap unsafe exceptions before serialization
+    e = wrap_exception(e)
+
+    stacktrace = getattr(e, 'stacktrace', None)
+    attributes = e.__dict__.copy()
+    if 'stacktrace' in attributes:
+        del attributes['stacktrace']
+    for attr_k in list(attributes.keys()):
+        attr_v = attributes[attr_k]
+        if isinstance(attr_v, types.TracebackType):
+            attributes[attr_k] = traceback.format_tb(attr_v)
+
+    data = {
+        'type': e.__class__.__name__,
+        'message': str(e),
+        'args': e.args,
+        'attributes': attributes,
+        'stacktrace': stacktrace,
+    }
+    if isinstance(e, SkyPilotExcludeArgsBaseException):
+        data['args'] = tuple()
+    return data
+
+
+def deserialize_exception(serialized: Dict[str, Any]) -> Exception:
+    """Deserialize the exception."""
+    exception_type = serialized['type']
+    if hasattr(builtins, exception_type):
+        exception_class = getattr(builtins, exception_type)
+    else:
+        exception_class = globals().get(exception_type, None)
+    if exception_class is None:
+        # Unknown exception type.
+        return Exception(f'{exception_type}: {serialized["message"]}')
+    e = exception_class(*serialized['args'], **serialized['attributes'])
+    if serialized['stacktrace'] is not None:
+        setattr(e, 'stacktrace', serialized['stacktrace'])
+    return e
+
+
+class CloudError(Exception):
+    """Wraps cloud-specific errors into a SkyPilot exception."""
+
+    def __init__(self, message: str, cloud_provider: str, error_type: str):
+        super().__init__(message)
+        self.cloud_provider = cloud_provider
+        self.error_type = error_type
+
+    def __str__(self):
+        return (f'{self.cloud_provider} error ({self.error_type}): '
+                f'{super().__str__()}')
+
+
+class InvalidSkyPilotConfigError(ValueError):
+    """Raised when the SkyPilot config is invalid."""
+    pass
+
+
 class ResourcesUnavailableError(Exception):
     """Raised when resources are unavailable.
 
@@ -94,7 +198,21 @@ class ResourcesMismatchError(Exception):
     pass
 
 
-class CommandError(Exception):
+class SkyPilotExcludeArgsBaseException(Exception):
+    """Base class for exceptions that don't need args while serialization.
+
+    Due to our serialization/deserialization logic, when an exception does
+    not take `args` as an argument in __init__, `args` should not be included
+    in the serialized exception.
+
+    This is useful when an exception needs to construct the error message based
+    on the arguments passed in instead of directly having the error message as
+    the first argument in __init__. Refer to `CommandError` for an example.
+    """
+    pass
+
+
+class CommandError(SkyPilotExcludeArgsBaseException):
     """Raised when a command fails.
 
     Args:
@@ -128,7 +246,7 @@ class ClusterNotUpError(Exception):
 
     def __init__(self,
                  message: str,
-                 cluster_status: Optional['status_lib.ClusterStatus'],
+                 cluster_status: Optional['status_lib.ClusterStatus'] = None,
                  handle: Optional['backend.ResourceHandle'] = None) -> None:
         super().__init__(message)
         self.cluster_status = cluster_status
@@ -265,7 +383,7 @@ class NoCloudAccessError(Exception):
     pass
 
 
-class AWSAzFetchingError(Exception):
+class AWSAzFetchingError(SkyPilotExcludeArgsBaseException):
     """Raised when fetching the AWS availability zone fails."""
 
     class Reason(enum.Enum):
@@ -314,3 +432,18 @@ class UserRequestRejectedByPolicy(Exception):
 class NoClusterLaunchedError(Exception):
     """No cluster launched, so cleanup can be skipped during failover."""
     pass
+
+
+class RequestCancelled(Exception):
+    """Raised when a request is cancelled."""
+    pass
+
+
+class ApiServerConnectionError(RuntimeError):
+    """Raised when the API server cannot be connected."""
+
+    def __init__(self, server_url: str):
+        super().__init__(
+            f'Could not connect to SkyPilot API server at {server_url}. '
+            f'Please ensure that the server is running. '
+            f'Try: curl {server_url}/api/health')