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

sky/backends/local_docker_backend.py

@@ -14,6 +14,7 @@ from sky.backends import backend_utils
 from sky.backends import docker_utils
 from sky.data import storage as storage_lib
 from sky.utils import rich_utils
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky import resources
@@ -130,13 +131,14 @@ class LocalDockerBackend(backends.Backend['LocalDockerResourceHandle']):
         pass
 
     def _provision(
-            self,
-            task: 'task_lib.Task',
-            to_provision: Optional['resources.Resources'],
-            dryrun: bool,
-            stream_logs: bool,
-            cluster_name: str,
-            retry_until_up: bool = False
+        self,
+        task: 'task_lib.Task',
+        to_provision: Optional['resources.Resources'],
+        dryrun: bool,
+        stream_logs: bool,
+        cluster_name: str,
+        retry_until_up: bool = False,
+        skip_unnecessary_provisioning: bool = False,
     ) -> Optional[LocalDockerResourceHandle]:
         """Builds docker image for the task and returns cluster name as handle.
 
@@ -152,6 +154,9 @@ class LocalDockerBackend(backends.Backend['LocalDockerResourceHandle']):
             logger.warning(
                 f'Retrying until up is not supported in backend: {self.NAME}. '
                 'Ignored the flag.')
+        if skip_unnecessary_provisioning:
+            logger.warning(f'skip_unnecessary_provisioning is not supported in '
+                           f'backend: {self.NAME}. Ignored the flag.')
         if stream_logs:
             logger.info(
                 'Streaming build logs is not supported in LocalDockerBackend. '
@@ -159,7 +164,8 @@ class LocalDockerBackend(backends.Backend['LocalDockerResourceHandle']):
         handle = LocalDockerResourceHandle(cluster_name)
         logger.info(f'Building docker image for task {task.name}. '
                     'This might take some time.')
-        with rich_utils.safe_status('[bold cyan]Building Docker image[/]'):
+        with rich_utils.safe_status(
+                ux_utils.spinner_message('Building Docker image')):
             image_tag, metadata = docker_utils.build_dockerimage_from_task(task)
         self.images[handle] = (image_tag, metadata)
         logger.info(f'Image {image_tag} built.')

sky/backends/wheel_utils.py

@@ -39,29 +39,30 @@ _WHEEL_PATTERN = (f'{_PACKAGE_WHEEL_NAME}-'
                   f'{version.parse(sky.__version__)}-*.whl')
 
 
-def _get_latest_wheel_and_remove_all_others() -> pathlib.Path:
-    wheel_name = (f'**/{_WHEEL_PATTERN}')
+def _remove_stale_wheels(latest_wheel_dir: pathlib.Path) -> None:
+    """Remove all wheels except the latest one."""
+    for f in WHEEL_DIR.iterdir():
+        if f != latest_wheel_dir:
+            if f.is_dir() and not f.is_symlink():
+                shutil.rmtree(f, ignore_errors=True)
+
+
+def _get_latest_wheel() -> pathlib.Path:
+    wheel_name = f'**/{_WHEEL_PATTERN}'
     try:
         latest_wheel = max(WHEEL_DIR.glob(wheel_name), key=os.path.getctime)
     except ValueError:
         raise FileNotFoundError(
             'Could not find built SkyPilot wheels with glob pattern '
             f'{wheel_name} under {WHEEL_DIR!r}') from None
-
-    latest_wheel_dir_name = latest_wheel.parent
-    # Cleanup older wheels.
-    for f in WHEEL_DIR.iterdir():
-        if f != latest_wheel_dir_name:
-            if f.is_dir() and not f.is_symlink():
-                shutil.rmtree(f, ignore_errors=True)
     return latest_wheel
 
 
-def _build_sky_wheel():
-    """Build a wheel for SkyPilot."""
-    with tempfile.TemporaryDirectory() as tmp_dir:
+def _build_sky_wheel() -> pathlib.Path:
+    """Build a wheel for SkyPilot and return the path to the wheel."""
+    with tempfile.TemporaryDirectory() as tmp_dir_str:
         # prepare files
-        tmp_dir = pathlib.Path(tmp_dir)
+        tmp_dir = pathlib.Path(tmp_dir_str)
         sky_tmp_dir = tmp_dir / 'sky'
         sky_tmp_dir.mkdir()
         for item in SKY_PACKAGE_PATH.iterdir():
@@ -128,7 +129,12 @@ def _build_sky_wheel():
 
         wheel_dir = WHEEL_DIR / hash_of_latest_wheel
         wheel_dir.mkdir(parents=True, exist_ok=True)
-        shutil.move(str(wheel_path), wheel_dir)
+        # shutil.move will fail when the file already exists and is being
+        # moved across filesystems.
+        if not os.path.exists(
+                os.path.join(wheel_dir, os.path.basename(wheel_path))):
+            shutil.move(str(wheel_path), wheel_dir)
+        return wheel_dir / wheel_path.name
 
 
 def build_sky_wheel() -> Tuple[pathlib.Path, str]:
@@ -147,7 +153,10 @@ def build_sky_wheel() -> Tuple[pathlib.Path, str]:
         if not path.exists():
             return -1.
         try:
-            return max(os.path.getmtime(root) for root, _, _ in os.walk(path))
+            return max(
+                os.path.getmtime(os.path.join(root, f))
+                for root, dirs, files in os.walk(path)
+                for f in (*dirs, *files))
         except ValueError:
             return -1.
 
@@ -161,13 +170,22 @@ def build_sky_wheel() -> Tuple[pathlib.Path, str]:
         last_modification_time = _get_latest_modification_time(SKY_PACKAGE_PATH)
         last_wheel_modification_time = _get_latest_modification_time(WHEEL_DIR)
 
-        # only build wheels if the wheel is outdated
-        if last_wheel_modification_time < last_modification_time:
+        # Only build wheels if the wheel is outdated or wheel does not exist
+        # for the requested version.
+        if (last_wheel_modification_time < last_modification_time) or not any(
+                WHEEL_DIR.glob(f'**/{_WHEEL_PATTERN}')):
             if not WHEEL_DIR.exists():
                 WHEEL_DIR.mkdir(parents=True, exist_ok=True)
-            _build_sky_wheel()
-
-        latest_wheel = _get_latest_wheel_and_remove_all_others()
+            latest_wheel = _build_sky_wheel()
+        else:
+            latest_wheel = _get_latest_wheel()
+
+        # We remove all wheels except the latest one for garbage collection.
+        # Otherwise stale wheels will accumulate over time.
+        # TODO(romilb): If the user switches versions every alternate launch,
+        #  the wheel will be rebuilt every time. At the risk of adding
+        #  complexity, we can consider TTL caching wheels by version here.
+        _remove_stale_wheels(latest_wheel.parent)
 
         wheel_hash = latest_wheel.parent.name
 

sky/benchmark/benchmark_utils.py

@@ -20,10 +20,11 @@ from rich import progress as rich_progress
 
 import sky
 from sky import backends
+from sky import clouds
 from sky import data
 from sky import global_user_state
+from sky import optimizer
 from sky import sky_logging
-from sky import status_lib
 from sky.backends import backend_utils
 from sky.benchmark import benchmark_state
 from sky.data import storage as storage_lib
@@ -33,6 +34,7 @@ from sky.skylet import log_lib
 from sky.utils import common_utils
 from sky.utils import log_utils
 from sky.utils import rich_utils
+from sky.utils import status_lib
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
@@ -99,7 +101,9 @@ def _get_optimized_resources(
             task = sky.Task()
             task.set_resources(resources)
 
-        dag = sky.optimize(dag, quiet=True)
+        # Do not use `sky.optimize` here, as this should be called on the API
+        # server side.
+        dag = optimizer.Optimizer.optimize(dag, quiet=True)
         task = dag.tasks[0]
         optimized_resources.append(task.best_resources)
     return optimized_resources
@@ -170,13 +174,19 @@ def _create_benchmark_bucket() -> Tuple[str, str]:
     # Select the bucket type.
     enabled_clouds = storage_lib.get_cached_enabled_storage_clouds_or_refresh(
         raise_if_no_cloud_access=True)
-    # Already checked by raise_if_no_cloud_access=True.
-    assert enabled_clouds
+    # Sky Benchmark only supports S3 (see _download_remote_dir and
+    # _delete_remote_dir).
+    enabled_clouds = [
+        cloud for cloud in enabled_clouds if cloud in [str(clouds.AWS())]
+    ]
+    assert enabled_clouds, ('No enabled cloud storage found. Sky Benchmark '
+                            'requires GCP or AWS to store logs.')
     bucket_type = data.StoreType.from_cloud(enabled_clouds[0]).value
 
     # Create a benchmark bucket.
     logger.info(f'Creating a bucket {bucket_name} to save the benchmark logs.')
     storage = data.Storage(bucket_name, source=None, persistent=True)
+    storage.construct()
     storage.add_store(bucket_type)
 
     # Save the bucket name and type to the config.
@@ -242,14 +252,8 @@ def _download_remote_dir(remote_dir: str, local_dir: str,
             stdout=subprocess.DEVNULL,
             stderr=subprocess.DEVNULL,
             check=True)
-    elif bucket_type == data.StoreType.GCS:
-        remote_dir = f'gs://{remote_dir}'
-        subprocess.run(['gsutil', '-m', 'cp', '-r', remote_dir, local_dir],
-                       stdout=subprocess.DEVNULL,
-                       stderr=subprocess.DEVNULL,
-                       check=True)
     else:
-        raise RuntimeError('Azure Blob Storage is not supported yet.')
+        raise RuntimeError(f'{bucket_type} is not supported yet.')
 
 
 def _delete_remote_dir(remote_dir: str, bucket_type: data.StoreType) -> None:
@@ -260,14 +264,8 @@ def _delete_remote_dir(remote_dir: str, bucket_type: data.StoreType) -> None:
                        stdout=subprocess.DEVNULL,
                        stderr=subprocess.DEVNULL,
                        check=True)
-    elif bucket_type == data.StoreType.GCS:
-        remote_dir = f'gs://{remote_dir}'
-        subprocess.run(['gsutil', '-m', 'rm', '-r', remote_dir],
-                       stdout=subprocess.DEVNULL,
-                       stderr=subprocess.DEVNULL,
-                       check=True)
     else:
-        raise RuntimeError('Azure Blob Storage is not supported yet.')
+        raise RuntimeError(f'{bucket_type} is not supported yet.')
 
 
 def _read_timestamp(path: str) -> float:
@@ -541,7 +539,7 @@ def launch_benchmark_clusters(benchmark: str, clusters: List[str],
                    for yaml_fd, cluster in zip(yaml_fds, clusters)]
 
     # Save stdout/stderr from cluster launches.
-    run_timestamp = backend_utils.get_run_timestamp()
+    run_timestamp = sky_logging.get_run_timestamp()
     log_dir = os.path.join(constants.SKY_LOGS_DIRECTORY, run_timestamp)
     log_dir = os.path.expanduser(log_dir)
     logger.info(
@@ -601,7 +599,8 @@ def update_benchmark_state(benchmark: str) -> None:
     remote_dir = os.path.join(bucket_name, benchmark)
     local_dir = os.path.join(_SKY_LOCAL_BENCHMARK_DIR, benchmark)
     os.makedirs(local_dir, exist_ok=True)
-    with rich_utils.safe_status('[bold cyan]Downloading benchmark logs[/]'):
+    with rich_utils.safe_status(
+            ux_utils.spinner_message('Downloading benchmark logs')):
         _download_remote_dir(remote_dir, local_dir, bucket_type)
 
     # Update the benchmark results in parallel.
@@ -610,9 +609,9 @@ def update_benchmark_state(benchmark: str) -> None:
     progress = rich_progress.Progress(transient=True,
                                       redirect_stdout=False,
                                       redirect_stderr=False)
-    task = progress.add_task(
-        f'[bold cyan]Processing {num_candidates} benchmark result{plural}[/]',
-        total=num_candidates)
+    task = progress.add_task(ux_utils.spinner_message(
+        f'Processing {num_candidates} benchmark result{plural}'),
+                             total=num_candidates)
 
     def _update_with_progress_bar(arg: Any) -> None:
         message = _update_benchmark_result(arg)

sky/check.py

@@ -1,26 +1,33 @@
 """Credential checks: check cloud credentials and enable clouds."""
+import os
 import traceback
 from types import ModuleType
 from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 import click
 import colorama
-import rich
 
 from sky import clouds as sky_clouds
 from sky import exceptions
 from sky import global_user_state
 from sky import skypilot_config
 from sky.adaptors import cloudflare
+from sky.utils import registry
+from sky.utils import rich_utils
 from sky.utils import ux_utils
 
+CHECK_MARK_EMOJI = '\U00002714'  # Heavy check mark unicode
+PARTY_POPPER_EMOJI = '\U0001F389'  # Party popper unicode
+
 
 def check(
     quiet: bool = False,
     verbose: bool = False,
     clouds: Optional[Iterable[str]] = None,
-) -> None:
-    echo = (lambda *_args, **_kwargs: None) if quiet else click.echo
+) -> List[str]:
+    echo = (lambda *_args, **_kwargs: None
+           ) if quiet else lambda *args, **kwargs: click.echo(
+               *args, **kwargs, color=True)
     echo('Checking credentials to enable clouds for SkyPilot.')
     enabled_clouds = []
     disabled_clouds = []
@@ -29,14 +36,13 @@ def check(
             cloud_tuple: Tuple[str, Union[sky_clouds.Cloud,
                                           ModuleType]]) -> None:
         cloud_repr, cloud = cloud_tuple
-        echo(f'  Checking {cloud_repr}...', nl=False)
-        try:
-            ok, reason = cloud.check_credentials()
-        except Exception:  # pylint: disable=broad-except
-            # Catch all exceptions to prevent a single cloud from blocking the
-            # check for other clouds.
-            ok, reason = False, traceback.format_exc()
-        echo('\r', nl=False)
+        with rich_utils.safe_status(f'Checking {cloud_repr}...'):
+            try:
+                ok, reason = cloud.check_credentials()
+            except Exception:  # pylint: disable=broad-except
+                # Catch all exceptions to prevent a single cloud from blocking
+                # the check for other clouds.
+                ok, reason = False, traceback.format_exc()
         status_msg = 'enabled' if ok else 'disabled'
         styles = {'fg': 'green', 'bold': False} if ok else {'dim': True}
         echo('  ' + click.style(f'{cloud_repr}: {status_msg}', **styles) +
@@ -44,7 +50,7 @@ def check(
         if ok:
             enabled_clouds.append(cloud_repr)
             if verbose and cloud is not cloudflare:
-                activated_account = cloud.get_current_user_identity_str()
+                activated_account = cloud.get_active_user_identity_str()
                 if activated_account is not None:
                     echo(f'    Activated account: {activated_account}')
             if reason is not None:
@@ -60,12 +66,12 @@ def check(
         if cloud_name.lower().startswith('cloudflare'):
             return cloudflare.SKY_CHECK_NAME, cloudflare
         else:
-            cloud_obj = sky_clouds.CLOUD_REGISTRY.from_str(cloud_name)
+            cloud_obj = registry.CLOUD_REGISTRY.from_str(cloud_name)
             assert cloud_obj is not None, f'Cloud {cloud_name!r} not found'
             return repr(cloud_obj), cloud_obj
 
     def get_all_clouds():
-        return tuple([repr(c) for c in sky_clouds.CLOUD_REGISTRY.values()] +
+        return tuple([repr(c) for c in registry.CLOUD_REGISTRY.values()] +
                      [cloudflare.SKY_CHECK_NAME])
 
     if clouds is not None:
@@ -77,8 +83,8 @@ def check(
     # Use allowed_clouds from config if it exists, otherwise check all clouds.
     # Also validate names with get_cloud_tuple.
     config_allowed_cloud_names = [
-        get_cloud_tuple(c)[0] for c in skypilot_config.get_nested(
-            ['allowed_clouds'], get_all_clouds())
+        get_cloud_tuple(c)[0] for c in skypilot_config.get_nested((
+            'allowed_clouds',), get_all_clouds())
     ]
     # Use disallowed_cloud_names for logging the clouds that will be disabled
     # because they are not included in allowed_clouds in config.yaml.
@@ -93,7 +99,7 @@ def check(
     for cloud_tuple in sorted(clouds_to_check):
         check_one_cloud(cloud_tuple)
 
-    # Cloudflare is not a real cloud in sky_clouds.CLOUD_REGISTRY, and should
+    # Cloudflare is not a real cloud in registry.CLOUD_REGISTRY, and should
     # not be inserted into the DB (otherwise `sky launch` and other code would
     # error out when it's trying to look it up in the registry).
     enabled_clouds_set = {
@@ -126,7 +132,7 @@ def check(
             '\nNote: The following clouds were disabled because they were not '
             'included in allowed_clouds in ~/.sky/config.yaml: '
             f'{", ".join([c for c in disallowed_cloud_names])}')
-    if len(all_enabled_clouds) == 0:
+    if not all_enabled_clouds:
         echo(
             click.style(
                 'No cloud is enabled. SkyPilot will not be able to run any '
@@ -145,7 +151,7 @@ def check(
                 dim=True) + click.style(f'sky check{clouds_arg}', bold=True) +
             '\n' + click.style(
                 'If any problems remain, refer to detailed docs at: '
-                'https://skypilot.readthedocs.io/en/latest/getting-started/installation.html',  # pylint: disable=line-too-long
+                'https://docs.skypilot.co/en/latest/getting-started/installation.html',  # pylint: disable=line-too-long
                 dim=True))
 
         if disallowed_clouds_hint:
@@ -153,10 +159,14 @@ def check(
 
         # Pretty print for UX.
         if not quiet:
-            enabled_clouds_str = '\n  :heavy_check_mark: '.join(
-                [''] + sorted(all_enabled_clouds))
-            rich.print('\n[green]:tada: Enabled clouds :tada:'
-                       f'{enabled_clouds_str}[/green]')
+            enabled_clouds_str = '\n  ' + '\n  '.join([
+                _format_enabled_cloud(cloud)
+                for cloud in sorted(all_enabled_clouds)
+            ])
+            echo(f'\n{colorama.Fore.GREEN}{PARTY_POPPER_EMOJI} '
+                 f'Enabled clouds {PARTY_POPPER_EMOJI}'
+                 f'{colorama.Style.RESET_ALL}{enabled_clouds_str}')
+    return enabled_clouds
 
 
 def get_cached_enabled_clouds_or_refresh(
@@ -194,19 +204,25 @@ def get_cached_enabled_clouds_or_refresh(
 def get_cloud_credential_file_mounts(
         excluded_clouds: Optional[Iterable[sky_clouds.Cloud]]
 ) -> Dict[str, str]:
-    """Returns the files necessary to access all enabled clouds.
+    """Returns the files necessary to access all clouds.
 
     Returns a dictionary that will be added to a task's file mounts
     and a list of patterns that will be excluded (used as rsync_exclude).
     """
-    enabled_clouds = get_cached_enabled_clouds_or_refresh()
+    # Uploading credentials for all clouds instead of only sky check
+    # enabled clouds because users may have partial credentials for some
+    # clouds to access their specific resources (e.g. cloud storage) but
+    # not have the complete credentials to pass sky check.
+    clouds = registry.CLOUD_REGISTRY.values()
     file_mounts = {}
-    for cloud in enabled_clouds:
+    for cloud in clouds:
         if (excluded_clouds is not None and
                 sky_clouds.cloud_in_iterable(cloud, excluded_clouds)):
             continue
         cloud_file_mounts = cloud.get_credential_file_mounts()
-        file_mounts.update(cloud_file_mounts)
+        for remote_path, local_path in cloud_file_mounts.items():
+            if os.path.exists(os.path.expanduser(local_path)):
+                file_mounts[remote_path] = local_path
     # Currently, get_cached_enabled_clouds_or_refresh() does not support r2 as
     # only clouds with computing instances are marked as enabled by skypilot.
     # This will be removed when cloudflare/r2 is added as a 'cloud'.
@@ -215,3 +231,36 @@ def get_cloud_credential_file_mounts(
         r2_credential_mounts = cloudflare.get_credential_file_mounts()
         file_mounts.update(r2_credential_mounts)
     return file_mounts
+
+
+def _format_enabled_cloud(cloud_name: str) -> str:
+
+    def _green_color(cloud_name: str) -> str:
+        return f'{colorama.Fore.GREEN}{cloud_name}{colorama.Style.RESET_ALL}'
+
+    if cloud_name == repr(sky_clouds.Kubernetes()):
+        # Get enabled contexts for Kubernetes
+        existing_contexts = sky_clouds.Kubernetes.existing_allowed_contexts()
+        if not existing_contexts:
+            return _green_color(cloud_name)
+
+        # Check if allowed_contexts is explicitly set in config
+        allowed_contexts = skypilot_config.get_nested(
+            ('kubernetes', 'allowed_contexts'), None)
+
+        # Format the context info with consistent styling
+        if allowed_contexts is not None:
+            contexts_formatted = []
+            for i, context in enumerate(existing_contexts):
+                symbol = (ux_utils.INDENT_LAST_SYMBOL
+                          if i == len(existing_contexts) -
+                          1 else ux_utils.INDENT_SYMBOL)
+                contexts_formatted.append(f'\n    {symbol}{context}')
+            context_info = f'Allowed contexts:{"".join(contexts_formatted)}'
+        else:
+            context_info = f'Active context: {existing_contexts[0]}'
+
+        return (f'{_green_color(cloud_name)}\n'
+                f'  {colorama.Style.DIM}{context_info}'
+                f'{colorama.Style.RESET_ALL}')
+    return _green_color(cloud_name)