from __future__ import annotations import contextlib import errno import pathlib import shutil import threading from collections.abc import Callable, Generator, Iterable from contextlib import contextmanager from typing import Any from toolz import concat from distributed.metrics import context_meter, thread_time from distributed.shuffle._buffer import ShardsBuffer from distributed.shuffle._exceptions import DataUnavailable, P2POutOfDiskError from distributed.shuffle._limiter import ResourceLimiter from distributed.shuffle._pickle import pickle_bytelist from distributed.utils import Deadline, empty_context, log_errors, nbytes class ReadWriteLock: _condition: threading.Condition _n_reads: int _write_pending: bool def __init__(self) -> None: self._condition = threading.Condition(threading.Lock()) self._n_reads = 0 self._write_pending = False self._write_active = False def acquire_write(self, timeout: float = -1) -> bool: deadline = Deadline.after(timeout if timeout >= 0 else None) with self._condition: result = self._condition.wait_for( lambda: not self._write_pending, timeout=deadline.remaining ) if result is False: return False self._write_pending = True result = self._condition.wait_for( lambda: self._n_reads == 0, timeout=deadline.remaining ) if result is False: self._write_pending = False self._condition.notify_all() return False self._write_active = True return True def release_write(self) -> None: with self._condition: if self._write_active is False: raise RuntimeError("Tried releasing unlocked write lock") self._write_pending = False self._write_active = False self._condition.notify_all() def acquire_read(self, timeout: float = -1) -> bool: deadline = Deadline.after(timeout if timeout >= 0 else None) with self._condition: result = self._condition.wait_for( lambda: not self._write_pending, timeout=deadline.remaining ) if result is False: return False self._n_reads += 1 return True def release_read(self) -> None: with self._condition: if self._n_reads == 0: raise RuntimeError("Tired releasing unlocked read lock") self._n_reads -= 1 if self._n_reads == 0: self._condition.notify_all() @contextmanager def write(self) -> Generator[None]: self.acquire_write() try: yield finally: self.release_write() @contextmanager def read(self) -> Generator[None]: self.acquire_read() try: yield finally: self.release_read() class DiskShardsBuffer(ShardsBuffer): """Accept, buffer, and write many small objects to many files This takes in lots of small objects, writes them to a local directory, and then reads them back when all writes are complete. It buffers these objects in memory so that it can optimize disk access for larger writes. **State** - shards: dict[str, list[bytes]] This is our in-memory buffer of data waiting to be written to files. - sizes: dict[str, int] The size of each list of shards. We find the largest and write data from that buffer Parameters ---------- directory : str or pathlib.Path Where to write and read data. Ideally points to fast disk. memory_limiter : ResourceLimiter Limiter for in-memory buffering (at most this much data) before writes to disk occur. If the incoming data that has yet to be processed exceeds this limit, then the buffer will block until below the threshold. See :meth:`.write` for the implementation of this scheme. """ def __init__( self, directory: str | pathlib.Path, read: Callable[[pathlib.Path], tuple[Any, int]], memory_limiter: ResourceLimiter, ): super().__init__( memory_limiter=memory_limiter, # Disk is not able to run concurrently atm concurrency_limit=1, ) self.directory = pathlib.Path(directory) self.directory.mkdir(exist_ok=True) self._closed = False self._read = read self._directory_lock = ReadWriteLock() @log_errors async def _process(self, id: str, shards: list[Any]) -> None: """Write one buffer to file This function was built to offload the disk IO, but since then we've decided to keep this within the event loop (disk bandwidth should be prioritized, and writes are typically small enough to not be a big deal). Most of the logic here is about possibly going back to a separate thread, or about diagnostics. If things don't change much in the future then we should consider simplifying this considerably and dropping the write into communicate above. """ frames: Iterable[bytes | bytearray | memoryview] if isinstance(shards[0], bytes): # Manually serialized dataframes frames = shards serialize_meter_ctx: Any = empty_context else: # Unserialized numpy arrays # Note: no calls to pickle_bytelist will happen until we actually start # writing to disk below. frames = concat(pickle_bytelist(shard) for shard in shards) serialize_meter_ctx = context_meter.meter("serialize", func=thread_time) with ( self._directory_lock.read(), context_meter.meter("disk-write"), serialize_meter_ctx, ): # Consider boosting total_size a bit here to account for duplication # We only need shared (i.e., read) access to the directory to write # to a file inside of it. if self._closed: raise RuntimeError("Already closed") try: self._write_frames(frames, id) except OSError as e: if e.errno == errno.ENOSPC: raise P2POutOfDiskError from e raise context_meter.digest_metric("disk-write", 1, "count") context_meter.digest_metric("disk-write", sum(map(nbytes, frames)), "bytes") def _write_frames( self, frames: Iterable[bytes | bytearray | memoryview], id: str ) -> None: with open(self.directory / str(id), mode="ab") as f: f.writelines(frames) def read(self, id: str) -> Any: """Read a complete file back into memory""" self.raise_on_exception() if not self._inputs_done: raise RuntimeError("Tried to read from file before done.") try: with self._directory_lock.read(): if self._closed: raise RuntimeError("Already closed") fname = (self.directory / str(id)).resolve() # Note: don't add `with context_meter.meter("p2p-disk-read"):` to # measure seconds here, as it would shadow "p2p-get-output-cpu" and # "p2p-get-output-noncpu". Also, for rechunk it would not measure # the whole disk access, as _read returns memory-mapped buffers. data, size = self._read(fname) context_meter.digest_metric("p2p-disk-read", 1, "count") context_meter.digest_metric("p2p-disk-read", size, "bytes") except FileNotFoundError: raise DataUnavailable(id) if data: self.bytes_read += size return data else: raise DataUnavailable(id) async def close(self) -> None: await super().close() with self._directory_lock.write(): self._closed = True with contextlib.suppress(FileNotFoundError): shutil.rmtree(self.directory)