diff --git a/src/braket/device_schema/device_service_properties_v1.py b/src/braket/device_schema/device_service_properties_v1.py index 3633dca..096ba8a 100644 --- a/src/braket/device_schema/device_service_properties_v1.py +++ b/src/braket/device_schema/device_service_properties_v1.py @@ -73,6 +73,9 @@ class DeviceServiceProperties(BraketSchemaBase): executionWindows (list[DeviceExecutionWindow]): List of the execution windows, it tells us which days the device can execute a task. shotsRange (tuple[int, int]): range of the shots for a given device. + reservationShotsRange (tuple[int, int] | None): per-circuit shots + range that overrides ``shotsRange`` when the task is submitted via + a Braket Direct reservation. Defaults to ``None``. deviceCost (Optional[DeviceCost]): cost of the device to run the quantum circuits deviceDocumentation (Optional[DeviceDocumentation]): provides device specific details like image, summary etc. @@ -119,6 +122,7 @@ class DeviceServiceProperties(BraketSchemaBase): braketSchemaHeader: BraketSchemaHeader = Field(default=_PROGRAM_HEADER, const=_PROGRAM_HEADER) executionWindows: list[DeviceExecutionWindow] shotsRange: tuple[int, int] + reservationShotsRange: tuple[int, int] | None = None deviceCost: DeviceCost | None deviceDocumentation: DeviceDocumentation | None deviceLocation: str | None diff --git a/src/braket/device_schema/ionq/ionq_device_capabilities_v1.py b/src/braket/device_schema/ionq/ionq_device_capabilities_v1.py index 662e613..7dfe65b 100644 --- a/src/braket/device_schema/ionq/ionq_device_capabilities_v1.py +++ b/src/braket/device_schema/ionq/ionq_device_capabilities_v1.py @@ -23,6 +23,9 @@ from braket.device_schema.ionq.ionq_provider_properties_v1 import IonqProviderProperties from braket.device_schema.jaqcd_device_action_properties import JaqcdDeviceActionProperties from braket.device_schema.openqasm_device_action_properties import OpenQASMDeviceActionProperties +from braket.device_schema.openqasm_program_set_device_action_properties import ( + OpenQASMProgramSetDeviceActionProperties, +) from braket.device_schema.standardized_gate_model_qpu_device_properties_v3 import ( StandardizedGateModelQpuDeviceProperties, ) @@ -126,7 +129,7 @@ class IonqDeviceCapabilities(BraketSchemaBase, DeviceCapabilities): braketSchemaHeader: BraketSchemaHeader = Field(default=_PROGRAM_HEADER, const=_PROGRAM_HEADER) action: dict[ DeviceActionType | str, - OpenQASMDeviceActionProperties | JaqcdDeviceActionProperties, + OpenQASMDeviceActionProperties | JaqcdDeviceActionProperties | OpenQASMProgramSetDeviceActionProperties, ] paradigm: GateModelQpuParadigmProperties provider: IonqProviderProperties | None diff --git a/src/braket/device_schema/openqasm_program_set_device_action_properties.py b/src/braket/device_schema/openqasm_program_set_device_action_properties.py index 6ffb649..8b5f00f 100644 --- a/src/braket/device_schema/openqasm_program_set_device_action_properties.py +++ b/src/braket/device_schema/openqasm_program_set_device_action_properties.py @@ -23,8 +23,15 @@ class OpenQASMProgramSetDeviceActionProperties(DeviceActionProperties): Attributes: maximumExecutables(int): The maximum number of executables that can be in a single program set. + maximumTotalShots(int): The maximum allowed total shots across all + executables in a single program set. + minimumTotalShots(int | None): The minimum allowed total shots + across all executables in a single program set. When unset, no + task-level minimum is enforced. Devices opt in by setting this + field to a value > 1. """ actionType: constr(regex=r"^braket\.ir\.openqasm\.program_set$") maximumExecutables: conint(ge=0) maximumTotalShots: conint(ge=0) + minimumTotalShots: conint(ge=1) | None = None diff --git a/test/unit_tests/braket/device_schema/ionq/test_ionq_device_capabilities_v1.py b/test/unit_tests/braket/device_schema/ionq/test_ionq_device_capabilities_v1.py index a956a86..dcef35a 100644 --- a/test/unit_tests/braket/device_schema/ionq/test_ionq_device_capabilities_v1.py +++ b/test/unit_tests/braket/device_schema/ionq/test_ionq_device_capabilities_v1.py @@ -18,6 +18,7 @@ from braket.device_schema.error_mitigation import Debias from braket.device_schema.ionq.ionq_device_capabilities_v1 import IonqDeviceCapabilities +from braket.device_schema.openqasm_program_set_device_action_properties import OpenQASMProgramSetDeviceActionProperties jaqcd_valid_input = { "braketSchemaHeader": { @@ -135,6 +136,22 @@ def test_valid(valid_input): assert result.braketSchemaHeader.name == "braket.device_schema.ionq.ionq_device_capabilities" +def test_valid_program_set(): + valid_input = dict(openqasm_valid_input) + valid_input["action"] = dict(valid_input["action"]) + valid_input["action"]["braket.ir.openqasm.program_set"] = { + "actionType": "braket.ir.openqasm.program_set", + "version": ["1"], + "maximumExecutables": 100, + "maximumTotalShots": 100000, + } + result = IonqDeviceCapabilities.parse_raw_schema(json.dumps(valid_input)) + action = result.action["braket.ir.openqasm.program_set"] + assert isinstance(action, OpenQASMProgramSetDeviceActionProperties) + assert action.maximumExecutables == 100 + assert action.maximumTotalShots == 100000 + + @pytest.mark.parametrize("valid_input", [openqasm_valid_input, jaqcd_valid_input]) def test_valid_provider(valid_input): valid_input = dict(valid_input) diff --git a/test/unit_tests/braket/device_schema/test_device_service_properties_v1.py b/test/unit_tests/braket/device_schema/test_device_service_properties_v1.py index 4c11c17..e102d53 100644 --- a/test/unit_tests/braket/device_schema/test_device_service_properties_v1.py +++ b/test/unit_tests/braket/device_schema/test_device_service_properties_v1.py @@ -82,3 +82,17 @@ def test_parse_valid_input_with_getTaskPollInterval(valid_input_with_getTaskPoll json.dumps(valid_input_with_getTaskPollInterval) ) assert service_props.getTaskPollIntervalMillis == 200 + + +@pytest.mark.parametrize( + "input_value, expected", + [ + (None, None), + ([1, 5000], (1, 5000)), + ], +) +def test_reservation_shots_range(valid_input, input_value, expected): + if input_value is not None: + valid_input["reservationShotsRange"] = input_value + service_props = DeviceServiceProperties.parse_raw_schema(json.dumps(valid_input)) + assert service_props.reservationShotsRange == expected diff --git a/test/unit_tests/braket/device_schema/test_openqasm_program_set_device_action_properties.py b/test/unit_tests/braket/device_schema/test_openqasm_program_set_device_action_properties.py index 8a8d297..570f991 100644 --- a/test/unit_tests/braket/device_schema/test_openqasm_program_set_device_action_properties.py +++ b/test/unit_tests/braket/device_schema/test_openqasm_program_set_device_action_properties.py @@ -57,3 +57,25 @@ def test_missing_field(valid_input, field): def test_negative_field(valid_input, field): valid_input[field] = -1 OpenQASMProgramSetDeviceActionProperties.parse_raw(json.dumps(valid_input)) + + +@pytest.mark.parametrize( + "input_value, expected", + [ + (None, None), + (1, 1), + (100, 100), + ], +) +def test_minimum_total_shots(valid_input, input_value, expected): + if input_value is not None: + valid_input["minimumTotalShots"] = input_value + result = OpenQASMProgramSetDeviceActionProperties.parse_raw(json.dumps(valid_input)) + assert result.minimumTotalShots == expected + + +@pytest.mark.parametrize("invalid_value", [0, -1]) +@pytest.mark.xfail(raises=ValidationError, reason="ensure this value is greater than or equal to 1") +def test_minimum_total_shots_invalid(valid_input, invalid_value): + valid_input["minimumTotalShots"] = invalid_value + OpenQASMProgramSetDeviceActionProperties.parse_raw(json.dumps(valid_input))