diff --git a/src/requirements.txt b/src/requirements.txt index 56a0b2bc..adcd96e2 100644 --- a/src/requirements.txt +++ b/src/requirements.txt @@ -1,10 +1,11 @@ -wheel==0.38.1 -docker==6.0.1 -fabric==2.5.0 -invoke==1.6.0 -pyfiglet==0.8.post1 -reprint==0.5.2 -ruamel.yaml==0.16.7 -boto3==1.28.53 -black==19.10b0 -junit-xml==1.9 +wheel +docker +fabric +invoke +pyfiglet +reprint +ruamel.yaml +boto3 +black +junit-xml +urllib3 \ No newline at end of file diff --git a/test/braket_tests/base/test_jobs_qaoa.py b/test/braket_tests/base/test_jobs_qaoa.py index 471ac57a..ee2533c9 100644 --- a/test/braket_tests/base/test_jobs_qaoa.py +++ b/test/braket_tests/base/test_jobs_qaoa.py @@ -12,24 +12,12 @@ # 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" - 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", **create_job_args) + job_test(account, role, s3_bucket, image_path, "base-qaoa", "autograd") diff --git a/test/braket_tests/common/braket_jobs_util.py b/test/braket_tests/common/braket_jobs_util.py index 0567005c..5f1c0653 100644 --- a/test/braket_tests/common/braket_jobs_util.py +++ b/test/braket_tests/common/braket_jobs_util.py @@ -12,20 +12,32 @@ # language governing permissions and limitations under the License. import io -import os +import sys import time from contextlib import redirect_stdout -import boto3 +from braket.aws import AwsSession +from braket.jobs import hybrid_job +from braket.devices import Devices -from braket.aws import AwsQuantumJob, AwsSession +from ...resources.qaoa_entry_point import entry_point -def job_test(account, role, s3_bucket, image_path, job_type, **kwargs): +def job_test(account, role, s3_bucket, image_path, job_type, interface): job_output = io.StringIO() with redirect_stdout(job_output): try: - create_job(account, role, s3_bucket, image_path, job_type, **kwargs) + job_args = { + "p": 2, + "seed": 1967, + "max_parallel": 10, + "num_iterations": 5, + "stepsize": 0.1, + "shots": 100, + "pl_interface": interface, + "start_time": time.time(), + } + create_job(account, role, s3_bucket, image_path, job_type, job_args) except Exception as e: print(e) output = job_output.getvalue() @@ -33,15 +45,20 @@ def job_test(account, role, s3_bucket, image_path, job_type, **kwargs): assert output.find("Braket Container Run Success") > 0, "Container did not run successfully" -def create_job(account, role, s3_bucket, image_path, job_type, **kwargs): +def create_job(account, role, s3_bucket, image_path, job_type, job_args): aws_session = AwsSession(default_bucket=s3_bucket) job_name = f"ContainerTest-{job_type}-{int(time.time())}" - AwsQuantumJob.create( + + @hybrid_job( aws_session=aws_session, job_name=job_name, - device="arn:aws:braket:::device/quantum-simulator/amazon/sv1", + device=Devices.Amazon.SV1, role_arn=f"arn:aws:iam::{account}:role/{role}", image_uri=image_path, wait_until_complete=True, - **kwargs + include_modules="test.resources", ) + def decorator_job(*args, **kwargs): + return entry_point(*args, **kwargs) + + decorator_job(**job_args) diff --git a/test/braket_tests/pytorch/test_jobs_qaoa.py b/test/braket_tests/pytorch/test_jobs_qaoa.py index cf180e7b..3ea4df31 100644 --- a/test/braket_tests/pytorch/test_jobs_qaoa.py +++ b/test/braket_tests/pytorch/test_jobs_qaoa.py @@ -17,19 +17,5 @@ def test_qaoa_circuit(account, role, s3_bucket, image_list): assert len(image_list) > 0, "Unable to find images for testing" - 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", **create_job_args) + job_test(account, role, s3_bucket, image_path, "pytorch-qaoa", "torch") diff --git a/test/braket_tests/tensorflow/test_jobs_qaoa.py b/test/braket_tests/tensorflow/test_jobs_qaoa.py index e5c07fdc..107b8ccc 100644 --- a/test/braket_tests/tensorflow/test_jobs_qaoa.py +++ b/test/braket_tests/tensorflow/test_jobs_qaoa.py @@ -17,19 +17,5 @@ def test_qaoa_circuit(account, role, s3_bucket, image_list): assert len(image_list) > 0, "Unable to find images for testing" - 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", **create_job_args) + job_test(account, role, s3_bucket, image_path, "tf-qaoa", "tf") diff --git a/test/resources/qaoa_entry_point.py b/test/resources/qaoa_entry_point.py index 32c823b5..11f589c8 100644 --- a/test/resources/qaoa_entry_point.py +++ b/test/resources/qaoa_entry_point.py @@ -11,8 +11,6 @@ # 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 @@ -20,7 +18,7 @@ from pennylane import numpy as np import pennylane as qml -from braket.jobs import save_job_checkpoint, save_job_result +from braket.jobs import get_job_device_arn, save_job_checkpoint, save_job_result from braket.jobs.metrics import log_metric from . import qaoa_utils @@ -60,22 +58,16 @@ def init_pl_device(device_arn, num_nodes, shots, max_parallel): ) -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"]) - +def entry_point( + p: int, + seed: int, + max_parallel: int, + num_iterations: int, + stepsize: float, + shots: int, + pl_interface: str, + start_time: float, +): record_test_metrics('Startup', start_time, pl_interface) interface = qaoa_utils.QAOAInterface.get_interface(pl_interface) @@ -95,7 +87,7 @@ def circuit(params, **kwargs): qml.Hadamard(wires=i) qml.layer(qaoa_layer, p, params[0], params[1]) - device_arn = os.environ["AMZN_BRAKET_DEVICE_ARN"] + device_arn = get_job_device_arn() dev = init_pl_device(device_arn, num_nodes, shots, max_parallel) np.random.seed(seed) @@ -158,7 +150,3 @@ def cost_function(params): record_test_metrics('Total', start_time, pl_interface) print("Braket Container Run Success") - - -if __name__ == "__main__": - start_function()