Repository: dask/distributed

Issue:
Set task key in Executor.map
If would be good to have an option to control task keys, created in `Executor.map`
Currently all tasks, created with this method, have keys of form `function-name-UUID-number`. 
For example, this method could accept one more argument with key format string.



Discussion:
Agreee that this would be useful

On Mon, Oct 10, 2016 at 4:57 AM, wl2776 notifications@github.com wrote:

> If would be good to have an option to control task keys, created in
> Executor.map
> Currently all tasks, created with this method, have keys of form
> function-name-UUID-number.
> For example, this method could accept one more argument with key format
> string.
> 
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> https://github.com/dask/distributed/issues/571, or mute the thread
> https://github.com/notifications/unsubscribe-auth/AASszLwC-w-3xulxKUBroQuxtjyp7Npdks5qyf4KgaJpZM4KSWV6
> .


Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -673,6 +673,8 @@ def map(self, func, *iterables, **kwargs):
         ----------
         func: callable
         iterables: Iterables, Iterators, or Queues
+        key: str, list
+            Prefix for task names if string.  Explicit names if list.
         pure: bool (defaults to True)
             Whether or not the function is pure.  Set ``pure=False`` for
             impure functions like ``np.random.random``.
@@ -709,6 +711,8 @@ def map(self, func, *iterables, **kwargs):
             else:
                 return queue_to_iterator(q_out)
 
+        key = kwargs.pop('key', None)
+        key = key or funcname(func)
         pure = kwargs.pop('pure', True)
         workers = kwargs.pop('workers', None)
         allow_other_workers = kwargs.pop('allow_other_workers', False)
@@ -717,13 +721,16 @@ def map(self, func, *iterables, **kwargs):
             raise ValueError("Only use allow_other_workers= if using workers=")
 
         iterables = list(zip(*zip(*iterables)))
-        if pure:
-            keys = [funcname(func) + '-' + tokenize(func, kwargs, *args)
-                    for args in zip(*iterables)]
+        if isinstance(key, list):
+            keys = key
         else:
-            uid = str(uuid.uuid4())
-            keys = [funcname(func) + '-' + uid + '-' + str(i)
-                    for i in range(min(map(len, iterables)))] if iterables else []
+            if pure:
+                keys = [key + '-' + tokenize(func, kwargs, *args)
+                        for args in zip(*iterables)]
+            else:
+                uid = str(uuid.uuid4())
+                keys = [key + '-' + uid + '-' + str(i)
+                        for i in range(min(map(len, iterables)))] if iterables else []
 
         if not kwargs:
             dsk = {key: (func,) + args


Test patch:
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
@@ -112,6 +112,20 @@ def test_map_empty(c, s, a, b):
     assert results == []
 
 
+@gen_cluster(client=True)
+def test_map_keynames(c, s, a, b):
+    futures = c.map(inc, range(4), key='INC')
+    assert all(f.key.startswith('INC') for f in futures)
+    assert isdistinct(f.key for f in futures)
+
+    futures2 = c.map(inc, [5, 6, 7, 8], key='INC')
+    assert [f.key for f in futures] != [f.key for f in futures2]
+
+    keys = ['inc-1', 'inc-2', 'inc-3', 'inc-4']
+    futures = c.map(inc, range(4), key=keys)
+    assert [f.key for f in futures] == keys
+
+
 @gen_cluster()
 def test_compatible_map(s, a, b):
     e = CompatibleExecutor((s.ip, s.port), start=False)