Merge branch 'main' into feature/pypi

.github/ISSUE_TEMPLATE/bug_report.md

@@ -2,7 +2,7 @@
 name: Bug report
 about: File a report to help us reproduce and fix the problem
 title: ''
-labels: ''
+labels: 'bug'
 assignees: ''
 
 ---

.github/ISSUE_TEMPLATE/documentation_request.md

@@ -2,7 +2,7 @@
 name: Documentation request
 about: Request improved documentation
 title: ''
-labels: ''
+labels: 'documentation'
 assignees: ''
 
 ---

.github/ISSUE_TEMPLATE/feature_request.md

@@ -2,7 +2,7 @@
 name: Feature request
 about: Suggest new functionality for this library
 title: ''
-labels: ''
+labels: 'feature'
 assignees: ''
 
 ---

.github/workflows/stale_issue.yml

@@ -0,0 +1,47 @@
+name: "Close stale issues"
+
+# Controls when the action will run.
+# This is every day at 10am
+on:
+  schedule:
+    - cron: "0 10 * * *"
+
+jobs:
+  cleanup:
+    runs-on: ubuntu-latest
+    name: Stale issue job
+    steps:
+      - uses: aws-actions/stale-issue-cleanup@v3
+        with:
+          # Setting messages to an empty string will cause the automation to skip
+          # that category
+          ancient-issue-message: Greetings! It looks like this issue hasn’t been active in longer than three years. We encourage you to check if this is still an issue in the latest release. Because it has been longer than three years since the last update on this, and in the absence of more information, we will be closing this issue soon. If you find that this is still a problem, please feel free to provide a comment to prevent automatic closure, or if the issue is already closed, please feel free to reopen it.
+          stale-issue-message: Greetings! It looks like this issue hasn’t been active in longer than a week. We encourage you to check if this is still an issue in the latest release. Because it has been longer than a week since the last update on this, and in the absence of more information, we will be closing this issue soon. If you find that this is still a problem, please feel free to provide a comment or add an upvote to prevent automatic closure, or if the issue is already closed, please feel free to open a new one.
+          stale-pr-message: Greetings! It looks like this PR hasn’t been active in longer than a week, add a comment or an upvote to prevent automatic closure, or if the issue is already closed, please feel free to open a new one.
+
+          # These labels are required
+          stale-issue-label: closing-soon
+          exempt-issue-label: auto-label-exempt
+          stale-pr-label: closing-soon
+          exempt-pr-label: pr/needs-review
+          response-requested-label: response-requested
+
+          # Don't set closed-for-staleness label to skip closing very old issues
+          # regardless of label
+          closed-for-staleness-label: closed-for-staleness
+
+          # Issue timing
+          days-before-stale: 7
+          days-before-close: 4
+          days-before-ancient: 1095
+
+          # If you don't want to mark a issue as being ancient based on a
+          # threshold of "upvotes", you can set this here. An "upvote" is
+          # the total number of +1, heart, hooray, and rocket reactions
+          # on an issue.
+          minimum-upvotes-to-exempt: 1
+
+          repo-token: ${{ secrets.GITHUB_TOKEN }}
+          loglevel: DEBUG
+          # Set dry-run to true to not perform label or close actions.
+          dry-run: false

src/braket/aws/aws_device.py

@@ -11,10 +11,13 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
+from __future__ import annotations
+
 from enum import Enum
-from typing import Optional, Union
+from typing import List, Optional, Set, Union
 
 import boto3
+from boltons.dictutils import FrozenDict
 from networkx import Graph, complete_graph, from_edgelist
 
 from braket.annealing.problem import Problem
@@ -41,18 +44,21 @@ class AwsDevice(Device):
     device.
     """
 
-    QPU_REGIONS = {
-        "rigetti": ["us-west-1"],
-        "ionq": ["us-east-1"],
-        "d-wave": ["us-west-2"],
-    }
+    DEVICE_REGIONS = FrozenDict(
+        {
+            "rigetti": ["us-west-1"],
+            "ionq": ["us-east-1"],
+            "d-wave": ["us-west-2"],
+            "amazon": ["us-west-2", "us-west-1", "us-east-1"],
+        }
+    )
 
     DEFAULT_SHOTS_QPU = 1000
     DEFAULT_SHOTS_SIMULATOR = 0
     DEFAULT_RESULTS_POLL_TIMEOUT = 432000
     DEFAULT_RESULTS_POLL_INTERVAL = 1
 
-    def __init__(self, arn: str, aws_session=None):
+    def __init__(self, arn: str, aws_session: Optional[AwsSession] = None):
         """
         Args:
             arn (str): The ARN of the device
@@ -64,8 +70,8 @@ def __init__(self, arn: str, aws_session=None):
             physically located. When this occurs, a cloned `aws_session` is created for the Region
             the QPU is located in.
 
-            See `braket.aws.aws_device.AwsQpu.QPU_REGIONS` for the AWS Regions the QPUs are located
-            in.
+            See `braket.aws.aws_device.AwsDevice.DEVICE_REGIONS` for the AWS Regions provider
+            devices are located in.
         """
         super().__init__(name=None, status=None)
         self._arn = arn
@@ -80,8 +86,8 @@ def run(
         task_specification: Union[Circuit, Problem],
         s3_destination_folder: AwsSession.S3DestinationFolder,
         shots: Optional[int] = None,
-        poll_timeout_seconds: int = DEFAULT_RESULTS_POLL_TIMEOUT,
-        poll_interval_seconds: int = DEFAULT_RESULTS_POLL_INTERVAL,
+        poll_timeout_seconds: Optional[int] = DEFAULT_RESULTS_POLL_TIMEOUT,
+        poll_interval_seconds: Optional[int] = DEFAULT_RESULTS_POLL_INTERVAL,
         *aws_quantum_task_args,
         **aws_quantum_task_kwargs,
     ) -> AwsQuantumTask:
@@ -230,38 +236,42 @@ def _construct_topology_graph(self) -> Graph:
             return None
 
     @staticmethod
-    def _aws_session_for_device(device_arn: str, aws_session: AwsSession) -> AwsSession:
+    def _aws_session_for_device(device_arn: str, aws_session: Optional[AwsSession]) -> AwsSession:
         """AwsSession: Returns an AwsSession for the device ARN. """
         if "qpu" in device_arn:
             return AwsDevice._aws_session_for_qpu(device_arn, aws_session)
         else:
             return aws_session or AwsSession()
 
     @staticmethod
-    def _aws_session_for_qpu(device_arn: str, aws_session: AwsSession) -> AwsSession:
+    def _aws_session_for_qpu(device_arn: str, aws_session: Optional[AwsSession]) -> AwsSession:
         """
         Get an AwsSession for the device ARN. QPUs are physically located in specific AWS Regions.
         The AWS sessions should connect to the Region that the QPU is located in.
 
-        See `braket.aws.aws_qpu.AwsDevice.QPU_REGIONS` for the AWS Regions the QPUs are located in.
+        See `braket.aws.aws_qpu.AwsDevice.DEVICE_REGIONS` for the
+        AWS Regions the devices are located in.
         """
         region_key = device_arn.split("/")[-2]
-        qpu_regions = AwsDevice.QPU_REGIONS.get(region_key, [])
+        qpu_regions = AwsDevice.DEVICE_REGIONS.get(region_key, [])
+        return AwsDevice._copy_aws_session(aws_session, qpu_regions)
 
+    @staticmethod
+    def _copy_aws_session(aws_session: Optional[AwsSession], regions: List[str]) -> AwsSession:
         if aws_session:
-            if aws_session.boto_session.region_name in qpu_regions:
+            if aws_session.boto_session.region_name in regions:
                 return aws_session
             else:
                 creds = aws_session.boto_session.get_credentials()
                 boto_session = boto3.Session(
                     aws_access_key_id=creds.access_key,
                     aws_secret_access_key=creds.secret_key,
                     aws_session_token=creds.token,
-                    region_name=qpu_regions[0],
+                    region_name=regions[0],
                 )
                 return AwsSession(boto_session=boto_session)
         else:
-            boto_session = boto3.Session(region_name=qpu_regions[0])
+            boto_session = boto3.Session(region_name=regions[0])
             return AwsSession(boto_session=boto_session)
 
     def __repr__(self):
@@ -271,3 +281,83 @@ def __eq__(self, other):
         if isinstance(other, AwsDevice):
             return self.arn == other.arn
         return NotImplemented
+
+    @staticmethod
+    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`
+            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 = None.
+
+        Returns:
+            List[AwsDevice]: list of AWS devices
+        """
+        order_by_list = ["arn", "name", "type", "provider_name", "status"]
+        if order_by not in order_by_list:
+            raise ValueError(f"order_by '{order_by}' must be in {order_by_list}")
+        device_regions_set = AwsDevice._get_devices_regions_set(arns, provider_names, types)
+        results = []
+        for region in device_regions_set:
+            aws_session = AwsDevice._copy_aws_session(aws_session, [region])
+            results.extend(
+                aws_session.search_devices(
+                    arns=arns,
+                    names=names,
+                    types=types,
+                    statuses=statuses,
+                    provider_names=provider_names,
+                )
+            )
+        arns = set([result["deviceArn"] for result in results])
+        devices = [AwsDevice(arn, aws_session) for arn in arns]
+        devices.sort(key=lambda x: getattr(x, order_by))
+        return devices
+
+    @staticmethod
+    def _get_devices_regions_set(
+        arns: Optional[List[str]],
+        provider_names: Optional[List[str]],
+        types: Optional[List[AwsDeviceType]],
+    ) -> Set[str]:
+        """Get the set of regions to call `SearchDevices` API given filters"""
+        device_regions_set = set(
+            AwsDevice.DEVICE_REGIONS[key][0] for key in AwsDevice.DEVICE_REGIONS
+        )
+        if provider_names:
+            provider_region_set = set()
+            for provider in provider_names:
+                for key in AwsDevice.DEVICE_REGIONS:
+                    if key in provider.lower():
+                        provider_region_set.add(AwsDevice.DEVICE_REGIONS[key][0])
+                        break
+            device_regions_set = device_regions_set.intersection(provider_region_set)
+        if arns:
+            arns_region_set = set([AwsDevice.DEVICE_REGIONS[arn.split("/")[-2]][0] for arn in arns])
+            device_regions_set = device_regions_set.intersection(arns_region_set)
+        if types and types == [AwsDeviceType.SIMULATOR]:
+            device_regions_set = device_regions_set.intersection(
+                [AwsDevice.DEVICE_REGIONS["amazon"][0]]
+            )
+        return device_regions_set

src/braket/aws/aws_session.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from typing import Any, Dict, NamedTuple
+from typing import Any, Dict, List, NamedTuple, Optional
 
 import backoff
 import boto3
@@ -113,6 +113,47 @@ def get_device(self, arn: str) -> Dict[str, Any]:
             arn (str): The ARN of the device
 
         Returns:
-            Dict[str, Any]: Device metadata
+            Dict[str, Any]: The response from the Amazon Braket `GetDevice` operation.
         """
         return self.braket_client.get_device(deviceArn=arn)
+
+    def search_devices(
+        self,
+        arns: Optional[List[str]] = None,
+        names: Optional[List[str]] = None,
+        types: Optional[List[str]] = None,
+        statuses: Optional[List[str]] = None,
+        provider_names: Optional[List[str]] = None,
+    ):
+        """
+        Get devices based on filters. The result is the AND of
+        all the filters `arns`, `names`, `types`, `statuses`, `provider_names`.
+
+        Args:
+            arns (List[str], optional): device ARN list, default is `None`
+            names (List[str], optional): device name list, default is `None`
+            types (List[str], optional): device type list, default is `None`
+            statuses (List[str], optional): device status list, default is `None`
+            provider_names (List[str], optional): provider name list, default is `None`
+
+        Returns:
+            List[Dict[str, Any]: The response from the Amazon Braket `SearchDevices` operation.
+        """
+        filters = []
+        if arns:
+            filters.append({"name": "deviceArn", "values": arns})
+        paginator = self.braket_client.get_paginator("search_devices")
+        page_iterator = paginator.paginate(filters=filters, PaginationConfig={"MaxItems": 100})
+        results = []
+        for page in page_iterator:
+            for result in page["devices"]:
+                if names and result["deviceName"] not in names:
+                    continue
+                if types and result["deviceType"] not in types:
+                    continue
+                if statuses and result["deviceStatus"] not in statuses:
+                    continue
+                if provider_names and result["providerName"] not in provider_names:
+                    continue
+                results.append(result)
+        return results

test/integ_tests/test_device_creation.py

@@ -36,3 +36,27 @@ def test_device_across_regions(aws_session):
     # assert QPUs across different regions can be created using the same aws_session
     AwsDevice(RIGETTI_ARN, aws_session)
     AwsDevice(IONQ_ARN, aws_session)
+
+
+@pytest.mark.parametrize("arn", [(RIGETTI_ARN), (IONQ_ARN), (DWAVE_ARN), (SIMULATOR_ARN)])
+def test_get_devices_arn(arn):
+    results = AwsDevice.get_devices(arns=[arn])
+    assert results[0].arn == arn
+
+
+def test_get_devices_others():
+    provider_names = ["Amazon Braket"]
+    types = ["SIMULATOR"]
+    statuses = ["ONLINE"]
+    results = AwsDevice.get_devices(provider_names=provider_names, types=types, statuses=statuses)
+    assert results
+    for result in results:
+        assert result.provider_name in provider_names
+        assert result.type in types
+        assert result.status in statuses
+
+
+def test_get_devices_all():
+    result_arns = [result.arn for result in AwsDevice.get_devices()]
+    for arn in [DWAVE_ARN, RIGETTI_ARN, IONQ_ARN, SIMULATOR_ARN]:
+        assert arn in result_arns