def _initialize_non_regional_device_session(aws_session: AwsSession,
device: AwsDevice,
logger: Logger) -> AwsSession:
original_region = aws_session.region
try:
aws_session.get_device(device)
return aws_session
except ClientError as e:
if e.response["Error"]["Code"] == "ResourceNotFoundException":
if "qpu" not in device:
raise ValueError(
f"Simulator '{device}' not found in '{original_region}'"
)
else:
raise e
for region in frozenset(AwsDevice.REGIONS) - {original_region}:
device_session = aws_session.copy_session(region=region)
try:
device_session.get_device(device)
logger.info(
f"Changed session region from '{original_region}' to '{device_session.region}'"
)
return device_session
except ClientError as e:
if e.response["Error"]["Code"] != "ResourceNotFoundException":
raise e
raise ValueError(f"QPU '{device}' not found.")
def run_batch(
self,
task_specifications: List[Union[Circuit, Problem, OpenQasmProgram, BlackbirdProgram]],
s3_destination_folder: Optional[AwsSession.S3DestinationFolder] = None,
shots: Optional[int] = None,
max_parallel: Optional[int] = None,
max_connections: int = AwsQuantumTaskBatch.MAX_CONNECTIONS_DEFAULT,
poll_timeout_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_TIMEOUT,
poll_interval_seconds: float = AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL,
*aws_quantum_task_args,
**aws_quantum_task_kwargs,
) -> AwsQuantumTaskBatch:
"""Executes a batch of tasks in parallel
Args:
task_specifications (List[Union[Circuit, Problem, OpenQasmProgram, BlackbirdProgram]]):
List of circuits or annealing problems to run on device.
s3_destination_folder (Optional[S3DestinationFolder]): The S3 location to
save the tasks' results to. Default is `<default_bucket>/tasks` if evoked
outside of a Braket Job, `<Job Bucket>/jobs/<job name>/tasks` if evoked inside of
a Braket Job.
shots (Optional[int]): The number of times to run the circuit or annealing problem.
Default is 1000 for QPUs and 0 for simulators.
max_parallel (Optional[int]): The maximum number of tasks to run on AWS in parallel.
Batch creation will fail if this value is greater than the maximum allowed
concurrent tasks on the device. Default: 10
max_connections (int): The maximum number of connections in the Boto3 connection pool.
Also the maximum number of thread pool workers for the batch. Default: 100
poll_timeout_seconds (float): The polling timeout for `AwsQuantumTask.result()`,
in seconds. Default: 5 days.
poll_interval_seconds (float): The polling interval for results in seconds.
Default: 1 second.
Returns:
AwsQuantumTaskBatch: A batch containing all of the tasks run
See Also:
`braket.aws.aws_quantum_task_batch.AwsQuantumTaskBatch`
"""
return AwsQuantumTaskBatch(
AwsSession.copy_session(self._aws_session, max_connections=max_connections),
self._arn,
task_specifications,
s3_destination_folder
or (
AwsSession.parse_s3_uri(os.environ.get("AMZN_BRAKET_TASK_RESULTS_S3_URI"))
if "AMZN_BRAKET_TASK_RESULTS_S3_URI" in os.environ
else None
)
or (self._aws_session.default_bucket(), "tasks"),
shots if shots is not None else self._default_shots,
max_parallel=max_parallel if max_parallel is not None else self._default_max_parallel,
max_workers=max_connections,
poll_timeout_seconds=poll_timeout_seconds,
poll_interval_seconds=poll_interval_seconds,
*aws_quantum_task_args,
**aws_quantum_task_kwargs,
)
def _get_regional_device_session(self, session: AwsSession) -> AwsSession:
device_region = AwsDevice.get_device_region(self._arn)
region_session = (
session
if session.region == device_region
else AwsSession.copy_session(session, device_region)
)
try:
self._populate_properties(region_session)
return region_session
except ClientError as e:
raise ValueError(f"'{self._arn}' not found") if e.response["Error"][
"Code"
] == "ResourceNotFoundException" else e
def _initialize_regional_device_session(aws_session: AwsSession,
device: AwsDevice,
logger: Logger) -> AwsSession:
device_region = AwsDevice.get_device_region(device)
current_region = aws_session.region
if current_region != device_region:
aws_session = aws_session.copy_session(region=device_region)
logger.info(
f"Changed session region from '{current_region}' to '{device_region}'"
)
try:
aws_session.get_device(device)
return aws_session
except ClientError as e:
raise ValueError(f"'{device}' not found.") if e.response["Error"][
"Code"] == "ResourceNotFoundException" else e
def _get_session_and_initialize(self, session):
current_region = session.region
try:
self._populate_properties(session)
return session
except ClientError as e:
if e.response["Error"]["Code"] == "ResourceNotFoundException":
if "qpu" not in self._arn:
raise ValueError(f"Simulator '{self._arn}' not found in '{current_region}'")
else:
raise e
# Search remaining regions for QPU
for region in frozenset(AwsDevice.REGIONS) - {current_region}:
region_session = AwsSession.copy_session(session, region)
try:
self._populate_properties(region_session)
return region_session
except ClientError as e:
if e.response["Error"]["Code"] != "ResourceNotFoundException":
raise e
raise ValueError(f"QPU '{self._arn}' not found")
def get_devices(
arns: Optional[List[str]] = None,
names: Optional[List[str]] = None,
types: Optional[List[AwsDeviceType]] = None,
statuses: Optional[List[str]] = None,
provider_names: Optional[List[str]] = None,
order_by: str = "name",
aws_session: Optional[AwsSession] = None,
) -> List[AwsDevice]:
"""
Get devices based on filters and desired ordering. The result is the AND of
all the filters `arns`, `names`, `types`, `statuses`, `provider_names`.
Examples:
>>> AwsDevice.get_devices(provider_names=['Rigetti'], statuses=['ONLINE'])
>>> AwsDevice.get_devices(order_by='provider_name')
>>> AwsDevice.get_devices(types=['SIMULATOR'])
Args:
arns (List[str], optional): device ARN list, default is `None`
names (List[str], optional): device name list, default is `None`
types (List[AwsDeviceType], optional): device type list, default is `None`
QPUs will be searched for all regions and simulators will only be
searched for the region of the current session.
statuses (List[str], optional): device status list, default is `None`
provider_names (List[str], optional): provider name list, default is `None`
order_by (str, optional): field to order result by, default is `name`.
Accepted values are ['arn', 'name', 'type', 'provider_name', 'status']
aws_session (AwsSession, optional) aws_session: An AWS session object.
Default is `None`.
Returns:
List[AwsDevice]: list of AWS devices
"""
if order_by not in AwsDevice._GET_DEVICES_ORDER_BY_KEYS:
raise ValueError(
f"order_by '{order_by}' must be in {AwsDevice._GET_DEVICES_ORDER_BY_KEYS}"
)
types = (
frozenset(types) if types else frozenset({device_type for device_type in AwsDeviceType})
)
aws_session = aws_session if aws_session else AwsSession()
device_map = {}
session_region = aws_session.boto_session.region_name
search_regions = (
(session_region,) if types == {AwsDeviceType.SIMULATOR} else AwsDevice.REGIONS
)
for region in search_regions:
session_for_region = (
aws_session
if region == session_region
else AwsSession.copy_session(aws_session, region)
)
# Simulators are only instantiated in the same region as the AWS session
types_for_region = sorted(
types if region == session_region else types - {AwsDeviceType.SIMULATOR}
)
region_device_arns = [
result["deviceArn"]
for result in session_for_region.search_devices(
arns=arns,
names=names,
types=types_for_region,
statuses=statuses,
provider_names=provider_names,
)
]
device_map.update(
{
arn: AwsDevice(arn, session_for_region)
for arn in region_device_arns
if arn not in device_map
}
)
devices = list(device_map.values())
devices.sort(key=lambda x: getattr(x, order_by))
return devices
