diff --git a/CMakeLists.txt b/CMakeLists.txt index cbd5583bbf..1ef9d596ae 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -140,6 +140,21 @@ if(Python_VERSION VERSION_GREATER_EQUAL "3.11") WITH_SOABI) endif() +# +# fs_io extension (pure CXX; must stay above the non-CUDA device branch +# so CPU builds define the target before the early return). +# GIL-releasing filesystem helpers for FileSystemTierManager. +# +if(Python_VERSION VERSION_GREATER_EQUAL "3.11") + define_extension_target( + fs_io_C + DESTINATION vllm + LANGUAGE CXX + SOURCES csrc/fs_io.cpp + USE_SABI 3.11 + WITH_SOABI) +endif() + # # Forward the non-CUDA device extensions to external CMake scripts. # diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp new file mode 100644 index 0000000000..fdf3e614e6 --- /dev/null +++ b/csrc/fs_io.cpp @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#include + +#include + +#include + +extern "C" { + +static void _batch_lookup(const std::vector& paths, + std::vector& exists_flags) { + for (size_t i = 0; i < paths.size(); i++) { + exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0; + } +} + +/// @brief Check file existence for a batch of paths. +/// @param paths list[str] – absolute paths to check. +/// @return list[bool] – True if the corresponding path exists, False otherwise. +/// @note Releases the GIL for the entire batch. File existence via access(2). +static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) { + PyObject* path_list; + if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) { + return nullptr; + } + + const Py_ssize_t n = PyList_Size(path_list); + std::vector paths(n); + for (Py_ssize_t i = 0; i < n; i++) { + paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr); + if (paths[i] == nullptr) { + return nullptr; + } + } + + std::vector exists_flags(n); + { + Py_BEGIN_ALLOW_THREADS _batch_lookup(paths, exists_flags); + Py_END_ALLOW_THREADS + } + + PyObject* result = PyList_New(n); + if (result == nullptr) { + return nullptr; + } + for (Py_ssize_t i = 0; i < n; i++) { + PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i])); + } + return result; +} + +static PyMethodDef fs_io_C_methods[] = { + {"batch_lookup", batch_lookup, METH_VARARGS, + "batch_lookup(paths: list[str]) -> list[bool]\n" + "\n" + "Check file existence for a batch of paths."}, + {nullptr, nullptr, 0, nullptr}, +}; + +static struct PyModuleDef fs_io_C_module = { + PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1, + fs_io_C_methods, +}; + +PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); } + +} // extern "C" diff --git a/setup.py b/setup.py index ad9aed07a3..b305fb1b00 100644 --- a/setup.py +++ b/setup.py @@ -777,6 +777,7 @@ class precompiled_wheel_utils: "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so", "vllm/cumem_allocator.abi3.so", "vllm/spinloop.abi3.so", + "vllm/fs_io_C.abi3.so", # ROCm-specific libraries "vllm/_rocm_C.abi3.so", } @@ -1104,6 +1105,7 @@ if _is_cuda() or _is_hip(): if sys.version_info >= (3, 11): ext_modules.append(CMakeExtension(name="vllm.spinloop")) + ext_modules.append(CMakeExtension(name="vllm.fs_io_C")) if _is_hip(): ext_modules.append(CMakeExtension(name="vllm._rocm_C")) diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 7245ae1ba7..0300fb5d4d 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -103,7 +103,7 @@ def lookup_and_wait( keys: list[OffloadKey], ctx: ReqContext = _CTX, timeout: float = 1.0, -) -> list[bool]: +) -> list[LookupResult]: """Perform a full async lookup cycle and return resolved results.""" for k in keys: tier.lookup(k, ctx) @@ -332,3 +332,66 @@ def test_wait_idle_blocks_until_tasks_complete(): gate.set() pool.shutdown(wait=True) waiter.join(timeout=5.0) + + +def test_batch_lookup_c_extension(tmp_path): + """Validates batch_lookup_C: empty, single, all-existing, all-missing, + mixed ordering, and input type validation.""" + try: + from vllm.fs_io_C import batch_lookup as batch_lookup_C + except ImportError: + pytest.skip("fs_io_C extension not built") + + # Setup + all_exist = [str(tmp_path / f"e{i}.bin") for i in range(3)] + for p in all_exist: + open(p, "w").close() + all_missing = [str(tmp_path / f"m{i}.bin") for i in range(3)] + + # Empty list + assert batch_lookup_C([]) == [] + + # Single existing / missing + assert batch_lookup_C([all_exist[0]]) == [True] + assert batch_lookup_C([all_missing[0]]) == [False] + + # All existing / all missing + assert batch_lookup_C(all_exist) == [True, True, True] + assert batch_lookup_C(all_missing) == [False, False, False] + + # Mixed — verifies index ordering is preserved + paths = [val for pair in zip(all_exist, all_missing) for val in pair] + assert batch_lookup_C(paths) == [True, False, True, False, True, False] + + # Input validation: non-list argument + with pytest.raises(TypeError): + batch_lookup_C(("/tmp/foo",)) + with pytest.raises(TypeError): + batch_lookup_C(None) + + # Input validation: non-str elements in list + with pytest.raises(TypeError): + batch_lookup_C([None]) + with pytest.raises(TypeError): + batch_lookup_C([b"/tmp/foo"]) + with pytest.raises(TypeError): + batch_lookup_C([42]) + with pytest.raises(TypeError): + batch_lookup_C([all_exist[0], None]) # valid first, invalid mid-list + + +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_batch_lookup_dispatch(fs_tier, monkeypatch, use_c_ext): + import vllm.v1.kv_offload.tiering.fs.manager as mgr_mod + + if use_c_ext and not mgr_mod._HAS_BATCH_LOOKUP_C: + pytest.skip("fs_io_C extension not built") + + monkeypatch.setattr(mgr_mod, "_HAS_BATCH_LOOKUP_C", use_c_ext) + + tier, _ = fs_tier + tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(tier)) + + results = lookup_and_wait(tier, [key(1), key(2)]) + assert results == [LookupResult.HIT, LookupResult.MISS] diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 329a24daf3..816e88c522 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -21,6 +21,13 @@ import os from collections.abc import Iterable from typing import TYPE_CHECKING +try: + from vllm.fs_io_C import batch_lookup as batch_lookup_C + + _HAS_BATCH_LOOKUP_C = True +except ImportError: + _HAS_BATCH_LOOKUP_C = False + from typing_extensions import override from vllm.logger import init_logger @@ -56,7 +63,11 @@ class FsAsyncLookupManager(AsyncLookupManager): def batch_lookup( self, keys: list[OffloadKey], req_context: ReqContext ) -> Iterable[bool]: - return (os.path.exists(self._tier.file_mapper.get_file_name(k)) for k in keys) + paths = [self._tier.file_mapper.get_file_name(k) for k in keys] + if _HAS_BATCH_LOOKUP_C: + # C extension: GIL released for the entire faccessat() batch. + return batch_lookup_C(paths) + return (os.path.exists(p) for p in paths) class FileSystemTierManager(SecondaryTierManager):