Pass GPU info to worker task (#29)

dask · Aug 27, 2019 · 00ffbc6 · 00ffbc6
1 parent e16a730
commit 00ffbc6
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/dask_cloudprovider/providers/aws/ecs.py b/dask_cloudprovider/providers/aws/ecs.py
@@ -322,15 +322,17 @@ class Worker(Task):
         Other kwargs to be passed to :class:`Task`.
     """
 
-    def __init__(self, scheduler: str, cpu: int, mem: int, **kwargs):
+    def __init__(self, scheduler: str, cpu: int, mem: int, gpu: int, **kwargs):
         super().__init__(**kwargs)
         self.task_type = "worker"
         self.scheduler = scheduler
         self._cpu = cpu
         self._mem = mem
+        self._gpu = gpu
         self._overrides = {
             "command": [
-                "dask-worker",
+                "dask-cuda-worker" if self._gpu else "dask-worker",
+                self.scheduler,
                 "--name",
                 str(self.name),
                 "--nthreads",
@@ -341,7 +343,6 @@ def __init__(self, scheduler: str, cpu: int, mem: int, **kwargs):
                 "60",
             ]
         }
-        self.environment["DASK_SCHEDULER_ADDRESS"] = self.scheduler
 
 
 class ECSCluster(SpecCluster):
@@ -695,6 +696,7 @@ async def _start(self,):
             "fargate": self._fargate_workers,
             "cpu": self._worker_cpu,
             "mem": self._worker_mem,
+            "gpu": self._worker_gpu,
             **options,
         }