Repository: skypilot-org/skypilot
Pull request: [CLI] Add sky ssh and port-forwarding
#87 
Enabled the cli: 
1. Attach and port forwarding to cluster: `sky ssh -p <port1> -p <port2> cluster_name` 
2. Port forwarding for gpu/cpunode: `sky gpunode -p <port1> -p <port2> ` 
3. `sky gpunode -c cluster_name` will carete the cluster with the name of cluster_name (rather than randomly generated one), if the provided cluster does not exist.


Problem statement:
`sky attach [--port-forward] <cluster_name>`
Can we have the CLI as the title?

Reason: The `sky status` only shows the <cluster_name>. If a user would like to connect to the cluster or do the port forwarding, they need to find the path of the config file by themselves. Directly using the `sky attach <cluster_name>` might be more intuitive.

Hints from issue discussion:
Yes, either `sky attach` or `sky ssh`.  Currently no one's working on it.
This will also help to expose less and less leaky messages (e.g., we
currently show "ray attach ..." and "sky down ..." as help messages).

On Sun, Dec 5, 2021 at 5:20 AM Zhanghao Wu ***@***.***> wrote:

> As the title.
>
> Reason: The sky status only shows the <cluster_name>. If a user would
> like to connect to the cluster or do the port forwarding, they need to find
> the path of the config file by themselves. Directly using the sky attach
> <cluster_name> might be more intuitive.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/concretevitamin/sky-experiments/issues/87>, or
> unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAEQWHTXDVHFNOJWP3MOQPLUPKA3FANCNFSM5JL7OQGQ>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub>.
>
>

Commit message: Fix Ctrl-D error

Patch to review:
diff --git a/prototype/sky/backends/backend_utils.py b/prototype/sky/backends/backend_utils.py
index c43d93e88f3..7144b014b30 100644
--- a/prototype/sky/backends/backend_utils.py
+++ b/prototype/sky/backends/backend_utils.py
@@ -351,7 +351,9 @@ def redirect_process_output(proc, log_path, stream_logs, start_streaming_at=''):
 
 
 def run(cmd, **kwargs):
-    return subprocess.run(cmd, shell=True, check=True, **kwargs)
+    check = kwargs.pop('check', True)
+    shell = kwargs.pop('shell', True)
+    return subprocess.run(cmd, shell=shell, check=check, **kwargs)
 
 
 def run_no_outputs(cmd, **kwargs):
diff --git a/prototype/sky/cli.py b/prototype/sky/cli.py
index e0e5730b125..8638d8b760c 100644
--- a/prototype/sky/cli.py
+++ b/prototype/sky/cli.py
@@ -25,6 +25,7 @@
 """
 import os
 import time
+from typing import List, Optional
 import yaml
 
 import click
@@ -65,7 +66,9 @@ def _get_region_zones_from_handle(handle):
 def _create_interactive_node(name,
                              resources,
                              cluster_handle=None,
-                             backend=cloud_vm_ray_backend.CloudVmRayBackend):
+                             backend=cloud_vm_ray_backend.CloudVmRayBackend,
+                             port_forward: Optional[List[int]] = None,
+                             cluster_name: Optional[str] = None):
     """Creates an interactive session.
 
     Args:
@@ -90,22 +93,27 @@ def _create_interactive_node(name,
     backend = backend()
     backend.register_info(dag=dag, optimize_target=sky.OptimizeTarget.COST)
 
-    dag = sky.optimize(dag, minimize=sky.OptimizeTarget.COST)
-    task = dag.tasks[0]
-
     handle = cluster_handle
     if handle is None:
+        dag = sky.optimize(dag, minimize=sky.OptimizeTarget.COST)
+        task = dag.tasks[0]
         handle = backend.provision(task,
                                    task.best_resources,
                                    dryrun=False,
-                                   stream_logs=True)
+                                   stream_logs=True,
+                                   cluster_name=cluster_name)
 
     task_id = global_user_state.add_task(task)
 
     # TODO: cd into workdir immediately on the VM
     # TODO: Delete the temporary cluster config yml (or figure out a way to
     # re-use it)
-    backend_utils.run(f'ray attach {handle}')
+    extra_cmd = ''
+    if port_forward is not None:
+        extra_cmd = ' '.join(
+            [f'--port-forward {port}' for port in port_forward])
+    # Disable check, since the returncode could be non-zero if the user Ctrl-D
+    backend_utils.run(f'ray attach {extra_cmd} {handle}', check=False)
 
     if cluster_handle is None:  # if this is a secondary
         backend_utils.run(f'ray down -y {handle}')
@@ -313,6 +321,49 @@ def down(cluster, all):  # pylint: disable=redefined-builtin
         click.secho(f'Tearing down cluster {name}...done.', fg='green')
 
 
+@cli.command()
+@click.argument('cluster', required=False)
+@click.option('--port-forward',
+              '-p',
+              multiple=True,
+              default=[],
+              type=int,
+              required=False,
+              help='Port to be forwarded. To forward multiple ports, '
+              'use this option multiple times.')
+def attach(cluster, port_forward):  # pylint: disable=redefined-builtin
+    """Attach to the cluster CLUSTER.
+
+    CLUSTER is the name of the cluster to attach.  If CLUSTER is not supplied,
+    the last cluster will be attached.
+
+    Examples:
+
+      \b
+      # Attach to a specific cluster.
+      sky attach cluster_name
+
+      \b
+      # Port forward.
+      sky attach -p 8080 -p 4650 cluster_name
+    """
+    # FIXME: make TPU part of handles; so that this kills TPUs too.
+    name = cluster
+
+    to_attach = name
+    if to_attach is None:
+        to_attach = sorted(global_user_state.get_clusters(),
+                           key=lambda x: x['launched_at'])[-1]
+    handle = global_user_state.get_handle_from_cluster_name(to_attach)
+    assert handle is not None, f'Cluster {name} is not found.'
+
+    # FIXME: Assumes a specific backend.
+    _create_interactive_node(name,
+                             sky.Resources(),
+                             handle,
+                             port_forward=port_forward)
+
+
 @cli.command()
 def status():
     """Show the status of all tasks and clusters."""
@@ -352,20 +403,29 @@ def status():
               default=None,
               type=str,
               help=_CLUSTER_FLAG_HELP)
-def gpunode(cluster):
+@click.option('--port-forward',
+              '-p',
+              multiple=True,
+              default=[],
+              type=int,
+              required=False,
+              help='Port to be forwarded. To forward multiple ports, '
+              'use this option multiple times.')
+def gpunode(cluster, port_forward):
     """Launches an interactive GPU node.
 
     Automatically syncs current working directory.
     """
     # TODO: Sync code files between local and interactive node (watch rsync?)
-    # TODO: Add port forwarding to allow access to localhost:PORT for jupyter
     handle = None
     if cluster is not None:
         handle = global_user_state.get_handle_from_cluster_name(cluster)
 
     _create_interactive_node('gpunode',
                              {sky.Resources(sky.AWS(), accelerators='V100')},
-                             cluster_handle=handle)
+                             cluster_handle=handle,
+                             port_forward=port_forward,
+                             cluster_name=cluster)
 
 
 @cli.command()
@@ -374,7 +434,15 @@ def gpunode(cluster):
               default=None,
               type=str,
               help=_CLUSTER_FLAG_HELP)
-def cpunode(cluster):
+@click.option('--port-forward',
+              '-p',
+              multiple=True,
+              default=[],
+              type=int,
+              required=False,
+              help='Port to be forwarded. To forward multiple ports, '
+              'use this option multiple times.')
+def cpunode(cluster, port_forward):
     """Launches an interactive CPU node.
 
     Automatically syncs current working directory.
@@ -384,7 +452,9 @@ def cpunode(cluster):
         handle = global_user_state.get_handle_from_cluster_name(cluster)
 
     _create_interactive_node('cpunode', {sky.Resources(sky.AWS())},
-                             cluster_handle=handle)
+                             cluster_handle=handle,
+                             port_forward=port_forward,
+                             cluster_name=cluster)
 
 
 def main():


Code review comments:
diff_hunk: @@ -90,22 +93,27 @@ def _create_interactive_node(name,
     backend = backend()
     backend.register_info(dag=dag, optimize_target=sky.OptimizeTarget.COST)
 
-    dag = sky.optimize(dag, minimize=sky.OptimizeTarget.COST)
-    task = dag.tasks[0]
-
     handle = cluster_handle
     if handle is None:
+        dag = sky.optimize(dag, minimize=sky.OptimizeTarget.COST)
+        task = dag.tasks[0]
         handle = backend.provision(task,
                                    task.best_resources,
                                    dryrun=False,
-                                   stream_logs=True)
+                                   stream_logs=True,
+                                   cluster_name=cluster_name)
 
     task_id = global_user_state.add_task(task)
 
     # TODO: cd into workdir immediately on the VM
     # TODO: Delete the temporary cluster config yml (or figure out a way to
     # re-use it)
-    backend_utils.run(f'ray attach {handle}')
+    extra_cmd = ''
+    if port_forward is not None:
+        extra_cmd = ' '.join(
+            [f'--port-forward {port}' for port in port_forward])
+    # Disable check, since the returncode could be non-zero if the user Ctrl-D
+    backend_utils.run(f'ray attach {extra_cmd} {handle}', check=False)
path: prototype/sky/cli.py
text: @user1:
Was confused for a moment, maybe `attach_options`.

diff_hunk: @@ -313,6 +321,49 @@ def down(cluster, all):  # pylint: disable=redefined-builtin
         click.secho(f'Tearing down cluster {name}...done.', fg='green')
 
 
+@cli.command()
+@click.argument('cluster', required=False)
+@click.option('--port-forward',
+              '-p',
+              multiple=True,
+              default=[],
+              type=int,
+              required=False,
+              help='Port to be forwarded. To forward multiple ports, '
+              'use this option multiple times.')
+def attach(cluster, port_forward):  # pylint: disable=redefined-builtin
+    """Attach to the cluster CLUSTER.
+
+    CLUSTER is the name of the cluster to attach.  If CLUSTER is not supplied,
+    the last cluster will be attached.
+
+    Examples:
+
+      \b
+      # Attach to a specific cluster.
+      sky attach cluster_name
+
+      \b
+      # Port forward.
+      sky attach -p 8080 -p 4650 cluster_name
+    """
+    # FIXME: make TPU part of handles; so that this kills TPUs too.
+    name = cluster
+
+    to_attach = name
+    if to_attach is None:
+        to_attach = sorted(global_user_state.get_clusters(),
+                           key=lambda x: x['launched_at'])[-1]
+    handle = global_user_state.get_handle_from_cluster_name(to_attach)
+    assert handle is not None, f'Cluster {name} is not found.'
path: prototype/sky/cli.py
text: @user1:
Probably want to raise an appropriate click exception like `sky run`; assert is for internal correctness

diff_hunk: @@ -313,6 +321,49 @@ def down(cluster, all):  # pylint: disable=redefined-builtin
         click.secho(f'Tearing down cluster {name}...done.', fg='green')
 
 
+@cli.command()
+@click.argument('cluster', required=False)
+@click.option('--port-forward',
+              '-p',
+              multiple=True,
+              default=[],
+              type=int,
+              required=False,
+              help='Port to be forwarded. To forward multiple ports, '
+              'use this option multiple times.')
+def attach(cluster, port_forward):  # pylint: disable=redefined-builtin
+    """Attach to the cluster CLUSTER.
+
+    CLUSTER is the name of the cluster to attach.  If CLUSTER is not supplied,
+    the last cluster will be attached.
+
+    Examples:
+
+      \b
+      # Attach to a specific cluster.
+      sky attach cluster_name
+
+      \b
+      # Port forward.
+      sky attach -p 8080 -p 4650 cluster_name
+    """
+    # FIXME: make TPU part of handles; so that this kills TPUs too.
+    name = cluster
+
+    to_attach = name
+    if to_attach is None:
+        to_attach = sorted(global_user_state.get_clusters(),
path: prototype/sky/cli.py
text: @user1:
Print an error when there are no clusters.

diff_hunk: @@ -313,6 +321,49 @@ def down(cluster, all):  # pylint: disable=redefined-builtin
         click.secho(f'Tearing down cluster {name}...done.', fg='green')
 
 
+@cli.command()
+@click.argument('cluster', required=False)
+@click.option('--port-forward',
+              '-p',
+              multiple=True,
+              default=[],
+              type=int,
+              required=False,
+              help='Port to be forwarded. To forward multiple ports, '
+              'use this option multiple times.')
+def attach(cluster, port_forward):  # pylint: disable=redefined-builtin
path: prototype/sky/cli.py
text: @user1:
Can do away without this pylint disable