Repository: dask/distributed

Issue:
Ensure one CI job runs without pyarrow but with numpy
P2P rechunking is an important new feature that requires numpy to be installed. We only have test jobs in our matrix that are installing pyarrow + numpy or neither but we should have a job running that only installs numpy.

This likely requires us to introduce a mindeps+numpy job since numpy is only a transitive dependency

Discussion:

Patch:
diff --git a/distributed/client.py b/distributed/client.py
index 2b47b000..4f3bc10b 100644
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -3500,8 +3500,8 @@ class Client(SyncMethodMixin):
 
     async def _restart_workers(
         self, workers: list[str], timeout: int | float | None = None
-    ):
-        results = await self.scheduler.broadcast(
+    ) -> dict[str, str]:
+        results: dict[str, str] = await self.scheduler.broadcast(
             msg={"op": "restart", "timeout": timeout}, workers=workers, nanny=True
         )
         timeout_workers = {
@@ -3511,8 +3511,11 @@ class Client(SyncMethodMixin):
             raise TimeoutError(
                 f"The following workers failed to restart with {timeout} seconds: {list(timeout_workers.keys())}"
             )
+        return results
 
-    def restart_workers(self, workers: list[str], timeout: int | float | None = None):
+    def restart_workers(
+        self, workers: list[str], timeout: int | float | None = None
+    ) -> dict[str, str]:
         """Restart a specified set of workers
 
         .. note::
@@ -3528,6 +3531,11 @@ class Client(SyncMethodMixin):
         timeout : int | float | None
             Number of seconds to wait
 
+        Returns
+        -------
+        dict[str, str]
+            Mapping of worker and restart status.
+
         Notes
         -----
         This method differs from :meth:`Client.restart` in that this method
diff --git a/distributed/spill.py b/distributed/spill.py
index e6789d50..47504ee5 100644
--- a/distributed/spill.py
+++ b/distributed/spill.py
@@ -248,6 +248,12 @@ class SpillBuffer(zict.Buffer):
         super().__delitem__(key)
         self.logged_pickle_errors.discard(key)
 
+    def pop(self, key: str, default: Any = None) -> Any:
+        raise NotImplementedError(
+            "Are you calling .pop(key, None) as a way to discard a key if it exists?"
+            "It may cause data to be read back from disk! Please use `del` instead."
+        )
+
     @property
     def memory(self) -> Mapping[str, Any]:
         """Key/value pairs stored in RAM. Alias of zict.Buffer.fast.
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
index bc09595c..9b63f65e 100644
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -1478,17 +1478,18 @@ class WorkerState:
         This is idempotent
         """
         logger.debug("Purge task: %s", ts)
-        key = ts.key
-        self.data.pop(key, None)
-        self.actors.pop(key, None)
+
+        # Do not use self.data.pop(key, None), as it could unspill the data from disk!
+        if ts.key in self.data:
+            del self.data[ts.key]
+        self.actors.pop(ts.key, None)
+        self.threads.pop(ts.key, None)
 
         for worker in ts.who_has:
             self.has_what[worker].discard(ts.key)
             self.data_needed[worker].discard(ts)
         ts.who_has.clear()
 
-        self.threads.pop(key, None)
-
         for d in ts.dependencies:
             ts.waiting_for_data.discard(d)
             d.waiters.discard(ts)


Test patch:
diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index 92b4489b..79a6620c 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -25,6 +25,7 @@ jobs:
         os: [ubuntu-latest, windows-latest, macos-latest]
         environment: [mindeps, "3.8", "3.9", "3.10", "3.11"]
         queuing: [queue]
+        extra_packages: [null]
         # Cherry-pick test modules to split the overall runtime roughly in half
         partition: [ci1, not ci1]
         exclude:
@@ -45,6 +46,15 @@ jobs:
             environment: "3.9"
             queuing: no_queue
             partition: "not ci1"
+          # P2P rechunking - requires numpy and nothing else beyond mindeps (no pyarrow)
+          - os: ubuntu-latest
+            environment: mindeps
+            extra_packages: [numpy]
+            partition: "ci1"
+          - os: ubuntu-latest
+            environment: mindeps
+            extra_packages: [numpy]
+            partition: "not ci1"
 
         # Uncomment to stress-test the test suite for random failures.
         # Must also change `export TEST_ID` in first step below.
@@ -118,6 +128,11 @@ jobs:
         run: |
           python -m pip install --no-deps -e .
 
+      - name: Extra Installs
+        if: ${{ matrix.extra_packages }}
+        shell: bash -l {0}
+        run: mamba install -y ${{ join(matrix.extra_packages, ' ') }}
+
       - name: mamba list
         shell: bash -l {0}
         run: mamba list
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
index f9dbb5e7..3e80abe0 100644
--- a/distributed/tests/test_client.py
+++ b/distributed/tests/test_client.py
@@ -4810,7 +4810,9 @@ async def test_restart_workers(c, s, a, b):
     assert await c.compute(x.sum()) == size
 
     # Restart a single worker
-    await c.restart_workers(workers=[a.worker_address])
+    a_worker_addr = a.worker_address
+    results = await c.restart_workers(workers=[a.worker_address])
+    assert results[a_worker_addr] == "OK"
     assert set(s.workers) == {a.worker_address, b.worker_address}
 
     # Make sure worker start times are as expected
diff --git a/distributed/tests/test_spill.py b/distributed/tests/test_spill.py
index 9c05b463..fd4def95 100644
--- a/distributed/tests/test_spill.py
+++ b/distributed/tests/test_spill.py
@@ -311,6 +311,12 @@ def test_spillbuffer_evict(tmp_path):
     assert_buf(buf, tmp_path, {"bad": bad}, {"a": a})
 
 
+def test_no_pop(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=100)
+    with pytest.raises(NotImplementedError):
+        buf.pop("x", None)
+
+
 class NoWeakRef:
     """A class which
     1. reports an arbitrary managed memory usage
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
index afe364ac..99c744ae 100644
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -210,7 +210,6 @@ def test_workerstate_fail_to_pickle_execute_1(ws_with_running_task):
     assert ws.tasks["x"].state == "error"
 
 
-@pytest.mark.xfail(reason="https://github.com/dask/distributed/issues/6705")
 def test_workerstate_fail_to_pickle_flight(ws):
     """Same as test_workerstate_fail_to_pickle_execute_1, but the task was
     computed on another host and for whatever reason it did not fail to pickle when it
@@ -247,6 +246,9 @@ def test_workerstate_fail_to_pickle_flight(ws):
     assert ws.tasks["x"].state == "error"
     assert ws.tasks["y"].state == "waiting"  # Not constrained
 
+    # FIXME https://github.com/dask/distributed/issues/6705
+    ws.validate = False
+
 
 @gen_cluster(
     client=True,
@@ -1207,3 +1209,24 @@ async def test_high_unmanaged_memory_warning(s, caplog):
         sum("Unmanaged memory use is high" in record.msg for record in caplog.records)
         == 1
     )  # Message is rate limited
+
+
+class WriteOnlyBuffer(UserDict):
+    def __getitem__(self, k):
+        raise AssertionError()
+
+
+@gen_cluster(client=True, nthreads=[("", 1)], worker_kwargs={"data": WriteOnlyBuffer})
+async def test_delete_spilled_keys(c, s, a):
+    """Test that freeing an in-memory key that has been spilled to disk does not
+    accidentally unspill it
+    """
+    x = c.submit(inc, 1, key="x")
+    await wait_for_state("x", "memory", a)
+    assert a.data.keys() == {"x"}
+    with pytest.raises(AssertionError):
+        a.data["x"]
+
+    x.release()
+    await async_wait_for(lambda: not a.data, timeout=2)
+    assert not a.state.tasks
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
index f5a7bc5b..db10434e 100644
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -1289,7 +1289,6 @@ def test_done_resumed_task_not_in_all_running_tasks(ws_with_running_task, done_e
     assert ts not in ws.all_running_tasks
 
 
-@pytest.mark.xfail(reason="https://github.com/dask/distributed/issues/6705")
 def test_gather_dep_failure(ws):
     """Simulate a task failing to unpickle when it reaches the destination worker after
     a flight.
@@ -1314,6 +1313,9 @@ def test_gather_dep_failure(ws):
     assert ws.transfer_incoming_count == 0
     assert ws.transfer_incoming_count_total == 1
 
+    # FIXME https://github.com/dask/distributed/issues/6705
+    ws.validate = False
+
 
 def test_transfer_incoming_metrics(ws):
     assert ws.transfer_incoming_bytes == 0
diff --git a/distributed/utils_test.py b/distributed/utils_test.py
index c2819afb..07b0e8ad 100644
--- a/distributed/utils_test.py
+++ b/distributed/utils_test.py
@@ -2411,7 +2411,8 @@ def ws():
     """An empty WorkerState"""
     state = WorkerState(address="127.0.0.1:1", transition_counter_max=50_000)
     yield state
-    state.validate_state()
+    if state.validate:
+        state.validate_state()
 
 
 @pytest.fixture(params=["executing", "long-running"])