Repository: dask/dask

Issue:
ResourceProfiler keeps running after `__exit__`
If you write something like:
```
    with ResourceProfiler(dt=0.01) as rprof:
        # ...
```

then the profiler process keeps running at the end of the `with` block. You have to call `rprof.close()` as well, which is unexpected.
(noticed in the distributed test suite)


Discussion:
jcrist: Ping @pitrou for review.

Patch:
diff --git a/dask/diagnostics/profile.py b/dask/diagnostics/profile.py
index 48e68a565..ec855173e 100644
--- a/dask/diagnostics/profile.py
+++ b/dask/diagnostics/profile.py
@@ -140,29 +140,35 @@ class ResourceProfiler(Callback):
     data will only be collected while a dask scheduler is active.
     """
     def __init__(self, dt=1):
-        self._tracker = _Tracker(dt)
-        self._tracker.start()
-        self.results = []
+        self._dt = dt
         self._entered = False
+        self._tracker = None
+        self.results = []
+
+    def _is_running(self):
+        return self._tracker is not None and self._tracker.is_alive()
 
     def _start_collect(self):
-        assert self._tracker.is_alive(), "Resource tracker is shutdown"
+        if not self._is_running():
+            self._tracker = _Tracker(self._dt)
+            self._tracker.start()
         self._tracker.parent_conn.send('collect')
 
     def _stop_collect(self):
-        if self._tracker.is_alive():
+        if self._is_running():
             self._tracker.parent_conn.send('send_data')
             self.results.extend(starmap(ResourceData, self._tracker.parent_conn.recv()))
 
     def __enter__(self):
-        self.clear()
         self._entered = True
+        self.clear()
         self._start_collect()
         return super(ResourceProfiler, self).__enter__()
 
     def __exit__(self, *args):
         self._entered = False
         self._stop_collect()
+        self.close()
         super(ResourceProfiler, self).__exit__(*args)
 
     def _start(self, dsk):
@@ -174,7 +180,9 @@ class ResourceProfiler(Callback):
 
     def close(self):
         """Shutdown the resource tracker process"""
-        self._tracker.shutdown()
+        if self._is_running():
+            self._tracker.shutdown()
+            self._tracker = None
 
     __del__ = close
 


Test patch:
diff --git a/dask/diagnostics/tests/test_profiler.py b/dask/diagnostics/tests/test_profiler.py
index 5c062640b..cf374d7a1 100644
--- a/dask/diagnostics/tests/test_profiler.py
+++ b/dask/diagnostics/tests/test_profiler.py
@@ -83,15 +83,20 @@ def test_resource_profiler():
     assert len(results) > 0
     assert all(isinstance(i, tuple) and len(i) == 3 for i in results)
 
+    # Tracker stopped on exit
+    assert not rprof._is_running()
+
     rprof.clear()
     assert rprof.results == []
 
+    # Close is idempotent
     rprof.close()
-    assert not rprof._tracker.is_alive()
+    assert not rprof._is_running()
 
-    with pytest.raises(AssertionError):
-        with rprof:
-            get(dsk, 'e')
+    # Restarts tracker if already closed
+    with rprof:
+        get(dsk2, 'c')
+    assert len(rprof.results) > 0
 
 
 @pytest.mark.skipif("not psutil")
@@ -114,7 +119,7 @@ def test_resource_profiler_multiple_gets():
     assert all(isinstance(i, tuple) and len(i) == 3 for i in results)
 
     rprof.close()
-    assert not rprof._tracker.is_alive()
+    assert not rprof._is_running()
 
 
 def test_cache_profiler():