From 7b55db49c9af633b1dacfe81df899f168f58a7c4 Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jtomlinson@nvidia.com>
Date: Wed, 17 Aug 2022 11:52:50 +0100
Subject: [PATCH 1/2] Expose setting multiple protocols via the dask-scheduler
 CLI

---
 distributed/cli/dask_scheduler.py            | 24 ++++++++++++++++++--
 distributed/cli/tests/test_dask_scheduler.py | 16 +++++++++++++
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/distributed/cli/dask_scheduler.py b/distributed/cli/dask_scheduler.py
index 8ef9a678ae1..96db9b83f37 100755
--- a/distributed/cli/dask_scheduler.py
+++ b/distributed/cli/dask_scheduler.py
@@ -27,7 +27,7 @@
 
 @click.command(context_settings=dict(ignore_unknown_options=True))
 @click.option("--host", type=str, default="", help="URI, IP or hostname of this server")
-@click.option("--port", type=int, default=None, help="Serving port")
+@click.option("--port", type=str, default=None, help="Serving port")
 @click.option(
     "--interface",
     type=str,
@@ -130,6 +130,7 @@
 def main(
     host,
     port,
+    protocol,
     bokeh_port,
     show,
     dashboard,
@@ -162,8 +163,26 @@ def main(
         )
         dashboard = bokeh
 
+    if protocol and "," in protocol:
+        protocol = protocol.split(",")
+
+    if port:
+        if "," in port:
+            port = [int(p) for p in port.split(",")]
+        else:
+            port = int(port)
+
     if port is None and (not host or not re.search(r":\d", host)):
-        port = 8786
+        if isinstance(protocol, list):
+            port = [8786] + [0] * (len(protocol) - 1)
+        else:
+            port = 8786
+
+    if isinstance(protocol, list) or isinstance(port, list):
+        if (not isinstance(protocol, list) or not isinstance(port, list)) or len(
+            port
+        ) != len(protocol):
+            raise ValueError("--protocol and --port must both be lists of equal length")
 
     sec = {
         k: v
@@ -202,6 +221,7 @@ async def run():
             security=sec,
             host=host,
             port=port,
+            protocol=protocol,
             dashboard=dashboard,
             dashboard_address=dashboard_address,
             http_prefix=dashboard_prefix,
diff --git a/distributed/cli/tests/test_dask_scheduler.py b/distributed/cli/tests/test_dask_scheduler.py
index 89376429aad..06303458eac 100644
--- a/distributed/cli/tests/test_dask_scheduler.py
+++ b/distributed/cli/tests/test_dask_scheduler.py
@@ -131,6 +131,22 @@ def test_dashboard_non_standard_ports(loop):
         requests.get(f"http://localhost:{port2}/status/")
 
 
+def test_multiple_protocols(loop):
+    port1 = open_port()
+    port2 = open_port()
+    with popen(
+        [
+            "dask-scheduler",
+            "--protocol=tcp,ws",
+            f"--port={port1},{port2}",
+        ]
+    ) as _:
+        with Client(f"tcp://127.0.0.1:{port1}", loop=loop):
+            pass
+        with Client(f"ws://127.0.0.1:{port2}", loop=loop):
+            pass
+
+
 @pytest.mark.skipif(not LINUX, reason="Need 127.0.0.2 to mean localhost")
 def test_dashboard_allowlist(loop):
     pytest.importorskip("bokeh")

From e4485c386afec1018649d8945e061affb522e4b2 Mon Sep 17 00:00:00 2001
From: Jacob Tomlinson <jtomlinson@nvidia.com>
Date: Wed, 17 Aug 2022 15:40:45 +0100
Subject: [PATCH 2/2] Also allow multiple options for interface

---
 distributed/cli/dask_scheduler.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/distributed/cli/dask_scheduler.py b/distributed/cli/dask_scheduler.py
index 96db9b83f37..a928dbfe83b 100755
--- a/distributed/cli/dask_scheduler.py
+++ b/distributed/cli/dask_scheduler.py
@@ -131,6 +131,7 @@ def main(
     host,
     port,
     protocol,
+    interface,
     bokeh_port,
     show,
     dashboard,
@@ -163,6 +164,9 @@ def main(
         )
         dashboard = bokeh
 
+    if interface and "," in interface:
+        interface = interface.split(",")
+
     if protocol and "," in protocol:
         protocol = protocol.split(",")
 
@@ -222,6 +226,7 @@ async def run():
             host=host,
             port=port,
             protocol=protocol,
+            interface=interface,
             dashboard=dashboard,
             dashboard_address=dashboard_address,
             http_prefix=dashboard_prefix,