Run your local code as a hybrid job
Amazon Braket Hybrid Jobs provides a fully managed orchestration of hybrid quantum-classical algorithms, combining Amazon EC2 compute resources with Amazon Braket Quantum Processing Unit (QPU) access. Quantum tasks created in a hybrid job have priority queueing over individual quantum tasks so that your algorithms won't be interrupted by fluctuations in the quantum task queue. Each QPU maintains a separate hybrid jobs queue, ensuring that only one hybrid job can run at any given time.
In this section:
Create a hybrid job from local Python code
You can run your local Python code as an Amazon Braket Hybrid Job. You can do this by annotating
your code with an @hybrid_job
decorator, as shown in the following code example.
For custom environments, you can opt to use a custom container
from Amazon Elastic Container Registry (ECR).
Note
Only Python 3.10 is supported by default.
You can use the @hybrid_job
decorator to annotate a function. Braket transforms
the code inside the decorator into a Braket hybrid job algorithm script.
The hybrid job then invokes the function inside the decorator on an Amazon EC2 instance. You can monitor
the progress of the job with job.state()
or with the Braket console. The following
code example shows how to run a sequence of five states on the State Vector Simulator (SV1) device.
from braket.aws import AwsDevice from braket.circuits import Circuit, FreeParameter, Observable from braket.devices import Devices from braket.jobs.hybrid_job import hybrid_job from braket.jobs.metrics import log_metric device_arn = Devices.Amazon.SV1 @hybrid_job(device=device_arn) # choose priority device def run_hybrid_job(num_tasks=1): device = AwsDevice(device_arn) # declare AwsDevice within the hybrid job # create a parametric circuit circ = Circuit() circ.rx(0, FreeParameter("theta")) circ.cnot(0, 1) circ.expectation(observable=Observable.X(), target=0) theta = 0.0 # initial parameter for i in range(num_tasks): task = device.run(circ, shots=100, inputs={"theta": theta}) # input parameters exp_val = task.result().values[0] theta += exp_val # modify the parameter (possibly gradient descent) log_metric(metric_name="exp_val", value=exp_val, iteration_number=i) return {"final_theta": theta, "final_exp_val": exp_val}
You create the hybrid job by invoking the function as you would normal Python functions.
However, the decorator function returns the hybrid job handle rather than the result of the function.
To retrieve the results after it has completed, use job.result()
.
job = run_hybrid_job(num_tasks=1) result = job.result()
The device argument in the @hybrid_job
decorator specifies the device that the
hybrid job has priority access to - in this case, the SV1 simulator. To get QPU priority,
you must ensure that the device ARN used within the function matches that specified in the decorator.
For convenience, you can use the helper function get_job_device_arn()
to capture the device
ARN declared in @hybrid_job
.
Note
Each hybrid job has at least a one minute startup time since it creates a containerized environment on Amazon EC2. So for very short workloads, such as a single circuit or a batch of circuits, it may suffice for you to use quantum tasks.
Hyperparameters
The run_hybrid_job()
function takes the argument num_tasks
to control the number of
quantum tasks created. The hybrid job automatically captures this as a hyperparameter.
Note
Hyperparameters are displayed in the Braket console as strings, that are limited to 2500 characters.
Metrics and logging
Within the run_hybrid_job()
function, metrics from iterative algorithms are recorded with
log_metrics
. Metrics are automatically plotted in the Braket console page under the hybrid job tab.
You can use metrics to track the quantum task costs in near-real time during the hybrid job run with the
Braket cost tracker. The example above uses the metric name “probability”
that records the first probability from the result type.
Retrieving results
After the hybrid job has completed, you use job.result()
to retrieve the hybrid jobs results.
Any objects in the return statement are automatically captured by Braket. Note that the objects returned by
the function must be a tuple with each element being serializable. For example, the following code shows a
working, and a failing example.
@hybrid_job(device=Devices.Amazon.SV1) def passing(): np_array = np.random.rand(5) return np_array # serializable @hybrid_job(device=Devices.Amazon.SV1) def failing(): return MyObject() # not serializable
Job name
By default, the name for this hybrid job is inferred from the function name. You may also specify a custom name up to 50 characters long. For example, in the following code the job name is "my-job-name".
@hybrid_job(device=Devices.Amazon.SV1, job_name="my-job-name") def function(): pass
Local mode
Local jobs are be created by adding the argument local=True
to the decorator. This runs the hybrid job in a containerized environment on your local compute environment, such as your
laptop. Local jobs do not have priority queueing for quantum tasks. For advanced cases
such as multi-node or MPI, local jobs may have access to the required Braket environment variables. The following code creates
a local hybrid job with the device as the SV1 simulator.
@hybrid_job(device=Devices.Amazon.SV1, local=True) def run_hybrid_job(num_tasks = 1): return ...
All other hybrid job options are supported. For a list of options see the
braket.jobs.quantum_job_creation module
Install additional Python packages and source code
You can customize your runtime environment to use your preferred Python packages. You can use
either a requirements.txt
file, a list of package names, or bring your own container (BYOC).
To customize a runtime environment using a requirements.txt
file, refer to the following code example.
@hybrid_job(device=Devices.Amazon.SV1, dependencies="requirements.txt") def run_hybrid_job(num_tasks = 1): return ...
For example, the requirements.txt
file may include other packages to install.
qiskit pennylane >= 0.31 mitiq == 0.29
Alternatively, you may supply the package names as a Python list as follows.
@hybrid_job(device=Devices.Amazon.SV1, dependencies=["qiskit", "pennylane>=0.31", "mitiq==0.29"]) def run_hybrid_job(num_tasks = 1): return ...
Additional source code can be specified either as a list of modules, or a single module as in the following code example.
@hybrid_job(device=Devices.Amazon.SV1, include_modules=["my_module1", "my_module2"]) def run_hybrid_job(num_tasks = 1): return ...
Save and load data into a hybrid job instance
Specifying input training data
When you create a hybrid job, you may provide an input training datasets by specifying an Amazon Simple Storage Service
(Amazon S3) bucket. You may also specify a local path, then Braket automatically uploads the
data to Amazon S3 at s3://<default_bucket_name>/jobs/<job_name>/<timestamp>/data/<channel_name>
.
If you specify a local path, the channel name defaults to “input”. The following code shows a numpy file from the local path data/file.npy
.
@hybrid_job(device=Devices.Amazon.SV1, input_data="data/file.npy") def run_hybrid_job(num_tasks = 1): data = np.load("data/file.npy") return ...
For S3, you must use the get_input_data_dir()
helper funciton.
s3_path = "s3://amazon-braket-us-west-1-961591465522/job-data/file.npy" @hybrid_job(device=None, input_data=s3_path) def job_s3_input(): np.load(get_input_data_dir() + "/file.npy") @hybrid_job(device=None, input_data={"channel": s3_path}) def job_s3_input_channel(): np.load(get_input_data_dir("channel") + "/file.npy")
You can specify multiple input data sources by providing a dictionary of channel values and S3 URIs or local paths.
input_data = { "input": "data/file.npy", "input_2": "s3://my-bucket/data.json" } @hybrid_job(device=None, input_data=input_data) def multiple_input_job(): np.load(get_input_data_dir("input") + "/file.npy") np.load(get_input_data_dir("input_2") + "/data.json")
Note
When the input data is large (>1GB), there is a long wait time before the job is created. This is due to the local input data when it is first uploaded to an S3 bucket, then the S3 path is added to the job request. Finally, the job request is submitted to the Braket service.
Saving results to S3
To save results not included in the return statement of the decorated function, you must append the correct directory to all file writing operations. The following example, shows saving a numpy array and a matplotlib figure.
@hybrid_job(device=Devices.Amazon.SV1) def run_hybrid_job(num_tasks = 1): result = np.random.rand(5) # save a numpy array np.save("result.npy", result) # save a matplotlib figure plt.plot(result) plt.savefig("fig.png") return ...
All results are compressed into a file named model.tar.gz
. You can download the results with the
Python function job.result()
, or by navigating to the results folder from the hybrid job
page in the Braket management console.
Saving and resuming from checkpoints
For long-running hybrid jobs, its recommended to periodically save the intermediate state of the algorithm.
You can use the built-in save_job_checkpoint()
helper function, or save files to the
AMZN_BRAKET_JOB_RESULTS_DIR
path. The later is available with the helper function get_job_results_dir()
.
The following is a minimal working example for saving and loading checkpoints with a hybrid job decorator:
from braket.jobs import save_job_checkpoint, load_job_checkpoint, hybrid_job @hybrid_job(device=None, wait_until_complete=True) def function(): save_job_checkpoint({"a": 1}) job = function() job_name = job.name job_arn = job.arn @hybrid_job(device=None, wait_until_complete=True, copy_checkpoints_from_job=job_arn) def continued_function(): load_job_checkpoint(job_name) continued_job = continued_function()
In the first hybrid job, save_job_checkpoint()
is called with a dictionary containing the data we want to save.
By default, every value must be serializable as text. For checkpointing more complex Python objects, such as numpy arrays, you
can set data_format = PersistedJobDataFormat.PICKLED_V4
. This code creates and overwrites a checkpoint file with
default name <jobname>.json
in your hybrid job artifacts under a subfolder called "checkpoints".
To create a new hybrid job to continue from the checkpoint, we need to pass copy_checkpoints_from_job=job_arn
where job_arn
is the hybrid job ARN of the previous job. Then we use load_job_checkpoint(job_name)
to
load from the checkpoint.
Best practices for hybrid job decorators
Embrace asynchronicity
Hybrid jobs created with the decorator annotation are asynchronous - they run once the classical and quantum resources are available. You monitor the progress of the algorithm using the Braket Management Console or Amazon CloudWatch. When you submit your algorithm to run, Braket runs your algorithm in a scalable containerized environment and results are retrieved when the algorithm is complete.
Run iterative variational algorithms
Hybrid jobs gives you the tools to run iterative quantum-classical algorithms. For purely quantum problems, use quantum tasks or a batch of quantum tasks. The priority access to certain QPUs is most beneficial for long-running variational algorithms requiring multiple iterative calls to the QPUs with classical processing in between.
Debug using local mode
Before you run a hybrid job on a QPU, its recommended to first run on the simulator SV1 to confirm it runs as expected. For small scale tests, you can run with local mode for rapid iteration and debugging.
Improve reproducibility with Bring your own container (BYOC)
Create a reproducible experiment by encapsulating your software and its dependencies within a containerized environment. By packaging all your code, dependencies, and settings in a container, you prevent potential conflicts and versioning issues.
Multi-instance distributed simulators
To run a large number of circuits, consider using built-in MPI support to run local simulators on multiple instances within a single hybrid job. For more information, see embedded simulators.
Use parametric circuits
Parametric circuits that you submit from a hybrid job are automatically compiled on certain QPUs using parametric compilation to improve the runtimes of your algorithms.
Checkpoint periodically
For long-running hybrid jobs, its recommended to periodically save the intermediate state of the algorithm.
For further examples, use cases, and best-practices, see
Amazon Braket examples GitHub