change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

change to context manager interface for open batches

pulser-core/pulser/backend/qpu.py

@@ -14,7 +14,8 @@
 """Defines the backend class for QPU execution."""
 from __future__ import annotations
 
-from typing import cast
+from types import TracebackType
+from typing import Self, Type, cast
 
 from pulser import Sequence
 from pulser.backend.remote import (
@@ -33,8 +34,6 @@ def run(
         self,
         job_params: list[JobParams] | None = None,
         wait: bool = False,
-        open_submission: bool = False,
-        submission_id: str | None = None,
     ) -> RemoteResults:
         """Runs the sequence on the remote QPU and returns the result.
 
@@ -52,23 +51,11 @@ def run(
             open_submission: A flag indicating whether or not the submission
                 should be for a single job or accept future jobs before
                 closing.
-            submission_id: If you have a preexisting submission and wish to add
-                more jobs to it then you can provide the ID here and 'run'
-                will attempt allocate jobs.
 
         Returns:
             The results, which can be accessed once all sequences have been
             successfully executed.
         """
-
-        # if submission is already open, we don't need an ID for a preexisting
-        # one
-        if open_submission and submission_id:
-            raise ValueError(
-                """Open submission can only be used for a new submission.
-                                Don't provide a preexisting submission_id"""
-            )
-
         suffix = " when executing a sequence on a real QPU."
         if not job_params:
             raise ValueError("'job_params' must be specified" + suffix)
@@ -88,11 +75,49 @@ def run(
             self._sequence,
             job_params=job_params,
             wait=wait,
-            open_submission=open_submission,
-            submission_id=submission_id,
+            submission_id=self.submission_id,
         )
         return cast(RemoteResults, results)
 
+    def open_submission(self, submission_id: str | None = None) -> Self:
+        """
+        If provided a submission_id, we can add new jobs to a batch within
+        the scope of the context manager otherwise an new open submission
+        is created that you can continuously add jobs to.
+
+        Args:
+            submission_id: An optional unique identifier for an already
+                 open submission
+
+        Returns:
+            self reference for current object
+        """
+        if submission_id:
+            self.submission_id = submission_id
+        else:
+            submission = cast(
+                RemoteResults, self._connection.submit(self._sequence)
+            )
+            self.submission_id = submission.submission_id
+        return self
+
+    def __enter__(self) -> Self:
+        # open returns an instance of self to use open_submission within a
+        # context manager
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        _, _, _ = exc_type, exc_value, traceback
+        # exiting the context of an open submission will
+        # make a remote call to close it.
+        if self.submission_id:
+            self.close_submission(self.submission_id)
+
     def validate_sequence(self, sequence: Sequence) -> None:
         """Validates a sequence prior to submission.
 
@@ -116,6 +141,5 @@ def close_submission(self, submission_id: str) -> SubmissionStatus | None:
 
         Returns:
             SubmissionStatus representing the new stats of the submission.
-
         """
         return self._connection.close_submission(submission_id)

pulser-core/pulser/backend/remote.py

@@ -110,13 +110,22 @@ def submit(
         self,
         sequence: Sequence,
         wait: bool = False,
-        open_submission: bool = False,
         submission_id: str | None = None,
         **kwargs: Any,
     ) -> RemoteResults | tuple[RemoteResults, ...]:
         """Submit a job for execution."""
         pass
 
+    def close_submission(self, submission_id: str) -> SubmissionStatus | None:
+        """
+        Close a submission and make it unavailable to submit any further jobs
+        to.
+        """
+        pass
+        # raise NotImplementedError(
+        #     "Unable to close submission through this remote connection"
+        # )
+
     @abstractmethod
     def _fetch_result(self, submission_id: str) -> typing.Sequence[Result]:
         """Fetches the results of a completed submission."""
@@ -131,14 +140,6 @@ def _get_submission_status(self, submission_id: str) -> SubmissionStatus:
         """
         pass
 
-    @abstractmethod
-    def close_submission(self, submission_id: str) -> SubmissionStatus | None:
-        """
-        Close a submission and make it unavailable to submit any further jobs
-        to.
-        """
-        pass
-
     def fetch_available_devices(self) -> dict[str, Device]:
         """Fetches the devices available through this connection."""
         raise NotImplementedError(  # pragma: no cover
@@ -169,3 +170,4 @@ def __init__(
                 "'connection' must be a valid RemoteConnection instance."
             )
         self._connection = connection
+        self.submission_id: str = ""

pulser-pasqal/pulser_pasqal/pasqal_cloud.py

@@ -102,7 +102,6 @@ def submit(
         self,
         sequence: Sequence,
         wait: bool = False,
-        open_submission: bool = False,
         submission_id: str | None = None,
         **kwargs: Any,
     ) -> RemoteResults:
@@ -157,7 +156,6 @@ def submit(
                 serialized_sequence=sequence.to_abstract_repr(),
                 jobs=job_params or [],  # type: ignore[arg-type]
                 emulator=emulator,
-                complete=open_submission,
                 configuration=configuration,
                 wait=wait,
             )

tests/test_backend.py

@@ -225,12 +225,14 @@ def submit(
         self,
         sequence,
         wait: bool = False,
-        open_submission=False,
         submission_id=None,
         **kwargs,
     ) -> RemoteResults:
         return RemoteResults("abcd", self)
 
+    def close_connection(self, submission_id: str) -> SubmissionStatus | None:
+        return None
+
     def _fetch_result(self, submission_id: str) -> typing.Sequence[Result]:
         return (
             SampledResult(
@@ -288,3 +290,26 @@ def test_qpu_backend(sequence):
 
     results = remote_results.results
     assert results[0].sampling_dist == {"00": 1.0}
+
+
+def test_qpu_back_with_context_manager(sequence):
+    connection = _MockConnection()
+    unique_test_id = "unique-id-string"
+
+    with pytest.warns(DeprecationWarning, match="From v0.18"):
+        seq = sequence.switch_device(replace(DigitalAnalogDevice, max_runs=10))
+
+    with QPUBackend(seq, connection).open_submission(unique_test_id) as qpub:
+        remote_results = qpub.run(job_params=[{"runs": 10}])
+        assert isinstance(remote_results, RemoteResults)
+        # confirms context manager assigns ID
+        assert qpub.submission_id == unique_test_id
+
+    with QPUBackend(seq, connection).open_submission() as qpub:
+        remote_results = qpub.run(job_params=[{"runs": 10}])
+        assert isinstance(remote_results, RemoteResults)
+        # confirms context manager assigns ID
+
+        assert qpub.submission_id is not None
+        # Assert new submission_id is equal to a predefined mock response
+        assert qpub.submission_id == "abcd"

tests/test_pasqal.py

@@ -212,7 +212,6 @@ def test_submit(fixt, parametrized, emulator, seq, mock_job):
             emulator=emulator,
             configuration=sdk_config,
             wait=False,
-            complete=False,
         )
     )
 
@@ -330,5 +329,4 @@ def test_emulators_run(fixt, seq, emu_cls, parametrized: bool):
         emulator=emulator_type,
         configuration=sdk_config,
         wait=False,
-        complete=False,
     )