infra: use a CUDA-Q requirements file

examples/nvidia_cuda_q/0_hello_cudaq_jobs.ipynb

@@ -270,7 +270,7 @@
     "\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",
+    "- `requirements-cuda.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",
@@ -309,17 +309,19 @@
     "device_arn = \"arn:aws:braket:eu-north-1::device/qpu/iqm/Garnet\"  # set device to IQM Garnet\n",
     "# device_arn = \"arn:aws:braket:::device/quantum-simulator/amazon/sv1\"  # set device to SV1\n",
     "\n",
-    "cudaq.set_target(\"braket\", machine=device_arn) \n",
+    "cudaq.set_target(\"braket\", machine=device_arn)\n",
     "\n",
-    "# Define our kernel \n",
-    "@cudaq.kernel \n",
-    "def ghz_state(qubit_count: int): # noqa: F821\n",
-    "    qubits = cudaq.qvector(qubit_count) \n",
-    "    h(qubits[0]) \n",
-    "    for q in range(qubit_count - 1): \n",
-    "        cx(qubits[0], qubits[q + 1]) \n",
+    "\n",
+    "# Define our kernel\n",
+    "@cudaq.kernel\n",
+    "def ghz_state(qubit_count: int):  # noqa: F821\n",
+    "    qubits = cudaq.qvector(qubit_count)\n",
+    "    h(qubits[0])\n",
+    "    for q in range(qubit_count - 1):\n",
+    "        cx(qubits[0], qubits[q + 1])\n",
     "    mz(qubits)\n",
     "\n",
+    "\n",
     "result = cudaq.sample(ghz_state, 8, shots_count=1000)\n",
     "measurement_probabilities = dict(result.items())\n",
     "print(measurement_probabilities)"

examples/nvidia_cuda_q/2_parallel_simulations.ipynb

@@ -190,9 +190,7 @@
     "\n",
     "    # Parallelize circuit simulation\n",
     "    t0 = time.time()\n",
-    "    result = cudaq.observe(\n",
-    "        kernel, hamiltonian, shots_count=n_shots, execution=cudaq.parallel.mpi\n",
-    "    )\n",
+    "    result = cudaq.observe(kernel, hamiltonian, shots_count=n_shots, execution=cudaq.parallel.mpi)\n",
     "    t1 = time.time()\n",
     "    total_time = t1 - t0\n",
     "    print(f\"rank={rank} | result: {result.expectation()} | time: {total_time}\")\n",
@@ -231,9 +229,7 @@
    "outputs": [],
    "source": [
     "parallel_obs_result = parallel_obs_job.result()\n",
-    "print(\n",
-    "    f\"result: {parallel_obs_result['expectation']} | time: {parallel_obs_result['total_time']}\"\n",
-    ")"
+    "print(f\"result: {parallel_obs_result['expectation']} | time: {parallel_obs_result['total_time']}\")"
    ]
   },
   {

examples/nvidia_cuda_q/container/Dockerfile

@@ -6,12 +6,9 @@ ARG SCRIPT_PATH
 ARG CUDAQ_PATH=/opt/conda/lib/python3.10/site-packages
 ENV MPI_PATH=/opt/amazon/openmpi
 
-RUN python3 -m pip install cudaq
+RUN python3 -m pip install --no-cache --upgrade -r requirements-cuda.txt
 RUN bash "${CUDAQ_PATH}/distributed_interfaces/activate_custom_mpi.sh"
 
-# install additional python dependencies
-RUN python3 -m pip install --no-cache --upgrade -r requirements.txt
-
 # Setup our entry point
 COPY "${SCRIPT_PATH}/braket_container.py" /opt/ml/code/braket_container.py
 ENV SAGEMAKER_PROGRAM=braket_container.py

examples/nvidia_cuda_q/container/braket_container.py

@@ -196,9 +196,7 @@ def get_code_setup_parameters() -> Tuple[str, str, str]:
             if not entry_point:
                 entry_point = hyperparameters.get("AMZN_BRAKET_SCRIPT_ENTRY_POINT")
             if not compression_type:
-                compression_type = hyperparameters.get(
-                    "AMZN_BRAKET_SCRIPT_COMPRESSION_TYPE"
-                )
+                compression_type = hyperparameters.get("AMZN_BRAKET_SCRIPT_COMPRESSION_TYPE")
         except Exception:
             log_failure_and_exit("Hyperparameters not specified in env")
     if not s3_uri:
@@ -210,13 +208,16 @@ def get_code_setup_parameters() -> Tuple[str, str, str]:
 
 def install_additional_requirements() -> None:
     """
-    Search for requirements from requirements.txt and install them.
+    Search for requirements from requirements-cuda.txt and install them.
+
+    The file requirements-cuda.txt contain the NBI requirements and the additional
+    CUDA-Q required libraries.
     """
     try:
         print("Checking for Additional Requirements")
         for root, _, files in os.walk(EXTRACTED_CUSTOMER_CODE_PATH):
-            if "requirements.txt" in files:
-                requirements_file_path = os.path.join(root, "requirements.txt")
+            if "requirements-cuda.txt" in files:
+                requirements_file_path = os.path.join(root, "requirements-cuda.txt")
                 subprocess.run(
                     ["python", "-m", "pip", "install", "-r", requirements_file_path],
                     cwd=EXTRACTED_CUSTOMER_CODE_PATH,
@@ -260,9 +261,7 @@ def wrapped_customer_code(**kwargs):
                 return customer_method(**kwargs)
         except Exception as e:
             exception_type = type(e).__name__
-            exception_string = (
-                exception_type if not str(e) else f"{exception_type}: {e}"
-            )
+            exception_string = exception_type if not str(e) else f"{exception_type}: {e}"
             _log_failure(exception_string, display=False)
             raise e
 

requirements-cuda.txt

@@ -0,0 +1,2 @@
+-r requirements.txt
+cuda-q