Repository: dask/distributed

Issue:
Ensuring a group of tasks are scheduled together.
Hi all,

This is a feature request for applications that depend on MPI like communication framework like [XGBoost](https://github.com/dmlc/xgboost).  I will use XGBoost as an example in the following description.

## Motivation
During training, XGBoost runs 1 process per-worker and needs to perform allreduce between all processes.  When training 1 model at a time that's fine, we just submit a task for each worker:

``` python
def train(...):
    futures = []
    # workers are obtained from input data
    for i, worker in enumerate(workers):
        f = client.submit(train_on_worker, ..., workers=[worker])
        futures.append(f)

    results = client.gather(futures)
    return results

# User side:
output = train(...)
```

This works as all `train_on_worker` will get a chance to run eventually.  But if users launch multiple training sessions simultaneously,  XGBoost might hang.  For example, given there are 2 available workers, and each call to `train` use all of them.

``` python
# user side:
output_0 = client.submit(train, Xy_0)
output_1 = client.submit(train, Xy_1)
# Notice that the sub-task launched in `train` from both `train` might interleave due to parallelization.
```

distributed might schedule 1 worker for each call to `train`.  Due to the MPI allreduce call, the scheduled worker will be waiting for unscheduled worker to synchronize.  And as both calls to `train` will be waiting (remember that we only have 2 workers), other tasks might never be scheduled, resulting into a hang.

## Feature request
Notice inside the `train` function where we submit tasks for each worker, is there a way to ensure tasks submitted there can be scheduled as a whole?


Discussion:
Maybe explict-comms in dask-cuda might be useful here. Using [`CommsContext.run()`](https://github.com/rapidsai/dask-cuda/blob/branch-0.18/dask_cuda/explicit_comms/comms.py#L161) you can schedule a task on each worker that are guaranteed to run in parallel.
Relaxing the worker pinning on the submit might help here.  I believe this can be done with `allow_other_workers`: https://distributed.dask.org/en/latest/locality.html#user-control .
Thanks for the reply @madsbk @quasiben I'm looking into the `CommsContext` class, seems quite complicated so might take some time to see if I can modify it enough to fit the use case.  Not sure if `allow_other_workers` is the right way to continue, relaxing worker pinning will make data copies.  Also in the example, there's no other available worker even if I don't pin any worker.  Let me look into it a bit more.

Resolving patch:
diff --git a/distributed/__init__.py b/distributed/__init__.py
--- a/distributed/__init__.py
+++ b/distributed/__init__.py
@@ -21,6 +21,7 @@
     get_task_metadata,
 )
 from .lock import Lock
+from .multi_lock import MultiLock
 from .nanny import Nanny
 from .pubsub import Pub, Sub
 from .queues import Queue
diff --git a/distributed/multi_lock.py b/distributed/multi_lock.py
new file mode 100644
--- /dev/null
+++ b/distributed/multi_lock.py
@@ -0,0 +1,238 @@
+import asyncio
+from collections import defaultdict
+import logging
+from typing import Hashable, List
+import uuid
+
+from .client import Client
+from .utils import log_errors, TimeoutError
+from .worker import get_worker
+from .utils import parse_timedelta
+
+logger = logging.getLogger(__name__)
+
+
+class MultiLockExtension:
+    """An extension for the scheduler to manage MultiLocks
+
+    This adds the following routes to the scheduler
+
+    *  multi_lock_acquire
+    *  multi_lock_release
+
+    The approach is to maintain `self.locks` that maps a lock (unique name given to
+    `MultiLock(names=, ...)` at creation) to a list of users (instances of `MultiLock`)
+    that "requests" the lock. Additionally, `self.requests` maps a user to its requested
+    locks and `self.requests_left` maps a user to the number of locks still need.
+
+    Every time a user `x` gets to the front in `self.locks[name] = [x, ...]` it means
+    that `x` now holds the lock `name` and when it holds all the requested locks
+    `acquire()` can return.
+
+    Finally, `self.events` contains all the events users are waiting on to finish.
+    """
+
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+        self.locks = defaultdict(list)  # lock -> users
+        self.requests = {}  # user -> locks
+        self.requests_left = {}  # user -> locks still needed
+        self.events = {}
+
+        self.scheduler.handlers.update(
+            {"multi_lock_acquire": self.acquire, "multi_lock_release": self.release}
+        )
+
+        self.scheduler.extensions["multi_locks"] = self
+
+    def _request_locks(self, locks: List[str], id: Hashable, num_locks: int):
+        """Request locks
+
+        Parameters
+        ----------
+        locks: List[str]
+            Names of the locks to request.
+        id: Hashable
+            Identifier of the `MultiLock` instance requesting the locks.
+        num_locks: int
+            Number of locks in `locks` requesting
+
+        Return
+        ------
+        result: bool
+            Whether `num_locks` requested locks are free immediately or not.
+        """
+        assert id not in self.requests
+        self.requests[id] = set(locks)
+        assert len(locks) >= num_locks and num_locks > 0
+        self.requests_left[id] = num_locks
+
+        locks = sorted(locks, key=lambda x: len(self.locks[x]))
+        for i, lock in enumerate(locks):
+            self.locks[lock].append(id)
+            if len(self.locks[lock]) == 1:  # The lock was free
+                self.requests_left[id] -= 1
+                if self.requests_left[id] == 0:  # Got all locks needed
+                    # Since we got all locks need, we can remove the rest of the requests
+                    self.requests[id] -= set(locks[i + 1 :])
+                    return True
+        return False
+
+    def _refain_locks(self, locks, id):
+        """Cancel/release previously requested/acquired locks
+
+        Parameters
+        ----------
+        locks: List[str]
+            Names of the locks to refain.
+        id: Hashable
+            Identifier of the `MultiLock` instance refraining the locks.
+        """
+        waiters_ready = set()
+        for lock in locks:
+            if self.locks[lock][0] == id:
+                self.locks[lock].pop(0)
+                if self.locks[lock]:
+                    new_first = self.locks[lock][0]
+                    self.requests_left[new_first] -= 1
+                    if self.requests_left[new_first] <= 0:
+                        # Notice, `self.requests_left[new_first]` might go below zero
+                        # if more locks are freed than requested.
+                        self.requests_left[new_first] = 0
+                        waiters_ready.add(new_first)
+            else:
+                self.locks[lock].remove(id)
+            assert id not in self.locks[lock]
+        del self.requests[id]
+        del self.requests_left[id]
+
+        for waiter in waiters_ready:
+            self.scheduler.loop.add_callback(self.events[waiter].set)
+
+    async def acquire(
+        self, comm=None, locks=None, id=None, timeout=None, num_locks=None
+    ):
+        with log_errors():
+            if not self._request_locks(locks, id, num_locks):
+                assert id not in self.events
+                event = asyncio.Event()
+                self.events[id] = event
+                future = event.wait()
+                if timeout is not None:
+                    future = asyncio.wait_for(future, timeout)
+                try:
+                    await future
+                except TimeoutError:
+                    self._refain_locks(locks, id)
+                    return False
+                finally:
+                    del self.events[id]
+            # At this point `id` acquired all `locks`
+            assert self.requests_left[id] == 0
+            return True
+
+    def release(self, comm=None, id=None):
+        with log_errors():
+            self._refain_locks(self.requests[id], id)
+
+
+class MultiLock:
+    """Distributed Centralized Lock
+
+    Parameters
+    ----------
+    names: List[str]
+        Names of the locks to acquire. Choosing the same name allows two
+        disconnected processes to coordinate a lock.
+    client: Client (optional)
+        Client to use for communication with the scheduler.  If not given, the
+        default global client will be used.
+
+    Examples
+    --------
+    >>> lock = MultiLock(['x', 'y'])  # doctest: +SKIP
+    >>> lock.acquire(timeout=1)  # doctest: +SKIP
+    >>> # do things with protected resource 'x' and 'y'
+    >>> lock.release()  # doctest: +SKIP
+    """
+
+    def __init__(self, names=[], client=None):
+        try:
+            self.client = client or Client.current()
+        except ValueError:
+            # Initialise new client
+            self.client = get_worker().client
+
+        self.names = names
+        self.id = uuid.uuid4().hex
+        self._locked = False
+
+    def acquire(self, blocking=True, timeout=None, num_locks=None):
+        """Acquire the lock
+
+        Parameters
+        ----------
+        blocking : bool, optional
+            If false, don't wait on the lock in the scheduler at all.
+        timeout : string or number or timedelta, optional
+            Seconds to wait on the lock in the scheduler.  This does not
+            include local coroutine time, network transfer time, etc..
+            It is forbidden to specify a timeout when blocking is false.
+            Instead of number of seconds, it is also possible to specify
+            a timedelta in string format, e.g. "200ms".
+        num_locks : int, optional
+            Number of locks needed. If None, all locks are needed
+
+        Examples
+        --------
+        >>> lock = MultiLock(['x', 'y'])  # doctest: +SKIP
+        >>> lock.acquire(timeout="1s")  # doctest: +SKIP
+
+        Returns
+        -------
+        True or False whether or not it successfully acquired the lock
+        """
+        timeout = parse_timedelta(timeout)
+
+        if not blocking:
+            if timeout is not None:
+                raise ValueError("can't specify a timeout for a non-blocking call")
+            timeout = 0
+
+        result = self.client.sync(
+            self.client.scheduler.multi_lock_acquire,
+            locks=self.names,
+            id=self.id,
+            timeout=timeout,
+            num_locks=num_locks or len(self.names),
+        )
+        self._locked = True
+        return result
+
+    def release(self):
+        """ Release the lock if already acquired """
+        if not self.locked():
+            raise ValueError("Lock is not yet acquired")
+        ret = self.client.sync(self.client.scheduler.multi_lock_release, id=self.id)
+        self._locked = False
+        return ret
+
+    def locked(self):
+        return self._locked
+
+    def __enter__(self):
+        self.acquire()
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        self.release()
+
+    async def __aenter__(self):
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, *args, **kwargs):
+        await self.release()
+
+    def __reduce__(self):
+        return (type(self), (self.names,))
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -81,6 +81,7 @@
 from .semaphore import SemaphoreExtension
 from .recreate_exceptions import ReplayExceptionScheduler
 from .lock import LockExtension
+from .multi_lock import MultiLockExtension
 from .event import EventExtension
 from .pubsub import PubSubSchedulerExtension
 from .stealing import WorkStealing
@@ -173,6 +174,7 @@ def nogil(func):
 
 DEFAULT_EXTENSIONS = [
     LockExtension,
+    MultiLockExtension,
     PublishExtension,
     ReplayExceptionScheduler,
     QueueExtension,


Test patch:
diff --git a/distributed/tests/test_multi_locks.py b/distributed/tests/test_multi_locks.py
new file mode 100644
--- /dev/null
+++ b/distributed/tests/test_multi_locks.py
@@ -0,0 +1,210 @@
+import asyncio
+from distributed.multi_lock import MultiLockExtension
+from time import sleep
+
+
+from distributed import MultiLock, get_client
+from distributed.metrics import time
+from distributed.utils_test import gen_cluster
+from distributed.utils_test import client, cluster_fixture, loop  # noqa F401
+
+
+@gen_cluster(client=True, nthreads=[("127.0.0.1", 8)] * 2)
+async def test_single_lock(c, s, a, b):
+    await c.set_metadata("locked", False)
+
+    def f(_):
+        client = get_client()
+        with MultiLock(names=["x"]):
+            assert client.get_metadata("locked") is False
+            client.set_metadata("locked", True)
+            sleep(0.05)
+            assert client.get_metadata("locked") is True
+            client.set_metadata("locked", False)
+
+    futures = c.map(f, range(20))
+    await c.gather(futures)
+    ext: MultiLockExtension = s.extensions["multi_locks"]
+    assert not ext.events
+    assert not ext.requests
+    assert not ext.requests_left
+    assert all(len(l) == 0 for l in ext.locks.values())
+
+
+@gen_cluster(client=True)
+async def test_timeout(c, s, a, b):
+    ext: MultiLockExtension = s.extensions["multi_locks"]
+    lock1 = MultiLock(names=["x"])
+    result = await lock1.acquire()
+    assert result is True
+    assert ext.requests_left[lock1.id] == 0
+    assert ext.locks["x"] == [lock1.id]
+    assert not ext.events
+
+    lock2 = MultiLock(names=["x"])
+    assert lock1.id != lock2.id
+
+    start = time()
+    result = await lock2.acquire(timeout=0.1)
+    stop = time()
+    assert stop - start < 0.3
+    assert result is False
+    assert ext.locks["x"] == [lock1.id]
+    assert not ext.events
+
+    await lock1.release()
+
+
+@gen_cluster(client=True)
+async def test_timeout_wake_waiter(c, s, a, b):
+    ext: MultiLockExtension = s.extensions["multi_locks"]
+    l1 = MultiLock(names=["x"])
+    l2 = MultiLock(names=["x", "y"])
+    l3 = MultiLock(names=["y"])
+    await l1.acquire()
+
+    l2_acquire = asyncio.ensure_future(l2.acquire(timeout=1))
+    try:
+        await asyncio.wait_for(asyncio.shield(l2_acquire), 0.1)
+    except asyncio.TimeoutError:
+        pass
+    else:
+        assert False
+
+    l3_acquire = asyncio.ensure_future(l3.acquire())
+    try:
+        await asyncio.wait_for(asyncio.shield(l3_acquire), 0.1)
+    except asyncio.TimeoutError:
+        pass
+    else:
+        assert False
+
+    assert await l2_acquire is False
+    assert await l3_acquire
+    l1.release()
+    l3.release()
+
+
+@gen_cluster(client=True)
+async def test_multiple_locks(c, s, a, b):
+    ext: MultiLockExtension = s.extensions["multi_locks"]
+    l1 = MultiLock(names=["l1"])
+    l2 = MultiLock(names=["l2"])
+    l3 = MultiLock(names=["l1", "l2"])
+
+    # Both `l1` and `l2` are free to acquire
+    assert await l1.acquire()
+    assert await l2.acquire()
+    assert list(ext.locks.keys()) == ["l1", "l2"]
+    assert list(ext.locks.values()) == [[l1.id], [l2.id]]
+    assert list(ext.requests.keys()) == [l1.id, l2.id]
+    assert list(ext.requests_left.values()) == [0, 0]
+    assert not ext.events
+
+    # Since `l3` requires both `l1` and `l2`, it isn't available immediately
+    l3_acquire = asyncio.ensure_future(l3.acquire())
+    try:
+        await asyncio.wait_for(asyncio.shield(l3_acquire), 0.1)
+    except asyncio.TimeoutError:
+        assert list(ext.locks.keys()) == ["l1", "l2"]
+        assert list(ext.locks.values()) == [[l1.id, l3.id], [l2.id, l3.id]]
+        assert ext.requests[l3.id] == {"l1", "l2"}
+        assert ext.requests_left[l3.id] == 2
+        assert l3.id in ext.events
+    else:
+        assert False  # We except a TimeoutError since `l3` isn't availabe
+
+    # Releasing `l1` isn't enough since `l3` also requires `l2`
+    await l1.release()
+    try:
+        await asyncio.wait_for(asyncio.shield(l3_acquire), 0.1)
+    except asyncio.TimeoutError:
+        # `l3` now only wait on `l2`
+        assert list(ext.locks.keys()) == ["l1", "l2"]
+        assert list(ext.locks.values()) == [[l3.id], [l2.id, l3.id]]
+        assert ext.requests[l3.id] == {"l1", "l2"}
+        assert ext.requests_left[l3.id] == 1
+        assert l3.id in ext.events
+    else:
+        assert False
+
+    # Releasing `l2` should make `l3` available
+    await l2.release()
+    assert list(ext.locks.keys()) == ["l1", "l2"]
+    assert list(ext.locks.values()) == [[l3.id], [l3.id]]
+    assert ext.requests[l3.id] == {"l1", "l2"}
+    assert ext.requests_left[l3.id] == 0
+
+    await l3.release()
+    assert not ext.events
+    assert not ext.requests
+    assert not ext.requests_left
+    assert all(len(l) == 0 for l in ext.locks.values())
+
+
+@gen_cluster(client=True)
+async def test_num_locks(c, s, a, b):
+    ext: MultiLockExtension = s.extensions["multi_locks"]
+    l1 = MultiLock(names=["l1", "l2", "l3"])
+    l2 = MultiLock(names=["l1", "l2", "l3"])
+    l3 = MultiLock(names=["l1", "l2", "l3", "l4"])
+
+    # Even though `l1` and `l2` uses the same lock names they
+    # only requires a subset of the locks
+    assert await l1.acquire(num_locks=1)
+    assert await l2.acquire(num_locks=2)
+    assert list(ext.locks.keys()) == ["l1", "l2", "l3"]
+    assert list(ext.locks.values()) == [[l1.id], [l2.id], [l2.id]]
+    assert list(ext.requests.keys()) == [l1.id, l2.id]
+    assert list(ext.requests_left.values()) == [0, 0]
+    assert not ext.events
+
+    # Since `l3` requires three out of four locks it has to wait
+    l3_acquire = asyncio.ensure_future(l3.acquire(num_locks=3))
+    try:
+        await asyncio.wait_for(asyncio.shield(l3_acquire), 0.1)
+    except asyncio.TimeoutError:
+        assert list(ext.locks.keys()) == ["l1", "l2", "l3", "l4"]
+        assert list(ext.locks.values()) == [
+            [l1.id, l3.id],
+            [l2.id, l3.id],
+            [l2.id, l3.id],
+            [l3.id],
+        ]
+        assert list(ext.requests_left.values()) == [0, 0, 2]
+        assert l3.id in ext.events
+    else:
+        assert False  # We except a TimeoutError since `l3` isn't availabe
+
+    # Releasing `l1` isn't enough since `l3` also requires three locks
+    await l1.release()
+    try:
+        await asyncio.wait_for(asyncio.shield(l3_acquire), 0.1)
+    except asyncio.TimeoutError:
+        assert list(ext.locks.keys()) == ["l1", "l2", "l3", "l4"]
+        assert list(ext.locks.values()) == [
+            [l3.id],
+            [l2.id, l3.id],
+            [l2.id, l3.id],
+            [l3.id],
+        ]
+        assert list(ext.requests.keys()) == [l2.id, l3.id]
+        assert list(ext.requests_left.values()) == [0, 1]
+        assert l3.id in ext.events
+    else:
+        assert False
+
+    # Releasing `l2` is enough to release `l3`
+    await l2.release()
+    await asyncio.sleep(0.1)  # Give `l3` a change to wake up and acquire its locks
+    assert list(ext.locks.keys()) == ["l1", "l2", "l3", "l4"]
+    assert list(ext.locks.values()) == [[l3.id], [l3.id], [l3.id], [l3.id]]
+    assert list(ext.requests.keys()) == [l3.id]
+    assert list(ext.requests_left.values()) == [0]
+    assert l3.id not in ext.events
+
+    await l3.release()
+    assert not ext.events
+    assert not ext.requests
+    assert not ext.requests_left
+    assert all(len(l) == 0 for l in ext.locks.values())