[KV Offload][CI] Fall back to buffered I/O without O_DIRECT; fix flaky api-server test (#49734)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Jul 24, 2026 at 21:56 UTC
89f6aa3a9e56a9228f55e8d722f23594223187cf
5 files changed
+109
-5
tests/entrypoints/unit_tests/_api_server_spawn_workers.py
new
+14
@@ -0,0 +1,14 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Spawn worker targets kept free of heavy imports.
4
+
5
+``multiprocessing`` with the ``spawn`` start method re-imports the module that
6
+defines a process target in the child. Housing these stubs in a stdlib-only
7
+module keeps child startup fast and deterministic, instead of paying a multi-
8
+second ``import vllm`` before the child can run.
9
+"""
10
+
11
+
12
+def exit_before_report_worker(listen_address, sock, args, client_config=None):
13
+ """Exit immediately without touching ``actual_address_pipe``."""
14
+ return
tests/entrypoints/unit_tests/test_api_server_process_manager.py
+7
-2
@@ -10,6 +10,9 @@ from unittest.mock import patch
10
import pytest
11
import zmq
12
13
+from tests.entrypoints.unit_tests._api_server_spawn_workers import (
14
+ exit_before_report_worker,
15
+)
16
from vllm.utils.network_utils import make_zmq_socket, split_zmq_path
17
from vllm.v1.utils import (
18
APIServerProcessManager,
@@ -228,6 +231,7 @@ def test_normal_completion(api_server_args):
231
def test_external_process_monitoring(api_server_args):
232
"""Test that wait_for_completion_or_failure handles additional processes."""
233
global WORKER_RUNTIME_SECONDS
234
+ prev_worker_runtime = WORKER_RUNTIME_SECONDS
235
WORKER_RUNTIME_SECONDS = 100
236
237
# Create and start the external process
@@ -307,6 +311,7 @@ def test_external_process_monitoring(api_server_args):
311
manager.shutdown()
312
mock_coordinator.shutdown()
313
time.sleep(0.2)
314
+ WORKER_RUNTIME_SECONDS = prev_worker_runtime
315
316
317
@pytest.mark.timeout(60)
@@ -383,10 +388,10 @@ def test_gather_actual_addresses_child_crash_before_report():
388
num_servers=num_servers,
389
input_addresses=placeholder_inputs,
390
output_addresses=placeholder_outputs,
386
- # mock_run_api_server_worker exits without touching
391
+ # exit_before_report_worker exits without touching
392
# ``actual_address_pipe`` — simulates a child that dies before
393
# reporting its bound addresses.
389
- target_server_fn=mock_run_api_server_worker,
394
+ target_server_fn=exit_before_report_worker,
395
)
396
try:
397
# Sentinel-first vs pipe-EOF-first both produce "reporting".
tests/v1/kv_offload/tiering/test_fs_tier.py
+37
@@ -379,6 +379,43 @@ def test_store_load_data_integrity(fs_tier):
379
)
380
381
382
+def test_store_load_roundtrip_without_o_direct(tmp_path, monkeypatch):
383
+ """Buffered fallback must round-trip data when O_DIRECT is unsupported.
384
+
385
+ Simulates filesystems (e.g. overlayfs, some NFS) that reject O_DIRECT by
386
+ forcing the capability probe to report it unavailable.
387
+ """
388
+ monkeypatch.setattr(
389
+ "vllm.v1.kv_offload.tiering.fs.manager.probe_o_direct",
390
+ lambda _dir: False,
391
+ )
392
+ tensor = _page_aligned_rand_tensor(4, _BLOCK_ELEMENTS)
393
+ tier = FileSystemTierManager(
394
+ offloading_spec=_MOCK_OFFLOADING_SPEC,
395
+ primary_kv_view=memoryview(tensor.numpy()),
396
+ tier_type="fs",
397
+ root_dir=str(tmp_path),
398
+ n_read_threads=4,
399
+ n_write_threads=4,
400
+ )
401
+ try:
402
+ assert tier._use_o_direct is False
403
+
404
+ keys = [key(0), key(1)]
405
+ expected = tensor[:2].clone()
406
+ tier.submit_store(make_job(1, keys, [0, 1]))
407
+ assert all(r.success for r in drain(tier))
408
+
409
+ tensor[:2] = 0.0
410
+ tier.submit_load(make_job(2, keys, [2, 3], is_promotion=True))
411
+ assert all(r.success for r in drain(tier))
412
+
413
+ for i, bid in enumerate([2, 3]):
414
+ assert torch.allclose(tensor[bid], expected[i])
415
+ finally:
416
+ tier.shutdown()
417
+
418
+
419
def test_wait_idle_blocks_until_tasks_complete():
420
"""wait_idle must not return while a task is still in flight."""
421
pool = DualQueueThreadPool(n_read_threads=1, n_write_threads=1)
vllm/v1/kv_offload/tiering/fs/io.py
+36
-2
@@ -1,7 +1,9 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import contextlib
5
import logging
6
+import mmap
7
import os
8
import random
9
import threading
@@ -24,6 +26,34 @@ def _get_tmp_suffix() -> str:
26
return _thread_local.tmp_suffix
27
28
29
+def probe_o_direct(directory: str) -> bool:
30
+ """Return whether ``O_DIRECT`` I/O works in *directory*.
31
+
32
+ ``O_DIRECT`` is unsupported on some filesystems (e.g. the overlayfs backing
33
+ a container ``/tmp``, older tmpfs, or some NFS mounts), where opening or
34
+ writing a file with it fails with ``EINVAL``. Probe once with an aligned
35
+ single-page write so callers can fall back to buffered I/O instead of
36
+ failing on every block.
37
+ """
38
+ if not O_DIRECT:
39
+ return False
40
+ path = os.path.join(directory, f".o_direct_probe{_get_tmp_suffix()}")
41
+ page = mmap.mmap(-1, mmap.PAGESIZE)
42
+ try:
43
+ fd = os.open(path, os.O_CREAT | os.O_WRONLY | os.O_TRUNC | O_DIRECT, 0o644)
44
+ try:
45
+ os.write(fd, page)
46
+ finally:
47
+ os.close(fd)
48
+ return True
49
+ except OSError:
50
+ return False
51
+ finally:
52
+ page.close()
53
+ with contextlib.suppress(OSError):
54
+ os.remove(path)
55
+
56
+
57
def _ensure_dirs(path: str) -> None:
58
"""Create parent directories of *path* if they don't exist."""
59
os.makedirs(os.path.dirname(path), exist_ok=True)
@@ -34,6 +64,7 @@ def store_block(
64
buffer: memoryview,
65
offset: int,
66
block_size: int,
67
+ use_o_direct: bool = True,
68
) -> None:
69
"""
70
Store callback: Writes to a temp file then atomically replaces the destination.
@@ -49,10 +80,11 @@ def store_block(
80
# Write block atomically. Cast to a flat byte view so the slice uses byte
81
# indices; the raw memoryview may be multi-dimensional with itemsize > 1.
82
view_slice = buffer.cast("B")[offset : offset + block_size]
83
+ o_direct = O_DIRECT if use_o_direct else 0
84
try:
85
fd = os.open(
86
tmp_path,
55
- os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC | O_DIRECT,
87
+ os.O_CREAT | os.O_EXCL | os.O_WRONLY | os.O_TRUNC | o_direct,
88
0o644,
89
)
90
try:
@@ -77,14 +109,16 @@ def load_block(
109
view: memoryview,
110
offset: int,
111
block_size: int,
112
+ use_o_direct: bool = True,
113
) -> None:
114
"""
115
Load callback: read one KV block from disk. Remove the file on failure.
116
"""
117
fd: int | None = None
118
view_slice = view.cast("B")[offset : offset + block_size]
119
+ o_direct = O_DIRECT if use_o_direct else 0
120
try:
87
- fd = os.open(source_path, os.O_RDONLY | O_DIRECT)
121
+ fd = os.open(source_path, os.O_RDONLY | o_direct)
122
bytes_read = os.readv(fd, [view_slice])
123
if bytes_read < block_size:
124
raise OSError(f"Short read: expected {block_size} bytes, read {bytes_read}")
vllm/v1/kv_offload/tiering/fs/manager.py
+15
-1
@@ -49,7 +49,7 @@ from vllm.v1.kv_offload.tiering.base import (
49
ScheduleEndContext,
50
SecondaryTierManager,
51
)
52
-from vllm.v1.kv_offload.tiering.fs.io import load_block, store_block
52
+from vllm.v1.kv_offload.tiering.fs.io import load_block, probe_o_direct, store_block
53
from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool
54
55
if TYPE_CHECKING:
@@ -168,6 +168,18 @@ class FileSystemTierManager(SecondaryTierManager):
168
self.file_mapper.get_run_config(), f, indent=2, sort_keys=True
169
)
170
171
+ # Prefer O_DIRECT to bypass the page cache, but fall back to buffered
172
+ # I/O on filesystems that reject it (e.g. overlayfs, some NFS mounts)
173
+ # rather than failing every block.
174
+ self._use_o_direct = probe_o_direct(os.path.dirname(config_path))
175
+ if not self._use_o_direct:
176
+ logger.warning(
177
+ "O_DIRECT is not supported at '%s'; falling back to buffered "
178
+ "I/O for the '%s' KV offload tier.",
179
+ root_dir,
180
+ tier_type,
181
+ )
182
+
183
self._pool = DualQueueThreadPool(
184
n_read_threads,
185
n_write_threads,
@@ -198,6 +210,7 @@ class FileSystemTierManager(SecondaryTierManager):
210
self._primary_kv_view,
211
int(bid) * self._block_size,
212
self._block_size,
213
+ self._use_o_direct,
214
)
215
for key, bid in zip(job_metadata.keys, job_metadata.block_ids)
216
)
@@ -212,6 +225,7 @@ class FileSystemTierManager(SecondaryTierManager):
225
self._primary_kv_view,
226
int(bid) * self._block_size,
227
self._block_size,
228
+ self._use_o_direct,
229
)
230
for key, bid in zip(job_metadata.keys, job_metadata.block_ids)
231
)