[braket] support reservation ARN (#2553)

* support reservation arn * DCO Remediation Commit for Tim <[email protected]> I, Tim <[email protected]>, hereby add my Signed-off-by to this commit: 955d837 Signed-off-by: Tim <[email protected]> * fix: run clang formatter Signed-off-by: Tim <[email protected]> --------- Signed-off-by: Tim <[email protected]>
NVIDIA · Feb 8, 2025 · 20e5869 · 20e5869
1 parent 3602887
commit 20e5869
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/runtime/common/BraketExecutor.h b/runtime/common/BraketExecutor.h
@@ -55,6 +55,7 @@ class BraketExecutor : public Executor {
 
   std::shared_future<std::string> defaultBucketFuture;
   char const *jobToken;
+  char const *reservationArn;
 
   std::chrono::microseconds pollingInterval = std::chrono::milliseconds{100};
 

diff --git a/runtime/cudaq/platform/default/rest/helpers/braket/BraketExecutor.cpp b/runtime/cudaq/platform/default/rest/helpers/braket/BraketExecutor.cpp
@@ -9,6 +9,8 @@
 #include "common/BraketExecutor.h"
 #include "common/BraketServerHelper.h"
 
+#include <aws/braket/model/Association.h>
+#include <aws/braket/model/AssociationType.h>
 #include <aws/braket/model/CreateQuantumTaskRequest.h>
 #include <aws/braket/model/GetQuantumTaskRequest.h>
 #include <aws/braket/model/QuantumTaskStatus.h>
@@ -100,7 +102,8 @@ void tryCreateBucket(Aws::S3Crt::S3CrtClient &client, std::string const &region,
 
 namespace cudaq {
 BraketExecutor::BraketExecutor()
-    : api(options), jobToken(std::getenv("AMZN_BRAKET_JOB_TOKEN")) {}
+    : api(options), jobToken(std::getenv("AMZN_BRAKET_JOB_TOKEN")),
+      reservationArn(std::getenv("AMZN_BRAKET_RESERVATION_TIME_WINDOW_ARN")) {}
 
 /// @brief Set the server helper
 void BraketExecutor::setServerHelper(ServerHelper *helper) {
@@ -190,6 +193,15 @@ BraketExecutor::execute(std::vector<KernelExecution> &codesToExecute,
     req.SetShots(message["shots"]);
     if (jobToken)
       req.SetJobToken(jobToken);
+
+    if (reservationArn) {
+      Aws::Braket::Model::Association assoc;
+      assoc.SetArn(reservationArn);
+      assoc.SetType(
+          Aws::Braket::Model::AssociationType::RESERVATION_TIME_WINDOW_ARN);
+      req.AddAssociations(std::move(assoc));
+    }
+
     req.SetOutputS3Bucket(defaultBucket);
     req.SetOutputS3KeyPrefix(defaultPrefix);