Repository: skypilot-org/skypilot
Pull request: Add `--retry-until-up`, `--region`, `--zone`, and `--idle-minutes-to-autostop` for interactive nodes
Add `--retry-until-up`, `--region`, and `--idle-minutes-to-autostop` for interactive nodes. This fixes #1131 and #1205.

Tested: `sky {cpunode, gpunode}` on Azure:
- [x] With and without `--idle-minutes-to-autostop`
- [x] With varying regions (`--region`)

I also used print statements to verify `_launch_with_confirm` is called with the correct `retry_until_up` argument. Not sure how else to test that on Azure.

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.
```

Hints from issue discussion:

Commit message: Add --idle-minutes-to-autostop flag for interactive nodes

Patch to review:
diff --git a/sky/cli.py b/sky/cli.py
index 0f36ad71cf2..307d427d012 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -207,15 +207,36 @@ def _interactive_node_cli_command(cli_func):
                               default=False,
                               required=False,
                               help='Skip confirmation prompt.')
+    retry_until_up = click.option('--retry-until-up',
+                                  '-r',
+                                  is_flag=True,
+                                  default=False,
+                                  required=False,
+                                  help='Retry until up.')
+    region_option = click.option('--region',
+                                 default=None,
+                                 type=str,
+                                 required=False,
+                                 help='The region to use.')
+    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.'))
 
     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,
         instance_type_option,
         *([gpus] if cli_func.__name__ == 'gpunode' else []),
         *([tpus] if cli_func.__name__ == 'tpunode' else []),
@@ -669,6 +690,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.
 
@@ -681,6 +704,9 @@ 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 this
+                                  many minutes of idleness.
+        retry_until_up: Retry until up when true.
     """
     assert node_type in _INTERACTIVE_NODE_TYPES, node_type
     assert session_manager in (None, 'screen', 'tmux'), session_manager
@@ -719,6 +745,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)
@@ -1957,10 +1985,11 @@ def _terminate_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],
+            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:
@@ -2008,6 +2037,7 @@ 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,
                               instance_type=instance_type,
                               accelerators=gpus,
                               use_spot=use_spot,
@@ -2021,6 +2051,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,
     )
 
 
@@ -2028,9 +2060,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],
+            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:
@@ -2075,6 +2109,7 @@ 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,
                               instance_type=instance_type,
                               use_spot=use_spot,
                               disk_size=disk_size)
@@ -2087,6 +2122,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,
     )
 
 
@@ -2094,10 +2131,11 @@ 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]],
-            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]):
+            region: 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],
+            idle_minutes_to_autostop: Optional[int], retry_until_up: bool):
     """Launch or attach to an interactive TPU node.
 
     Examples:
@@ -2147,6 +2185,7 @@ 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,
                               instance_type=instance_type,
                               accelerators=tpus,
                               accelerator_args=accelerator_args,
@@ -2161,6 +2200,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: @@ -207,15 +207,36 @@ def _interactive_node_cli_command(cli_func):
                               default=False,
                               required=False,
                               help='Skip confirmation prompt.')
+    retry_until_up = click.option('--retry-until-up',
+                                  '-r',
+                                  is_flag=True,
+                                  default=False,
+                                  required=False,
+                                  help='Retry until up.')
path: sky/cli.py
text: @user2:
Better description, such as:
```
@user1.option(
    '--retry-until-up',
    '-r',
    default=False,
    is_flag=True,
    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.'))
```

@author:
Sounds good. Will do! Thanks for the comments

diff_hunk: @@ -207,15 +207,36 @@ def _interactive_node_cli_command(cli_func):
                               default=False,
                               required=False,
                               help='Skip confirmation prompt.')
+    retry_until_up = click.option('--retry-until-up',
+                                  '-r',
+                                  is_flag=True,
+                                  default=False,
+                                  required=False,
+                                  help='Retry until up.')
+    region_option = click.option('--region',
+                                 default=None,
+                                 type=str,
+                                 required=False,
+                                 help='The region to use.')
+    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.'))
path: sky/cli.py
text: @user2:
```
@user1.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. '
          '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.'))
```