[Hack] Let's try more than once to load the dataset (#83)

Author: Yoric

.github/workflows/test.yml

@@ -33,22 +33,6 @@ jobs:
     - name: Install Hatch
       run: |
         pip install hatch
-    - name: Pre-download dataset
-        # On macOS (and only on macOS), we encounter a strange issue:
-        #
-        # - loading the dataset causes a FileNotFoundException to be raised
-        # - checking afterwards, the dataset *is* on the disk
-        # - rerunning the process, the dataset is found.
-        #
-        # The running hypothesis is that pytorch-geometric downloads the
-        # dataset asynchronously, but fails to wait until the download is
-        # complete.
-        #
-        # With this small script, we force the pipeline to wait until the
-        # download is complete. So far, this seems to solve the issue.
-      if: ${{ matrix.os == 'macos-latest' }}
-      run: |
-        hatch -v run before_tests
     - name: Run tests
       run: |
         hatch -v run test
@@ -77,22 +61,6 @@ jobs:
       - name: Install Hatch
         run: |
           pip install hatch
-      - name: Pre-download dataset
-        if: ${{ matrix.os == 'macos-latest' }}
-        # On macOS (and only on macOS), we encounter a strange issue:
-        #
-        # - loading the dataset causes a FileNotFoundException to be raised
-        # - checking afterwards, the dataset *is* on the disk
-        # - rerunning the process, the dataset is found.
-        #
-        # The running hypothesis is that pytorch-geometric downloads the
-        # dataset asynchronously, but fails to wait until the download is
-        # complete.
-        #
-        # With this small script, we force the pipeline to wait until the
-        # download is complete. So far, this seems to solve the issue.
-        run: |
-          hatch -v run before_tests
       - name: Copy samples
         run: |
           cp examples/ptcfm_processed_dataset.json .

examples/tutorial 1 - Using a Quantum Device to Extract Machine-Learning Features.ipynb

@@ -36,8 +36,11 @@
    "source": [
     "import torch_geometric.datasets as pyg_dataset\n",
     "\n",
-    "# Load the original PTC-FM dataset\n",
-    "og_ptcfm = pyg_dataset.TUDataset(root=\"dataset\", name=\"PTC_FM\")\n",
+    "from qek.shared.retrier import PygRetrier\n",
+    "\n",
+    "# Load the original PTC-FM dataset.\n",
+    "# We use PygRetrier to retry the download if it fails.\n",
+    "og_ptcfm = PygRetrier().insist(pyg_dataset.TUDataset, root=\"dataset\", name=\"PTC_FM\")\n",
     "\n",
     "display(\"Loaded %s samples\" % (len(og_ptcfm), ))"
    ]

examples/tutorial 1a - Using a Quantum Device to Extract Machine-Learning Features - low-level.ipynb

@@ -37,7 +37,10 @@
    "source": [
     "# Load the original PTC-FM dataset\n",
     "import torch_geometric.datasets as pyg_dataset\n",
-    "og_ptcfm = pyg_dataset.TUDataset(root=\"dataset\", name=\"PTC_FM\")\n",
+    "from qek.shared.retrier import PygRetrier\n",
+    "\n",
+    "# We use PygRetrier to retry the download if it fails.\n",
+    "og_ptcfm = PygRetrier().insist(pyg_dataset.TUDataset, root=\"dataset\", name=\"PTC_FM\")\n",
     "\n",
     "display(\"Loaded %s samples\" % (len(og_ptcfm), ))"
    ]

qek/shared/retrier.py

@@ -0,0 +1,56 @@
+"""
+Backoff-and-retry utilities.
+"""
+
+import logging
+from time import sleep
+from typing import Any, Final, Type
+
+from torch_geometric.data import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+class PygRetrier:
+    """
+    Our test harness attempts to run tests concurrently, but the pyg dataset loader does not
+    work well with concurrency.
+
+    We work around this by simply retrying the loads a few times, until it succeeds.
+    """
+
+    def __init__(self, max_attempts: int = 3, name: str = "PygRetrier"):
+        """
+        Create a PygRetrier
+
+        Arguments:
+            max_attempts (optional): The max number of attempts to undertake before
+                giving up. Defaults to 3.
+            name (optional): A name to use during logging.
+        """
+        self._max_attempts: Final[int] = max_attempts
+        self.name: Final[str] = name
+
+    def insist(self, callback: Type[Dataset], **kwargs: Any) -> Dataset:
+        """
+        Attempt to call a function or constructor repeatedly until, hopefully,
+        it works.
+        """
+        exn: FileNotFoundError | RuntimeError | OSError | None = None
+        result = None
+        for i in range(self._max_attempts):
+            sleep(i * i)
+            try:
+                logger.debug("%s: attempt %s", self.name, i + 1)
+                result = callback(**kwargs)  # type: ignore
+                logger.debug("%s: attempt %s succeeded", self.name, i + 1)
+                exn = None
+                break
+            except (FileNotFoundError, RuntimeError, OSError) as e:
+                logger.warning("%s: attempt %s failed: %s", self.name, i + 1, e)
+                exn = e
+        if exn is not None:
+            logger.warning("%s: all attempts failed, bailing out", self.name)
+            raise exn
+        assert result is not None
+        return result

tests/conftest.py

@@ -3,11 +3,11 @@
 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
+from qek.shared.retrier import PygRetrier
 
 if os.name == "posix":
     # As of this writing, emu-mps only works under Unix.
@@ -19,11 +19,11 @@ 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")
+        cast(pyg_data.Data, d)
+        for d in PygRetrier().insist(pyg_dataset.TUDataset, root="dataset", name="PTC_FM")
     ]
 
     compiled: list[tuple[qek_graphs.BaseGraph, pl.Register, pl.Pulse]] = []

tests/test_backends.py

@@ -4,8 +4,8 @@
 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
+from qek.shared.retrier import PygRetrier
 
 if os.name == "posix":
     # As of this writing, emu-mps only works under Unix.
@@ -19,11 +19,10 @@ async def test_async_emulators() -> None:
     """
     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")
+        cast(pyg_data.Data, d)
+        for d in PygRetrier().insist(pyg_dataset.TUDataset, root="dataset", name="PTC_FM")
     ]
     MAX_NUMBER_OF_SAMPLES = 5
     MAX_NUMBER_OF_QUBITS = 5

tests/test_extractors.py

@@ -6,7 +6,6 @@
 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,
@@ -16,13 +15,12 @@
     NXGraphCompiler,
     NXWithPos,
 )
+from qek.shared.retrier import PygRetrier
 
 
 def test_graph_init() -> None:
-    conftest.preload_dataset()
-
     # Load dataset
-    original_ptcfm_data = pyg_dataset.TUDataset(root="dataset", name="PTC_FM")
+    original_ptcfm_data = PygRetrier().insist(pyg_dataset.TUDataset, root="dataset", name="PTC_FM")
 
     # Check that `add_graph_coord` doesn't break with this dataset.