Repository: dask/distributed

Issue:
`SpillBuffer.spilled_total` appears to return incorrect results 
When spilling data to disk, the `SpillBuffer` appears to return the incorrect size of the data on disk. For example, when spilling an ~8 MB random matrix onto disk, the `data` file created by the `SpillBuffer` is also ~8 MB in size, yet the SpillBuffer only returns ~1 MB.

Reproducer:
```python3
from __future__ import annotations

import os
import tempfile

import numpy as np

from distributed.spill import SpillBuffer


def test_spill_size():
    tmpdir = tempfile.mkdtemp()
    buf = SpillBuffer(spill_directory=tmpdir, target=0, max_spill=False)
    data = np.random.random((1024, 1024))
    buf["data"] = data
    spill_size = os.stat(os.path.join(tmpdir,  "data")).st_size
    assert buf.spilled_total.disk == spill_size
```

fails with
```
>       assert buf.spilled_total.disk == spill_size
E       assert 1048808 == 8388840
E        +  where 1048808 = SpilledSize(memory=8388608, disk=1048808).disk
E        +    where SpilledSize(memory=8388608, disk=1048808) = Buffer<<LRU: 0/0 on dict>, <zict.cache.Cache object at 0x11f2413d0>>.spilled_total
```


Discussion:
A couple of comments here, it looks a bit weird to me that the disk size calculated as `spill_size` is very close to the value we are calculating as the `buf.spilled_total.memory`. I would expect the `spill_size` to be bigger.  However, I wonder if this is related to how we calculate spilled_total. See:
https://github.com/dask/distributed/blob/55cc1a50f6fc7c10815bed93681eba970c9b63f0/distributed/spill.py#L30-L31



The size of spill_size feels fine to me. Just looking at the way we serialize things, the overhead for this array appears to be 232 bytes (`sys.getsizeof(serialize(data)[0])`), which matches perfectly. There might be a clue in the fact that `spill_size` is always _almost_ 8x the size of `buf.spilled_total.disk`, regardless of the size of `data`.
It looks like the calculation of the memory size is wrong:
https://github.com/dask/distributed/blob/55cc1a50f6fc7c10815bed93681eba970c9b63f0/distributed/spill.py#L290
uses `len` but we should probably use `nbytes` on `memoryview` objects (https://docs.python.org/3/library/stdtypes.html#memoryview.nbytes). I'll file a PR for this.
That line was suggested by Guido when we were working on this. I'm sure he'll have an idea of what could be happening here. I know he is on PTO, but he might be a good person to review this. He will be back next week. 

Resolving patch:
diff --git a/distributed/spill.py b/distributed/spill.py
--- a/distributed/spill.py
+++ b/distributed/spill.py
@@ -287,7 +287,10 @@ def __setitem__(self, key: str, value: Any) -> None:
             # which will then unwrap it.
             raise PickleError(key, e)
 
-        pickled_size = sum(len(frame) for frame in pickled)
+        pickled_size = sum(
+            frame.nbytes if isinstance(frame, memoryview) else len(frame)
+            for frame in pickled
+        )
 
         if has_zict_210:
             # Thanks to Buffer.__setitem__, we never update existing


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,8 +1,11 @@
 from __future__ import annotations
 
+import array
 import logging
 import os
+import random
 import uuid
+from pathlib import Path
 
 import pytest
 
@@ -10,7 +13,6 @@
 
 from distributed import profile
 from distributed.compatibility import WINDOWS
-from distributed.protocol import serialize_bytelist
 from distributed.spill import SpillBuffer, has_zict_210, has_zict_220
 from distributed.utils_test import captured_logger
 
@@ -24,14 +26,16 @@
 )
 
 
-def psize(*objs: object) -> tuple[int, int]:
+def psize(tmp_path: Path, **objs: object) -> tuple[int, int]:
     return (
-        sum(sizeof(o) for o in objs),
-        sum(len(frame) for obj in objs for frame in serialize_bytelist(obj)),
+        sum(sizeof(o) for o in objs.values()),
+        sum(os.stat(tmp_path / k).st_size for k in objs),
     )
 
 
-def assert_buf(buf: SpillBuffer, expect_fast: dict, expect_slow: dict) -> None:
+def assert_buf(
+    buf: SpillBuffer, tmp_path: Path, 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()}
@@ -42,91 +46,111 @@ def assert_buf(buf: SpillBuffer, expect_fast: dict, expect_slow: dict) -> None:
     # 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.weight_by_key == {
+        k: psize(tmp_path, **{k: v}) for k, v in expect_slow.items()
+    }
+    total_weight = psize(tmp_path, **expect_slow)
     assert slow.total_weight == total_weight
     assert buf.spilled_total == total_weight
 
 
-def test_spillbuffer(tmpdir):
-    buf = SpillBuffer(str(tmpdir), target=300)
+def test_psize(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=0)
+    a = "a" * 100
+    assert 100 < sizeof(a) < 200
+    buf["a"] = a
+    memory_size, disk_size = psize(tmp_path, a=a)
+    assert memory_size != disk_size
+
+
+def test_spillbuffer(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=300)
     # Convenience aliases
     assert buf.memory is buf.fast
     assert buf.disk is buf.slow
 
-    assert_buf(buf, {}, {})
+    assert_buf(buf, tmp_path, {}, {})
 
     a, b, c, d = "a" * 100, "b" * 99, "c" * 98, "d" * 97
 
     # Test assumption made by this test, mostly for non CPython implementations
     assert 100 < sizeof(a) < 200
-    assert psize(a)[0] != psize(a)[1]
 
     buf["a"] = a
-    assert_buf(buf, {"a": a}, {})
+    assert_buf(buf, tmp_path, {"a": a}, {})
     assert buf["a"] == a
 
     buf["b"] = b
-    assert_buf(buf, {"a": a, "b": b}, {})
+    assert_buf(buf, tmp_path, {"a": a, "b": b}, {})
 
     buf["c"] = c
-    assert_buf(buf, {"b": b, "c": c}, {"a": a})
+    assert_buf(buf, tmp_path, {"b": b, "c": c}, {"a": a})
 
     assert buf["a"] == a
-    assert_buf(buf, {"a": a, "c": c}, {"b": b})
+    assert_buf(buf, tmp_path, {"a": a, "c": c}, {"b": b})
 
     buf["d"] = d
-    assert_buf(buf, {"a": a, "d": d}, {"b": b, "c": c})
+    assert_buf(buf, tmp_path, {"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_buf(buf, {"d": d}, {"b": b, "c": c})
+    assert_buf(buf, tmp_path, {"d": d}, {"b": b, "c": c})
     with pytest.raises(KeyError):
         buf["a"]
 
     # Deleting a spilled key updates the metadata
     del buf["b"]
-    assert_buf(buf, {"d": d}, {"c": c})
+    assert_buf(buf, tmp_path, {"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
     c2 = c * 2
     buf["c"] = c2
-    assert_buf(buf, {"c": c2}, {"d": d})
+    assert_buf(buf, tmp_path, {"c": c2}, {"d": d})
 
     # Single key is larger than target and goes directly into slow
     e = "e" * 500
 
     buf["e"] = e
-    assert_buf(buf, {"c": c2}, {"d": d, "e": e})
+    assert_buf(buf, tmp_path, {"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_buf(buf, {"c": c2}, {"d": d, "e": e})
+    assert_buf(buf, tmp_path, {"c": c2}, {"d": d, "e": e})
+
+
+def test_disk_size_calculation(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=0)
+    a = "a" * 100
+    b = array.array("d", (random.random() for _ in range(100)))
+    buf["a"] = a
+    buf["b"] = b
+    assert_buf(buf, tmp_path, {}, {"a": a, "b": b})
 
 
 @requires_zict_210
-def test_spillbuffer_maxlim(tmpdir):
-    buf = SpillBuffer(str(tmpdir), target=200, max_spill=600, min_log_interval=0)
+def test_spillbuffer_maxlim(tmp_path_factory):
+    buf_dir = tmp_path_factory.mktemp("buf")
+    buf = SpillBuffer(str(buf_dir), target=200, max_spill=600, min_log_interval=0)
 
     a, b, c, d, e = "a" * 200, "b" * 100, "c" * 99, "d" * 199, "e" * 98
 
     # size of a is bigger than target and is smaller than max_spill;
     # key should be in slow
     buf["a"] = a
-    assert_buf(buf, {}, {"a": a})
+    assert_buf(buf, buf_dir, {}, {"a": a})
     assert buf["a"] == a
 
     # size of b is smaller than target key should be in fast
     buf["b"] = b
-    assert_buf(buf, {"b": b}, {"a": a})
+    assert_buf(buf, buf_dir, {"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_buf(buf, {"c": c}, {"a": a, "b": b})
+    assert_buf(buf, buf_dir, {"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
@@ -135,7 +159,7 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["e"] = e
 
     assert "disk reached capacity" in logs_e.getvalue()
-    assert_buf(buf, {"c": c, "e": e}, {"a": a, "b": b})
+    assert_buf(buf, buf_dir, {"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
@@ -144,19 +168,24 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["d"] = d
 
     assert "disk reached capacity" in logs_d.getvalue()
-    assert_buf(buf, {"c": c, "d": d, "e": e}, {"a": a, "b": b})
+    assert_buf(buf, buf_dir, {"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
 
     a_large = "a" * 500
-    assert psize(a_large)[1] > 600  # size of max_spill
+
+    # Assert precondition that a_large is larger than max_spill when written to disk
+    unlimited_buf_dir = tmp_path_factory.mktemp("unlimited_buf")
+    unlimited_buf = SpillBuffer(unlimited_buf_dir, target=0)
+    unlimited_buf["a_large"] = a_large
+    assert psize(unlimited_buf_dir, a_large=a_large)[1] > 600
 
     with captured_logger(logging.getLogger("distributed.spill")) as logs_alarge:
         buf["a"] = a_large
 
     assert "disk reached capacity" in logs_alarge.getvalue()
-    assert_buf(buf, {"a": a_large, "d": d, "e": e}, {"b": b, "c": c})
+    assert_buf(buf, buf_dir, {"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
@@ -166,7 +195,7 @@ def test_spillbuffer_maxlim(tmpdir):
         buf["d"] = d_large
 
     assert "disk reached capacity" in logs_dlarge.getvalue()
-    assert_buf(buf, {"a": a_large, "d": d_large, "e": e}, {"b": b, "c": c})
+    assert_buf(buf, buf_dir, {"a": a_large, "d": d_large, "e": e}, {"b": b, "c": c})
 
 
 class MyError(Exception):
@@ -185,8 +214,8 @@ def __sizeof__(self):
 
 
 @requires_zict_210
-def test_spillbuffer_fail_to_serialize(tmpdir):
-    buf = SpillBuffer(str(tmpdir), target=200, max_spill=600, min_log_interval=0)
+def test_spillbuffer_fail_to_serialize(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=200, max_spill=600, min_log_interval=0)
 
     # bad data individually larger than spill threshold target 200
     a = Bad(size=201)
@@ -198,12 +227,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_buf(buf, {}, {})
+    assert_buf(buf, tmp_path, {}, {})
 
     b = Bad(size=100)  # this is small enough to fit in memory/fast
 
     buf["b"] = b
-    assert_buf(buf, {"b": b}, {})
+    assert_buf(buf, tmp_path, {"b": b}, {})
 
     c = "c" * 100
     with captured_logger(logging.getLogger("distributed.spill")) as logs_bad_key_mem:
@@ -215,13 +244,13 @@ 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_buf(buf, {"b": b, "c": c}, {})
+    assert_buf(buf, tmp_path, {"b": b, "c": c}, {})
 
 
 @requires_zict_210
 @pytest.mark.skipif(WINDOWS, reason="Needs chmod")
-def test_spillbuffer_oserror(tmpdir):
-    buf = SpillBuffer(str(tmpdir), target=200, max_spill=800, min_log_interval=0)
+def test_spillbuffer_oserror(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=200, max_spill=800, min_log_interval=0)
 
     a, b, c, d = (
         "a" * 200,
@@ -233,21 +262,21 @@ def test_spillbuffer_oserror(tmpdir):
     # let's have something in fast and something in slow
     buf["a"] = a
     buf["b"] = b
-    assert_buf(buf, {"b": b}, {"a": a})
+    assert_buf(buf, tmp_path, {"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.
-    os.chmod(tmpdir, 0o555)
+    os.chmod(tmp_path, 0o555)
 
     # Add key > than target
     with captured_logger(logging.getLogger("distributed.spill")) as logs_oserror_slow:
         buf["c"] = c
 
     assert "Spill to disk failed" in logs_oserror_slow.getvalue()
-    assert_buf(buf, {"b": b, "c": c}, {"a": a})
+    assert_buf(buf, tmp_path, {"b": b, "c": c}, {"a": a})
 
     del buf["c"]
-    assert_buf(buf, {"b": b}, {"a": a})
+    assert_buf(buf, tmp_path, {"b": b}, {"a": a})
 
     # add key to fast which is smaller than target but when added it triggers spill,
     # which triggers OSError
@@ -255,26 +284,26 @@ def test_spillbuffer_oserror(tmpdir):
         buf["d"] = d
 
     assert "Spill to disk failed" in logs_oserror_evict.getvalue()
-    assert_buf(buf, {"b": b, "d": d}, {"a": a})
+    assert_buf(buf, tmp_path, {"b": b, "d": d}, {"a": a})
 
 
 @requires_zict_210
-def test_spillbuffer_evict(tmpdir):
-    buf = SpillBuffer(str(tmpdir), target=300, min_log_interval=0)
+def test_spillbuffer_evict(tmp_path):
+    buf = SpillBuffer(str(tmp_path), target=300, min_log_interval=0)
 
     bad = Bad(size=100)
     a = "a" * 100
 
     buf["a"] = a
-    assert_buf(buf, {"a": a}, {})
+    assert_buf(buf, tmp_path, {"a": a}, {})
 
     # successful eviction
     weight = buf.evict()
     assert weight == sizeof(a)
-    assert_buf(buf, {}, {"a": a})
+    assert_buf(buf, tmp_path, {}, {"a": a})
 
     buf["bad"] = bad
-    assert_buf(buf, {"bad": bad}, {"a": a})
+    assert_buf(buf, tmp_path, {"bad": bad}, {"a": a})
 
     # unsuccessful eviction
     with captured_logger(logging.getLogger("distributed.spill")) as logs_evict_key:
@@ -283,7 +312,7 @@ def test_spillbuffer_evict(tmpdir):
 
     assert "Failed to pickle" in logs_evict_key.getvalue()
     # bad keys stays in fast
-    assert_buf(buf, {"bad": bad}, {"a": a})
+    assert_buf(buf, tmp_path, {"bad": bad}, {"a": a})
 
 
 class NoWeakRef:
@@ -318,8 +347,8 @@ class SupportsWeakRef(NoWeakRef):
     ],
 )
 @pytest.mark.parametrize("size", [60, 110])
-def test_weakref_cache(tmpdir, cls, expect_cached, size):
-    buf = SpillBuffer(str(tmpdir), target=100)
+def test_weakref_cache(tmp_path, cls, expect_cached, size):
+    buf = SpillBuffer(str(tmp_path), target=100)
 
     # Run this test twice:
     # - x is smaller than target and is evicted by y;