Repository: skypilot-org/skypilot
Pull request: [Core] Support choosing cloud for Spot controller
<!-- Describe the changes in this PR -->

Addresses the new comment under #3231.

<!-- 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):

- [x] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
     - [x]  `pytest tests/unit_tests/test_controller_utils.py` with #3457
     - [x] The following YAML and check cloud resources is in `{AWS, GCP}`/`{AWS, GCP, Azure}`

```yaml
service:
  readiness_probe: /
  replicas: 1

resources:
  ports: 8000
  any_of:
    - cloud: aws
    - cloud: gcp
    - cloud: kubernetes

run: python -m http.server 8000
```

```yaml
resources:
  any_of:
    - cloud: aws
    - cloud: gcp
run: echo "Hello, World!"
---
resources:
  any_of:
    - cloud: aws
    - cloud: runpod
run: echo "Hello, World!"
---
resources:
  any_of:
    - cloud: azure
    - cloud: gcp
    - cloud: kubernetes
run: echo "Hello, World!"
```

- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Serve] Fail to launch controller when no resources is specified for a service when firstly used
<!-- Describe the bug report / feature request here -->

This issue affects all first-time SkyServe user who don't have serve controller started before

To reproduce:
1. with no serve controller
2. start the following service

```yaml
resources:
  accelerators: L4
  ports: 8081

run: |
  echo hi
```

`sky serve up -n vllm task.yaml`


The command will fail with the following error:
```
Launching controller for 'vllm'...
sky.exceptions.ResourcesUnavailableError: Catalog does not contain any instances satisfying the request:
Task<name=vllm>(run='# Start sky serve se...')
  resources: default instances.

To fix: relax or change the resource requirements.

Hint: sky show-gpus to list available accelerators.
      sky check to check the enabled clouds.
```

Hints from issue discussion:

Commit message: fix unittest

Patch to review:
diff --git a/sky/check.py b/sky/check.py
index 7e129e4fde5..70d561fba15 100644
--- a/sky/check.py
+++ b/sky/check.py
@@ -124,7 +124,7 @@ def get_cloud_credential_file_mounts(
     file_mounts = {}
     for cloud in enabled_clouds:
         if (excluded_clouds is not None and
-                clouds.cloud_in_list(cloud, excluded_clouds)):
+                clouds.cloud_in_iterable(cloud, excluded_clouds)):
             continue
         cloud_file_mounts = cloud.get_credential_file_mounts()
         file_mounts.update(cloud_file_mounts)
diff --git a/sky/clouds/__init__.py b/sky/clouds/__init__.py
index a00d04e8688..f526c29749b 100644
--- a/sky/clouds/__init__.py
+++ b/sky/clouds/__init__.py
@@ -1,6 +1,6 @@
 """Clouds in Sky."""
 from sky.clouds.cloud import Cloud
-from sky.clouds.cloud import cloud_in_list
+from sky.clouds.cloud import cloud_in_iterable
 from sky.clouds.cloud import CloudImplementationFeatures
 from sky.clouds.cloud import ProvisionerVersion
 from sky.clouds.cloud import Region
@@ -44,5 +44,5 @@
     'StatusVersion',
     'Fluidstack',
     # Utility functions
-    'cloud_in_list',
+    'cloud_in_iterable',
 ]
diff --git a/sky/clouds/cloud.py b/sky/clouds/cloud.py
index ed1ff0235bb..5322a906481 100644
--- a/sky/clouds/cloud.py
+++ b/sky/clouds/cloud.py
@@ -744,6 +744,6 @@ def __getstate__(self):
 
 
 # === Helper functions ===
-def cloud_in_list(cloud: Cloud, cloud_list: Iterable[Cloud]) -> bool:
+def cloud_in_iterable(cloud: Cloud, cloud_list: Iterable[Cloud]) -> bool:
     """Returns whether the cloud is in the given cloud list."""
     return any(cloud.is_same_cloud(c) for c in cloud_list)
diff --git a/sky/execution.py b/sky/execution.py
index 7b83a57077f..cc2102fef59 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -666,9 +666,15 @@ def spot_launch(
         prefix = spot.SPOT_TASK_YAML_PREFIX
         remote_user_yaml_path = f'{prefix}/{dag.name}-{dag_uuid}.yaml'
         remote_user_config_path = f'{prefix}/{dag.name}-{dag_uuid}.config_yaml'
-        controller_resources = (controller_utils.get_controller_resources(
+        controller_resources = controller_utils.get_controller_resources(
             controller_type='spot',
-            controller_resources_config=spot.constants.CONTROLLER_RESOURCES))
+            controller_name=controller_name,
+            controller_resources_config=spot.constants.CONTROLLER_RESOURCES,
+            task_resources=sum(
+                [list(task_.resources) for task_ in dag.tasks],
+                [],
+            ),
+        )
 
         vars_to_fill = {
             'remote_user_yaml_path': remote_user_yaml_path,
diff --git a/sky/optimizer.py b/sky/optimizer.py
index 6b007266d59..e034e7c5e16 100644
--- a/sky/optimizer.py
+++ b/sky/optimizer.py
@@ -334,7 +334,8 @@ def _estimate_nodes_cost_or_time(
                 # in the error message.
                 enabled_clouds = (
                     sky_check.get_cached_enabled_clouds_or_refresh())
-                if clouds.cloud_in_list(clouds.Kubernetes(), enabled_clouds):
+                if clouds.cloud_in_iterable(clouds.Kubernetes(),
+                                            enabled_clouds):
                     if any(orig_resources.cloud is None
                            for orig_resources in node.resources):
                         source_hint = 'catalog and kubernetes cluster'
@@ -1166,7 +1167,7 @@ def _fill_in_launchable_resources(
         blocked_resources = []
     for resources in task.resources:
         if (resources.cloud is not None and
-                not clouds.cloud_in_list(resources.cloud, enabled_clouds)):
+                not clouds.cloud_in_iterable(resources.cloud, enabled_clouds)):
             if try_fix_with_sky_check:
                 # Explicitly check again to update the enabled cloud list.
                 sky_check.check(quiet=True)
diff --git a/sky/serve/core.py b/sky/serve/core.py
index 1ebfaeba067..7d5fe719b50 100644
--- a/sky/serve/core.py
+++ b/sky/serve/core.py
@@ -1,15 +1,13 @@
 """SkyServe core APIs."""
 import re
 import tempfile
-import typing
-from typing import Any, Dict, List, Optional, Set, Union
+from typing import Any, Dict, List, Optional, Union
 
 import colorama
 
 import sky
 from sky import backends
 from sky import exceptions
-from sky import global_user_state
 from sky import sky_logging
 from sky import task as task_lib
 from sky.backends import backend_utils
@@ -25,9 +23,6 @@
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
-if typing.TYPE_CHECKING:
-    from sky import clouds
-
 logger = sky_logging.init_logger(__name__)
 
 
@@ -126,14 +121,6 @@ def up(
     controller_utils.maybe_translate_local_file_mounts_and_sync_up(task,
                                                                    path='serve')
 
-    # If the controller and replicas are from the same cloud, it should
-    # provide better connectivity. We will let the controller choose from
-    # the clouds of the resources if the controller does not exist.
-    requested_clouds: Set['clouds.Cloud'] = set()
-    for resources in task.resources:
-        if resources.cloud is not None:
-            requested_clouds.add(resources.cloud)
-
     with tempfile.NamedTemporaryFile(
             prefix=f'service-task-{service_name}-',
             mode='w',
@@ -150,11 +137,12 @@ def up(
             serve_utils.generate_remote_config_yaml_file_name(service_name))
         controller_log_file = (
             serve_utils.generate_remote_controller_log_file_name(service_name))
-        controller_resources_in_config = (
-            controller_utils.get_controller_resources(
-                controller_type='serve',
-                controller_resources_config=serve_constants.CONTROLLER_RESOURCES
-            ))
+        controller_resources = controller_utils.get_controller_resources(
+            controller_type='serve',
+            controller_name=controller_name,
+            controller_resources_config=serve_constants.CONTROLLER_RESOURCES,
+            task_resources=task.resources,
+        )
 
         vars_to_fill = {
             'remote_task_yaml_path': remote_tmp_task_yaml_path,
@@ -171,23 +159,13 @@ def up(
                                    vars_to_fill,
                                    output_path=controller_file.name)
         controller_task = task_lib.Task.from_yaml(controller_file.name)
-        controller_exist = (
-            global_user_state.get_cluster_from_name(controller_name)
-            is not None)
-        if (not controller_exist and
-                controller_resources_in_config.cloud is None):
-            controller_clouds = requested_clouds
-        else:
-            controller_clouds = {controller_resources_in_config.cloud}
         # TODO(tian): Probably run another sky.launch after we get the load
         # balancer port from the controller? So we don't need to open so many
         # ports here. Or, we should have a nginx traffic control to refuse
         # any connection to the unregistered ports.
         controller_resources = {
-            controller_resources_in_config.copy(
-                cloud=controller_cloud,
-                ports=[serve_constants.LOAD_BALANCER_PORT_RANGE])
-            for controller_cloud in controller_clouds
+            res.copy(ports=[serve_constants.LOAD_BALANCER_PORT_RANGE])
+            for res in controller_resources
         }
         controller_task.set_resources(controller_resources)
 
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
index d8ca2abbba9..7ed5a722e05 100644
--- a/sky/utils/controller_utils.py
+++ b/sky/utils/controller_utils.py
@@ -6,13 +6,14 @@
 import os
 import tempfile
 import typing
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Iterable, List, Optional, Set
 
 import colorama
 
 from sky import check as sky_check
 from sky import clouds
 from sky import exceptions
+from sky import global_user_state
 from sky import resources
 from sky import sky_logging
 from sky import skypilot_config
@@ -261,15 +262,18 @@ def shared_controller_vars_to_fill(
 
 def get_controller_resources(
     controller_type: str,
+    controller_name: str,
     controller_resources_config: Dict[str, Any],
-) -> 'resources.Resources':
+    task_resources: Iterable['resources.Resources'],
+) -> Set['resources.Resources']:
     """Read the skypilot config and setup the controller resources.
 
     Returns:
-        A tuple of (vars_to_fill, controller_resources_config). `var_to_fill`
-        is a dict of variables that will be filled in the controller template.
-        The controller_resources_config is the resources config that will be
-        used to launch the controller.
+        A set of controller resources that will be used to launch the
+        controller. All fields are the same except for the cloud. If no
+        controller exists and the controller resources has no cloud
+        specified, the controller will be launched on one of the clouds
+        of the task resources for better connectivity.
     """
     controller_resources_config_copied: Dict[str, Any] = copy.copy(
         controller_resources_config)
@@ -300,7 +304,32 @@ def get_controller_resources(
                     err=f'Expected exactly one resource, got '
                     f'{len(controller_resources)} resources: '
                     f'{controller_resources}'))
-    return list(controller_resources)[0]
+    controller_resources_to_use: resources.Resources = list(
+        controller_resources)[0]
+
+    controller_exist = (global_user_state.get_cluster_from_name(controller_name)
+                        is not None)
+    if controller_exist or controller_resources_to_use.cloud is not None:
+        return {controller_resources_to_use}
+
+    # If the controller and replicas are from the same cloud, it should
+    # provide better connectivity. We will let the controller choose from
+    # the clouds of the resources if the controller does not exist.
+    # TODO(tian): Consider respecting the regions/zones specified for the
+    # resources as well.
+    requested_clouds: Set['clouds.Cloud'] = set()
+    for res in task_resources:
+        # cloud is an object and will not be able to be distinguished by set.
+        # Here we manually check if the cloud is in the set.
+        if res.cloud is not None and not clouds.cloud_in_iterable(
+                res.cloud, requested_clouds):
+            requested_clouds.add(res.cloud)
+    if not requested_clouds:
+        return {controller_resources_to_use}
+    return {
+        controller_resources_to_use.copy(cloud=controller_cloud)
+        for controller_cloud in requested_clouds
+    }
 
 
 def _setup_proxy_command_on_controller(
diff --git a/tests/unit_tests/test_controller_utils.py b/tests/unit_tests/test_controller_utils.py
index 57490139f8b..673d208a642 100644
--- a/tests/unit_tests/test_controller_utils.py
+++ b/tests/unit_tests/test_controller_utils.py
@@ -52,8 +52,12 @@ def get_custom_controller_resources(keys, default):
     monkeypatch.setattr('sky.skypilot_config.get_nested',
                         get_custom_controller_resources)
 
-    controller_resources = controller_utils.get_controller_resources(
-        controller_type, controller_resources_config)
+    controller_resources = list(
+        controller_utils.get_controller_resources(
+            controller_type=controller_type,
+            controller_name='dummy',
+            controller_resources_config=controller_resources_config,
+            task_resources=[]))[0]
     controller_resources_config = controller_resources.to_yaml_config()
     for k, v in expected.items():
         assert controller_resources_config[k] == v, (


Code review comments:
diff_hunk: @@ -150,11 +137,12 @@ def up(
             serve_utils.generate_remote_config_yaml_file_name(service_name))
         controller_log_file = (
             serve_utils.generate_remote_controller_log_file_name(service_name))
-        controller_resources_in_config = (
-            controller_utils.get_controller_resources(
-                controller_type='serve',
-                controller_resources_config=serve_constants.CONTROLLER_RESOURCES
-            ))
+        controller_resources = controller_utils.get_controller_resources(
+            controller_type='serve',
+            controller_name=controller_name,
+            controller_resources_config=serve_constants.CONTROLLER_RESOURCES,
path: sky/serve/core.py
text: @user1:
Similarly, should this `controller_resources_config` be inferred by the `controller_type` in the `controller_utils`?

diff_hunk: @@ -150,11 +137,12 @@ def up(
             serve_utils.generate_remote_config_yaml_file_name(service_name))
         controller_log_file = (
             serve_utils.generate_remote_controller_log_file_name(service_name))
-        controller_resources_in_config = (
-            controller_utils.get_controller_resources(
-                controller_type='serve',
-                controller_resources_config=serve_constants.CONTROLLER_RESOURCES
-            ))
+        controller_resources = controller_utils.get_controller_resources(
+            controller_type='serve',
+            controller_name=controller_name,
path: sky/serve/core.py
text: @user1:
nit: Should the `controller_name` be inferred from the `controller_type`, so as to reduce the argument list?

@author:
Done. Thanks!