Repository: skypilot-org/skypilot
Pull request #5536 (2025-05-09 05:17:14)

[WIP] workspaces implementation
<!-- Describe the changes in this PR -->
Sample server-side config:
```
workspaces:
  my_workspace:
    gcp:
      project_id: <project>
      disabled: true
    aws:
      disabled: true
    fluidstack:
      disabled: true
```

Sample client-side config:
```
...
active_workspace: my_workspace
```

If no workspace is specified, `default` workspace is used.

Enabled clouds are now cached per workspace.

If a non-default workspace is specified and the corresponding workspace does not exist in the server, request errors out.

Each workspace can disable a cloud specified in `allowed_clouds` using `disabled` field.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->



Patch:
diff --git a/sky/check.py b/sky/check.py
index 65a3f9236..730f33550 100644
--- a/sky/check.py
+++ b/sky/check.py
@@ -82,6 +82,19 @@ def check_capabilities(
         get_cloud_tuple(c)[0] for c in skypilot_config.get_nested((
             'allowed_clouds',), get_all_clouds())
     ])
+
+    # filter out the clouds that are disabled in the workspace config
+    workspace_disabled_clouds = []
+    for cloud in config_allowed_cloud_names:
+        cloud_config = skypilot_config.get_workspace_cloud(cloud)
+        cloud_disabled = cloud_config.get('disabled', False)
+        if cloud_disabled:
+            workspace_disabled_clouds.append(cloud)
+    config_allowed_cloud_names = [
+        c for c in config_allowed_cloud_names
+        if c not in workspace_disabled_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 = [
@@ -141,13 +154,15 @@ def check_capabilities(
         }
         previously_enabled_clouds_set = {
             repr(cloud)
-            for cloud in global_user_state.get_cached_enabled_clouds(capability)
+            for cloud in global_user_state.get_cached_enabled_clouds(
+                capability, skypilot_config.get_active_workspace())
         }
         enabled_clouds_for_capability = (config_allowed_clouds_set & (
             (previously_enabled_clouds_set | enabled_clouds_set) -
             disabled_clouds_set))
         global_user_state.set_enabled_clouds(
-            list(enabled_clouds_for_capability), capability)
+            list(enabled_clouds_for_capability), capability,
+            skypilot_config.get_active_workspace())
         all_enabled_clouds = all_enabled_clouds.union(
             enabled_clouds_for_capability)
     disallowed_clouds_hint = None
@@ -234,7 +249,7 @@ def get_cached_enabled_clouds_or_refresh(
             raise_if_no_cloud_access is set to True.
     """
     cached_enabled_clouds = global_user_state.get_cached_enabled_clouds(
-        capability)
+        capability, skypilot_config.get_active_workspace())
     if not cached_enabled_clouds:
         try:
             check_capability(sky_cloud.CloudCapability.COMPUTE, quiet=True)
@@ -244,7 +259,7 @@ def get_cached_enabled_clouds_or_refresh(
             # raise_if_no_cloud_access is set to True.
             pass
         cached_enabled_clouds = global_user_state.get_cached_enabled_clouds(
-            capability)
+            capability, skypilot_config.get_active_workspace())
     if raise_if_no_cloud_access and not cached_enabled_clouds:
         with ux_utils.print_exception_no_traceback():
             raise exceptions.NoCloudAccessError(
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 1cfa23826..8fd79156e 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -999,6 +999,10 @@ class GCP(clouds.Cloud):
             return 'dryrun-project-id'
         # pylint: disable=import-outside-toplevel
         from google import auth  # type: ignore
+        config_project_id = skypilot_config.get_workspace_cloud('gcp').get(
+            'project_id', None)
+        if config_project_id:
+            return config_project_id
         _, project_id = auth.default()
         if project_id is None:
             raise exceptions.CloudUserIdentityError(
diff --git a/sky/core.py b/sky/core.py
index 2e5e4900c..c1e8b3748 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -17,6 +17,7 @@ from sky import global_user_state
 from sky import models
 from sky import optimizer
 from sky import sky_logging
+from sky import skypilot_config
 from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.clouds import cloud as sky_cloud
@@ -1004,7 +1005,8 @@ def storage_delete(name: str) -> None:
 @usage_lib.entrypoint
 def enabled_clouds() -> List[clouds.Cloud]:
     return global_user_state.get_cached_enabled_clouds(
-        sky_cloud.CloudCapability.COMPUTE)
+        sky_cloud.CloudCapability.COMPUTE,
+        workspace=skypilot_config.get_active_workspace())
 
 
 @usage_lib.entrypoint
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index e1c9ce21d..66635aa7f 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -802,11 +802,12 @@ def get_cluster_names_start_with(starts_with: str) -> List[str]:
     return [row[0] for row in rows]
 
 
-def get_cached_enabled_clouds(
-        cloud_capability: 'cloud.CloudCapability') -> List['clouds.Cloud']:
+def get_cached_enabled_clouds(cloud_capability: 'cloud.CloudCapability',
+                              workspace: str) -> List['clouds.Cloud']:
 
-    rows = _DB.cursor.execute('SELECT value FROM config WHERE key = ?',
-                              (_get_capability_key(cloud_capability),))
+    rows = _DB.cursor.execute(
+        'SELECT value FROM config WHERE key = ?',
+        (_get_enabled_clouds_key(cloud_capability, workspace),))
     ret = []
     for (value,) in rows:
         ret = json.loads(value)
@@ -827,15 +828,17 @@ def get_cached_enabled_clouds(
 
 
 def set_enabled_clouds(enabled_clouds: List[str],
-                       cloud_capability: 'cloud.CloudCapability') -> None:
-    _DB.cursor.execute(
-        'INSERT OR REPLACE INTO config VALUES (?, ?)',
-        (_get_capability_key(cloud_capability), json.dumps(enabled_clouds)))
+                       cloud_capability: 'cloud.CloudCapability',
+                       workspace: str) -> None:
+    _DB.cursor.execute('INSERT OR REPLACE INTO config VALUES (?, ?)',
+                       (_get_enabled_clouds_key(cloud_capability, workspace),
+                        json.dumps(enabled_clouds)))
     _DB.conn.commit()
 
 
-def _get_capability_key(cloud_capability: 'cloud.CloudCapability') -> str:
-    return _ENABLED_CLOUDS_KEY_PREFIX + cloud_capability.value
+def _get_enabled_clouds_key(cloud_capability: 'cloud.CloudCapability',
+                            workspace: str) -> str:
+    return _ENABLED_CLOUDS_KEY_PREFIX + workspace + '_' + cloud_capability.value
 
 
 def add_or_update_storage(storage_name: str,
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
index 9514c8625..3a9952184 100644
--- a/sky/skylet/constants.py
+++ b/sky/skylet/constants.py
@@ -375,7 +375,8 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
 # we skip the following keys because they are meant to be client-side configs.
 SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('admin_policy',),
                                                        ('api_server',),
-                                                       ('allowed_clouds',)]
+                                                       ('allowed_clouds',),
+                                                       ('workspaces',)]
 
 # Constants for Azure blob storage
 WAIT_FOR_STORAGE_ACCOUNT_CREATION = 60
diff --git a/sky/skypilot_config.py b/sky/skypilot_config.py
index 80f02eca3..f2d816e28 100644
--- a/sky/skypilot_config.py
+++ b/sky/skypilot_config.py
@@ -139,8 +139,7 @@ def get_user_config() -> config_utils.Config:
 
     # load the user config file
     if os.path.exists(user_config_path):
-        user_config = parse_config_file(user_config_path)
-        _validate_config(user_config, user_config_path)
+        user_config = parse_and_validate_config_file(user_config_path)
     else:
         user_config = config_utils.Config()
     return user_config
@@ -168,8 +167,7 @@ def _get_project_config() -> config_utils.Config:
 
     # load the project config file
     if os.path.exists(project_config_path):
-        project_config = parse_config_file(project_config_path)
-        _validate_config(project_config, project_config_path)
+        project_config = parse_and_validate_config_file(project_config_path)
     else:
         project_config = config_utils.Config()
     return project_config
@@ -197,8 +195,7 @@ def get_server_config() -> config_utils.Config:
 
     # load the server config file
     if os.path.exists(server_config_path):
-        server_config = parse_config_file(server_config_path)
-        _validate_config(server_config, server_config_path)
+        server_config = parse_and_validate_config_file(server_config_path)
     else:
         server_config = config_utils.Config()
     return server_config
@@ -232,6 +229,25 @@ def get_nested(keys: Tuple[str, ...],
         disallowed_override_keys=None)
 
 
+def get_workspace_cloud(cloud: str) -> config_utils.Config:
+    """Returns the workspace config."""
+    current_workspace = get_active_workspace()
+    if current_workspace is None:
+        return config_utils.Config()
+    clouds = _dict.get_nested(keys=(
+        'workspaces',
+        current_workspace,
+    ),
+                              default_value=None)
+    if clouds is None:
+        return config_utils.Config()
+    return clouds.get(cloud.lower(), config_utils.Config())
+
+
+def get_active_workspace() -> str:
+    return _dict.get_nested(keys=('active_workspace',), default_value='default')
+
+
 def set_nested(keys: Tuple[str, ...], value: Any) -> Dict[str, Any]:
     """Returns a deep-copied config with the nested key set to value.
 
@@ -302,7 +318,7 @@ def _reload_config() -> None:
         _reload_config_as_client()
 
 
-def parse_config_file(config_path: str) -> config_utils.Config:
+def parse_and_validate_config_file(config_path: str) -> config_utils.Config:
     config = config_utils.Config()
     try:
         config_dict = common_utils.read_yaml(config_path)
@@ -359,7 +375,7 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
                 'exist. Please double check the path or unset the env var: '
                 f'unset {ENV_VAR_SKYPILOT_CONFIG}')
     logger.debug(f'Using config path: {config_path}')
-    _dict = parse_config_file(config_path)
+    _dict = parse_and_validate_config_file(config_path)
     _loaded_config_path = config_path
 
 
@@ -461,6 +477,12 @@ def override_skypilot_config(
         override_configs=dict(override_configs),
         allowed_override_keys=None,
         disallowed_override_keys=constants.SKIPPED_CLIENT_OVERRIDE_KEYS)
+    workspace = config.get_nested(keys=('active_workspace',),
+                                  default_value=None)
+    if workspace and workspace not in get_nested(keys=('workspaces',),
+                                                 default_value={}):
+        raise ValueError(f'Workspace {workspace} does not exist. '
+                         'Please check your config and try again.')
     try:
         common_utils.validate_schema(
             config,
@@ -506,7 +528,7 @@ def _compose_cli_config(cli_config: Optional[List[str]]) -> config_utils.Config:
                     'Cannot use multiple --config flags with a config file.')
             config_source = maybe_config_path
             # cli_config is a path to a config file
-            parsed_config = parse_config_file(maybe_config_path)
+            parsed_config = parse_and_validate_config_file(maybe_config_path)
         else:  # cli_config is a comma-separated list of key-value pairs
             parsed_config = _parse_dotlist(cli_config)
         _validate_config(parsed_config, config_source)
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 3b201917e..60b08d5e7 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -1041,6 +1041,43 @@ def get_config_schema():
         }
     }
 
+    workspace_schema = {'type': 'string'}
+
+    allowed_workspace_names = list(service_catalog.ALL_CLOUDS) + ['cloudflare']
+    allowed_workspace_regex = '|'.join(allowed_workspace_names)
+    workspaces_schema = {
+        'type': 'object',
+        'required': [],
+        # each key is a workspace name
+        'additionalProperties': {
+            'type': 'object',
+            'additionalProperties': False,
+            'patternProperties': {
+                allowed_workspace_regex: {
+                    'type': 'object',
+                    'properties': {
+                        'disabled': {
+                            'type': 'boolean'
+                        }
+                    },
+                },
+            },
+            'properties': {
+                'gcp': {
+                    'type': 'object',
+                    'properties': {
+                        'project_id': {
+                            'type': 'string'
+                        },
+                        'disabled': {
+                            'type': 'boolean'
+                        }
+                    },
+                },
+            },
+        },
+    }
+
     for cloud, config in cloud_configs.items():
         if cloud == 'aws':
             config['properties'].update({
@@ -1063,6 +1100,8 @@ def get_config_schema():
             'docker': docker_configs,
             'nvidia_gpus': gpu_configs,
             'api_server': api_server,
+            'active_workspace': workspace_schema,
+            'workspaces': workspaces_schema,
             **cloud_configs,
         },
     }


Test patch:
diff --git a/tests/test_global_user_state.py b/tests/test_global_user_state.py
index 8107c9d3c..7189e02a3 100644
--- a/tests/test_global_user_state.py
+++ b/tests/test_global_user_state.py
@@ -9,4 +9,4 @@ import sky
 def test_enabled_clouds_empty():
     # In test environment, no cloud should be enabled.
     assert sky.global_user_state.get_cached_enabled_clouds(
-        sky.clouds.cloud.CloudCapability.COMPUTE) == []
+        sky.clouds.cloud.CloudCapability.COMPUTE, workspace='default') == []