[core] Add allowed_clouds to config to check only specific clouds (#…

…3556)

* candidate_clouds

* Working allowed_clouds

* Working allowed_clouds

* comments

* lint

* change skipped clouds to disabled clouds

* lint

docs/source/reference/config.rst

@@ -27,6 +27,19 @@ Available fields and semantics:
  cpus: 4+ # number of vCPUs, max concurrent spot jobs = 2 * cpus
  disk_size: 100
 
+ # Allow list for clouds to be used in `sky check`
+ #
+ # This field is used to restrict the clouds that SkyPilot will check and use
+ # when running `sky check`. Any cloud already enabled but not specified here
+ # will be disabled on the next `sky check` run.
+ # If this field is not set, SkyPilot will check and use all supported clouds.
+ #
+ # Default: null (use all supported clouds).
+ allowed_clouds:
+ - aws
+ - gcp
+ - kubernetes
+
  # Advanced AWS configurations (optional).
  # Apply to all new instances but not existing ones.
  aws:

sky/adaptors/cloudflare.py

@@ -23,6 +23,7 @@
 R2_PROFILE_NAME = 'r2'
 _INDENT_PREFIX = ' '
 NAME = 'Cloudflare'
+SKY_CHECK_NAME = 'Cloudflare (for R2 object store)'
 
 
 @contextlib.contextmanager

sky/check.py

@@ -1,7 +1,7 @@
 """Credential checks: check cloud credentials and enable clouds."""
 import traceback
 from types import ModuleType
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Dict, Iterable, List, Optional, Tuple, Union
 
 import click
 import colorama
@@ -10,6 +10,7 @@
 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 ux_utils
 
@@ -52,20 +53,42 @@ def check_one_cloud(
  disabled_clouds.append(cloud_repr)
  echo(f' Reason: {reason}')
 
+ def get_cloud_tuple(
+ cloud_name: str) -> Tuple[str, Union[sky_clouds.Cloud, ModuleType]]:
+ # Validates cloud_name and returns a tuple of the cloud's name and
+ # the cloud object. Includes special handling for Cloudflare.
+ if cloud_name.lower().startswith('cloudflare'):
+ return cloudflare.SKY_CHECK_NAME, cloudflare
+ else:
+ cloud_obj = sky_clouds.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()] +
+ [cloudflare.SKY_CHECK_NAME])
+
  if clouds is not None:
- clouds_to_check: List[Tuple[str, Any]] = []
- for cloud in clouds:
- if cloud.lower() == 'cloudflare':
- clouds_to_check.append(
- ('Cloudflare, for R2 object store', cloudflare))
- else:
- cloud_obj = sky_clouds.CLOUD_REGISTRY.from_str(cloud)
- assert cloud_obj is not None, f'Cloud {cloud!r} not found'
- clouds_to_check.append((repr(cloud_obj), cloud_obj))
+ cloud_list = clouds
  else:
- clouds_to_check = [(repr(cloud_obj), cloud_obj)
- for cloud_obj in sky_clouds.CLOUD_REGISTRY.values()]
- clouds_to_check.append(('Cloudflare, for R2 object store', cloudflare))
+ cloud_list = get_all_clouds()
+ clouds_to_check = [get_cloud_tuple(c) for c in cloud_list]
+
+ # 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())
+ ]
+ # Use disallowed_cloud_names for logging the clouds that will be disabled
+ # because they are not included in allowed_clouds in config.yaml.
+ disallowed_cloud_names = [
+ c for c in get_all_clouds() if c not in config_allowed_cloud_names
+ ]
+ # Check only the clouds which are allowed in the config.
+ clouds_to_check = [
+ c for c in clouds_to_check if c[0] in config_allowed_cloud_names
+ ]
 
  for cloud_tuple in sorted(clouds_to_check):
  check_one_cloud(cloud_tuple)
@@ -79,23 +102,39 @@ def check_one_cloud(
  disabled_clouds_set = {
  cloud for cloud in disabled_clouds if not cloud.startswith('Cloudflare')
  }
+ config_allowed_clouds_set = {
+ cloud for cloud in config_allowed_cloud_names
+ if not cloud.startswith('Cloudflare')
+ }
  previously_enabled_clouds_set = {
  repr(cloud) for cloud in global_user_state.get_cached_enabled_clouds()
  }
 
- # Determine the set of enabled clouds: previously enabled clouds + newly
- # enabled clouds - newly disabled clouds.
- all_enabled_clouds = ((previously_enabled_clouds_set | enabled_clouds_set) -
- disabled_clouds_set)
+ # Determine the set of enabled clouds: (previously enabled clouds + newly
+ # enabled clouds - newly disabled clouds) intersected with
+ # config_allowed_clouds, if specified in config.yaml.
+ # This means that if a cloud is already enabled and is not included in
+ # allowed_clouds in config.yaml, it will be disabled.
+ all_enabled_clouds = (config_allowed_clouds_set & (
+ (previously_enabled_clouds_set | enabled_clouds_set) -
+ disabled_clouds_set))
  global_user_state.set_enabled_clouds(list(all_enabled_clouds))
 
+ disallowed_clouds_hint = None
+ if disallowed_cloud_names:
+ disallowed_clouds_hint = (
+ '\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:
  echo(
  click.style(
  'No cloud is enabled. SkyPilot will not be able to run any '
  'task. Run `sky check` for more info.',
  fg='red',
  bold=True))
+ if disallowed_clouds_hint:
+ echo(click.style(disallowed_clouds_hint, dim=True))
  raise SystemExit()
  else:
  clouds_arg = (' ' +
@@ -109,6 +148,9 @@ def check_one_cloud(
  'https://skypilot.readthedocs.io/en/latest/getting-started/installation.html', # pylint: disable=line-too-long
  dim=True))
 
+ if disallowed_clouds_hint:
+ echo(click.style(disallowed_clouds_hint, dim=True))
+
  # Pretty print for UX.
  if not quiet:
  enabled_clouds_str = '\n :heavy_check_mark: '.join(

sky/utils/schemas.py

@@ -587,6 +587,7 @@ def get_default_remote_identity(cloud: str) -> str:
 
 def get_config_schema():
  # pylint: disable=import-outside-toplevel
+ from sky.clouds import service_catalog
  from sky.utils import kubernetes_enums
 
  resources_schema = {
@@ -722,6 +723,16 @@ def get_config_schema():
  },
  }
 
+ allowed_clouds = {
+ # A list of cloud names that are allowed to be used
+ 'type': 'array',
+ 'items': {
+ 'type': 'string',
+ 'case_insensitive_enum':
+ (list(service_catalog.ALL_CLOUDS) + ['cloudflare'])
+ }
+ }
+
  for cloud, config in cloud_configs.items():
  if cloud == 'aws':
  config['properties'].update(_REMOTE_IDENTITY_SCHEMA_AWS)
@@ -738,6 +749,7 @@ def get_config_schema():
  'jobs': controller_resources_schema,
  'spot': controller_resources_schema,
  'serve': controller_resources_schema,
+ 'allowed_clouds': allowed_clouds,
  **cloud_configs,
  },
  # Avoid spot and jobs being present at the same time.