1# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2# SPDX-License-Identifier: Apache-2.0
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""S3-over-RDMA (NVIDIA cuObject) storage provider.
17
18Subclasses :py:class:`S3StorageProvider` and overrides only the object
19GET/PUT data plane to move payloads over RDMA via cuObject, leaving the base
20``s3`` provider (and all its metadata/list/credentials/error handling)
21untouched. Kept separate so the RDMA data path -- which cannot be exercised
22in CI without an RDMA NIC and an RDMA-capable endpoint -- can never affect the
23base ``s3`` provider.
24
25See https://docs.nvidia.com/gpudirect-storage/cuobject.
26"""
27
28import base64
29import io
30import logging
31import os
32import struct
33from collections.abc import Callable, Iterator
34from typing import IO, Any, cast
35
36from botocore.exceptions import ClientError
37
38from ..types import Range
39from ..utils import split_path, validate_attributes
40from ._cuobj import (
41 RDMA_MAX_MEMORY_REG_SIZE,
42 CuObjEngine,
43 CuObjError,
44 parse_rdma_bytes_transferred,
45 raise_if_declined,
46)
47from .s3 import EXPRESS_ONEZONE_STORAGE_CLASS, MiB, S3StorageProvider
48
49logger = logging.getLogger(__name__)
50
51PROVIDER = "s3_cuobject"
52
53# cuObject transfers the whole registered buffer in a single shot, so the boto
54# multipart threshold is raised past any practical object size and every
55# transfer takes the single-shot path. Transfers larger than
56# ``rdma.multipart_chunksize`` are split by the RDMA paths below: uploads into
57# multipart parts, downloads into ranged GETs.
58RDMA_SINGLE_SHOT_THRESHOLD = 1 << 62
59
60# Default RDMA part size. Uploads larger than this are sent as an RDMA multipart
61# upload (one registered buffer + token + CRC64NVME per part); downloads larger
62# than this are fetched as ranged RDMA GETs of this size.
63RDMA_MULTIPART_CHUNKSIZE = 512 * MiB
64
65
66def _http_status(error: ClientError) -> int | None:
67 return error.response.get("ResponseMetadata", {}).get("HTTPStatusCode")
68
69
[docs]
70class S3CuObjectStorageProvider(S3StorageProvider):
71 """
72 A concrete implementation of the :py:class:`multistorageclient.types.StorageProvider` for interacting with
73 S3 via NVIDIA cuObject (S3-over-RDMA).
74
75 https://docs.nvidia.com/gpudirect-storage/cuobject
76 """
77
78 def __init__(self, *args: Any, **kwargs: Any) -> None:
79 if "rust_client" in kwargs:
80 raise ValueError("The 's3_cuobject' provider is mutually exclusive with 'rust_client'.")
81 if kwargs.get("checksum_algorithm") is not None:
82 raise ValueError("checksum_algorithm is not supported for the 's3_cuobject' provider.")
83
84 self._rdma_options: dict[str, Any] = kwargs.get("rdma") or {}
85
86 # Force the empty-body, unsigned-payload wire contract the RDMA endpoint
87 # expects onto the boto client before it is constructed. Only the
88 # empty-body contract is enforced; addressing style stays user-controlled
89 # via the `s3` option.
90 overrides = CuObjEngine.client_config_overrides()
91 kwargs["request_checksum_calculation"] = overrides["request_checksum_calculation"]
92 kwargs["response_checksum_validation"] = overrides["response_checksum_validation"]
93 kwargs["s3"] = {**(kwargs.get("s3") or {}), **overrides["s3"]}
94
95 super().__init__(*args, **kwargs)
96
97 # Override the provider name from "s3".
98 self._provider_name = PROVIDER
99
100 self._checksum_algorithm = None
101 self._multipart_threshold = RDMA_SINGLE_SHOT_THRESHOLD
102 self._rdma_multipart_chunksize = int(self._rdma_options.get("multipart_chunksize", RDMA_MULTIPART_CHUNKSIZE))
103 if not 1 <= self._rdma_multipart_chunksize <= RDMA_MAX_MEMORY_REG_SIZE:
104 raise ValueError(
105 f"rdma.multipart_chunksize must be between 1 and {RDMA_MAX_MEMORY_REG_SIZE}, "
106 f"got {self._rdma_multipart_chunksize}"
107 )
108
109 self._rdma_engine = CuObjEngine()
110 self._rdma_engine.install_hooks(self._s3_client)
111
112 def _get_object(self, path: str, byte_range: Range | None = None) -> bytes:
113 bucket, key = split_path(path)
114 return bytes(self._rdma_get(path, bucket, key, byte_range))
115
116 def _put_object(
117 self,
118 path: str,
119 body: bytes,
120 if_match: str | None = None,
121 if_none_match: str | None = None,
122 attributes: dict[str, str] | None = None,
123 content_type: str | None = None,
124 ) -> int:
125 bucket, key = split_path(path)
126 size = memoryview(body).nbytes
127 if size > self._rdma_multipart_chunksize:
128 extra = self._rdma_create_extra(bucket, attributes, content_type)
129 return self._rdma_upload_multipart(
130 bucket, key, self._buffer_parts(body), size, extra, if_match, if_none_match
131 )
132
133 def _invoke_api() -> int:
134 kwargs: dict[str, Any] = {"Bucket": bucket, "Key": key, "Body": body}
135 kwargs.update(self._rdma_create_extra(bucket, attributes, content_type))
136 if if_match:
137 kwargs["IfMatch"] = if_match
138 if if_none_match:
139 kwargs["IfNoneMatch"] = if_none_match
140 return self._rdma_put(kwargs, body)
141
142 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
143
144 def _upload_file(
145 self,
146 remote_path: str,
147 f: str | IO,
148 attributes: dict[str, str] | None = None,
149 content_type: str | None = None,
150 ) -> int:
151 return self._rdma_upload(remote_path, f, attributes, content_type)
152
153 @staticmethod
154 def _rdma_checksum(buffer: bytes | bytearray | memoryview) -> str:
155 """Base64 CRC64NVME of ``buffer`` for the ``x-amz-checksum-crc64nvme`` header.
156
157 CRC64NVME is hardware-accelerated and computed via ``awscrt`` (the same
158 implementation botocore uses for this algorithm).
159 """
160 try:
161 from awscrt import checksums
162 except ImportError as error:
163 raise RuntimeError(
164 "RDMA PUT computes a CRC64NVME checksum and requires the 'awscrt' package (pip install awscrt)."
165 ) from error
166 return base64.b64encode(struct.pack(">Q", checksums.crc64nvme(cast(bytes, buffer)))).decode("ascii")
167
168 def _rdma_call(self, operation: Callable[..., Any], is_put: bool, **kwargs: Any) -> Any:
169 """Issue an RDMA request and validate its ``x-amz-rdma-reply``.
170
171 A decline arrives as an HTTP error status, which botocore raises as a
172 ``ClientError``; it is surfaced as a :class:`CuObjError` instead.
173 """
174 engine = self._rdma_engine
175 assert engine is not None
176 try:
177 response = operation(**kwargs)
178 except ClientError as error:
179 raise_if_declined(error.response)
180 raise
181 engine.check_reply(response, is_put=is_put)
182 return response
183
184 def _rdma_put(self, kwargs: dict[str, Any], body: bytes | bytearray | memoryview) -> int:
185 """Single-shot RDMA PUT: cuObject transfers the registered buffer; the HTTP body is empty.
186
187 A CRC64NVME checksum of the payload is computed on the client and sent as
188 ``x-amz-checksum-crc64nvme``. On an RDMA-capable endpoint it is validated
189 against the bytes delivered over RDMA; on an endpoint that ignores the
190 RDMA token it is validated against the (empty) HTTP body and fails with a
191 checksum mismatch, so the upload is rejected rather than silently storing
192 a 0-byte object.
193 """
194 engine = self._rdma_engine
195 assert engine is not None
196 # cuObject pins a writable region: reuse the caller's buffer when it is
197 # writable, and copy only a read-only (immutable) body such as bytes.
198 buffer = bytearray(body) if memoryview(body).readonly else cast(bytearray | memoryview, body)
199 if len(buffer) == 0:
200 self._s3_client.put_object(**kwargs)
201 return 0
202 kwargs = {**kwargs, "Body": b"", "ChecksumCRC64NVME": self._rdma_checksum(buffer)}
203 with engine.transfer(buffer, is_put=True):
204 self._rdma_call(self._s3_client.put_object, True, **kwargs)
205 return memoryview(buffer).nbytes
206
207 def _rdma_get(self, path: str, bucket: str, key: str, byte_range: Range | None) -> bytearray:
208 """RDMA GET into a registered buffer, split into ranged GETs above the part size.
209
210 Returns the buffer truncated to ``x-amz-rdma-bytes-transferred``, which
211 can be shorter than requested for a range past the end of the object. A
212 full-object GET must deliver exactly the size the HEAD reported.
213 """
214 if_match: str | None = None
215 if byte_range is not None:
216 offset, size = byte_range.offset, byte_range.size
217 else:
218 metadata = self._get_object_metadata(path)
219 offset, size = 0, metadata.content_length
220 # Bind the GET to the object version the buffer was sized against; if
221 # the object is replaced between the HEAD and the GET the endpoint
222 # returns 412 instead of delivering bytes into a mismatched buffer.
223 if_match = metadata.etag
224 part_size = self._rdma_multipart_chunksize
225 ranged = byte_range is not None or size > part_size
226
227 def _invoke_api() -> bytearray:
228 buffer = bytearray(size)
229 part_if_match = if_match
230 pos = 0
231 while pos < size:
232 n = min(part_size, size - pos)
233 bytes_range = f"bytes={offset + pos}-{offset + pos + n - 1}" if ranged else None
234 try:
235 with memoryview(buffer)[pos : pos + n] as view:
236 transferred, etag = self._rdma_get_part(view, bucket, key, bytes_range, part_if_match)
237 except ClientError as error:
238 # A ranged read whose earlier parts ended exactly at the end
239 # of the object gets 416 for the next part: the read is done.
240 if byte_range is not None and pos > 0 and _http_status(error) == 416:
241 break
242 raise
243 # Pin the remaining parts to the version the first part read, so
244 # a split read never mixes bytes from two object versions.
245 part_if_match = part_if_match or etag
246 pos += transferred
247 if transferred < n:
248 break
249 if byte_range is None and pos != size:
250 raise CuObjError(f"RDMA GET of {path} delivered {pos} of {size} bytes")
251 del buffer[pos:]
252 return buffer
253
254 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
255
256 def _rdma_get_part(
257 self, view: memoryview, bucket: str, key: str, bytes_range: str | None, if_match: str | None
258 ) -> tuple[int, str | None]:
259 """RDMA GET of one part into ``view``; returns the bytes transferred and the object ETag."""
260 engine = self._rdma_engine
261 assert engine is not None
262 get_kwargs: dict[str, Any] = {"Bucket": bucket, "Key": key}
263 if bytes_range is not None:
264 get_kwargs["Range"] = bytes_range
265 if if_match:
266 get_kwargs["IfMatch"] = if_match
267 with engine.transfer(view, is_put=False):
268 response = self._rdma_call(self._s3_client.get_object, False, **get_kwargs)
269 response["Body"].read()
270 return parse_rdma_bytes_transferred(response, view.nbytes), response.get("ETag")
271
272 def _rdma_create_extra(
273 self, bucket: str, attributes: dict[str, str] | None, content_type: str | None
274 ) -> dict[str, Any]:
275 extra: dict[str, Any] = {}
276 if content_type:
277 extra["ContentType"] = content_type
278 if self._is_directory_bucket(bucket):
279 extra["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS
280 validated = validate_attributes(attributes)
281 if validated:
282 extra["Metadata"] = validated
283 return extra
284
285 def _rdma_upload(
286 self,
287 remote_path: str,
288 f: str | IO,
289 attributes: dict[str, str] | None,
290 content_type: str | None,
291 ) -> int:
292 """RDMA upload entry point: single-shot below the part size, multipart above it."""
293 bucket, key = split_path(remote_path)
294 if isinstance(f, str):
295 size = os.path.getsize(f)
296 with open(f, "rb") as fp:
297 if size > self._rdma_multipart_chunksize:
298 extra = self._rdma_create_extra(bucket, attributes, content_type)
299 return self._rdma_upload_multipart(bucket, key, self._stream_parts(fp, size), size, extra)
300 return self._put_object(remote_path, fp.read(), attributes=attributes, content_type=content_type)
301
302 f.seek(0, io.SEEK_END)
303 size = f.tell()
304 f.seek(0)
305 # Multipart reads raw chunks into a bytearray, so only binary streams
306 # can take it; any text-mode stream (StringIO, TextIOWrapper, open in
307 # "r") reads str and falls through to the single-shot encode path.
308 if size > self._rdma_multipart_chunksize and not isinstance(f, io.TextIOBase):
309 extra = self._rdma_create_extra(bucket, attributes, content_type)
310 return self._rdma_upload_multipart(bucket, key, self._stream_parts(f, size), size, extra)
311 data = f.read()
312 if isinstance(data, str):
313 data = data.encode("utf-8")
314 return self._put_object(remote_path, data, attributes=attributes, content_type=content_type)
315
316 def _stream_parts(self, fp: IO, size: int) -> Iterator[bytearray]:
317 """Read ``size`` bytes from ``fp`` as part-sized writable buffers."""
318 remaining = size
319 while remaining > 0:
320 n = min(self._rdma_multipart_chunksize, remaining)
321 chunk = bytearray()
322 while len(chunk) < n:
323 data = fp.read(n - len(chunk))
324 if not data:
325 raise RuntimeError(f"unexpected end of input after {size - remaining + len(chunk)} bytes")
326 chunk.extend(data)
327 yield chunk
328 remaining -= n
329
330 def _buffer_parts(self, body: bytes | bytearray | memoryview) -> Iterator[bytearray | memoryview]:
331 """Slice an in-memory body into part-sized writable buffers, copying only read-only parts."""
332 view = memoryview(body).cast("B")
333 for pos in range(0, view.nbytes, self._rdma_multipart_chunksize):
334 part = view[pos : pos + self._rdma_multipart_chunksize]
335 yield bytearray(part) if part.readonly else part
336
337 def _rdma_upload_multipart(
338 self,
339 bucket: str,
340 key: str,
341 chunks: Iterator[bytearray | memoryview],
342 size: int,
343 extra: dict[str, Any],
344 if_match: str | None = None,
345 if_none_match: str | None = None,
346 ) -> int:
347 """RDMA multipart upload: each part is transferred as its own registered buffer.
348
349 cuObject transfers one part-sized buffer per ``UploadPart`` (empty HTTP
350 body + RDMA token), each carrying a CRC64NVME the endpoint validates
351 against the RDMA-delivered bytes -- so the 0-byte-save guard applies per
352 part. Multipart is required past the single-PutObject size limit and
353 bounds the size of any one registered buffer / RDMA transfer.
354 """
355 engine = self._rdma_engine
356 assert engine is not None
357
358 def _invoke_api() -> int:
359 upload_id = self._s3_client.create_multipart_upload(Bucket=bucket, Key=key, **extra)["UploadId"]
360 parts: list[dict[str, Any]] = []
361 try:
362 for part_number, chunk in enumerate(chunks, start=1):
363 checksum = self._rdma_checksum(chunk)
364 with engine.transfer(chunk, is_put=True):
365 response = self._rdma_call(
366 self._s3_client.upload_part,
367 True,
368 Bucket=bucket,
369 Key=key,
370 UploadId=upload_id,
371 PartNumber=part_number,
372 Body=b"",
373 ChecksumCRC64NVME=checksum,
374 )
375 parts.append({"PartNumber": part_number, "ETag": response["ETag"]})
376 complete: dict[str, Any] = {}
377 if if_match:
378 complete["IfMatch"] = if_match
379 if if_none_match:
380 complete["IfNoneMatch"] = if_none_match
381 self._s3_client.complete_multipart_upload(
382 Bucket=bucket, Key=key, UploadId=upload_id, MultipartUpload={"Parts": parts}, **complete
383 )
384 except BaseException:
385 try:
386 self._s3_client.abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
387 except Exception as error:
388 logger.warning("Failed to abort RDMA multipart upload %s: %s", upload_id, error)
389 raise
390 return size
391
392 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)