diff --git a/src/plugins/kv/kv_backend.cpp b/src/plugins/kv/kv_backend.cpp new file mode 100644 index 0000000000..aa8b9205a5 --- /dev/null +++ b/src/plugins/kv/kv_backend.cpp @@ -0,0 +1,122 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "kv_backend.h" + +#include +#include +#include + +#include "common/nixl_log.h" +#include "nixl_descriptors.h" +#include "nixl_types.h" + +// ── registerMem ───────────────────────────────────────────────────────────── + +nixl_status_t +nixlKVBackendBase::registerMem(const nixlBlobDesc &mem, + const nixl_mem_t &nixl_mem, + nixlBackendMD *&out) { + const auto supported = getSupportedMems(); + if (std::find(supported.begin(), supported.end(), nixl_mem) == supported.end()) { + NIXL_ERROR << "KV backend: unsupported memory type " << nixl_mem; + out = nullptr; + return NIXL_ERR_NOT_SUPPORTED; + } + + // Key policy: use metaInfo if non-empty, otherwise fall back to devId string. + std::string kv_key = mem.metaInfo.empty() ? std::to_string(mem.devId) : mem.metaInfo; + + auto md = std::make_unique(nixl_mem, mem.devId, kv_key); + devIdToKey_[mem.devId] = kv_key; + out = md.release(); + return NIXL_SUCCESS; +} + +// ── deregisterMem ─────────────────────────────────────────────────────────── + +nixl_status_t +nixlKVBackendBase::deregisterMem(nixlBackendMD *meta) { + auto *kv_md = static_cast(meta); + if (kv_md) { + std::unique_ptr owned(kv_md); + devIdToKey_.erase(kv_md->devId); + } + return NIXL_SUCCESS; +} + +// ── resolveKey ────────────────────────────────────────────────────────────── + +bool +nixlKVBackendBase::resolveKey(const nixlMetaDesc &desc, std::string &out_key) const { + // Step 1: try the per-descriptor metadata pointer. + auto *kv_md = dynamic_cast(desc.metadataP); + if (kv_md) { + out_key = kv_md->key; + return true; + } + + // Step 2: fall back to the devId → key map populated at registerMem time. + auto it = devIdToKey_.find(desc.devId); + if (it != devIdToKey_.end()) { + out_key = it->second; + return true; + } + + return false; +} + +// ── checkXfer ─────────────────────────────────────────────────────────────── + +nixl_status_t +nixlKVBackendBase::checkXfer(nixlBackendReqH *handle) const { + auto *req = static_cast(handle); + if (!req) { + return NIXL_ERR_INVALID_PARAM; + } + + if (req->pending.load(std::memory_order_acquire) > 0) { + return NIXL_ERR_NOT_POSTED; // IN_PROG sentinel used by NIXL + } + + if (req->first_error.load(std::memory_order_relaxed) != 0) { + return NIXL_ERR_BACKEND; + } + + return NIXL_SUCCESS; +} + +// ── releaseReqH ───────────────────────────────────────────────────────────── + +nixl_status_t +nixlKVBackendBase::releaseReqH(nixlBackendReqH *handle) const { + auto *req = static_cast(handle); + if (!req) { + return NIXL_ERR_INVALID_PARAM; + } + + // Block until all in-flight operations have called nixlKVXferCallback. + { + std::unique_lock lk(req->mu); + req->cv.wait(lk, [req] { + return req->pending.load(std::memory_order_acquire) == 0; + }); + } + + delete req; + return NIXL_SUCCESS; +} diff --git a/src/plugins/kv/kv_backend.h b/src/plugins/kv/kv_backend.h new file mode 100644 index 0000000000..f6ca7a35bb --- /dev/null +++ b/src/plugins/kv/kv_backend.h @@ -0,0 +1,187 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_PLUGINS_KV_KV_BACKEND_H +#define NIXL_SRC_PLUGINS_KV_KV_BACKEND_H + +#include +#include +#include + +#include "backend/backend_engine.h" +#include "kv_req_handle.h" + +/** + * @brief Abstract base class for NIXL key-value storage backend plugins. + * + * nixlKVBackendBase factors out the boilerplate that is identical across all + * KV-family plugins (REDIS, WQSKV, and future backends): + * + * - Capability flags: supportsRemote/Local/Notif all return fixed values. + * - getSupportedMems(): defaults to {DRAM_SEG}; override to add OBJ_SEG etc. + * - Key management: metaInfo-or-devId resolution, devIdToKey_ map, + * registerMem / deregisterMem, resolveKey helper. + * - Async request handle lifecycle: checkXfer (poll pending/first_error) and + * releaseReqH (block on CV until done) operating on nixlKVReqH. + * - No-op lifecycle stubs: connect, disconnect, unloadMD, loadLocalMD. + * + * Concrete backends must implement: + * - prepXfer — allocate a nixlKVReqH (or subclass) and validate inputs + * - postXfer — set pending = n, dispatch vendor async ops, return immediately + * - queryMem — backend-specific memory query (strong vs. local-only semantics) + * + * Each backend plugin registers itself via nixlBackendPluginCreator. + * + * Directory convention: + * src/plugins/kv/ ← this shared base + * src/plugins/kv/redis/ ← REDIS plugin (links kv_base) + * src/plugins/kv/wqskv/ ← WQSKV plugin (links kv_base) + * src/plugins/kv// ← additional KV backends + */ +class nixlKVBackendBase : public nixlBackendEngine { +public: + explicit nixlKVBackendBase(const nixlBackendInitParams *init_params) + : nixlBackendEngine(init_params) {} + + ~nixlKVBackendBase() override = default; + + // ── Capability flags ─────────────────────────────────────────────────── + bool supportsRemote() const override { return false; } + bool supportsLocal() const override { return true; } + bool supportsNotif() const override { return false; } + + /** + * @brief Returns {DRAM_SEG} by default. + * + * Backends that also accept OBJ_SEG (e.g., REDIS) override this. + */ + nixl_mem_list_t getSupportedMems() const override { return {DRAM_SEG}; } + + // ── Key management ───────────────────────────────────────────────────── + + /** + * @brief Register a memory region and resolve its KV key. + * + * Key policy (shared across all KV backends): + * key = mem.metaInfo (if non-empty) + * else std::to_string(mem.devId) + * + * Stores the mapping in devIdToKey_ and creates a nixlKVMetadata output. + * Returns NIXL_ERR_NOT_SUPPORTED if nixl_mem is not in getSupportedMems(). + */ + nixl_status_t registerMem(const nixlBlobDesc &mem, + const nixl_mem_t &nixl_mem, + nixlBackendMD *&out) override; + + /** + * @brief Deregister a memory region and remove its key mapping. + */ + nixl_status_t deregisterMem(nixlBackendMD *meta) override; + + // ── Async request handle lifecycle ───────────────────────────────────── + + /** + * @brief Non-blocking transfer status poll. + * + * Returns NIXL_IN_PROG while pending > 0, NIXL_ERR_BACKEND on first error, + * NIXL_SUCCESS when all operations complete. + * + * Expects handle to be a nixlKVReqH (or subclass). + */ + nixl_status_t checkXfer(nixlBackendReqH *handle) const override; + + /** + * @brief Block until all in-flight ops finish, then delete the handle. + * + * Waits on nixlKVReqH::cv until pending == 0. + */ + nixl_status_t releaseReqH(nixlBackendReqH *handle) const override; + + // ── No-op lifecycle stubs ─────────────────────────────────────────────── + + nixl_status_t connect(const std::string &) override { return NIXL_SUCCESS; } + nixl_status_t disconnect(const std::string &) override { return NIXL_SUCCESS; } + nixl_status_t unloadMD(nixlBackendMD *) override { return NIXL_SUCCESS; } + + nixl_status_t loadLocalMD(nixlBackendMD *input, nixlBackendMD *&output) override { + if (!input) { + output = nullptr; + return NIXL_ERR_INVALID_PARAM; + } + output = input; + return NIXL_SUCCESS; + } + + // ── Pure virtuals that each backend must provide ──────────────────────── + + nixl_status_t prepXfer(const nixl_xfer_op_t &operation, + const nixl_meta_dlist_t &local, + const nixl_meta_dlist_t &remote, + const std::string &remote_agent, + nixlBackendReqH *&handle, + const nixl_opt_b_args_t *opt_args = nullptr) const override = 0; + + nixl_status_t postXfer(const nixl_xfer_op_t &operation, + const nixl_meta_dlist_t &local, + const nixl_meta_dlist_t &remote, + const std::string &remote_agent, + nixlBackendReqH *&handle, + const nixl_opt_b_args_t *opt_args = nullptr) const override = 0; + + nixl_status_t queryMem(const nixl_reg_dlist_t &descs, + std::vector &resp) const override = 0; + +protected: + /** + * @brief Shared metadata type for all KV backends. + * + * Holds the resolved KV key and the devId used for map lookups. + * Backends that need additional fields may subclass this. + */ + class nixlKVMetadata : public nixlBackendMD { + public: + nixlKVMetadata(nixl_mem_t mem_type, uint64_t dev_id, std::string kv_key) + : nixlBackendMD(/*isPrivate=*/true), + memType(mem_type), + devId(dev_id), + key(std::move(kv_key)) {} + + nixl_mem_t memType; + uint64_t devId; + std::string key; + }; + + /** + * @brief Map from devId → resolved KV key, populated by registerMem. + * + * Used as a fallback when metadataP cannot be cast to nixlKVMetadata + * (e.g., for remote descriptors that arrive without local metadata). + */ + std::unordered_map devIdToKey_; + + /** + * @brief Resolve the KV key for a descriptor. + * + * Two-step lookup: + * 1. Cast desc.metadataP to nixlKVMetadata — use its key if valid. + * 2. Fall back to devIdToKey_[desc.devId]. + * + * @return true if a key was found, false if the descriptor is unmapped. + */ + bool resolveKey(const nixlMetaDesc &desc, std::string &out_key) const; +}; + +#endif // NIXL_SRC_PLUGINS_KV_KV_BACKEND_H diff --git a/src/plugins/kv/kv_req_handle.h b/src/plugins/kv/kv_req_handle.h new file mode 100644 index 0000000000..b2611c2552 --- /dev/null +++ b/src/plugins/kv/kv_req_handle.h @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef NIXL_SRC_PLUGINS_KV_KV_REQ_HANDLE_H +#define NIXL_SRC_PLUGINS_KV_KV_REQ_HANDLE_H + +#include +#include +#include +#include +#include + +#include "backend/backend_aux.h" + +/** + * @brief Shared async request handle for KV-family backend plugins. + * + * Tracks in-flight descriptors via an atomic pending counter and a condition + * variable. The C-style callback nixlKVXferCallback() decrements the counter + * and signals the CV when all operations complete, making it compatible with + * both C++ std::async and C vendor callback APIs. + * + * Lifecycle contract (same as the NIXL async transfer protocol): + * postXfer — sets pending = n, dispatches n async ops, returns immediately + * checkXfer — reads pending / first_error atomically (non-blocking) + * releaseReqH — waits on CV until pending == 0, then deletes the handle + */ +class nixlKVReqH : public nixlBackendReqH { +public: + /// Number of async operations still in flight. + std::atomic pending{0}; + + /// First non-zero error code returned by any operation (set via CAS). + std::atomic first_error{0}; + + /// CV + mutex used by releaseReqH to block until pending reaches zero. + std::mutex mu; + std::condition_variable cv; + + /// Per-descriptor state kept alive for the duration of the async call. + struct DescState { + std::string key; + std::vector ioVec; + }; + std::vector descStates; + + nixlKVReqH() = default; + ~nixlKVReqH() override = default; +}; + +/** + * @brief Standard C-style completion callback for KV async operations. + * + * Backends pass this (or a thin wrapper) as the completion callback to their + * vendor API. @p rc is the vendor return code (0 = success), @p arg is the + * nixlKVReqH pointer cast to void*. + * + * The callback: + * 1. Records the first non-zero rc via compare-exchange. + * 2. Decrements pending. + * 3. Notifies the CV when the last operation completes (pending reaches 0). + */ +inline void +nixlKVXferCallback(int rc, void *arg) { + auto *req = static_cast(arg); + if (rc != 0) { + int expected = 0; + req->first_error.compare_exchange_strong(expected, rc, std::memory_order_relaxed); + } + if (req->pending.fetch_sub(1, std::memory_order_acq_rel) == 1) { + // Last in-flight operation completed — wake up releaseReqH. + std::lock_guard lk(req->mu); + req->cv.notify_all(); + } +} + +#endif // NIXL_SRC_PLUGINS_KV_KV_REQ_HANDLE_H diff --git a/src/plugins/kv/meson.build b/src/plugins/kv/meson.build new file mode 100644 index 0000000000..fc7e3126c5 --- /dev/null +++ b/src/plugins/kv/meson.build @@ -0,0 +1,70 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ── Shared KV backend abstraction layer ──────────────────────────────────── +# +# nixl_kv_base is a static utility library that provides the shared base class +# (nixlKVBackendBase), shared request handle (nixlKVReqH), and the standard +# completion callback (nixlKVXferCallback) for all KV-family backend plugins. +# +# Plugins link against nixl_kv_base_dep: +# +# src/plugins/kv/redis/ REDIS plugin links nixl_kv_base_dep +# src/plugins/kv/wqskv/ WQSKV plugin links nixl_kv_base_dep +# src/plugins/kv// future backends links nixl_kv_base_dep +# +# Each plugin subdirectory is added via its own conditional subdir() call below. +# The subdir() calls are intentionally commented out until each PR lands. + +kv_base_sources = [ + 'kv_backend.cpp', + 'kv_backend.h', + 'kv_req_handle.h', +] + +kv_base_compile_flags = [] +if is_variable('compile_flags') + kv_base_compile_flags = compile_flags +endif + +nixl_kv_base_lib = static_library( + 'nixl_kv_base', + kv_base_sources, + dependencies: [nixl_infra, nixl_common_dep], + cpp_args: kv_base_compile_flags, + include_directories: [nixl_inc_dirs, utils_inc_dirs, + include_directories('.')], + install: false, +) + +# Dependency object consumed by plugin subdirectories. +nixl_kv_base_dep = declare_dependency( + link_with: nixl_kv_base_lib, + dependencies: [nixl_infra, nixl_common_dep], + include_directories: [nixl_inc_dirs, utils_inc_dirs, + include_directories('.')], +) + +# ── Plugin subdirectories ────────────────────────────────────────────────── +# Uncomment each line when the corresponding PR is merged and the plugin +# directory is moved under src/plugins/kv/. +# +# if enabled_plugins.get('REDIS') +# subdir('redis') +# endif +# +# if enabled_plugins.get('WQSKV') +# subdir('wqskv') +# endif diff --git a/src/plugins/meson.build b/src/plugins/meson.build index 055ed17c56..5bf43db5e6 100644 --- a/src/plugins/meson.build +++ b/src/plugins/meson.build @@ -29,6 +29,10 @@ if enabled_plugins.get('OBJ') subdir('obj') endif +if enabled_plugins.get('KV') + subdir('kv') +endif + if enabled_plugins.get('AZURE_BLOB') subdir('azure_blob') endif