Bring your own container (BYOC) - Amazon Braket

Bring your own container (BYOC)

Amazon Braket Hybrid Jobs provides three pre-built containers for running code in different environments that are described in the Define the environment for your algorithm script topic. If one of these containers supports your use case, you only have to provide your algorithm script when you create a job. Minor missing dependencies can be added from your algorithm script using pip.

Note

The Strawberry Fields plugin is not natively included in the Amazon Braket hybrid jobs containers. If you wish to run a hybrid workload using the Xanadu device, you have to BYOC with the Strawberry Fields dependencies included.

If none of these containers support your use case or if you wish to expand on them, Amazon Braket Hybrid Jobs supports running jobs with your own custom Docker container image. This BYOC capability enables you to run code in an environment that has software installed that you specify. The container starts running a specified entrypoint script, which you may then use to access custom user code and data. For examples, see the Amazon Braket Hybrid Jobs tutorials.

To run Amazon Braket Hybrid Jobs in your own container:
  1. Set up prerequisites: To build and upload your custom container, you must have Docker installed. Amazon Braket notebooks come pre-installed with Docker.

  2. Create your container: Create a new directory with a Dockerfile to install and set up the software and the environment your script needs to run. You can download an example file by following the instructions in the Define the environment for your algorithm script topic. You may need to add additional packages to support your container. This file may have more software than you need, in which case you may speed up the build process by removing the unnecessary components. Note that the sagemaker-training component is necessary and should not be removed. You must specify the initial script to run when your container starts with the environment variable ENV SAGEMAKER_PROGRAM your_file.py, where your_file.py is a file that exists in the directory /opt/ml/code inside of your container. When creating a job, the user is able to specify an Amazon S3 location with code and specify an entry point to run. Braket Jobs creates environment variables for your script to find user input code to run. The Amazon S3 location and entry point to the user code are exposed to the initial script through the environment variables AMZN_BRAKET_SCRIPT_S3_URI and AMZN_BRAKET_SCRIPT_ENTRY_POINT respectively. The user input may be compressed and a compression type can be accessed through the variable AMZN_BRAKET_SCRIPT_COMPRESSION_TYPE. It is the responsibility of your script, your_file.py, to use this information to run the user code. An example script that runs user code on user data is braket_container.py in the examples you can download as described in the Define the environment for your algorithm script topic.

  3. Create an Amazon ECR repository: The repository you create must be private. If you specify a public image, the Amazon Braket Hybrid Jobs client throws an error when it attempts to validate the image location. Also, the AWS Region of the container image must match the AWS Region where you run the job. To set up your repository, access the Amazon ECR console. Then, select Create repository, choose private for visibility setting, and give your new repository a name. For a full user guide on setting up a repository, see Getting started with Amazon ECR using the AWS Management Console. Note that the standard role created for jobs only has permissions to download container images beginning with the prefix "amazon-braket-". You will need to add inline permissions to your job role if you wish to use container images with a different name (see current policy restrictions here). If you plan to run jobs from an account that is different than the account you used to create the repository, you must allow the account used to run the jobs to pull the images you have created from the account used to create them. The actions that you must allow on your repository are ecr:DescribeImages and ecr:BatchGetImage. To add these permissions to the image from the Amazon ECR console, select your image and choose permissions in the menu. You can add a rule to allow these two actions for the account you wish to allow to run jobs with your container.

  4. Build the image, and push it to the repository: You must authenticate Docker to upload to Amazon ECR. You can retrieve an authentication token and authenticate your Docker client to your registry with the AWS CLI using the following code.

aws ecr get-login-password --region ${your_region} | docker login --username AWS --password-stdin ${aws_account_id}.dkr.ecr.${your_region}.amazonaws.com

Now you should be able to build your image and push it to your repository.

cd ${your_docker_directory} docker build -t ${your_job_container} . docker tag ${your_job_container}:latest ${aws_account_id}.dkr.ecr.${your_region}.amazonaws.com/${your_job_container}:latest docker push ${aws_account_id}.dkr.ecr.${your_region}.amazonaws.com/${your_job_container}:latest

To create a job with your own container, call AwsQuantumJob.create with the argument image_uri specified. You can use a QPU, an on-demand simulator, or run your code locally on the classical processor available with Amazon Braket Hybrid jobs. To run your code on the classical processor, specify the instanceType and the instanceCount you wish to use by updating the InstanceConfig. Note that if you specify an instance_count > 1, you will need to make sure that your code has the ability to run across multiple hosts. The upper limit for the number of instances you can choose is 5.

Let’s look at an example.

job = AwsQuantumJob.create( source_module="source_dir", entry_point="source_dir.algorithm_script:start_here", image_uri="111122223333.dkr.ecr.us-west-2.amazonaws.com/my-byoc-container:latest", instance_config=InstanceConfig(instanceType="ml.p3.8xlarge", instanceCount=3), device="local:braket/braket.local.qubit", # ... )
Note

The device arn allows you to track the simulator you used as job metadata. Acceptable values must follow the format device = "local:<provider>/<simulator_name>". Remember that <provider> and <simulator_name> must consist only of letters, numbers, _, - and .. The string is limited to 256 characters.

If you plan to use BYOC and you are not using the Amazon Braket SDK to create tasks, you should pass the value of the environmental variable AMZN_BRAKET_JOB_TOKEN to the jobToken parameter in the CreateQuantumTask request. Failing to do so causes the tasks to not get priority and to be billed as regular standalone tasks.