OBJ: move RDMA objects larger than one cuObject request in parts

src/plugins/obj/s3_accel/client.cpp

@@ -87,7 +87,7 @@ awsS3AccelClient::putObjectAsync(std::string_view key,
         nixl_obj_rdma::S3RdmaClientCtx ctx;
         ctx.bucket = bucket;
         ctx.object = k;
-        const ssize_t r = nixl_obj_rdma::rdmaPutWithRetry(
+        const ssize_t r = nixl_obj_rdma::rdmaPutObject(
             *rdma, *cp, ctx, reinterpret_cast<void *>(data_ptr), data_len);
         // Success is a complete transfer: the descriptor length is a promise, so a
         // short count is a failure, not a partially-written object.
@@ -130,7 +130,7 @@ awsS3AccelClient::getObjectAsync(std::string_view key,
         nixl_obj_rdma::S3RdmaClientCtx ctx;
         ctx.bucket = bucket;
         ctx.object = k;
-        const ssize_t r = nixl_obj_rdma::rdmaGetWithRetry(
+        const ssize_t r = nixl_obj_rdma::rdmaGetObject(
             *rdma, *cp, ctx, reinterpret_cast<void *>(data_ptr), data_len, offset);
         // A full read is required: the server clamps to min(requested, servable),
         // so a short count (r < data_len) leaves the buffer tail unfilled and must

src/utils/object/rdma/rdmaRetry.cpp

@@ -5,6 +5,10 @@
 
 #include "rdmaRetry.h"
 
+#include <algorithm>
+#include <string>
+#include <vector>
+
 #include "common/nixl_log.h"
 
 namespace nixl_obj_rdma {
@@ -68,4 +72,73 @@ rdmaGetWithRetry(SharedCuObjClient &rdma,
     return ret;
 }
 
+ssize_t
+rdmaPutObject(SharedCuObjClient &rdma,
+              S3RdmaControlPlane &cp,
+              S3RdmaClientCtx &ctx,
+              void *buf,
+              uint64_t size) {
+    if (size <= rdma_max_request_size) {
+        return rdmaPutWithRetry(rdma, cp, ctx, buf, size);
+    }
+    ctx.uploadId = cp.createMultipartUpload(ctx);
+    if (ctx.uploadId.empty()) {
+        return rdma_error;
+    }
+
+    struct abortUnlessCompleted {
+        S3RdmaControlPlane &cp;
+        const S3RdmaClientCtx &ctx;
+        bool completed = false;
+
+        ~abortUnlessCompleted() {
+            if (!completed) {
+                cp.abortMultipartUpload(ctx);
+            }
+        }
+    } upload{cp, ctx};
+
+    std::vector<std::string> etags;
+    for (uint64_t done = 0; done < size; done += rdma_max_request_size) {
+        const uint64_t len = std::min(rdma_max_request_size, size - done);
+        ctx.partNumber = static_cast<uint32_t>(etags.size() + 1);
+        const ssize_t r = rdmaPutWithRetry(rdma, cp, ctx, static_cast<char *>(buf) + done, len);
+        if (r != static_cast<ssize_t>(len)) {
+            NIXL_ERROR << "RDMA multipart upload part " << ctx.partNumber << " failed (" << r
+                       << " of " << len << ") for key=" << ctx.object;
+            return rdma_error;
+        }
+        etags.push_back(ctx.etag);
+    }
+    upload.completed = cp.completeMultipartUpload(ctx, etags);
+    return upload.completed ? static_cast<ssize_t>(size) : rdma_error;
+}
+
+ssize_t
+rdmaGetObject(SharedCuObjClient &rdma,
+              S3RdmaControlPlane &cp,
+              S3RdmaClientCtx &ctx,
+              void *buf,
+              uint64_t size,
+              uint64_t offset) {
+    if (size == 0) {
+        return rdmaGetWithRetry(rdma, cp, ctx, buf, size, offset);
+    }
+    uint64_t done = 0;
+    while (done < size) {
+        const uint64_t len = std::min(rdma_max_request_size, size - done);
+        const ssize_t r =
+            rdmaGetWithRetry(rdma, cp, ctx, static_cast<char *>(buf) + done, len, offset + done);
+        if (r < 0) {
+            return r;
+        }
+        done += static_cast<uint64_t>(r);
+        if (r != static_cast<ssize_t>(len)) {
+            break;
+        }
+        ctx.ifMatch = ctx.etag;
+    }
+    return static_cast<ssize_t>(done);
+}
+
 } // namespace nixl_obj_rdma

src/utils/object/rdma/rdmaRetry.h

@@ -64,6 +64,38 @@ rdmaGetWithRetry(SharedCuObjClient &rdma,
                  uint64_t size,
                  uint64_t offset);
 
+// cuObject mints a token for at most this many bytes, even from a larger
+// registration, so a single S3 request moves at most this much.
+inline constexpr uint64_t rdma_max_request_size = CUOBJ_MAX_MEMORY_REG_SIZE;
+
+/**
+ * @brief Write buf to ctx.bucket/ctx.object over RDMA. An object larger than
+ *        rdma_max_request_size is uploaded as a multipart upload, one RDMA part
+ *        per rdma_max_request_size bytes.
+ * @return size on success, or a negative rdma_* code.
+ */
+[[nodiscard]] ssize_t
+rdmaPutObject(SharedCuObjClient &rdma,
+              S3RdmaControlPlane &cp,
+              S3RdmaClientCtx &ctx,
+              void *buf,
+              uint64_t size);
+
+/**
+ * @brief Read size bytes of ctx.bucket/ctx.object from offset into buf over
+ *        RDMA, as ranged GETs of at most rdma_max_request_size bytes each.
+ *        Every range after the first must match the first range's ETag, so a
+ *        read never mixes two versions of the object.
+ * @return Bytes read, or a negative rdma_* code.
+ */
+[[nodiscard]] ssize_t
+rdmaGetObject(SharedCuObjClient &rdma,
+              S3RdmaControlPlane &cp,
+              S3RdmaClientCtx &ctx,
+              void *buf,
+              uint64_t size,
+              uint64_t offset);
+
 } // namespace nixl_obj_rdma
 
 #endif // NIXL_SRC_UTILS_OBJECT_RDMA_RDMARETRY_H

src/utils/object/rdma/s3_control_plane_http.cpp

@@ -25,6 +25,10 @@
 #include <aws/core/utils/DateTime.h>
 #include <aws/core/utils/HashingUtils.h>
 #include <aws/core/utils/StringUtils.h>
+#include <aws/s3/S3Client.h>
+#include <aws/s3/model/AbortMultipartUploadRequest.h>
+#include <aws/s3/model/CompleteMultipartUploadRequest.h>
+#include <aws/s3/model/CreateMultipartUploadRequest.h>
 
 #include "object/s3/utils.h"
 #include "object/s3/aws_sdk_init.h"
@@ -173,6 +177,7 @@ struct S3RdmaControlPlane::Impl {
     Aws::String access_key;
     Aws::String secret_key;
     Aws::String session_token;
+    std::unique_ptr<Aws::S3::S3Client> s3;
 
     // Build the token-carrying control-plane request, apply the op-specific
     // headers, SigV4-sign it, and send it. The common prologue (rdma token +
@@ -270,8 +275,14 @@ S3RdmaControlPlane::S3RdmaControlPlane(const nixl_b_params_t *custom_params) : i
         config.connectTimeoutMs = rdma_connect_timeout_secs * 1000;
         config.requestTimeoutMs = rdma_timeout_secs * 1000;
         impl_->http = Aws::Http::CreateHttpClient(config);
-
-        valid_ = impl_->http != nullptr && !impl_->access_key.empty() && !impl_->secret_key.empty();
+        impl_->s3 = std::make_unique<Aws::S3::S3Client>(
+            creds,
+            config,
+            Aws::Client::AWSAuthV4Signer::PayloadSigningPolicy::RequestDependent,
+            impl_->virtual_addressing);
+
+        valid_ = impl_->http != nullptr && impl_->s3 != nullptr && !impl_->access_key.empty() &&
+            !impl_->secret_key.empty();
     }
     catch (const std::exception &e) {
         NIXL_WARN << "S3 RDMA control plane init failed: " << e.what();
@@ -378,8 +389,11 @@ S3RdmaControlPlane::rdmaGet(S3RdmaClientCtx &ctx,
             "bytes=" + std::to_string(offset) + "-" + std::to_string(offset + size - 1);
 
         auto resp = impl_->sendRdmaRequest(
-            Aws::Http::HttpMethod::HTTP_GET, uri, token, [&range](Aws::Http::HttpRequest &req) {
+            Aws::Http::HttpMethod::HTTP_GET, uri, token, [&](Aws::Http::HttpRequest &req) {
                 req.SetHeaderValue("range", range.c_str());
+                if (!ctx.ifMatch.empty()) {
+                    req.SetHeaderValue("if-match", ("\"" + ctx.ifMatch + "\"").c_str());
+                }
             });
         if (!resp) {
             NIXL_ERROR << "rdmaGet: MakeRequest returned null for key=" << ctx.object;
@@ -411,6 +425,10 @@ S3RdmaControlPlane::rdmaGet(S3RdmaClientCtx &ctx,
 
         const std::string etag =
             resp->HasHeader("etag") ? stripQuotes(resp->GetHeader("etag").c_str()) : "";
+        if (etag.empty()) {
+            NIXL_ERROR << "rdmaGet: response has no ETag for key=" << ctx.object;
+            return rdma_error;
+        }
 
         // The server reports the actual transferred byte count in
         // x-amz-rdma-bytes-transferred: it is authoritative and already clamped
@@ -443,4 +461,59 @@ S3RdmaControlPlane::rdmaGet(S3RdmaClientCtx &ctx,
     }
 }
 
+std::string
+S3RdmaControlPlane::createMultipartUpload(const S3RdmaClientCtx &ctx) {
+    Aws::S3::Model::CreateMultipartUploadRequest req;
+    req.SetBucket(ctx.bucket.c_str());
+    req.SetKey(ctx.object.c_str());
+    auto outcome = impl_->s3->CreateMultipartUpload(req);
+    if (!outcome.IsSuccess()) {
+        NIXL_ERROR << "CreateMultipartUpload failed for key=" << ctx.object << ": "
+                   << outcome.GetError().GetMessage();
+        return "";
+    }
+    return outcome.GetResult().GetUploadId().c_str();
+}
+
+bool
+S3RdmaControlPlane::completeMultipartUpload(const S3RdmaClientCtx &ctx,
+                                            const std::vector<std::string> &etags) {
+    Aws::S3::Model::CompletedMultipartUpload upload;
+    for (size_t i = 0; i < etags.size(); ++i) {
+        upload.AddParts(Aws::S3::Model::CompletedPart()
+                            .WithPartNumber(static_cast<int>(i + 1))
+                            .WithETag(("\"" + etags[i] + "\"").c_str()));
+    }
+    Aws::S3::Model::CompleteMultipartUploadRequest req;
+    req.SetBucket(ctx.bucket.c_str());
+    req.SetKey(ctx.object.c_str());
+    req.SetUploadId(ctx.uploadId.c_str());
+    req.SetMultipartUpload(std::move(upload));
+    auto outcome = impl_->s3->CompleteMultipartUpload(req);
+    if (!outcome.IsSuccess()) {
+        NIXL_ERROR << "CompleteMultipartUpload failed for key=" << ctx.object << ": "
+                   << outcome.GetError().GetMessage();
+        return false;
+    }
+    return true;
+}
+
+void
+S3RdmaControlPlane::abortMultipartUpload(const S3RdmaClientCtx &ctx) noexcept {
+    try {
+        Aws::S3::Model::AbortMultipartUploadRequest req;
+        req.SetBucket(ctx.bucket.c_str());
+        req.SetKey(ctx.object.c_str());
+        req.SetUploadId(ctx.uploadId.c_str());
+        auto outcome = impl_->s3->AbortMultipartUpload(req);
+        if (!outcome.IsSuccess()) {
+            NIXL_WARN << "AbortMultipartUpload failed for key=" << ctx.object << ": "
+                      << outcome.GetError().GetMessage();
+        }
+    }
+    catch (const std::exception &e) {
+        NIXL_WARN << "AbortMultipartUpload failed for key=" << ctx.object << ": " << e.what();
+    }
+}
+
 } // namespace nixl_obj_rdma

src/utils/object/rdma/s3_control_plane_http.h

@@ -16,6 +16,7 @@
 #include <cstdint>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include <sys/types.h> // ssize_t
 
@@ -38,6 +39,7 @@ struct S3RdmaClientCtx {
     std::string uploadId; ///< Multipart upload id; empty for single-shot.
     uint32_t partNumber = 0; ///< Part number 1..10000 when uploadId is set.
     std::string etag; ///< ETag returned by the server; populated on success.
+    std::string ifMatch; ///< When set, a GET fails unless the object's ETag matches.
 };
 
 /**
@@ -96,6 +98,26 @@ class S3RdmaControlPlane {
     [[nodiscard]] ssize_t
     rdmaGet(S3RdmaClientCtx &ctx, const char *token, uint64_t size, uint64_t offset);
 
+    /**
+     * @brief Start a multipart upload for ctx.bucket/ctx.object.
+     * @return The upload id, or an empty string on failure.
+     */
+    [[nodiscard]] std::string
+    createMultipartUpload(const S3RdmaClientCtx &ctx);
+
+    /**
+     * @brief Complete ctx.uploadId from its parts' ETags, in part order.
+     * @return true if the server assembled the object.
+     */
+    [[nodiscard]] bool
+    completeMultipartUpload(const S3RdmaClientCtx &ctx, const std::vector<std::string> &etags);
+
+    /**
+     * @brief Abort ctx.uploadId and discard the parts uploaded so far.
+     */
+    void
+    abortMultipartUpload(const S3RdmaClientCtx &ctx) noexcept;
+
 private:
     struct Impl;
     std::unique_ptr<Impl> impl_;