Compatibility fixes with distributed 1.21.3

dask_drmaa/adaptive.py

@@ -6,6 +6,7 @@
 from distributed import Scheduler
 from distributed.utils import log_errors
 from distributed.deploy import adaptive
+from six import string_types
 from tornado import gen
 
 from .core import get_session
@@ -34,7 +35,7 @@ class Adaptive(adaptive.Adaptive):
     ...        """ Remove worker addresses from cluster """
     '''
     def __init__(self, cluster=None, scheduler=None, interval=1000,
-                 startup_cost=1, scale_factor=2):
+                 startup_cost=1, scale_factor=2, **kwargs):
         if cluster is None:
             raise TypeError("`Adaptive.__init__() missing required argument: "
                             "`cluster`")
@@ -50,7 +51,8 @@ def __init__(self, cluster=None, scheduler=None, interval=1000,
 
         super(Adaptive, self).__init__(scheduler, cluster, interval,
                                        startup_cost=startup_cost,
-                                       scale_factor=scale_factor)
+                                       scale_factor=scale_factor,
+                                       **kwargs)
 
     def get_busy_workers(self):
         s = self.scheduler
@@ -77,10 +79,20 @@ def get_scale_up_kwargs(self):
         kwargs = {'n': max(instances, len(self.get_busy_workers()))}
         memory = []
         if self.scheduler.unrunnable:
-            for key in self.scheduler.unrunnable:
+            for task in self.scheduler.unrunnable:
+                if isinstance(task, string_types):
+                    # Backwards compatibility for distributed pre-1.21.0
+                    key = task
+                    prefix = key
+                else:
+                    # In distributed==1.21.0, the scheduler now stores TaskState objects
+                    # instead of string keys in its task collections:
+                    # https://github.com/dask/distributed/pull/1594
+                    key = task.key
+                    prefix = task.prefix
                 duration = 0
                 memory = []
-                duration += self.scheduler.task_duration.get(key, 0.1)
+                duration += self.scheduler.task_duration.get(prefix, 0.1)
 
                 if key in self.scheduler.resource_restrictions:
                     m = self.scheduler.resource_restrictions[key].get('memory')
@@ -93,7 +105,17 @@ def get_scale_up_kwargs(self):
         return kwargs
 
     @gen.coroutine
-    def _retire_workers(self):
+    def _retire_workers(self, workers=None):
+        if workers is None:
+            workers = self.workers_to_close()
+        if not workers:
+            raise gen.Return(workers)
         with log_errors():
-            workers = yield self.scheduler.retire_workers(close_workers=True)
-            logger.info("Retiring workers {}".format(workers))
+            result = yield self.scheduler.retire_workers(workers,
+                                                         remove=True,
+                                                         close_workers=True)
+            if result:
+                logger.info("Retiring workers {}".format(result))
+            # Diverges from distributed.Adaptive here:
+            # ref c51a15a35a8a64c21c1182bfd9209cb6b7d95380
+        raise gen.Return(result)

dask_drmaa/sge.py

@@ -42,6 +42,7 @@ def create_job_template(self, nativeSpecification='', cpus=1, memory=None,
         if memory:
             args = args + ['--memory-limit', str(memory * (1 - memory_fraction))]
             args = args + ['--resources', 'memory=%f' % (memory * memory_fraction)]
+            # h_vmem is SGE-specific
             ns += ' -l h_vmem=%dG' % int(memory / 1e9) # / cpus
         if cpus:
             args = args + ['--nprocs', '1', '--nthreads', str(cpus)]

dask_drmaa/tests/test_adaptive.py

@@ -8,8 +8,6 @@
 from distributed import Client
 from distributed.utils_test import loop, inc, slowinc
 
-
-@pytest.mark.skip(reason="currently times out for an unknown reason")
 def test_adaptive_memory(loop):
     with SGECluster(scheduler_port=0, cleanup_interval=100) as cluster:
         adapt = Adaptive(cluster, cluster.scheduler)