[FIX] A few fixes to QPUExtractor (#34)

Author: Yoric

examples/tutorial.ipynb

@@ -398,7 +398,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "train_kernel = kernel.create_train_kernel_matrix(processed_dataset)\n",
+    "train_kernel = kernel.fit_transform(processed_dataset)\n",
     "y_tot = [data.target for data in processed_dataset]"
    ]
   },

qek/data/extractors.py

@@ -5,9 +5,9 @@
 import logging
 from math import ceil
 from typing import Any, Callable, Generic, Sequence, TypeVar, cast
+from uuid import UUID
 import emu_mps
-from pasqal_cloud import SDK
-from pasqal_cloud.batch import Batch
+from pasqal_cloud import SDK, Batch, BatchFilters
 import pulser as pl
 from pulser.devices import Device
 from pulser.json.abstract_repr.deserializer import deserialize_device
@@ -341,7 +341,7 @@ def __init__(
         username: str,
         password: str | None = None,
         device_name: str = "FRESNEL",
-        batch_id: list[str] | None = None,
+        batch_ids: list[str] | None = None,
     ):
         sdk = SDK(username=username, project_id=project_id, password=password)
 
@@ -351,65 +351,79 @@ def __init__(
 
         super().__init__(path=path, device=device, compiler=compiler)
         self._sdk = sdk
-        self._batch_id = batch_id
+        self._batch_ids: list[str] | None = batch_ids
 
     @property
     def batch_ids(self) -> list[str] | None:
-        return self._batch_id
+        return self._batch_ids
 
     async def run(self) -> list[ProcessedData]:
         if len(self.sequences) == 0:
             logger.warning("No sequences to run, did you forget to call compile()?")
             return []
 
-        batches: list[Batch] = []
-        if self._batch_id is None:
+        device: pl.devices.Device = self.sequences[0].sequence.device
+        # The API doesn't support run longer than 500 jobs.
+        # If we want to add more runs, we'll need to split them across several jobs.
+        max_runs = device.max_runs if isinstance(device.max_runs, int) else 500
+
+        if self._batch_ids is None:
             # Enqueue jobs.
             self._batch_ids = []
             for compiled in self.sequences:
-                logger.debug("Executing compiled graph #%s", id)
+                logger.debug("Enqueuing execution of compiled graph #%s", compiled.graph.id)
                 batch = self._sdk.create_batch(
                     compiled.sequence.to_abstract_repr(),
-                    jobs=[{"runs": 1000}],
+                    jobs=[{"runs": max_runs}],
                     wait=False,
                 )
                 logger.info(
                     "Remote execution of compiled graph #%s starting, batched with id %s",
-                    id,
+                    compiled.graph.id,
                     batch.id,
                 )
-                batches.append(batch)
                 self._batch_ids.append(batch.id)
             logger.info(
                 "All %s jobs enqueued for remote execution, with ids %s",
-                len(batches),
+                len(self._batch_ids),
                 self._batch_ids,
             )
-        else:
-            # Get jobs back from the cloud API.
-            for batch_id in self._batch_id:
-                batches.append(self._sdk.get_batch(batch_id))
+        assert len(self._batch_ids) == len(self.sequences)
 
         # Now wait until all batches are complete.
-        waiting = True
-        while waiting:
-            waiting = False
-            for batch in batches:
-                if batch.status in {"PENDING", "RUNNING"}:
-                    # At least one job is pending, let's wait.
-                    await sleep(2)
-                    logger.debug("Job %s is still incomplete")
-                    waiting = True
-
-        logger.info("All jobs complete, %s sequences executed", len(batches))
+        pending_batch_ids: set[str] = set(self._batch_ids)
+        completed_batches: dict[str, Batch] = {}
+
+        while len(pending_batch_ids) > 0:
+            await sleep(delay=2)
+            # We can check up to 100 batches in a single query with the SDK, so let's do that.
+            MAX_BATCH_LEN = 100
+            check_ids: list[str | UUID] = [cast(str | UUID, id) for id in pending_batch_ids][
+                :MAX_BATCH_LEN
+            ]
+            # Update their status.
+            check_batches = self._sdk.get_batches(
+                filters=BatchFilters(id=check_ids)
+            )  # Ideally, this should be async, see https://github.com/pasqal-io/pasqal-cloud/issues/162.
+            for batch in check_batches.results:
+                assert isinstance(batch, Batch)
+                if batch.status not in {"PENDING", "RUNNING"}:
+                    logger.debug("Job %s is now complete", batch.id)
+                    pending_batch_ids.discard(batch.id)
+                    completed_batches[batch.id] = batch
+
+        logger.info("All jobs complete, %s sequences executed", len(completed_batches))
 
         # At this point, all batches are complete.
-        # Now collect data. We rely upon the fact
-        # that we enqueued exactly one batch per sequence, in the same order.
+        # Now, collect data.
+        #
+        # We rely upon the fact that for any `i`,
+        # `self._batch_id[i]` is the batch for `self.sequences[i]`.
         processed_data: list[ProcessedData] = []
-        for i, batch in enumerate(batches):
-            # Note: There's only one job per batch.
+        for i, id in enumerate(self._batch_ids):
+            batch = completed_batches[id]
             assert len(batch.jobs) == 1
+
             for _, job in batch.jobs.items():
                 if job.status == "DONE":
                     state_dict = job.result