Repository: dask/distributed

Issue:
Setting resource locality when calling Executor.submit
I can think of a few use cases where certain tasks need to be executed on particular nodes:
- Partitioning of input files onto different nodes of a cluster (possibly with some redundancy) means that initial file read tasks must run on a subset of the workers.
- Heterogeneous clusters where some tasks can run on low memory systems, but a few tasks need to run on a high memory system.
- Similar to above, but some nodes have GPUs and others do not.

One possible API for this could look like:

``` python
executor = Executor(('127.0.0.1', 8787)) 
a = executor.submit(add, 1, 2, _allowed_workers=['192.168.1.100:8787', '192.168.1.101:8787']) 
```

The `_allowed_workers` kwarg is has a leading underscore to attempt to avoid collision with any kwargs of the user function being submitted.

As a related feature, it would be nice if workers could register themselves with unique identifiers that could be used with `_allowed_workers` instead of the full IP address and port.



Discussion:
What are your thoughts on handling this by restricting data rather than restricting computations.

E.g.

``` python
>>> a = send('192.168.1.100:8787', 1)
>>> b = send('192.168.1.102:8787', 2)
>>> type(a)
RemoteData

>>> c = executor.submit(add, a, b)
```

So in the case of my file example, I would send the filename to only the workers that have it?  Is there a way to indicate that the data is replicated on 2 or 3 workers?

Restricting the data wouldn't handle the situation where special computing resources are required, but that use case is more theoretical.  (I've had fewer situations in the past where that would have been helpful.)

Yeah, we'd probably do something like the following:

``` python
scatter({'filename.1.txt': ['192.168.1.100:8787', '192.168.2.100:8787'],
         'filename.2.txt': ['192.168.2.100:8787', '192.168.3.100:8787'],
         'filename.3.txt': ['192.168.3.100:8787', '192.168.4.100:8787']}, type='all')
```

Putting data explicitly in place jives much more strongly with the existing scheduler system and is a fairly trivial operation.  Not to say that we couldn't also add computation locality constraints.  They're both quite doable.

OK, that seems like a reasonable approach.


Resolving patch:
diff --git a/distributed/executor.py b/distributed/executor.py
--- a/distributed/executor.py
+++ b/distributed/executor.py
@@ -107,6 +107,7 @@ def __init__(self, center, start=True, delete_batch_time=1):
         self.futures = dict()
         self.refcount = defaultdict(lambda: 0)
         self.dask = dict()
+        self.restrictions = dict()
         self.loop = IOLoop()
         self.report_queue = Queue()
         self.scheduler_queue = Queue()
@@ -196,7 +197,7 @@ def _go(self):
             self.report(),
             scheduler(self.scheduler_queue, self.report_queue, worker_queues,
                       delete_queue, self.who_has, self.has_what, self.ncores,
-                      self.dask),
+                      self.dask, self.restrictions),
             delete(self.scheduler_queue, delete_queue,
                    self.center.ip, self.center.port, self._delete_batch_time)]
          + [worker(self.scheduler_queue, worker_queues[w], w, n)
@@ -208,6 +209,18 @@ def _go(self):
     def submit(self, func, *args, **kwargs):
         """ Submit a function application to the scheduler
 
+        Parameters
+        ----------
+        func: callable
+        *args:
+        **kwargs:
+        pure: bool (defaults to True)
+            Whether or not the function is pure.  Set ``pure=False`` for
+            impure functions like ``np.random.random``.
+        workers: set, iterable of sets
+            A set of workers on which these computations may be performed.
+            Defaults to all workers
+
         Examples
         --------
         >>> c = executor.submit(add, a, b)  # doctest: +SKIP
@@ -225,6 +238,7 @@ def submit(self, func, *args, **kwargs):
 
         key = kwargs.pop('key', None)
         pure = kwargs.pop('pure', True)
+        workers = kwargs.pop('workers', None)
 
         if key is None:
             if pure:
@@ -240,13 +254,19 @@ def submit(self, func, *args, **kwargs):
         else:
             task = (func,) + args
 
+        if workers is not None:
+            restrictions = {key: workers}
+        else:
+            restrictions = {}
+
         if key not in self.futures:
             self.futures[key] = {'event': Event(), 'status': 'waiting'}
 
         logger.debug("Submit %s(...), %s", funcname(func), key)
         self.scheduler_queue.put_nowait({'op': 'update-graph',
                                          'dsk': {key: task},
-                                         'keys': [key]})
+                                         'keys': [key],
+                                         'restrictions': restrictions})
 
         return Future(key, self)
 
@@ -262,6 +282,9 @@ def map(self, func, *iterables, **kwargs):
         pure: bool (defaults to True)
             Whether or not the function is pure.  Set ``pure=False`` for
             impure functions like ``np.random.random``.
+        workers: set, iterable of sets
+            A set of workers on which these computations may be performed.
+            Defaults to all workers
 
         Examples
         --------
@@ -276,6 +299,7 @@ def map(self, func, *iterables, **kwargs):
         distributed.executor.Executor.submit
         """
         pure = kwargs.get('pure', True)
+        workers = kwargs.get('workers', None)
         if not callable(func):
             raise TypeError("First input to map must be a callable function")
         iterables = [list(it) for it in iterables]
@@ -293,10 +317,24 @@ def map(self, func, *iterables, **kwargs):
             if key not in self.futures:
                 self.futures[key] = {'event': Event(), 'status': 'waiting'}
 
+        if isinstance(workers, (list, set)):
+            if workers and isinstance(first(workers), (list, set)):
+                if len(workers) != len(keys):
+                    raise ValueError("You only provided %d worker restrictions"
+                    " for a sequence of length %d" % (len(workers), len(keys)))
+                restrictions = dict(zip(keys, workers))
+            else:
+                restrictions = {key: workers for key in keys}
+        elif workers is None:
+            restrictions = {}
+        else:
+            raise TypeError("Workers must be a list or set of workers or None")
+
         logger.debug("map(%s, ...)", funcname(func))
         self.scheduler_queue.put_nowait({'op': 'update-graph',
                                          'dsk': dsk,
-                                         'keys': keys})
+                                         'keys': keys,
+                                         'restrictions': restrictions})
 
         return [Future(key, self) for key in keys]
 
@@ -340,7 +378,7 @@ def gather(self, futures):
         return sync(self.loop, self._gather, futures)
 
     @gen.coroutine
-    def _get(self, dsk, keys):
+    def _get(self, dsk, keys, restrictions=None):
         flatkeys = list(flatten(keys))
         for key in flatkeys:
             if key not in self.futures:
@@ -349,16 +387,23 @@ def _get(self, dsk, keys):
 
         self.scheduler_queue.put_nowait({'op': 'update-graph',
                                          'dsk': dsk,
-                                         'keys': flatkeys})
+                                         'keys': flatkeys,
+                                         'restrictions': restrictions or {}})
 
         packed = pack_data(keys, futures)
         result = yield self._gather(packed)
         raise gen.Return(result)
 
-    def get(self, dsk, keys):
+    def get(self, dsk, keys, **kwargs):
         """ Gather futures from distributed memory
 
-        Accepts a future or any nested core container of futures
+        Parameters
+        ----------
+        dsk: dict
+        keys: object, or nested lists of objects
+        restrictions: dict (optional)
+            A mapping of {key: {set of workers}} that restricts where jobs can
+            take place
 
         Examples
         --------
@@ -367,7 +412,7 @@ def get(self, dsk, keys):
         >>> e.get({'x': (add, 1, 2)}, 'x')  # doctest: +SKIP
         3
         """
-        return sync(self.loop, self._get, dsk, keys)
+        return sync(self.loop, self._get, dsk, keys, **kwargs)
 
 
 @gen.coroutine
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -151,7 +151,7 @@ def delete(scheduler_queue, delete_queue, ip, port, batch_time=1):
 
 @gen.coroutine
 def scheduler(scheduler_queue, report_queue, worker_queues, delete_queue,
-              who_has, has_what, ncores, dsk=None):
+              who_has, has_what, ncores, dsk=None, restrictions=None):
     """ The scheduler coroutine for dask scheduling
 
     This coroutine manages interactions with all worker cores and with the
@@ -181,6 +181,8 @@ def scheduler(scheduler_queue, report_queue, worker_queues, delete_queue,
     held_data = set()
     if dsk is None:
         dsk = dict()
+    if restrictions is None:
+        restrictions = dict()
     dependencies = dict()
     dependents = dict()
     waiting = dict()
@@ -206,7 +208,8 @@ def mark_ready_to_run(key):
         """ Send task to an appropriate worker, trigger worker if idle """
         if key in waiting:
             del waiting[key]
-        new_worker = decide_worker(dependencies, stacks, who_has, key)
+
+        new_worker = decide_worker(dependencies, stacks, who_has, restrictions, key)
         stacks[new_worker].append(key)
         ensure_occupied(new_worker)
 
@@ -223,7 +226,8 @@ def ensure_occupied(worker):
 
     def seed_ready_tasks():
         """ Distribute leaves among workers """
-        new_stacks = assign_many_tasks(dependencies, waiting, keyorder, who_has, stacks,
+        new_stacks = assign_many_tasks(dependencies, waiting, keyorder, who_has,
+                                       stacks, restrictions,
                                        [k for k, deps in waiting.items() if not deps])
         for worker, stack in new_stacks.items():
             if stack:
@@ -257,6 +261,8 @@ def release_key(key):
                          set(who_has), in_play,
                          waiting, waiting_data, new_dsk, new_keys)
 
+            restrictions.update(msg.get('restrictions', {}))
+
             new_keyorder = order(new_dsk)
             for key in new_keyorder:
                 if key not in keyorder:
@@ -609,34 +615,50 @@ def unrunnable(key):
     return output
 
 
-def decide_worker(dependencies, stacks, who_has, key):
+def decide_worker(dependencies, stacks, who_has, restrictions, key):
     """ Decide which worker should take task
 
     >>> dependencies = {'c': {'b'}, 'b': {'a'}}
     >>> stacks = {'alice': ['z'], 'bob': []}
     >>> who_has = {'a': {'alice'}}
+    >>> restrictions = {}
 
     We choose the worker that has the data on which 'b' depends (alice has 'a')
 
-    >>> decide_worker(dependencies, stacks, who_has, 'b')
+    >>> decide_worker(dependencies, stacks, who_has, restrictions, 'b')
     'alice'
 
     If both Alice and Bob have dependencies then we choose the less-busy worker
 
     >>> who_has = {'a': {'alice', 'bob'}}
-    >>> decide_worker(dependencies, stacks, who_has, 'b')
+    >>> decide_worker(dependencies, stacks, who_has, restrictions, 'b')
     'bob'
+
+    Optionally provide restrictions of where jobs are allowed to occur
+
+    >>> restrictions = {'b': {'alice', 'charile'}}
+    >>> decide_worker(dependencies, stacks, who_has, restrictions, 'b')
+    'alice'
     """
     deps = dependencies[key]
     workers = frequencies(w for dep in deps
                             for w in who_has[dep])
     if not workers:
         workers = stacks
+    if key in restrictions:
+        r = restrictions[key]
+        workers = {w for w in r if w in workers}
+        if not workers:
+            workers = {w for w in r if w in stacks}
+            if not workers:
+                raise ValueError("Task has no valid workers", key, r)
+
     worker = min(workers, key=lambda w: len(stacks[w]))
     return worker
 
 
-def assign_many_tasks(dependencies, waiting, keyorder, who_has, stacks, keys):
+def assign_many_tasks(dependencies, waiting, keyorder, who_has, stacks,
+        restrictions, keys):
     """ Assign many new ready tasks to workers
 
     Often at the beginning of computation we have to assign many new leaves to
@@ -656,7 +678,7 @@ def assign_many_tasks(dependencies, waiting, keyorder, who_has, stacks, keys):
 
     for k in keys:
         assert not waiting.pop(k)
-        if not dependencies[k]:
+        if not dependencies[k] and k not in restrictions:
             leaves.append(k)
         else:
             ready.append(k)
@@ -670,7 +692,7 @@ def assign_many_tasks(dependencies, waiting, keyorder, who_has, stacks, keys):
         stacks[worker].extend(keys)
 
     for key in ready:
-        worker = decide_worker(dependencies, stacks, who_has, key)
+        worker = decide_worker(dependencies, stacks, who_has, restrictions, key)
         new_stacks[worker].append(key)
         stacks[worker].append(key)
 


Test patch:
diff --git a/distributed/tests/test_executor.py b/distributed/tests/test_executor.py
--- a/distributed/tests/test_executor.py
+++ b/distributed/tests/test_executor.py
@@ -522,3 +522,87 @@ def test_tokenize_on_futures():
     e.futures[x.key]['status'] = 'finished'
 
     assert tok == tokenize(y)
+
+
+def test_restrictions_submit():
+    @gen.coroutine
+    def f(c, a, b):
+        e = Executor((c.ip, c.port), start=False)
+        IOLoop.current().spawn_callback(e._go)
+
+        x = e.submit(inc, 1, workers={a.address})
+        y = e.submit(inc, x, workers={b.address})
+        yield _wait([x, y])
+
+        assert e.restrictions[x.key] == {a.address}
+        assert x.key in a.data
+
+        assert e.restrictions[y.key] == {b.address}
+        assert y.key in b.data
+
+        yield e._shutdown()
+    _test_cluster(f)
+
+
+def test_restrictions_map():
+    @gen.coroutine
+    def f(c, a, b):
+        e = Executor((c.ip, c.port), start=False)
+        IOLoop.current().spawn_callback(e._go)
+
+        L = e.map(inc, range(5), workers={a.address})
+        yield _wait(L)
+
+        assert set(a.data) == {x.key for x in L}
+        assert not b.data
+        for x in L:
+            assert e.restrictions[x.key] == {a.address}
+
+        L = e.map(inc, [10, 11, 12], workers=[{a.address},
+                                              {a.address, b.address},
+                                              {b.address}])
+        yield _wait(L)
+
+        assert e.restrictions[L[0].key] == {a.address}
+        assert e.restrictions[L[1].key] == {a.address, b.address}
+        assert e.restrictions[L[2].key] == {b.address}
+
+        with pytest.raises(ValueError):
+            e.map(inc, [10, 11, 12], workers=[{a.address}])
+
+        yield e._shutdown()
+    _test_cluster(f)
+
+def test_restrictions_get():
+    @gen.coroutine
+    def f(c, a, b):
+        e = Executor((c.ip, c.port), start=False)
+        IOLoop.current().spawn_callback(e._go)
+
+        dsk = {'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')}
+        restrictions = {'y': {a.address}, 'z': {b.address}}
+
+        result = yield e._get(dsk, 'z', restrictions)
+        assert result == 3
+        assert 'y' in a.data
+        assert 'z' in b.data
+
+        yield e._shutdown()
+    _test_cluster(f)
+
+def dont_test_bad_restrictions_raise_exception():
+    @gen.coroutine
+    def f(c, a, b):
+        e = Executor((c.ip, c.port), start=False)
+        IOLoop.current().spawn_callback(e._go)
+
+        z = e.submit(inc, 2, workers={'bad-address'})
+        try:
+            yield z._result()
+            assert False
+        except ValueError as e:
+            assert 'bad-address' in str(e)
+            assert z.key in str(e)
+
+        yield e._shutdown()
+    _test_cluster(f)
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
@@ -281,7 +281,7 @@ def test_decide_worker_with_many_independent_leaves():
                     {('x', i * 2 + 1): {'bob'} for i in range(50)})
 
     for key in dsk:
-        worker = decide_worker(dependencies, stacks, who_has, key)
+        worker = decide_worker(dependencies, stacks, who_has, {}, key)
         stacks[worker].append(key)
 
     nhits = (len([k for k in stacks['alice'] if 'alice' in who_has[('x', k[1])]])
@@ -290,6 +290,24 @@ def test_decide_worker_with_many_independent_leaves():
     assert nhits > 90
 
 
+def test_decide_worker_with_restrictions():
+    dependencies = {'x': set()}
+    stacks = {'alice': [], 'bob': [], 'charlie': []}
+    who_has = {}
+    restrictions = {'x': {'alice', 'charlie'}}
+    result = decide_worker(dependencies, stacks, who_has, restrictions, 'x')
+    assert result in {'alice', 'charlie'}
+
+    stacks = {'alice': [1, 2, 3], 'bob': [], 'charlie': [4, 5, 6]}
+    result = decide_worker(dependencies, stacks, who_has, restrictions, 'x')
+    assert result in {'alice', 'charlie'}
+
+    dependencies = {'x': {'y'}}
+    who_has = {'y': {'bob'}}
+    result = decide_worker(dependencies, stacks, who_has, restrictions, 'x')
+    assert result in {'alice', 'charlie'}
+
+
 def test_validate_state():
     dsk = {'x': 1, 'y': (inc, 'x')}
     dependencies = {'x': set(), 'y': {'x'}}
@@ -348,10 +366,11 @@ def test_assign_many_tasks():
     keyorder = {c: 1 for c in 'abxy'}
     who_has = {'x': {'alice'}}
     stacks = {'alice': [], 'bob': []}
+    restrictions = {}
     keys = ['y', 'a']
 
     new_stacks = assign_many_tasks(dependencies, waiting, keyorder, who_has,
-                                   stacks, ['y', 'a'])
+                                   stacks, restrictions, ['y', 'a'])
 
     assert 'y' in stacks['alice']
     assert 'a' in stacks['alice'] + stacks['bob']
@@ -361,6 +380,26 @@ def test_assign_many_tasks():
     assert set(concat(new_stacks.values())) == set(concat(stacks.values()))
 
 
+def test_assign_many_tasks_with_restrictions():
+    dependencies = {'y': {'x'}, 'b': {'a'}, 'x': set(), 'a': set()}
+    waiting = {'y': set(), 'b': {'a'}, 'a': set()}
+    keyorder = {c: 1 for c in 'abxy'}
+    who_has = {'x': {'alice'}}
+    stacks = {'alice': [], 'bob': []}
+    restrictions = {'y': {'bob'}, 'a': {'alice'}}
+    keys = ['y', 'a']
+
+    new_stacks = assign_many_tasks(dependencies, waiting, keyorder, who_has,
+                                   stacks, restrictions, ['y', 'a'])
+
+    assert 'y' in stacks['bob']
+    assert 'a' in stacks['alice']
+    assert 'a' not in waiting
+    assert 'y' not in waiting
+
+    assert set(concat(new_stacks.values())) == set(concat(stacks.values()))
+
+
 def test_fill_missing_data():
     dsk = {'x': 1, 'y': (inc, 'x'), 'z': (inc, 'y')}
     dependencies, dependents = get_deps(dsk)