Expose --jupyter on the scheduler (#802)

* Expose --jupyter on the scheduler

* Add KubeCluster.jupyter_link property

dask_kubernetes/cli/cli.py

@@ -55,6 +55,13 @@ def gen():
     default=None,
     help="Service type for scheduler (default 'ClusterIP')",
 )
+@click.option(
+    "--jupyter",
+    type=bool,
+    default=False,
+    is_flag=True,
+    help="Start Jupyter on the scheduler (default 'False')",
+)
 def cluster(**kwargs):
     if "resources" in kwargs and kwargs["resources"] is not None:
         kwargs["resources"] = json.loads(kwargs["resources"])
@@ -89,6 +96,8 @@ def port_forward(cluster):
     try:
         console.print(f"Scheduler at: [magenta][not bold]{kcluster.scheduler_address}")
         console.print(f"Dashboard at: [cyan][not bold]{kcluster.dashboard_link}")
+        if kcluster.jupyter:
+            console.print(f"Jupyter at: [orange3][not bold]{kcluster.jupyter_link}")
         console.print("Press ctrl+c to exit", style="bright_black")
         while True:
             time.sleep(0.1)

dask_kubernetes/kubernetes.yaml

@@ -21,6 +21,7 @@ kubernetes:
   resource-timeout: 60
   custom-cluster-spec: null
   scheduler-forward-port: null
+  scheduler-jupyter: false
 
   # Classic KubeCluster options
   host: "0.0.0.0"

dask_kubernetes/operator/kubecluster/kubecluster.py

@@ -116,6 +116,8 @@ class KubeCluster(Cluster):
         Defaults to ``60`` seconds.
     scheduler_service_type: str (optional)
         Kubernetes service type to use for the scheduler. Defaults to ``ClusterIP``.
+    jupyter: bool (optional)
+        Start Jupyter on the scheduler node.
     custom_cluster_spec: str | dict (optional)
         Path to a YAML manifest or a dictionary representation of a ``DaskCluster`` resource object which will be
         used to create the cluster instead of generating one from the other keyword arguments.
@@ -179,6 +181,7 @@ def __init__(
         scheduler_service_type=None,
         custom_cluster_spec=None,
         scheduler_forward_port=None,
+        jupyter=False,
         loop=None,
         asynchronous=False,
         **kwargs,
@@ -226,6 +229,9 @@ def __init__(
         self.scheduler_forward_port = dask.config.get(
             "kubernetes.scheduler-forward-port", override_with=scheduler_forward_port
         )
+        self.jupyter = dask.config.get(
+            "kubernetes.scheduler-jupyter", override_with=jupyter
+        )
         self.idle_timeout = dask.config.get(
             "kubernetes.idle-timeout", override_with=idle_timeout
         )
@@ -345,6 +351,7 @@ async def _create_cluster(self):
                 image=self.image,
                 scheduler_service_type=self.scheduler_service_type,
                 idle_timeout=self.idle_timeout,
+                jupyter=self.jupyter,
             )
         else:
             data = self._custom_cluster_spec
@@ -420,6 +427,7 @@ async def _connect_cluster(self):
             self.env = container_spec.env
         else:
             self.env = {}
+        self.jupyter = "--jupyter" in cluster.spec.scheduler.spec.containers[0].args
         service_name = f"{cluster.name}-scheduler"
         self._log("Waiting for scheduler pod")
         await wait_for_scheduler(
@@ -810,6 +818,12 @@ def from_name(cls, name, **kwargs):
             **kwargs,
         )
 
+    @property
+    def jupyter_link(self):
+        if self.jupyter:
+            return self.dashboard_link.replace("/status", "/jupyter/lab")
+        raise RuntimeError("KubeCluster not started with jupyter enabled")
+
 
 def make_cluster_spec(
     name,
@@ -820,6 +834,7 @@ def make_cluster_spec(
     worker_command="dask-worker",
     scheduler_service_type="ClusterIP",
     idle_timeout=0,
+    jupyter=False,
 ):
     """Generate a ``DaskCluster`` kubernetes resource.
 
@@ -841,6 +856,8 @@ def make_cluster_spec(
         Worker command to use when starting the workers
     idle_timeout: int (optional)
         Timeout to cleanup idle cluster
+    jupyter: bool (optional)
+        Start Jupyter on the Dask scheduler
     """
     return {
         "apiVersion": "kubernetes.dask.org/v1",
@@ -861,6 +878,7 @@ def make_cluster_spec(
                 resources=resources,
                 image=image,
                 scheduler_service_type=scheduler_service_type,
+                jupyter=jupyter,
             ),
         },
     }
@@ -919,6 +937,7 @@ def make_scheduler_spec(
     resources=None,
     image="ghcr.io/dask/dask:latest",
     scheduler_service_type="ClusterIP",
+    jupyter=False,
 ):
     # TODO: Take the values provided in the current class constructor
     # and build a DaskWorker compatible dict
@@ -927,14 +946,17 @@ def make_scheduler_spec(
     else:
         # If they gave us a list, assume its a list of dicts and already ready to go
         env = env
+    args = ["dask-scheduler", "--host", "0.0.0.0"]
+    if jupyter:
+        args.append("--jupyter")
 
     return {
         "spec": {
             "containers": [
                 {
                     "name": "scheduler",
                     "image": image,
-                    "args": ["dask-scheduler", "--host", "0.0.0.0"],
+                    "args": args,
                     "env": env,
                     "resources": resources,
                     "ports": [