Add port-forward CLI utility (#714)

dask_kubernetes/cli/cli.py

@@ -1,8 +1,12 @@
 import click
 import yaml
 import json
+import time
+from rich.console import Console
 
-from dask_kubernetes.operator import make_cluster_spec
+from dask_kubernetes.operator import make_cluster_spec, KubeCluster
+
+console = Console()
 
 
 class NoAliasDumper(yaml.SafeDumper):
@@ -70,3 +74,24 @@ def cluster(**kwargs):
             Dumper=NoAliasDumper,
         )
     )
+
+
+@main.command(help="Port-forward the scheduler of a DaskCluster resource")
+@click.argument("cluster")
+def port_forward(cluster):
+    with console.status(f"Connecting to cluster {cluster}") as status:
+        try:
+            kcluster = KubeCluster.from_name(
+                cluster, shutdown_on_close=False, quiet=True
+            )
+        except ValueError:
+            raise click.ClickException(f"No such cluster {cluster}")
+    try:
+        console.print(f"Scheduler at: [magenta][not bold]{kcluster.scheduler_address}")
+        console.print(f"Dashboard at: [cyan][not bold]{kcluster.dashboard_link}")
+        console.print("Press ctrl+c to exit", style="bright_black")
+        while True:
+            time.sleep(0.1)
+    except KeyboardInterrupt:
+        console.print("Shutting down port-forward")
+        kcluster.close()

dask_kubernetes/operator/kubecluster/kubecluster.py

@@ -275,7 +275,8 @@ async def _start(self):
             watch_component_status_task = asyncio.create_task(
                 self._watch_component_status()
             )
-            show_rich_output_task = asyncio.create_task(self._show_rich_output())
+            if not self.quiet:
+                show_rich_output_task = asyncio.create_task(self._show_rich_output())
             await ClusterAuth.load_first(self.auth)
             cluster_exists = (await self._get_cluster()) is not None
 
@@ -299,7 +300,8 @@ async def _start(self):
             self._log(f"Ready, dashboard available at {self.dashboard_link}")
         finally:
             watch_component_status_task.cancel()
-            show_rich_output_task.cancel()
+            if not self.quiet:
+                show_rich_output_task.cancel()
 
     def __await__(self):
         async def _():