examples/nvidia_cuda_q/0_hello_cudaq_jobs.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "57cc0f78-56cc-4a24-b5af-7e156a17a002",
   "metadata": {},
   "source": [
    "# Hello CUDA-Q with Braket Hybrid Jobs\n",
    "[CUDA-Q](https://nvidia.github.io/cuda-quantum/latest/index.html) offers a unified programming model designed for hybrid workloads that run on CPUs, GPUs and QPUs. In this notebook, you will learn how to run CUDA-Q programs on Amazon Braket. Specifically, we will be using the Bring Your Own Container (BYOC) feature of Braket Hybrid Jobs. BYOC enables you to configure the environment that you want to use in your jobs. This notebook assumes basic knowledge of Amazon Braket Hybrid Jobs. You can learn about Hybrid Jobs from [this page](https://docs.aws.amazon.com/braket/latest/developerguide/braket-what-is-hybrid-job.html) in the Amazon Braket Developer Guide and the notebooks in [Amazon Braket Examples](https://github.com/amazon-braket/amazon-braket-examples/tree/main/examples/hybrid_jobs)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fbd971e-d6c1-4180-80e3-1844d6c8afc1",
   "metadata": {},
   "source": [
    "The shell script, `container_build_and_push.sh`, in the \"container\" folder will build a Docker container and upload the container image to your [ECR](https://aws.amazon.com/ecr/) repository with the name and region that you specify.\n",
    "\n",
    "By default, the permission policy `AmazonBraketFullAccess` only grants read access to ECR container images which name start with `amazon-braket`. Because we are going to create, push a ECR repository and  access a repository which is not prefixed with `amazon-braket`, you need to attach additional permissions to your IAM identity. If you are running this notebook on an Amazon Braket notebook instance, you may attach the `AmazonEC2ContainerRegistryFullAccess` policy to the IAM role you specified when creating the notebook instance.\n",
    "\n",
    "After you attach the additional ECR policy, we can use the provided shell script to build a container image! First, use `chmod` to allow the execution of the shell script. Then, choose a image name and an AWS region name to store the image, and then run this script:\n",
    "```\n",
    "container/container_build_and_push.sh <image-name> <region-name>\n",
    "```\n",
    "The container image will persist in your ECR repository until you update it. The ARN of the image will take the format: `<aws-account-id>.dkr.ecr.<region-name>.amazonaws.com/<container-image-name>:latest`. To learn more about what's under the hood when you call the shell script and how to configure the environment of the job container, you can view the [Appendix](#Appendix:-Procedure-of-the-container-build) of this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8abe18f4-0cf1-4d4e-96da-521aac2ab3fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "!chmod +x container/container_build_and_push.sh\n",
    "!container/container_build_and_push.sh cudaq-job us-west-2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b2ebf4-24f0-4706-960e-2e3898322cf1",
   "metadata": {},
   "source": [
    "## Running CUDA-Q with Braket\n",
    "Now we have prepared the environment for CUDA-Q in a container image. Let's run our first CUDA-Q job! First, we start with the necessary imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d93602c9-c064-452c-85d4-3d2debd38cfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from braket.jobs import hybrid_job\n",
    "from braket.jobs.environment_variables import get_job_device_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baaa87e5-2b94-4892-ad4e-b59b785ab53b",
   "metadata": {},
   "source": [
    "We also prepare the URI of the container image. Fill the proper value of `aws_account_id`, `region_name` and `container_image_name` in the cell below. For example, with the shell command above, `region_name=\"us-west-2\"` and `container_image_name=\"cudaq-job\"`. The cell below prints out the image URI. When you use a container image to run a job, it ensures that your code is run in the same environment every time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "047c78b0-5814-4eb2-8604-a13b43c3179d",
   "metadata": {},
   "outputs": [],
   "source": [
    "aws_account_id = \"<aws-account-id>\"\n",
    "region_name = \"<region-name>\"\n",
    "container_image_name = \"<container-image-name>\"\n",
    "\n",
    "image_uri = f\"{aws_account_id}.dkr.ecr.{region_name}.amazonaws.com/{container_image_name}:latest\"\n",
    "print(image_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a902c96e-7da5-4ad0-b20b-ef712732e6ec",
   "metadata": {},
   "source": [
    "## Test your CUDA-Q job locally\n",
    "Before submitting a job, it is recommended to test with a local job. A local job runs scripts in a container locally. It is a good way to test your code with a small problem size before scaling up. Note, running a local hybrid job requires Docker installed on the local computer. \n",
    "\n",
    "Here, let's use a Bell circuit with CUDA-Q for the local job. This example does not require a GPU to run. The `hello_quantum` function in the code snippet below defines an experiment to sample a Bell circuit. The string `qpp-cpu` in the `device` keyword argument of the decorator is the name of a CUDA-Q CPU simulator. You can view the tutorial of CUDA-Q and the available backends in the [CUDA-Q documentation](https://nvidia.github.io/cuda-quantum/latest/index.html).\n",
    "\n",
    "When called, the decorated `hello_quantum` function starts a local job because of the keyword `local=True` in the `hybrid_job` decorator. The code inside the `hello_quantum` function will run locally in the environment defined by the container image that you have built earlier."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09d6bd61-6fb2-4fbc-9d60-93d4161a724e",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hybrid_job(device=\"local:cudaq/qpp-cpu\", image_uri=image_uri, local=True)\n",
    "def hello_quantum():\n",
    "    import cudaq\n",
    "\n",
    "    # define the backend\n",
    "    device = get_job_device_arn()\n",
    "    cudaq.set_target(device.split(\"/\")[-1])\n",
    "    print(\"CUDA-Q backend: \", cudaq.get_target())\n",
    "\n",
    "    @cudaq.kernel \n",
    "    def bell_state():\n",
    "        qubits = cudaq.qvector(2) \n",
    "        h(qubits[0]) \n",
    "        cx(qubits[0], qubits[1]) \n",
    "\n",
    "    # sample the Bell circuit\n",
    "    result = cudaq.sample(bell_state, shots_count=1000)\n",
    "    measurement_probabilities = dict(result.items())\n",
    "    print(\"Samples: \", measurement_probabilities)\n",
    "\n",
    "    return measurement_probabilities"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5c90cb7-baa6-4c51-b828-e1f334b3eb5c",
   "metadata": {},
   "source": [
    "Let's test your CUDA-Q job!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d402d3b9-1623-4d39-a8e6-f8f69d40ba7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "hello_quantum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4871d290-8173-454a-a69c-d9625b80acc7",
   "metadata": {},
   "source": [
    "## Run your CUDA-Q job\n",
    "After testing locally that your code works correctly in the container environment, you can remove the `local=True` keyword argument (or, equivalently, set `local=False`), so the next job will run on an AWS-managed compute instance. This is great to scale up to larger instances or parallelize your workload over multiple instances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1329a6f9-8d2e-4067-905a-e16590edfa50",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hybrid_job(device=\"local:cudaq/qpp-cpu\", image_uri=image_uri)\n",
    "def hello_quantum():\n",
    "    import cudaq\n",
    "\n",
    "    device = get_job_device_arn()\n",
    "    cudaq.set_target(device.split(\"/\")[-1])\n",
    "    print(cudaq.get_target())\n",
    "\n",
    "    @cudaq.kernel \n",
    "    def bell_state():\n",
    "        qubits = cudaq.qvector(2) \n",
    "        h(qubits[0]) \n",
    "        cx(qubits[0], qubits[1]) \n",
    "\n",
    "    result = cudaq.sample(bell_state, shots_count=1000)\n",
    "    measurement_probabilities = dict(result.items())\n",
    "    print(\"Samples: \", measurement_probabilities)\n",
    "\n",
    "    return measurement_probabilities\n",
    "\n",
    "\n",
    "job = hello_quantum()\n",
    "print(job.arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4d56392-77a9-46cd-8512-bd0ded291b8b",
   "metadata": {},
   "source": [
    "When called, the decorated `hello_quantum` function will create a Braket Hybrid Job on AWS, running the code defined in the decorated function with the environment specified by the container image built above. You can view the progress of your job with `job.state()` or in the \"Hybrid jobs\" tab of the Amazon Braket Console."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4611f86a-4e23-4aaf-8ab8-df33b7eb050a",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = job.result()\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d9cef6b-295e-4360-a193-64b4ed945daa",
   "metadata": {},
   "source": [
    "## Run CUDA-Q jobs on Braket devices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee3996aa-195b-45d7-a007-077a0d557b9f",
   "metadata": {},
   "source": [
    "Now, let’s learn how to run CUDA-Q programs on Braket devices via CUDA-Q in Braket Hybrid Jobs. All you need to do is to configure the CUDA-Q target to a Braket device. In the code snippet below, we run a Bell circuit on Braket SV1 using CUDA-Q and Braket Hybrid Jobs. When you finish testing on simulators and are ready to run experiments on QPUs, you can switch the target to Braket QPUs such as IQM, IonQ, and Rigetti devices by changing `device_arn` in the example below. The circuits run with a hybrid job receive higher-priority access to the target Braket QPUs, which not only reduces the run time of your experiments, but also minimizes the impact of hardware drift to your algorithms. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "820a1a20-4eeb-41a0-874d-44ad96969496",
   "metadata": {},
   "outputs": [],
   "source": [
    "device_arn = \"arn:aws:braket:::device/quantum-simulator/amazon/sv1\"  # set device to SV1\n",
    "# device_arn = \"arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet\" # set device to IQM Garnet\n",
    "\n",
    "\n",
    "@hybrid_job(device=device_arn, image_uri=image_uri)\n",
    "def job_with_braket_device():\n",
    "    import cudaq\n",
    "\n",
    "    # define the backend\n",
    "    device = get_job_device_arn()\n",
    "    cudaq.set_target(\"braket\", machine=device)\n",
    "\n",
    "    # define the Bell circuit\n",
    "    @cudaq.kernel \n",
    "    def bell_state():\n",
    "        qubits = cudaq.qvector(2) \n",
    "        h(qubits[0]) \n",
    "        cx(qubits[0], qubits[1])\n",
    "        mz(qubits)\n",
    "\n",
    "    # sample the Bell circuit\n",
    "    result = cudaq.sample(bell_state, shots_count=1000)\n",
    "    measurement_probabilities = dict(result.items())\n",
    "\n",
    "    return measurement_probabilities\n",
    "\n",
    "\n",
    "job = job_with_braket_device()\n",
    "print(job.arn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a76eade-78ff-417d-8246-04ccb28b8944",
   "metadata": {},
   "source": [
    "## Summary\n",
    "This notebook shows you how to run your first CUDA-Q program with Amazon Braket Hybrid Jobs. Using the BYOC feature of Amazon Braket and a shell script we provide, you can create a CUDA-Q environment with a few lines of code. Once you have registered your CUDA-Q container image, you can run CUDA-Q programs with Braket Hybrid Jobs and scale your workloads up and out with the range of compute options provided by AWS. In the following tutorials, we will show you how to run CUDA-Q simulations on GPUs ([notebook](1_simulation_with_GPUs.ipynb)) and distribute workloads across multiple instances ([notebook](2_parallel_simulations.ipynb))."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eba7f189-1036-4585-a641-127d1dbeed88",
   "metadata": {},
   "source": [
    "## Appendix: Procedure for building the container\n",
    "\n",
    "When the shell script `container_build_and_push.sh` is called, a Docker container is built with CUDA-Q and other GPU related settings are configured. The procedure for BYOC is presented in [this page from the Braket Developer Guide](https://docs.aws.amazon.com/braket/latest/developerguide/braket-jobs-byoc.html). The required files for building a container with CUDA-Q are in the \"container\" folder, including\n",
    "- `Dockerfile`: Describes how the container is built.\n",
    "- `requirements.txt`: Additional Python dependencies to include.\n",
    "- `braket_container.py`: The start-up script of a job container.\n",
    "\n",
    "These files include basic settings to use CUDA-Q in jobs. You can modify these files to suit your needs. For example, you can add more Python dependencies to `requirements.txt` and other dependencies to `Dockerfile`.\n",
    "\n",
    "The shell script `container_build_and_push.sh` automates the procedure of building on top of a Braket managed container and pushing your custom container image. The shell script requires Docker to be installed on your local machine. If you are using Braket NBI, Docker is already installed. If you don't have Docker in your local environment, you can follow the [instructions on this page](https://www.docker.com/) to install it. The shell script takes two parameters:\n",
    "- image-name\n",
    "- region-name\n",
    "\n",
    "The shell script can be called with the following syntax:\n",
    "```\n",
    "container/container_build_and_push.sh <image-name> <region-name>\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a7dd19-489c-4eb6-857c-918ad98ee003",
   "metadata": {},
   "source": [
    "## Appendix: Using job submission script\n",
    "\n",
    "The `@hybrid_job` decorator provides an convenient interface to submit a hybrid job, but it limits the availability of the source code to the inner functions defined in the decorator. The source code is critical for `@cudaq.kernel`. For more complex workloads, if you encounter an error related to the source code availability, you can choose to submit job without the `@hybrid_job` decorator and use `AwsQuantumJob.create` instead.\n",
    "\n",
    "To create a hybrid job using `AwsQuantumJob.create`, first you need to write your CUDA-Q program as a separate `.py` file. We will call this file the \"algorithm script\". For demo purpose, we have prepared an example algorithm script, `algorithm_script_getting_started.py`. Then, you can run the following code snippet to create the job. This interface of creating a hybrid job will not have any source code error. To learn more about creating hybrid job this way, you can read this [documentation page](https://docs.aws.amazon.com/braket/latest/developerguide/braket-jobs-first.html#braket-jobs-first-create) and this [example notebook](https://github.com/amazon-braket/amazon-braket-examples/tree/main/examples/hybrid_jobs/8_Creating_Hybrid_Job_Scripts)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51d4e1c8-5d7c-443a-a0ae-db8ee10255e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a hybrid job\n",
    "job = AwsQuantumJob.create(\n",
    "    device=Devices.Amazon.SV1,\n",
    "    source_module=\"algorithm_script_getting_started.py\",\n",
    "    image_uri=image_uri,\n",
    ")\n",
    "\n",
    "# view the ARN and the status of the job\n",
    "print(\"ARN of the job: \", job.arn)\n",
    "print(\"Status of the job: \", job.status())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}