update: default no longer returning RETIRED devices from get_devices (#…

…796)
amazon-braket · Nov 15, 2023 · f1a549a · f1a549a
1 parent 7f28109
commit f1a549a
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 10 deletions.
diff --git a/src/braket/aws/aws_device.py b/src/braket/aws/aws_device.py
@@ -564,13 +564,15 @@ def get_devices(
             >>> AwsDevice.get_devices(types=['SIMULATOR'])
 
         Args:
-            arns (Optional[list[str]]): device ARN list, default is `None`
-            names (Optional[list[str]]): device name list, default is `None`
-            types (Optional[list[AwsDeviceType]]): device type list, default is `None`
+            arns (Optional[list[str]]): device ARN filter, default is `None`
+            names (Optional[list[str]]): device name filter, default is `None`
+            types (Optional[list[AwsDeviceType]]): device type filter, default is `None`
                 QPUs will be searched for all regions and simulators will only be
                 searched for the region of the current session.
-            statuses (Optional[list[str]]): device status list, default is `None`
-            provider_names (Optional[list[str]]): provider name list, default is `None`
+            statuses (Optional[list[str]]): device status filter, default is `None`. When `None`
+                is used, RETIRED devices will not be returned. To include RETIRED devices in
+                the results, use a filter that includes "RETIRED" for this parameter.
+            provider_names (Optional[list[str]]): provider name filter, default is `None`
             order_by (str): field to order result by, default is `name`.
                 Accepted values are ['arn', 'name', 'type', 'provider_name', 'status']
             aws_session (Optional[AwsSession]): An AWS session object.

diff --git a/src/braket/aws/aws_session.py b/src/braket/aws/aws_session.py
@@ -622,10 +622,12 @@ def search_devices(
         all the filters `arns`, `names`, `types`, `statuses`, `provider_names`.
 
         Args:
-            arns (Optional[list[str]]): device ARN list, default is `None`.
-            names (Optional[list[str]]): device name list, default is `None`.
-            types (Optional[list[str]]): device type list, default is `None`.
-            statuses (Optional[list[str]]): device status list, default is `None`.
+            arns (Optional[list[str]]): device ARN filter, default is `None`.
+            names (Optional[list[str]]): device name filter, default is `None`.
+            types (Optional[list[str]]): device type filter, default is `None`.
+            statuses (Optional[list[str]]): device status filter, default is `None`. When `None`
+                is used, RETIRED devices will not be returned. To include RETIRED devices in
+                the results, use a filter that includes "RETIRED" for this parameter.
             provider_names (Optional[list[str]]): provider name list, default is `None`.
 
         Returns:
@@ -645,6 +647,8 @@ def search_devices(
                     continue
                 if statuses and result["deviceStatus"] not in statuses:
                     continue
+                if statuses is None and result["deviceStatus"] == "RETIRED":
+                    continue
                 if provider_names and result["providerName"] not in provider_names:
                     continue
                 results.append(result)

diff --git a/test/integ_tests/test_device_creation.py b/test/integ_tests/test_device_creation.py
@@ -18,7 +18,7 @@
 from braket.aws import AwsDevice
 from braket.devices import Devices
 
-RIGETTI_ARN = "arn:aws:braket:::device/qpu/rigetti/Aspen-10"
+RIGETTI_ARN = "arn:aws:braket:us-west-1::device/qpu/rigetti/Aspen-M-3"
 IONQ_ARN = "arn:aws:braket:us-east-1::device/qpu/ionq/Harmony"
 SIMULATOR_ARN = "arn:aws:braket:::device/quantum-simulator/amazon/sv1"
 OQC_ARN = "arn:aws:braket:eu-west-2::device/qpu/oqc/Lucy"

diff --git a/test/unit_tests/braket/aws/test_aws_session.py b/test/unit_tests/braket/aws/test_aws_session.py
@@ -673,6 +673,18 @@ def test_cancel_job_surfaces_errors(exception_type, aws_session):
                 },
             ],
         ),
+        (
+            {"statuses": ["RETIRED"]},
+            [
+                {
+                    "deviceArn": "arn4",
+                    "deviceName": "name4",
+                    "deviceType": "QPU",
+                    "deviceStatus": "RETIRED",
+                    "providerName": "pname3",
+                },
+            ],
+        ),
         (
             {"provider_names": ["pname2"]},
             [
@@ -745,6 +757,13 @@ def test_search_devices(input, output, aws_session):
                     "deviceStatus": "ONLINE",
                     "providerName": "pname2",
                 },
+                {
+                    "deviceArn": "arn4",
+                    "deviceName": "name4",
+                    "deviceType": "QPU",
+                    "deviceStatus": "RETIRED",
+                    "providerName": "pname3",
+                },
             ]
         }
     ]