Repository: dask/distributed

Issue:
Explicitly retry failed tasks
Currently if a task fails it is marked with an exception and that state is final.  If the user then changes an external reason for the failure they may want to rerun that failed task and restart all dependent tasks.

It would be convenient to have a `retry` operation that explicitly enabled this behavior.


Discussion:
It will also be convenient to set the `retry `operation depending on the exception.
cc @pitrou if you want to get into adding things into the scheduler itself, this would be a good issue.  This is one of our more-requested features.  

If I were to do this I would do the following:

-  Create a dict on the scheduler that mapped keys to the number of retries available (no key in the dict means no retries allowed)
-  In Scheduler.stimulus_task_errored we check this dict and, if it contains a positive number then we transition the task to `waiting` rather than `erred`
-  In the client add `retries=int` keywords to submit, map, and maybe compute/persist in a similar way to how it is implemented for `resources=`
-  Be sure to clean up the state at appropriate times, notably in `remove_key`  and `clear_task_state`
So the retry operation would be implicit rather than explicit? Or would it be desirable to allow the user to retry a future as well?
People have asked for both.  I get the implicit request more often.

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -10,9 +10,10 @@
 import errno
 from functools import partial
 from glob import glob
+import itertools
 import json
 import logging
-from numbers import Number
+from numbers import Number, Integral
 import os
 import sys
 from time import sleep
@@ -1049,6 +1050,8 @@ def submit(self, func, *args, **kwargs):
         allow_other_workers: bool (defaults to False)
             Used with `workers`. Inidicates whether or not the computations
             may be performed on workers that are not in the `workers` set(s).
+        retries: int (default to 0)
+            Number of allowed automatic retries if the task fails
 
         Examples
         --------
@@ -1069,6 +1072,7 @@ def submit(self, func, *args, **kwargs):
         pure = kwargs.pop('pure', True)
         workers = kwargs.pop('workers', None)
         resources = kwargs.pop('resources', None)
+        retries = kwargs.pop('retries', None)
         allow_other_workers = kwargs.pop('allow_other_workers', False)
 
         if allow_other_workers not in (True, False, None):
@@ -1105,7 +1109,8 @@ def submit(self, func, *args, **kwargs):
 
         futures = self._graph_to_futures(dsk, [skey], restrictions,
                                          loose_restrictions, priority={skey: 0},
-                                         resources={skey: resources} if resources else None)
+                                         resources={skey: resources} if resources else None,
+                                         retries={skey: retries} if retries else None)
 
         logger.debug("Submit %s(...), %s", funcname(func), key)
 
@@ -1146,6 +1151,8 @@ def map(self, func, *iterables, **kwargs):
         workers: set, iterable of sets
             A set of worker hostnames on which computations may be performed.
             Leave empty to default to all workers (common case)
+        retries: int (default to 0)
+            Number of allowed automatic retries if a task fails
 
         Examples
         --------
@@ -1182,6 +1189,7 @@ def map(self, func, *iterables, **kwargs):
         key = key or funcname(func)
         pure = kwargs.pop('pure', True)
         workers = kwargs.pop('workers', None)
+        retries = kwargs.pop('retries', None)
         resources = kwargs.pop('resources', None)
         allow_other_workers = kwargs.pop('allow_other_workers', False)
 
@@ -1228,6 +1236,11 @@ def map(self, func, *iterables, **kwargs):
         else:
             loose_restrictions = set()
 
+        if retries:
+            retries = {k: retries for k in keys}
+        else:
+            retries = None
+
         priority = dict(zip(keys, range(len(keys))))
 
         if resources:
@@ -1236,7 +1249,8 @@ def map(self, func, *iterables, **kwargs):
             resources = None
 
         futures = self._graph_to_futures(dsk, keys, restrictions,
-                                         loose_restrictions, priority=priority, resources=resources)
+                                         loose_restrictions, priority=priority,
+                                         resources=resources, retries=retries)
         logger.debug("map(%s, ...)", funcname(func))
 
         return [futures[tokey(k)] for k in keys]
@@ -1889,8 +1903,8 @@ def run_coroutine(self, function, *args, **kwargs):
         return self.sync(self._run_coroutine, function, *args, **kwargs)
 
     def _graph_to_futures(self, dsk, keys, restrictions=None,
-                          loose_restrictions=None, allow_other_workers=True, priority=None,
-                          resources=None):
+                          loose_restrictions=None, priority=None,
+                          resources=None, retries=None):
         with self._lock:
             keyset = set(keys)
             flatkeys = list(map(tokey, keys))
@@ -1936,7 +1950,9 @@ def _graph_to_futures(self, dsk, keys, restrictions=None,
                                      'loose_restrictions': loose_restrictions,
                                      'priority': priority,
                                      'resources': resources,
-                                     'submitting_task': getattr(thread_state, 'key', None)})
+                                     'submitting_task': getattr(thread_state, 'key', None),
+                                     'retries': retries,
+                                     })
             return futures
 
     def get(self, dsk, keys, restrictions=None, loose_restrictions=None,
@@ -2039,7 +2055,7 @@ def normalize_collection(self, collection):
 
     def compute(self, collections, sync=False, optimize_graph=True,
                 workers=None, allow_other_workers=False, resources=None,
-                **kwargs):
+                retries=0, **kwargs):
         """ Compute dask collections on cluster
 
         Parameters
@@ -2059,6 +2075,8 @@ def compute(self, collections, sync=False, optimize_graph=True,
         allow_other_workers: bool, list
             If True then all restrictions in workers= are considered loose
             If a list then only the keys for the listed collections are loose
+        retries: int (default to 0)
+            Number of allowed automatic retries if computing a result fails
         **kwargs:
             Options to pass to the graph optimize calls
 
@@ -2113,11 +2131,19 @@ def compute(self, collections, sync=False, optimize_graph=True,
                                                                  workers, allow_other_workers)
 
         if resources:
-            resources = self.expand_resources(resources)
+            resources = self._expand_resources(resources,
+                                               all_keys=itertools.chain(dsk, dsk2))
+
+        if retries:
+            retries = self._expand_retries(retries,
+                                           all_keys=itertools.chain(dsk, dsk2))
+        else:
+            retries = None
 
         futures_dict = self._graph_to_futures(merge(dsk2, dsk), names,
                                               restrictions, loose_restrictions,
-                                              resources=resources)
+                                              resources=resources,
+                                              retries=retries)
 
         i = 0
         futures = []
@@ -2139,7 +2165,8 @@ def compute(self, collections, sync=False, optimize_graph=True,
             return result
 
     def persist(self, collections, optimize_graph=True, workers=None,
-                allow_other_workers=None, resources=None, **kwargs):
+                allow_other_workers=None, resources=None, retries=None,
+                **kwargs):
         """ Persist dask collections on cluster
 
         Starts computation of the collection on the cluster in the background.
@@ -2161,6 +2188,8 @@ def persist(self, collections, optimize_graph=True, workers=None,
         allow_other_workers: bool, list
             If True then all restrictions in workers= are considered loose
             If a list then only the keys for the listed collections are loose
+        retries: int (default to 0)
+            Number of allowed automatic retries if computing a result fails
         kwargs:
             Options to pass to the graph optimize calls
 
@@ -2193,10 +2222,18 @@ def persist(self, collections, optimize_graph=True, workers=None,
                                                                  workers, allow_other_workers)
 
         if resources:
-            resources = self.expand_resources(resources)
+            resources = self._expand_resources(resources,
+                                               all_keys=itertools.chain(dsk, names))
+
+        if retries:
+            retries = self._expand_retries(retries,
+                                           all_keys=itertools.chain(dsk, names))
+        else:
+            retries = None
 
         futures = self._graph_to_futures(dsk, names, restrictions,
-                                         loose_restrictions, resources=resources)
+                                         loose_restrictions,
+                                         resources=resources, retries=retries)
 
         postpersists = [c.__dask_postpersist__() for c in collections]
         result = [func({k: futures[k] for k in flatten(c.__dask_keys__())}, *args)
@@ -2963,23 +3000,70 @@ def start_ipython_scheduler(self, magic_name='scheduler_if_ipython',
                               extra_args=qtconsole_args,)
         return info
 
-    @staticmethod
-    def expand_resources(resources):
-        assert isinstance(resources, dict)
-        out = {}
+    @classmethod
+    def _expand_key(cls, k):
+        """
+        Expand a user-provided task key specification, e.g. in a resources
+        or retries dictionary.
+        """
+        if not isinstance(k, tuple):
+            k = (k,)
+        for kk in k:
+            if dask.is_dask_collection(kk):
+                for kkk in kk.__dask_keys__():
+                    yield tokey(kkk)
+            else:
+                yield tokey(kk)
+
+    @classmethod
+    def _expand_retries(cls, retries, all_keys):
+        """
+        Expand the user-provided "retries" specification
+        to a {task key: Integral} dictionary.
+        """
+        if retries and isinstance(retries, dict):
+            return {name: value
+                    for key, value in retries.items()
+                    for name in cls._expand_key(key)}
+        elif isinstance(retries, Integral):
+            # Each task unit may potentially fail, allow retrying all of them
+            return {name: retries for name in all_keys}
+        else:
+            raise TypeError("`retries` should be an integer or dict, got %r"
+                            % (type(retries,)))
+
+    def _expand_resources(cls, resources, all_keys):
+        """
+        Expand the user-provided "resources" specification
+        to a {task key: {resource name: Number}} dictionary.
+        """
+        # Resources can either be a single dict such as {'GPU': 2},
+        # indicating a requirement for all keys, or a nested dict
+        # such as {'x': {'GPU': 1}, 'y': {'SSD': 4}} indicating
+        # per-key requirements
+        if not isinstance(resources, dict):
+            raise TypeError("`retries` should be a dict, got %r"
+                            % (type(retries,)))
+
+        per_key_reqs = {}
+        global_reqs = {}
+        all_keys = list(all_keys)
         for k, v in resources.items():
-            if not isinstance(k, tuple):
-                k = (k,)
-            for kk in k:
-                if dask.is_dask_collection(kk):
-                    for kkk in kk.__dask_keys__():
-                        out[tokey(kkk)] = v
-                else:
-                    out[tokey(kk)] = v
-        return out
+            if isinstance(v, dict):
+                # It's a per-key requirement
+                per_key_reqs.update((kk, v) for kk in cls._expand_key(k))
+            else:
+                # It's a global requirement
+                global_reqs.update((kk, {k: v}) for kk in all_keys)
 
-    @staticmethod
-    def get_restrictions(collections, workers, allow_other_workers):
+        if global_reqs and per_key_reqs:
+            raise ValueError("cannot have both per-key and all-key requirements "
+                             "in resources dict %r" % (resources,))
+        return global_reqs or per_key_reqs
+
+
+    @classmethod
+    def get_restrictions(cls, collections, workers, allow_other_workers):
         """ Get restrictions from inputs to compute/persist """
         if isinstance(workers, (str, tuple, list)):
             workers = {tuple(collections): workers}
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -170,6 +170,8 @@ class Scheduler(ServerNode):
         A dict mapping a key to another key on which it depends that has failed
     * **suspicious_tasks:** ``{key: int}``
         Number of times a task has been involved in a worker failure
+    * **retries:** ``{key: int}``
+        Number of times a task may be automatically retried after failing
     * **deleted_keys:** ``{key: {workers}}``
         Locations of workers that have keys that should be deleted
     * **wants_what:** ``{client: {key}}``:
@@ -261,6 +263,7 @@ def __init__(
         self.worker_restrictions = dict()
         self.resource_restrictions = dict()
         self.loose_restrictions = set()
+        self.retries = dict()
         self.suspicious_tasks = defaultdict(lambda: 0)
         self.waiting = dict()
         self.waiting_data = dict()
@@ -287,7 +290,7 @@ def __init__(
                                   self.host_restrictions, self.worker_restrictions,
                                   self.loose_restrictions, self.ready, self.who_wants,
                                   self.wants_what, self.unknown_durations, self.rprocessing,
-                                  self.resource_restrictions]
+                                  self.resource_restrictions, self.retries]
 
         # Worker state
         self.ncores = dict()
@@ -691,7 +694,7 @@ def add_worker(self, comm=None, address=None, keys=(), ncores=None,
     def update_graph(self, client=None, tasks=None, keys=None,
                      dependencies=None, restrictions=None, priority=None,
                      loose_restrictions=None, resources=None,
-                     submitting_task=None):
+                     submitting_task=None, retries=None):
         """
         Add new computations to the internal dask graph
 
@@ -786,6 +789,9 @@ def update_graph(self, client=None, tasks=None, keys=None,
         if resources:
             self.resource_restrictions.update(resources)
 
+        if retries:
+            self.retries.update(retries)
+
         for key in sorted(touched | keys, key=self.priority.get):
             if self.task_state[key] == 'released':
                 recommendations[key] = 'waiting'
@@ -854,9 +860,17 @@ def stimulus_task_erred(self, key=None, worker=None,
             return {}
 
         if self.task_state[key] == 'processing':
-            recommendations = self.transition(key, 'erred', cause=key,
-                                              exception=exception, traceback=traceback, worker=worker,
-                                              **kwargs)
+            retries = self.retries.get(key, 0)
+            if retries > 0:
+                self.retries[key] = retries - 1
+                recommendations = self.transition(key, 'waiting')
+            else:
+                recommendations = self.transition(key, 'erred',
+                                                  cause=key,
+                                                  exception=exception,
+                                                  traceback=traceback,
+                                                  worker=worker,
+                                                  **kwargs)
         else:
             recommendations = {}
 
@@ -2868,12 +2882,16 @@ def remove_key(self, key):
             del self.exceptions[key]
         if key in self.exceptions_blame:
             del self.exceptions_blame[key]
+        if key in self.tracebacks:
+            del self.tracebacks[key]
         if key in self.released:
             self.released.remove(key)
         if key in self.waiting_data:
             del self.waiting_data[key]
         if key in self.suspicious_tasks:
             del self.suspicious_tasks[key]
+        if key in self.retries:
+            del self.retries[key]
         if key in self.nbytes:
             del self.nbytes[key]
         if key in self.resource_restrictions:


Test patch:
diff --git a/distributed/deploy/tests/test_local.py b/distributed/deploy/tests/test_local.py
--- a/distributed/deploy/tests/test_local.py
+++ b/distributed/deploy/tests/test_local.py
@@ -323,7 +323,7 @@ def test_diagnostics_available_at_localhost(loop, processes):
         try:
             with LocalCluster(2, scheduler_port=0, processes=processes,
                               silence_logs=False, diagnostics_port=port, loop=loop) as c:
-                requests.get('http://localhost:%d/status/' % port, timeout=1)
+                requests.get('http://localhost:%d/status/' % port, timeout=3)
                 requests.get('http://127.0.0.1:%d/status/' % port, timeout=1)
         except OSError:
             pass
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -44,7 +44,8 @@
 from distributed.utils_test import (cluster, slow, slowinc, slowadd, slowdec,
                                     randominc, inc, dec, div, throws, geninc, asyncinc,
                                     gen_cluster, gen_test, double, deep, popen,
-                                    captured_logger, wait_for, async_wait_for)
+                                    captured_logger, varying, map_varying,
+                                    wait_for, async_wait_for)
 from distributed.utils_test import loop, loop_in_thread, nodebug  # flake8: noqa
 
 
@@ -135,6 +136,139 @@ def test_map_keynames(c, s, a, b):
     assert [f.key for f in futures] == keys
 
 
+@gen_cluster(client=True)
+def test_map_retries(c, s, a, b):
+    args = [[ZeroDivisionError("one"), 2, 3],
+            [4, 5, 6],
+            [ZeroDivisionError("seven"), ZeroDivisionError("eight"), 9]]
+
+    x, y, z = c.map(*map_varying(args), retries=2)
+    assert (yield x) == 2
+    assert (yield y) == 4
+    assert (yield z) == 9
+
+    x, y, z = c.map(*map_varying(args), retries=1, pure=False)
+    assert (yield x) == 2
+    assert (yield y) == 4
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield z
+    exc_info.match("eight")
+
+    x, y, z = c.map(*map_varying(args), retries=0, pure=False)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield x
+    exc_info.match("one")
+    assert (yield y) == 4
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield z
+    exc_info.match("seven")
+
+
+@gen_cluster(client=True)
+def test_compute_retries(c, s, a, b):
+    args = [ZeroDivisionError("one"), ZeroDivisionError("two"), 3]
+
+    # Sanity check for varying() use
+    x = c.compute(delayed(varying(args))())
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield x
+    exc_info.match("one")
+
+    # Same retries for all
+    x = c.compute(delayed(varying(args))(), retries=1)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield x
+    exc_info.match("two")
+
+    x = c.compute(delayed(varying(args))(), retries=2)
+    assert (yield x) == 3
+
+    args.append(4)
+    x = c.compute(delayed(varying(args))(), retries=2)
+    assert (yield x) == 3
+
+    # Per-future retries
+    xargs = [ZeroDivisionError("one"), ZeroDivisionError("two"), 30, 40]
+    yargs = [ZeroDivisionError("five"), ZeroDivisionError("six"), 70]
+    zargs = [80, 90, 100]
+
+    x, y = [delayed(varying(args))() for args in (xargs, yargs)]
+    x, y = c.compute([x, y], retries={x: 2})
+    gc.collect()
+
+    assert (yield x) == 30
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield y
+    exc_info.match("five")
+
+    x, y, z = [delayed(varying(args))() for args in (xargs, yargs, zargs)]
+    x, y, z = c.compute([x, y, z], retries={(y, z): 2})
+
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield x
+    exc_info.match("one")
+    assert (yield y) == 70
+    assert (yield z) == 80
+
+
+@gen_cluster(client=True)
+def test_compute_persisted_retries(c, s, a, b):
+    args = [ZeroDivisionError("one"), ZeroDivisionError("two"), 3]
+
+    # Sanity check
+    x = c.persist(delayed(varying(args))())
+    fut = c.compute(x)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield fut
+    exc_info.match("one")
+
+    x = c.persist(delayed(varying(args))())
+    fut = c.compute(x, retries=1)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield fut
+    exc_info.match("two")
+
+    x = c.persist(delayed(varying(args))())
+    fut = c.compute(x, retries=2)
+    assert (yield fut) == 3
+
+    args.append(4)
+    x = c.persist(delayed(varying(args))())
+    fut = c.compute(x, retries=3)
+    assert (yield fut) == 3
+
+
+@gen_cluster(client=True)
+def test_persist_retries(c, s, a, b):
+    # Same retries for all
+    args = [ZeroDivisionError("one"), ZeroDivisionError("two"), 3]
+
+    x = c.persist(delayed(varying(args))(), retries=1)
+    x = c.compute(x)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield x
+    exc_info.match("two")
+
+    x = c.persist(delayed(varying(args))(), retries=2)
+    x = c.compute(x)
+    assert (yield x) == 3
+
+    # Per-key retries
+    xargs = [ZeroDivisionError("one"), ZeroDivisionError("two"), 30, 40]
+    yargs = [ZeroDivisionError("five"), ZeroDivisionError("six"), 70]
+    zargs = [80, 90, 100]
+
+    x, y, z = [delayed(varying(args))() for args in (xargs, yargs, zargs)]
+    x, y, z = c.persist([x, y, z], retries={(y, z): 2})
+    x, y, z = c.compute([x, y, z])
+
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        yield x
+    exc_info.match("one")
+    assert (yield y) == 70
+    assert (yield z) == 80
+
+
 @gen_cluster(client=True)
 def test_future_repr(c, s, a, b):
     for func in [repr, lambda x: x._repr_html_()]:
@@ -1738,7 +1872,7 @@ def test_repr(loop):
 
 @gen_cluster(client=True)
 def test_forget_simple(c, s, a, b):
-    x = c.submit(inc, 1)
+    x = c.submit(inc, 1, retries=2)
     y = c.submit(inc, 2)
     z = c.submit(add, x, y, workers=[a.ip], allow_other_workers=True)
 
@@ -1754,8 +1888,9 @@ def test_forget_simple(c, s, a, b):
     for coll in [s.tasks, s.dependencies, s.dependents, s.waiting,
                  s.waiting_data, s.who_has, s.worker_restrictions,
                  s.host_restrictions, s.loose_restrictions,
-                 s.released, s.priority, s.exceptions, s.who_wants,
-                 s.exceptions_blame, s.nbytes, s.task_state]:
+                 s.released, s.priority, s.exceptions, s.tracebacks,
+                 s.who_wants, s.exceptions_blame, s.nbytes, s.task_state,
+                 s.retries]:
         assert x.key not in coll
         assert z.key not in coll
 
diff --git a/distributed/tests/test_resources.py b/distributed/tests/test_resources.py
--- a/distributed/tests/test_resources.py
+++ b/distributed/tests/test_resources.py
@@ -118,7 +118,7 @@ def test_persist(c, s, a, b):
 
 
 @gen_cluster(client=True, ncores=[('127.0.0.1', 1, {'resources': {'A': 1}}),
-                                  ('127.0.0.1', 1, {'resources': {'B': 1}})])
+                                  ('127.0.0.1', 1, {'resources': {'B': 11}})])
 def test_compute(c, s, a, b):
     x = delayed(inc)(1)
     y = delayed(inc)(x)
@@ -128,6 +128,12 @@ def test_compute(c, s, a, b):
 
     assert b.data
 
+    xs = [delayed(inc)(i) for i in range(10, 20)]
+    xxs = c.compute(xs, resources={'B': 1})
+    yield wait(xxs)
+
+    assert len(b.data) > 10
+
 
 @gen_cluster(client=True, ncores=[('127.0.0.1', 1, {'resources': {'A': 1}}),
                                   ('127.0.0.1', 1, {'resources': {'B': 1}})])
diff --git a/distributed/tests/test_scheduler.py b/distributed/tests/test_scheduler.py
--- a/distributed/tests/test_scheduler.py
+++ b/distributed/tests/test_scheduler.py
@@ -20,11 +20,11 @@
 from distributed.metrics import time
 from distributed.protocol.pickle import dumps
 from distributed.worker import dumps_function, dumps_task
+from distributed.utils import tmpfile
 from distributed.utils_test import (inc, dec, gen_cluster, gen_test, readone,
-                                    slowinc, slowadd, slowdec, cluster, div)
+                                    slowinc, slowadd, slowdec, cluster, div,
+                                    varying, slow)
 from distributed.utils_test import loop, nodebug  # flake8: noqa
-from distributed.utils import tmpfile
-from distributed.utils_test import slow
 from dask.compatibility import apply
 
 
@@ -1221,3 +1221,30 @@ def test_deque_handler():
     msg = deque_handler.deque[-1]
     assert 'distributed.scheduler' in deque_handler.format(msg)
     assert any(msg.msg == 'foo123' for msg in deque_handler.deque)
+
+
+@gen_cluster(client=True)
+def test_retries(c, s, a, b):
+    args = [ZeroDivisionError("one"), ZeroDivisionError("two"), 42]
+
+    future = c.submit(varying(args), retries=3)
+    result = yield future
+    assert result == 42
+    assert s.retries[future.key] == 1
+    assert future.key not in s.exceptions
+
+    future = c.submit(varying(args), retries=2, pure=False)
+    result = yield future
+    assert result == 42
+    assert s.retries[future.key] == 0
+    assert future.key not in s.exceptions
+
+    future = c.submit(varying(args), retries=1, pure=False)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        res = yield future
+    exc_info.match("two")
+
+    future = c.submit(varying(args), retries=0, pure=False)
+    with pytest.raises(ZeroDivisionError) as exc_info:
+        res = yield future
+    exc_info.match("one")
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -5,6 +5,7 @@
 import functools
 import gc
 from glob import glob
+import itertools
 import inspect
 import logging
 import logging.config
@@ -270,6 +271,59 @@ def slowidentity(*args, **kwargs):
         return args
 
 
+# This dict grows at every varying() invocation
+_varying_dict = {}
+_varying_key_gen = itertools.count()
+
+class _ModuleSlot(object):
+    def __init__(self, modname, slotname):
+        self.modname = modname
+        self.slotname = slotname
+
+    def get(self):
+        return getattr(sys.modules[self.modname], self.slotname)
+
+
+def varying(items):
+    """
+    Return a function that returns a result (or raises an exception)
+    from *items* at each call.
+    """
+    # cloudpickle would serialize the *values* of all globals
+    # used by *func* below, so we can't use `global <something>`.
+    # Instead look up the module by name to get the original namespace
+    # and not a copy.
+    slot = _ModuleSlot(__name__, '_varying_dict')
+    key = next(_varying_key_gen)
+    _varying_dict[key] = 0
+
+    def func():
+        dct = slot.get()
+        i = dct[key]
+        if i == len(items):
+            raise IndexError
+        else:
+            x = items[i]
+            dct[key] = i + 1
+            if isinstance(x, Exception):
+                raise x
+            else:
+                return x
+
+    return func
+
+
+def map_varying(itemslists):
+    """
+    Like *varying*, but return the full specification for a map() call
+    on multiple items lists.
+    """
+    def apply(func, *args, **kwargs):
+        return func(*args, **kwargs)
+
+    return apply, map(varying, itemslists)
+
+
 @gen.coroutine
 def geninc(x, delay=0.02):
     yield gen.sleep(delay)