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, SuffixRange
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 _supports_suffix_range(self) -> bool:
113 return False
114
115 def _get_object(self, path: str, byte_range: Range | SuffixRange | None = None) -> bytes:
116 bucket, key = split_path(path)
117 return bytes(self._rdma_get(path, bucket, key, cast(Range | None, byte_range)))
118
119 def _put_object(
120 self,
121 path: str,
122 body: bytes,
123 if_match: str | None = None,
124 if_none_match: str | None = None,
125 attributes: dict[str, str] | None = None,
126 content_type: str | None = None,
127 ) -> int:
128 bucket, key = split_path(path)
129 size = memoryview(body).nbytes
130 if size > self._rdma_multipart_chunksize:
131 extra = self._rdma_create_extra(bucket, attributes, content_type)
132 return self._rdma_upload_multipart(
133 bucket, key, self._buffer_parts(body), size, extra, if_match, if_none_match
134 )
135
136 def _invoke_api() -> int:
137 kwargs: dict[str, Any] = {"Bucket": bucket, "Key": key, "Body": body}
138 kwargs.update(self._rdma_create_extra(bucket, attributes, content_type))
139 if if_match:
140 kwargs["IfMatch"] = if_match
141 if if_none_match:
142 kwargs["IfNoneMatch"] = if_none_match
143 return self._rdma_put(kwargs, body)
144
145 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)
146
147 def _upload_file(
148 self,
149 remote_path: str,
150 f: str | IO,
151 attributes: dict[str, str] | None = None,
152 content_type: str | None = None,
153 ) -> int:
154 return self._rdma_upload(remote_path, f, attributes, content_type)
155
156 @staticmethod
157 def _rdma_checksum(buffer: bytes | bytearray | memoryview) -> str:
158 """Base64 CRC64NVME of ``buffer`` for the ``x-amz-checksum-crc64nvme`` header.
159
160 CRC64NVME is hardware-accelerated and computed via ``awscrt`` (the same
161 implementation botocore uses for this algorithm).
162 """
163 try:
164 from awscrt import checksums
165 except ImportError as error:
166 raise RuntimeError(
167 "RDMA PUT computes a CRC64NVME checksum and requires the 'awscrt' package (pip install awscrt)."
168 ) from error
169 return base64.b64encode(struct.pack(">Q", checksums.crc64nvme(cast(bytes, buffer)))).decode("ascii")
170
171 def _rdma_call(self, operation: Callable[..., Any], is_put: bool, **kwargs: Any) -> Any:
172 """Issue an RDMA request and validate its ``x-amz-rdma-reply``.
173
174 A decline arrives as an HTTP error status, which botocore raises as a
175 ``ClientError``; it is surfaced as a :class:`CuObjError` instead.
176 """
177 engine = self._rdma_engine
178 assert engine is not None
179 try:
180 response = operation(**kwargs)
181 except ClientError as error:
182 raise_if_declined(error.response)
183 raise
184 engine.check_reply(response, is_put=is_put)
185 return response
186
187 def _rdma_put(self, kwargs: dict[str, Any], body: bytes | bytearray | memoryview) -> int:
188 """Single-shot RDMA PUT: cuObject transfers the registered buffer; the HTTP body is empty.
189
190 A CRC64NVME checksum of the payload is computed on the client and sent as
191 ``x-amz-checksum-crc64nvme``. On an RDMA-capable endpoint it is validated
192 against the bytes delivered over RDMA; on an endpoint that ignores the
193 RDMA token it is validated against the (empty) HTTP body and fails with a
194 checksum mismatch, so the upload is rejected rather than silently storing
195 a 0-byte object.
196 """
197 engine = self._rdma_engine
198 assert engine is not None
199 # cuObject pins a writable region: reuse the caller's buffer when it is
200 # writable, and copy only a read-only (immutable) body such as bytes.
201 buffer = bytearray(body) if memoryview(body).readonly else cast(bytearray | memoryview, body)
202 if len(buffer) == 0:
203 self._s3_client.put_object(**kwargs)
204 return 0
205 kwargs = {**kwargs, "Body": b"", "ChecksumCRC64NVME": self._rdma_checksum(buffer)}
206 with engine.transfer(buffer, is_put=True):
207 self._rdma_call(self._s3_client.put_object, True, **kwargs)
208 return memoryview(buffer).nbytes
209
210 def _rdma_get(self, path: str, bucket: str, key: str, byte_range: Range | None) -> bytearray:
211 """RDMA GET into a registered buffer, split into ranged GETs above the part size.
212
213 Returns the buffer truncated to ``x-amz-rdma-bytes-transferred``, which
214 can be shorter than requested for a range past the end of the object. A
215 full-object GET must deliver exactly the size the HEAD reported.
216 """
217 if_match: str | None = None
218 if byte_range is not None:
219 offset, size = byte_range.offset, byte_range.size
220 else:
221 metadata = self._get_object_metadata(path)
222 offset, size = 0, metadata.content_length
223 # Bind the GET to the object version the buffer was sized against; if
224 # the object is replaced between the HEAD and the GET the endpoint
225 # returns 412 instead of delivering bytes into a mismatched buffer.
226 if_match = metadata.etag
227 part_size = self._rdma_multipart_chunksize
228 ranged = byte_range is not None or size > part_size
229
230 def _invoke_api() -> bytearray:
231 buffer = bytearray(size)
232 part_if_match = if_match
233 pos = 0
234 while pos < size:
235 n = min(part_size, size - pos)
236 bytes_range = f"bytes={offset + pos}-{offset + pos + n - 1}" if ranged else None
237 try:
238 with memoryview(buffer)[pos : pos + n] as view:
239 transferred, etag = self._rdma_get_part(view, bucket, key, bytes_range, part_if_match)
240 except ClientError as error:
241 # A ranged read whose earlier parts ended exactly at the end
242 # of the object gets 416 for the next part: the read is done.
243 if byte_range is not None and pos > 0 and _http_status(error) == 416:
244 break
245 raise
246 # Pin the remaining parts to the version the first part read, so
247 # a split read never mixes bytes from two object versions.
248 part_if_match = part_if_match or etag
249 pos += transferred
250 if transferred < n:
251 break
252 if byte_range is None and pos != size:
253 raise CuObjError(f"RDMA GET of {path} delivered {pos} of {size} bytes")
254 del buffer[pos:]
255 return buffer
256
257 return self._translate_errors(_invoke_api, operation="GET", bucket=bucket, key=key)
258
259 def _rdma_get_part(
260 self, view: memoryview, bucket: str, key: str, bytes_range: str | None, if_match: str | None
261 ) -> tuple[int, str | None]:
262 """RDMA GET of one part into ``view``; returns the bytes transferred and the object ETag."""
263 engine = self._rdma_engine
264 assert engine is not None
265 get_kwargs: dict[str, Any] = {"Bucket": bucket, "Key": key}
266 if bytes_range is not None:
267 get_kwargs["Range"] = bytes_range
268 if if_match:
269 get_kwargs["IfMatch"] = if_match
270 with engine.transfer(view, is_put=False):
271 response = self._rdma_call(self._s3_client.get_object, False, **get_kwargs)
272 response["Body"].read()
273 return parse_rdma_bytes_transferred(response, view.nbytes), response.get("ETag")
274
275 def _rdma_create_extra(
276 self, bucket: str, attributes: dict[str, str] | None, content_type: str | None
277 ) -> dict[str, Any]:
278 extra: dict[str, Any] = {}
279 if content_type:
280 extra["ContentType"] = content_type
281 if self._is_directory_bucket(bucket):
282 extra["StorageClass"] = EXPRESS_ONEZONE_STORAGE_CLASS
283 validated = validate_attributes(attributes)
284 if validated:
285 extra["Metadata"] = validated
286 return extra
287
288 def _rdma_upload(
289 self,
290 remote_path: str,
291 f: str | IO,
292 attributes: dict[str, str] | None,
293 content_type: str | None,
294 ) -> int:
295 """RDMA upload entry point: single-shot below the part size, multipart above it."""
296 bucket, key = split_path(remote_path)
297 if isinstance(f, str):
298 size = os.path.getsize(f)
299 with open(f, "rb") as fp:
300 if size > self._rdma_multipart_chunksize:
301 extra = self._rdma_create_extra(bucket, attributes, content_type)
302 return self._rdma_upload_multipart(bucket, key, self._stream_parts(fp, size), size, extra)
303 return self._put_object(remote_path, fp.read(), attributes=attributes, content_type=content_type)
304
305 f.seek(0, io.SEEK_END)
306 size = f.tell()
307 f.seek(0)
308 # Multipart reads raw chunks into a bytearray, so only binary streams
309 # can take it; any text-mode stream (StringIO, TextIOWrapper, open in
310 # "r") reads str and falls through to the single-shot encode path.
311 if size > self._rdma_multipart_chunksize and not isinstance(f, io.TextIOBase):
312 extra = self._rdma_create_extra(bucket, attributes, content_type)
313 return self._rdma_upload_multipart(bucket, key, self._stream_parts(f, size), size, extra)
314 data = f.read()
315 if isinstance(data, str):
316 data = data.encode("utf-8")
317 return self._put_object(remote_path, data, attributes=attributes, content_type=content_type)
318
319 def _stream_parts(self, fp: IO, size: int) -> Iterator[bytearray]:
320 """Read ``size`` bytes from ``fp`` as part-sized writable buffers."""
321 remaining = size
322 while remaining > 0:
323 n = min(self._rdma_multipart_chunksize, remaining)
324 chunk = bytearray()
325 while len(chunk) < n:
326 data = fp.read(n - len(chunk))
327 if not data:
328 raise RuntimeError(f"unexpected end of input after {size - remaining + len(chunk)} bytes")
329 chunk.extend(data)
330 yield chunk
331 remaining -= n
332
333 def _buffer_parts(self, body: bytes | bytearray | memoryview) -> Iterator[bytearray | memoryview]:
334 """Slice an in-memory body into part-sized writable buffers, copying only read-only parts."""
335 view = memoryview(body).cast("B")
336 for pos in range(0, view.nbytes, self._rdma_multipart_chunksize):
337 part = view[pos : pos + self._rdma_multipart_chunksize]
338 yield bytearray(part) if part.readonly else part
339
340 def _rdma_upload_multipart(
341 self,
342 bucket: str,
343 key: str,
344 chunks: Iterator[bytearray | memoryview],
345 size: int,
346 extra: dict[str, Any],
347 if_match: str | None = None,
348 if_none_match: str | None = None,
349 ) -> int:
350 """RDMA multipart upload: each part is transferred as its own registered buffer.
351
352 cuObject transfers one part-sized buffer per ``UploadPart`` (empty HTTP
353 body + RDMA token), each carrying a CRC64NVME the endpoint validates
354 against the RDMA-delivered bytes -- so the 0-byte-save guard applies per
355 part. Multipart is required past the single-PutObject size limit and
356 bounds the size of any one registered buffer / RDMA transfer.
357 """
358 engine = self._rdma_engine
359 assert engine is not None
360
361 def _invoke_api() -> int:
362 upload_id = self._s3_client.create_multipart_upload(Bucket=bucket, Key=key, **extra)["UploadId"]
363 parts: list[dict[str, Any]] = []
364 try:
365 for part_number, chunk in enumerate(chunks, start=1):
366 checksum = self._rdma_checksum(chunk)
367 with engine.transfer(chunk, is_put=True):
368 response = self._rdma_call(
369 self._s3_client.upload_part,
370 True,
371 Bucket=bucket,
372 Key=key,
373 UploadId=upload_id,
374 PartNumber=part_number,
375 Body=b"",
376 ChecksumCRC64NVME=checksum,
377 )
378 parts.append({"PartNumber": part_number, "ETag": response["ETag"]})
379 complete: dict[str, Any] = {}
380 if if_match:
381 complete["IfMatch"] = if_match
382 if if_none_match:
383 complete["IfNoneMatch"] = if_none_match
384 self._s3_client.complete_multipart_upload(
385 Bucket=bucket, Key=key, UploadId=upload_id, MultipartUpload={"Parts": parts}, **complete
386 )
387 except BaseException:
388 try:
389 self._s3_client.abort_multipart_upload(Bucket=bucket, Key=key, UploadId=upload_id)
390 except Exception as error:
391 logger.warning("Failed to abort RDMA multipart upload %s: %s", upload_id, error)
392 raise
393 return size
394
395 return self._translate_errors(_invoke_api, operation="PUT", bucket=bucket, key=key)