[Testing] Let's test the local backends (#60)

Author: Yoric

tests/conftest.py

@@ -0,0 +1,24 @@
+from time import sleep
+import torch_geometric.datasets as pyg_dataset
+
+
+def preload_dataset() -> None:
+    # Attempt to force-download the PTC_FM dataset, which has difficulties on some
+    # platforms. We suspect that it's a race condition somewhere in pytorch geometric.
+    #
+    # Sometimes, this raises FileNotFoundError, sometimes RuntimeError, reinforcing
+    # the suggestion that something is happening in the background.
+    exn: FileNotFoundError | RuntimeError | None = None
+    for i in range(0, 10):
+        sleep(i * i)
+        try:
+            print(f"Attempt {i+1} to download dataset")
+            pyg_dataset.TUDataset(root="dataset", name="PTC_FM")
+            print(f"Attempt {i+1} to download dataset succeeded")
+            exn = None
+            break
+        except (FileNotFoundError, RuntimeError) as e:
+            print(f"Attempt {i+1} to download failed: {e}")
+            exn = e
+    if exn is not None:
+        raise exn

tests/test_backends.py

@@ -0,0 +1,61 @@
+from typing import cast
+
+import os
+import pulser as pl
+import pytest
+import conftest
+import torch_geometric.data as pyg_data
+import torch_geometric.datasets as pyg_dataset
+from qek.backends import CompilationError, QutipBackend, BaseBackend
+import qek.data.graphs as qek_graphs
+
+if os.name == "posix":
+    # As of this writing, emu-mps only works under Unix.
+    from qek.backends import EmuMPSBackend
+
+
+@pytest.mark.asyncio
+async def test_async_emulators() -> None:
+    """
+    Test that backends based on emulators can execute without exploding (async).
+    """
+    conftest.preload_dataset()
+
+    # Load dataset
+    original_ptcfm_data = [
+        cast(pyg_data.Data, d) for d in pyg_dataset.TUDataset(root="dataset", name="PTC_FM")
+    ]
+
+    compiled: list[tuple[qek_graphs.BaseGraph, pl.Register, pl.Pulse]] = []
+    for i, data in enumerate(original_ptcfm_data):
+        graph = qek_graphs.PTCFMGraph(data=data, device=pl.AnalogDevice, id=i)
+        try:
+            register = graph.compile_register()
+            pulse = graph.compile_pulse()
+            if len(register.qubits) >= 5:
+                # This will be too slow to execute, skip.
+                continue
+        except CompilationError:
+            # Let's just skip graphs that cannot be computed.
+            continue
+        compiled.append((graph, register, pulse))
+        if len(compiled) >= 5:
+            # We only need a few samples.
+            break
+
+    assert len(compiled) >= 5
+
+    backends: list[BaseBackend] = [QutipBackend(pl.AnalogDevice)]
+    if os.name == "posix":
+        backends.append(EmuMPSBackend(pl.AnalogDevice))
+
+    for backend in backends:
+        for g, register, pulse in compiled:
+            result = await backend.run(register, pulse)
+            # The only thing we can test from the result is that the values make _some_ kind of sense.
+            assert isinstance(result, dict)
+            for k, v in result.items():
+                assert isinstance(k, str)
+                assert v >= 0
+                for c in k:
+                    assert c in {"0", "1"}

tests/test_extractors.py

@@ -4,6 +4,7 @@
 import pytest
 import torch_geometric.data as pyg_data
 import torch_geometric.datasets as pyg_dataset
+import conftest
 from qek.data.extractors import QutipExtractor, BaseExtracted
 
 if os.name == "posix":
@@ -16,8 +17,10 @@
 @pytest.mark.asyncio
 async def test_async_emulators() -> None:
     """
-    Test that emulators can execute without exploding (both sync and async).
+    Test that extractors emulators can execute without exploding (both sync and async).
     """
+    conftest.preload_dataset()
+
     # Load dataset
     original_ptcfm_data = [
         cast(pyg_data.Data, d) for d in pyg_dataset.TUDataset(root="dataset", name="PTC_FM")

tests/test_graphs.py

@@ -6,6 +6,7 @@
 import torch_geometric.datasets as pyg_dataset
 import torch_geometric.utils as pyg_utils
 from torch_geometric.data import Data
+import conftest
 
 from qek.data.graphs import (
     BaseGraph,
@@ -18,6 +19,8 @@
 
 
 def test_graph_init() -> None:
+    conftest.preload_dataset()
+
     # Load dataset
     original_ptcfm_data = pyg_dataset.TUDataset(root="dataset", name="PTC_FM")