diff --git a/runtime/common/BraketExecutor.h b/runtime/common/BraketExecutor.h index 258e5c9425..abd451307d 100644 --- a/runtime/common/BraketExecutor.h +++ b/runtime/common/BraketExecutor.h @@ -55,6 +55,7 @@ class BraketExecutor : public Executor { std::shared_future 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 index 7a46a7b0da..2ea8eea8ba 100644 --- 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 +#include #include #include #include @@ -100,7 +102,8 @@ void tryCreateBucket(Aws::S3Crt::S3CrtClient &client, std::string const ®ion, 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 &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);