Repository: skypilot-org/skypilot (python)

Issue:
[Feature] Autostop definition in YAML?
Is it possible to add an `autostop` field in the YAML for creating a cluster, please? I do not see any such way in https://github.com/skypilot-org/skypilot/blob/master/docs/source/reference/yaml-spec.rst 

The current way to implement autostop seems to require running a command after-launch. Guess who forgot to run it a month ago and just got hit by a multi-hundred dollar bill from EC2? 💸 💸 💸 

Adding `autostop` in the cluster definition would save me (and probably others) from shooting themselves in the foot again. Thanks !


Pull request description:
allow specifying autostop in resources
Fixes #3953.
<!-- Describe the changes in this PR -->
Add an `autostop` field to the `resources` spec in the SkyPilot YAML, much like the config added in #5182.

Examples:
```
resources:
  # Enable default autostop config
  autostop: true
```
```
resources:
  # Shorthand to set autostop timeout
  autostop: 10
```
```
resources:
  # Set autodown with a custom time
  autostop:
    down: true
    idle_minutes: 3
```
```
resources:
  # Force disable autostop
  autostop: false
```

Issues:
- multiple different autostop configs for different resources
- sky start will forget autostop config

<!-- 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/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 8ce9691c4..48eb2ac7e 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -4504,6 +4504,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
 
                     # down = False is the default, but warn the user in case
                     # they have explicitly specified it.
+                    # TODO(cooperc): Fix for new autostop stuff.
                     config_override_down = skypilot_config.get_nested(
                         (controller.value.controller_type, 'controller',
                          'autostop', 'down'), None)
diff --git a/sky/core.py b/sky/core.py
index 356e66e5c..c74b1264d 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -369,9 +369,18 @@ def _start(
                 'supported when starting SkyPilot controllers. To '
                 f'fix: omit the {arguments_str} to use the '
                 f'default autostop settings from config.')
-        idle_minutes_to_autostop, down = (
-            controller_utils.get_controller_autostop_config(
-                controller=controller))
+
+        # Get the autostop resources, from which we extract the correct autostop
+        # config.
+        controller_resources = controller_utils.get_controller_resources(
+            controller, [])
+        # All resources should have the same autostop config.
+        controller_autostop_config = list(
+            controller_resources)[0].autostop_config
+        if (controller_autostop_config is not None and
+                controller_autostop_config.enabled):
+            idle_minutes_to_autostop = controller_autostop_config.idle_minutes
+            down = controller_autostop_config.down
 
     usage_lib.record_cluster_name_for_current_operation(cluster_name)
 
diff --git a/sky/execution.py b/sky/execution.py
index 0f9a0cf23..9d42ac116 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -226,11 +226,44 @@ def _execute(
     requested_features |= task.get_required_cloud_features()
 
     backend = backend if backend is not None else backends.CloudVmRayBackend()
+    # Figure out autostop config.
+    # Note: Ideally this can happen after provisioning, so we can check the
+    # autostop config from the launched resources. Before provisioning,
+    # we aren't sure which resources will be launched, and different
+    # resources may have different autostop configs.
     if isinstance(backend, backends.CloudVmRayBackend):
         if down and idle_minutes_to_autostop is None:
             # Use auto{stop,down} to terminate the cluster after the task is
             # done.
             idle_minutes_to_autostop = 0
+        elif not down and idle_minutes_to_autostop is None:
+            # No autostop config specified on command line, use the
+            # config from resources.
+            # TODO(cooperc): This should be done after provisioning, in order to
+            # support different autostop configs for different resources.
+            # Blockers:
+            # - Need autostop config to set requested_features before
+            #   provisioning.
+            # - Need to send info message about idle_minutes_to_autostop==0 here
+            # - Need to check if autostop is supported by the backend.
+            resources = list(task.resources)
+            for resource in resources:
+                if resource.autostop_config != resources[0].autostop_config:
+                    raise ValueError(
+                        'All resources must have the same autostop config.')
+            resource_autostop_config = resources[0].autostop_config
+
+            if resource_autostop_config is not None:
+                if resource_autostop_config.enabled:
+                    idle_minutes_to_autostop = (
+                        resource_autostop_config.idle_minutes)
+                    down = resource_autostop_config.down
+                else:
+                    # Autostop is explicitly disabled, so cancel it if it's
+                    # already set.
+                    assert not resource_autostop_config.enabled
+                    idle_minutes_to_autostop = -1
+                    down = False
         if idle_minutes_to_autostop is not None:
             if idle_minutes_to_autostop == 0:
                 # idle_minutes_to_autostop=0 can cause the following problem:
diff --git a/sky/jobs/server/core.py b/sky/jobs/server/core.py
index dc178fe27..09080c8a0 100644
--- a/sky/jobs/server/core.py
+++ b/sky/jobs/server/core.py
@@ -172,9 +172,6 @@ def launch(
         controller_resources = controller_utils.get_controller_resources(
             controller=controller,
             task_resources=sum([list(t.resources) for t in dag.tasks], []))
-        controller_idle_minutes_to_autostop, controller_down = (
-            controller_utils.get_controller_autostop_config(
-                controller=controller))
 
         vars_to_fill = {
             'remote_user_yaml_path': remote_user_yaml_path,
@@ -216,15 +213,12 @@ def launch(
         # Launch with the api server's user hash, so that sky status does not
         # show the owner of the controller as whatever user launched it first.
         with common.with_server_user_hash():
-            return execution.launch(
-                task=controller_task,
-                cluster_name=controller_name,
-                stream_logs=stream_logs,
-                idle_minutes_to_autostop=controller_idle_minutes_to_autostop,
-                down=controller_down,
-                retry_until_up=True,
-                fast=True,
-                _disable_controller_check=True)
+            return execution.launch(task=controller_task,
+                                    cluster_name=controller_name,
+                                    stream_logs=stream_logs,
+                                    retry_until_up=True,
+                                    fast=True,
+                                    _disable_controller_check=True)
 
 
 def queue_from_kubernetes_pod(
diff --git a/sky/resources.py b/sky/resources.py
index 6e9d72643..f988c0005 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -2,7 +2,7 @@
 import dataclasses
 import textwrap
 import typing
-from typing import Any, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union
 
 import colorama
 
@@ -35,6 +35,50 @@ RESOURCE_CONFIG_ALIASES = {
 }
 
 
+@dataclasses.dataclass
+class AutostopConfig:
+    """Configuration for autostop."""
+    # enabled isn't present in the yaml config, but it's needed for this class
+    # to be complete.
+    enabled: bool
+    # If enabled is False, these values are ignored.
+    idle_minutes: int = 5
+    down: bool = False
+
+    def to_yaml_config(self) -> Union[Literal[False], Dict[str, Any]]:
+        if not self.enabled:
+            return False
+        return {
+            'idle_minutes': self.idle_minutes,
+            'down': self.down,
+        }
+
+    @classmethod
+    def from_yaml_config(
+        cls, config: Union[bool, int, Dict[str, Any], None]
+    ) -> Optional['AutostopConfig']:
+        if isinstance(config, bool):
+            if config:
+                return cls(enabled=True)
+            else:
+                return cls(enabled=False)
+
+        if isinstance(config, int):
+            return cls(idle_minutes=config, down=False, enabled=True)
+
+        if isinstance(config, dict):
+            # If we have a dict, autostop is enabled. (Only way to disable is
+            # with `false`, a bool.)
+            autostop_config = cls(enabled=True)
+            if 'idle_minutes' in config:
+                autostop_config.idle_minutes = config['idle_minutes']
+            if 'down' in config:
+                autostop_config.down = config['down']
+            return autostop_config
+
+        return None
+
+
 class Resources:
     """Resources: compute requirements of Tasks.
 
@@ -52,7 +96,7 @@ class Resources:
     """
     # If any fields changed, increment the version. For backward compatibility,
     # modify the __setstate__ method to handle the old version.
-    _VERSION = 22
+    _VERSION = 23
 
     def __init__(
         self,
@@ -72,6 +116,7 @@ class Resources:
         disk_tier: Optional[Union[str, resources_utils.DiskTier]] = None,
         ports: Optional[Union[int, str, List[str], Tuple[str]]] = None,
         labels: Optional[Dict[str, str]] = None,
+        autostop: Union[bool, int, Dict[str, Any], None] = None,
         # Internal use only.
         # pylint: disable=invalid-name
         _docker_login_config: Optional[docker_utils.DockerLoginConfig] = None,
@@ -154,6 +199,8 @@ class Resources:
             instance tags. On GCP, labels map to instance labels. On
             Kubernetes, labels map to pod labels. On other clouds, labels are
             not supported and will be ignored.
+          autostop: the autostop configuration to use. For launched resources,
+            may or may not correspond to the actual current autostop config.
           _docker_login_config: the docker configuration to use. This includes
             the docker username, password, and registry server. If None, skip
             docker login.
@@ -261,6 +308,7 @@ class Resources:
         self._set_cpus(cpus)
         self._set_memory(memory)
         self._set_accelerators(accelerators, accelerator_args)
+        self._set_autostop_config(autostop)
 
     def validate(self):
         """Validate the resources and infer the missing fields if possible."""
@@ -485,6 +533,16 @@ class Resources:
     def labels(self) -> Optional[Dict[str, str]]:
         return self._labels
 
+    @property
+    def autostop_config(self) -> Optional[AutostopConfig]:
+        """The requested autostop config.
+
+        Warning: This is the autostop config that was originally used to
+        launch the resources. It may not correspond to the actual current
+        autostop config.
+        """
+        return self._autostop_config
+
     @property
     def is_image_managed(self) -> Optional[bool]:
         return self._is_image_managed
@@ -662,6 +720,12 @@ class Resources:
                                                      float]]] = accelerators
         self._accelerator_args: Optional[Dict[str, Any]] = accelerator_args
 
+    def _set_autostop_config(
+        self,
+        autostop: Union[bool, int, Dict[str, Any], None],
+    ) -> None:
+        self._autostop_config = AutostopConfig.from_yaml_config(autostop)
+
     def is_launchable(self) -> bool:
         """Returns whether the resource is launchable."""
         return self.cloud is not None and self._instance_type is not None
@@ -1363,6 +1427,10 @@ class Resources:
             if elem is not None:
                 override_configs.set_nested(key, elem)
 
+        current_autostop_config = None
+        if self.autostop_config is not None:
+            current_autostop_config = self.autostop_config.to_yaml_config()
+
         override_configs = dict(override_configs) if override_configs else None
         resources = Resources(
             cloud=override.pop('cloud', self.cloud),
@@ -1381,6 +1449,7 @@ class Resources:
             disk_tier=override.pop('disk_tier', self.disk_tier),
             ports=override.pop('ports', self.ports),
             labels=override.pop('labels', self.labels),
+            autostop=override.pop('autostop', current_autostop_config),
             _docker_login_config=override.pop('_docker_login_config',
                                               self._docker_login_config),
             _docker_username_for_runpod=override.pop(
@@ -1578,6 +1647,7 @@ class Resources:
         resources_fields['disk_tier'] = config.pop('disk_tier', None)
         resources_fields['ports'] = config.pop('ports', None)
         resources_fields['labels'] = config.pop('labels', None)
+        resources_fields['autostop'] = config.pop('autostop', None)
         resources_fields['_docker_login_config'] = config.pop(
             '_docker_login_config', None)
         resources_fields['_docker_username_for_runpod'] = config.pop(
@@ -1627,6 +1697,8 @@ class Resources:
             config['disk_tier'] = self.disk_tier.value
         add_if_not_none('ports', self.ports)
         add_if_not_none('labels', self.labels)
+        if self._autostop_config is not None:
+            config['autostop'] = self._autostop_config.to_yaml_config()
         if self._docker_login_config is not None:
             config['_docker_login_config'] = dataclasses.asdict(
                 self._docker_login_config)
@@ -1782,6 +1854,9 @@ class Resources:
             self._docker_username_for_runpod = state.pop(
                 '_docker_username_for_runpod', None)
 
+        if version < 23:
+            self._autostop_config = None
+
         self.__dict__.update(state)
 
 
diff --git a/sky/serve/server/core.py b/sky/serve/server/core.py
index 1b66bdbe1..7aeaa0272 100644
--- a/sky/serve/server/core.py
+++ b/sky/serve/server/core.py
@@ -219,17 +219,12 @@ def up(
         # whether the service is already running. If the id is the same
         # with the current job id, we know the service is up and running
         # for the first time; otherwise it is a name conflict.
-        controller_idle_minutes_to_autostop, controller_down = (
-            controller_utils.get_controller_autostop_config(
-                controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER))
         # Since the controller may be shared among multiple users, launch the
         # controller with the API server's user hash.
         with common.with_server_user_hash():
             controller_job_id, controller_handle = execution.launch(
                 task=controller_task,
                 cluster_name=controller_name,
-                idle_minutes_to_autostop=controller_idle_minutes_to_autostop,
-                down=controller_down,
                 retry_until_up=True,
                 _disable_controller_check=True,
             )
diff --git a/sky/utils/cli_utils/status_utils.py b/sky/utils/cli_utils/status_utils.py
index f6ea56ca8..6e267770b 100644
--- a/sky/utils/cli_utils/status_utils.py
+++ b/sky/utils/cli_utils/status_utils.py
@@ -201,12 +201,13 @@ def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
             controller = controller_utils.Controllers.from_name(controller_name)
             if controller is None:
                 raise ValueError(f'Controller {controller_name} not found.')
-            autostop_minutes, _ = (
-                controller_utils.get_controller_autostop_config(
-                    controller=controller))
-            if autostop_minutes is not None:
+            controller_handle: backends.CloudVmRayResourceHandle = (
+                cluster_records[0]['handle'])
+            autostop_config = (
+                controller_handle.launched_resources.autostop_config)
+            if autostop_config is not None:
                 autostop_str = (f'{colorama.Style.DIM} (will be autostopped if '
-                                f'idle for {autostop_minutes}min)'
+                                f'idle for {autostop_config.idle_minutes}min)'
                                 f'{colorama.Style.RESET_ALL}')
             click.echo(f'\n{colorama.Fore.CYAN}{colorama.Style.BRIGHT}'
                        f'{controller_name}{colorama.Style.RESET_ALL}'
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
index dcbf46e05..50db197ae 100644
--- a/sky/utils/controller_utils.py
+++ b/sky/utils/controller_utils.py
@@ -6,7 +6,7 @@ import getpass
 import os
 import tempfile
 import typing
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Set
 import uuid
 
 import colorama
@@ -517,6 +517,30 @@ def get_controller_resources(
         if custom_controller_resources_config is not None:
             controller_resources_config_copied.update(
                 custom_controller_resources_config)
+        # Compatibility with the old way of specifying the controller autostop
+        # config. TODO(cooperc): Remove this before 0.12.0.
+        custom_controller_autostop_config = skypilot_config.get_nested(
+            (controller.value.controller_type, 'controller', 'autostop'), None)
+        if custom_controller_autostop_config is not None:
+            logger.warning(
+                f'{colorama.Fore.YELLOW}Warning: Config value '
+                f'`{controller.value.controller_type}.controller.autostop` '
+                'is deprecated. Please use '
+                f'`{controller.value.controller_type}.controller.resources.'
+                f'autostop` instead.{colorama.Style.RESET_ALL}')
+            # Only set the autostop config if it is not already specified.
+            if controller_resources_config_copied.get('autostop') is None:
+                controller_resources_config_copied['autostop'] = (
+                    custom_controller_autostop_config)
+            else:
+                logger.warning(f'{colorama.Fore.YELLOW}Ignoring the old '
+                               'config, since it is already specified in '
+                               f'resources.{colorama.Style.RESET_ALL}')
+    # Set the default autostop config for the controller, if not already
+    # specified.
+    if controller_resources_config_copied.get('autostop') is None:
+        controller_resources_config_copied['autostop'] = (
+            controller.value.default_autostop_config)
 
     try:
         controller_resources = resources.Resources.from_yaml_config(
@@ -547,7 +571,10 @@ def get_controller_resources(
     if controller_record is not None:
         handle = controller_record.get('handle', None)
         if handle is not None:
-            controller_resources_to_use = handle.launched_resources
+            # Use the existing resources, but override the autostop config with
+            # the one currently specified in the config.
+            controller_resources_to_use = handle.launched_resources.copy(
+                autostop=controller_resources_config_copied.get('autostop'))
 
     # If the controller and replicas are from the same cloud (and region/zone),
     # it should provide better connectivity. We will let the controller choose
@@ -644,40 +671,6 @@ def get_controller_resources(
     return result
 
 
-def get_controller_autostop_config(
-        controller: Controllers) -> Tuple[Optional[int], bool]:
-    """Get the autostop config for the controller.
-
-    Returns:
-      A tuple of (idle_minutes_to_autostop, down), which correspond to the
-      values passed to execution.launch().
-    """
-    controller_autostop_config_copied: Dict[str, Any] = copy.copy(
-        controller.value.default_autostop_config)
-    if skypilot_config.loaded():
-        custom_controller_autostop_config = skypilot_config.get_nested(
-            (controller.value.controller_type, 'controller', 'autostop'), None)
-        if custom_controller_autostop_config is False:
-            # Disabled with `autostop: false` in config.
-            # To indicate autostop is disabled, we return None for
-            # idle_minutes_to_autostop.
-            return None, False
-        elif custom_controller_autostop_config is True:
-            # Enabled with default values. There is no change in behavior, but
-            # this is included by for completeness, since `False` is valid.
-            pass
-        elif custom_controller_autostop_config is not None:
-            # We have specific config values.
-            # Override the controller autostop config with the ones specified in
-            # the config.
-            assert isinstance(custom_controller_autostop_config, dict)
-            controller_autostop_config_copied.update(
-                custom_controller_autostop_config)
-
-    return (controller_autostop_config_copied['idle_minutes'],
-            controller_autostop_config_copied['down'])
-
-
 def _setup_proxy_command_on_controller(
         controller_launched_cloud: 'clouds.Cloud',
         user_config: Dict[str, Any]) -> config_utils.Config:
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 4ed551e08..479c23cb7 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -33,6 +33,37 @@ def _check_not_both_fields_present(field1: str, field2: str):
     }
 
 
+_AUTOSTOP_SCHEMA = {
+    'anyOf': [
+        {
+            # Use boolean to disable autostop completely, e.g.
+            #   autostop: false
+            'type': 'boolean',
+        },
+        {
+            # Shorthand to set idle_minutes by directly specifying, e.g.
+            #   autostop: 5
+            'type': 'integer',
+            'minimum': 0,
+        },
+        {
+            'type': 'object',
+            'required': [],
+            'additionalProperties': False,
+            'properties': {
+                'idle_minutes': {
+                    'type': 'integer',
+                    'minimum': 0,
+                },
+                'down': {
+                    'type': 'boolean',
+                },
+            },
+        },
+    ],
+}
+
+
 def _get_single_resources_schema():
     """Schema for a single resource in a resources list."""
     # To avoid circular imports, only import when needed.
@@ -165,6 +196,7 @@ def _get_single_resources_schema():
                     'type': 'null',
                 }]
             },
+            'autostop': _AUTOSTOP_SCHEMA,
             # The following fields are for internal use only. Should not be
             # specified in the task config.
             '_docker_login_config': {
@@ -725,29 +757,6 @@ def get_config_schema():
         if k != '$schema'
     }
     resources_schema['properties'].pop('ports')
-    autostop_schema = {
-        'anyOf': [
-            {
-                # Use boolean to disable autostop completely, e.g.
-                #   autostop: false
-                'type': 'boolean',
-            },
-            {
-                'type': 'object',
-                'required': [],
-                'additionalProperties': False,
-                'properties': {
-                    'idle_minutes': {
-                        'type': 'integer',
-                        'minimum': 0,
-                    },
-                    'down': {
-                        'type': 'boolean',
-                    },
-                },
-            },
-        ],
-    }
     controller_resources_schema = {
         'type': 'object',
         'required': [],
@@ -762,7 +771,7 @@ def get_config_schema():
                     'high_availability': {
                         'type': 'boolean',
                     },
-                    'autostop': autostop_schema,
+                    'autostop': _AUTOSTOP_SCHEMA,
                 }
             },
             'bucket': {


Test patch:
diff --git a/tests/unit_tests/test_controller_utils.py b/tests/unit_tests/test_controller_utils.py
index 28bfc6e3f..d3704afcd 100644
--- a/tests/unit_tests/test_controller_utils.py
+++ b/tests/unit_tests/test_controller_utils.py
@@ -8,32 +8,41 @@ from sky.jobs import constants as managed_job_constants
 from sky.serve import constants as serve_constants
 from sky.utils import controller_utils
 
+_DEFAULT_AUTOSTOP = {
+    'down': False,
+    'idle_minutes': 10,
+}
+
 
 @pytest.mark.parametrize(
     ('controller_type', 'custom_controller_resources_config', 'expected'), [
         ('jobs', {}, {
             'cpus': '4+',
             'memory': '8x',
-            'disk_size': 50
+            'disk_size': 50,
+            'autostop': _DEFAULT_AUTOSTOP,
         }),
         ('jobs', {
             'cpus': '4+',
-            'disk_size': 100
+            'disk_size': 100,
         }, {
             'cpus': '4+',
             'memory': '8x',
-            'disk_size': 100
+            'disk_size': 100,
+            'autostop': _DEFAULT_AUTOSTOP,
         }),
         ('serve', {}, {
             'cpus': '4+',
-            'disk_size': 200
+            'disk_size': 200,
+            'autostop': _DEFAULT_AUTOSTOP,
         }),
         ('serve', {
             'memory': '32+',
         }, {
             'cpus': '4+',
             'memory': '32+',
-            'disk_size': 200
+            'disk_size': 200,
+            'autostop': _DEFAULT_AUTOSTOP,
         }),
     ])
 def test_get_controller_resources(controller_type: str,
@@ -82,8 +91,14 @@ def _check_controller_resources(
 
 
 @pytest.mark.parametrize(('controller_type', 'default_controller_resources'), [
-    ('jobs', managed_job_constants.CONTROLLER_RESOURCES),
-    ('serve', serve_constants.CONTROLLER_RESOURCES),
+    ('jobs', {
+        **managed_job_constants.CONTROLLER_RESOURCES,
+        'autostop': _DEFAULT_AUTOSTOP,
+    }),
+    ('serve', {
+        **serve_constants.CONTROLLER_RESOURCES,
+        'autostop': _DEFAULT_AUTOSTOP,
+    }),
 ])
 def test_get_controller_resources_with_task_resources(
         controller_type: str, default_controller_resources: Dict[str, Any],