From fd44999fe01f424f4c6ed27d22d27b7b0d2b2b38 Mon Sep 17 00:00:00 2001
From: Romil Bhardwaj <romil.bhardwaj@gmail.com>
Date: Thu, 16 May 2024 20:13:11 -0700
Subject: [PATCH] Working allowed_clouds

---
 docs/source/reference/config.rst | 13 +++++++++++++
 sky/check.py                     | 16 ++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
index 53c983edfad..dce0ce1f643 100644
--- a/docs/source/reference/config.rst
+++ b/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:
diff --git a/sky/check.py b/sky/check.py
index fc506562748..64cde9a31f1 100644
--- a/sky/check.py
+++ b/sky/check.py
@@ -100,14 +100,22 @@ def get_all_clouds():
     disabled_clouds_set = {
         cloud for cloud in disabled_clouds if not cloud.startswith('Cloudflare')
     }
+    config_allowed_clouds_set = {
+        cloud for cloud in config_allowed_clouds
+        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))
 
     skipped_clouds_hint = None