Update condition.py,test_condition.py

Author: nadzhou

distributed/condition.py

@@ -27,6 +27,7 @@ def __init__(self, scheduler):
                 "condition_notify": self.notify,
                 "condition_acquire": self.acquire,
                 "condition_release": self.release,
+                "condition_notify_all": self.notify_all,
             }
         )
 
@@ -37,13 +38,15 @@ def _get_lock(self, name):
 
     @log_errors
     async def acquire(self, name=None, id=None):
+        """Acquire the underlying lock"""
         lock = self._get_lock(name)
         await lock.acquire()
         self._lock_holders[name] = id
         return True
 
     @log_errors
     async def release(self, name=None, id=None):
+        """Release the underlying lock"""
         if self._lock_holders.get(name) != id:
             return False
 
@@ -59,6 +62,7 @@ async def release(self, name=None, id=None):
 
     @log_errors
     async def wait(self, name=None, id=None, timeout=None):
+        """Wait on condition"""
         # Verify lock is held by this client
         if self._lock_holders.get(name) != id:
             raise RuntimeError("wait() called without holding the lock")
@@ -99,6 +103,7 @@ async def wait(self, name=None, id=None, timeout=None):
 
     @log_errors
     def notify(self, name=None, n=1):
+        """Notify n waiters"""
         if self._lock_holders.get(name) is None:
             raise RuntimeError("notify() called without holding the lock")
 
@@ -111,6 +116,7 @@ def notify(self, name=None, n=1):
 
     @log_errors
     def notify_all(self, name=None):
+        """Notify all waiters"""
         if self._lock_holders.get(name) is None:
             raise RuntimeError("notify_all() called without holding the lock")
 
@@ -123,18 +129,27 @@ def notify_all(self, name=None):
 class Condition(SyncMethodMixin):
     """Distributed Condition Variable
 
+    Mimics asyncio.Condition API. Allows coordination between
+    distributed workers using wait/notify pattern.
+
     Parameters
     ----------
-    name: str, optional
+    name : str, optional
         Name of the condition. Same name = shared state.
-    client: Client, optional
+    client : Client, optional
         Client for scheduler communication.
 
     Examples
     --------
+    >>> from distributed import Condition
+    >>> condition = Condition('my-condition')
+    >>> async with condition:
+    ...     await condition.wait()  # Wait for notification
+
+    >>> # In another worker/client
     >>> condition = Condition('my-condition')
     >>> async with condition:
-    ...     await condition.wait()
+    ...     condition.notify()  # Wake one waiter
     """
 
     def __init__(self, name=None, client=None):
@@ -152,11 +167,20 @@ def client(self):
                 pass
         return self._client
 
+    @property
+    def loop(self):
+        return self.client.loop if self.client else None
+
     def _verify_running(self):
         if not self.client:
-            raise RuntimeError(f"{type(self)} object not properly initialized.")
+            raise RuntimeError(
+                f"{type(self)} object not properly initialized. This can happen"
+                " if the object is being deserialized outside of the context of"
+                " a Client or Worker."
+            )
 
     async def acquire(self):
+        """Acquire underlying lock"""
         self._verify_running()
         result = await self.client.scheduler.condition_acquire(
             name=self.name, id=self.id
@@ -165,13 +189,29 @@ async def acquire(self):
         return result
 
     async def release(self):
+        """Release underlying lock"""
         if not self._locked:
             raise RuntimeError("Cannot release un-acquired lock")
         self._verify_running()
         await self.client.scheduler.condition_release(name=self.name, id=self.id)
         self._locked = False
 
     async def wait(self, timeout=None):
+        """Wait until notified
+
+        Must be called while lock is held. Releases lock and waits
+        for notify(), then reacquires lock before returning.
+
+        Parameters
+        ----------
+        timeout : number or string or timedelta, optional
+            Seconds to wait on the condition in the scheduler.
+
+        Returns
+        -------
+        bool
+            True if notified, False if timeout occurred
+        """
         if not self._locked:
             raise RuntimeError("wait() called without holding the lock")
 
@@ -182,19 +222,43 @@ async def wait(self, timeout=None):
         )
         return result
 
-    async def notify(self, n=1):
+    def notify(self, n=1):
+        """Wake up one or more waiters
+
+        Parameters
+        ----------
+        n : int, optional
+            Number of waiters to wake. Default is 1.
+
+        Returns
+        -------
+        int
+            Number of waiters notified
+        """
         if not self._locked:
             raise RuntimeError("Cannot notify without holding the lock")
         self._verify_running()
-        return await self.client.scheduler.condition_notify(name=self.name, n=n)
+        return self.client.sync(
+            self.client.scheduler.condition_notify, name=self.name, n=n
+        )
+
+    def notify_all(self):
+        """Wake up all waiters
 
-    async def notify_all(self):
+        Returns
+        -------
+        int
+            Number of waiters notified
+        """
         if not self._locked:
             raise RuntimeError("Cannot notify without holding the lock")
         self._verify_running()
-        return await self.client.scheduler.condition_notify_all(name=self.name)
+        return self.client.sync(
+            self.client.scheduler.condition_notify_all, name=self.name
+        )
 
     def locked(self):
+        """Return True if lock is held"""
         return self._locked
 
     async def __aenter__(self):

distributed/tests/test_condition.py

@@ -11,7 +11,6 @@
 async def test_condition_acquire_release(c, s, a, b):
     """Test basic lock acquire/release"""
     condition = Condition("test-lock")
-
     assert not condition.locked()
     await condition.acquire()
     assert condition.locked()
@@ -23,7 +22,6 @@ async def test_condition_acquire_release(c, s, a, b):
 async def test_condition_context_manager(c, s, a, b):
     """Test context manager interface"""
     condition = Condition("test-context")
-
     assert not condition.locked()
     async with condition:
         assert condition.locked()
@@ -119,7 +117,6 @@ async def waiter():
         async with condition:
             result = await condition.wait(timeout=0.2)
             results.append(f"timeout: {result}")
-
         async with condition:
             result = await condition.wait()
             results.append(f"notified: {result}")
@@ -142,10 +139,10 @@ async def test_condition_error_without_lock(c, s, a, b):
         await condition.wait()
 
     with pytest.raises(RuntimeError, match="Cannot notify"):
-        await condition.notify()
+        condition.notify()
 
     with pytest.raises(RuntimeError, match="Cannot notify"):
-        await condition.notify_all()
+        condition.notify_all()
 
 
 @gen_cluster(client=True)
@@ -184,7 +181,6 @@ async def consumer():
 
     await prod_task
     results = await cons_task
-
     assert results == [0, 1, 2, 3, 4]
 
 
@@ -211,7 +207,6 @@ async def consumer():
         return results
 
     results = await asyncio.gather(producer(0), producer(10), consumer(), consumer())
-
     # Last two results are from consumers
     consumed = results[2] + results[3]
     assert sorted(consumed) == [0, 1, 2, 10, 11, 12]
@@ -222,41 +217,27 @@ async def test_condition_from_worker(c, s, a, b):
     """Test condition accessed from worker tasks"""
 
     def wait_on_condition(name):
-
         from distributed import Condition
 
-        async def _wait():
-            condition = Condition(name)
-            async with condition:
-                await condition.wait()
-                return "worker_notified"
-
-        from distributed.worker import get_worker
-
-        worker = get_worker()
-        return worker.loop.run_until_complete(_wait())
+        condition = Condition(name)
+        with condition:
+            condition.wait()
+        return "worker_notified"
 
     def notify_condition(name):
-        import asyncio
+        import time
 
         from distributed import Condition
 
-        async def _notify():
-            await asyncio.sleep(0.2)
-            condition = Condition(name)
-            async with condition:
-                condition.notify()
-                return "notified"
-
-        from distributed.worker import get_worker
-
-        worker = get_worker()
-        return worker.loop.run_until_complete(_notify())
+        time.sleep(0.2)
+        condition = Condition(name)
+        with condition:
+            condition.notify()
+        return "notified"
 
     name = "worker-condition"
     f1 = c.submit(wait_on_condition, name, workers=[a.address])
     f2 = c.submit(notify_condition, name, workers=[b.address])
-
     results = await c.gather([f1, f2])
     assert results == ["worker_notified", "notified"]
 
@@ -267,7 +248,6 @@ async def test_condition_same_name_different_instances(c, s, a, b):
     name = "shared-condition"
     cond1 = Condition(name)
     cond2 = Condition(name)
-
     results = []
 
     async def waiter():
@@ -336,7 +316,6 @@ async def worker(i):
         return f"worker-{i}-done"
 
     results = await asyncio.gather(worker(0), worker(1), worker(2))
-
     assert sorted(results) == ["worker-0-done", "worker-1-done", "worker-2-done"]
     assert len(arrived) == 3