fix: reverting decorator change, but bumping python version to 3.10 (#…

…126)

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