Merge branch 'main' into pl_33

.github/workflows/code-freeze.yml

@@ -0,0 +1,47 @@
+name: Code Freeze
+
+on:
+  pull_request:
+    branches:
+      - main
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+env:
+  FROZEN: ${{ vars.FROZEN }}
+  UNFROZEN_PREFIX: ${{ vars.UNFROZEN_PREFIX }}
+
+jobs:
+  check-pr-frozen-status:
+    runs-on: ubuntu-latest
+    steps:
+    - name: Fetch PR data and check if merge allowed
+      if: env.FROZEN == 'true'
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      run: |
+        PR_DATA=$(curl -s \
+          -H "Authorization: Bearer $GITHUB_TOKEN" \
+          -H "Accept: application/vnd.github.v3+json" \
+          https://api.github.com/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }})
+        BRANCH_NAME=$(echo $PR_DATA | jq .head.ref -r)
+        PR_TITLE=$(echo $PR_DATA | jq .title -r)
+        echo $BRANCH_NAME
+        echo $PR_TITLE
+        # if it's not a critical fix
+        if ! [[ "$PR_TITLE" == fix\(critical\):* ]]; then
+          # and there's an unfrozen prefix
+          if ! [[ -z $UNFROZEN_PREFIX ]]; then
+            # check if the branch matches unfrozen prefix
+            if [[ "$BRANCH_NAME" != $UNFROZEN_PREFIX* ]]; then
+              echo "Error: You can only merge from branches that start with '$UNFROZEN_PREFIX', or PRs titled with prefix 'fix(critical): '."
+              exit 1
+            fi
+          # repo is fully frozen
+          else
+            echo "Error: You can only merge PRs titled with prefix 'fix(critical): '."
+            exit 1
+          fi
+        fi

.github/workflows/python-package.yml

@@ -21,7 +21,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

buildspec.yml

@@ -9,7 +9,7 @@ phases:
   pre_build:
     commands:
       - echo Logging in to Amazon ECR...
-      - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION)
+      - aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com
       - pip install -r src/requirements.txt
       - echo The following is only required if we use DLC images as our base 'FROM' images.
       - aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.us-east-1.amazonaws.com

test/braket_tests/base/test_jobs_qaoa.py

@@ -12,21 +12,24 @@
 # language governing permissions and limitations under the License.
 
 import time
-
 from ..common.braket_jobs_util import job_test
 
 
 def test_qaoa_circuit(account, role, s3_bucket, image_list):
     assert len(image_list) > 0, "Unable to find images for testing"
-    job_args = {
-        "p": 2,
-        "seed": 1967,
-        "max_parallel": 10,
-        "num_iterations": 5,
-        "stepsize": 0.1,
-        "shots": 100,
-        "pl_interface": "autograd",
-        "start_time": time.time(),
+    create_job_args = {
+        "source_module": "./test/resources/",
+        "entry_point": "resources.qaoa_entry_point",
+        "hyperparameters": {
+            "p": "2",
+            "seed": "1967",
+            "max_parallel": "10",
+            "num_iterations": "5",
+            "stepsize": "0.1",
+            "shots": "100",
+            "interface": "autograd",
+            "start_time": time.time(),
+        }
     }
     for image_path in image_list:
-        job_test(account, role, s3_bucket, image_path, "base-qaoa", job_args)
+        job_test(account, role, s3_bucket, image_path, "base-qaoa", **create_job_args)

test/braket_tests/common/braket_jobs_util.py

@@ -12,43 +12,36 @@
 # language governing permissions and limitations under the License.
 
 import io
-import sys
+import os
 import time
 from contextlib import redirect_stdout
 
-from braket.aws import AwsSession
-from braket.jobs import hybrid_job
-from braket.devices import Devices
+import boto3
 
-from ...resources.qaoa_entry_point import entry_point
+from braket.aws import AwsQuantumJob, AwsSession
 
 
-def job_test(account, role, s3_bucket, image_path, job_type, job_args):
+def job_test(account, role, s3_bucket, image_path, job_type, **kwargs):
     job_output = io.StringIO()
     with redirect_stdout(job_output):
         try:
-            create_job(account, role, s3_bucket, image_path, job_type, job_args)
+            create_job(account, role, s3_bucket, image_path, job_type, **kwargs)
         except Exception as e:
             print(e)
     output = job_output.getvalue()
     print(output)
     assert output.find("Braket Container Run Success") > 0, "Container did not run successfully"
 
 
-def create_job(account, role, s3_bucket, image_path, job_type, job_args):
+def create_job(account, role, s3_bucket, image_path, job_type, **kwargs):
     aws_session = AwsSession(default_bucket=s3_bucket)
     job_name = f"ContainerTest-{job_type}-{int(time.time())}"
-
-    @hybrid_job(
+    AwsQuantumJob.create(
         aws_session=aws_session,
         job_name=job_name,
-        device=Devices.Amazon.SV1,
+        device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
         role_arn=f"arn:aws:iam::{account}:role/{role}",
         image_uri=image_path,
         wait_until_complete=True,
-        include_modules="test.resources",
+        **kwargs
     )
-    def decorator_job(*args, **kwargs):
-        return entry_point(*args, **kwargs)
-
-    decorator_job(**job_args)

test/braket_tests/pytorch/test_jobs_qaoa.py

@@ -17,15 +17,19 @@
 
 def test_qaoa_circuit(account, role, s3_bucket, image_list):
     assert len(image_list) > 0, "Unable to find images for testing"
-    job_args = {
-        "p": 2,
-        "seed": 1967,
-        "max_parallel": 10,
-        "num_iterations": 5,
-        "stepsize": 0.1,
-        "shots": 100,
-        "pl_interface": "torch",
-        "start_time": time.time(),
+    create_job_args = {
+        "source_module": "./test/resources/",
+        "entry_point": "resources.qaoa_entry_point",
+        "hyperparameters": {
+            "p": "2",
+            "seed": "1967",
+            "max_parallel": "10",
+            "num_iterations": "5",
+            "stepsize": "0.1",
+            "shots": "100",
+            "interface": "torch",
+            "start_time": time.time(),
+        }
     }
     for image_path in image_list:
-        job_test(account, role, s3_bucket, image_path, "pytorch-qaoa", job_args)
+        job_test(account, role, s3_bucket, image_path, "pytorch-qaoa", **create_job_args)

test/braket_tests/tensorflow/test_jobs_qaoa.py

@@ -17,15 +17,19 @@
 
 def test_qaoa_circuit(account, role, s3_bucket, image_list):
     assert len(image_list) > 0, "Unable to find images for testing"
-    job_args = {
-        "p": 2,
-        "seed": 1967,
-        "max_parallel": 10,
-        "num_iterations": 5,
-        "stepsize": 0.1,
-        "shots": 100,
-        "pl_interface": "torch",
-        "start_time": time.time(),
+    create_job_args = {
+        "source_module": "./test/resources/",
+        "entry_point": "resources.qaoa_entry_point",
+        "hyperparameters": {
+            "p": "2",
+            "seed": "1967",
+            "max_parallel": "10",
+            "num_iterations": "5",
+            "stepsize": "0.1",
+            "shots": "100",
+            "interface": "tf",
+            "start_time": time.time(),
+        }
     }
     for image_path in image_list:
-        job_test(account, role, s3_bucket, image_path, "tf-qaoa", job_args)
+        job_test(account, role, s3_bucket, image_path, "tf-qaoa", **create_job_args)

test/resources/qaoa_entry_point.py

@@ -11,14 +11,16 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
+import json
+import os
 import time
 import boto3
 
 import networkx as nx
 from pennylane import numpy as np
 import pennylane as qml
 
-from braket.jobs import get_job_device_arn, save_job_checkpoint, save_job_result
+from braket.jobs import save_job_checkpoint, save_job_result
 from braket.jobs.metrics import log_metric
 
 from . import qaoa_utils
@@ -58,16 +60,22 @@ def init_pl_device(device_arn, num_nodes, shots, max_parallel):
     )
 
 
-def entry_point(
-    p: int,
-    seed: int,
-    max_parallel: int,
-    num_iterations: int,
-    stepsize: float,
-    shots: int,
-    pl_interface: str,
-    start_time: float,
-):
+def start_function():
+    # Read the hyperparameters
+    hp_file = os.environ["AMZN_BRAKET_HP_FILE"]
+    with open(hp_file, "r") as f:
+        hyperparams = json.load(f)
+    print(hyperparams)
+
+    p = int(hyperparams["p"])
+    seed = int(hyperparams["seed"])
+    max_parallel = int(hyperparams["max_parallel"])
+    num_iterations = int(hyperparams["num_iterations"])
+    stepsize = float(hyperparams["stepsize"])
+    shots = int(hyperparams["shots"])
+    pl_interface = hyperparams["interface"]
+    start_time = float(hyperparams["start_time"])
+
     record_test_metrics('Startup', start_time, pl_interface)
 
     interface = qaoa_utils.QAOAInterface.get_interface(pl_interface)
@@ -87,7 +95,7 @@ def circuit(params, **kwargs):
             qml.Hadamard(wires=i)
         qml.layer(qaoa_layer, p, params[0], params[1])
 
-    device_arn = get_job_device_arn()
+    device_arn = os.environ["AMZN_BRAKET_DEVICE_ARN"]
     dev = init_pl_device(device_arn, num_nodes, shots, max_parallel)
 
     np.random.seed(seed)
@@ -150,3 +158,7 @@ def cost_function(params):
 
     record_test_metrics('Total', start_time, pl_interface)
     print("Braket Container Run Success")
+
+
+if __name__ == "__main__":
+    start_function()

testspec.yml

@@ -8,7 +8,7 @@ phases:
   pre_build:
     commands:
       - echo Logging in to Amazon ECR...
-      - $(aws ecr get-login --no-include-email --region $AWS_DEFAULT_REGION)
+      - aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com
       - pip install -r test/requirements.txt
       - echo Performing additional setup...
       - python test/perform_additional_setup.py