diff --git a/.gitignore b/.gitignore index 3028d191..5f7e1004 100644 --- a/.gitignore +++ b/.gitignore @@ -57,3 +57,6 @@ multi-storage-client/src/multistorageclient/explorer/static/ .windsurf/ CLAUDE.md .mcp.json + +# Local cuObject SDK copy for the optional MSC `rdma` feature (never committed) +multi-storage-client/rust/vendor/ diff --git a/.release_notes/.unreleased.md b/.release_notes/.unreleased.md index f4197c84..c93527c5 100644 --- a/.release_notes/.unreleased.md +++ b/.release_notes/.unreleased.md @@ -4,4 +4,6 @@ ## New Features +- Add an optional S3-over-RDMA data plane to the `s3` provider (`rdma` option), backed by NVIDIA cuObject for direct buffer transfers to RDMA-capable endpoints. + ## Bug Fixes diff --git a/multi-storage-client/examples/rdma_roundtrip.py b/multi-storage-client/examples/rdma_roundtrip.py new file mode 100644 index 00000000..fc96b8a6 --- /dev/null +++ b/multi-storage-client/examples/rdma_roundtrip.py @@ -0,0 +1,94 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +""" +End-to-end validation of the S3-over-RDMA (NVIDIA cuObject) data plane for the +``s3`` storage provider, mirroring the PyTorch cuObject checkpoint roundtrip. + +This requires an RDMA-capable S3 endpoint and the cuObject runtime, so it is a +standalone script rather than a CI test. Run it on a cluster client (e.g. coe09) +against the RDMA MinIO endpoint (e.g. coe01:9000): + + # Build the wheel once on a host with the cuObject runtime, enabling the + # crate's `rdma` feature so the cuObject shim is compiled into the extension: + # maturin develop --features rdma # into the active venv, or + # maturin build --features rdma # a distributable wheel + + export CUFILE_ENV_PATH_JSON=/path/to/cuobj.json # rdma_dev_addr_list, use_pci_p2pdma, ... + export LD_LIBRARY_PATH=/path/to/sdklib:$LD_LIBRARY_PATH # version-matched libcufile/libcuobjclient + export S3_ENDPOINT=http://coe01:9000 + export S3_BUCKET=rdma-test + export AWS_ACCESS_KEY_ID=minioadmin AWS_SECRET_ACCESS_KEY=minioadmin + + python examples/rdma_roundtrip.py +""" + +import os + +import multistorageclient as msc + +_ENDPOINT = os.environ["S3_ENDPOINT"] +_BUCKET = os.environ["S3_BUCKET"] +_PROFILE = "rdma-test" + + +def _config() -> dict: + return { + "profiles": { + _PROFILE: { + "storage_provider": { + # The `s3_cuobject` provider moves payloads over RDMA via + # cuObject; it subclasses `s3` and forces the empty-body, + # unsigned-payload wire contract. + "type": "s3_cuobject", + "options": { + "base_path": _BUCKET, + "endpoint_url": _ENDPOINT, + "region_name": "us-east-1", + # MinIO AIStor RDMA endpoints are path-addressed; the + # provider leaves addressing style to the user. + "s3": {"addressing_style": "path"}, + # Optional: tune the RDMA multipart part size. + # "rdma": {"multipart_chunksize": 536870912}, + }, + } + } + } + } + + +def main() -> None: + config = msc.StorageClientConfig.from_dict(_config(), profile=_PROFILE) + client = msc.StorageClient(config) + + # A few sizes spanning the old multipart threshold (64 MiB); all must take + # the single-shot RDMA path. + for mib in (1, 64, 256): + size = mib * 1024 * 1024 + payload = os.urandom(size) + key = f"rdma_roundtrip/blob_{mib}mib.bin" + + client.write(key, payload) + reloaded = bytes(client.read(key)) + + assert len(reloaded) == size, f"{key}: size {len(reloaded)} != {size}" + assert reloaded == payload, f"{key}: byte mismatch after RDMA roundtrip" + print(f"OK {key} ({mib} MiB) byte-identical over RDMA") + + print("\nAll cuObject RDMA roundtrips passed.") + + +if __name__ == "__main__": + main() diff --git a/multi-storage-client/rust/Cargo.lock b/multi-storage-client/rust/Cargo.lock index 4b211082..45805ccc 100644 --- a/multi-storage-client/rust/Cargo.lock +++ b/multi-storage-client/rust/Cargo.lock @@ -1238,8 +1238,10 @@ dependencies = [ "aws-credential-types", "aws-smithy-http-client", "bytes", + "cc", "chrono", "http 1.4.2", + "libc", "object_store", "pyo3", "pyo3-async-runtimes", diff --git a/multi-storage-client/rust/Cargo.toml b/multi-storage-client/rust/Cargo.toml index 633ba9e5..002080d2 100644 --- a/multi-storage-client/rust/Cargo.toml +++ b/multi-storage-client/rust/Cargo.toml @@ -13,6 +13,17 @@ publish = false name = "multistorageclient_rust" crate-type = ["cdylib"] +[features] +default = [] +# NVIDIA cuObject (S3-over-RDMA) data plane. Off by default; the default build +# needs no C++ compiler or cuObject SDK. When enabled, build.rs compiles the +# extern "C" shim (csrc/cuobj_shim.cc) against the vendored cuObject headers/libs +# (vendor/cuobj/) and links libcuobjclient/libcufile. Linux only. +rdma = ["dep:libc"] + +[build-dependencies] +cc = "1.2" + [dependencies] pyo3 = "0.29" pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } @@ -31,3 +42,4 @@ async-trait = "0.1.89" aws-config = { version = "1.8.18", default-features = false, features = ["rt-tokio", "credentials-process", "sso"] } aws-credential-types = "1.2.14" aws-smithy-http-client = { version = "1.1.13", default-features = false, features = ["rustls-ring"] } +libc = { version = "0.2", optional = true } diff --git a/multi-storage-client/rust/build.rs b/multi-storage-client/rust/build.rs new file mode 100644 index 00000000..f399fadf --- /dev/null +++ b/multi-storage-client/rust/build.rs @@ -0,0 +1,125 @@ +// SPDX-FileCopyrightText: Copyright (c) 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. + +use std::env; +use std::path::PathBuf; + +fn main() { + println!("cargo:rerun-if-changed=build.rs"); + println!("cargo:rerun-if-changed=csrc/cuobj_shim.cc"); + println!("cargo:rerun-if-changed=csrc/cuobj_shim.h"); + println!("cargo:rerun-if-env-changed=MSC_CUOBJ_HOME"); + println!("cargo:rerun-if-env-changed=CUDA_HOME"); + println!("cargo:rerun-if-env-changed=CUDA_PATH"); + println!("cargo:rerun-if-env-changed=CUDA_ROOT"); + + // The cuObject (RDMA) data plane is opt-in. Without the `rdma` feature the + // crate builds as a pure-Rust cdylib with no C++ compiler or cuObject SDK, + // so build.rs is a no-op. + if env::var("CARGO_FEATURE_RDMA").is_err() { + return; + } + + if env::var("CARGO_CFG_TARGET_OS").as_deref() != Ok("linux") { + panic!("`rdma` feature is only supported on Linux"); + } + + let arch = match env::var("CARGO_CFG_TARGET_ARCH").as_deref() { + Ok("x86_64") => "x86_64", + Ok("aarch64") => "aarch64", + Ok(other) => panic!("`rdma` feature: unsupported target arch `{other}`"), + Err(_) => panic!("CARGO_CFG_TARGET_ARCH not set"), + }; + + // Resolve the cuObject SDK (headers + libs) at build time. Nothing is + // vendored into the repository: the RDMA feature expects the NVIDIA + // cuObject SDK (CUDA Toolkit >= 13.1) to be installed. Resolution order: + // 1. MSC_CUOBJ_HOME -> {include, lib} + // 2. CUDA_HOME/CUDA_PATH -> targets//{include, lib} + // 3. ./vendor/cuobj -> {include, lib/} (local dev only) + let (include_dir, lib_dir) = resolve_cuobj_dirs(arch) + .expect( + "`rdma` feature: could not locate the cuObject SDK. Set MSC_CUOBJ_HOME (a directory \ + with include/ and lib/), or CUDA_HOME/CUDA_PATH to a CUDA Toolkit (>= 13.1) that \ + ships cuobjclient.h + libcuobjclient.", + ); + + cc::Build::new() + .cpp(true) + .std("c++17") + .file("csrc/cuobj_shim.cc") + .include(&include_dir) + .flag_if_supported("-Wno-unused-parameter") + .flag_if_supported("-fvisibility=hidden") + .compile("miniors_cuobj_shim"); + + println!("cargo:rustc-link-search=native={}", lib_dir.display()); + println!("cargo:rustc-link-arg=-Wl,-rpath,{}", lib_dir.display()); + + for lib in &[ + "cuobjclient", + "cufile", + "ibverbs", + "rdmacm", + "numa", + "pthread", + "dl", + "rt", + ] { + println!("cargo:rustc-link-lib=dylib={lib}"); + } + println!("cargo:rustc-link-lib=dylib=stdc++"); +} + +fn resolve_cuobj_dirs(arch: &str) -> Option<(PathBuf, PathBuf)> { + let has_header = |inc: &PathBuf| inc.join("cuobjclient.h").exists(); + + if let Ok(home) = env::var("MSC_CUOBJ_HOME") { + let base = PathBuf::from(home); + let inc = base.join("include"); + let lib = base.join("lib"); + if has_header(&inc) { + return Some((inc, lib)); + } + } + + let triple = match arch { + "x86_64" => "x86_64-linux", + "aarch64" => "sbsa-linux", + _ => "x86_64-linux", + }; + for var in ["CUDA_HOME", "CUDA_PATH", "CUDA_ROOT"] { + if let Ok(cuda) = env::var(var) { + let base = PathBuf::from(cuda).join("targets").join(triple); + let inc = base.join("include"); + let lib = base.join("lib"); + if has_header(&inc) { + return Some((inc, lib)); + } + } + } + + // Local dev fallback: a vendored SDK copy (git-ignored, never committed). + let vendor = PathBuf::from(env::var("CARGO_MANIFEST_DIR").ok()?) + .join("vendor") + .join("cuobj"); + let inc = vendor.join("include"); + let lib = vendor.join("lib").join(arch); + if has_header(&inc) { + return Some((inc, lib)); + } + + None +} diff --git a/multi-storage-client/rust/csrc/cuobj_shim.cc b/multi-storage-client/rust/csrc/cuobj_shim.cc new file mode 100644 index 00000000..7c7570db --- /dev/null +++ b/multi-storage-client/rust/csrc/cuobj_shim.cc @@ -0,0 +1,141 @@ +// MinIO Rust Library for Amazon S3 Compatible Cloud Storage +// Copyright 2026 MinIO, Inc. +// +// 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 "cuobj_shim.h" + +#include + +// cuObject client API from the installed NVIDIA cuObject SDK (CUDA Toolkit +// >= 13.1). It transitively pulls in cufile.h + cuobjtelem.h from the same +// include dir, so only the SDK include path is needed to build the shim. +#include + +namespace { + +cuObjOpType_t map_op(int op) { + return op == MINIORS_CUOBJ_OP_PUT ? CUOBJ_PUT : CUOBJ_GET; +} + +int map_memory_type(cuObjMemoryType_t t) { + switch (t) { + case CUOBJ_MEMORY_SYSTEM: + return MINIORS_CUOBJ_MEM_SYSTEM; + case CUOBJ_MEMORY_CUDA_MANAGED: + return MINIORS_CUOBJ_MEM_CUDA_MANAGED; + case CUOBJ_MEMORY_CUDA_DEVICE: + return MINIORS_CUOBJ_MEM_CUDA_DEVICE; + default: + return MINIORS_CUOBJ_MEM_UNKNOWN; + } +} + +struct ClientHolder { + CUObjIOOps ops{}; + cuObjClient client; + ClientHolder() : client(ops, CUOBJ_PROTO_RDMA_DC_V1) {} +}; + +} // namespace + +extern "C" { + +miniors_cuobj_client* miniors_cuobj_client_new(void) { + try { + return reinterpret_cast(new ClientHolder()); + } catch (...) { + return nullptr; + } +} + +void miniors_cuobj_client_free(miniors_cuobj_client* client) { + delete reinterpret_cast(client); +} + +int miniors_cuobj_is_connected(miniors_cuobj_client* client) { + if (client == nullptr) return 0; + auto* holder = reinterpret_cast(client); + try { + return holder->client.isConnected() ? 1 : 0; + } catch (...) { + return 0; + } +} + +int miniors_cuobj_get_descriptor(miniors_cuobj_client* client, void* ptr, + size_t size) { + if (client == nullptr) return MINIORS_CUOBJ_FAIL; + auto* holder = reinterpret_cast(client); + try { + return holder->client.cuMemObjGetDescriptor(ptr, size) == CU_OBJ_SUCCESS + ? MINIORS_CUOBJ_SUCCESS + : MINIORS_CUOBJ_FAIL; + } catch (...) { + return MINIORS_CUOBJ_FAIL; + } +} + +int miniors_cuobj_put_descriptor(miniors_cuobj_client* client, void* ptr) { + if (client == nullptr) return MINIORS_CUOBJ_FAIL; + auto* holder = reinterpret_cast(client); + try { + return holder->client.cuMemObjPutDescriptor(ptr) == CU_OBJ_SUCCESS + ? MINIORS_CUOBJ_SUCCESS + : MINIORS_CUOBJ_FAIL; + } catch (...) { + return MINIORS_CUOBJ_FAIL; + } +} + +int miniors_cuobj_get_rdma_token(miniors_cuobj_client* client, void* ptr, + size_t size, size_t offset, int op, + char** token_out) { + if (client == nullptr || token_out == nullptr) return MINIORS_CUOBJ_FAIL; + *token_out = nullptr; + auto* holder = reinterpret_cast(client); + try { + char* token = nullptr; + cuObjErr_t err = holder->client.cuMemObjGetRDMAToken(ptr, size, offset, + map_op(op), &token); + if (err != CU_OBJ_SUCCESS || token == nullptr) { + return MINIORS_CUOBJ_FAIL; + } + *token_out = token; + return MINIORS_CUOBJ_SUCCESS; + } catch (...) { + return MINIORS_CUOBJ_FAIL; + } +} + +int miniors_cuobj_put_rdma_token(miniors_cuobj_client* client, char* token) { + if (client == nullptr || token == nullptr) return MINIORS_CUOBJ_FAIL; + auto* holder = reinterpret_cast(client); + try { + return holder->client.cuMemObjPutRDMAToken(token) == CU_OBJ_SUCCESS + ? MINIORS_CUOBJ_SUCCESS + : MINIORS_CUOBJ_FAIL; + } catch (...) { + return MINIORS_CUOBJ_FAIL; + } +} + +int miniors_cuobj_memory_type(const void* ptr) { + try { + return map_memory_type(cuObjClient::getMemoryType(ptr)); + } catch (...) { + return MINIORS_CUOBJ_MEM_UNKNOWN; + } +} + +} // extern "C" diff --git a/multi-storage-client/rust/csrc/cuobj_shim.h b/multi-storage-client/rust/csrc/cuobj_shim.h new file mode 100644 index 00000000..af4e4f21 --- /dev/null +++ b/multi-storage-client/rust/csrc/cuobj_shim.h @@ -0,0 +1,65 @@ +// MinIO Rust Library for Amazon S3 Compatible Cloud Storage +// Copyright 2026 MinIO, Inc. +// +// 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. +// +// C ABI wrapping NVIDIA libcuobjclient's `cuObjClient` C++ class so the Rust +// side can bind to it without a C++ compiler. Mirrors the surface used by +// minio-cpp's RDMA path: descriptor lifecycle, token mint/release, +// connectivity probe, memory-type detection. + +#ifndef MINIORS_CUOBJ_SHIM_H +#define MINIORS_CUOBJ_SHIM_H + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct miniors_cuobj_client miniors_cuobj_client; + +#define MINIORS_CUOBJ_SUCCESS 0 +#define MINIORS_CUOBJ_FAIL 1 + +#define MINIORS_CUOBJ_OP_GET 0 +#define MINIORS_CUOBJ_OP_PUT 1 + +#define MINIORS_CUOBJ_MEM_SYSTEM 0 +#define MINIORS_CUOBJ_MEM_CUDA_MANAGED 1 +#define MINIORS_CUOBJ_MEM_CUDA_DEVICE 2 +#define MINIORS_CUOBJ_MEM_UNKNOWN 3 + +miniors_cuobj_client* miniors_cuobj_client_new(void); +void miniors_cuobj_client_free(miniors_cuobj_client* client); + +int miniors_cuobj_is_connected(miniors_cuobj_client* client); + +int miniors_cuobj_get_descriptor(miniors_cuobj_client* client, void* ptr, + size_t size); +int miniors_cuobj_put_descriptor(miniors_cuobj_client* client, void* ptr); + +int miniors_cuobj_get_rdma_token(miniors_cuobj_client* client, void* ptr, + size_t size, size_t offset, int op, + char** token_out); +int miniors_cuobj_put_rdma_token(miniors_cuobj_client* client, char* token); + +int miniors_cuobj_memory_type(const void* ptr); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/multi-storage-client/rust/src/cuobj.rs b/multi-storage-client/rust/src/cuobj.rs new file mode 100644 index 00000000..1cb25ca9 --- /dev/null +++ b/multi-storage-client/rust/src/cuobj.rs @@ -0,0 +1,244 @@ +// SPDX-FileCopyrightText: Copyright (c) 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. + +//! PyO3 bindings for the NVIDIA cuObject (S3-over-RDMA) token API. +//! +//! Compiled only under the `rdma` cargo feature. The C++ `cuObjClient` class is +//! reached through the `miniors_cuobj_*` C ABI shim (`csrc/cuobj_shim.cc`), +//! which this module declares by hand -- no bindgen. The safe wrappers mirror +//! the minio-rs S3 RDMA client; the token registry (descriptor string -> owning +//! pointer) lets Python mint a token in one call and release it in another, +//! matching the manual-token pattern the Python `_cuobj.py` control plane drives. + +#![allow(non_camel_case_types)] + +use std::collections::HashMap; +use std::ffi::{c_void, CStr}; +use std::ptr::NonNull; +use std::sync::{LazyLock, Mutex, OnceLock}; + +use libc::{c_char, c_int, size_t}; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::*; +use pyo3::types::PyModule; + +#[repr(C)] +struct miniors_cuobj_client { + _private: [u8; 0], +} + +const MINIORS_CUOBJ_SUCCESS: c_int = 0; +const MINIORS_CUOBJ_OP_GET: c_int = 0; +const MINIORS_CUOBJ_OP_PUT: c_int = 1; + +extern "C" { + fn miniors_cuobj_client_new() -> *mut miniors_cuobj_client; + fn miniors_cuobj_client_free(client: *mut miniors_cuobj_client); + fn miniors_cuobj_is_connected(client: *mut miniors_cuobj_client) -> c_int; + fn miniors_cuobj_get_descriptor( + client: *mut miniors_cuobj_client, + ptr: *mut c_void, + size: size_t, + ) -> c_int; + fn miniors_cuobj_put_descriptor(client: *mut miniors_cuobj_client, ptr: *mut c_void) -> c_int; + fn miniors_cuobj_get_rdma_token( + client: *mut miniors_cuobj_client, + ptr: *mut c_void, + size: size_t, + offset: size_t, + op: c_int, + token_out: *mut *mut c_char, + ) -> c_int; + fn miniors_cuobj_put_rdma_token(client: *mut miniors_cuobj_client, token: *mut c_char) + -> c_int; +} + +/// Safe handle to the NVIDIA `cuObjClient` C++ instance owned by the shim. +struct CuObjClient { + raw: NonNull, +} + +// The underlying cuObjClient is driven through a process-wide singleton whose +// individual calls are serialized by cuObject's own internal locking. +unsafe impl Send for CuObjClient {} +unsafe impl Sync for CuObjClient {} + +impl CuObjClient { + fn new() -> Option { + let raw = unsafe { miniors_cuobj_client_new() }; + NonNull::new(raw).map(|raw| Self { raw }) + } + + fn is_connected(&self) -> bool { + unsafe { miniors_cuobj_is_connected(self.raw.as_ptr()) != 0 } + } + + unsafe fn get_descriptor(&self, ptr: *mut c_void, size: usize) -> bool { + miniors_cuobj_get_descriptor(self.raw.as_ptr(), ptr, size) == MINIORS_CUOBJ_SUCCESS + } + + unsafe fn put_descriptor(&self, ptr: *mut c_void) -> bool { + miniors_cuobj_put_descriptor(self.raw.as_ptr(), ptr) == MINIORS_CUOBJ_SUCCESS + } + + unsafe fn get_rdma_token( + &self, + ptr: *mut c_void, + size: usize, + offset: usize, + op: c_int, + ) -> Option<*mut c_char> { + let mut token: *mut c_char = std::ptr::null_mut(); + let rc = miniors_cuobj_get_rdma_token(self.raw.as_ptr(), ptr, size, offset, op, &mut token); + if rc != MINIORS_CUOBJ_SUCCESS || token.is_null() { + None + } else { + Some(token) + } + } + + unsafe fn put_rdma_token(&self, token: *mut c_char) -> bool { + miniors_cuobj_put_rdma_token(self.raw.as_ptr(), token) == MINIORS_CUOBJ_SUCCESS + } +} + +impl Drop for CuObjClient { + fn drop(&mut self) { + unsafe { miniors_cuobj_client_free(self.raw.as_ptr()) }; + } +} + +/// Process-wide shared `cuObjClient`. Constructing per-call is racy and corrupts +/// malloc state under concurrent workers, so one instance is shared per process. +fn shared() -> Option<&'static CuObjClient> { + static INSTANCE: OnceLock> = OnceLock::new(); + INSTANCE.get_or_init(CuObjClient::new).as_ref() +} + +fn client() -> PyResult<&'static CuObjClient> { + shared().ok_or_else(|| { + PyRuntimeError::new_err( + "cuObject client unavailable: cuObjClient construction failed (missing RDMA NIC, \ + cuFile/cuObject config, or version-matched libcuobjclient/libcufile)", + ) + }) +} + +/// Descriptors minted by `cuMemObjGetRDMAToken` are owned by cuObject and must be +/// released by the original pointer. Keyed by descriptor string so Python can +/// mint a token in one call and release it (by value) in a later one; the value +/// is the owning `*mut c_char` stored as `usize` to keep the map `Send`/`Sync`. +static TOKEN_REGISTRY: LazyLock>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + +/// Return whether cuObject is usable: the client constructed and connected to an +/// RDMA fabric. +#[pyfunction] +fn cuobj_available() -> bool { + shared().map(CuObjClient::is_connected).unwrap_or(false) +} + +/// Register a contiguous buffer with cuObject for RDMA. Raises on failure. +#[pyfunction] +fn cuobj_register_buffer(addr: usize, size: usize) -> PyResult<()> { + let client = client()?; + if unsafe { client.get_descriptor(addr as *mut c_void, size) } { + Ok(()) + } else { + Err(PyRuntimeError::new_err(format!( + "cuMemObjGetDescriptor failed for buffer at 0x{addr:x} ({size} bytes)" + ))) + } +} + +/// Deregister a buffer previously passed to `cuobj_register_buffer`. Raises on +/// failure. +#[pyfunction] +fn cuobj_deregister_buffer(addr: usize) -> PyResult<()> { + let client = client()?; + if unsafe { client.put_descriptor(addr as *mut c_void) } { + Ok(()) + } else { + Err(PyRuntimeError::new_err(format!( + "cuMemObjPutDescriptor failed for buffer at 0x{addr:x}" + ))) + } +} + +/// Mint an RDMA descriptor for a region of a registered buffer. `is_put` selects +/// PUT (server reads) vs GET (server writes). Release it with +/// `cuobj_put_rdma_token`. Raises on failure. +#[pyfunction] +fn cuobj_get_rdma_token(addr: usize, size: usize, offset: usize, is_put: bool) -> PyResult { + let client = client()?; + let op = if is_put { + MINIORS_CUOBJ_OP_PUT + } else { + MINIORS_CUOBJ_OP_GET + }; + let ptr = unsafe { client.get_rdma_token(addr as *mut c_void, size, offset, op) }.ok_or_else( + || { + PyRuntimeError::new_err(format!( + "cuMemObjGetRDMAToken failed for buffer at 0x{addr:x} ({size} bytes)" + )) + }, + )?; + match unsafe { CStr::from_ptr(ptr) }.to_str() { + Ok(s) => { + let desc = s.to_owned(); + TOKEN_REGISTRY + .lock() + .unwrap() + .insert(desc.clone(), ptr as usize); + Ok(desc) + } + Err(e) => { + // Release the descriptor we cannot represent, so a non-ASCII token + // never leaks the pinned registration. + unsafe { client.put_rdma_token(ptr) }; + Err(PyRuntimeError::new_err(format!( + "cuObject returned a non-UTF-8 RDMA descriptor: {e}" + ))) + } + } +} + +/// Release an RDMA descriptor returned by `cuobj_get_rdma_token`. Unknown or +/// already-released tokens are a no-op. +#[pyfunction] +fn cuobj_put_rdma_token(token: String) -> PyResult<()> { + let client = client()?; + // Claim the descriptor by removing it under the lock: `remove` returns the + // pointer to exactly one caller, so concurrent releases of the same token + // can't double-free the pinned registration. + let ptr = match TOKEN_REGISTRY.lock().unwrap().remove(&token) { + Some(ptr) => ptr, + None => return Ok(()), + }; + if !unsafe { client.put_rdma_token(ptr as *mut c_char) } { + return Err(PyRuntimeError::new_err("cuMemObjPutRDMAToken failed")); + } + Ok(()) +} + +/// Register the cuObject token-API functions on the extension module. +pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(cuobj_available, m)?)?; + m.add_function(wrap_pyfunction!(cuobj_register_buffer, m)?)?; + m.add_function(wrap_pyfunction!(cuobj_deregister_buffer, m)?)?; + m.add_function(wrap_pyfunction!(cuobj_get_rdma_token, m)?)?; + m.add_function(wrap_pyfunction!(cuobj_put_rdma_token, m)?)?; + Ok(()) +} diff --git a/multi-storage-client/rust/src/lib.rs b/multi-storage-client/rust/src/lib.rs index f33745ec..a6f65e94 100644 --- a/multi-storage-client/rust/src/lib.rs +++ b/multi-storage-client/rust/src/lib.rs @@ -44,6 +44,8 @@ use aws_smithy_http_client::{tls, Builder}; use aws_config::BehaviorVersion; mod credentials; +#[cfg(feature = "rdma")] +mod cuobj; mod types; use credentials::{AwsCredentialsProvider, AwsSdkCredentialsProvider, GcpCredentialsProvider}; @@ -1083,6 +1085,8 @@ fn multistorageclient_rust(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> m.add_class::()?; m.add("RustRetryableError", _py.get_type::())?; m.add("RustClientError", _py.get_type::())?; + #[cfg(feature = "rdma")] + cuobj::register(m)?; Ok(()) } diff --git a/multi-storage-client/src/multistorageclient/config.py b/multi-storage-client/src/multistorageclient/config.py index fa8dc308..a0476b86 100644 --- a/multi-storage-client/src/multistorageclient/config.py +++ b/multi-storage-client/src/multistorageclient/config.py @@ -91,6 +91,7 @@ def create_implicit_profile_config(profile_name: str, protocol: str, base_path: "ais_s3": "AIStoreS3StorageProvider", "s8k": "S8KStorageProvider", "gcs_s3": "GoogleS3StorageProvider", + "s3_cuobject": "S3CuObjectStorageProvider", "huggingface": "HuggingFaceStorageProvider", } diff --git a/multi-storage-client/src/multistorageclient/providers/__init__.py b/multi-storage-client/src/multistorageclient/providers/__init__.py index ee07c39e..9ee5ea9c 100644 --- a/multi-storage-client/src/multistorageclient/providers/__init__.py +++ b/multi-storage-client/src/multistorageclient/providers/__init__.py @@ -47,6 +47,7 @@ def __getattr__(name: str) -> Any: # S3 "S3StorageProvider": ".s3", "StaticS3CredentialsProvider": ".s3", + "S3CuObjectStorageProvider": ".s3_cuobject", # S8K "S8KStorageProvider": ".s8k", # AIS @@ -75,6 +76,7 @@ def __getattr__(name: str) -> Any: ".gcs_s3": ["boto3"], ".oci": ["oci"], ".s3": ["boto3"], + ".s3_cuobject": ["boto3"], ".s8k": ["boto3"], ".ais": ["aistore"], ".ais_s3": ["boto3", "aistore"], @@ -90,6 +92,7 @@ def __getattr__(name: str) -> Any: ".gcs_s3": "Google Cloud Storage with S3 API", ".oci": "Oracle Cloud Infrastructure", ".s3": "Amazon S3 or other S3-compatible storage", + ".s3_cuobject": "S3-over-RDMA (NVIDIA cuObject)", ".s8k": "S8K storage", ".ais": "NVIDIA AIStore", ".ais_s3": "NVIDIA AIStore with S3 API", diff --git a/multi-storage-client/src/multistorageclient/providers/_cuobj.py b/multi-storage-client/src/multistorageclient/providers/_cuobj.py new file mode 100644 index 00000000..a3140066 --- /dev/null +++ b/multi-storage-client/src/multistorageclient/providers/_cuobj.py @@ -0,0 +1,245 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +""" +S3-over-RDMA data plane for the S3 storage provider, backed by NVIDIA cuObject. + +cuObject (``libcuobjclient``) registers a contiguous host (or, in a future +revision, device) buffer for RDMA and mints an RDMA descriptor (token). The +descriptor is carried to an RDMA-capable S3 endpoint as the signed +``x-amz-rdma-token`` header; the endpoint then transfers the object payload +directly into or out of the registered buffer over RDMA, leaving the HTTP body +empty. This offloads the bulk transfer from the CPU and the HTTP/TLS path. + +This mirrors the PyTorch cuObject checkpoint backend (``torch.cuda.cuobj`` plus +``torch.distributed.checkpoint._cuobj_rdma_storage``): a thin set of token-API +primitives (:func:`is_available`, :func:`register_buffer`, +:func:`deregister_buffer`, :func:`get_rdma_token`, :func:`put_rdma_token`) and a +:class:`CuObjEngine` control plane that formats the descriptor and carries it on +the boto3 request -- the MSC equivalent of ``BotoCuObjClient``. + +cuObject is a C++ library whose client object requires an I/O-ops callback table +at construction, so it cannot be driven directly from ctypes. The token API is +instead reached through the ``multistorageclient_rust`` extension: an +``extern "C"`` shim over ``cuObjClient`` (``rust/csrc/cuobj_shim.cc``) bound by a +PyO3 module (``rust/src/cuobj.rs``) and compiled into the wheel by Maturin. The +shim is behind the crate's optional ``rdma`` feature, so it is present only when +the wheel is built with cuObject support:: + + maturin develop --features rdma # into an existing venv, or + maturin build --features rdma # a distributable wheel + +built on a host with the cuObject runtime (``libcufile``/``libcuobjclient``). + +This module is import-safe when the extension was built without the ``rdma`` +feature (the ``cuobj_*`` symbols are absent): the five primitives raise +:class:`CuObjError` and :func:`is_available` returns ``False``, exactly like the +``_dummy_fn`` fallback in ``torch/cuda/cuobj.py``. The S3 provider only +instantiates :class:`CuObjEngine` when the ``rdma`` option is configured. +""" + +import ctypes +import threading +from contextlib import contextmanager +from typing import Iterator, Optional, Union + +# The cuObject token API lives in the multistorageclient_rust extension behind +# the crate's `rdma` feature. Import the compiled module defensively: a default +# (non-rdma) wheel omits the cuobj_* functions entirely, and a source checkout +# may not have the extension built at all. +try: + from multistorageclient_rust import multistorageclient_rust as _rust_ext +except Exception: # pragma: no cover - extension unbuilt + _rust_ext = None + +_HAS_CUOBJ = _rust_ext is not None and hasattr(_rust_ext, "cuobj_available") + +# The boto3 ``before-sign`` hook runs per request, possibly on transfer-manager +# worker threads, so the token in flight for the current request is kept in +# thread-local state rather than on the client (mirrors the thread-local token +# in BotoCuObjClient). +_thread_state = threading.local() + + +class CuObjError(RuntimeError): + """Raised when a cuObject token-API call fails.""" + + +def _require_cuobj() -> None: + if not _HAS_CUOBJ: + raise CuObjError( + "cuObject support is not available: the multistorageclient_rust extension was built " + "without the 'rdma' feature (or is not built). Rebuild the wheel with cuObject support " + "on a host with the cuObject runtime, e.g. `maturin develop --features rdma`." + ) + + +def is_available() -> bool: + """Return whether NVIDIA cuObject (S3-over-RDMA) support is usable. + + ``True`` only when the extension was built with the ``rdma`` feature and a + cuObject client connection can be established (RDMA-capable NIC and a + reachable RDMA S3 endpoint). + """ + if not _HAS_CUOBJ: + return False + try: + return bool(_rust_ext.cuobj_available()) + except Exception: + return False + + +def register_buffer(addr: int, size: int) -> None: + """Register a contiguous buffer with cuObject for RDMA transfers. + + Registration is required before requesting an RDMA token for the buffer. + """ + _require_cuobj() + try: + _rust_ext.cuobj_register_buffer(addr, size) + except Exception as error: + raise CuObjError(f"cuMemObjGetDescriptor failed for buffer at 0x{addr:x} ({size} bytes)") from error + + +def deregister_buffer(addr: int) -> None: + """Deregister a buffer previously passed to :func:`register_buffer`.""" + _require_cuobj() + try: + _rust_ext.cuobj_deregister_buffer(addr) + except Exception as error: + raise CuObjError(f"cuMemObjPutDescriptor failed for buffer at 0x{addr:x}") from error + + +def get_rdma_token(addr: int, size: int, offset: int = 0, is_put: bool = True) -> str: + """Return an RDMA descriptor for a region of a registered buffer. + + ``is_put`` is ``True`` for a PUT (the server reads from the buffer), ``False`` + for a GET (the server writes into it). Release the descriptor with + :func:`put_rdma_token` once the request finishes. + """ + _require_cuobj() + try: + return _rust_ext.cuobj_get_rdma_token(addr, size, offset, is_put) + except Exception as error: + raise CuObjError(f"cuMemObjGetRDMAToken failed for buffer at 0x{addr:x} ({size} bytes)") from error + + +def put_rdma_token(token: str) -> None: + """Release an RDMA descriptor returned by :func:`get_rdma_token`.""" + _require_cuobj() + try: + _rust_ext.cuobj_put_rdma_token(token) + except Exception as error: + raise CuObjError("cuMemObjPutRDMAToken failed") from error + + +def _buffer_address(buffer: Union[bytearray, memoryview], nbytes: int) -> int: + """Return the address of a writable, contiguous buffer for RDMA registration. + + A writable buffer is required: GET delivers payload into it over RDMA, and a + PUT source is copied into one by the caller so cuObject can pin a stable, + non-immutable region. ``nbytes`` is the byte length (``memoryview.nbytes``), + which differs from ``len()`` for multi-byte item formats. + """ + array = (ctypes.c_char * nbytes).from_buffer(buffer) + return ctypes.addressof(array) + + +class CuObjEngine: + """cuObject RDMA control plane for a single S3 provider instance. + + The MSC analog of ``BotoCuObjClient``: it owns the per-request token + lifecycle and the boto3 hooks that carry the descriptor on the wire. The S3 + provider supplies the buffer and issues the (body-less) ``PutObject`` / + ``GetObject`` inside :meth:`transfer`. + """ + + def __init__(self) -> None: + if not is_available(): + raise CuObjError( + "cuObject client is not connected to an RDMA fabric. Check the RDMA NIC, the " + "cuFile/cuObject JSON config (CUFILE_ENV_PATH_JSON), and version-matched " + "libcufile/libcuobjclient libraries." + ) + + @staticmethod + def client_config_overrides() -> dict: + """botocore ``Config`` keys the S3 client must use for the RDMA wire contract. + + The payload travels over RDMA, so the HTTP body is empty and must not be + signed or checksummed; otherwise SigV4 / content checksums computed over + the empty body are rejected by the endpoint. + """ + return { + "request_checksum_calculation": "when_required", + "response_checksum_validation": "when_required", + "s3": {"payload_signing_enabled": False}, + } + + def install_hooks(self, s3_client) -> None: + events = s3_client.meta.events + events.register("before-sign.s3.PutObject", self._inject_token) + events.register("before-sign.s3.GetObject", self._inject_token) + events.register("before-sign.s3.UploadPart", self._inject_token) + + @staticmethod + def _inject_token(request, **kwargs) -> None: + token = getattr(_thread_state, "rdma_token", None) + if token is not None: + # SigV4 signs every x-amz-* header, so the token must be present + # before signing (before-sign), not after. + request.headers["x-amz-rdma-token"] = token + + @staticmethod + def check_reply(response) -> None: + """Fail loudly when an endpoint did not honor the RDMA request. + + With ``rdma`` explicitly enabled there is no silent TCP fallback: a + missing or ``501`` reply means the payload did not move over RDMA. + """ + headers = response["ResponseMetadata"]["HTTPHeaders"] + reply = headers.get("x-amz-rdma-reply") + if not reply or reply == "501": + raise CuObjError( + f"S3 endpoint declined RDMA (x-amz-rdma-reply={reply!r}); the endpoint is not " + "RDMA-capable. Disable the 'rdma' option to use the standard TCP data plane." + ) + + @contextmanager + def transfer(self, buffer: Union[bytearray, memoryview], is_put: bool) -> Iterator[None]: + """Register ``buffer``, publish its RDMA token for the wrapped request, then clean up. + + The descriptor is formatted ``::`` + so the endpoint can locate the exact registered region. + """ + nbytes = memoryview(buffer).nbytes + addr = _buffer_address(buffer, nbytes) + register_buffer(addr, nbytes) + token: Optional[str] = None + try: + token = get_rdma_token(addr, nbytes, 0, is_put) + _thread_state.rdma_token = f"{token}:{addr:016x}:{nbytes:016x}" + try: + yield + finally: + _thread_state.rdma_token = None + finally: + # Deregister the buffer even if releasing the token raises, so a + # failed release never leaks the pinned region. + try: + if token is not None: + put_rdma_token(token) + finally: + deregister_buffer(addr) diff --git a/multi-storage-client/src/multistorageclient/providers/s3_cuobject.py b/multi-storage-client/src/multistorageclient/providers/s3_cuobject.py new file mode 100644 index 00000000..c299df72 --- /dev/null +++ b/multi-storage-client/src/multistorageclient/providers/s3_cuobject.py @@ -0,0 +1,303 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +"""S3-over-RDMA (NVIDIA cuObject) storage provider. + +Subclasses :py:class:`S3StorageProvider` and overrides only the object +GET/PUT data plane to move payloads over RDMA via cuObject, leaving the base +``s3`` provider (and all its metadata/list/credentials/error handling) +untouched. Kept separate so the RDMA data path -- which cannot be exercised +in CI without an RDMA NIC and an RDMA-capable endpoint -- can never affect the +base ``s3`` provider. + +See https://docs.nvidia.com/gpudirect-storage/cuobject. +""" + +import base64 +import io +import os +import struct +from typing import IO, Any, Optional, Union + +from ..types import Range +from ..utils import split_path, validate_attributes +from ._cuobj import CuObjEngine +from .s3 import EXPRESS_ONEZONE_STORAGE_CLASS, MiB, S3StorageProvider + +PROVIDER = "s3_cuobject" + +# cuObject transfers the whole registered buffer in a single shot, so the boto +# multipart threshold is raised past any practical object size and every +# transfer takes the single-shot path. Uploads larger than +# ``rdma.multipart_chunksize`` are split by the RDMA multipart path below. +RDMA_SINGLE_SHOT_THRESHOLD = 1 << 62 + +# Default RDMA multipart part size. Uploads larger than this are sent as an +# RDMA multipart upload (one registered buffer + token + CRC64NVME per part). +RDMA_MULTIPART_CHUNKSIZE = 512 * MiB + + +class S3CuObjectStorageProvider(S3StorageProvider): + """ + A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with + S3 via NVIDIA cuObject (S3-over-RDMA). + + https://docs.nvidia.com/gpudirect-storage/cuobject + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + if "rust_client" in kwargs: + raise ValueError("The 's3_cuobject' provider is mutually exclusive with 'rust_client'.") + if kwargs.get("checksum_algorithm") is not None: + raise ValueError("checksum_algorithm is not supported for the 's3_cuobject' provider.") + + self._rdma_options: dict[str, Any] = kwargs.get("rdma") or {} + + # Force the empty-body, unsigned-payload wire contract the RDMA endpoint + # expects onto the boto client before it is constructed. Only the + # empty-body contract is enforced; addressing style stays user-controlled + # via the `s3` option. + overrides = CuObjEngine.client_config_overrides() + kwargs["request_checksum_calculation"] = overrides["request_checksum_calculation"] + kwargs["response_checksum_validation"] = overrides["response_checksum_validation"] + kwargs["s3"] = {**(kwargs.get("s3") or {}), **overrides["s3"]} + + super().__init__(*args, **kwargs) + + # Override the provider name from "s3". + self._provider_name = PROVIDER + + self._checksum_algorithm = None + self._multipart_threshold = RDMA_SINGLE_SHOT_THRESHOLD + self._rdma_multipart_chunksize = int(self._rdma_options.get("multipart_chunksize", RDMA_MULTIPART_CHUNKSIZE)) + if self._rdma_multipart_chunksize < 1: + raise ValueError( + f"rdma.multipart_chunksize must be a positive integer, got {self._rdma_multipart_chunksize}" + ) + + self._rdma_engine = CuObjEngine() + self._rdma_engine.install_hooks(self._s3_client) + + def _get_object(self, path: str, byte_range: Optional[Range] = None) -> bytes: + bucket, key = split_path(path) + return self._rdma_get(path, bucket, key, byte_range) + + def _put_object( + self, + path: str, + body: bytes, + if_match: Optional[str] = None, + if_none_match: Optional[str] = None, + attributes: Optional[dict[str, str]] = None, + content_type: Optional[str] = None, + ) -> int: + bucket, key = split_path(path) + + def _invoke_api() -> int: + kwargs: dict[str, Any] = {"Bucket": bucket, "Key": key, "Body": body} + if content_type: + kwargs["ContentType"] = content_type + if self._is_directory_bucket(bucket): + kwargs["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS + if if_match: + kwargs["IfMatch"] = if_match + if if_none_match: + kwargs["IfNoneMatch"] = if_none_match + validated_attributes = validate_attributes(attributes) + if validated_attributes: + kwargs["Metadata"] = validated_attributes + return self._rdma_put(kwargs, body) + + return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key) + + def _upload_file( + self, + remote_path: str, + f: Union[str, IO], + attributes: Optional[dict[str, str]] = None, + content_type: Optional[str] = None, + ) -> int: + return self._rdma_upload(remote_path, f, attributes, content_type) + + @staticmethod + def _rdma_checksum(buffer) -> str: + """Base64 CRC64NVME of ``buffer`` for the ``x-amz-checksum-crc64nvme`` header. + + CRC64NVME is hardware-accelerated and computed via ``awscrt`` (the same + implementation botocore uses for this algorithm). + """ + try: + from awscrt import checksums + except ImportError as error: + raise RuntimeError( + "RDMA PUT computes a CRC64NVME checksum and requires the 'awscrt' package (pip install awscrt)." + ) from error + return base64.b64encode(struct.pack(">Q", checksums.crc64nvme(buffer))).decode("ascii") + + def _rdma_put(self, kwargs: dict[str, Any], body: bytes) -> int: + """Single-shot RDMA PUT: cuObject transfers the registered buffer; the HTTP body is empty. + + A CRC64NVME checksum of the payload is computed on the client and sent as + ``x-amz-checksum-crc64nvme``. On an RDMA-capable endpoint it is validated + against the bytes delivered over RDMA; on an endpoint that ignores the + RDMA token it is validated against the (empty) HTTP body and fails with a + checksum mismatch, so the upload is rejected rather than silently storing + a 0-byte object. + """ + engine = self._rdma_engine + assert engine is not None + # cuObject pins a writable region: reuse the caller's buffer when it is + # writable, and copy only a read-only (immutable) body such as bytes. + buffer = bytearray(body) if memoryview(body).readonly else body + if len(buffer) == 0: + self._s3_client.put_object(**kwargs) + return 0 + kwargs = {**kwargs, "Body": b"", "ChecksumCRC64NVME": self._rdma_checksum(buffer)} + with engine.transfer(buffer, is_put=True): + response = self._s3_client.put_object(**kwargs) + engine.check_reply(response) + return len(buffer) + + def _rdma_get(self, path: str, bucket: str, key: str, byte_range: Optional[Range]) -> bytearray: + """Single-shot RDMA GET into a registered buffer; returns the buffer.""" + engine = self._rdma_engine + assert engine is not None + if_match: Optional[str] = None + if byte_range is not None: + size = byte_range.size + bytes_range: Optional[str] = f"bytes={byte_range.offset}-{byte_range.offset + byte_range.size - 1}" + else: + metadata = self._get_object_metadata(path) + size = metadata.content_length + bytes_range = None + # Bind the GET to the object version the buffer was sized against; if + # the object is replaced between the HEAD and the GET the endpoint + # returns 412 instead of delivering bytes into a mismatched buffer. + if_match = metadata.etag + + def _invoke_api() -> bytearray: + buffer = bytearray(size) + if size == 0: + return buffer + get_kwargs: dict[str, Any] = {"Bucket": bucket, "Key": key} + if bytes_range is not None: + get_kwargs["Range"] = bytes_range + if if_match: + get_kwargs["IfMatch"] = if_match + with engine.transfer(buffer, is_put=False): + response = self._s3_client.get_object(**get_kwargs) + response["Body"].read() + engine.check_reply(response) + return buffer + + return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key) + + def _rdma_create_extra( + self, bucket: str, attributes: Optional[dict[str, str]], content_type: Optional[str] + ) -> dict[str, Any]: + extra: dict[str, Any] = {} + if content_type: + extra["ContentType"] = content_type + if self._is_directory_bucket(bucket): + extra["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS + validated = validate_attributes(attributes) + if validated: + extra["Metadata"] = validated + return extra + + def _rdma_upload( + self, + remote_path: str, + f: Union[str, IO], + attributes: Optional[dict[str, str]], + content_type: Optional[str], + ) -> int: + """RDMA upload entry point: single-shot below the part size, multipart above it.""" + bucket, key = split_path(remote_path) + if isinstance(f, str): + size = os.path.getsize(f) + with open(f, "rb") as fp: + if size > self._rdma_multipart_chunksize: + extra = self._rdma_create_extra(bucket, attributes, content_type) + return self._rdma_upload_multipart(bucket, key, fp, size, extra) + return self._put_object(remote_path, fp.read(), attributes=attributes, content_type=content_type) + + f.seek(0, io.SEEK_END) + size = f.tell() + f.seek(0) + # Multipart reads raw chunks into a bytearray, so only binary streams + # can take it; any text-mode stream (StringIO, TextIOWrapper, open in + # "r") reads str and falls through to the single-shot encode path. + if size > self._rdma_multipart_chunksize and not isinstance(f, io.TextIOBase): + extra = self._rdma_create_extra(bucket, attributes, content_type) + return self._rdma_upload_multipart(bucket, key, f, size, extra) + data = f.read() + if isinstance(data, str): + data = data.encode("utf-8") + return self._put_object(remote_path, data, attributes=attributes, content_type=content_type) + + def _rdma_upload_multipart(self, bucket: str, key: str, fp: IO, size: int, extra: dict[str, Any]) -> int: + """RDMA multipart upload: each part is transferred as its own registered buffer. + + cuObject transfers one part-sized buffer per ``UploadPart`` (empty HTTP + body + RDMA token), each carrying a CRC64NVME the endpoint validates + against the RDMA-delivered bytes -- so the 0-byte-save guard applies per + part. Multipart is required past the single-PutObject size limit and + bounds the size of any one registered buffer / RDMA transfer. + """ + engine = self._rdma_engine + assert engine is not None + part_size = self._rdma_multipart_chunksize + + def _invoke_api() -> int: + upload_id = self._s3_client.create_multipart_upload(Bucket=bucket, Key=key, **extra)["UploadId"] + parts: list[dict[str, Any]] = [] + try: + part_number = 1 + remaining = size + while remaining > 0: + n = min(part_size, remaining) + chunk = bytearray() + while len(chunk) < n: + data = fp.read(n - len(chunk)) + if not data: + raise RuntimeError(f"unexpected end of input for {bucket}/{key} at part {part_number}") + chunk.extend(data) + checksum = self._rdma_checksum(chunk) + with engine.transfer(chunk, is_put=True): + response = self._s3_client.upload_part( + Bucket=bucket, + Key=key, + UploadId=upload_id, + PartNumber=part_number, + Body=b"", + ChecksumCRC64NVME=checksum, + ) + engine.check_reply(response) + parts.append({"PartNumber": part_number, "ETag": response["ETag"]}) + remaining -= n + part_number += 1 + self._s3_client.complete_multipart_upload( + Bucket=bucket, Key=key, UploadId=upload_id, MultipartUpload={"Parts": parts} + ) + except BaseException: + try: + self._s3_client.abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id) + except Exception: + pass + raise + return size + + return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key) diff --git a/multi-storage-client/src/multistorageclient/schema.py b/multi-storage-client/src/multistorageclient/schema.py index 6fa8bb90..ad87a95c 100644 --- a/multi-storage-client/src/multistorageclient/schema.py +++ b/multi-storage-client/src/multistorageclient/schema.py @@ -108,13 +108,31 @@ "properties": { "type": { "type": "string", - "enum": ["ais", "ais_s3", "azure", "file", "gcs", "gcs_s3", "oci", "s3", "s8k", "huggingface"], + "enum": [ + "ais", + "ais_s3", + "azure", + "file", + "gcs", + "gcs_s3", + "oci", + "s3", + "s3_cuobject", + "s8k", + "huggingface", + ], }, "options": { "type": "object", "properties": { "base_path": {"type": "string", "minLength": 0}, "rust_client": {"type": "object"}, + "rdma": { + "type": "object", + "properties": { + "multipart_chunksize": {"type": "integer", "minimum": 1}, + }, + }, }, "required": ["base_path"], }, diff --git a/multi-storage-client/tests/test_multistorageclient/unit/providers/test_s3_rdma.py b/multi-storage-client/tests/test_multistorageclient/unit/providers/test_s3_rdma.py new file mode 100644 index 00000000..03ecc1d7 --- /dev/null +++ b/multi-storage-client/tests/test_multistorageclient/unit/providers/test_s3_rdma.py @@ -0,0 +1,317 @@ +# SPDX-FileCopyrightText: Copyright (c) 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. + +""" +Unit tests for the S3-over-RDMA (cuObject) data plane wiring. + +The native cuObject engine is mocked, so these run anywhere -- they verify the +provider plumbing (option parsing, wire-contract config, single-shot routing, +empty-body PUT / sized GET), not the RDMA transfer itself. The transfer is +covered end-to-end against a live RDMA endpoint by ``examples/rdma_roundtrip.py``. +""" + +import base64 +import io +import struct +from array import array +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest + +from multistorageclient.providers._cuobj import CuObjEngine as _RealCuObjEngine +from multistorageclient.providers.s3 import StaticS3CredentialsProvider +from multistorageclient.providers.s3_cuobject import ( + RDMA_SINGLE_SHOT_THRESHOLD, + S3CuObjectStorageProvider, +) +from multistorageclient.types import Range + +_FAKE_CHECKSUM = "ZmFrZWNyYzY0" + + +def _make_rdma_provider(engine_cls: MagicMock, **extra: Any) -> S3CuObjectStorageProvider: + """Construct an RDMA-enabled provider with the cuObject engine mocked out.""" + engine_cls.client_config_overrides.return_value = _RealCuObjEngine.client_config_overrides() + return S3CuObjectStorageProvider( + region_name="us-east-1", + endpoint_url="https://s3.example.com", + base_path="test-bucket", + credentials_provider=StaticS3CredentialsProvider(access_key="test", secret_key="test"), + rdma={}, + **extra, + ) + + +def test_rdma_and_rust_client_are_mutually_exclusive(): + with pytest.raises(ValueError, match="mutually exclusive"): + S3CuObjectStorageProvider( + region_name="us-east-1", + endpoint_url="https://s3.example.com", + base_path="test-bucket", + credentials_provider=StaticS3CredentialsProvider(access_key="a", secret_key="b"), + rdma={}, + rust_client={}, + ) + + +def test_client_config_overrides_enforce_empty_body_contract(): + overrides = _RealCuObjEngine.client_config_overrides() + assert overrides["request_checksum_calculation"] == "when_required" + assert overrides["response_checksum_validation"] == "when_required" + assert overrides["s3"]["payload_signing_enabled"] is False + + +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_enables_single_shot_and_installs_hooks(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + + assert provider._rdma_engine is engine_cls.return_value + assert provider._rust_client is None + assert provider._checksum_algorithm is None + assert provider._multipart_threshold == RDMA_SINGLE_SHOT_THRESHOLD + engine_cls.return_value.install_hooks.assert_called_once_with(provider._s3_client) + + +@patch.object(S3CuObjectStorageProvider, "_rdma_checksum", staticmethod(lambda buffer: _FAKE_CHECKSUM)) +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_put_sends_empty_body_checksum_and_registers_buffer(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + engine = engine_cls.return_value + + written = provider._put_object(path="test-bucket/key.bin", body=b"hello world") + + assert written == len("hello world") + assert engine.transfer.call_args.kwargs["is_put"] is True + _, put_kwargs = provider._s3_client.put_object.call_args + assert put_kwargs["Body"] == b"" + # Precomputed CRC64NVME sent so a non-RDMA endpoint rejects the empty body + # instead of storing a 0-byte object. + assert put_kwargs["ChecksumCRC64NVME"] == _FAKE_CHECKSUM + engine.check_reply.assert_called_once() + + +@patch.object(S3CuObjectStorageProvider, "_rdma_checksum", staticmethod(lambda buffer: _FAKE_CHECKSUM)) +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_put_reuses_writable_buffer_and_copies_readonly(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + engine = engine_cls.return_value + + # Writable buffers (bytearray, writable memoryview) are registered in place. + writable = bytearray(b"writable payload") + provider._put_object(path="test-bucket/k1", body=writable) + assert engine.transfer.call_args.args[0] is writable + + view = memoryview(bytearray(b"view payload")) + provider._put_object(path="test-bucket/k2", body=view) + assert engine.transfer.call_args.args[0] is view + + # Read-only bytes are copied into a writable bytearray (cannot be pinned). + provider._put_object(path="test-bucket/k3", body=b"immutable payload") + copied = engine.transfer.call_args.args[0] + assert isinstance(copied, bytearray) + assert bytes(copied) == b"immutable payload" + + +def test_rdma_checksum_matches_awscrt(): + checksums = pytest.importorskip("awscrt.checksums") + data = b"the quick brown fox" * 1000 + expected = base64.b64encode(struct.pack(">Q", checksums.crc64nvme(data))).decode("ascii") + assert S3CuObjectStorageProvider._rdma_checksum(data) == expected + + +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_put_empty_payload_skips_rdma(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + engine = engine_cls.return_value + + written = provider._put_object(path="test-bucket/empty", body=b"") + + assert written == 0 + engine.transfer.assert_not_called() + provider._s3_client.put_object.assert_called_once() + + +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_get_byte_range_sizes_buffer_and_passes_range(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + engine = engine_cls.return_value + + result = provider._get_object(path="test-bucket/key.bin", byte_range=Range(offset=10, size=32)) + + assert isinstance(result, bytearray) + assert len(result) == 32 + assert engine.transfer.call_args.kwargs["is_put"] is False + _, get_kwargs = provider._s3_client.get_object.call_args + assert get_kwargs["Range"] == "bytes=10-41" + + +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_get_full_object_heads_for_size(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + engine = engine_cls.return_value + + metadata = MagicMock() + metadata.content_length = 128 + with patch.object(provider, "_get_object_metadata", return_value=metadata) as head: + result = provider._get_object(path="test-bucket/key.bin") + + head.assert_called_once() + assert len(result) == 128 + assert engine.transfer.call_args.kwargs["is_put"] is False + _, get_kwargs = provider._s3_client.get_object.call_args + assert "Range" not in get_kwargs + + +@patch.object(S3CuObjectStorageProvider, "_rdma_checksum", staticmethod(lambda buffer: _FAKE_CHECKSUM)) +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_upload_small_uses_single_shot(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + provider._rdma_multipart_chunksize = 16 + + provider._upload_file(remote_path="test-bucket/small.bin", f=io.BytesIO(b"x" * 10)) + + provider._s3_client.create_multipart_upload.assert_not_called() + provider._s3_client.put_object.assert_called_once() + + +@patch.object(S3CuObjectStorageProvider, "_rdma_checksum", staticmethod(lambda buffer: _FAKE_CHECKSUM)) +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_upload_multipart_splits_and_completes(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + provider._rdma_multipart_chunksize = 16 + provider._s3_client.create_multipart_upload.return_value = {"UploadId": "uid"} + provider._s3_client.upload_part.side_effect = [{"ETag": f"etag{i}"} for i in range(1, 4)] + engine = engine_cls.return_value + + written = provider._upload_file(remote_path="test-bucket/big.bin", f=io.BytesIO(b"a" * 40)) + + assert written == 40 + # 40 bytes / 16 => parts of 16, 16, 8. + assert provider._s3_client.upload_part.call_count == 3 + assert engine.transfer.call_count == 3 + for call in provider._s3_client.upload_part.call_args_list: + assert call.kwargs["Body"] == b"" + assert call.kwargs["ChecksumCRC64NVME"] == _FAKE_CHECKSUM + part_numbers = [c.kwargs["PartNumber"] for c in provider._s3_client.upload_part.call_args_list] + assert part_numbers == [1, 2, 3] + _, complete_kwargs = provider._s3_client.complete_multipart_upload.call_args + assert complete_kwargs["MultipartUpload"]["Parts"] == [ + {"PartNumber": 1, "ETag": "etag1"}, + {"PartNumber": 2, "ETag": "etag2"}, + {"PartNumber": 3, "ETag": "etag3"}, + ] + + +@patch.object(S3CuObjectStorageProvider, "_rdma_checksum", staticmethod(lambda buffer: _FAKE_CHECKSUM)) +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_upload_multipart_aborts_on_failure(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + provider._rdma_multipart_chunksize = 16 + provider._s3_client.create_multipart_upload.return_value = {"UploadId": "uid"} + provider._s3_client.upload_part.side_effect = [{"ETag": "etag1"}, RuntimeError("part failed")] + + with pytest.raises(RuntimeError): + provider._upload_file(remote_path="test-bucket/big.bin", f=io.BytesIO(b"a" * 40)) + + provider._s3_client.abort_multipart_upload.assert_called_once() + provider._s3_client.complete_multipart_upload.assert_not_called() + + +def test_install_hooks_registers_token_for_put_get_and_upload_part(): + engine = object.__new__(_RealCuObjEngine) + s3_client = MagicMock() + + engine.install_hooks(s3_client) + + registered = {call.args[0] for call in s3_client.meta.events.register.call_args_list} + assert registered == { + "before-sign.s3.PutObject", + "before-sign.s3.GetObject", + "before-sign.s3.UploadPart", + } + + +def test_transfer_registers_full_nbytes_for_multibyte_memoryview(): + import multistorageclient.providers._cuobj as cuobj + + engine = object.__new__(_RealCuObjEngine) + buffer = memoryview(array("H", [0x1111, 0x2222, 0x3333, 0x4444])) # 4 items, 8 bytes + assert len(buffer) == 4 and buffer.nbytes == 8 + + with ( + patch.object(cuobj, "register_buffer") as register, + patch.object(cuobj, "get_rdma_token", return_value="tok") as get_token, + patch.object(cuobj, "put_rdma_token"), + patch.object(cuobj, "deregister_buffer"), + ): + with engine.transfer(buffer, is_put=False): + pass + + assert register.call_args.args[1] == 8 # nbytes, not len() == 4 + assert get_token.call_args.args[1] == 8 + + +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_get_full_object_binds_ifmatch_to_head_version(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + + metadata = MagicMock() + metadata.content_length = 64 + metadata.etag = '"abc123"' + with patch.object(provider, "_get_object_metadata", return_value=metadata): + provider._get_object(path="test-bucket/key.bin") + + _, get_kwargs = provider._s3_client.get_object.call_args + assert get_kwargs["IfMatch"] == '"abc123"' + + +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_multipart_chunksize_must_be_positive(engine_cls: MagicMock): + engine_cls.client_config_overrides.return_value = _RealCuObjEngine.client_config_overrides() + with pytest.raises(ValueError, match="multipart_chunksize"): + S3CuObjectStorageProvider( + region_name="us-east-1", + endpoint_url="https://s3.example.com", + base_path="test-bucket", + credentials_provider=StaticS3CredentialsProvider(access_key="a", secret_key="b"), + rdma={"multipart_chunksize": 0}, + ) + + +@patch.object(S3CuObjectStorageProvider, "_rdma_checksum", staticmethod(lambda buffer: _FAKE_CHECKSUM)) +@patch("multistorageclient.providers.s3_cuobject.CuObjEngine") +def test_rdma_upload_text_stream_uses_single_shot(engine_cls: MagicMock): + provider = _make_rdma_provider(engine_cls) + provider._s3_client = MagicMock() + provider._rdma_multipart_chunksize = 16 + + # A text-mode stream larger than the part size (and not a StringIO) must not + # take the multipart path -- its chunks are str and would crash the raw + # bytearray reader -- so it falls through to the single-shot encode path. + text_stream = io.TextIOWrapper(io.BytesIO(b"a" * 40)) + provider._upload_file(remote_path="test-bucket/text.bin", f=text_stream) + + provider._s3_client.create_multipart_upload.assert_not_called() + provider._s3_client.put_object.assert_called_once()