class DiskBackend:
"""Async disk offload backend with pipelined GPU DMA and interleaved IO.
Architecture:
- Separate coordinator threads for store and load (never block each other)
- Interleaved pipeline: DMA slot N while preadv/pwritev slot N-1
- O_DIRECT by default; page cache is opt-in via use_page_cache
Same launch_copy interface as DmaCopyBackend so the worker can swap
backends without changing calling code.
"""
def __init__(self) -> None:
self._store_params: BatchMemcpyParams | None = None
self._load_params: BatchMemcpyParams | None = None
self._load_stream: torch.cuda.Stream | None = None
self._store_stream: torch.cuda.Stream | None = None
self._store_queue: queue.SimpleQueue = queue.SimpleQueue()
self._load_queue: queue.SimpleQueue = queue.SimpleQueue()
self._store_thread: threading.Thread | None = None
self._load_thread: threading.Thread | None = None
self._shutdown: bool = False
self._fd: int = -1
self._disk_path: str = ""
self._total_block_bytes: int = 0
self._store_buffer_caches: dict[str, torch.Tensor] = {}
self._load_buffer_caches: dict[str, torch.Tensor] = {}
self._store_slot_views: list[list[memoryview]] = []
self._load_slot_views: list[list[memoryview]] = []
self._per_tensor_bpb: list[int] = []
self._tensor_names: list[str] = []
def init(
self,
gpu_caches: dict[str, torch.Tensor],
device: torch.device,
load_stream: torch.cuda.Stream,
store_stream: torch.cuda.Stream,
disk_path: str,
num_disk_slots: int,
total_block_bytes: int,
num_buffer_slots: int = 2,
use_page_cache: bool = False,
) -> None:
self._load_stream = load_stream
self._store_stream = store_stream
self._total_block_bytes = total_block_bytes
self._num_buffer_slots = num_buffer_slots
self._tensor_names = list(gpu_caches.keys())
self._per_tensor_bpb = [
t.stride(0) * t.element_size() for t in gpu_caches.values()
]
assert total_block_bytes % _ALIGNMENT == 0, (
f"total_block_bytes={total_block_bytes} not aligned to {_ALIGNMENT}"
)
# Separate buffer pools for store and load threads
self._store_buffer_caches = {}
self._load_buffer_caches = {}
for name, gpu_t in gpu_caches.items():
bpb = gpu_t.stride(0) * gpu_t.element_size()
store_buf = _alloc_aligned(num_buffer_slots, bpb)
pin_tensor(store_buf)
self._store_buffer_caches[name] = store_buf
load_buf = _alloc_aligned(num_buffer_slots, bpb)
pin_tensor(load_buf)
self._load_buffer_caches[name] = load_buf
# Pre-built iovec views per slot (avoid per-transfer .numpy() calls)
self._store_slot_views = [
[
memoryview(self._store_buffer_caches[name][slot].numpy())
for name in self._tensor_names
]
for slot in range(num_buffer_slots)
]
self._load_slot_views = [
[
memoryview(self._load_buffer_caches[name][slot].numpy())
for name in self._tensor_names
]
for slot in range(num_buffer_slots)
]
self._store_params = build_params(
gpu_caches,
self._store_buffer_caches,
store_stream,
src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_STREAM,
)
self._load_params = build_params(
self._load_buffer_caches,
gpu_caches,
load_stream,
src_access_order=CU_MEMCPY_SRC_ACCESS_ORDER_ANY,
)
os.makedirs(os.path.dirname(disk_path) or ".", exist_ok=True)
# Slot contents never outlive the process, so unlink then O_EXCL rather
# than reopening: a pre-existing file would otherwise keep its own
# (possibly world-readable) mode, and blocks may encode user prompts.
with contextlib.suppress(FileNotFoundError):
os.unlink(disk_path)
# O_DIRECT by default: page cache would consume the very host DRAM this
# backend exists to conserve, and doubles the copy on the store path.
flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW
if not use_page_cache:
flags |= O_DIRECT
self._fd = os.open(disk_path, flags, 0o600)
self._disk_path = disk_path
os.ftruncate(self._fd, num_disk_slots * total_block_bytes)
logger.info(
"DiskBackend: path=%s, slots=%d, total=%.2f GB, buf=%dx%d bytes"
" (page_cache=%s)",
disk_path,
num_disk_slots,
(num_disk_slots * total_block_bytes) / (1024**3),
num_buffer_slots,
total_block_bytes,
use_page_cache,
)
self._store_thread = threading.Thread(
target=self._store_loop,
args=(device, store_stream),
daemon=True,
)
self._load_thread = threading.Thread(
target=self._load_loop,
args=(device, load_stream),
daemon=True,
)
self._store_thread.start()
self._load_thread.start()
def launch_copy(
self,
src_blocks: list[int],
dst_blocks: list[int],
is_store: bool,
event_idx: int,
events_list: list[tuple[int, torch.Event]],
wait_event: torch.Event | None = None,
) -> None:
q = self._store_queue if is_store else self._load_queue
q.put((src_blocks, dst_blocks, event_idx, events_list, wait_event))
def shutdown(self) -> None:
if self._shutdown:
return
self._shutdown = True
self._store_queue.put(None)
self._load_queue.put(None)
if self._store_thread is not None:
self._store_thread.join(timeout=10.0)
if self._load_thread is not None:
self._load_thread.join(timeout=10.0)
if self._fd < 0:
return
# Slot contents can encode user prompts, so drop the name now rather
# than leaving them readable until the next run overwrites the file.
# Unlinking only removes the directory entry: any thread still holding
# the fd keeps writing to the (now anonymous) inode, which the kernel
# frees once the last fd goes away.
with contextlib.suppress(OSError):
os.unlink(self._disk_path)
# Closing under a still-running IO thread would let the fd number be
# reused by an unrelated open(), turning its next pwritev into a write
# into that file. Leaking one fd for the remaining process lifetime is
# the cheaper failure.
if any(
t is not None and t.is_alive()
for t in (self._store_thread, self._load_thread)
):
logger.warning(
"IO thread still running after shutdown timeout; leaking fd %d",
self._fd,
)
return
os.close(self._fd)
self._fd = -1
def _store_loop(
self,
device: torch.device,
stream: torch.cuda.Stream,
) -> None:
current_platform.set_device(device)
while True:
item = self._store_queue.get()
if item is None:
return
(src_blocks, dst_blocks, event_idx, events_list, wait_event) = item
if wait_event is not None:
stream.wait_event(wait_event)
self._do_store(src_blocks, dst_blocks, stream)
event = torch.Event()
event.record(stream)
events_list.append((event_idx, event))
def _writev_slot(self, buf_slot: int, file_offset: int) -> None:
written = os.pwritev(self._fd, self._store_slot_views[buf_slot], file_offset)
if written < self._total_block_bytes:
raise OSError(
f"Short write: expected {self._total_block_bytes} bytes, "
f"wrote {written}"
)
def _readv_slot(self, buf_slot: int, file_offset: int) -> None:
bytes_read = os.preadv(self._fd, self._load_slot_views[buf_slot], file_offset)
if bytes_read < self._total_block_bytes:
raise OSError(
f"Short read: expected {self._total_block_bytes} bytes, "
f"read {bytes_read}"
)
def _load_loop(
self,
device: torch.device,
stream: torch.cuda.Stream,
) -> None:
current_platform.set_device(device)
while True:
item = self._load_queue.get()
if item is None:
return
(src_blocks, dst_blocks, event_idx, events_list, wait_event) = item
if wait_event is not None:
stream.wait_event(wait_event)
self._do_load(src_blocks, dst_blocks, stream)
event = torch.Event()
event.record(stream)
events_list.append((event_idx, event))
def _do_store(
self,
gpu_blocks: list[int],
disk_slots: list[int],
stream: torch.cuda.Stream,
) -> None:
"""GPU -> buffer (DMA) -> disk (pwritev), interleaved double-buffer."""
assert self._store_params is not None
n = self._num_buffer_slots
# (DMA event, file offset) of the block already staged in each slot.
pending: list[tuple[torch.Event, int] | None] = [None] * n
for i, (gpu_blk, disk_slot) in enumerate(zip(gpu_blocks, disk_slots)):
buf_slot = i % n
prev = pending[buf_slot]
if prev is not None:
prev[0].synchronize()
self._writev_slot(buf_slot, prev[1])
copy_blocks([gpu_blk], [buf_slot], self._store_params)
ev = torch.Event()
ev.record(stream)
pending[buf_slot] = (ev, disk_slot * self._total_block_bytes)
for slot, last in enumerate(pending):
if last is not None:
last[0].synchronize()
self._writev_slot(slot, last[1])
def _do_load(
self,
disk_slots: list[int],
gpu_blocks: list[int],
stream: torch.cuda.Stream,
) -> None:
"""Disk (preadv) -> buffer -> GPU (DMA), interleaved double-buffer."""
assert self._load_params is not None
n = self._num_buffer_slots
prev_dma_events: list[torch.Event | None] = [None] * n
for i, (disk_slot, gpu_blk) in enumerate(zip(disk_slots, gpu_blocks)):
buf_slot = i % n
prev = prev_dma_events[buf_slot]
if prev is not None:
prev.synchronize()
self._readv_slot(buf_slot, disk_slot * self._total_block_bytes)
copy_blocks([buf_slot], [gpu_blk], self._load_params)
ev = torch.Event()
ev.record(stream)
prev_dma_events[buf_slot] = ev