skypilot-nightly 1.0.0.dev20250320__py3-none-any.whl → 1.0.0.dev20250322__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (50) hide show
  1. sky/__init__.py +2 -2
  2. sky/adaptors/cloudflare.py +16 -5
  3. sky/adaptors/kubernetes.py +2 -1
  4. sky/adaptors/nebius.py +128 -6
  5. sky/backends/cloud_vm_ray_backend.py +3 -1
  6. sky/benchmark/benchmark_utils.py +3 -2
  7. sky/check.py +114 -114
  8. sky/cloud_stores.py +66 -0
  9. sky/clouds/aws.py +14 -7
  10. sky/clouds/azure.py +13 -6
  11. sky/clouds/cloud.py +34 -10
  12. sky/clouds/cudo.py +3 -2
  13. sky/clouds/do.py +3 -2
  14. sky/clouds/fluidstack.py +3 -2
  15. sky/clouds/gcp.py +8 -9
  16. sky/clouds/ibm.py +15 -6
  17. sky/clouds/kubernetes.py +3 -1
  18. sky/clouds/lambda_cloud.py +3 -1
  19. sky/clouds/nebius.py +59 -11
  20. sky/clouds/oci.py +15 -6
  21. sky/clouds/paperspace.py +3 -2
  22. sky/clouds/runpod.py +7 -1
  23. sky/clouds/scp.py +3 -1
  24. sky/clouds/service_catalog/kubernetes_catalog.py +3 -1
  25. sky/clouds/vast.py +3 -2
  26. sky/clouds/vsphere.py +3 -2
  27. sky/core.py +6 -4
  28. sky/data/data_transfer.py +75 -0
  29. sky/data/data_utils.py +34 -0
  30. sky/data/mounting_utils.py +18 -0
  31. sky/data/storage.py +540 -10
  32. sky/data/storage_utils.py +102 -84
  33. sky/exceptions.py +2 -0
  34. sky/global_user_state.py +12 -33
  35. sky/jobs/server/core.py +1 -1
  36. sky/jobs/utils.py +5 -0
  37. sky/optimizer.py +10 -5
  38. sky/resources.py +6 -1
  39. sky/setup_files/dependencies.py +3 -1
  40. sky/task.py +16 -5
  41. sky/utils/command_runner.py +2 -0
  42. sky/utils/controller_utils.py +8 -5
  43. sky/utils/kubernetes/gpu_labeler.py +4 -4
  44. sky/utils/kubernetes/kubernetes_deploy_utils.py +4 -3
  45. {skypilot_nightly-1.0.0.dev20250320.dist-info → skypilot_nightly-1.0.0.dev20250322.dist-info}/METADATA +16 -7
  46. {skypilot_nightly-1.0.0.dev20250320.dist-info → skypilot_nightly-1.0.0.dev20250322.dist-info}/RECORD +50 -50
  47. {skypilot_nightly-1.0.0.dev20250320.dist-info → skypilot_nightly-1.0.0.dev20250322.dist-info}/WHEEL +1 -1
  48. {skypilot_nightly-1.0.0.dev20250320.dist-info → skypilot_nightly-1.0.0.dev20250322.dist-info}/entry_points.txt +0 -0
  49. {skypilot_nightly-1.0.0.dev20250320.dist-info → skypilot_nightly-1.0.0.dev20250322.dist-info}/licenses/LICENSE +0 -0
  50. {skypilot_nightly-1.0.0.dev20250320.dist-info → skypilot_nightly-1.0.0.dev20250322.dist-info}/top_level.txt +0 -0
sky/__init__.py CHANGED
@@ -5,7 +5,7 @@ from typing import Optional
5
5
  import urllib.request
6
6
 
7
7
  # Replaced with the current commit when building the wheels.
8
- _SKYPILOT_COMMIT_SHA = 'a480f342522afcd17a3b30a20086f28333ddb7b5'
8
+ _SKYPILOT_COMMIT_SHA = '139a09e3445956740743049d352cd1cb6d202479'
9
9
 
10
10
 
11
11
  def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
35
35
 
36
36
 
37
37
  __commit__ = _get_git_commit()
38
- __version__ = '1.0.0.dev20250320'
38
+ __version__ = '1.0.0.dev20250322'
39
39
  __root_dir__ = os.path.dirname(os.path.abspath(__file__))
40
40
 
41
41
 
@@ -6,7 +6,9 @@ import os
6
6
  import threading
7
7
  from typing import Dict, Optional, Tuple
8
8
 
9
+ from sky import exceptions
9
10
  from sky.adaptors import common
11
+ from sky.clouds import cloud
10
12
  from sky.utils import annotations
11
13
  from sky.utils import ux_utils
12
14
 
@@ -23,7 +25,6 @@ R2_CREDENTIALS_PATH = '~/.cloudflare/r2.credentials'
23
25
  R2_PROFILE_NAME = 'r2'
24
26
  _INDENT_PREFIX = ' '
25
27
  NAME = 'Cloudflare'
26
- SKY_CHECK_NAME = 'Cloudflare (for R2 object store)'
27
28
 
28
29
 
29
30
  @contextlib.contextmanager
@@ -130,8 +131,8 @@ def client(service_name: str, region):
130
131
  @common.load_lazy_modules(_LAZY_MODULES)
131
132
  def botocore_exceptions():
132
133
  """AWS botocore exception."""
133
- from botocore import exceptions
134
- return exceptions
134
+ from botocore import exceptions as boto_exceptions
135
+ return boto_exceptions
135
136
 
136
137
 
137
138
  def create_endpoint():
@@ -148,8 +149,18 @@ def create_endpoint():
148
149
  return endpoint
149
150
 
150
151
 
151
- def check_credentials() -> Tuple[bool, Optional[str]]:
152
- return check_storage_credentials()
152
+ def check_credentials(
153
+ cloud_capability: cloud.CloudCapability) -> Tuple[bool, Optional[str]]:
154
+ if cloud_capability == cloud.CloudCapability.COMPUTE:
155
+ # for backward compatibility,
156
+ # we check storage credentials for compute.
157
+ # TODO(seungjin): properly return not supported error for compute.
158
+ return check_storage_credentials()
159
+ elif cloud_capability == cloud.CloudCapability.STORAGE:
160
+ return check_storage_credentials()
161
+ else:
162
+ raise exceptions.NotSupportedError(
163
+ f'{NAME} does not support {cloud_capability}.')
153
164
 
154
165
 
155
166
  def check_storage_credentials() -> Tuple[bool, Optional[str]]:
@@ -79,10 +79,11 @@ def _load_config(context: Optional[str] = None):
79
79
  ' If you were running a local Kubernetes '
80
80
  'cluster, run `sky local up` to start the cluster.')
81
81
  else:
82
+ kubeconfig_path = os.environ.get('KUBECONFIG', '~/.kube/config')
82
83
  err_str = (
83
84
  f'Failed to load Kubernetes configuration for {context!r}. '
84
85
  'Please check if your kubeconfig file exists at '
85
- f'~/.kube/config and is valid.\n{suffix}')
86
+ f'{kubeconfig_path} and is valid.\n{suffix}')
86
87
  err_str += '\nTo disable Kubernetes for SkyPilot: run `sky check`.'
87
88
  with ux_utils.print_exception_no_traceback():
88
89
  raise ValueError(err_str) from None
sky/adaptors/nebius.py CHANGED
@@ -1,7 +1,11 @@
1
1
  """Nebius cloud adaptor."""
2
2
  import os
3
+ import threading
4
+ from typing import Optional
3
5
 
4
6
  from sky.adaptors import common
7
+ from sky.utils import annotations
8
+ from sky.utils import ux_utils
5
9
 
6
10
  NEBIUS_TENANT_ID_FILENAME = 'NEBIUS_TENANT_ID.txt'
7
11
  NEBIUS_IAM_TOKEN_FILENAME = 'NEBIUS_IAM_TOKEN.txt'
@@ -12,6 +16,10 @@ NEBIUS_IAM_TOKEN_PATH = '~/.nebius/' + NEBIUS_IAM_TOKEN_FILENAME
12
16
  NEBIUS_PROJECT_ID_PATH = '~/.nebius/' + NEBIUS_PROJECT_ID_FILENAME
13
17
  NEBIUS_CREDENTIALS_PATH = '~/.nebius/' + NEBIUS_CREDENTIALS_FILENAME
14
18
 
19
+ DEFAULT_REGION = 'eu-north1'
20
+
21
+ NEBIUS_PROFILE_NAME = 'nebius'
22
+
15
23
  MAX_RETRIES_TO_DISK_CREATE = 120
16
24
  MAX_RETRIES_TO_INSTANCE_STOP = 120
17
25
  MAX_RETRIES_TO_INSTANCE_START = 120
@@ -23,15 +31,27 @@ MAX_RETRIES_TO_INSTANCE_WAIT = 120 # Maximum number of retries
23
31
  POLL_INTERVAL = 5
24
32
 
25
33
  _iam_token = None
34
+ _sdk = None
26
35
  _tenant_id = None
27
36
  _project_id = None
28
37
 
38
+ _IMPORT_ERROR_MESSAGE = ('Failed to import dependencies for Nebius AI Cloud.'
39
+ 'Try pip install "skypilot[nebius]"')
40
+
29
41
  nebius = common.LazyImport(
30
42
  'nebius',
31
- import_error_message='Failed to import dependencies for Nebius AI Cloud. '
32
- 'Try running: pip install "skypilot[nebius]"',
43
+ import_error_message=_IMPORT_ERROR_MESSAGE,
33
44
  # https://github.com/grpc/grpc/issues/37642 to avoid spam in console
34
45
  set_loggers=lambda: os.environ.update({'GRPC_VERBOSITY': 'NONE'}))
46
+ boto3 = common.LazyImport('boto3', import_error_message=_IMPORT_ERROR_MESSAGE)
47
+ botocore = common.LazyImport('botocore',
48
+ import_error_message=_IMPORT_ERROR_MESSAGE)
49
+
50
+ _LAZY_MODULES = (boto3, botocore, nebius)
51
+ _session_creation_lock = threading.RLock()
52
+ _INDENT_PREFIX = ' '
53
+ NAME = 'Nebius'
54
+ SKY_CHECK_NAME = 'Nebius (for Nebius Object Storae)'
35
55
 
36
56
 
37
57
  def request_error():
@@ -104,7 +124,109 @@ def get_tenant_id():
104
124
 
105
125
 
106
126
  def sdk():
107
- if get_iam_token() is not None:
108
- return nebius.sdk.SDK(credentials=get_iam_token())
109
- return nebius.sdk.SDK(
110
- credentials_file_name=os.path.expanduser(NEBIUS_CREDENTIALS_PATH))
127
+ global _sdk
128
+ if _sdk is None:
129
+ if get_iam_token() is not None:
130
+ _sdk = nebius.sdk.SDK(credentials=get_iam_token())
131
+ return _sdk
132
+ _sdk = nebius.sdk.SDK(
133
+ credentials_file_name=os.path.expanduser(NEBIUS_CREDENTIALS_PATH))
134
+ return _sdk
135
+
136
+
137
+ def get_nebius_credentials(boto3_session):
138
+ """Gets the Nebius credentials from the boto3 session object.
139
+
140
+ Args:
141
+ boto3_session: The boto3 session object.
142
+ Returns:
143
+ botocore.credentials.ReadOnlyCredentials object with the R2 credentials.
144
+ """
145
+ nebius_credentials = boto3_session.get_credentials()
146
+ if nebius_credentials is None:
147
+ with ux_utils.print_exception_no_traceback():
148
+ raise ValueError('Nebius credentials not found. Run '
149
+ '`sky check` to verify credentials are '
150
+ 'correctly set up.')
151
+ return nebius_credentials.get_frozen_credentials()
152
+
153
+
154
+ # lru_cache() is thread-safe and it will return the same session object
155
+ # for different threads.
156
+ # Reference: https://docs.python.org/3/library/functools.html#functools.lru_cache # pylint: disable=line-too-long
157
+ @annotations.lru_cache(scope='global')
158
+ def session():
159
+ """Create an AWS session."""
160
+ # Creating the session object is not thread-safe for boto3,
161
+ # so we add a reentrant lock to synchronize the session creation.
162
+ # Reference: https://github.com/boto/boto3/issues/1592
163
+ # However, the session object itself is thread-safe, so we are
164
+ # able to use lru_cache() to cache the session object.
165
+ with _session_creation_lock:
166
+ session_ = boto3.session.Session(profile_name=NEBIUS_PROFILE_NAME)
167
+ return session_
168
+
169
+
170
+ @annotations.lru_cache(scope='global')
171
+ def resource(resource_name: str, region: str = DEFAULT_REGION, **kwargs):
172
+ """Create a Nebius resource.
173
+
174
+ Args:
175
+ resource_name: Nebius resource name (e.g., 's3').
176
+ kwargs: Other options.
177
+ """
178
+ # Need to use the resource retrieved from the per-thread session
179
+ # to avoid thread-safety issues (Directly creating the client
180
+ # with boto3.resource() is not thread-safe).
181
+ # Reference: https://stackoverflow.com/a/59635814
182
+
183
+ session_ = session()
184
+ nebius_credentials = get_nebius_credentials(session_)
185
+ endpoint = create_endpoint(region)
186
+
187
+ return session_.resource(
188
+ resource_name,
189
+ endpoint_url=endpoint,
190
+ aws_access_key_id=nebius_credentials.access_key,
191
+ aws_secret_access_key=nebius_credentials.secret_key,
192
+ region_name=region,
193
+ **kwargs)
194
+
195
+
196
+ @annotations.lru_cache(scope='global')
197
+ def client(service_name: str, region):
198
+ """Create an Nebius client of a certain service.
199
+
200
+ Args:
201
+ service_name: Nebius service name (e.g., 's3').
202
+ kwargs: Other options.
203
+ """
204
+ # Need to use the client retrieved from the per-thread session
205
+ # to avoid thread-safety issues (Directly creating the client
206
+ # with boto3.client() is not thread-safe).
207
+ # Reference: https://stackoverflow.com/a/59635814
208
+
209
+ session_ = session()
210
+ nebius_credentials = get_nebius_credentials(session_)
211
+ endpoint = create_endpoint(region)
212
+
213
+ return session_.client(service_name,
214
+ endpoint_url=endpoint,
215
+ aws_access_key_id=nebius_credentials.access_key,
216
+ aws_secret_access_key=nebius_credentials.secret_key,
217
+ region_name=region)
218
+
219
+
220
+ @common.load_lazy_modules(_LAZY_MODULES)
221
+ def botocore_exceptions():
222
+ """AWS botocore exception."""
223
+ # pylint: disable=import-outside-toplevel
224
+ from botocore import exceptions
225
+ return exceptions
226
+
227
+
228
+ def create_endpoint(region: Optional[str] = DEFAULT_REGION) -> str:
229
+ """Reads accountid necessary to interact with Nebius Object Storage"""
230
+ if region is None:
231
+ region = DEFAULT_REGION
232
+ return f'https://storage.{region}.nebius.cloud:443'
@@ -38,6 +38,7 @@ from sky import sky_logging
38
38
  from sky import task as task_lib
39
39
  from sky.backends import backend_utils
40
40
  from sky.backends import wheel_utils
41
+ from sky.clouds import cloud as sky_cloud
41
42
  from sky.clouds import service_catalog
42
43
  from sky.clouds.utils import gcp_utils
43
44
  from sky.data import data_utils
@@ -1981,7 +1982,8 @@ class RetryingVmProvisioner(object):
1981
1982
  # is running. Here we check the enabled clouds and expiring credentials
1982
1983
  # and raise a warning to the user.
1983
1984
  if task.is_controller_task():
1984
- enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh()
1985
+ enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh(
1986
+ sky_cloud.CloudCapability.COMPUTE)
1985
1987
  expirable_clouds = backend_utils.get_expirable_clouds(
1986
1988
  enabled_clouds)
1987
1989
 
@@ -172,8 +172,9 @@ def _create_benchmark_bucket() -> Tuple[str, str]:
172
172
  bucket_name = f'sky-bench-{uuid.uuid4().hex[:4]}-{getpass.getuser()}'
173
173
 
174
174
  # Select the bucket type.
175
- enabled_clouds = storage_lib.get_cached_enabled_storage_clouds_or_refresh(
176
- raise_if_no_cloud_access=True)
175
+ enabled_clouds = (
176
+ storage_lib.get_cached_enabled_storage_cloud_names_or_refresh(
177
+ raise_if_no_cloud_access=True))
177
178
  # Sky Benchmark only supports S3 (see _download_remote_dir and
178
179
  # _delete_remote_dir).
179
180
  enabled_clouds = [
sky/check.py CHANGED
@@ -1,9 +1,9 @@
1
1
  """Credential checks: check cloud credentials and enable clouds."""
2
- import enum
3
2
  import os
4
3
  import traceback
5
4
  from types import ModuleType
6
- from typing import Dict, Iterable, List, Optional, Set, Tuple, Union
5
+ from typing import (Any, Callable, Dict, Iterable, List, Optional, Set, Tuple,
6
+ Union)
7
7
 
8
8
  import click
9
9
  import colorama
@@ -13,6 +13,7 @@ from sky import exceptions
13
13
  from sky import global_user_state
14
14
  from sky import skypilot_config
15
15
  from sky.adaptors import cloudflare
16
+ from sky.clouds import cloud as sky_cloud
16
17
  from sky.utils import registry
17
18
  from sky.utils import rich_utils
18
19
  from sky.utils import ux_utils
@@ -21,98 +22,55 @@ CHECK_MARK_EMOJI = '\U00002714' # Heavy check mark unicode
21
22
  PARTY_POPPER_EMOJI = '\U0001F389' # Party popper unicode
22
23
 
23
24
 
24
- # Declaring CloudCapability as a subclass of str
25
- # allows it to be JSON serializable.
26
- class CloudCapability(str, enum.Enum):
27
- # Compute capability.
28
- COMPUTE = 'compute'
29
- # Storage capability.
30
- STORAGE = 'storage'
31
-
32
-
33
- ALL_CAPABILITIES = [CloudCapability.COMPUTE, CloudCapability.STORAGE]
34
-
35
-
36
25
  def check_capabilities(
37
26
  quiet: bool = False,
38
27
  verbose: bool = False,
39
28
  clouds: Optional[Iterable[str]] = None,
40
- capabilities: Optional[List[CloudCapability]] = None,
41
- ) -> Dict[str, List[CloudCapability]]:
29
+ capabilities: Optional[List[sky_cloud.CloudCapability]] = None,
30
+ ) -> Dict[str, List[sky_cloud.CloudCapability]]:
42
31
  echo = (lambda *_args, **_kwargs: None
43
32
  ) if quiet else lambda *args, **kwargs: click.echo(
44
33
  *args, **kwargs, color=True)
45
34
  echo('Checking credentials to enable clouds for SkyPilot.')
46
35
  if capabilities is None:
47
- capabilities = ALL_CAPABILITIES
36
+ capabilities = sky_cloud.ALL_CAPABILITIES
48
37
  assert capabilities is not None
49
- enabled_clouds: Dict[str, List[CloudCapability]] = {}
50
- disabled_clouds: Dict[str, List[CloudCapability]] = {}
51
-
52
- def check_credentials(
53
- cloud: Union[sky_clouds.Cloud, ModuleType],
54
- capability: CloudCapability) -> Tuple[bool, Optional[str]]:
55
- if capability == CloudCapability.COMPUTE:
56
- return cloud.check_credentials()
57
- elif capability == CloudCapability.STORAGE:
58
- return cloud.check_storage_credentials()
59
- else:
60
- raise ValueError(f'Invalid capability: {capability}')
61
-
62
- def get_cached_state(capability: CloudCapability) -> List[sky_clouds.Cloud]:
63
- if capability == CloudCapability.COMPUTE:
64
- return global_user_state.get_cached_enabled_clouds()
65
- elif capability == CloudCapability.STORAGE:
66
- return global_user_state.get_cached_enabled_storage_clouds()
67
- else:
68
- raise ValueError(f'Invalid capability: {capability}')
69
-
70
- def set_cached_state(clouds: List[str],
71
- capability: CloudCapability) -> None:
72
- if capability == CloudCapability.COMPUTE:
73
- global_user_state.set_enabled_clouds(clouds)
74
- elif capability == CloudCapability.STORAGE:
75
- global_user_state.set_enabled_storage_clouds(clouds)
76
- else:
77
- raise ValueError(f'Invalid capability: {capability}')
38
+ enabled_clouds: Dict[str, List[sky_cloud.CloudCapability]] = {}
39
+ disabled_clouds: Dict[str, List[sky_cloud.CloudCapability]] = {}
78
40
 
79
41
  def check_one_cloud(
80
42
  cloud_tuple: Tuple[str, Union[sky_clouds.Cloud,
81
43
  ModuleType]]) -> None:
82
44
  cloud_repr, cloud = cloud_tuple
83
45
  assert capabilities is not None
46
+ # cloud_capabilities is a list of (capability, ok, reason)
47
+ # where ok is True if the cloud credentials are valid for the capability
48
+ cloud_capabilities: List[Tuple[sky_cloud.CloudCapability, bool,
49
+ Optional[str]]] = []
84
50
  for capability in capabilities:
85
51
  with rich_utils.safe_status(f'Checking {cloud_repr}...'):
86
52
  try:
87
- ok, reason = check_credentials(cloud, capability)
53
+ ok, reason = cloud.check_credentials(capability)
88
54
  except exceptions.NotSupportedError:
89
55
  continue
90
56
  except Exception: # pylint: disable=broad-except
91
57
  # Catch all exceptions to prevent a single cloud
92
58
  # from blocking the check for other clouds.
93
59
  ok, reason = False, traceback.format_exc()
94
- status_msg = ('enabled' if ok else 'disabled')
95
- styles = {'fg': 'green', 'bold': False} if ok else {'dim': True}
96
- echo(' ' + click.style(f'{cloud_repr}: {status_msg}', **styles) +
97
- ' ' * 30)
60
+ cloud_capabilities.append(
61
+ (capability, ok, reason.strip() if reason else None))
98
62
  if ok:
99
63
  enabled_clouds.setdefault(cloud_repr, []).append(capability)
100
- if verbose and cloud is not cloudflare:
101
- activated_account = cloud.get_active_user_identity_str()
102
- if activated_account is not None:
103
- echo(f' Activated account: {activated_account}')
104
- if reason is not None:
105
- echo(f' Hint: {reason}')
106
64
  else:
107
65
  disabled_clouds.setdefault(cloud_repr, []).append(capability)
108
- echo(f' Reason: {reason}')
66
+ _print_checked_cloud(echo, verbose, cloud_tuple, cloud_capabilities)
109
67
 
110
68
  def get_cloud_tuple(
111
69
  cloud_name: str) -> Tuple[str, Union[sky_clouds.Cloud, ModuleType]]:
112
70
  # Validates cloud_name and returns a tuple of the cloud's name and
113
71
  # the cloud object. Includes special handling for Cloudflare.
114
72
  if cloud_name.lower().startswith('cloudflare'):
115
- return cloudflare.SKY_CHECK_NAME, cloudflare
73
+ return cloudflare.NAME, cloudflare
116
74
  else:
117
75
  cloud_obj = registry.CLOUD_REGISTRY.from_str(cloud_name)
118
76
  assert cloud_obj is not None, f'Cloud {cloud_name!r} not found'
@@ -120,7 +78,7 @@ def check_capabilities(
120
78
 
121
79
  def get_all_clouds():
122
80
  return tuple([repr(c) for c in registry.CLOUD_REGISTRY.values()] +
123
- [cloudflare.SKY_CHECK_NAME])
81
+ [cloudflare.NAME])
124
82
 
125
83
  if clouds is not None:
126
84
  cloud_list = clouds
@@ -170,12 +128,14 @@ def check_capabilities(
170
128
  if not cloud.startswith('Cloudflare')
171
129
  }
172
130
  previously_enabled_clouds_set = {
173
- repr(cloud) for cloud in get_cached_state(capability)
131
+ repr(cloud)
132
+ for cloud in global_user_state.get_cached_enabled_clouds(capability)
174
133
  }
175
134
  enabled_clouds_for_capability = (config_allowed_clouds_set & (
176
135
  (previously_enabled_clouds_set | enabled_clouds_set) -
177
136
  disabled_clouds_set))
178
- set_cached_state(list(enabled_clouds_for_capability), capability)
137
+ global_user_state.set_enabled_clouds(
138
+ list(enabled_clouds_for_capability), capability)
179
139
  all_enabled_clouds = all_enabled_clouds.union(
180
140
  enabled_clouds_for_capability)
181
141
  disallowed_clouds_hint = None
@@ -195,8 +155,8 @@ def check_capabilities(
195
155
  echo(click.style(disallowed_clouds_hint, dim=True))
196
156
  raise SystemExit()
197
157
  else:
198
- clouds_arg = (' ' +
199
- ' '.join(disabled_clouds) if clouds is not None else '')
158
+ clouds_arg = (f' {" ".join(disabled_clouds).lower()}'
159
+ if clouds is not None else '')
200
160
  echo(
201
161
  click.style(
202
162
  '\nTo enable a cloud, follow the hints above and rerun: ',
@@ -212,7 +172,8 @@ def check_capabilities(
212
172
  # Pretty print for UX.
213
173
  if not quiet:
214
174
  enabled_clouds_str = '\n ' + '\n '.join([
215
- _format_enabled_cloud(cloud) for cloud in sorted(enabled_clouds)
175
+ _format_enabled_cloud(cloud, capabilities)
176
+ for cloud, capabilities in enabled_clouds.items()
216
177
  ])
217
178
  echo(f'\n{colorama.Fore.GREEN}{PARTY_POPPER_EMOJI} '
218
179
  f'Enabled clouds {PARTY_POPPER_EMOJI}'
@@ -220,14 +181,11 @@ def check_capabilities(
220
181
  return enabled_clouds
221
182
 
222
183
 
223
- # 'sky check' command and associated '/check' server endpoint
224
- # only checks compute capability for backward compatibility.
225
- # This necessitates setting default capability to CloudCapability.COMPUTE.
226
- def check(
184
+ def check_capability(
185
+ capability: sky_cloud.CloudCapability,
227
186
  quiet: bool = False,
228
187
  verbose: bool = False,
229
188
  clouds: Optional[Iterable[str]] = None,
230
- capability: CloudCapability = CloudCapability.COMPUTE,
231
189
  ) -> List[str]:
232
190
  clouds_with_capability = []
233
191
  enabled_clouds = check_capabilities(quiet, verbose, clouds, [capability])
@@ -237,7 +195,18 @@ def check(
237
195
  return clouds_with_capability
238
196
 
239
197
 
198
+ def check(
199
+ quiet: bool = False,
200
+ verbose: bool = False,
201
+ clouds: Optional[Iterable[str]] = None,
202
+ ) -> List[str]:
203
+ return list(
204
+ check_capabilities(quiet, verbose, clouds,
205
+ sky_cloud.ALL_CAPABILITIES).keys())
206
+
207
+
240
208
  def get_cached_enabled_clouds_or_refresh(
209
+ capability: sky_cloud.CloudCapability,
241
210
  raise_if_no_cloud_access: bool = False) -> List[sky_clouds.Cloud]:
242
211
  """Returns cached enabled clouds and if no cloud is enabled, refresh.
243
212
 
@@ -251,16 +220,18 @@ def get_cached_enabled_clouds_or_refresh(
251
220
  exceptions.NoCloudAccessError: if no public cloud is enabled and
252
221
  raise_if_no_cloud_access is set to True.
253
222
  """
254
- cached_enabled_clouds = global_user_state.get_cached_enabled_clouds()
223
+ cached_enabled_clouds = global_user_state.get_cached_enabled_clouds(
224
+ capability)
255
225
  if not cached_enabled_clouds:
256
226
  try:
257
- check(quiet=True, capability=CloudCapability.COMPUTE)
227
+ check_capability(sky_cloud.CloudCapability.COMPUTE, quiet=True)
258
228
  except SystemExit:
259
229
  # If no cloud is enabled, check() will raise SystemExit.
260
230
  # Here we catch it and raise the exception later only if
261
231
  # raise_if_no_cloud_access is set to True.
262
232
  pass
263
- cached_enabled_clouds = global_user_state.get_cached_enabled_clouds()
233
+ cached_enabled_clouds = global_user_state.get_cached_enabled_clouds(
234
+ capability)
264
235
  if raise_if_no_cloud_access and not cached_enabled_clouds:
265
236
  with ux_utils.print_exception_no_traceback():
266
237
  raise exceptions.NoCloudAccessError(
@@ -269,41 +240,6 @@ def get_cached_enabled_clouds_or_refresh(
269
240
  return cached_enabled_clouds
270
241
 
271
242
 
272
- def get_cached_enabled_storage_clouds_or_refresh(
273
- raise_if_no_cloud_access: bool = False) -> List[sky_clouds.Cloud]:
274
- """Returns cached enabled storage clouds and if no cloud is enabled,
275
- refresh.
276
-
277
- This function will perform a refresh if no public cloud is enabled.
278
-
279
- Args:
280
- raise_if_no_cloud_access: if True, raise an exception if no public
281
- cloud is enabled.
282
-
283
- Raises:
284
- exceptions.NoCloudAccessError: if no public cloud is enabled and
285
- raise_if_no_cloud_access is set to True.
286
- """
287
- cached_enabled_storage_clouds = (
288
- global_user_state.get_cached_enabled_storage_clouds())
289
- if not cached_enabled_storage_clouds:
290
- try:
291
- check(quiet=True, capability=CloudCapability.STORAGE)
292
- except SystemExit:
293
- # If no cloud is enabled, check() will raise SystemExit.
294
- # Here we catch it and raise the exception later only if
295
- # raise_if_no_cloud_access is set to True.
296
- pass
297
- cached_enabled_storage_clouds = (
298
- global_user_state.get_cached_enabled_storage_clouds())
299
- if raise_if_no_cloud_access and not cached_enabled_storage_clouds:
300
- with ux_utils.print_exception_no_traceback():
301
- raise exceptions.NoCloudAccessError(
302
- 'Cloud access is not set up. Run: '
303
- f'{colorama.Style.BRIGHT}sky check{colorama.Style.RESET_ALL}')
304
- return cached_enabled_storage_clouds
305
-
306
-
307
243
  def get_cloud_credential_file_mounts(
308
244
  excluded_clouds: Optional[Iterable[sky_clouds.Cloud]]
309
245
  ) -> Dict[str, str]:
@@ -336,16 +272,80 @@ def get_cloud_credential_file_mounts(
336
272
  return file_mounts
337
273
 
338
274
 
339
- def _format_enabled_cloud(cloud_name: str) -> str:
275
+ def _print_checked_cloud(
276
+ echo: Callable,
277
+ verbose: bool,
278
+ cloud_tuple: Tuple[str, Union[sky_clouds.Cloud, ModuleType]],
279
+ cloud_capabilities: List[Tuple[sky_cloud.CloudCapability, bool,
280
+ Optional[str]]],
281
+ ) -> None:
282
+ """Prints whether a cloud is enabled, and the capabilities that are enabled.
283
+ If any hints (for enabled capabilities) or
284
+ reasons (for disabled capabilities) are provided, they will be printed.
285
+
286
+ Args:
287
+ echo: The function to use to print the message.
288
+ verbose: Whether to print the verbose output.
289
+ cloud_tuple: The cloud to print the capabilities for.
290
+ cloud_capabilities: The capabilities for the cloud.
291
+ """
292
+ cloud_repr, cloud = cloud_tuple
293
+ # Print the capabilities for the cloud.
294
+ # consider cloud enabled if any capability is enabled.
295
+ enabled_capabilities: List[sky_cloud.CloudCapability] = []
296
+ hints_to_capabilities: Dict[str, List[sky_cloud.CloudCapability]] = {}
297
+ reasons_to_capabilities: Dict[str, List[sky_cloud.CloudCapability]] = {}
298
+ for capability, ok, reason in cloud_capabilities:
299
+ if ok:
300
+ enabled_capabilities.append(capability)
301
+ if reason is not None:
302
+ hints_to_capabilities.setdefault(reason, []).append(capability)
303
+ elif reason is not None:
304
+ reasons_to_capabilities.setdefault(reason, []).append(capability)
305
+ status_msg: str = 'disabled'
306
+ styles: Dict[str, Any] = {'dim': True}
307
+ capability_string: str = ''
308
+ activated_account: Optional[str] = None
309
+ if enabled_capabilities:
310
+ status_msg = 'enabled'
311
+ styles = {'fg': 'green', 'bold': False}
312
+ capability_string = f'[{", ".join(enabled_capabilities)}]'
313
+ if verbose and cloud is not cloudflare:
314
+ activated_account = cloud.get_active_user_identity_str()
315
+
316
+ echo(
317
+ click.style(f' {cloud_repr}: {status_msg} {capability_string}',
318
+ **styles))
319
+ if activated_account is not None:
320
+ echo(f' Activated account: {activated_account}')
321
+ for reason, caps in hints_to_capabilities.items():
322
+ echo(f' Hint [{", ".join(caps)}]: {reason}')
323
+ for reason, caps in reasons_to_capabilities.items():
324
+ echo(f' Reason [{", ".join(caps)}]: {reason}')
325
+
326
+
327
+ def _format_enabled_cloud(cloud_name: str,
328
+ capabilities: List[sky_cloud.CloudCapability]) -> str:
329
+ """Format the summary of enabled cloud and its enabled capabilities.
330
+
331
+ Args:
332
+ cloud_name: The name of the cloud.
333
+ capabilities: The capabilities of the cloud.
334
+
335
+ Returns:
336
+ A string of the formatted cloud and capabilities.
337
+ """
338
+ cloud_and_capabilities = f'{cloud_name} [{", ".join(capabilities)}]'
340
339
 
341
- def _green_color(cloud_name: str) -> str:
342
- return f'{colorama.Fore.GREEN}{cloud_name}{colorama.Style.RESET_ALL}'
340
+ def _green_color(str_to_format: str) -> str:
341
+ return (
342
+ f'{colorama.Fore.GREEN}{str_to_format}{colorama.Style.RESET_ALL}')
343
343
 
344
344
  if cloud_name == repr(sky_clouds.Kubernetes()):
345
345
  # Get enabled contexts for Kubernetes
346
346
  existing_contexts = sky_clouds.Kubernetes.existing_allowed_contexts()
347
347
  if not existing_contexts:
348
- return _green_color(cloud_name)
348
+ return _green_color(cloud_and_capabilities)
349
349
 
350
350
  # Check if allowed_contexts is explicitly set in config
351
351
  allowed_contexts = skypilot_config.get_nested(
@@ -363,7 +363,7 @@ def _format_enabled_cloud(cloud_name: str) -> str:
363
363
  else:
364
364
  context_info = f'Active context: {existing_contexts[0]}'
365
365
 
366
- return (f'{_green_color(cloud_name)}\n'
366
+ return (f'{_green_color(cloud_and_capabilities)}\n'
367
367
  f' {colorama.Style.DIM}{context_info}'
368
368
  f'{colorama.Style.RESET_ALL}')
369
- return _green_color(cloud_name)
369
+ return _green_color(cloud_and_capabilities)