diff --git a/daser/transfer/iouring/layer.py b/daser/transfer/iouring/layer.py index 7a31ed5..a23ba44 100644 --- a/daser/transfer/iouring/layer.py +++ b/daser/transfer/iouring/layer.py @@ -19,6 +19,7 @@ logger = init_logger(__name__) _DIRECT_IO_ALIGNMENT = 4096 +_IO_URING_ENTRIES = 64 class TieredIOUringTransferLayer(TransferLayer): @@ -65,7 +66,10 @@ def __init__( self._path = path self._fd = os.open(path, os.O_RDWR | os.O_DIRECT) - self._urings = [NativeIOUring(entries=64) for _ in range(io_workers)] + self._urings = [ + NativeIOUring(entries=_IO_URING_ENTRIES) for _ in range(io_workers) + ] + self._l2_read_batch_spans = io_workers * _IO_URING_ENTRIES self._uring_lock = threading.Lock() self._next_uring_index = 0 self._io_executor = concurrent.futures.ThreadPoolExecutor( @@ -405,6 +409,17 @@ def _read_l2_into( """Blocking io_uring L2 read into pinned memory.""" return uring.read_into(self._fd, file_offset, dst.view()) + def _readv_l2_into( + self, + reads: list[tuple[int, PinnedMemorySlice]], + uring: NativeIOUring, + ) -> int: + """Blocking batched io_uring L2 reads into pinned memory.""" + return uring.readv_into( + self._fd, + [(file_offset, pinned.view()) for file_offset, pinned in reads], + ) + def _write_l2( self, file_offset: int, @@ -537,16 +552,20 @@ async def _load_l2_miss_batch( pinned = await self._reserve_l1_buffer(key, nbytes) reads.append((span, pinned)) + read_pairs = [(int(span["file_offset"]), pinned) for span, pinned in reads] await asyncio.gather( *( loop.run_in_executor( self._io_executor, - self._read_l2_into, - int(span["file_offset"]), - pinned, + self._readv_l2_into, + read_pairs[start : start + _IO_URING_ENTRIES], self._next_uring(), ) - for span, pinned in reads + for start in range( + 0, + len(read_pairs), + _IO_URING_ENTRIES, + ) ) ) @@ -589,7 +608,8 @@ def _next_l2_miss_batch( for span in misses[start:]: nbytes = int(span["nbytes"]) if batch and ( - batch_bytes + nbytes > self._l1_bytes or len(batch) >= len(self._urings) + batch_bytes + nbytes > self._l1_bytes + or len(batch) >= self._l2_read_batch_spans ): break batch.append(span) diff --git a/daser/transfer/iouring/native.py b/daser/transfer/iouring/native.py index bc02ea9..1cda8a5 100644 --- a/daser/transfer/iouring/native.py +++ b/daser/transfer/iouring/native.py @@ -239,6 +239,74 @@ def read_into(self, fd: int, file_offset: int, dst: memoryview) -> int: cursor += chunk return total + def readv_into( + self, + fd: int, + reads: list[tuple[int, memoryview]], + ) -> int: + """Read multiple positioned ranges into writable buffers. + + Args: + fd: Open file descriptor. + reads: ``(file_offset, dst)`` pairs. Each destination is read fully. + + Returns: + Total number of bytes read. + + Thread-safety: + Serialized by the wrapper lock. Destination buffers must not be + mutated by another thread until this call returns. + """ + total = 0 + pending: list[tuple[int, memoryview]] = [] + for file_offset, dst in reads: + view = memoryview(dst).cast("B") + if view.readonly: + raise ValueError("dst must be writable") + cursor = 0 + while cursor < len(view): + chunk = min(_MAX_RW_COUNT, len(view) - cursor) + pending.append((file_offset + cursor, view[cursor : cursor + chunk])) + total += chunk + cursor += chunk + if not pending: + return 0 + + results: list[int] = [] + batch_limit = int(self._params.sq_entries) + with self._lock: + for batch_start in range(0, len(pending), batch_limit): + batch = pending[batch_start : batch_start + batch_limit] + start_user_data = self._next_user_data + self._next_user_data += len(batch) + for idx, (file_offset, view) in enumerate(batch): + self._submit_locked( + _IORING_OP_READ, + fd, + file_offset, + view, + len(view), + start_user_data + idx, + ) + self._enter( + len(batch), + len(batch), + _IORING_ENTER_GETEVENTS, + ) + results.extend( + self._wait_completions_locked( + start_user_data, + [len(view) for _offset, view in batch], + ) + ) + for idx, res in enumerate(results): + if res < 0: + raise OSError(-res, os.strerror(-res)) + expected = len(pending[idx][1]) + if res != expected: + raise IOError(f"short io_uring result: {res} != {expected}") + return total + def write(self, fd: int, file_offset: int, data: bytes | memoryview) -> int: """Write bytes at a file offset through io_uring. @@ -305,15 +373,16 @@ def _submit_and_wait( with self._lock: user_data = self._next_user_data self._next_user_data += 1 - self._submit(opcode, fd, file_offset, buf, nbytes, user_data) - res = self._wait_completion(user_data) + self._submit_locked(opcode, fd, file_offset, buf, nbytes, user_data) + self._enter(1, 1, _IORING_ENTER_GETEVENTS) + res = self._wait_completion_locked(user_data) if res < 0: raise OSError(-res, os.strerror(-res)) if res != nbytes: raise IOError(f"short io_uring result: {res} != {nbytes}") return res - def _submit( + def _submit_locked( self, opcode: int, fd: int, @@ -350,12 +419,14 @@ def _submit( ) self._write_u32(self._sq_ring, self._params.sq_off.tail, sq_tail + 1) + def _enter(self, to_submit: int, min_complete: int, flags: int) -> None: + """Enter the kernel to submit SQEs and optionally wait for CQEs.""" ret = self._libc.syscall( _SYS_IO_URING_ENTER, ctypes.c_int(self._ring_fd), - ctypes.c_uint32(1), - ctypes.c_uint32(0), - ctypes.c_uint32(0), + ctypes.c_uint32(to_submit), + ctypes.c_uint32(min_complete), + ctypes.c_uint32(flags), ctypes.c_void_p(0), ctypes.c_size_t(0), ) @@ -363,8 +434,18 @@ def _submit( errno = ctypes.get_errno() raise OSError(errno, os.strerror(errno)) - def _wait_completion(self, user_data: int) -> int: + def _wait_completion_locked(self, user_data: int) -> int: """Wait for a CQE matching ``user_data`` and return its result.""" + return self._wait_completions_locked(user_data, [0])[0] + + def _wait_completions_locked( + self, + start_user_data: int, + expected_sizes: list[int], + ) -> list[int]: + """Wait for a contiguous user-data range and return ordered results.""" + remaining = len(expected_sizes) + results: list[int | None] = [None] * remaining while True: cq_head = self._read_u32(self._cq_ring, self._params.cq_off.head) cq_tail = self._read_u32(self._cq_ring, self._params.cq_off.tail) @@ -379,25 +460,20 @@ def _wait_completion(self, user_data: int) -> int: res = int(cqe.res) seen_user_data = int(cqe.user_data) self._write_u32(self._cq_ring, self._params.cq_off.head, cq_head + 1) - if seen_user_data != user_data: + index = seen_user_data - start_user_data + if index < 0 or index >= len(expected_sizes): raise RuntimeError( "unexpected io_uring completion " - f"{seen_user_data}, expected {user_data}" + f"{seen_user_data}, expected range " + f"[{start_user_data}, {start_user_data + len(expected_sizes)})" ) - return res - - ret = self._libc.syscall( - _SYS_IO_URING_ENTER, - ctypes.c_int(self._ring_fd), - ctypes.c_uint32(0), - ctypes.c_uint32(1), - ctypes.c_uint32(_IORING_ENTER_GETEVENTS), - ctypes.c_void_p(0), - ctypes.c_size_t(0), - ) - if ret < 0: - errno = ctypes.get_errno() - raise OSError(errno, os.strerror(errno)) + results[index] = res + remaining -= 1 + if remaining == 0: + return [int(result) for result in results] + continue + + self._enter(0, 1, _IORING_ENTER_GETEVENTS) def _read_u32(self, buf: mmap.mmap, offset: int) -> int: """Read a uint32 from a ring mapping without exporting the buffer.""" diff --git a/tests/transfer/test_tiered_iouring_transfer.py b/tests/transfer/test_tiered_iouring_transfer.py index 9e3f162..c8b4326 100644 --- a/tests/transfer/test_tiered_iouring_transfer.py +++ b/tests/transfer/test_tiered_iouring_transfer.py @@ -72,6 +72,27 @@ def _copy_grouped_to_dst( super()._copy_grouped_to_dst(dst, chunks) +class MissBatchProbe(TieredIOUringTransferLayer): + """Test transfer layer that records grouped L2 miss batch sizes.""" + + def __init__(self, path: str, l1_bytes: int, l2_bytes: int) -> None: + super().__init__( + path=path, + l1_bytes=l1_bytes, + l2_bytes=l2_bytes, + ) + self.miss_batches: list[int] = [] + + async def _load_l2_miss_batch( + self, + dst: object, + misses: list[dict[str, int]], + ) -> None: + """Record miss batch size before delegating to production logic.""" + self.miss_batches.append(len(misses)) + await super()._load_l2_miss_batch(dst, misses) + + def _run(coro: object) -> object: """Run a coroutine on the current test event loop.""" return asyncio.get_event_loop().run_until_complete(coro) @@ -192,6 +213,57 @@ def test_native_iouring_read_into(tmp_path) -> None: os.close(fd) +def test_native_iouring_readv_into_batches_submit(tmp_path, monkeypatch) -> None: + """Native io_uring submits multiple positioned reads in one enter call.""" + path = tmp_path / "readv.store" + path.write_bytes(b"abcdefghijklmnop") + fd = os.open(path, os.O_RDWR) + calls: list[tuple[int, int, int]] = [] + original_cdll = native_iouring.ctypes.CDLL + + def as_int(value: object) -> int: + """Return an int from a ctypes scalar or plain Python value.""" + return int(getattr(value, "value", value)) + + class TrackingCDLL: + """Proxy libc calls while recording io_uring_enter arguments.""" + + def __init__(self, *args: object, **kwargs: object) -> None: + self._real = original_cdll(*args, **kwargs) + + def syscall(self, number: int, *args: object) -> int: + """Record io_uring_enter calls and delegate to libc.""" + if number == 426: + calls.append((as_int(args[1]), as_int(args[2]), as_int(args[3]))) + return self._real.syscall(number, *args) + + monkeypatch.setattr(native_iouring.ctypes, "CDLL", TrackingCDLL) + + uring = NativeIOUring(entries=8) + one = bytearray(4) + two = bytearray(4) + three = bytearray(4) + try: + assert ( + uring.readv_into( + fd, + [ + (0, memoryview(one)), + (4, memoryview(two)), + (8, memoryview(three)), + ], + ) + == 12 + ) + assert bytes(one) == b"abcd" + assert bytes(two) == b"efgh" + assert bytes(three) == b"ijkl" + assert calls == [(3, 3, 1)] + finally: + uring.close() + os.close(fd) + + def test_iouring_direct_io_aligned_roundtrip(tmp_path) -> None: """Production io_uring mode uses O_DIRECT for aligned L2 ranges.""" path = str(tmp_path / "direct.store") @@ -424,6 +496,55 @@ async def scenario() -> None: _run(scenario()) +def test_iouring_grouped_l2_misses_batch_beyond_ring_count(tmp_path) -> None: + """Grouped L2 misses are batched by queue capacity, not just ring count.""" + + async def scenario() -> None: + path = str(tmp_path / "daser.store") + block_count = 16 + layer = TieredIOUringTransferLayer( + path=path, + l1_bytes=ALIGNMENT * block_count, + l2_bytes=ALIGNMENT * (block_count + 1), + ) + try: + for idx in range(block_count): + await layer.store_bytes( + _block(bytes([idx])), + file_offset=idx * ALIGNMENT, + nbytes=ALIGNMENT, + ) + await layer.drain() + finally: + layer.close() + + layer = MissBatchProbe( + path=path, + l1_bytes=ALIGNMENT * block_count, + l2_bytes=ALIGNMENT * (block_count + 1), + ) + try: + dst = bytearray(ALIGNMENT * block_count) + loaded = await layer.load_bytes_grouped( + dst, + [ + { + "target_offset": idx * ALIGNMENT, + "file_offset": idx * ALIGNMENT, + "nbytes": ALIGNMENT, + } + for idx in range(block_count) + ], + ) + + assert loaded == ALIGNMENT * block_count + assert layer.miss_batches == [block_count] + finally: + layer.close() + + _run(scenario()) + + def test_iouring_parallel_l2_loads_use_independent_offsets(tmp_path) -> None: """Concurrent L2 loads read their requested byte ranges exactly."""