Repository: dask/distributed

Issue:
Spill to disk may cause data duplication
In aggressive spill-to-disk scenarios I observed that distributed may spill all the data it has in memory while still complaining with the following message that there is no more data to spill

> Memory use is high but worker has no data "
                        "to store to disk.  Perhaps some other process "
                        "is leaking memory?
Side note: In our setup every worker runs in an isolated container so the chance of another process interfering with it is virtually zero.

It is true that these workers do indeed still hold on to a lot of memory without it being very transparent about where the data is. GC is useless, i.e. it is **not** related to https://github.com/dask/zict/issues/19

Investigating this issue let me realise that we may try to spill data which is actually currently in use.
The most common example for this is probably once the worker collected the dependencies and schedules the execution of the task, see [here](https://github.com/dask/distributed/blob/c8ee44bebe40e4db4495be1e2e0340fafbb41a37/distributed/worker.py#L2485). Another one would be when data is [requested form a worker](https://github.com/dask/distributed/blob/c8ee44bebe40e4db4495be1e2e0340fafbb41a37/distributed/worker.py#L1232) and we're still serializing/submitting it.
I could _prove_ this assumption by patching the worker code, tracking the keys in-execution/in-spilling with some logging and it turns out that for some jobs I hold on to multiple GBs of memory although it was supposedly already spilled.

If we spill data which is currently still in use, this is not only misleading to the user, since the data is still in memory, but it may also cause heavy data duplication if the spilled-but-still-in-memory dependency is requested by another worker since the original data is still in memory but the buffer would fetch the key from the slow store and materialise it a second time since it doesn't know of the original object anymore. If this piece of data is requested by more than one worker, the data could be duplicated multiple times even.

In non-aggressive spill-to-disk scenarios we should be protected from this by the LRU in the buffer but if memory pressure is high, the spilling might actually worsen the situation in these cases.

My practical approach to this would be to introduce something like a `(un)lock_key_in_store` method to the buffer which protects the key from spilling and manually (un)setting this in the distributed code. If there is a smarter approach, I'd be glad to hear about it.

Also, If my reasoning is somewhere flawed, I'd appreciate feedback since, so far, I could only prove that we try to spill data currently in use but, so far, the duplication is just theory.

Related issues:
https://github.com/dask/dask/issues/2456


Discussion:
I think that your understanding that we can both spill and hold onto data is correct.  I would expect this inconsistency to be short-lived though, only as long as the task execution or data communication.  

Given the LRU policy, I would expect that this inconsistency would not happen until all other data had been spilled.  At this point my guess is that the impact of this inconsistency is usually small (the workload was probably going to run out of memory anyway) but that's just a guess.  It looks like you have more hands-on experience here.
> I would expect this inconsistency to be short-lived though, only as long as the task execution or data communication.

I would also think that it is short lived. I don't think we actually hold on permanently to references

> Given the LRU policy, I would expect that this inconsistency would not happen until all other data had been spilled. At this point my guess is that the impact of this inconsistency is usually small (the workload was probably going to run out of memory anyway) but that's just a guess. It looks like you have more hands-on experience here.

Yes, this can only happen once the LRU is essentially depleted.

I'll try to collect a bit more data to see if this actually has a huge impact. From what I've seen, a very busy cluster can communicate/execute for minutes. We may need to decrease task sizes but still, if the network is busy we cannot rely on very fast transfers. A minute or two with strong data duplication may already be enough to drive workers OOM. However, if it is marginal the complexity is hardly worth the effort.
Following up with a few observations:


# Context

The computation I'm running in this example is a huge `groupby.apply` where the function reduces to essentially no data.
I estimate the total data in memory to be here about ~1.5TB of in-memory pandas dataframe data. The cluster I'm running this on has 50 workers, 1CPU each with 32GB RAM each, in total that makes 1.6TB RAM. Obviously, for this to succeed we need to rely on spill-to-disk.

The dask version I am using is based on https://github.com/dask/dask/pull/6154 since otherwise the spill-to-disk effectively is useless and first workers go OOM after about 10mins.

The distributed version is more or less 2.14.0 with some additional patches to observe and count the keys duplicated by _get_data_ and _execute_

# Observations

The following is a grafana dashboard showing the task states, CPU and memory usage statistics of the workers. The memory usage per worker remains consistently stable due to passive spill-to-disk (result is finished, put into buffer, buffer evicts results if necessary to maintain the threshold). Once the first user functions are evaluated (at around 18:30) we can see that the memory profile is a bit more volatile since the function itself collects huge chunks of data and works on these. This is also the time where we can see first logs about active spill-to-disk by the memory-monitor (70% threshold hit, evict data until back at target)

![image](https://user-images.githubusercontent.com/8629629/80964239-c8573200-8e10-11ea-886b-4ed67cb60307.png)

The following plot shows all data currently held by the workers which is not part of the `self.data.fast` dictionary. This is either data currently used by executing code or data currently submitting but not part of the fast dict (data based on logs emitted by a PeriodicCallback which compares the keys-in-use and keys-in-store). The left axis (blue) shows the number of keys currently in memory but not tracked while the right hand side shows the memory used by these keys, normalised to the worker memory limit.
We can see that even early on, already up to 18% of a workers memory is occupied by this untracked memory. Especially at a later point in the computation this peaks at about 40% of the memory limit, with duplication factors up to 11.

![image](https://user-images.githubusercontent.com/8629629/80979459-a917cf00-8e27-11ea-897e-b57d330a51c3.png)

# Summary

I was surprised to see that this duplication kicks in even at early stages of the pipeline and is more or less continuously an issue over the entire run. The impact is mostly below 5% memory which is probably not worth the effort but that it peaks at 40% is a bit concerning.

In the past I've often seen workers with up to 80% mem usage while reporting the "Memory use is high but worker has no data" and now I'm wondering if this is/was the issue (in combination to in-flight data but that's another topic).

Is this something we should address or am I paranoid?
I wonder whether this is still the case with PR ( https://github.com/dask/distributed/pull/3849 ). For clarity, that would impact spilling that uses pickle in addition to data transmission. IOW out-of-band buffers would be collected, but not copied. Meaning one would expect the overall memory usage to be lower when pickling objects that support out-of-band buffer extraction. This is because the same memory could be shared between the spilled and non-spilled forms.
Once the data is spilled and the reference to it is lost, there is no way to connect the serialized object to its original, is there? I'm thinking of a scenario like this

* task finishes; put data of `key` in buffer on worker A
* spill `key` to disk on worker A since memory threshold was hit
* **Loose reference and free memory**
* Data for `key` is requested by worker B from worker A but worker A is at memory limit, i.e. `fast` data is "full"
* Worker A still loads data into memory *but* it doesn't store it in the fast dict since the memory threshold was breached. Therefore we get an entirely new copy without central reference to the object; Let's call this `key1`
* In the meantime Worker C asks for the same key and the same thing happens but in a different coroutine. Since the data was never registered in the fast dict, the worker doesn't know of it and loads the data from disk, deserializes it and we have an entirely new copy `key2`
* Now worker A wants to execute some code and requires the data of `key` but the data is still not available (data submission of the above two tasks not completed, yet). It loads data again from disk... now we have `key3`

In this scenario we have three _distinct_ copies of the same data, don't we? In my naive understanding, pickle5 would help if we were trying to spill/submit/pickle concurrently but I don't believe this is causing this issue.
Just for clarity, what I'm meaning is that data pickled with protocol 5 will take views of the data as part of pickling (as opposed to copying the data into `bytes`). Here's a code example for context:

```python
import pickle5 as pickle
import numpy as np

a = np.arange(6)
b = []
d = pickle.dumps(a, protocol=5, buffer_callback=b.append)

assert np.may_share_memory(a, b[0])
```

This would mean the high-water mark would remain a bit lower during the run. Though it's worth noting compression during spilling may affect this as well. So more exploration of exactly how this would affect things is still needed.

That all being said, I've not dug into this issue deeply. Perhaps pickle protocol 5 is not relevant.
There may also be other explanations for high memory usage. For example `merge_frames` (part of the serialization/deserialization pipeline) was causing data to get copied in some cases. PR ( https://github.com/dask/distributed/pull/3960 ) fixes that issue.
@gjoseph92 am wondering if PR ( https://github.com/dask/distributed/pull/5208 ) would help here as well. Do you know? If not, no worries (mostly curious).
I have no idea unfortunately. As I understood, this issue was more about a bookkeeping error, so I'd doubt it. It might help make the case where we re-materialize data in order to send it to a worker slightly less bad, since at least there wouldn't be _another_ copy during the frames merge, but there would still be the data re-read from disk either way.
From my understanding the zero-copy fix doesn't help since the duplication I am referring to here is not caused by memcopies but rather by having distinct python objects. I'll try to explain with some simplified pseudo worker/zict buffer

```python
class ZictSpillData:
    def __getitem__(self, key):
        if key in self.memory:
            return self.memory[key]
        else:
            data =  load_from_store(key)
            if sizeof(data) < self.memory_target:
                self.memory[key] = data
            return data


class Worker:
    async def execute(self, ts, ...):
        assert isinstance(self.data, ZictSpillData)
        data = {}
        for key in ts.dependencies:
            data[key] = self.data[key]
            assert key not in self.data.memory  # data was too large, the buffer forgot it again
        run_task(ts, data)
    
    def get_data(self, keys, ...):
        data = {}
        for key in ts.dependencies:
            data[key] = self.data[key]
            assert key not in self.data.memory  # data was too large, the buffer forgot it again
        return data
```

Let's assume, for simplicity, the data we're concerned about is large and will be spilled by the buffer immediately because it is beyond a given threshold `memory_target = memory_limit * memory_target_fraction`. In this case, whenever the data is accessed, the buffer will read it from storage and will create a new data instance without storing a ref to it in it's data. This instance is returned to the caller and every subsequent call to the buffer will do the exact same thing. Load data, create new instance, return without holding a ref.

In this extreme situation this is hopefully easy to follow. In more realisitc scenarios we're spilling data concurrently to other operations. For instance, data is being used as part of an execute but the buffer spills the data. This results in the data no longer being tracked by the buffer (in our new terminology this means the data is "unmanaged memory") but it is still in memory since it is being used by the task execution. If another worker then requests this piece of data we just spilled, it will load it from store and create a new instance and will start serialization. Simultaneously X other workers could do the same, resulting in X (get_data) + 1 (executing) copies of the same data because our buffer doesn't know about the keys still being used.


While typing, I'm wondering if we can't fix this by being a bit smart with weakrefs but I currently don't have time to test this idea.

Edit: Just in case I haven't made myself clear, yet. I believe this issue is caused by the `zict.Buffer`, not our serialization protocol or any other thing distributed is doing. This issue might be better off in the zict repo but I believe the visibility here is more important (and we have "our own" `Buffer` subclass, by now, anyways https://github.com/dask/distributed/blob/8c73a18ba5c2ffe61977ea936da8ffeacb815c61/distributed/spill.py#L16)
A **very naive** worst case calculation about how severe this issue is.

We're targeting `get_data` requests of about 50MB (`Worker.target_message_size`). Let's take this as a reference size for a keys size. In realistic workloads I wouldn't be surprised to see much larger data keys. The original report here described a shuffle for which 50MB splits are reasonable (input 2GB / 32 splits ~ 62MB).

The worker limits concurrent incoming `get_data` requests to `Worker.total_in_connections == 10 (default)` but allows this to double for same-host connections, i.e. `#key_size * #total_in_commections * #same_host_mult`

That yields a 500-1000MB data duplication for `get_data` assuming the LRU is missed consistently which I assume is approximately true for a shuffle.

The duplication on `Worker.execute` side scales linearly with the number of threads and the number of input dependents, i.e. `#key_size * #dependencies * #num_threads`. Again using 50MB per key and assuming 32 dependencies for a shuffle, that yields `1600MB` per thread.

Let's have a look at a worst case, all workers on the same machine, max branching of 32 in a shuffle, key size at 50MB and 2 threads per worker (the two is pretty random, I have to admit), that's up to 84 duplications (10 * 2 + 32 * 2) of a keys data or `4200MB` unmanaged memory caused by our spillage layer. various data copies and overhead of the UDF is on top of that.

I guess this worst case is not too far off from a LocalCluster "deployment" on a big VM. However, most of the numbers are mere guessing and this can easily be scaled up or down so please take this with a healthy pinch of salt.
Either way, duplicates also mean duplicate effort in deserialization and disk IO so there is a case to be made to avoid these regardless of the memory footprint.
This looks like something that could be solved with weakref and a few extra lines in spill.py. I'll take a stab at it after #5543 is merged.
> This looks like something that could be solved with weakref and a few extra lines in spill.py.

That's what I thought as well.

Resolving patch:
diff --git a/distributed/spill.py b/distributed/spill.py
--- a/distributed/spill.py
+++ b/distributed/spill.py
@@ -2,7 +2,7 @@
 
 import logging
 import time
-from collections.abc import Mapping
+from collections.abc import Mapping, MutableMapping
 from contextlib import contextmanager
 from functools import partial
 from typing import Any, Literal, NamedTuple, cast
@@ -14,7 +14,9 @@
 from distributed.sizeof import safe_sizeof
 
 logger = logging.getLogger(__name__)
-has_zict_210 = parse_version(zict.__version__) > parse_version("2.0.0")
+has_zict_210 = parse_version(zict.__version__) >= parse_version("2.1.0")
+# At the moment of writing, zict 2.2.0 has not been released yet. Support git tip.
+has_zict_220 = parse_version(zict.__version__) >= parse_version("2.2.0.dev2")
 
 
 class SpilledSize(NamedTuple):
@@ -38,7 +40,7 @@ class SpillBuffer(zict.Buffer):
     the total size of the stored data exceeds the target. If max_spill is provided the
     key/value pairs won't be spilled once this threshold has been reached.
 
-    Paramaters
+    Parameters
     ----------
     spill_directory: str
         Location on disk to write the spill files to
@@ -63,14 +65,15 @@ def __init__(
     ):
 
         if max_spill is not False and not has_zict_210:
-            raise ValueError("zict > 2.0.0 required to set max_weight")
+            raise ValueError("zict >= 2.1.0 required to set max-spill")
 
-        super().__init__(
-            fast={},
-            slow=Slow(spill_directory, max_spill),
-            n=target,
-            weight=_in_memory_weight,
-        )
+        slow: MutableMapping[str, Any] = Slow(spill_directory, max_spill)
+        if has_zict_220:
+            # If a value is still in use somewhere on the worker since the last time it
+            # was unspilled, don't duplicate it
+            slow = zict.Cache(slow, zict.WeakValueMapping())
+
+        super().__init__(fast={}, slow=slow, n=target, weight=_in_memory_weight)
         self.last_logged = 0
         self.min_log_interval = min_log_interval
         self.logged_pickle_errors = set()  # keys logged with pickle error
@@ -204,7 +207,8 @@ def spilled_total(self) -> SpilledSize:
         The two may differ substantially, e.g. if sizeof() is inaccurate or in case of
         compression.
         """
-        return cast(Slow, self.slow).total_weight
+        slow = cast(zict.Cache, self.slow).data if has_zict_220 else self.slow
+        return cast(Slow, slow).total_weight
 
 
 def _in_memory_weight(key: str, value: Any) -> int:
@@ -224,6 +228,7 @@ class HandledError(Exception):
     pass
 
 
+# zict.Func[str, Any] requires zict >= 2.2.0
 class Slow(zict.Func):
     max_weight: int | Literal[False]
     weight_by_key: dict[str, SpilledSize]


Test patch:
diff --git a/distributed/tests/test_spill.py b/distributed/tests/test_spill.py
--- a/distributed/tests/test_spill.py
+++ b/distributed/tests/test_spill.py
@@ -1,20 +1,27 @@
 from __future__ import annotations
 
+import gc
 import logging
 import os
 
 import pytest
 
-zict = pytest.importorskip("zict")
-from packaging.version import parse as parse_version
-
 from dask.sizeof import sizeof
 
 from distributed.compatibility import WINDOWS
 from distributed.protocol import serialize_bytelist
-from distributed.spill import SpillBuffer
+from distributed.spill import SpillBuffer, has_zict_210, has_zict_220
 from distributed.utils_test import captured_logger
 
+requires_zict_210 = pytest.mark.skipif(
+    not has_zict_210,
+    reason="requires zict version >= 2.1.0",
+)
+requires_zict_220 = pytest.mark.skipif(
+    not has_zict_220,
+    reason="requires zict version >= 2.2.0",
+)
+
 
 def psize(*objs) -> tuple[int, int]:
     return (
@@ -23,15 +30,30 @@ def psize(*objs) -> tuple[int, int]:
     )
 
 
+def assert_buf(buf: SpillBuffer, expect_fast: dict, expect_slow: dict) -> None:
+    # assertions on fast
+    assert dict(buf.fast) == expect_fast
+    assert buf.fast.weights == {k: sizeof(v) for k, v in expect_fast.items()}
+    assert buf.fast.total_weight == sum(sizeof(v) for v in expect_fast.values())
+    for k, v in buf.fast.items():
+        assert buf[k] is v
+
+    # assertions on slow
+    assert set(buf.slow) == expect_slow.keys()
+    slow = buf.slow.data if has_zict_220 else buf.slow  # type: ignore
+    assert slow.weight_by_key == {k: psize(v) for k, v in expect_slow.items()}
+    total_weight = psize(*expect_slow.values())
+    assert slow.total_weight == total_weight
+    assert buf.spilled_total == total_weight
+
+
 def test_spillbuffer(tmpdir):
     buf = SpillBuffer(str(tmpdir), target=300)
     # Convenience aliases
     assert buf.memory is buf.fast
     assert buf.disk is buf.slow
 
-    assert not buf.slow.weight_by_key
-    assert buf.slow.total_weight == (0, 0)
-    assert buf.spilled_total == (0, 0)
+    assert_buf(buf, {}, {})
 
     a, b, c, d = "a" * 100, "b" * 99, "c" * 98, "d" * 97
 
@@ -40,75 +62,48 @@ def test_spillbuffer(tmpdir):
     assert psize(a)[0] != psize(a)[1]
 
     buf["a"] = a
-    assert not buf.slow
-    assert buf.fast.weights == {"a": sizeof(a)}
-    assert buf.fast.total_weight == sizeof(a)
-    assert buf.slow.weight_by_key == {}
-    assert buf.slow.total_weight == (0, 0)
+    assert_buf(buf, {"a": a}, {})
     assert buf["a"] == a
 
     buf["b"] = b
-    assert not buf.slow
-    assert not buf.slow.weight_by_key
-    assert buf.slow.total_weight == (0, 0)
+    assert_buf(buf, {"a": a, "b": b}, {})
 
     buf["c"] = c
-    assert set(buf.slow) == {"a"}
-    assert buf.slow.weight_by_key == {"a": psize(a)}
-    assert buf.slow.total_weight == psize(a)
+    assert_buf(buf, {"b": b, "c": c}, {"a": a})
 
     assert buf["a"] == a
-    assert set(buf.slow) == {"b"}
-    assert buf.slow.weight_by_key == {"b": psize(b)}
-    assert buf.slow.total_weight == psize(b)
+    assert_buf(buf, {"a": a, "c": c}, {"b": b})
 
     buf["d"] = d
-    assert set(buf.slow) == {"b", "c"}
-    assert buf.slow.weight_by_key == {"b": psize(b), "c": psize(c)}
-    assert buf.slow.total_weight == psize(b, c)
+    assert_buf(buf, {"a": a, "d": d}, {"b": b, "c": c})
 
     # Deleting an in-memory key does not automatically move spilled keys back to memory
     del buf["a"]
-    assert set(buf.slow) == {"b", "c"}
-    assert buf.slow.weight_by_key == {"b": psize(b), "c": psize(c)}
-    assert buf.slow.total_weight == psize(b, c)
+    assert_buf(buf, {"d": d}, {"b": b, "c": c})
     with pytest.raises(KeyError):
         buf["a"]
 
     # Deleting a spilled key updates the metadata
     del buf["b"]
-    assert set(buf.slow) == {"c"}
-    assert buf.slow.weight_by_key == {"c": psize(c)}
-    assert buf.slow.total_weight == psize(c)
+    assert_buf(buf, {"d": d}, {"c": c})
     with pytest.raises(KeyError):
         buf["b"]
 
     # Updating a spilled key moves it to the top of the LRU and to memory
-    buf["c"] = c * 2
-    assert set(buf.slow) == {"d"}
-    assert buf.slow.weight_by_key == {"d": psize(d)}
-    assert buf.slow.total_weight == psize(d)
+    c2 = c * 2
+    buf["c"] = c2
+    assert_buf(buf, {"c": c2}, {"d": d})
 
     # Single key is larger than target and goes directly into slow
     e = "e" * 500
 
     buf["e"] = e
-    assert set(buf.slow) == {"d", "e"}
-    assert buf.slow.weight_by_key == {"d": psize(d), "e": psize(e)}
-    assert buf.slow.total_weight == psize(d, e)
+    assert_buf(buf, {"c": c2}, {"d": d, "e": e})
 
     # Updating a spilled key with another larger than target updates slow directly
     d = "d" * 500
     buf["d"] = d
-    assert set(buf.slow) == {"d", "e"}
-    assert buf.slow.weight_by_key == {"d": psize(d), "e": psize(e)}
-    assert buf.slow.total_weight == psize(d, e)
-
-
-requires_zict_210 = pytest.mark.skipif(
-    parse_version(zict.__version__) <= parse_version("2.0.0"),
-    reason="requires zict version > 2.0.0",
-)
+    assert_buf(buf, {"c": c2}, {"d": d, "e": e})
 
 
 @requires_zict_210
@@ -120,32 +115,17 @@ def test_spillbuffer_maxlim(tmpdir):
     # size of a is bigger than target and is smaller than max_spill;
     # key should be in slow
     buf["a"] = a
-    assert not buf.fast
-    assert not buf.fast.weights
-    assert set(buf.slow) == {"a"}
-    assert buf.slow.weight_by_key == {"a": psize(a)}
-    assert buf.slow.total_weight == psize(a)
+    assert_buf(buf, {}, {"a": a})
     assert buf["a"] == a
 
     # size of b is smaller than target key should be in fast
     buf["b"] = b
-    assert set(buf.fast) == {"b"}
-    assert buf.fast.weights == {"b": sizeof(b)}
-    assert buf["b"] == b
-    assert buf.fast.total_weight == sizeof(b)
+    assert_buf(buf, {"b": b}, {"a": a})
 
     # size of c is smaller than target but b+c > target, c should stay in fast and b
     # move to slow since the max_spill limit has not been reached yet
-
     buf["c"] = c
-    assert set(buf.fast) == {"c"}
-    assert buf.fast.weights == {"c": sizeof(c)}
-    assert buf["c"] == c
-    assert buf.fast.total_weight == sizeof(c)
-
-    assert set(buf.slow) == {"a", "b"}
-    assert buf.slow.weight_by_key == {"a": psize(a), "b": psize(b)}
-    assert buf.slow.total_weight == psize(a, b)
+    assert_buf(buf, {"c": c}, {"a": a, "b": b})
 
     # size of e < target but e+c > target, this will trigger movement of c to slow
     # but the max spill limit prevents it. Resulting in e remaining in fast
@@ -154,15 +134,7 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["e"] = e
 
     assert "disk reached capacity" in logs_e.getvalue()
-
-    assert set(buf.fast) == {"c", "e"}
-    assert buf.fast.weights == {"c": sizeof(c), "e": sizeof(e)}
-    assert buf["e"] == e
-    assert buf.fast.total_weight == sizeof(c) + sizeof(e)
-
-    assert set(buf.slow) == {"a", "b"}
-    assert buf.slow.weight_by_key == {"a": psize(a), "b": psize(b)}
-    assert buf.slow.total_weight == psize(a, b)
+    assert_buf(buf, {"c": c, "e": e}, {"a": a, "b": b})
 
     # size of d > target, d should go to slow but slow reached the max_spill limit then
     # d will end up on fast with c (which can't be move to slow because it won't fit
@@ -171,15 +143,7 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["d"] = d
 
     assert "disk reached capacity" in logs_d.getvalue()
-
-    assert set(buf.fast) == {"c", "d", "e"}
-    assert buf.fast.weights == {"c": sizeof(c), "d": sizeof(d), "e": sizeof(e)}
-    assert buf["d"] == d
-    assert buf.fast.total_weight == sizeof(c) + sizeof(d) + sizeof(e)
-
-    assert set(buf.slow) == {"a", "b"}
-    assert buf.slow.weight_by_key == {"a": psize(a), "b": psize(b)}
-    assert buf.slow.total_weight == psize(a, b)
+    assert_buf(buf, {"c": c, "d": d, "e": e}, {"a": a, "b": b})
 
     # Overwrite a key that was in slow, but the size of the new key is larger than
     # max_spill
@@ -191,11 +155,7 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["a"] = a_large
 
     assert "disk reached capacity" in logs_alarge.getvalue()
-
-    assert set(buf.fast) == {"a", "d", "e"}
-    assert set(buf.slow) == {"b", "c"}
-    assert buf.fast.total_weight == sizeof(d) + sizeof(a_large) + sizeof(e)
-    assert buf.slow.total_weight == psize(b, c)
+    assert_buf(buf, {"a": a_large, "d": d, "e": e}, {"b": b, "c": c})
 
     # Overwrite a key that was in fast, but the size of the new key is larger than
     # max_spill
@@ -205,11 +165,7 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["d"] = d_large
 
     assert "disk reached capacity" in logs_dlarge.getvalue()
-
-    assert set(buf.fast) == {"a", "d", "e"}
-    assert set(buf.slow) == {"b", "c"}
-    assert buf.fast.total_weight == sizeof(a_large) + sizeof(d_large) + sizeof(e)
-    assert buf.slow.total_weight == psize(b, c)
+    assert_buf(buf, {"a": a_large, "d": d_large, "e": e}, {"b": b, "c": c})
 
 
 class MyError(Exception):
@@ -241,13 +197,12 @@ def test_spillbuffer_fail_to_serialize(tmpdir):
 
     # spill.py must remain silent because we're already logging in worker.py
     assert not logs_bad_key.getvalue()
-    assert not set(buf.fast)
-    assert not set(buf.slow)
+    assert_buf(buf, {}, {})
 
     b = Bad(size=100)  # this is small enough to fit in memory/fast
 
     buf["b"] = b
-    assert set(buf.fast) == {"b"}
+    assert_buf(buf, {"b": b}, {})
 
     c = "c" * 100
     with captured_logger(logging.getLogger("distributed.spill")) as logs_bad_key_mem:
@@ -259,9 +214,7 @@ def test_spillbuffer_fail_to_serialize(tmpdir):
     logs_value = logs_bad_key_mem.getvalue()
     assert "Failed to pickle" in logs_value  # from distributed.spill
     assert "Traceback" in logs_value  # from distributed.spill
-    assert set(buf.fast) == {"b", "c"}
-    assert buf.fast.total_weight == sizeof(b) + sizeof(c)
-    assert not set(buf.slow)
+    assert_buf(buf, {"b": b, "c": c}, {})
 
 
 @requires_zict_210
@@ -279,8 +232,7 @@ def test_spillbuffer_oserror(tmpdir):
     # let's have something in fast and something in slow
     buf["a"] = a
     buf["b"] = b
-    assert set(buf.fast) == {"b"}
-    assert set(buf.slow) == {"a"}
+    assert_buf(buf, {"b": b}, {"a": a})
 
     # modify permissions of disk to be read only.
     # This causes writes to raise OSError, just like in case of disk full.
@@ -291,15 +243,10 @@ def test_spillbuffer_oserror(tmpdir):
         buf["c"] = c
 
     assert "Spill to disk failed" in logs_oserror_slow.getvalue()
-    assert set(buf.fast) == {"b", "c"}
-    assert set(buf.slow) == {"a"}
-
-    assert buf.slow.weight_by_key == {"a": psize(a)}
-    assert buf.fast.weights == {"b": sizeof(b), "c": sizeof(c)}
+    assert_buf(buf, {"b": b, "c": c}, {"a": a})
 
     del buf["c"]
-    assert set(buf.fast) == {"b"}
-    assert set(buf.slow) == {"a"}
+    assert_buf(buf, {"b": b}, {"a": a})
 
     # add key to fast which is smaller than target but when added it triggers spill,
     # which triggers OSError
@@ -307,40 +254,26 @@ def test_spillbuffer_oserror(tmpdir):
         buf["d"] = d
 
     assert "Spill to disk failed" in logs_oserror_evict.getvalue()
-    assert set(buf.fast) == {"b", "d"}
-    assert set(buf.slow) == {"a"}
-
-    assert buf.slow.weight_by_key == {"a": psize(a)}
-    assert buf.fast.weights == {"b": sizeof(b), "d": sizeof(d)}
+    assert_buf(buf, {"b": b, "d": d}, {"a": a})
 
 
 @requires_zict_210
 def test_spillbuffer_evict(tmpdir):
     buf = SpillBuffer(str(tmpdir), target=300, min_log_interval=0)
 
-    a_bad = Bad(size=100)
+    bad = Bad(size=100)
     a = "a" * 100
 
     buf["a"] = a
-
-    assert set(buf.fast) == {"a"}
-    assert not set(buf.slow)
-    assert buf.fast.weights == {"a": sizeof(a)}
+    assert_buf(buf, {"a": a}, {})
 
     # successful eviction
     weight = buf.evict()
     assert weight == sizeof(a)
+    assert_buf(buf, {}, {"a": a})
 
-    assert not buf.fast
-    assert set(buf.slow) == {"a"}
-    assert buf.slow.weight_by_key == {"a": psize(a)}
-
-    buf["a_bad"] = a_bad
-
-    assert set(buf.fast) == {"a_bad"}
-    assert buf.fast.weights == {"a_bad": sizeof(a_bad)}
-    assert set(buf.slow) == {"a"}
-    assert buf.slow.weight_by_key == {"a": psize(a)}
+    buf["bad"] = bad
+    assert_buf(buf, {"bad": bad}, {"a": a})
 
     # unsuccessful eviction
     with captured_logger(logging.getLogger("distributed.spill")) as logs_evict_key:
@@ -349,7 +282,63 @@ def test_spillbuffer_evict(tmpdir):
 
     assert "Failed to pickle" in logs_evict_key.getvalue()
     # bad keys stays in fast
-    assert set(buf.fast) == {"a_bad"}
-    assert buf.fast.weights == {"a_bad": sizeof(a_bad)}
-    assert set(buf.slow) == {"a"}
-    assert buf.slow.weight_by_key == {"a": psize(a)}
+    assert_buf(buf, {"bad": bad}, {"a": a})
+
+
+class SupportsWeakRef:
+    def __init__(self, n):
+        self.n = n
+
+    def __sizeof__(self):
+        return self.n
+
+
+class NoWeakRef:
+    __slots__ = ("n",)
+
+    def __init__(self, n):
+        self.n = n
+
+    def __sizeof__(self):
+        return self.n
+
+
+@pytest.mark.parametrize(
+    "cls,expect_cached",
+    [
+        (SupportsWeakRef, has_zict_220),
+        (NoWeakRef, False),
+    ],
+)
+@pytest.mark.parametrize("size", [60, 110])
+def test_weakref_cache(tmpdir, cls, expect_cached, size):
+    buf = SpillBuffer(str(tmpdir), target=100)
+
+    # Run this test twice:
+    # - x is smaller than target and is evicted by y;
+    # - x is individually larger than target and it never touches fast
+    x = cls(size)
+    buf["x"] = x
+    if size < 100:
+        buf["y"] = cls(60)  # spill x
+    assert "x" in buf.slow
+
+    # Test that we update the weakref cache on setitem
+    assert (buf["x"] is x) == expect_cached
+
+    id_x = id(x)
+    del x
+    gc.collect()  # Only needed on pypy
+
+    if size < 100:
+        buf["y"]
+    assert "x" in buf.slow
+
+    x2 = buf["x"]
+    assert id(x2) != id_x
+    if size < 100:
+        buf["y"]
+    assert "x" in buf.slow
+
+    # Test that we update the weakref cache on getitem
+    assert (buf["x"] is x2) == expect_cached
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -16,7 +16,6 @@
 
 import psutil
 import pytest
-from packaging.version import parse as parse_version
 from tlz import first, pluck, sliding_window
 
 import dask
@@ -43,6 +42,7 @@
 from distributed.metrics import time
 from distributed.protocol import pickle
 from distributed.scheduler import Scheduler
+from distributed.spill import has_zict_210
 from distributed.utils import TimeoutError
 from distributed.utils_test import (
     TaskStateMetadataPlugin,
@@ -63,15 +63,9 @@
 
 pytestmark = pytest.mark.ci1
 
-try:
-    import zict
-except ImportError:
-    zict = None  # type: ignore
-
-requires_zict = pytest.mark.skipif(not zict, reason="requires zict")
 requires_zict_210 = pytest.mark.skipif(
-    not zict or parse_version(zict.__version__) <= parse_version("2.0.0"),
-    reason="requires zict version > 2.0.0",
+    not has_zict_210,
+    reason="requires zict version >= 2.1.0",
 )
 
 
@@ -924,7 +918,6 @@ async def assert_basic_futures(c: Client) -> None:
     assert results == list(map(inc, range(10)))
 
 
-@requires_zict
 @gen_cluster(client=True)
 async def test_fail_write_to_disk_target_1(c, s, a, b):
     """Test failure to spill triggered by key which is individually larger
@@ -942,7 +935,6 @@ async def test_fail_write_to_disk_target_1(c, s, a, b):
     await assert_basic_futures(c)
 
 
-@requires_zict
 @gen_cluster(
     client=True,
     nthreads=[("", 1)],
@@ -965,10 +957,8 @@ async def test_fail_write_to_disk_target_2(c, s, a):
 
     y = c.submit(lambda: "y" * 256, key="y")
     await wait(y)
-    if parse_version(zict.__version__) <= parse_version("2.0.0"):
-        assert set(a.data.memory) == {"y"}
-    else:
-        assert set(a.data.memory) == {"x", "y"}
+
+    assert set(a.data.memory) == {"x", "y"} if has_zict_210 else {"y"}
     assert not a.data.disk
 
     await assert_basic_futures(c)
@@ -1187,7 +1177,6 @@ async def test_statistical_profiling_2(c, s, a, b):
             break
 
 
-@requires_zict
 @gen_cluster(
     client=True,
     nthreads=[("", 1)],
@@ -1277,7 +1266,6 @@ async def test_spill_constrained(c, s, w):
     assert set(w.data.disk) == {x.key}
 
 
-@requires_zict
 @gen_cluster(
     nthreads=[("", 1)],
     client=True,
@@ -1301,7 +1289,6 @@ async def test_spill_spill_threshold(c, s, a):
     assert await x == 1
 
 
-@requires_zict
 @pytest.mark.parametrize(
     "memory_target_fraction,managed,expect_spilled",
     [