change: handling capability upgrade issues (#1075)

Author: krneta

src/braket/aws/aws_device.py

@@ -22,6 +22,7 @@
 from enum import Enum
 from typing import Any, ClassVar, Optional
 
+import pydantic
 from botocore.errorfactory import ClientError
 from networkx import DiGraph, complete_graph, from_edgelist
 
@@ -379,13 +380,21 @@ def _populate_properties(self, session: AwsSession) -> None:
         self._status = metadata.get("deviceStatus")
         self._type = AwsDeviceType(metadata.get("deviceType"))
         self._provider_name = metadata.get("providerName")
-        self._properties = BraketSchemaBase.parse_raw_schema(metadata.get("deviceCapabilities"))
-        device_poll_interval = self._properties.service.getTaskPollIntervalMillis
-        self._poll_interval_seconds = (
-            device_poll_interval / 1000.0
-            if device_poll_interval
-            else AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL
-        )
+        try:
+            self._properties = BraketSchemaBase.parse_raw_schema(metadata.get("deviceCapabilities"))
+            device_poll_interval = self._properties.service.getTaskPollIntervalMillis
+            self._poll_interval_seconds = (
+                device_poll_interval / 1000.0
+                if device_poll_interval
+                else AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL
+            )
+        except (pydantic.v1.ValidationError, pydantic.ValidationError):
+            warnings.warn(
+                f"Unable to determine device capabilities for '{self._arn}'."
+                " Please make sure you are using the latest version of amazon-braket-schemas.",
+                stacklevel=1,
+            )
+            self._poll_interval_seconds = AwsQuantumTask.DEFAULT_RESULTS_POLL_INTERVAL
         self._topology_graph = None
         self._frames = None
         self._ports = None

test/unit_tests/braket/aws/test_aws_device.py

@@ -466,6 +466,14 @@ def test_gate_model_sim_schema():
     "tasks",
 )
 
+MOCK_GATE_MODEL_INVALID_CAPABILITIES_QPU = {
+    "deviceName": "Aspen-10",
+    "deviceType": "QPU",
+    "providerName": "Rigetti",
+    "deviceStatus": "OFFLINE",
+    "deviceCapabilities": {},
+}
+
 
 @pytest.fixture(
     params=[
@@ -677,6 +685,15 @@ def test_device_refresh_metadata(arn):
     _assert_device_fields(device, MOCK_GATE_MODEL_QPU_CAPABILITIES_2, MOCK_GATE_MODEL_QPU_2)
 
 
+def test_get_device_invalid_capabilities(arn):
+    mock_session = Mock()
+    mock_session.get_device.return_value = MOCK_GATE_MODEL_INVALID_CAPABILITIES_QPU
+    mock_session.region = RIGETTI_REGION
+    device = AwsDevice(arn, mock_session)
+    assert device is not None
+    assert device.properties is None
+
+
 MOCK_PULSE_MODEL_QPU_PULSE_CAPABILITIES_JSON_1 = {
     "braketSchemaHeader": {
         "name": "braket.device_schema.pulse.pulse_device_action_properties",