Merge branch 'main' into fix_search_device

.github/workflows/check-format.yml

@@ -16,7 +16,7 @@ jobs:
   check-code-format:
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0
+    - uses: actions/checkout@8ade135a41bc03ea155e62e844d188df1ea18608 # v4.1.0
     - name: Set up Python
       uses: actions/setup-python@61a6322f88396a6271a6ee3565807d608ecaddd1 # v4.7.0
       with:

.github/workflows/dependent-tests.yml

@@ -21,7 +21,7 @@ jobs:
           - amazon-braket-pennylane-plugin-python
 
     steps:
-    - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0
+    - uses: actions/checkout@8ade135a41bc03ea155e62e844d188df1ea18608 # v4.1.0
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@61a6322f88396a6271a6ee3565807d608ecaddd1 # v4.7.0
       with:

.github/workflows/publish-to-pypi.yml

@@ -12,7 +12,7 @@ jobs:
     name: Build and publish distribution to PyPi
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0
+      - uses: actions/checkout@8ade135a41bc03ea155e62e844d188df1ea18608 # v4.1.0
       - name: Set up Python
         uses: actions/setup-python@61a6322f88396a6271a6ee3565807d608ecaddd1 # v4.7.0
         with:

.github/workflows/python-package.yml

@@ -24,7 +24,7 @@ jobs:
         python-version: ["3.8", "3.9", "3.10", "3.11"]
 
     steps:
-    - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0
+    - uses: actions/checkout@8ade135a41bc03ea155e62e844d188df1ea18608 # v4.1.0
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@61a6322f88396a6271a6ee3565807d608ecaddd1 # v4.7.0
       with:

.github/workflows/twine-check.yml

@@ -14,7 +14,7 @@ jobs:
     name: Check long description
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@3df4ab11eba7bda6032a0b82a6bb43b11571feac # v4.0.0
+      - uses: actions/checkout@8ade135a41bc03ea155e62e844d188df1ea18608 # v4.1.0
       - name: Set up Python
         uses: actions/setup-python@61a6322f88396a6271a6ee3565807d608ecaddd1 # v4.7.0
         with:

CHANGELOG.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## v1.56.0 (2023-09-26)
+
+### Features
+
+ * add queue visibility information
+
 ## v1.55.1.post0 (2023-09-18)
 
 ### Documentation Changes

setup.py

@@ -33,7 +33,7 @@
         "setuptools",
         "backoff",
         "boltons",
-        "boto3>=1.22.3",
+        "boto3>=1.28.53",
         "nest-asyncio",
         "networkx",
         "numpy",

src/braket/_sdk/_version.py

@@ -15,4 +15,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "1.55.2.dev0"
+__version__ = "1.56.1.dev0"

src/braket/aws/aws_device.py

@@ -30,6 +30,7 @@
 from braket.aws.aws_quantum_task import AwsQuantumTask
 from braket.aws.aws_quantum_task_batch import AwsQuantumTaskBatch
 from braket.aws.aws_session import AwsSession
+from braket.aws.queue_information import QueueDepthInfo, QueueType
 from braket.circuits import Circuit, Gate, QubitSet
 from braket.circuits.gate_calibrations import GateCalibrations
 from braket.device_schema import DeviceCapabilities, ExecutionDay, GateModelQpuParadigmProperties
@@ -677,6 +678,54 @@ def get_device_region(device_arn: str) -> str:
                 "see 'https://docs.aws.amazon.com/braket/latest/developerguide/braket-devices.html'"
             )
 
+    def queue_depth(self) -> QueueDepthInfo:
+        """
+        Task queue depth refers to the total number of quantum tasks currently waiting
+        to run on a particular device.
+
+        Returns:
+            QueueDepthInfo: Instance of the QueueDepth class representing queue depth
+            information for quantum tasks and hybrid jobs.
+            Queue depth refers to the number of quantum tasks and hybrid jobs queued on a particular
+            device. The normal tasks refers to the quantum tasks not submitted via Hybrid Jobs.
+            Whereas, the priority tasks refers to the total number of quantum tasks waiting to run
+            submitted through Amazon Braket Hybrid Jobs. These tasks run before the normal tasks.
+            If the queue depth for normal or priority quantum tasks is greater than 4000, we display
+            their respective queue depth as '>4000'. Similarly, for hybrid jobs if there are more
+            than 1000 jobs queued on a device, display the hybrid jobs queue depth as '>1000'.
+            Additionally, for QPUs if hybrid jobs queue depth is 0, we display information about
+            priority and count of the running hybrid job.
+
+        Example:
+            Queue depth information for a running job.
+            >>> device = AwsDevice(Device.Amazon.SV1)
+            >>> print(device.queue_depth())
+            QueueDepthInfo(quantum_tasks={<QueueType.NORMAL: 'Normal'>: '0',
+            <QueueType.PRIORITY: 'Priority'>: '1'}, jobs='0 (1 prioritized job(s) running)')
+
+            If more than 4000 quantum tasks queued on a device.
+            >>> device = AwsDevice(Device.Amazon.DM1)
+            >>> print(device.queue_depth())
+            QueueDepthInfo(quantum_tasks={<QueueType.NORMAL: 'Normal'>: '>4000',
+            <QueueType.PRIORITY: 'Priority'>: '2000'}, jobs='100')
+        """
+        metadata = self.aws_session.get_device(arn=self.arn)
+        queue_metadata = metadata.get("deviceQueueInfo")
+        queue_info = {}
+
+        for response in queue_metadata:
+            queue_name = response.get("queue")
+            queue_priority = response.get("queuePriority")
+            queue_size = response.get("queueSize")
+
+            if queue_name == "QUANTUM_TASKS_QUEUE":
+                priority_enum = QueueType(queue_priority)
+                queue_info.setdefault("quantum_tasks", {})[priority_enum] = queue_size
+            else:
+                queue_info["jobs"] = queue_size
+
+        return QueueDepthInfo(**queue_info)
+
     def refresh_gate_calibrations(self) -> Optional[GateCalibrations]:
         """
         Refreshes the gate calibration data upon request.

src/braket/aws/aws_quantum_job.py

@@ -27,6 +27,7 @@
 
 from braket.aws import AwsDevice
 from braket.aws.aws_session import AwsSession
+from braket.aws.queue_information import HybridJobQueueInfo
 from braket.jobs import logs
 from braket.jobs.config import (
     CheckpointConfig,
@@ -278,6 +279,38 @@ def state(self, use_cached_value: bool = False) -> str:
         """
         return self.metadata(use_cached_value).get("status")
 
+    def queue_position(self) -> HybridJobQueueInfo:
+        """
+        The queue position details for the hybrid job.
+
+        Returns:
+            HybridJobQueueInfo: Instance of HybridJobQueueInfo class representing
+            the queue position information for the hybrid job. The queue_position is
+            only returned when the hybrid job is not in RUNNING/CANCELLING/TERMINAL states,
+            else queue_position is returned as None. If the queue position of the hybrid
+            job is greater than 15, we return '>15' as the queue_position return value.
+
+        Examples:
+            job status = QUEUED and position is 2 in the queue.
+            >>> job.queue_position()
+            HybridJobQueueInfo(queue_position='2', message=None)
+
+            job status = QUEUED and position is 18 in the queue.
+            >>> job.queue_position()
+            HybridJobQueueInfo(queue_position='>15', message=None)
+
+            job status = COMPLETED
+            >>> job.queue_position()
+            HybridJobQueueInfo(queue_position=None,
+            message='Job is in COMPLETED status. AmazonBraket does
+                        not show queue position for this status.')
+        """
+        response = self.metadata()["queueInfo"]
+        queue_position = None if response.get("position") == "None" else response.get("position")
+        message = response.get("message")
+
+        return HybridJobQueueInfo(queue_position=queue_position, message=message)
+
     def logs(self, wait: bool = False, poll_interval_seconds: int = 5) -> None:
         """Display logs for a given hybrid job, optionally tailing them until hybrid job is
            complete.

src/braket/aws/aws_quantum_task.py

@@ -24,6 +24,7 @@
 from braket.ahs.analog_hamiltonian_simulation import AnalogHamiltonianSimulation
 from braket.annealing.problem import Problem
 from braket.aws.aws_session import AwsSession
+from braket.aws.queue_information import QuantumTaskQueueInfo, QueueType
 from braket.circuits import Instruction
 from braket.circuits.circuit import Circuit, Gate, QubitSet
 from braket.circuits.circuit_helpers import validate_circuit_and_shots
@@ -314,6 +315,40 @@ def state(self, use_cached_value: bool = False) -> str:
         """
         return self._status(use_cached_value)
 
+    def queue_position(self) -> QuantumTaskQueueInfo:
+        """
+        The queue position details for the quantum task.
+
+        Returns:
+            QuantumTaskQueueInfo: Instance of QuantumTaskQueueInfo class
+            representing the queue position information for the quantum task.
+            The queue_position is only returned when quantum task is not in
+            RUNNING/CANCELLING/TERMINAL states, else queue_position is returned as None.
+            The normal tasks refers to the quantum tasks not submitted via Hybrid Jobs.
+            Whereas, the priority tasks refers to the total number of quantum tasks waiting to run
+            submitted through Amazon Braket Hybrid Jobs. These tasks run before the normal tasks.
+            If the queue position for normal or priority quantum tasks is greater than 2000,
+            we display their respective queue position as '>2000'.
+
+        Examples:
+            task status = QUEUED and queue position is 2050
+            >>> task.queue_position()
+            QuantumTaskQueueInfo(queue_type=<QueueType.NORMAL: 'Normal'>,
+            queue_position='>2000', message=None)
+
+            task status = COMPLETED
+            >>> task.queue_position()
+            QuantumTaskQueueInfo(queue_type=<QueueType.NORMAL: 'Normal'>,
+            queue_position=None, message='Task is in COMPLETED status. AmazonBraket does
+            not show queue position for this status.')
+        """
+        response = self.metadata()["queueInfo"]
+        queue_type = QueueType(response["queuePriority"])
+        queue_position = None if response.get("position") == "None" else response.get("position")
+        message = response.get("message")
+
+        return QuantumTaskQueueInfo(queue_type, queue_position, message)
+
     def _status(self, use_cached_value: bool = False) -> str:
         metadata = self.metadata(use_cached_value)
         status = metadata.get("status")

src/braket/aws/aws_session.py

@@ -279,7 +279,9 @@ def get_quantum_task(self, arn: str) -> Dict[str, Any]:
         Returns:
             Dict[str, Any]: The response from the Amazon Braket `GetQuantumTask` operation.
         """
-        response = self.braket_client.get_quantum_task(quantumTaskArn=arn)
+        response = self.braket_client.get_quantum_task(
+            quantumTaskArn=arn, additionalAttributeNames=["QueueInfo"]
+        )
         broadcast_event(_TaskStatusEvent(arn=response["quantumTaskArn"], status=response["status"]))
         return response
 
@@ -324,7 +326,7 @@ def get_job(self, arn: str) -> Dict[str, Any]:
         Returns:
             Dict[str, Any]: The response from the Amazon Braket `GetQuantumJob` operation.
         """
-        return self.braket_client.get_job(jobArn=arn)
+        return self.braket_client.get_job(jobArn=arn, additionalAttributeNames=["QueueInfo"])
 
     def cancel_job(self, arn: str) -> Dict[str, Any]:
         """

src/braket/aws/queue_information.py

@@ -0,0 +1,85 @@
+# Copyright Amazon.com Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Dict, Optional
+
+
+class QueueType(str, Enum):
+    """
+    Enumerates the possible priorities for the queue.
+
+    Values:
+        NORMAL: Represents normal queue for the device.
+        PRIORITY: Represents priority queue for the device.
+    """
+
+    NORMAL = "Normal"
+    PRIORITY = "Priority"
+
+
+@dataclass()
+class QueueDepthInfo:
+    """
+    Represents quantum tasks and hybrid jobs queue depth information.
+
+    Attributes:
+        quantum_tasks (Dict[QueueType, str]): number of quantum tasks waiting
+            to run on a device. This includes both 'Normal' and 'Priority' tasks.
+            For Example, {'quantum_tasks': {QueueType.NORMAL: '7', QueueType.PRIORITY: '3'}}
+        jobs (str): number of hybrid jobs waiting to run on a device. Additionally, for QPUs if
+            hybrid jobs queue depth is 0, we display information about priority and count of the
+            running hybrid jobs. Example, 'jobs': '0 (1 prioritized job(s) running)'
+    """
+
+    quantum_tasks: Dict[QueueType, str]
+    jobs: str
+
+
+@dataclass
+class QuantumTaskQueueInfo:
+    """
+    Represents quantum tasks queue information.
+
+    Attributes:
+        queue_type (QueueType): type of the quantum_task queue either 'Normal'
+            or 'Priority'.
+        queue_position (Optional[str]): current position of your quantum task within a respective
+            device queue. This value can be None based on the state of the task. Default: None.
+        message (Optional[str]): Additional message information. This key is present only
+            if 'queue_position' is None. Default: None.
+    """
+
+    queue_type: QueueType
+    queue_position: Optional[str] = None
+    message: Optional[str] = None
+
+
+@dataclass
+class HybridJobQueueInfo:
+    """
+    Represents hybrid job queue information.
+
+    Attributes:
+        queue_position (Optional[str]): current position of your hybrid job within a respective
+            device queue. If the queue position of the hybrid job is greater than 15, we
+            return '>15' as the queue_position return value. The queue_position is only
+            returned when hybrid job is not in RUNNING/CANCELLING/TERMINAL states, else
+            queue_position is returned as None.
+        message (Optional[str]): Additional message information. This key is present only
+            if 'queue_position' is None. Default: None.
+    """
+
+    queue_position: Optional[str] = None
+    message: Optional[str] = None