| 1 | // SPDX-License-Identifier: Apache-2.0 |
| 2 | // SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | #include <Python.h> |
| 5 | |
| 6 | #include <errno.h> |
| 7 | #include <fcntl.h> |
| 8 | #include <unistd.h> |
| 9 | |
| 10 | #include <filesystem> |
| 11 | #include <string> |
| 12 | #include <vector> |
| 13 | |
| 14 | #if defined(O_DIRECT) |
| 15 | constexpr int kODirectFlag = O_DIRECT; |
| 16 | #else |
| 17 | constexpr int kODirectFlag = 0; |
| 18 | #endif |
| 19 | |
| 20 | extern "C" { |
| 21 | |
| 22 | namespace { |
| 23 | |
| 24 | // Returns 0 on success, or the std::error_code's POSIX-compatible value on |
| 25 | // failure, mirroring the errno convention used by the syscalls below. |
| 26 | inline int ensure_parent_dirs(const std::string& path) { |
| 27 | const auto parent = std::filesystem::path(path).parent_path(); |
| 28 | if (parent.empty()) { |
| 29 | return 0; |
| 30 | } |
| 31 | std::error_code ec; |
| 32 | std::filesystem::create_directories(parent, ec); |
| 33 | return ec ? ec.value() : 0; |
| 34 | } |
| 35 | |
| 36 | // Core single-block store: src/size are raw pointer + byte count. Returns 0 |
| 37 | // on success, or the errno of the failing step on failure -- captured |
| 38 | // before any subsequent cleanup call can overwrite it. On failure, the temp |
| 39 | // file is removed. |
| 40 | inline int _store_block(const char* tmp_path, const char* dest_path, |
| 41 | const char* src, size_t size, bool use_o_direct) { |
| 42 | if (access(dest_path, F_OK) == 0) { |
| 43 | return 0; // Already present. |
| 44 | } |
| 45 | |
| 46 | if (const int err = ensure_parent_dirs(dest_path); err != 0) { |
| 47 | return err; |
| 48 | } |
| 49 | |
| 50 | const int o_direct_flag = use_o_direct ? kODirectFlag : 0; |
| 51 | const int fd = open( |
| 52 | tmp_path, O_CREAT | O_EXCL | O_WRONLY | O_TRUNC | o_direct_flag, 0644); |
| 53 | if (fd < 0) { |
| 54 | return errno; |
| 55 | } |
| 56 | |
| 57 | const ssize_t written = write(fd, src, size); |
| 58 | if (written < 0 || static_cast<size_t>(written) != size) { |
| 59 | const int err = written < 0 ? errno : EIO; |
| 60 | close(fd); // Best-effort cleanup; the real error is already captured. |
| 61 | unlink(tmp_path); |
| 62 | return err; |
| 63 | } |
| 64 | |
| 65 | if (close(fd) != 0) { |
| 66 | const int err = errno; |
| 67 | unlink(tmp_path); |
| 68 | return err; |
| 69 | } |
| 70 | |
| 71 | if (rename(tmp_path, dest_path) != 0) { |
| 72 | const int err = errno; |
| 73 | unlink(tmp_path); |
| 74 | return err; |
| 75 | } |
| 76 | |
| 77 | return 0; |
| 78 | } |
| 79 | |
| 80 | // Core single-block load: dst/size are raw pointer + byte count. Returns 0 |
| 81 | // on success, or the errno of the failing step on failure. Removes the source |
| 82 | // file ONLY on a provable short read (the read completed but returned fewer |
| 83 | // bytes than requested): stores are atomic, so a too-short file is genuine |
| 84 | // corruption. Open failures and read errors (bytes_read < 0) are |
| 85 | // transient/ambiguous and leave the file untouched; a close failure after a |
| 86 | // full read is harmless and does not fail the load. |
| 87 | inline int _load_block(const char* source_path, char* dst, size_t size, |
| 88 | bool use_o_direct) { |
| 89 | const int o_direct_flag = use_o_direct ? kODirectFlag : 0; |
| 90 | const int fd = open(source_path, O_RDONLY | o_direct_flag, 0); |
| 91 | if (fd < 0) { |
| 92 | return errno; |
| 93 | } |
| 94 | |
| 95 | const ssize_t bytes_read = read(fd, dst, size); |
| 96 | if (bytes_read < 0) { |
| 97 | // Transient read error: leave the file untouched. |
| 98 | const int err = errno; |
| 99 | close(fd); |
| 100 | return err; |
| 101 | } |
| 102 | if (static_cast<size_t>(bytes_read) < size) { |
| 103 | // Provable short read: the block is genuinely corrupt, so remove it. |
| 104 | close(fd); |
| 105 | unlink(source_path); |
| 106 | return EIO; |
| 107 | } |
| 108 | |
| 109 | // A close error after a successful full read is harmless: the data is |
| 110 | // already in the destination buffer, so the load succeeds. |
| 111 | close(fd); |
| 112 | return 0; |
| 113 | } |
| 114 | |
| 115 | inline void _batch_lookup(const std::vector<const char*>& paths, |
| 116 | std::vector<int>& exists_flags) { |
| 117 | for (size_t i = 0; i < paths.size(); i++) { |
| 118 | exists_flags[i] = (access(paths[i], F_OK) == 0) ? 1 : 0; |
| 119 | } |
| 120 | } |
| 121 | |
| 122 | // Helper: extract a list[str] of length n into a vector<const char*>. |
| 123 | // Returns false and sets a Python exception on error. |
| 124 | inline bool extract_str_list(PyObject* list, Py_ssize_t n, |
| 125 | std::vector<const char*>& out) { |
| 126 | for (Py_ssize_t i = 0; i < n; i++) { |
| 127 | out[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(list, i), nullptr); |
| 128 | if (out[i] == nullptr) { |
| 129 | return false; |
| 130 | } |
| 131 | } |
| 132 | return true; |
| 133 | } |
| 134 | |
| 135 | // Helper: extract a Py_buffer per element of a list[bytes-like] of length n. |
| 136 | // On success, `out` holds n acquired buffers (caller must PyBuffer_Release |
| 137 | // each). On failure, any buffers already acquired are released before |
| 138 | // returning false, and a Python exception is set. |
| 139 | inline bool extract_buffer_list(PyObject* list, Py_ssize_t n, int flags, |
| 140 | std::vector<Py_buffer>& out) { |
| 141 | for (Py_ssize_t i = 0; i < n; i++) { |
| 142 | if (PyObject_GetBuffer(PyList_GetItem(list, i), &out[i], flags) != 0) { |
| 143 | for (Py_ssize_t j = 0; j < i; j++) { |
| 144 | PyBuffer_Release(&out[j]); |
| 145 | } |
| 146 | return false; |
| 147 | } |
| 148 | } |
| 149 | return true; |
| 150 | } |
| 151 | |
| 152 | inline void release_buffer_list(std::vector<Py_buffer>& buffers) { |
| 153 | for (auto& buf : buffers) { |
| 154 | PyBuffer_Release(&buf); |
| 155 | } |
| 156 | } |
| 157 | |
| 158 | } // namespace |
| 159 | |
| 160 | /// @brief Check file existence for a batch of paths. |
| 161 | /// @param paths list[str] – absolute paths to check. |
| 162 | /// @return list[bool] – True if the corresponding path exists, False otherwise. |
| 163 | /// @note Releases the GIL for the entire batch. File existence via access(2). |
| 164 | static PyObject* batch_lookup(PyObject* /*self*/, PyObject* args) { |
| 165 | PyObject* path_list; |
| 166 | if (!PyArg_ParseTuple(args, "O!", &PyList_Type, &path_list)) { |
| 167 | return nullptr; |
| 168 | } |
| 169 | |
| 170 | const Py_ssize_t n = PyList_Size(path_list); |
| 171 | std::vector<const char*> paths(n); |
| 172 | for (Py_ssize_t i = 0; i < n; i++) { |
| 173 | paths[i] = PyUnicode_AsUTF8AndSize(PyList_GetItem(path_list, i), nullptr); |
| 174 | if (paths[i] == nullptr) { |
| 175 | return nullptr; |
| 176 | } |
| 177 | } |
| 178 | |
| 179 | std::vector<int> exists_flags(n); |
| 180 | { |
| 181 | Py_BEGIN_ALLOW_THREADS _batch_lookup(paths, exists_flags); |
| 182 | Py_END_ALLOW_THREADS |
| 183 | } |
| 184 | |
| 185 | PyObject* result = PyList_New(n); |
| 186 | if (result == nullptr) { |
| 187 | return nullptr; |
| 188 | } |
| 189 | for (Py_ssize_t i = 0; i < n; i++) { |
| 190 | PyList_SetItem(result, i, PyBool_FromLong(exists_flags[i])); |
| 191 | } |
| 192 | return result; |
| 193 | } |
| 194 | |
| 195 | /// @brief Store a batch of blocks, each from its own buffer, to disk. |
| 196 | /// @param tmp_paths list[str] – one temp path per block. |
| 197 | /// @param dest_paths list[str] – one destination path per block. |
| 198 | /// @param buffers list[bytes-like] – one source buffer per block. |
| 199 | /// @param use_o_direct bool – whether to open files with O_DIRECT |
| 200 | /// (default True). Ignored where O_DIRECT is unsupported |
| 201 | /// by the platform. |
| 202 | /// @note Releases the GIL for the entire batch. Raises on first error. |
| 203 | static PyObject* batch_store_block(PyObject* /*self*/, PyObject* args) { |
| 204 | PyObject* tmp_paths_obj = nullptr; |
| 205 | PyObject* dest_paths_obj = nullptr; |
| 206 | PyObject* buffers_obj = nullptr; |
| 207 | int use_o_direct = 1; |
| 208 | |
| 209 | if (!PyArg_ParseTuple(args, "O!O!O!|p", &PyList_Type, &tmp_paths_obj, |
| 210 | &PyList_Type, &dest_paths_obj, &PyList_Type, |
| 211 | &buffers_obj, &use_o_direct)) { |
| 212 | return nullptr; |
| 213 | } |
| 214 | |
| 215 | const Py_ssize_t n = PyList_Size(tmp_paths_obj); |
| 216 | if (PyList_Size(dest_paths_obj) != n || PyList_Size(buffers_obj) != n) { |
| 217 | PyErr_SetString( |
| 218 | PyExc_ValueError, |
| 219 | "tmp_paths, dest_paths and buffers must have the same length"); |
| 220 | return nullptr; |
| 221 | } |
| 222 | |
| 223 | std::vector<const char*> tmp_paths(n); |
| 224 | std::vector<const char*> dest_paths(n); |
| 225 | |
| 226 | if (!extract_str_list(tmp_paths_obj, n, tmp_paths)) return nullptr; |
| 227 | if (!extract_str_list(dest_paths_obj, n, dest_paths)) return nullptr; |
| 228 | |
| 229 | std::vector<Py_buffer> buffers(n); |
| 230 | if (!extract_buffer_list(buffers_obj, n, PyBUF_SIMPLE, buffers)) { |
| 231 | return nullptr; |
| 232 | } |
| 233 | |
| 234 | Py_ssize_t failed_index = -1; |
| 235 | int failure_errno = 0; |
| 236 | |
| 237 | { |
| 238 | Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) { |
| 239 | const char* buf = static_cast<const char*>(buffers[i].buf); |
| 240 | const int err = |
| 241 | _store_block(tmp_paths[i], dest_paths[i], buf, |
| 242 | static_cast<size_t>(buffers[i].len), use_o_direct); |
| 243 | if (err != 0) { |
| 244 | failed_index = i; |
| 245 | failure_errno = err; |
| 246 | break; |
| 247 | } |
| 248 | } |
| 249 | Py_END_ALLOW_THREADS |
| 250 | } |
| 251 | |
| 252 | release_buffer_list(buffers); |
| 253 | |
| 254 | if (failed_index >= 0) { |
| 255 | // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. |
| 256 | errno = failure_errno; |
| 257 | return PyErr_SetFromErrnoWithFilename(PyExc_OSError, |
| 258 | dest_paths[failed_index]); |
| 259 | } |
| 260 | |
| 261 | Py_RETURN_NONE; |
| 262 | } |
| 263 | |
| 264 | /// @brief Load a batch of blocks from disk, each into its own buffer. |
| 265 | /// @param source_paths list[str] – one source path per block. |
| 266 | /// @param buffers list[writable bytes-like] – one destination buffer |
| 267 | /// per block. |
| 268 | /// @param use_o_direct bool – whether to open files with O_DIRECT |
| 269 | /// (default True). Ignored where O_DIRECT is unsupported |
| 270 | /// by the platform. |
| 271 | /// @note Releases the GIL for the entire batch. Raises on first error. |
| 272 | static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) { |
| 273 | PyObject* source_paths_obj = nullptr; |
| 274 | PyObject* buffers_obj = nullptr; |
| 275 | int use_o_direct = 1; |
| 276 | |
| 277 | if (!PyArg_ParseTuple(args, "O!O!|p", &PyList_Type, &source_paths_obj, |
| 278 | &PyList_Type, &buffers_obj, &use_o_direct)) { |
| 279 | return nullptr; |
| 280 | } |
| 281 | |
| 282 | const Py_ssize_t n = PyList_Size(source_paths_obj); |
| 283 | if (PyList_Size(buffers_obj) != n) { |
| 284 | PyErr_SetString(PyExc_ValueError, |
| 285 | "source_paths and buffers must have the same length"); |
| 286 | return nullptr; |
| 287 | } |
| 288 | |
| 289 | std::vector<const char*> source_paths(n); |
| 290 | if (!extract_str_list(source_paths_obj, n, source_paths)) return nullptr; |
| 291 | |
| 292 | std::vector<Py_buffer> buffers(n); |
| 293 | if (!extract_buffer_list(buffers_obj, n, PyBUF_WRITABLE, buffers)) { |
| 294 | return nullptr; |
| 295 | } |
| 296 | |
| 297 | Py_ssize_t failed_index = -1; |
| 298 | int failure_errno = 0; |
| 299 | |
| 300 | { |
| 301 | Py_BEGIN_ALLOW_THREADS for (Py_ssize_t i = 0; i < n; i++) { |
| 302 | char* buf = static_cast<char*>(buffers[i].buf); |
| 303 | const int err = |
| 304 | _load_block(source_paths[i], buf, static_cast<size_t>(buffers[i].len), |
| 305 | use_o_direct); |
| 306 | if (err != 0) { |
| 307 | failed_index = i; |
| 308 | failure_errno = err; |
| 309 | break; |
| 310 | } |
| 311 | } |
| 312 | Py_END_ALLOW_THREADS |
| 313 | } |
| 314 | |
| 315 | release_buffer_list(buffers); |
| 316 | |
| 317 | if (failed_index >= 0) { |
| 318 | // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. |
| 319 | errno = failure_errno; |
| 320 | PyErr_SetFromErrnoWithFilename(PyExc_OSError, source_paths[failed_index]); |
| 321 | // Attach the number of blocks that loaded before the failure so the tier |
| 322 | // can keep them (partial success). failed_index == count of blocks read OK. |
| 323 | PyObject *etype, *evalue, *etb; |
| 324 | PyErr_Fetch(&etype, &evalue, &etb); |
| 325 | PyErr_NormalizeException(&etype, &evalue, &etb); |
| 326 | if (evalue != nullptr) { |
| 327 | PyObject* num = PyLong_FromSsize_t(failed_index); |
| 328 | if (num != nullptr) { |
| 329 | PyObject_SetAttrString(evalue, "num_succeeded", num); |
| 330 | Py_DECREF(num); |
| 331 | } |
| 332 | } |
| 333 | PyErr_Restore(etype, evalue, etb); |
| 334 | return nullptr; |
| 335 | } |
| 336 | |
| 337 | Py_RETURN_NONE; |
| 338 | } |
| 339 | |
| 340 | static PyMethodDef fs_io_C_methods[] = { |
| 341 | {"batch_lookup", batch_lookup, METH_VARARGS, |
| 342 | "batch_lookup(paths: list[str]) -> list[bool]\n" |
| 343 | "\n" |
| 344 | "Check file existence for a batch of paths."}, |
| 345 | {"batch_store_block", batch_store_block, METH_VARARGS, |
| 346 | "batch_store_block(tmp_paths: list[str], dest_paths: list[str],\n" |
| 347 | " buffers: list[bytes-like],\n" |
| 348 | " use_o_direct: bool = True) -> None\n" |
| 349 | "\n" |
| 350 | "Store a batch of blocks, each from its own buffer, to disk. Raises on " |
| 351 | "first error."}, |
| 352 | {"batch_load_block", batch_load_block, METH_VARARGS, |
| 353 | "batch_load_block(source_paths: list[str],\n" |
| 354 | " buffers: list[writable bytes-like],\n" |
| 355 | " use_o_direct: bool = True) -> None\n" |
| 356 | "\n" |
| 357 | "Load a batch of blocks from disk into corresponding buffers. " |
| 358 | "Raises on first error."}, |
| 359 | {nullptr, nullptr, 0, nullptr}, |
| 360 | }; |
| 361 | |
| 362 | static struct PyModuleDef fs_io_C_module = { |
| 363 | PyModuleDef_HEAD_INIT, "fs_io_C", "Filesystem helpers for KV offload", -1, |
| 364 | fs_io_C_methods, |
| 365 | }; |
| 366 | |
| 367 | PyMODINIT_FUNC PyInit_fs_io_C(void) { return PyModule_Create(&fs_io_C_module); } |
| 368 | |
| 369 | } // extern "C" |