Merge pull request #5 from oqc-community/feature/qpu_info_endpoints

CSU-10 Add QPU information endpoints to ZMQ server

.github/CODEOWNERS

@@ -1,3 +1,3 @@
 # For now everyone gets added to all PRs by default. Fine while the PR velocity is small, we'll change if
 # necessary as time goes on.
-* @jfriel-oqc @keriksson-rosenqvist @owen-oqc @hamidelmaazouz @chemix-lunacy @bgsach
+* @keriksson-rosenqvist @owen-oqc @hamidelmaazouz @bgsach @daria-oqc @lcauser-oqc

.gitignore

@@ -13,6 +13,9 @@ logs/
 # PyCharm
 .idea/
 
+# VS Code
+.vscode
+
 # C extensions
 *.so
 

README.md

@@ -12,7 +12,7 @@ The QAT-RPC package has not yet been released to PyPI.
 ### Building from Source
 
 We use [poetry](https://python-poetry.org/) for dependency management and run on
-[Python 3.8+](https://www.python.org/downloads/).
+[Python 3.10+](https://www.python.org/downloads/).
 Once both of these are installed run this in the root folder to install all the dependencies that you need: `poetry install`
 
 > If you are contributing to the project we recommend that you also run

pyproject.toml

@@ -11,9 +11,9 @@ packages = [
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.13"
-qat-compiler = "^2.0.0"
 pyzmq = "^26.1.0"
 prometheus-client="^0.20.0"
+qat-compiler = "^2.3.0"
 
 [tool.poetry.group.dev.dependencies]
 coverage = "^6.3.2"

src/QAT_RPC/qat_rpc/zmq/wrappers.py

@@ -1,16 +1,27 @@
+from enum import Enum
+from importlib.metadata import version
 from time import time
 from typing import Optional, Union
 
 import zmq
+from compiler_config.config import CompilerConfig
 from qat.purr.backends.echo import get_default_echo_hardware
-from qat.purr.compiler.config import CompilerConfig
 from qat.purr.compiler.hardware_models import QuantumHardwareModel
 from qat.purr.compiler.runtime import get_runtime
+from qat.purr.integrations.features import OpenPulseFeatures
 from qat.qat import execute_with_metrics
 
 from qat_rpc.utils.metrics import MetricExporter
 
 
+class Messages(Enum):
+    PROGRAM = "program"
+    VERSION = "version"
+    COUPLINGS = "couplings"
+    QUBIT_INFO = "qubit_info"
+    QPU_INFO = "qpu_info"
+
+
 class ZMQBase:
     def __init__(self, socket_type: zmq.SocketType):
         self._context = zmq.Context()
@@ -38,7 +49,7 @@ def _send(self, message) -> None:
             try:
                 self._socket.send_pyobj(message, zmq.NOBLOCK)
                 sent = True
-            except zmq.ZMQError as e:
+            except zmq.ZMQError:
                 if time() > t0 + self._timeout:
                     raise TimeoutError(
                         "Sending %s on %s timedout" % (message, self.address)
@@ -73,6 +84,53 @@ def __init__(
     def address(self):
         return f"{self._protocol}://*:{self._port}"
 
+    def _program(self, program, config):
+        program = program
+        config = CompilerConfig.create_from_json(config)
+        result, metrics = execute_with_metrics(program, self._engine, config)
+        return {"results": result, "execution_metrics": metrics}
+
+    def _version(self):
+        return {"qat_rpc_version": str(version("qat_rpc"))}
+
+    def _couplings(self):
+        coupling = [
+            coupled.direction for coupled in self._hardware.qubit_direction_couplings
+        ]
+        return {"couplings": coupling}
+
+    def _qubit_info(self):
+        raise NotImplementedError(
+            "Individual qubit information not implented, pending hardware model changes."
+        )
+
+    def _qpu_info(self):
+        features = OpenPulseFeatures()
+        features.for_hardware(self._hardware)
+        qpu_info = features.to_json_dict()
+        return {"qpu_info": qpu_info}
+
+    def _interpret_message(self, message):
+        match message[0]:
+            case Messages.PROGRAM.value:
+                print(message)
+                if len(message) != 3:
+                    raise ValueError(
+                        f"Program message should be of length 3, not {len(message)}"
+                    )
+                return self._program(message[1], message[2])
+            case Messages.VERSION.value:
+                return self._version()
+            case Messages.COUPLINGS.value:
+                return self._couplings()
+            case Messages.QUBIT_INFO:
+                return self._qubit_info()
+            case Messages.QPU_INFO.value:
+                return self._qpu_info()
+
+            case _:
+                return self._program(message[0], message[1])
+
     def run(self):
         self._running = True
         with self._metric.receiver_status() as metric:
@@ -81,10 +139,7 @@ def run(self):
             msg = self._check_recieved()
             if msg is not None:
                 try:
-                    program = msg[0]
-                    config = CompilerConfig.create_from_json(msg[1])
-                    result, metrics = execute_with_metrics(program, self._engine, config)
-                    reply = {"results": result, "execution_metrics": metrics}
+                    reply = self._interpret_message(message=msg)
                     with self._metric.executed_messages() as executed:
                         executed.increment()
                 except Exception as e:
@@ -105,17 +160,32 @@ def __init__(self):
         self._socket.setsockopt(zmq.LINGER, 0)
         self._socket.connect(self.address)
 
-    def execute_task(self, program: str, config: Union[CompilerConfig, str] = None):
-        self.result = None
-        if isinstance(config, str):
-            # Verify config string is valid before submitting.
-            config = CompilerConfig.create_from_json(config)
-        cfg = config or CompilerConfig()
-        self._send((program, cfg.to_json()))
+    def _send(self, message):
+        super()._send(message=message)
         return self._await_results()
 
     def _await_results(self):
         result = None
         while result is None:
             result = self._check_recieved()
         return result
+
+    def execute_task(self, program: str, config: Union[CompilerConfig, str] = None):
+        self.result = None
+        if isinstance(config, str):
+            # Verify config string is valid before submitting.
+            config = CompilerConfig.create_from_json(config)
+        cfg = config or CompilerConfig()
+        return self._send((Messages.PROGRAM.value, program, cfg.to_json()))
+
+    def api_version(self):
+        return self._send((Messages.VERSION.value,))
+
+    def qpu_couplings(self):
+        return self._send((Messages.COUPLINGS.value,))
+
+    def qubit_info(self):
+        return self._send((Messages.QUBIT_INFO.value,))
+
+    def qpu_info(self):
+        return self._send((Messages.QPU_INFO.value,))

src/tests/zmq/test_wrappers.py

@@ -1,17 +1,28 @@
 import threading
+from importlib.metadata import version
 
 import pytest
-from qat.purr.compiler.config import CompilerConfig
+from compiler_config.config import CompilerConfig
+from qat.purr.backends.echo import (
+    add_direction_couplings_to_hardware,
+    get_default_echo_hardware,
+)
 
 from qat_rpc.utils.constants import PROMETHEUS_PORT
 from qat_rpc.utils.metrics import MetricExporter, PrometheusReceiver
 from qat_rpc.zmq.wrappers import ZMQClient, ZMQServer
 
+qubit_count = 8
+qpu_couplings = [(i, j) for i in range(qubit_count) for j in range(qubit_count) if i != j]
+
 
 @pytest.fixture(scope="module", autouse=True)
 def server():
+    hardware = get_default_echo_hardware(qubit_count=qubit_count)
+    hardware = add_direction_couplings_to_hardware(hardware, qpu_couplings)
     server = ZMQServer(
-        metric_exporter=MetricExporter(backend=PrometheusReceiver(PROMETHEUS_PORT))
+        hardware=hardware,
+        metric_exporter=MetricExporter(backend=PrometheusReceiver(PROMETHEUS_PORT)),
     )
     server_thread = threading.Thread(target=server.run, daemon=True)
     server_thread.start()
@@ -27,17 +38,6 @@ def server():
 """
 
 
-def test_zmq_flow():
-    client = ZMQClient()
-
-    config = CompilerConfig()
-    config.results_format.binary_count()
-    config.repeats = 100
-
-    response = client.execute_task(program, config)
-    assert response["results"]["c"]["00"] == 100
-
-
 def test_zmq_exception():
     client = ZMQClient()
 
@@ -84,3 +84,50 @@ def test_two_zmq_clients():
     thread00.join()
     thread01.join()
     thread10.join()
+
+
+def test_program():
+    client = ZMQClient()
+
+    config = CompilerConfig()
+    config.results_format.binary_count()
+    config.repeats = 100
+
+    response = client.execute_task(program, config)
+    assert response["results"]["c"]["00"] == 100
+
+
+def test_program_backwards_compatible():
+    client = ZMQClient()
+
+    config = CompilerConfig()
+    config.results_format.binary_count()
+    config.repeats = 100
+
+    response = client._send((program, config.to_json()))
+    print(response)
+    assert response["results"]["c"]["00"] == 100
+
+
+def test_api_version():
+    client = ZMQClient()
+    api_version = client.api_version()
+    assert api_version["qat_rpc_version"] == version("qat_rpc")
+
+
+def test_couplings():
+    client = ZMQClient()
+    couplings = client.qpu_couplings()
+    assert couplings["couplings"] == qpu_couplings
+
+
+def test_qubit_info():
+    client = ZMQClient()
+    qubit_info = client.qubit_info()
+    assert qubit_info["Exception"] is not None
+
+
+def test_qpu_info():
+    client = ZMQClient()
+    qpu_info = client.qpu_info()
+    assert qpu_info["qpu_info"] is not None