Skip to content

Commit

Permalink
Add Braket backend
Browse files Browse the repository at this point in the history
Signed-off-by: Ryan <3620100+rmshaffer@users.noreply.github.com>
  • Loading branch information
rmshaffer committed Oct 31, 2024
1 parent bc17b54 commit 79089cd
Show file tree
Hide file tree
Showing 47 changed files with 1,943 additions and 7 deletions.
3 changes: 3 additions & 0 deletions .gitmodules
Original file line number Diff line number Diff line change
Expand Up @@ -48,3 +48,6 @@
[submodule "tpls/Stim"]
path = tpls/Stim
url = https://github.com/quantumlib/Stim
[submodule "tpls/aws-sdk-cpp"]
path = tpls/aws-sdk-cpp
url = https://github.com/aws/aws-sdk-cpp.git
12 changes: 11 additions & 1 deletion CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -124,7 +124,8 @@ endif()
# ==============================================================================
option(CUDAQ_BUILD_TESTS "Build cudaq tests" ON)
option(CUDAQ_ENABLE_RPC_LOGGING "Enable verbose printout for client/server qpud connection." OFF)
option(CUDAQ_TEST_MOCK_SERVERS "Enable Remote QPU Tests via Mock Servers." OFF)
option(CUDAQ_TEST_MOCK_SERVERS "Enable Remote QPU Tests via Mock Servers." ON)
option(CUDAQ_INTEGRATION_TEST_BRAKET "Run integration tests on Braket backend." ON)
option(CUDAQ_DISABLE_RUNTIME "Build without the CUDA-Q runtime, just the compiler toolchain." OFF)
option(CUDAQ_SKIP_MPI "Do not build with MPI, even when it is available." OFF)
option(CUDAQ_DISABLE_CPP_FRONTEND "Build without the CUDA-Q C++ Clang-based Frontend." OFF)
Expand Down Expand Up @@ -478,6 +479,15 @@ if (OPENSSL_FOUND AND CUDAQ_ENABLE_REST)
add_subdirectory(tpls/Crow)
endif()

# AWS SDK for C++
if (CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64")
set(crypto_LIBRARY "$ENV{OPENSSL_INSTALL_PREFIX}/lib64/libcrypto.a" CACHE STRING "" FORCE)
else()
set(crypto_LIBRARY "$ENV{OPENSSL_INSTALL_PREFIX}/lib/libcrypto.a" CACHE STRING "" FORCE)
endif()
set(crypto_INCLUDE_DIR "$ENV{OPENSSL_INSTALL_PREFIX}/include" CACHE STRING "" FORCE)
set(SERVICE_COMPONENTS braket s3-crt sts)
find_package(AWSSDK REQUIRED COMPONENTS ${SERVICE_COMPONENTS})

# Check for CUDA Support
# ==============================================================================
Expand Down
9 changes: 6 additions & 3 deletions runtime/common/BaseRemoteRESTQPU.h
Original file line number Diff line number Diff line change
Expand Up @@ -144,8 +144,6 @@ class BaseRemoteRESTQPU : public cudaq::QPU {
BaseRemoteRESTQPU() : QPU() {
std::filesystem::path cudaqLibPath{cudaq::getCUDAQLibraryPath()};
platformPath = cudaqLibPath.parent_path().parent_path() / "targets";
// Default is to run sampling via the remote rest call
executor = std::make_unique<cudaq::Executor>();
}

BaseRemoteRESTQPU(BaseRemoteRESTQPU &&) = delete;
Expand Down Expand Up @@ -315,14 +313,19 @@ class BaseRemoteRESTQPU : public cudaq::QPU {

// Set the qpu name
qpuName = mutableBackend;

// Create the ServerHelper for this QPU and give it the backend config
serverHelper = cudaq::registry::get<cudaq::ServerHelper>(qpuName);
if (!serverHelper) {
throw std::runtime_error("ServerHelper not found for target");
}
serverHelper->initialize(backendConfig);
serverHelper->updatePassPipeline(platformPath, passPipelineConfig);
cudaq::info("Retrieving executor with name {}", qpuName);
cudaq::info("Is this executor registered? {}",
cudaq::registry::isRegistered<cudaq::Executor>(qpuName));
executor = cudaq::registry::isRegistered<cudaq::Executor>(qpuName)
? cudaq::registry::get<cudaq::Executor>(qpuName)
: std::make_unique<cudaq::Executor>();

// Give the server helper to the executor
executor->setServerHelper(serverHelper.get());
Expand Down
2 changes: 2 additions & 0 deletions runtime/common/Executor.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -56,3 +56,5 @@ Executor::execute(std::vector<KernelExecution> &codesToExecute) {
return details::future(ids, name, config);
}
} // namespace cudaq

LLVM_INSTANTIATE_REGISTRY(cudaq::Executor::RegistryType)
2 changes: 1 addition & 1 deletion runtime/common/Executor.h
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ class Executor : public registry::RegisteredType<Executor> {

/// @brief Execute the provided quantum codes and return a future object
/// The caller can make this synchronous by just immediately calling .get().
details::future execute(std::vector<KernelExecution> &codesToExecute);
virtual details::future execute(std::vector<KernelExecution> &codesToExecute);
};

} // namespace cudaq
6 changes: 5 additions & 1 deletion runtime/cudaq/platform/default/rest/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ target_include_directories(cudaq-rest-qpu PRIVATE .
$<BUILD_INTERFACE:${CMAKE_SOURCE_DIR}/runtime>
$<INSTALL_INTERFACE:include>)

set(SERVICE_COMPONENTS braket s3-crt sts)
find_package(AWSSDK REQUIRED COMPONENTS ${SERVICE_COMPONENTS})

target_link_libraries(cudaq-rest-qpu
PUBLIC
cudaq-spin
Expand All @@ -28,7 +31,8 @@ target_link_libraries(cudaq-rest-qpu
MLIRTranslateLib
fmt::fmt-header-only
cudaq
cudaq-platform-default)
cudaq-platform-default
${AWSSDK_LINK_LIBRARIES})

install(TARGETS cudaq-rest-qpu DESTINATION lib)

Expand Down
1 change: 1 addition & 0 deletions runtime/cudaq/platform/default/rest/helpers/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
# the terms of the Apache License 2.0 which accompanies this distribution. #
# ============================================================================ #
add_subdirectory(anyon)
add_subdirectory(braket)
add_subdirectory(oqc)
add_subdirectory(ionq)
add_subdirectory(quantinuum)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
/*******************************************************************************
* Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. *
* All rights reserved. *
* *
* This source code and the accompanying materials are made available under *
* the terms of the Apache License 2.0 which accompanies this distribution. *
******************************************************************************/

#include "BraketExecutor.h"
#include "BraketServerHelper.h"

namespace cudaq {

details::future
BraketExecutor::execute(std::vector<KernelExecution> &codesToExecute) {
auto braketServerHelper = dynamic_cast<BraketServerHelper *>(serverHelper);
assert(braketServerHelper);
braketServerHelper->setShots(shots);

auto [dummy1, dummy2, messages] =
braketServerHelper->createJob(codesToExecute);

cudaq::debug("messages[0] = {}", messages[0].dump());

auto const &message = messages[0];

CountsDictionary counts;

std::string const defaultBucket = defaultBucketFuture.get();
std::string const defaultPrefix = "tasks-cudaq";

Aws::Braket::Model::CreateQuantumTaskRequest req;

auto config = braketServerHelper->getConfig();
cudaq::info("Backend config: {}, shots {}", config, shots);
config.insert({"shots", std::to_string(shots)});

req.SetAction(message["action"]);
req.SetDeviceArn(message["deviceArn"]);
req.SetShots(message["shots"]);
req.SetOutputS3Bucket(defaultBucket);
req.SetOutputS3KeyPrefix(defaultPrefix);

auto response = braketClient.CreateQuantumTask(req);

if (response.IsSuccess()) {
std::string taskArn = response.GetResult().GetQuantumTaskArn();
cudaq::info("Created {}", taskArn);
Aws::Braket::Model::QuantumTaskStatus taskStatus;
Aws::Braket::Model::GetQuantumTaskRequest getTaskReq;
getTaskReq.SetQuantumTaskArn(taskArn);

auto r = braketClient.GetQuantumTask(getTaskReq);
do {
std::this_thread::sleep_for(std::chrono::milliseconds{100});
r = braketClient.GetQuantumTask(getTaskReq);
taskStatus = r.GetResult().GetStatus();
} while (taskStatus != Aws::Braket::Model::QuantumTaskStatus::COMPLETED &&
taskStatus != Aws::Braket::Model::QuantumTaskStatus::FAILED &&
taskStatus != Aws::Braket::Model::QuantumTaskStatus::CANCELLED);

std::string outBucket = r.GetResult().GetOutputS3Bucket();
std::string outPrefix = r.GetResult().GetOutputS3Directory();

cudaq::info("results at {}/{}", outBucket, outPrefix);
Aws::S3Crt::Model::GetObjectRequest resultRequest;
resultRequest.SetBucket(outBucket);
resultRequest.SetKey(fmt::format("{}/results.json", outPrefix));

auto results = s3Client.GetObject(resultRequest);

auto parsedResult =
nlohmann::json::parse(results.GetResultWithOwnership().GetBody());

auto measurements = parsedResult.at("measurements");

for (auto const &m : measurements) {
std::string bitString = "";
for (int bit : m) {
bitString += std::to_string(bit);
}
counts[bitString] += 1;
}

} else {
std::cout << "Create error\n" << response.GetError() << "\n";
}

ExecutionResult ex_r{counts};
sample_result result{ex_r};

std::promise<sample_result> p;
p.set_value(result);
return {p.get_future()};
};
} // namespace cudaq

CUDAQ_REGISTER_TYPE(cudaq::Executor, cudaq::BraketExecutor, braket);
106 changes: 106 additions & 0 deletions runtime/cudaq/platform/default/rest/helpers/braket/BraketExecutor.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
/*******************************************************************************
* Copyright (c) 2022 - 2024 NVIDIA Corporation & Affiliates. *
* All rights reserved. *
* *
* This source code and the accompanying materials are made available under *
* the terms of the Apache License 2.0 which accompanies this distribution. *
******************************************************************************/

#pragma once
#include "common/Executor.h"
#include "common/FmtCore.h"
#include "common/MeasureCounts.h"
#include "cudaq.h"

#include <chrono>
#include <iostream>

#include <aws/core/Aws.h>

#include <aws/braket/BraketClient.h>
#include <aws/braket/model/CreateQuantumTaskRequest.h>
#include <aws/braket/model/GetQuantumTaskRequest.h>
#include <aws/braket/model/QuantumTaskStatus.h>

#include <aws/sts/STSClient.h>

#include <aws/s3-crt/S3CrtClient.h>
#include <aws/s3-crt/model/GetObjectRequest.h>

#include <aws/core/utils/logging/AWSLogging.h>
#include <aws/core/utils/logging/ConsoleLogSystem.h>
#include <aws/core/utils/logging/LogLevel.h>

#include "BraketServerHelper.h"
#include "common/Logger.h"

#include <nlohmann/json.hpp>
#include <regex>
#include <string>
#include <thread>

namespace cudaq {
/// @brief The Executor subclass for Amazon Braket
class BraketExecutor : public Executor {
Aws::SDKOptions options;

std::string region;
std::string accountId;

class ScopedApi {
Aws::SDKOptions &options;

public:
ScopedApi(Aws::SDKOptions &options) : options(options) {
Aws::InitAPI(options);
}
~ScopedApi() { Aws::ShutdownAPI(options); }
};

ScopedApi api;
Aws::Braket::BraketClient braketClient;
Aws::STS::STSClient stsClient;
Aws::S3Crt::S3CrtClient s3Client;

std::future<std::string> defaultBucketFuture;

static auto getClientConfig() {
Aws::Client::ClientConfiguration clientConfig;
clientConfig.verifySSL = false;
return clientConfig;
}

static auto getS3ClientConfig() {
Aws::S3Crt::ClientConfiguration clientConfig;
clientConfig.verifySSL = false;
return clientConfig;
}

public:
BraketExecutor()
: api(options), braketClient(getClientConfig()),
stsClient(getClientConfig()), s3Client(getS3ClientConfig()) {
cudaq::debug("Creating BraketExecutor");

defaultBucketFuture = std::async(std::launch::async, [&] {
auto response = stsClient.GetCallerIdentity();
std::string bucketName;
if (response.IsSuccess()) {
bucketName =
fmt::format("amazon-braket-{}-{}", getClientConfig().region,
response.GetResult().GetAccount());
cudaq::info("Task results will use S3 bucket \"{}\"", bucketName);
return bucketName;
} else {
throw response.GetError();
}
});
}

~BraketExecutor() = default;

/// @brief Execute the provided Braket task
details::future
execute(std::vector<KernelExecution> &codesToExecute) override;
};
} // namespace cudaq
Loading

0 comments on commit 79089cd

Please sign in to comment.