qoro-divi 0.2.0b1__py3-none-any.whl → 0.5.0__py3-none-any.whl

divi/qprog/batch.py

@@ -2,36 +2,35 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import atexit
 import traceback
+import warnings
 from abc import ABC, abstractmethod
-from concurrent.futures import ProcessPoolExecutor, as_completed
-from multiprocessing import Event, Manager
-from multiprocessing.synchronize import Event as EventClass
+from concurrent.futures import Future, ThreadPoolExecutor, as_completed
 from queue import Empty, Queue
-from threading import Lock, Thread
+from threading import Event, Lock, Thread
+from typing import Any
 from warnings import warn
 
 from rich.console import Console
 from rich.progress import Progress, TaskID
 
-from divi._pbar import make_progress_bar
-from divi.interfaces import CircuitRunner
-from divi.parallel_simulator import ParallelSimulator
-from divi.qlogger import disable_logging
+from divi.backends import CircuitRunner
 from divi.qprog.quantum_program import QuantumProgram
+from divi.reporting import disable_logging, make_progress_bar
 
 
-def queue_listener(
+def _queue_listener(
     queue: Queue,
     progress_bar: Progress,
     pb_task_map: dict[QuantumProgram, TaskID],
-    done_event: EventClass,
+    done_event: Event,
     is_jupyter: bool,
     lock: Lock,
 ):
     while not done_event.is_set():
         try:
-            msg = queue.get(timeout=0.1)
+            msg: dict[str, Any] = queue.get(timeout=0.1)
         except Empty:
             continue
         except Exception as e:
@@ -41,14 +40,30 @@ def queue_listener(
         with lock:
             task_id = pb_task_map[msg["job_id"]]
 
-        progress_bar.update(
-            task_id,
-            advance=msg["progress"],
-            poll_attempt=msg.get("poll_attempt", 0),
-            message=msg.get("message", ""),
-            final_status=msg.get("final_status", ""),
-            refresh=is_jupyter,
-        )
+        # Prepare update arguments, starting with progress.
+        update_args = {"advance": msg["progress"]}
+
+        if "poll_attempt" in msg:
+            update_args["poll_attempt"] = msg.get("poll_attempt", 0)
+        if "max_retries" in msg:
+            update_args["max_retries"] = msg.get("max_retries")
+        if "service_job_id" in msg:
+            update_args["service_job_id"] = msg.get("service_job_id")
+        if "job_status" in msg:
+            update_args["job_status"] = msg.get("job_status")
+        if msg.get("message"):
+            update_args["message"] = msg.get("message")
+        if "final_status" in msg:
+            update_args["final_status"] = msg.get("final_status", "")
+
+        update_args["refresh"] = is_jupyter
+
+        progress_bar.update(task_id, **update_args)
+        queue.task_done()
+
+
+def _default_task_function(program: QuantumProgram):
+    return program.run()
 
 
 class ProgramBatch(ABC):
@@ -74,12 +89,12 @@ class ProgramBatch(ABC):
 
         self.backend = backend
         self._executor = None
-        self.programs = {}
+        self._task_fn = _default_task_function
+        self._programs = {}
 
         self._total_circuit_count = 0
         self._total_run_time = 0.0
 
-        self._is_local = isinstance(backend, ParallelSimulator)
         self._is_jupyter = Console().is_jupyter
 
         # Disable logging since we already have the bars to track progress
@@ -87,22 +102,100 @@ class ProgramBatch(ABC):
 
     @property
     def total_circuit_count(self):
+        """
+        Get the total number of circuits executed across all programs in the batch.
+
+        Returns:
+            int: Cumulative count of circuits submitted by all programs.
+        """
         return self._total_circuit_count
 
     @property
     def total_run_time(self):
+        """
+        Get the total runtime across all programs in the batch.
+
+        Returns:
+            float: Cumulative execution time in seconds across all programs.
+        """
         return self._total_run_time
 
+    @property
+    def programs(self) -> dict:
+        """
+        Get a copy of the programs dictionary.
+
+        Returns:
+            dict: Copy of the programs dictionary mapping program IDs to
+                QuantumProgram instances. Modifications to this dict will not
+                affect the internal state.
+        """
+        return self._programs.copy()
+
+    @programs.setter
+    def programs(self, value: dict):
+        """Set the programs dictionary."""
+        self._programs = value
+
     @abstractmethod
     def create_programs(self):
-        self._manager = Manager()
-        self._queue = self._manager.Queue()
+        """Generate and populate the programs dictionary for batch execution.
+
+        This method must be implemented by subclasses to create the quantum programs
+        that will be executed as part of the batch. The method operates via side effects:
+        it populates `self._programs` (or `self.programs`) with a dictionary mapping
+        program identifiers to `QuantumProgram` instances.
+
+        Implementation Notes:
+            - Subclasses should call `super().create_programs()` first to initialize
+              internal state (queue, events, etc.) and validate that no programs
+              already exist.
+            - After calling super(), subclasses should populate `self.programs` or
+              `self._programs` with their program instances.
+            - Program identifiers can be any hashable type (e.g., strings, tuples).
+              Common patterns include strings like "program_1", "program_2" or tuples like
+              ('A', 5) for partitioned problems.
+
+        Side Effects:
+            - Populates `self._programs` with program instances.
+            - Initializes `self._queue` for progress reporting.
+            - Initializes `self._done_event` if `max_iterations` attribute exists.
+
+        Raises:
+            RuntimeError: If programs already exist (should call `reset()` first).
+
+        Example:
+            >>> def create_programs(self):
+            ...     super().create_programs()
+            ...     self.programs = {
+            ...         "prog1": QAOA(...),
+            ...         "prog2": QAOA(...),
+            ...     }
+        """
+        if len(self._programs) > 0:
+            raise RuntimeError(
+                "Some programs already exist. "
+                "Clear the program dictionary before creating new ones by using batch.reset()."
+            )
+
+        self._queue = Queue()
 
         if hasattr(self, "max_iterations"):
             self._done_event = Event()
 
     def reset(self):
-        self.programs.clear()
+        """
+        Reset the batch to its initial state.
+
+        Clears all programs, stops any running executors, terminates listener threads,
+        and stops progress bars. This allows the batch to be reused for a new set of
+        programs.
+
+        Note:
+            Any running programs will be forcefully stopped. Results from incomplete
+            programs will be lost.
+        """
+        self._programs.clear()
 
         # Stop any active executor
         if self._executor is not None:
@@ -118,66 +211,111 @@ class ProgramBatch(ABC):
         if getattr(self, "_listener_thread", None) is not None:
             self._listener_thread.join(timeout=1)
             if self._listener_thread.is_alive():
-                warn(
-                    "Listener thread did not terminate within timeout and may still be running."
-                )
-            else:
-                self._listener_thread = None
-            self._queue.close()
-            self._queue.join_thread()
-            self._queue = None
+                warn("Listener thread did not terminate within timeout.")
+            self._listener_thread = None
 
         # Stop the progress bar if it's still active
         if getattr(self, "_progress_bar", None) is not None:
             try:
                 self._progress_bar.stop()
-                self._progress_bar = None
             except Exception:
                 pass  # Already stopped or not running
-
+            self._progress_bar = None
             self._pb_task_map.clear()
 
-    def add_program_to_executor(self, program):
-        self.futures.append(self._executor.submit(program.run))
+    def _atexit_cleanup_hook(self):
+        # This hook is only registered for non-blocking runs.
+        if self._executor is not None:
+            warn(
+                "A non-blocking ProgramBatch run was not explicitly closed with "
+                "'join()'. The batch was cleaned up automatically on exit.",
+                UserWarning,
+            )
+            self.reset()
+
+    def _add_program_to_executor(self, program: QuantumProgram) -> Future:
+        """
+        Add a quantum program to the thread pool executor for execution.
+
+        Sets up the program with cancellation support and progress tracking, then
+        submits it for execution in a separate thread.
+
+        Args:
+            program (QuantumProgram): The quantum program to execute.
+
+        Returns:
+            Future: A Future object representing the program's execution.
+        """
+        if hasattr(program, "_set_cancellation_event"):
+            program._set_cancellation_event(self._cancellation_event)
 
         if self._progress_bar is not None:
             with self._pb_lock:
-                self._pb_task_map[program.job_id] = self._progress_bar.add_task(
+                self._pb_task_map[program.program_id] = self._progress_bar.add_task(
                     "",
-                    job_name=f"Job {program.job_id}",
+                    job_name=f"Program {program.program_id}",
                     total=self.max_iterations,
                     completed=0,
-                    poll_attempt=0,
                     message="",
-                    final_status="",
-                    mode=("simulation" if self._is_local else "network"),
                 )
 
-    def run(self):
+        return self._executor.submit(self._task_fn, program)
+
+    def run(self, blocking: bool = False):
+        """
+        Execute all programs in the batch.
+
+        Starts all quantum programs in parallel using a thread pool. Can run in
+        blocking or non-blocking mode.
+
+        Args:
+            blocking (bool, optional): If True, waits for all programs to complete
+                before returning. If False, returns immediately and programs run in
+                the background. Defaults to False.
+
+        Returns:
+            ProgramBatch: Returns self for method chaining.
+
+        Raises:
+            RuntimeError: If a batch is already running or if no programs have been
+                created.
+
+        Note:
+            In non-blocking mode, call `join()` later to wait for completion and
+            collect results.
+        """
         if self._executor is not None:
             raise RuntimeError("A batch is already being run.")
 
-        if len(self.programs) == 0:
+        if len(self._programs) == 0:
             raise RuntimeError("No programs to run.")
 
         self._progress_bar = (
-            make_progress_bar(
-                max_retries=None if self._is_local else self.backend.max_retries,
-                is_jupyter=self._is_jupyter,
-            )
+            make_progress_bar(is_jupyter=self._is_jupyter)
             if hasattr(self, "max_iterations")
             else None
         )
 
-        self._executor = ProcessPoolExecutor()
+        # Validate that all program instances are unique to prevent thread-safety issues
+        program_instances = list(self._programs.values())
+        if len(set(program_instances)) != len(program_instances):
+            raise RuntimeError(
+                "Duplicate program instances detected in batch. "
+                "QuantumProgram instances are stateful and NOT thread-safe. "
+                "You must provide a unique instance for each program ID."
+            )
+
+        self._executor = ThreadPoolExecutor()
+        self._cancellation_event = Event()
         self.futures = []
+        self._future_to_program = {}
         self._pb_task_map = {}
         self._pb_lock = Lock()
 
         if self._progress_bar is not None:
             self._progress_bar.start()
             self._listener_thread = Thread(
-                target=queue_listener,
+                target=_queue_listener,
                 args=(
                     self._queue,
                     self._progress_bar,
@@ -190,46 +328,218 @@ class ProgramBatch(ABC):
             )
             self._listener_thread.start()
 
-        for program in self.programs.values():
-            self.add_program_to_executor(program)
+        for program in self._programs.values():
+            future = self._add_program_to_executor(program)
+            self.futures.append(future)
+            self._future_to_program[future] = program
+
+        if not blocking:
+            # Arm safety net
+            atexit.register(self._atexit_cleanup_hook)
+        else:
+            self.join()
 
-    def check_all_done(self):
+        return self
+
+    def check_all_done(self) -> bool:
+        """
+        Check if all programs in the batch have completed execution.
+
+        Returns:
+            bool: True if all programs are finished (successfully or with errors),
+                False if any are still running.
+        """
         return all(future.done() for future in self.futures)
 
-    def wait_for_all(self):
+    def _collect_completed_results(self, completed_futures: list):
+        """
+        Collects results from any futures that have completed successfully.
+        Appends (circuit_count, run_time) tuples to the completed_futures list.
+
+        Args:
+            completed_futures: List to append results to
+        """
+        for future in self.futures:
+            if future.done() and not future.cancelled():
+                try:
+                    completed_futures.append(future.result())
+                except Exception:
+                    pass  # Skip failed futures
+
+    def _handle_cancellation(self):
+        """
+        Handles cancellation gracefully, providing accurate feedback by checking
+        the result of future.cancel().
+        """
+        self._cancellation_event.set()
+
+        successfully_cancelled = []
+        unstoppable_futures = []
+
+        # --- Phase 1: Attempt to cancel all non-finished tasks ---
+        for future, program in self._future_to_program.items():
+            if future.done():
+                continue
+
+            task_id = self._pb_task_map.get(program.program_id)
+            if self._progress_bar and task_id is not None:
+                cancel_result = future.cancel()
+                if cancel_result:
+                    # The task was pending and was successfully cancelled.
+                    successfully_cancelled.append(program)
+                else:
+                    # The task is already running and cannot be stopped.
+                    # Attempt to cancel the cloud job to allow polling loop to exit.
+                    program.cancel_unfinished_job()
+                    unstoppable_futures.append(future)
+                    self._progress_bar.update(
+                        task_id,
+                        message="Finishing... ⏳",
+                        refresh=self._is_jupyter,
+                    )
+
+        # --- Phase 2: Immediately mark the successfully cancelled tasks ---
+        for program in successfully_cancelled:
+            task_id = self._pb_task_map.get(program.program_id)
+            if self._progress_bar and task_id is not None:
+                self._progress_bar.update(
+                    task_id,
+                    final_status="Cancelled",
+                    message="Cancelled by user",
+                    refresh=self._is_jupyter,
+                )
+
+        # --- Phase 3: Wait for the unstoppable tasks to finish ---
+        if unstoppable_futures:
+            for future in as_completed(unstoppable_futures):
+                program = self._future_to_program[future]
+                task_id = self._pb_task_map.get(program.program_id)
+                if self._progress_bar and task_id is not None:
+                    self._progress_bar.update(
+                        task_id,
+                        final_status="Aborted",
+                        message="Completed during cancellation",
+                        refresh=self._is_jupyter,
+                    )
+
+    def join(self):
+        """
+        Wait for all programs in the batch to complete and collect results.
+
+        Blocks until all programs finish execution, aggregating their circuit counts
+        and run times. Handles keyboard interrupts gracefully by attempting to cancel
+        remaining programs.
+
+        Returns:
+            bool or None: Returns False if interrupted by KeyboardInterrupt, None otherwise.
+
+        Raises:
+            RuntimeError: If any program fails with an exception, after cancelling
+                remaining programs.
+
+        Note:
+            This method should be called after `run(blocking=False)` to wait for
+            completion. It's automatically called when using `run(blocking=True)`.
+        """
         if self._executor is None:
             return
 
-        exceptions = []
+        completed_futures = []
         try:
-            # Ensure all futures are completed and handle exceptions.
+            # The as_completed iterator will yield futures as they finish.
+            # If a task fails, future.result() will raise the exception immediately.
             for future in as_completed(self.futures):
-                try:
-                    future.result()  # Raises an exception if the task failed.
-                except Exception as e:
-                    exceptions.append(e)
+                completed_futures.append(future.result())
+
+        except KeyboardInterrupt:
+
+            if self._progress_bar is not None:
+                self._progress_bar.console.print(
+                    "[bold yellow]Shutdown signal received, waiting for programs to finish current iteration...[/bold yellow]"
+                )
+                self._handle_cancellation()
+
+            # Collect results from any futures that completed before/during cancellation
+            self._collect_completed_results(completed_futures)
+
+            return False
+
+        except Exception as e:
+            # A task has failed. Print the error and cancel the rest.
+            print(f"A task failed with an exception. Cancelling remaining tasks...")
+            traceback.print_exception(type(e), e, e.__traceback__)
+
+            # Collect results from any futures that completed before the failure
+            self._collect_completed_results(completed_futures)
+
+            # Cancel all other futures that have not yet completed.
+            for f in self.futures:
+                f.cancel()
+
+            # Re-raise a new error to indicate the batch failed.
+            raise RuntimeError("Batch execution failed and was cancelled.") from e
 
         finally:
-            self._executor.shutdown(wait=True, cancel_futures=False)
+            # Aggregate results from completed futures
+            if completed_futures:
+                self._total_circuit_count += sum(
+                    result[0] for result in completed_futures
+                )
+                self._total_run_time += sum(result[1] for result in completed_futures)
+                self.futures.clear()
+
+            self._executor.shutdown(wait=False)
             self._executor = None
 
             if self._progress_bar is not None:
+                self._queue.join()
                 self._done_event.set()
                 self._listener_thread.join()
+                self._progress_bar.stop()
 
-        if exceptions:
-            for i, exc in enumerate(exceptions, 1):
-                print(f"Task {i} failed with exception:")
-                traceback.print_exception(type(exc), exc, exc.__traceback__)
-            raise RuntimeError("One or more tasks failed. Check logs for details.")
-
-        if self._progress_bar is not None:
-            self._progress_bar.stop()
-
-        self._total_circuit_count += sum(future.result()[0] for future in self.futures)
-        self._total_run_time += sum(future.result()[1] for future in self.futures)
-        self.futures = []
+        # After successful cleanup, try to unregister the hook.
+        try:
+            atexit.unregister(self._atexit_cleanup_hook)
+        except TypeError:
+            pass
 
     @abstractmethod
     def aggregate_results(self):
-        raise NotImplementedError
+        """
+        Aggregate results from all programs in the batch after execution.
+
+        This is an abstract method that must be implemented by subclasses. The base
+        implementation performs validation checks:
+        - Ensures programs have been created
+        - Waits for any running programs to complete (calls join() if needed)
+        - Verifies that all programs have completed execution (non-empty losses_history)
+
+        Subclasses should call super().aggregate_results() first, then implement
+        their own aggregation logic to combine results from all programs. The
+        aggregation should handle different result formats (counts dictionary,
+        expectation values, etc.) as appropriate for the specific use case.
+
+        Returns:
+            The aggregated result, format depends on the subclass implementation.
+
+        Raises:
+            RuntimeError: If no programs exist, or if programs haven't completed
+                execution (empty losses_history).
+        """
+        if len(self._programs) == 0:
+            raise RuntimeError("No programs to aggregate. Run create_programs() first.")
+
+        if self._executor is not None:
+            self.join()
+
+        # Suppress warnings when checking for empty losses_history for cleanliness sake
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                "ignore", category=UserWarning, message=".*losses_history is empty.*"
+            )
+            if any(
+                len(program.losses_history) == 0 for program in self._programs.values()
+            ):
+                raise RuntimeError(
+                    "Some/All programs have empty losses. Did you call run()?"
+                )