def test_initialize_session_for_valid_non_regional_device(aws_session, caplog):
device_arn = "arn:aws:braket:::device/qpu/test/device-name"
first_region = aws_session.region
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = [
ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
),
ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
),
device_arn,
]
caplog.set_level(logging.INFO)
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
assert f"Changed session region from '{first_region}' to '{aws_session.region}'" in caplog.text
def test_create_fake_arg():
unexpected_kwarg = "create\\(\\) got an unexpected keyword argument 'fake_arg'"
with pytest.raises(TypeError, match=unexpected_kwarg):
AwsQuantumJob.create(
device="device",
source_module="source",
fake_arg="fake_value",
)
def test_cancel_job(quantum_job_arn, aws_session,
generate_cancel_job_response):
cancellation_status = "CANCELLING"
aws_session.cancel_job.return_value = generate_cancel_job_response(
cancellationStatus=cancellation_status)
quantum_job = AwsQuantumJob(quantum_job_arn, aws_session)
status = quantum_job.cancel()
aws_session.cancel_job.assert_called_with(quantum_job_arn)
assert status == cancellation_status
def test_initialize_session_local_device(mock_new_session, aws_session):
logger = logging.getLogger(__name__)
device = "local:provider.device.name"
# don't change a provided AwsSession
assert AwsQuantumJob._initialize_session(aws_session, device,
logger) == aws_session
# otherwise, create an AwsSession with the profile defaults
assert AwsQuantumJob._initialize_session(None, device,
logger) == mock_new_session()
def test_regional_device_raises_error(get_device_side_effect,
expected_exception, aws_session, caplog):
device_arn = f"arn:aws:braket:{aws_session.region}::device/qpu/test/device-name"
aws_session.get_device.side_effect = get_device_side_effect
logger = logging.getLogger(__name__)
caplog.set_level(logging.INFO)
with pytest.raises(expected_exception):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
aws_session.get_device.assert_called_with(device_arn)
assert not caplog.text
def test_bad_arn_format(aws_session):
logger = logging.getLogger(__name__)
device_not_found = (
"Device ARN is not a valid format: bad-arn-format. For valid Braket ARNs, "
"see 'https://docs.aws.amazon.com/braket/latest/developerguide/braket-devices.html'"
)
with pytest.raises(ValueError, match=device_not_found):
AwsQuantumJob._initialize_session(aws_session, "bad-arn-format",
logger)
def test_initialize_session_for_invalid_device(aws_session, device_arn):
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
)
device_not_found = "QPU 'arn:aws:braket:::device/qpu/test/device-name' not found."
with pytest.raises(ValueError, match=device_not_found):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_exception_in_credentials_session_region(device_arn, aws_session):
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = ClientError(
{"Error": {
"Code": "SomeOtherErrorMessage",
}},
"getDevice",
)
error_message = ("An error occurred \\(SomeOtherErrorMessage\\) "
"when calling the getDevice operation: Unknown")
with pytest.raises(ClientError, match=error_message):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_no_region_routing_simulator(aws_session):
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
)
device_arn = "arn:aws:braket:::device/simulator/test/device-name"
device_not_found = f"Simulator '{device_arn}' not found in 'us-test-1'"
with pytest.raises(ValueError, match=device_not_found):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_quantum_job_constructor_explicit_session(mock_session,
quantum_job_arn, job_region):
aws_session_mock = Mock(braket_client=Mock(meta=Mock(
region_name=job_region)))
job = AwsQuantumJob(quantum_job_arn, aws_session_mock)
assert job._aws_session == aws_session_mock
assert job.arn == quantum_job_arn
mock_session.assert_not_called()
def test_equality(quantum_job_arn, aws_session, job_region):
new_aws_session = Mock(braket_client=Mock(meta=Mock(
region_name=job_region)))
quantum_job_1 = AwsQuantumJob(quantum_job_arn, aws_session)
quantum_job_2 = AwsQuantumJob(quantum_job_arn, aws_session)
quantum_job_3 = AwsQuantumJob(quantum_job_arn, new_aws_session)
other_quantum_job = AwsQuantumJob(
"arn:aws:braket:us-west-2:875981177017:job/other-job", aws_session)
non_quantum_job = quantum_job_1.arn
assert quantum_job_1 == quantum_job_2
assert quantum_job_1 == quantum_job_3
assert quantum_job_1 is not quantum_job_2
assert quantum_job_1 is not quantum_job_3
assert quantum_job_1 is quantum_job_1
assert quantum_job_1 != other_quantum_job
assert quantum_job_1 != non_quantum_job
def test_quantum_job_constructor_default_session(aws_session_mock,
mock_session, arn,
expected_region):
mock_boto_session = Mock()
aws_session_mock.return_value = Mock()
mock_session.return_value = mock_boto_session
job = AwsQuantumJob(arn)
mock_session.assert_called_with(region_name=expected_region)
aws_session_mock.assert_called_with(boto_session=mock_boto_session)
assert job.arn == arn
assert job._aws_session == aws_session_mock.return_value
def test_exceptions_in_all_device_regions(aws_session):
device_arn = "arn:aws:braket:::device/qpu/test/device-name"
logger = logging.getLogger(__name__)
aws_session.get_device.side_effect = [
ClientError(
{"Error": {
"Code": "ResourceNotFoundException",
}},
"getDevice",
),
ClientError(
{"Error": {
"Code": "SomeOtherErrorMessage",
}},
"getDevice",
),
]
error_message = ("An error occurred \\(SomeOtherErrorMessage\\) "
"when calling the getDevice operation: Unknown")
with pytest.raises(ClientError, match=error_message):
AwsQuantumJob._initialize_session(aws_session, device_arn, logger)
def test_regional_device_switches(aws_session, caplog):
original_region = aws_session.region
device_region = "us-east-1"
device_arn = f"arn:aws:braket:{device_region}::device/qpu/test/device-name"
mock_session = Mock()
mock_session.get_device.side_effect = device_arn
aws_session.copy_session.side_effect = [mock_session]
logger = logging.getLogger(__name__)
caplog.set_level(logging.INFO)
assert mock_session == AwsQuantumJob._initialize_session(
aws_session, device_arn, logger)
aws_session.copy_session.assert_called_with(region=device_region)
mock_session.get_device.assert_called_with(device_arn)
assert f"Changed session region from '{original_region}' to '{device_region}'" in caplog.text
def test_create_job(
mock_prepare_quantum_job,
mock_logs,
aws_session,
prepare_job_args,
quantum_job_arn,
wait_until_complete,
):
test_response_args = {"testArgs": "MyTestArg"}
mock_prepare_quantum_job.return_value = test_response_args
job = AwsQuantumJob.create(wait_until_complete=wait_until_complete,
**prepare_job_args)
mock_prepare_quantum_job.assert_called_with(**prepare_job_args)
aws_session.create_job.assert_called_with(**test_response_args)
if wait_until_complete:
mock_logs.assert_called_once()
else:
mock_logs.assert_not_called()
assert job.arn == quantum_job_arn
def test_quantum_job_constructor_invalid_region(aws_session):
region_mismatch = "The aws session region does not match the region for the supplied arn."
arn = "arn:aws:braket:unknown-region:875981177017:job/quantum_job_name"
with pytest.raises(ValueError, match=region_mismatch):
AwsQuantumJob(arn, aws_session)
def test_name(quantum_job_arn, quantum_job_name, aws_session):
quantum_job = AwsQuantumJob(quantum_job_arn, aws_session)
assert quantum_job.name == quantum_job_name
def test_arn(quantum_job_arn, aws_session):
quantum_job = AwsQuantumJob(quantum_job_arn, aws_session)
assert quantum_job.arn == quantum_job_arn
def test_quantum_job_constructor_invalid_region(aws_session):
arn = "arn:aws:braket:unknown-region:875981177017:job/quantum_job_name"
AwsQuantumJob(arn, aws_session)
from braket.aws import AwsDevice, AwsQuantumJob
from braket.circuits import Circuit
from braket.jobs import save_job_result
def run_job():
device = AwsDevice(os.environ.get("AMZN_BRAKET_DEVICE_ARN"))
bell = Circuit().h(0).cnot(0, 1)
num_tasks = 10
results = []
for i in range(num_tasks):
task = device.run(bell, shots=100)
result = task.result().measurement_counts
results.append(result)
print(f"iter {i}: {result}")
save_job_result({"results": results})
if __name__ == "__main__":
job = AwsQuantumJob.create(
device="arn:aws:braket:::device/quantum-simulator/amazon/sv1",
source_module="job.py",
entry_point="job:run_job",
wait_until_complete=True,
)
print(job.result())
def quantum_job(quantum_job_arn, aws_session):
return AwsQuantumJob(quantum_job_arn, aws_session)