Add cli for opening code against a cluster

sky/cli.py

@@ -29,6 +29,7 @@
 import multiprocessing
 import os
 import shlex
+import shutil
 import signal
 import subprocess
 import sys
@@ -5228,6 +5229,56 @@ def local_down():
             f'{colorama.Fore.GREEN}Local cluster removed.{style.RESET_ALL}')
 
 
+@cli.command('code', cls=_DocumentedCodeCommand)
+@click.option('--cluster',
+              '-c',
+              default=None,
+              type=str,
+              **_get_shell_complete_args(_complete_cluster_name),
+              help=_CLUSTER_FLAG_HELP)
+@click.option("--new-window/--no-new-window",
+              default=True,
+              help="Open Code in a new window.")
+@click.option("--profile",
+              type=str,
+              default="",
+              help="Open with a VSCode profile")
+# TODO: A way to default to the sky_workdir would simplify this
+@click.option("--workdir",
+              type=str,
+              default="",
+              help="The working directory to open in VSCode.")
+@usage_lib.entrypoint
+def code(
+    cluster: Optional[str],
+    new_window: bool,
+    profile: Optional[str],
+    workdir: str,
+):
+    """Open and attach VSCode to a cluster."""
+
+    if not cluster:
+        raise click.UsageError('Please specify a cluster with ---cluster/-c.')
+
+    code_bin = shutil.which("code")
+    if not code_bin:
+        raise RuntimeError(
+            '"code" is not available in your path. Ensure VSCode in installed and the shell command is available on '
+            'your path https://code.visualstudio.com/docs/setup/mac#_launching-from-the-command-line'
+        )
+
+    command = [
+        code_bin, "--disable-workspace-trust",
+        f"--folder-uri=vscode-remote://ssh-remote+{cluster}{workdir}"
+    ]
+    if new_window:
+        command.append("--new-window")
+    if profile:
+        command += ["--profile", profile]
+
+    subprocess.check_call(command)
+
+
 def main():
     return cli()