mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-04 05:31:03 +08:00
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
"""Positioned, off-loop file writes.
|
|
|
|
Network I/O stays on the event loop; every blocking disk op (preallocate,
|
|
positioned write, fsync) is run in a bounded thread pool via
|
|
``run_in_executor`` so downloads never stall inference or the web server.
|
|
|
|
A single file descriptor is opened for the whole download. Segments write to
|
|
their own offsets with ``os.pwrite`` — which is offset-addressed and atomic
|
|
per call, so concurrent segment writers need no extra locking. Per-chunk
|
|
fsync is avoided; we fsync once at completion.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from typing import Optional
|
|
|
|
# One shared, bounded pool for all download disk I/O.
|
|
_EXECUTOR = ThreadPoolExecutor(max_workers=8, thread_name_prefix="dl-writer")
|
|
|
|
|
|
class FileWriter:
|
|
"""Owns the ``.part`` file descriptor for one download."""
|
|
|
|
def __init__(self, path: str) -> None:
|
|
self.path = path
|
|
self._fd: Optional[int] = None
|
|
|
|
def _open(self) -> None:
|
|
os.makedirs(os.path.dirname(self.path), exist_ok=True)
|
|
self._fd = os.open(self.path, os.O_RDWR | os.O_CREAT, 0o644)
|
|
|
|
async def open(self) -> None:
|
|
await asyncio.get_running_loop().run_in_executor(_EXECUTOR, self._open)
|
|
|
|
async def preallocate(self, size: int) -> None:
|
|
"""Grow the file to ``size`` so segments write to their offsets."""
|
|
if self._fd is None or size <= 0:
|
|
return
|
|
await asyncio.get_running_loop().run_in_executor(
|
|
_EXECUTOR, os.ftruncate, self._fd, size
|
|
)
|
|
|
|
async def truncate(self, size: int = 0) -> None:
|
|
"""Truncate the file to ``size`` bytes (default: empty it)."""
|
|
if self._fd is None:
|
|
return
|
|
await asyncio.get_running_loop().run_in_executor(
|
|
_EXECUTOR, os.ftruncate, self._fd, size
|
|
)
|
|
|
|
def _pwrite_all(self, data: bytes, offset: int) -> None:
|
|
"""``os.pwrite`` may write fewer bytes than requested (signal
|
|
interruption, near-ENOSPC); loop until every byte lands so we never
|
|
leave a gap while the caller advances by the full chunk length."""
|
|
assert self._fd is not None, "writer not opened"
|
|
view = memoryview(data)
|
|
written = 0
|
|
total = len(view)
|
|
while written < total:
|
|
n = os.pwrite(self._fd, view[written:], offset + written)
|
|
if n == 0:
|
|
raise OSError(
|
|
f"os.pwrite wrote 0 bytes at offset {offset + written} "
|
|
f"({written}/{total} bytes written)"
|
|
)
|
|
written += n
|
|
|
|
async def write_at(self, offset: int, data: bytes) -> None:
|
|
assert self._fd is not None, "writer not opened"
|
|
await asyncio.get_running_loop().run_in_executor(
|
|
_EXECUTOR, self._pwrite_all, data, offset
|
|
)
|
|
|
|
async def flush(self) -> None:
|
|
if self._fd is None:
|
|
return
|
|
await asyncio.get_running_loop().run_in_executor(_EXECUTOR, os.fsync, self._fd)
|
|
|
|
async def close(self) -> None:
|
|
if self._fd is None:
|
|
return
|
|
fd, self._fd = self._fd, None
|
|
await asyncio.get_running_loop().run_in_executor(_EXECUTOR, os.close, fd)
|