Remove automatic persist from await call

dask/base.py

@@ -157,15 +157,19 @@ def compute(self, **kwargs):
         return result
 
     def __await__(self):
-        from distributed import wait, futures_of
+        try:
+            from distributed import wait, futures_of
+        except ImportError:
+            raise ImportError(
+                "Using async/await with dask requires the `distributed` package"
+            )
         from tornado import gen
 
         @gen.coroutine
         def f():
-            x = self.persist()
-            if futures_of(x):
-                yield wait(x)
-            raise gen.Return(x)
+            if futures_of(self):
+                yield wait(self)
+            raise gen.Return(self)
 
         return f().__await__()
 

dask/tests/py3_test_await.py

@@ -8,7 +8,7 @@
 @gen_cluster(client=True)
 async def test_await(c, s, a, b):
     x = dask.delayed(inc)(1)
-    x = await x
+    x = await x.persist()
     assert x.key in s.tasks
     assert a.data or b.data
     assert all(f.done() for f in futures_of(x))
@@ -18,7 +18,7 @@ def test_local_scheduler():
     async def f():
         x = dask.delayed(inc)(1)
         y = x + 1
-        z = await y
+        z = await y.persist()
         assert len(z.dask) == 1
 
     asyncio.get_event_loop().run_until_complete(f())