Repository: skypilot-org/skypilot
Pull request: Add `--retry-until-up`, `--region`, `--zone`, and `--idle-minutes-to-autostop` for interactive nodes (v2)
This pull request builds on https://github.com/skypilot-org/skypilot/pull/1207 by setting `user_requested_resources` appropriately and handling the autostop edge case.

Tested:

`sky {cpunode, gpunode}` on GCP:

- [x] With and without -`-idle-minutes-to-autostop`
- [x] With varying regions and zones (`--region`, `--zone`)

I used print statements to verify `_launch_with_confirm` is called with the correct `retry_until_up` argument. 
I tested the autostop edge case by running:

- [x] `sky {cpunode, gpunode}` followed by `sky {cpunode, gpunode} -i 1`

I tested that `user_requested_resources` was set correctly by running:

- [x] `sky {cpunode, gpunode}` followed by `sky {cpunode, gpunode}` with  `--region` and `--zone` flags

@concretevitamin @Michaelvll 

Problem statement:
Support autostop option for interactive node CLIs
I found myself wanting to use autostopping when launching `sky gpunode`.  We can add the following option from `sky launch` to the interactive node CLIs:
```
  -i, --idle-minutes-to-autostop INTEGER
                                  Automatically stop the cluster after this
                                  many minutes of idleness, i.e., no running
                                  or pending jobs in the cluster's job queue.
                                  Idleness starts counting after
                                  setup/file_mounts are done; the clock gets
                                  reset whenever there are running/pending
                                  jobs in the job queue. Setting this flag is
                                  equivalent to running ``sky launch -d ...``
                                  and then ``sky autostop -i <minutes>``. If
                                  not set, the cluster will not be auto-
                                  stopped.
```


Interactive node CLIs should have feature parity with YAML specification
The interactive node CLIs (gpunode, tpunode) are missing key functionality I'd like to use:
* specification of a single region and zone to use for provisioning
* `--retry-until-up`

Overriding the Skypilot optimizer with a single region is helpful since it doesn't consider factors such as bucket locations. I'd prefer to have Skypilot wait until it can allocate a VM in a single region until successful. Currently, it will search all regions for a VM and then fail as there is zero on-demand A100 capacity across supported Skypilot AWS regions (US only).

Hints from issue discussion:
Reopening, we can close this when https://github.com/skypilot-org/skypilot/pull/1207 is resubmitted.
Thanks for the report @parasj. This is indeed very useful.

A workaround (as I'm sure you know; posting for others) is `sky launch --gpus A100 --retry-until-up --region <region>`.
Working on it :)

Commit message: Update user_requested_resources

Patch to review:
diff --git a/sky/cli.py b/sky/cli.py
index 705e2500a7a..45702514b06 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -208,15 +208,53 @@ def _interactive_node_cli_command(cli_func):
                               default=False,
                               required=False,
                               help='Skip confirmation prompt.')
+    idle_autostop = click.option('--idle-minutes-to-autostop',
+                                 '-i',
+                                 default=None,
+                                 type=int,
+                                 required=False,
+                                 help=('Automatically stop the cluster after '
+                                       'this many minutes of idleness, i.e. '
+                                       'no running or pending jobs in the '
+                                       'cluster\'s job queue. Idleness starts '
+                                       'counting after setup/file_mounts are '
+                                       'done; the clock gets reset whenever '
+                                       'there are running/pending jobs in the '
+                                       'job queue. If not set, the cluster '
+                                       'will not be auto-stopped.'))
+    retry_until_up = click.option('--retry-until-up',
+                                  '-r',
+                                  is_flag=True,
+                                  default=False,
+                                  required=False,
+                                  help=('Whether to retry provisioning '
+                                        'infinitely until the cluster is up '
+                                        'if we fail to launch the cluster on '
+                                        'any possible region/cloud due to '
+                                        'unavailability errors.'))
+    region_option = click.option('--region',
+                                 default=None,
+                                 type=str,
+                                 required=False,
+                                 help='The region to use.')
+    zone_option = click.option('--zone',
+                               default=None,
+                               type=str,
+                               required=False,
+                               help='The zone to use.')
 
     click_decorators = [
         cli.command(cls=_DocumentedCodeCommand),
         cluster_option,
         no_confirm,
         port_forward_option,
+        idle_autostop,
+        retry_until_up,
 
         # Resource options
         *([cloud_option] if cli_func.__name__ != 'tpunode' else []),
+        region_option,
+        zone_option,
         instance_type_option,
         *([gpus] if cli_func.__name__ == 'gpunode' else []),
         *([tpus] if cli_func.__name__ == 'tpunode' else []),
@@ -657,6 +695,13 @@ def _launch_with_confirm(
     elif not confirm_shown:
         click.secho(f'Running task on cluster {cluster}...', fg='yellow')
 
+    if node_type is not None and idle_minutes_to_autostop is not None \
+            and maybe_status == global_user_state.ClusterStatus.UP:
+        # We do not sky.launch if interactive node is already up, so we need
+        # to update idle timeout here.
+        handle = global_user_state.get_handle_from_cluster_name(cluster)
+        backend.set_autostop(handle, idle_minutes_to_autostop)
+
     if node_type is None or maybe_status != global_user_state.ClusterStatus.UP:
         # No need to sky.launch again when interactive node is already up.
         sky.launch(
@@ -683,6 +728,8 @@ def _create_and_ssh_into_node(
     session_manager: Optional[str] = None,
     user_requested_resources: Optional[bool] = False,
     no_confirm: bool = False,
+    idle_minutes_to_autostop: Optional[int] = None,
+    retry_until_up: bool = False,
 ):
     """Creates and attaches to an interactive node.
 
@@ -695,6 +742,14 @@ def _create_and_ssh_into_node(
         session_manager: Attach session manager: { 'screen', 'tmux' }.
         user_requested_resources: If true, user requested resources explicitly.
         no_confirm: If true, skips confirmation prompt presented to user.
+        idle_minutes_to_autostop: Automatically stop the cluster after
+                                  specified minutes of idleness. Idleness
+                                  starts counting after setup/file_mounts are
+                                  done; the clock gets reset whenever there
+                                  are running/pending jobs in the job queue.
+        retry_until_up: Whether to retry provisioning infinitely until the
+                        cluster is up if we fail to launch due to
+                        unavailability errors.
     """
     assert node_type in _INTERACTIVE_NODE_TYPES, node_type
     assert session_manager in (None, 'screen', 'tmux'), session_manager
@@ -733,6 +788,8 @@ def _create_and_ssh_into_node(
         dryrun=False,
         detach_run=True,
         no_confirm=no_confirm,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+        retry_until_up=retry_until_up,
         node_type=node_type,
     )
     handle = global_user_state.get_handle_from_cluster_name(cluster_name)
@@ -2126,10 +2183,11 @@ def _down_or_stop(name: str):
 @usage_lib.entrypoint
 # pylint: disable=redefined-outer-name
 def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
-            cloud: Optional[str], instance_type: Optional[str],
-            gpus: Optional[str], use_spot: Optional[bool],
-            screen: Optional[bool], tmux: Optional[bool],
-            disk_size: Optional[int]):
+            cloud: Optional[str], region: Optional[str], zone: Optional[str],
+            instance_type: Optional[str], gpus: Optional[str],
+            use_spot: Optional[bool], screen: Optional[bool],
+            tmux: Optional[bool], disk_size: Optional[int],
+            idle_minutes_to_autostop: Optional[int], retry_until_up: bool):
     """Launch or attach to an interactive GPU node.
 
     Examples:
@@ -2166,7 +2224,8 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
     if name is None:
         name = _default_interactive_node_name('gpunode')
 
-    user_requested_resources = not (cloud is None and instance_type is None and
+    user_requested_resources = not (cloud is None and region is None and
+                                    zone is None and instance_type is None and
                                     gpus is None and use_spot is None)
     default_resources = _INTERACTIVE_NODE_DEFAULT_RESOURCES['gpunode']
     cloud_provider = clouds.CLOUD_REGISTRY.from_str(cloud)
@@ -2177,6 +2236,8 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
     if use_spot is None:
         use_spot = default_resources.use_spot
     resources = sky.Resources(cloud=cloud_provider,
+                              region=region,
+                              zone=zone,
                               instance_type=instance_type,
                               accelerators=gpus,
                               use_spot=use_spot,
@@ -2190,6 +2251,8 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
         session_manager=session_manager,
         user_requested_resources=user_requested_resources,
         no_confirm=yes,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+        retry_until_up=retry_until_up,
     )
 
 
@@ -2197,9 +2260,11 @@ def gpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 @usage_lib.entrypoint
 # pylint: disable=redefined-outer-name
 def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
-            cloud: Optional[str], instance_type: Optional[str],
-            use_spot: Optional[bool], screen: Optional[bool],
-            tmux: Optional[bool], disk_size: Optional[int]):
+            cloud: Optional[str], region: Optional[str], zone: Optional[str],
+            instance_type: Optional[str], use_spot: Optional[bool],
+            screen: Optional[bool], tmux: Optional[bool],
+            disk_size: Optional[int], idle_minutes_to_autostop: Optional[int],
+            retry_until_up: bool):
     """Launch or attach to an interactive CPU node.
 
     Examples:
@@ -2235,7 +2300,8 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
     if name is None:
         name = _default_interactive_node_name('cpunode')
 
-    user_requested_resources = not (cloud is None and instance_type is None and
+    user_requested_resources = not (cloud is None and region is None and
+                                    zone is None and instance_type is None and
                                     use_spot is None)
     default_resources = _INTERACTIVE_NODE_DEFAULT_RESOURCES['cpunode']
     cloud_provider = clouds.CLOUD_REGISTRY.from_str(cloud)
@@ -2244,6 +2310,8 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
     if use_spot is None:
         use_spot = default_resources.use_spot
     resources = sky.Resources(cloud=cloud_provider,
+                              region=region,
+                              zone=zone,
                               instance_type=instance_type,
                               use_spot=use_spot,
                               disk_size=disk_size)
@@ -2256,6 +2324,8 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
         session_manager=session_manager,
         user_requested_resources=user_requested_resources,
         no_confirm=yes,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+        retry_until_up=retry_until_up,
     )
 
 
@@ -2263,10 +2333,12 @@ def cpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
 @usage_lib.entrypoint
 # pylint: disable=redefined-outer-name
 def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
+            region: Optional[str], zone: Optional[str],
             instance_type: Optional[str], tpus: Optional[str],
             use_spot: Optional[bool], tpu_vm: Optional[bool],
             screen: Optional[bool], tmux: Optional[bool],
-            disk_size: Optional[int]):
+            disk_size: Optional[int], idle_minutes_to_autostop: Optional[int],
+            retry_until_up: bool):
     """Launch or attach to an interactive TPU node.
 
     Examples:
@@ -2302,7 +2374,8 @@ def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
     if name is None:
         name = _default_interactive_node_name('tpunode')
 
-    user_requested_resources = not (instance_type is None and tpus is None and
+    user_requested_resources = not (region is None and zone is None and
+                                    instance_type is None and tpus is None and
                                     use_spot is None)
     default_resources = _INTERACTIVE_NODE_DEFAULT_RESOURCES['tpunode']
     accelerator_args = default_resources.accelerator_args
@@ -2316,6 +2389,8 @@ def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
     if use_spot is None:
         use_spot = default_resources.use_spot
     resources = sky.Resources(cloud=sky.GCP(),
+                              region=region,
+                              zone=zone,
                               instance_type=instance_type,
                               accelerators=tpus,
                               accelerator_args=accelerator_args,
@@ -2330,6 +2405,8 @@ def tpunode(cluster: str, yes: bool, port_forward: Optional[List[int]],
         session_manager=session_manager,
         user_requested_resources=user_requested_resources,
         no_confirm=yes,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+        retry_until_up=retry_until_up,
     )
 
 


Code review comments:
diff_hunk: @@ -657,6 +695,13 @@ def _launch_with_confirm(
     elif not confirm_shown:
         click.secho(f'Running task on cluster {cluster}...', fg='yellow')
 
+    if node_type is not None and idle_minutes_to_autostop is not None \
+            and maybe_status == global_user_state.ClusterStatus.UP:
+        # We do not sky.launch if interactive node is already up, so we need
+        # to update idle timeout here.
+        handle = global_user_state.get_handle_from_cluster_name(cluster)
+        backend.set_autostop(handle, idle_minutes_to_autostop)
path: sky/cli.py
text: @user1:
Great catch! How about we use the programmatic API `core.autostop` instead?

Also, should we make this an `elif idle_minutes_to_autostop is not None` after L718?

@author:
Good point!