🔌 AWS Braket Executor#
Covalent is a Pythonic workflow tool used to execute tasks on advanced computing hardware.
This plugin allows executing quantum circuits and quantum-classical hybrid jobs in Amazon Braket when you use Covalent.
1. Installation#
To use this plugin with Covalent, simply install it using pip
:
pip install covalent-awsbraket-plugin
Note
Users will also need to have Docker installed on their local machine to use this plugin.
2. Usage Example#
The following toy example executes a simple quantum circuit on one qubit that prepares a uniform superposition of the standard basis states and then measures the state. We use the Pennylane framework.
import covalent as ct
from covalent_braket_plugin.braket import BraketExecutor
import os
# AWS resources to pass to the executor
credentials = "~/.aws/credentials"
profile = "default"
region = "us-east-1"
s3_bucket_name = "braket_s3_bucket"
ecr_repo_name = "braket_ecr_repo"
iam_role_name = "covalent-braket-iam-role"
# Instantiate the executor
ex = BraketExecutor(
credentials=credentials,
profile=profile,
region=region,
s3_bucket_name=s3_bucket_name,
ecr_repo_name=ecr_repo_name,
braket_job_execution_role_name=iam_role_name,
quantum_device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
classical_device="ml.m5.large",
storage=30,
time_limit=300,
)
# Execute the following circuit:
# |0> - H - Measure
@ct.electron(executor=ex)
def simple_quantum_task(num_qubits: int):
import pennylane as qml
# These are passed to the Hybrid Jobs container at runtime
device_arn = os.environ["AMZN_BRAKET_DEVICE_ARN"]
s3_bucket = os.environ["AMZN_BRAKET_OUT_S3_BUCKET"]
s3_task_dir = os.environ["AMZN_BRAKET_TASK_RESULTS_S3_URI"].split(s3_bucket)[1]
device = qml.device(
"braket.aws.qubit",
device_arn=device_arn,
s3_destination_folder=(s3_bucket, s3_task_dir),
wires=num_qubits,
)
@qml.qnode(device=device)
def simple_circuit():
qml.Hadamard(wires=[0])
return qml.expval(qml.PauliZ(wires=[0]))
res = simple_circuit().numpy()
return res
@ct.lattice
def simple_quantum_workflow(num_qubits: int):
return simple_quantum_task(num_qubits=num_qubits)
dispatch_id = ct.dispatch(simple_quantum_workflow)(1)
result_object = ct.get_result(dispatch_id, wait=True)
# We expect 0 as the result
print("Result:", result_object.result)
During the execution of the workflow one can navigate to the UI to see the status of the workflow, once completed however the above script should also output a value with the output of the quantum measurement.
>>> Result: 0
3. Overview of Configuration#
Config Key |
Is Required |
Default |
Description |
---|---|---|---|
credentials |
No |
“~/.aws/credentials” |
The path to the AWS credentials file |
braket_job_execution_role_name |
Yes |
“CovalentBraketJobsExecutionRole” |
The name of the IAM role that Braket will assume during task execution. |
profile |
No |
“default” |
Named AWS profile used for authentication |
region |
Yes |
:code`AWS_DEFAULT_REGION` environment variable |
AWS Region to use to for client calls to AWS |
s3_bucket_name |
Yes |
amazon-braket-covalent-job-resources |
The S3 bucket where Covalent will store input and output files for the task. |
ecr_repo_name |
No |
“covalent-braket-job-images” |
The ECR repo where the job container will be uploaded |
quantum_device |
No |
“arn:aws:braket:::device/quantum-simulator/amazon/sv1” |
The ARN of the quantum device to use |
classical_device |
No |
“ml.m5.large” |
Instance type for the classical device to use |
storage |
No |
30 |
Storage size in GB for the classical device |
time_limit |
No |
300 |
Max running time in seconds for the Braket job |
poll_freq |
No |
30 |
How often (in seconds) to poll Braket for the job status |
cache_dir |
No |
“/tmp/covalent” |
Location for storing temporary files generated by the Covalent server |
This plugin can be configured in one of two ways:
Configuration options can be passed in as constructor keys to the executor class
ct.executor.BraketExecutor
By modifying the covalent configuration file under the section
[executors.braket]
The following shows an example of how a user might modify their covalent configuration file to support this plugin:
[executors.braket]
quantum_device = "arn:aws:braket:::device/qpu/ionq/ionQdevice"
time_limit = 3600
4. Required Cloud Resources#
The Braket executor requires some resources to be provisioned on AWS. Precisely, users will need an S3 bucket, an ECR repo, and an IAM role with the appropriate permissions to be passed to Braket.
Resource |
Is Required |
Config Key |
Description |
---|---|---|---|
IAM role |
Yes |
|
An IAM role granting permissions to Braket, S3, ECR, and a few other resources. |
ECR repository |
Yes |
|
An ECR repository for storing container images to be run by Braket. |
S3 bucket |
Yes |
|
An S3 bucket for storing task-specific data, such as Braket outputs or function inputs. |
Braket jobs are packaged and shipped in containers together with some supporting packages; for more context, see the “Bring your own container” documentation on AWS. This is why an ECR repo is needed. For more information on configuring an ECR repository, consult the AWS documentation.
The AWS documentation on S3 details how to configure an S3 bucket.
The permissions required for the the IAM role are documented in the article “managing access to Amazon Braket”. The following policy is attached to the default role “CovalentBraketJobsExecutionRole”:
Sample IAM policy for Braket’s execution role
- {
“Version”: “2012-10-17”, “Statement”: [
- {
“Sid”: “VisualEditor0”, “Effect”: “Allow”, “Action”: “cloudwatch:PutMetricData”, “Resource”: “*”, “Condition”: {
- “StringEquals”: {
“cloudwatch:namespace”: “/aws/braket”
}
}
}, {
“Sid”: “VisualEditor1”, “Effect”: “Allow”, “Action”: [
“logs:CreateLogStream”, “logs:DescribeLogStreams”, “ecr:GetDownloadUrlForLayer”, “ecr:BatchGetImage”, “logs:StartQuery”, “logs:GetLogEvents”, “logs:CreateLogGroup”, “logs:PutLogEvents”, “ecr:BatchCheckLayerAvailability”
], “Resource”: [
“arn:aws:ecr::348041629502:repository/”, “arn:aws:logs:::log-group:/aws/braket*”
]
}, {
“Sid”: “VisualEditor2”, “Effect”: “Allow”, “Action”: “iam:PassRole”, “Resource”: “arn:aws:iam::348041629502:role/CovalentBraketJobsExecutionRole”, “Condition”: {
- “StringLike”: {
“iam:PassedToService”: “braket.amazonaws.com”
}
}
}, {
“Sid”: “VisualEditor3”, “Effect”: “Allow”, “Action”: [
“braket:SearchDevices”, “s3:CreateBucket”, “ecr:BatchDeleteImage”, “ecr:BatchGetRepositoryScanningConfiguration”, “ecr:DeleteRepository”, “ecr:TagResource”, “ecr:BatchCheckLayerAvailability”, “ecr:GetLifecyclePolicy”, “braket:CreateJob”, “ecr:DescribeImageScanFindings”, “braket:GetJob”, “ecr:CreateRepository”, “ecr:PutImageScanningConfiguration”, “ecr:GetDownloadUrlForLayer”, “ecr:DescribePullThroughCacheRules”, “ecr:GetAuthorizationToken”, “ecr:DeleteLifecyclePolicy”, “braket:ListTagsForResource”, “ecr:PutImage”, “s3:PutObject”, “s3:GetObject”, “braket:GetDevice”, “ecr:UntagResource”, “ecr:BatchGetImage”, “ecr:DescribeImages”, “braket:CancelQuantumTask”, “ecr:StartLifecyclePolicyPreview”, “braket:CancelJob”, “ecr:InitiateLayerUpload”, “ecr:PutImageTagMutability”, “ecr:StartImageScan”, “ecr:DescribeImageReplicationStatus”, “ecr:ListTagsForResource”, “s3:ListBucket”, “ecr:UploadLayerPart”, “ecr:CreatePullThroughCacheRule”, “ecr:ListImages”, “ecr:GetRegistryScanningConfiguration”, “braket:TagResource”, “ecr:CompleteLayerUpload”, “ecr:DescribeRepositories”, “ecr:ReplicateImage”, “ecr:GetRegistryPolicy”, “ecr:PutLifecyclePolicy”, “s3:PutBucketPublicAccessBlock”, “ecr:GetLifecyclePolicyPreview”, “ecr:DescribeRegistry”, “braket:SearchJobs”, “braket:CreateQuantumTask”, “iam:ListRoles”, “ecr:PutRegistryScanningConfiguration”, “ecr:DeletePullThroughCacheRule”, “braket:UntagResource”, “ecr:BatchImportUpstreamImage”, “braket:GetQuantumTask”, “s3:PutBucketPolicy”, “braket:SearchQuantumTasks”, “ecr:GetRepositoryPolicy”, “ecr:PutReplicationConfiguration”
], “Resource”: “*”
}, {
“Sid”: “VisualEditor4”, “Effect”: “Allow”, “Action”: “logs:GetQueryResults”, “Resource”: “arn:aws:logs:::log-group:*”
}, {
“Sid”: “VisualEditor5”, “Effect”: “Allow”, “Action”: “logs:StopQuery”, “Resource”: “arn:aws:logs:::log-group:/aws/braket*”
}
]
}