From 394fe6a717c645c29828fc4863a5b154e36754d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Henrique=20Silv=C3=A9rio?=
 <henrique.silverio@tecnico.ulisboa.pt>
Date: Wed, 18 Oct 2023 10:50:35 +0200
Subject: [PATCH] Add `wait` option to remote backends (#598)

* Add `wait` option to remote backends

* Updating docstring description
---
 pulser-core/pulser/backend/qpu.py           | 10 ++++++++--
 pulser-core/pulser/backend/remote.py        |  2 +-
 pulser-pasqal/pulser_pasqal/backends.py     |  7 ++++++-
 pulser-pasqal/pulser_pasqal/pasqal_cloud.py |  6 ++++--
 pulser-pasqal/requirements.txt              |  2 +-
 tests/test_backend.py                       |  2 +-
 6 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/pulser-core/pulser/backend/qpu.py b/pulser-core/pulser/backend/qpu.py
index 925d6d0d2..e7c30f89f 100644
--- a/pulser-core/pulser/backend/qpu.py
+++ b/pulser-core/pulser/backend/qpu.py
@@ -24,7 +24,9 @@
 class QPUBackend(RemoteBackend):
     """Backend for sequence execution on a QPU."""
 
-    def run(self, job_params: list[JobParams] | None = None) -> RemoteResults:
+    def run(
+        self, job_params: list[JobParams] | None = None, wait: bool = False
+    ) -> RemoteResults:
         """Runs the sequence on the remote QPU and returns the result.
 
         Args:
@@ -34,6 +36,10 @@ def run(self, job_params: list[JobParams] | None = None) -> RemoteResults:
                 is parametrized, the values for all the variables necessary
                 to build the sequence must be given in it's own mapping, for
                 each job, under the 'variables' field.
+            wait: Whether to wait until the results of the jobs become
+                available.  If set to False, the call is non-blocking and the
+                obtained results' status can be checked using their `status`
+                property.
 
         Returns:
             The results, which can be accessed once all sequences have been
@@ -55,7 +61,7 @@ def run(self, job_params: list[JobParams] | None = None) -> RemoteResults:
                     f"device ({max_runs})" + suffix
                 )
         results = self._connection.submit(
-            self._sequence, job_params=job_params
+            self._sequence, job_params=job_params, wait=wait
         )
         return cast(RemoteResults, results)
 
diff --git a/pulser-core/pulser/backend/remote.py b/pulser-core/pulser/backend/remote.py
index f15af5d8d..aa07dccfe 100644
--- a/pulser-core/pulser/backend/remote.py
+++ b/pulser-core/pulser/backend/remote.py
@@ -96,7 +96,7 @@ class RemoteConnection(ABC):
 
     @abstractmethod
     def submit(
-        self, sequence: Sequence, **kwargs: Any
+        self, sequence: Sequence, wait: bool = False, **kwargs: Any
     ) -> RemoteResults | tuple[RemoteResults, ...]:
         """Submit a job for execution."""
         pass
diff --git a/pulser-pasqal/pulser_pasqal/backends.py b/pulser-pasqal/pulser_pasqal/backends.py
index 43d72333c..73706511b 100644
--- a/pulser-pasqal/pulser_pasqal/backends.py
+++ b/pulser-pasqal/pulser_pasqal/backends.py
@@ -57,7 +57,7 @@ def __init__(
             )
 
     def run(
-        self, job_params: list[JobParams] | None = None
+        self, job_params: list[JobParams] | None = None, wait: bool = False
     ) -> RemoteResults | tuple[RemoteResults, ...]:
         """Executes on the emulator backend through the Pasqal Cloud.
 
@@ -68,6 +68,10 @@ def run(
                 is parametrized, the values for all the variables necessary
                 to build the sequence must be given in it's own mapping, for
                 each job, under the 'variables' field.
+            wait: Whether to wait until the results of the jobs become
+                available.  If set to False, the call is non-blocking and the
+                obtained results' status can be checked using their `status`
+                property.
 
         Returns:
             The results, which can be accessed once all sequences have been
@@ -87,6 +91,7 @@ def run(
             job_params=job_params,
             emulator=self.emulator,
             config=self._config,
+            wait=wait,
         )
 
     def _validate_config(self, config: EmulatorConfig) -> None:
diff --git a/pulser-pasqal/pulser_pasqal/pasqal_cloud.py b/pulser-pasqal/pulser_pasqal/pasqal_cloud.py
index b64592b89..d821cc014 100644
--- a/pulser-pasqal/pulser_pasqal/pasqal_cloud.py
+++ b/pulser-pasqal/pulser_pasqal/pasqal_cloud.py
@@ -98,7 +98,9 @@ def __init__(
             **kwargs,
         )
 
-    def submit(self, sequence: Sequence, **kwargs: Any) -> RemoteResults:
+    def submit(
+        self, sequence: Sequence, wait: bool = False, **kwargs: Any
+    ) -> RemoteResults:
         """Submits the sequence for execution on a remote Pasqal backend."""
         if not sequence.is_measured():
             bases = sequence.get_addressed_bases()
@@ -141,7 +143,7 @@ def submit(self, sequence: Sequence, **kwargs: Any) -> RemoteResults:
             jobs=job_params or [],  # type: ignore[arg-type]
             emulator=emulator,
             configuration=configuration,
-            wait=False,
+            wait=wait,
         )
 
         return RemoteResults(batch.id, self)
diff --git a/pulser-pasqal/requirements.txt b/pulser-pasqal/requirements.txt
index 495317ef1..3d5108f84 100644
--- a/pulser-pasqal/requirements.txt
+++ b/pulser-pasqal/requirements.txt
@@ -1,2 +1,2 @@
-pasqal-cloud ~= 0.3.5
+pasqal-cloud ~= 0.4.0
 backoff ~= 2.2
\ No newline at end of file
diff --git a/tests/test_backend.py b/tests/test_backend.py
index b21cfd23f..a2e96db23 100644
--- a/tests/test_backend.py
+++ b/tests/test_backend.py
@@ -219,7 +219,7 @@ class _MockConnection(RemoteConnection):
     def __init__(self):
         self._status_calls = 0
 
-    def submit(self, sequence, **kwargs) -> RemoteResults:
+    def submit(self, sequence, wait: bool = False, **kwargsn) -> RemoteResults:
         return RemoteResults("abcd", self)
 
     def _fetch_result(self, submission_id: str) -> typing.Sequence[Result]: