gcp: configure scheduler port (#342)

* gcp: configure scheduler port * Refactor for flake8 Co-authored-by: Jacob Tomlinson <jtomlinson@nvidia.com>
dask · Apr 14, 2022 · ba15c3a · ba15c3a
1 parent a6249fc
commit ba15c3a
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/dask_cloudprovider/gcp/instances.py b/dask_cloudprovider/gcp/instances.py
@@ -308,18 +308,19 @@ async def start_scheduler(self):
             # scheduler must be publicly available, and firewall
             # needs to be in place to allow access to 8786 on
             # the external IP
-            self.address = f"{self.cluster.protocol}://{self.external_ip}:8786"
+            self.address = f"{self.cluster.protocol}://{self.external_ip}:{self.port}"
         else:
             # if the client is running inside GCE environment
             # it's better to use internal IP, which doesn't
             # require firewall setup
-            self.address = f"{self.cluster.protocol}://{self.internal_ip}:8786"
+            self.address = f"{self.cluster.protocol}://{self.internal_ip}:{self.port}"
         await self.wait_for_scheduler()
 
         # need to reserve internal IP for workers
         # gcp docker containers can't see resolve ip address
         self.cluster.scheduler_internal_ip = self.internal_ip
         self.cluster.scheduler_external_ip = self.external_ip
+        self.cluster.scheduler_port = self.port
 
 
 class GCPWorker(GCPInstance):
@@ -337,9 +338,12 @@ def __init__(
         self.scheduler = scheduler
         self.worker_class = worker_class
         self.name = f"dask-{self.cluster.uuid}-worker-{str(uuid.uuid4())[:8]}"
-        internal_scheduler = (
-            f"{self.cluster.protocol}://{self.cluster.scheduler_internal_ip}:8786"
+        proto, ip, port = (
+            self.cluster.protocol,
+            self.cluster.scheduler_internal_ip,
+            self.cluster.scheduler_port,
         )
+        internal_scheduler = f"{proto}://{ip}:{port}"
         self.command = " ".join(
             [
                 self.set_env,

diff --git a/dask_cloudprovider/generic/vmcluster.py b/dask_cloudprovider/generic/vmcluster.py
@@ -72,6 +72,7 @@ def __init__(
     ):
         super().__init__(*args, **kwargs)
         self.name = f"dask-{self.cluster.uuid}-scheduler"
+        self.port = scheduler_options.get("port", 8786)
         self.command = " ".join(
             [
                 self.set_env,
@@ -85,7 +86,7 @@ def __init__(
     async def start(self):
         self.cluster._log("Creating scheduler instance")
         ip = await self.create_vm()
-        self.address = f"{self.cluster.protocol}://{ip}:8786"
+        self.address = f"{self.cluster.protocol}://{ip}:{self.port}"
         await self.wait_for_scheduler()
         await super().start()