diff --git a/docs/source/api_reference/api/merlin.core.merlin_processor.rst b/docs/source/api_reference/api/merlin.core.merlin_processor.rst index fedbfd25..c85a77b2 100644 --- a/docs/source/api_reference/api/merlin.core.merlin_processor.rst +++ b/docs/source/api_reference/api/merlin.core.merlin_processor.rst @@ -21,8 +21,8 @@ classical layers local. It supports these backend entry points: All execution paths support batched execution, per-call/global timeouts, cooperative cancellation, and a Torch-friendly async interface returning -:class:`torch.futures.Future`. Remote backends additionally support chunking -with limited intra-leaf concurrency. +:class:`MerlinFuture` (a :class:`torch.futures.Future` subclass). Remote +backends additionally support chunking with limited intra-leaf concurrency. Key Capabilities ---------------- @@ -91,6 +91,44 @@ BackendCapabilities print(f"Platform: {caps.name}") print(f"Supports: {caps.available_commands}") +MerlinFuture +------------ +.. class:: MerlinFuture(call_state, cancel_all) + + Typed async handle returned by :meth:`forward_async`. Subclasses + :class:`torch.futures.Future`, so all inherited behavior (``wait``, + ``done``, ``then``, ``value``, ...) is unchanged, and adds the + Merlin-specific async contract that was previously monkey-patched onto + plain Future instances at runtime. + + Instances are constructed by :meth:`forward_async`; user code only + consumes them. + + **Attributes** + + .. attribute:: job_ids + + ``list[str]`` — Remote job IDs accumulated across chunk jobs, in + observation order. This is a live view: IDs recorded while chunks are + polling appear immediately. + + **Methods** + + .. method:: status() -> dict + + Current progress and state of the call: + ``{"state", "progress", "message", "chunks_total", "chunks_done", + "active_chunks"}``. ``state`` is ``"IDLE"`` before the first backend + poll, the backend-reported state while polling, and ``"COMPLETE"`` + once the future resolves without a recorded backend status. + + .. method:: cancel_remote() -> None + + Cooperatively cancel the call and its in-flight remote jobs. If the + future is not already done, it resolves with + :class:`concurrent.futures.CancelledError`, which ``wait()`` then + raises. + MerlinProcessor --------------- .. class:: MerlinProcessor(remote_processor=None, session=None, microbatch_size=32, timeout=3600.0, max_shots_per_call=None, chunk_concurrency=1, token=None, *, processor=None) @@ -224,19 +262,12 @@ Execution APIs :raises concurrent.futures.CancelledError: If the call is cooperatively cancelled via the async API. -.. method:: forward_async(module, input, *, nsample=None, timeout=None) -> torch.futures.Future - - Asynchronous execution. Returns a :class:`torch.futures.Future` with extra - helpers attached: - - **Future extensions** +.. method:: forward_async(module, input, *, nsample=None, timeout=None) -> MerlinFuture - * ``future.job_ids: list[str]`` — accumulates remote job IDs across chunk jobs. - * ``future.status() -> dict`` — current state/progress/message plus chunk - counters: ``{"chunks_total", "chunks_done", "active_chunks"}``. - * ``future.cancel_remote() -> None`` — cooperative cancel; in-flight jobs are - best-effort cancelled and ``future.wait()`` raises - ``CancelledError``. + Asynchronous execution. Returns a :class:`MerlinFuture` — a + :class:`torch.futures.Future` subclass declaring ``job_ids``, + ``status()``, and ``cancel_remote()`` on the class (see + :class:`MerlinFuture` for details). :param module: See :meth:`forward`. :param input: See :meth:`forward`. diff --git a/merlin/__init__.py b/merlin/__init__.py index 66b37057..1bac00f2 100644 --- a/merlin/__init__.py +++ b/merlin/__init__.py @@ -49,7 +49,7 @@ # Advanced components (for power users) from .core.computation_space import ComputationSpace from .core.encoding_space import EncodingSpace -from .core.merlin_processor import MerlinProcessor +from .core.merlin_processor import MerlinFuture, MerlinProcessor from .core.process import ComputationProcess from .core.state import StatePattern, generate_state from .measurement import ( @@ -106,6 +106,7 @@ "resolve_detectors", "ModeExpectations", "MerlinProcessor", + "MerlinFuture", "Amplitudes", "LexGrouping", "ModGrouping", diff --git a/merlin/core/execution.py b/merlin/core/execution.py new file mode 100644 index 00000000..78405270 --- /dev/null +++ b/merlin/core/execution.py @@ -0,0 +1,444 @@ +"""Execution units extracted from MerlinProcessor (PML-305). + +This module hosts the focused components that own remote chunk execution: + +- :class:`BatchChunker` — splits Merlin-level batches into microbatch chunks, + runs them with bounded concurrency, and stitches the outputs back together. +- :class:`RemoteJobRunner` — executes one remote chunk end to end: builds a + fresh remote processor, prepares sampler iterations, submits a job, polls it, + and maps the raw results into a tensor. + +Both units receive their dependencies (fresh-processor factory, backend +capabilities, result mapping, job tracking) as injected callables so they are +independently testable with no-cloud fakes. :class:`~merlin.core.merlin_processor.MerlinProcessor` +remains the public entry point and coordinates these units without owning the +execution details itself. + +The threading semantics (daemon chunk threads, cooperative cancellation, +deadline checks, polling backoff) are intentionally identical to the +pre-extraction MerlinProcessor implementation. +""" + +from __future__ import annotations + +import logging +import threading +import time +import zlib +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch +from perceval.runtime import RemoteJob, RemoteProcessor + +from .perceval_adapter import PercevalAdapter, RemoteJobFailedError + +if TYPE_CHECKING: + from ..algorithms.module import MerlinModule + from .merlin_processor import CallState, ValidatedLayerConfig + +logger = logging.getLogger(__name__) + + +class BatchChunker: + """Split input batches into chunks and run them with bounded concurrency. + + Parameters + ---------- + run_chunk : Callable + Callable executing one chunk with the signature + ``(layer, config, input_chunk, nsample, state, deadline, job_base_label)`` + and returning a ``torch.Tensor``. Injected so chunk orchestration is + testable with fakes and so monkeypatched processor methods stay + observable. + chunk_concurrency : int + Maximum number of chunk jobs in flight at once. + cancel_all : Callable[[], None] + Invoked when the deadline elapses so in-flight remote jobs are + cancelled best-effort before raising ``TimeoutError``. + """ + + def __init__( + self, + *, + run_chunk: Callable[..., torch.Tensor], + chunk_concurrency: int, + cancel_all: Callable[[], None], + ) -> None: + self._run_chunk = run_chunk + self._chunk_concurrency = max(1, int(chunk_concurrency)) + self._cancel_all = cancel_all + + @staticmethod + def split_batch(batch_size: int, microbatch_size: int) -> list[tuple[int, int]]: + """Split ``batch_size`` rows into ``[start, end)`` microbatch chunks.""" + chunks: list[tuple[int, int]] = [] + start = 0 + while start < batch_size: + end = min(start + microbatch_size, batch_size) + chunks.append((start, end)) + start = end + return chunks + + def run_chunks( + self, + layer: MerlinModule, + config: ValidatedLayerConfig, + input_tensor: torch.Tensor, + chunks: list[tuple[int, int]], + nsample: int | None, + state: CallState, + deadline: float | None, + ) -> torch.Tensor: + """Submit chunk jobs with limited concurrency and stitch results.""" + state.add_planned_chunks(len(chunks)) + outputs: list[torch.Tensor | None] = [None] * len(chunks) + errors: list[BaseException] = [] + + total_chunks = len(chunks) + layer_name = getattr(layer, "name", layer.__class__.__name__) + + def _call(s: int, e: int, idx: int): + try: + base_label = ( + f"mer:{layer_name}:{state.call_id}:{idx + 1}/{total_chunks}" + ) + t = self._run_chunk( + layer, + config, + input_tensor[s:e], + nsample, + state, + deadline, + job_base_label=base_label, + ) + outputs[idx] = t + except BaseException as ex: + errors.append(ex) + + in_flight = 0 + idx = 0 + futures: list[threading.Thread] = [] + while idx < len(chunks) or in_flight > 0: + while idx < len(chunks) and in_flight < self._chunk_concurrency: + s, e = chunks[idx] + state.mark_chunk_started() + th = threading.Thread(target=_call, args=(s, e, idx), daemon=True) + th.start() + futures.append(th) + idx += 1 + in_flight += 1 + + for th in list(futures): + if not th.is_alive(): + futures.remove(th) + in_flight -= 1 + state.mark_chunk_finished() + + if deadline is not None and time.time() >= deadline: + self._cancel_all() + raise TimeoutError("Remote call timed out (remote cancel issued)") + + time.sleep(0.01) + + if errors: + raise errors[0] + + return torch.cat(outputs, dim=0) # type: ignore[arg-type] + + +class RemoteJobRunner: + """Execute a single remote chunk: fresh processor, submit, poll, map. + + Parameters + ---------- + create_processor : Callable[[], RemoteProcessor] + Factory returning a fresh, independent RemoteProcessor per attempt. + get_available_commands : Callable[[], tuple[str, ...]] + Returns the backend command snapshot driving probs-vs-sampling. + effective_sample_count : Callable[[int | None], int] + Maps a requested ``nsample`` to the capped shot count to submit. + get_max_shots_per_call : Callable[[], int | None] + Returns the current hard cap on shots per sampler call. + default_shots_per_call : int + Fallback shots value used when the cap is unset. + map_results : Callable + Maps a raw results dict to a tensor with the signature + ``(raw_results, batch_size, layer, nsample, is_probability)``. + register_job : Callable[[RemoteJob], None] + Records a submitted job for cancellation tracking and history. + unregister_job : Callable[[RemoteJob], None] + Removes a job from active-cancellation tracking. + get_microbatch_limit : Callable[[], int | None] + Returns the per-chunk size guard, or ``None`` when chunk sizes are + not bounded (session backends). + max_retries : int + Number of submission attempts per chunk. + job_name_max : int + Maximum length of remote job names. + """ + + def __init__( + self, + *, + create_processor: Callable[[], RemoteProcessor], + get_available_commands: Callable[[], tuple[str, ...]], + effective_sample_count: Callable[[int | None], int], + get_max_shots_per_call: Callable[[], int | None], + default_shots_per_call: int, + map_results: Callable[..., torch.Tensor], + register_job: Callable[[RemoteJob], None], + unregister_job: Callable[[RemoteJob], None], + get_microbatch_limit: Callable[[], int | None], + max_retries: int = 3, + job_name_max: int = 50, + ) -> None: + self._create_processor = create_processor + self._get_available_commands = get_available_commands + self._effective_sample_count = effective_sample_count + self._get_max_shots_per_call = get_max_shots_per_call + self._default_shots_per_call = default_shots_per_call + self._map_results = map_results + self._register_job = register_job + self._unregister_job = unregister_job + self._get_microbatch_limit = get_microbatch_limit + self._max_retries = max_retries + self._job_name_max = job_name_max + + def run_chunk( + self, + layer: MerlinModule, + config: ValidatedLayerConfig, + input_chunk: torch.Tensor, + nsample: int | None, + state: CallState, + deadline: float | None, + job_base_label: str | None = None, + ) -> torch.Tensor: + """Submit a single chunk job with retries and return the mapped tensor.""" + from concurrent.futures import CancelledError + + batch_size = input_chunk.shape[0] + microbatch_limit = self._get_microbatch_limit() + if microbatch_limit is not None and batch_size > microbatch_limit: + raise ValueError( + f"Chunk size {batch_size} exceeds microbatch {microbatch_limit}. " + "Please report this bug." + ) + + input_param_names = list(config.input_param_order) + input_np = input_chunk.detach().cpu().numpy() + + # Pre-compute iteration params (cheap, only done once). + iteration_params: list[dict[str, float]] = [] + for i in range(batch_size): + circuit_params = {} + for j, param_name in enumerate(input_param_names): + circuit_params[param_name] = ( + float(input_np[i, j]) if j < input_chunk.shape[1] else 0.0 + ) + iteration_params.append(circuit_params) + + last_error: BaseException | None = None + for attempt in range(self._max_retries): + if state.cancel_requested: + raise CancelledError("Remote call was cancelled") + if deadline is not None and time.time() >= deadline: + raise TimeoutError("Remote call timed out (remote cancel issued)") + + # Build a fresh RemoteProcessor and Sampler on each attempt so that + # a corrupted RP doesn't poison retries. + rp = self._create_processor() + PercevalAdapter.configure_processor(rp, config.circuit, config.input_state) + + max_shots_per_call = self._get_max_shots_per_call() + max_shots_arg = ( + self._default_shots_per_call + if max_shots_per_call is None + else int(max_shots_per_call) + ) + sampler = PercevalAdapter.create_sampler( + rp, max_shots_arg, iteration_params + ) + + job = None + try: + job, is_probability = self.submit_job( + sampler, nsample, job_base_label, self._capped_name + ) + self._register_job(job) + + return self.poll_job( + job, state, deadline, batch_size, layer, nsample, is_probability + ) + except (CancelledError, TimeoutError, KeyboardInterrupt): + raise + except Exception as exc: + last_error = exc + if job is not None: + self._unregister_job(job) + logger.warning( + "Chunk attempt %d/%d failed: %s", + attempt + 1, + self._max_retries, + exc, + ) + if attempt < self._max_retries - 1: + time.sleep(min(1.0 * (2**attempt), 5.0)) + + raise RuntimeError( + f"Chunk failed after {self._max_retries} attempts" + ) from last_error + + def _capped_name(self, base: str, cmd: str) -> str: + """Return a sanitized remote job name capped at ``job_name_max``.""" + name = f"{base}:{cmd}" + name = "".join(ch if ch.isalnum() or ch in "-_:/=." else "_" for ch in name) + if len(name) <= self._job_name_max: + return name + h = f"{zlib.adler32(name.encode()):08x}" + keep = self._job_name_max - 1 - len(h) + if keep < 1: + return h[: self._job_name_max] + return name[:keep] + "~" + h + + def submit_job(self, sampler, nsample, job_base_label, capped_name): + """Submit a job to the sampler, selecting command based on backend capabilities. + + **Command Selection Strategy** + + 1. **Exact Probabilities** (``"probs"`` command): + - Used if backend exposes ``"probs"`` AND (``nsample`` is None or ``nsample <= 0``). + - Returns normalized probability distribution; ``nsample`` is ignored. + + 2. **Sampling** (``"sample_count"`` or ``"samples"`` commands): + - Used if exact probabilities are not available or ``nsample > 0``. + - Tries ``"sample_count"`` first, falls back to ``"samples"``. + - Number of samples = ``effective_sample_count(nsample)``. + + Parameters + ---------- + sampler : Sampler + Perceval Sampler instance configured with circuit and iterations. + nsample : int | None + Number of samples requested. If ``None`` or ``<= 0``, triggers + exact probability computation (if available). + job_base_label : str | None + Base label for the remote job name. + capped_name : callable + Function to cap and format job names. + + Returns + ------- + tuple[RemoteJob, bool] + The submitted job handle and the ``is_probability`` execution flag. + """ + available_commands = self._get_available_commands() + is_probability = ("probs" in available_commands) and ( + nsample is None or int(nsample) <= 0 + ) + + if is_probability: + cmd = "probs" + max_samples = None + else: + if "sample_count" in available_commands: + cmd = "sample_count" + elif "samples" in available_commands: + cmd = "samples" + else: + cmd = "sample_count" + max_samples = self._effective_sample_count(nsample) + + name = capped_name(job_base_label, cmd) if job_base_label else None + job = PercevalAdapter.submit_async( + sampler, cmd, name=name, max_samples=max_samples + ) + return job, is_probability + + def poll_job( + self, + job: RemoteJob, + state: CallState, + deadline: float | None, + batch_size: int, + layer: MerlinModule, + nsample: int | None, + is_probability: bool = False, + ) -> torch.Tensor: + """Poll a submitted job until complete/failed/timeout and return results. + + Continuously polls the job status, updating call state and handling + timeouts, cancellation, and failures. Upon completion, maps results to + a tensor through the injected ``map_results`` dependency. + """ + from concurrent.futures import CancelledError + + _MAX_NON_DICT_RETRIES = 60 # 60 * 0.1s = 6s + non_dict_retries = 0 + sleep_ms = 50 + while True: + if state.cancel_requested: + PercevalAdapter.cancel_job(job) + raise CancelledError("Remote call was cancelled") + + if deadline is not None and time.time() >= deadline: + PercevalAdapter.cancel_job(job) + raise TimeoutError("Remote call timed out (remote cancel issued)") + + snapshot = PercevalAdapter.job_snapshot(job) + state.set_current_status( + state=snapshot.state, + progress=snapshot.progress, + message=snapshot.stop_message, + ) + + job_id = snapshot.job_id + if job_id is not None: + state.record_job_id(job_id) + + if snapshot.is_failed: + msg = snapshot.stop_message + if msg and "Cancel requested" in str(msg): + self._unregister_job(job) + raise CancelledError("Remote call was cancelled") + self._unregister_job(job) + raise RemoteJobFailedError( + f"Remote job failed: {msg or 'unknown error'} (job_id={job_id!r})" + ) + + if snapshot.is_complete: + try: + raw = PercevalAdapter.get_results(job) + except RuntimeError as ex: + msg = str(ex) + if "Results are not available" in msg: + time.sleep(0.05) + continue + if "Cancel requested" in msg: + self._unregister_job(job) + raise CancelledError("Remote call was cancelled") + raise + + if isinstance(raw, dict): + self._unregister_job(job) + return self._map_results( + raw, batch_size, layer, nsample, is_probability + ) + + # The backend sometimes reports completion before the dict + # payload is actually available. Re-poll the same job for a + # bounded window before giving up to the outer retry loop. + non_dict_retries += 1 + if non_dict_retries >= _MAX_NON_DICT_RETRIES: + self._unregister_job(job) + raise RuntimeError( + f"Job complete but results were not a dict after " + f"{_MAX_NON_DICT_RETRIES} re-polls; " + f"job_id={job_id!r}, type={type(raw)}, value={raw!r}" + ) + time.sleep(0.1) + continue + + time.sleep(sleep_ms / 1000.0) + sleep_ms = min(sleep_ms * 2, 400) diff --git a/merlin/core/merlin_processor.py b/merlin/core/merlin_processor.py index 2815528e..00fd43ce 100644 --- a/merlin/core/merlin_processor.py +++ b/merlin/core/merlin_processor.py @@ -1,12 +1,9 @@ -import copy import logging import threading import time import uuid import warnings -import zlib -from collections.abc import Iterable, Sequence -from contextlib import suppress +from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass from numbers import Integral from typing import Any, Protocol, cast, runtime_checkable @@ -15,13 +12,19 @@ import perceval as pcvl import torch import torch.nn as nn -from perceval.algorithm import Sampler -from perceval.runtime import AProcessor, Processor, RemoteJob, RemoteProcessor +from perceval.runtime import AProcessor, RemoteJob, RemoteProcessor from perceval.runtime.session import ISession from torch.futures import Future from ..algorithms.module import MerlinModule from ..utils.combinadics import Combinadics +from .execution import BatchChunker, RemoteJobRunner +from .perceval_adapter import ( + LOCAL_EXPERIMENT_SNAPSHOT_ATTR, + LocalExperimentSnapshot, + PercevalAdapter, + TokenExtractionError, +) logger = logging.getLogger(__name__) @@ -43,31 +46,268 @@ class BackendCapabilities: @dataclass(frozen=True) -class _LocalExperimentMetadata: - """Experiment-level state that must survive local circuit replacement.""" - - circuit_size: int - in_ports: tuple[tuple[Any, tuple[int, ...]], ...] - out_ports: tuple[tuple[Any, tuple[int, ...]], ...] - detectors: tuple[Any | None, ...] - detectors_injected: tuple[int, ...] - in_mode_type: tuple[Any, ...] - out_mode_type: tuple[Any, ...] - anon_herald_num: int - postselection: Any +class JobStatus: + """Immutable snapshot of the most recently observed backend job status. + + Attributes + ---------- + state : Any + Backend-reported job state (e.g. ``"RUNNING"``), or ``None`` when the + backend did not expose one. + progress : Any + Backend-reported progress value, or ``None``. + message : Any + Backend-reported stop/status message, or ``None``. + """ + + state: Any = None + progress: Any = None + message: Any = None + + +class CallState: + """Typed per-call execution state for one :meth:`MerlinProcessor.forward_async` call. + + Replaces the anonymous mutable ``state`` dict previously threaded through + ``forward_async()``, chunk orchestration, chunk execution, and job polling. + All per-call runtime state is owned by this object and mutated only through + named helpers, making the contract explicit and searchable. + + **Thread ownership** + + - ``call_id``: immutable after creation; readable from any thread. + - ``cancel_requested``: set (never cleared) by the caller thread through + :meth:`request_cancel`; read cooperatively by the pipeline, chunk, and + polling threads. + - ``current_status``: written by polling threads through + :meth:`set_current_status`; read by :meth:`status_snapshot`. + - ``job_ids``: appended (with deduplication) by polling threads through + :meth:`record_job_id`. The list object is intentionally shared with + ``future.job_ids`` so recorded ids appear on the future as they arrive. + - Chunk counters (``chunks_total``, ``chunks_done``, ``active_chunks``): + mutated by the chunk orchestration thread through + :meth:`add_planned_chunks`, :meth:`mark_chunk_started`, and + :meth:`mark_chunk_finished` under the internal lock. + """ + + def __init__(self, call_id: str) -> None: + """Initialize an empty call state. + + Parameters + ---------- + call_id : str + Short identifier embedded in remote job names for traceability. + Use :meth:`new` to generate one automatically. + """ + self.call_id = call_id + self.job_ids: list[str] = [] + self._lock = threading.Lock() + self._cancel_requested = False + self._current_status: JobStatus | None = None + self._chunks_total = 0 + self._chunks_done = 0 + self._active_chunks = 0 + + @classmethod + def new(cls) -> "CallState": + """Create a fresh call state with a short unique call identifier. + + Returns + ------- + CallState + Empty state carrying an 8-character hex ``call_id``. + """ + return cls(call_id=uuid.uuid4().hex[:8]) + + # ---- cancellation ---- + + @property + def cancel_requested(self) -> bool: + """Whether cooperative cancellation has been requested for this call.""" + return self._cancel_requested + + def request_cancel(self) -> None: + """Request cooperative cancellation of this call (irreversible).""" + self._cancel_requested = True + + # ---- job ids ---- + + def record_job_id(self, job_id: str) -> None: + """Record a remote job id, deduplicating repeated observations. + + Parameters + ---------- + job_id : str + Identifier reported by the backend for a submitted job. Re-observing + an already-recorded id (each polling cycle re-reads it) is a no-op. + """ + with self._lock: + if job_id not in self.job_ids: + self.job_ids.append(job_id) + + # ---- backend job status ---- @property - def has_mode_metadata(self) -> bool: - """Return whether metadata is tied to a concrete circuit mode layout.""" - - return ( - bool(self.in_ports) - or bool(self.out_ports) - or any(detector is not None for detector in self.detectors) - or bool(self.detectors_injected) - or self.postselection != pcvl.PostSelect() + def current_status(self) -> JobStatus | None: + """Most recent backend job status, or ``None`` before first poll.""" + return self._current_status + + def set_current_status( + self, *, state: Any = None, progress: Any = None, message: Any = None + ) -> None: + """Record the latest backend job status observed while polling. + + Parameters + ---------- + state : Any + Backend-reported job state (e.g. ``"RUNNING"``), or ``None``. + Default value is None. + progress : Any + Backend-reported progress value, or ``None``. Default value is None. + message : Any + Backend-reported stop/status message, or ``None``. Default value + is None. + """ + self._current_status = JobStatus( + state=state, progress=progress, message=message ) + # ---- chunk counters ---- + + @property + def chunks_total(self) -> int: + """Total number of chunks planned so far for this call.""" + return self._chunks_total + + @property + def chunks_done(self) -> int: + """Number of chunks that finished (successfully or not).""" + return self._chunks_done + + @property + def active_chunks(self) -> int: + """Number of chunk jobs currently in flight.""" + return self._active_chunks + + def add_planned_chunks(self, count: int) -> None: + """Register additional chunks planned for submission. + + Parameters + ---------- + count : int + Number of chunks about to be submitted for one quantum leaf. + """ + with self._lock: + self._chunks_total += count + + def mark_chunk_started(self) -> None: + """Mark one chunk job as submitted and in flight.""" + with self._lock: + self._active_chunks += 1 + + def mark_chunk_finished(self) -> None: + """Mark one in-flight chunk job as finished.""" + with self._lock: + self._active_chunks = max(0, self._active_chunks - 1) + self._chunks_done += 1 + + # ---- snapshots ---- + + def status_snapshot(self, future_done: bool = False) -> dict: + """Return the public status dict exposed through ``future.status()``. + + Parameters + ---------- + future_done : bool + Whether the owning future has already resolved. A resolved future + with no recorded backend status reports state ``"COMPLETE"``. + + Returns + ------- + dict + ``{"state", "progress", "message", "chunks_total", "chunks_done", + "active_chunks"}`` with the same semantics as before the CallState + refactor. + """ + js = self._current_status + return { + "state": ( + "COMPLETE" + if future_done and js is None + else (js.state if js else "IDLE") + ), + "progress": js.progress if js else 0.0, + "message": js.message if js else None, + "chunks_total": self._chunks_total, + "chunks_done": self._chunks_done, + "active_chunks": self._active_chunks, + } + + +class MerlinFuture(Future): + """Typed async handle returned by :meth:`MerlinProcessor.forward_async`. + + Extends ``torch.futures.Future[torch.Tensor]`` with the Merlin-specific + async contract that was previously monkey-patched onto plain Future + instances at runtime: remote-job visibility (:attr:`job_ids`), progress + reporting (:meth:`status`), and cooperative cancellation + (:meth:`cancel_remote`). All inherited Future behavior (``wait``, + ``done``, ``then``, ``value``, ...) is unchanged. + + Parameters + ---------- + call_state : CallState + Typed per-call state backing this handle. Job ids, chunk counters, + and backend status recorded during execution are read live from it. + cancel_all : Callable[[], None] + Processor-level callback cancelling all in-flight remote jobs; used + by :meth:`cancel_remote`. + """ + + def __init__(self, call_state: CallState, cancel_all: Callable[[], None]) -> None: + super().__init__() + self._call_state = call_state + self._cancel_all = cancel_all + + @property + def job_ids(self) -> list[str]: + """Remote job ids accumulated across chunks, in observation order. + + This is a live view of the underlying :class:`CallState` list: ids + recorded while chunks are polling appear here immediately. + """ + return self._call_state.job_ids + + def status(self) -> dict: + """Return the current progress and state of this call. + + Returns + ------- + dict + ``{"state", "progress", "message", "chunks_total", "chunks_done", + "active_chunks"}``. ``state`` is ``"IDLE"`` before the first + backend poll, the backend-reported state while polling, and + ``"COMPLETE"`` once the future resolves without a recorded + backend status. + """ + return self._call_state.status_snapshot(future_done=self.done()) + + def cancel_remote(self) -> None: + """Cooperatively cancel this call and its in-flight remote jobs. + + Requests cancellation on the per-call state (observed by chunk and + polling threads), cancels all active remote jobs best-effort, and + resolves this future with ``concurrent.futures.CancelledError`` if + it is not already done. Awaiting the future afterwards raises that + error. + """ + from concurrent.futures import CancelledError + + self._call_state.request_cancel() + self._cancel_all() + if not self.done(): + self.set_exception(CancelledError("Remote call was cancelled")) + _ALLOWED_STATE_TYPES = ( pcvl.StateVector, @@ -566,7 +806,7 @@ def __init__( # Build ONE initial processor to extract metadata (backend name, available commands). # Fresh processors will be created per chunk via _create_fresh_rp(). - _init_rp = self.session.build_remote_processor() + _init_rp = PercevalAdapter.build_from_session(self.session) remote_processor = _init_rp else: self.backend_kind = "remote_processor" @@ -586,11 +826,12 @@ def __init__( assert capability_processor is not None # Extract backend capabilities (name and available commands) - backend_name = capability_processor.name - available_cmds = capability_processor.available_commands + backend_name, available_cmds = PercevalAdapter.get_backend_capabilities( + capability_processor + ) self.backend_capabilities = BackendCapabilities( name=backend_name, - available_commands=tuple(available_cmds), + available_commands=available_cmds, ) # Check if commands list is empty and warn @@ -613,7 +854,7 @@ def __init__( self._token = self._extract_rp_token(remote_processor) if self._token is None: - raise ValueError( + raise TokenExtractionError( "Could not extract auth token from RemoteProcessor. " "Either pass token= to MerlinProcessor or call " "RemoteConfig.set_token() before constructing the " @@ -722,10 +963,7 @@ def cancel_all(self) -> None: with self._lock: jobs = list(self._active_jobs) for job in jobs: - cancel = getattr(job, "cancel", None) - if callable(cancel): - with suppress(Exception): - cancel() + PercevalAdapter.cancel_job(job) def forward( self, @@ -793,10 +1031,11 @@ def forward_async( *, nsample: int | None = None, timeout: float | None = None, - ) -> Future: + ) -> MerlinFuture: """Asynchronously execute a module against the configured Perceval backend. - Returns a ``torch.futures.Future`` that resolves to the output tensor. + Returns a :class:`MerlinFuture` (a ``torch.futures.Future`` subclass) + that resolves to the output tensor. Remote batches are automatically chunked and submitted with limited concurrency. Local processor inputs are kept as one Merlin-level batch and represented as Perceval sampler iterations using an isolated @@ -832,8 +1071,8 @@ def forward_async( Returns ------- - Future - ``torch.futures.Future[torch.Tensor]`` with extra attributes: + MerlinFuture + Typed ``torch.futures.Future[torch.Tensor]`` subclass exposing: - ``future.job_ids: list[str]`` — accumulates remote job IDs across chunks. @@ -885,48 +1124,8 @@ def forward_async( original_dtype = input.dtype layers: list[Any] = list(self._iter_layers_in_order(module)) - fut: Future = Future() - state = { - "cancel_requested": False, - "current_status": None, - "job_ids": [], - "chunks_total": 0, - "chunks_done": 0, - "active_chunks": 0, - "call_id": uuid.uuid4().hex[:8], - } - - def _cancel_remote(): - state["cancel_requested"] = True - self.cancel_all() - if not fut.done(): - try: - from concurrent.futures import CancelledError - except Exception: # pragma: no cover - - class CancelledError(RuntimeError): - pass - - fut.set_exception(CancelledError("Remote call was cancelled")) - - def _status(): - js = state.get("current_status") - return { - "state": ( - "COMPLETE" - if fut.done() and not js - else (js.get("state") if js else "IDLE") - ), - "progress": js.get("progress") if js else 0.0, - "message": js.get("message") if js else None, - "chunks_total": state["chunks_total"], - "chunks_done": state["chunks_done"], - "active_chunks": state["active_chunks"], - } - - fut.cancel_remote = _cancel_remote # type: ignore[attr-defined] - fut.status = _status # type: ignore[attr-defined] - fut.job_ids = state["job_ids"] # type: ignore[attr-defined] + state = CallState.new() + fut = MerlinFuture(state, self.cancel_all) def _run_pipeline(): try: @@ -943,7 +1142,7 @@ def _run_pipeline(): else: should_offload = False - if state["cancel_requested"]: + if state.cancel_requested: raise self._cancelled_error() if should_offload: @@ -970,7 +1169,7 @@ def _offload_quantum_layer_with_chunking( layer: MerlinModule, input_tensor: torch.Tensor, nsample: int | None, - state: dict, + state: CallState, deadline: float | None, ) -> torch.Tensor: """Execute a quantum layer through the selected backend route. @@ -1002,12 +1201,7 @@ def _offload_quantum_layer_with_chunking( layer, config, input_tensor, nsample, state, deadline ) - chunks: list[tuple[int, int]] = [] - start = 0 - while start < B: - end = min(start + self.microbatch_size, B) - chunks.append((start, end)) - start = end + chunks = BatchChunker.split_batch(B, self.microbatch_size) return self._run_chunks_pooled( layer, config, input_tensor, chunks, nsample, state, deadline ) @@ -1019,67 +1213,67 @@ def _run_chunks_pooled( input_tensor: torch.Tensor, chunks: list[tuple[int, int]], nsample: int | None, - state: dict, + state: CallState, deadline: float | None, ) -> torch.Tensor: - """Submit chunk jobs with limited concurrency and stitch results.""" - state["chunks_total"] += len(chunks) - outputs: list[torch.Tensor | None] = [None] * len(chunks) - errors: list[BaseException] = [] + """Submit chunk jobs with limited concurrency and stitch results. - total_chunks = len(chunks) - layer_name = getattr(layer, "name", layer.__class__.__name__) + Delegates chunk orchestration to :class:`BatchChunker`; kept as the + processor-level entry point so callers and tests interact with the + coordinator rather than the execution unit directly. + """ + return self._make_batch_chunker().run_chunks( + layer, config, input_tensor, chunks, nsample, state, deadline + ) - def _call(s: int, e: int, idx: int): - try: - base_label = ( - f"mer:{layer_name}:{state['call_id']}:{idx + 1}/{total_chunks}" - ) - t = self._run_chunk( - layer, - config, - input_tensor[s:e], - nsample, - state, - deadline, - job_base_label=base_label, - ) - outputs[idx] = t - except BaseException as ex: - errors.append(ex) - - in_flight = 0 - idx = 0 - futures: list[threading.Thread] = [] - while idx < len(chunks) or in_flight > 0: - while idx < len(chunks) and in_flight < self.chunk_concurrency: - s, e = chunks[idx] - with self._lock: - state["active_chunks"] += 1 - th = threading.Thread(target=_call, args=(s, e, idx), daemon=True) - th.start() - futures.append(th) - idx += 1 - in_flight += 1 - - for th in list(futures): - if not th.is_alive(): - futures.remove(th) - in_flight -= 1 - with self._lock: - state["active_chunks"] = max(0, state["active_chunks"] - 1) - state["chunks_done"] += 1 - - if deadline is not None and time.time() >= deadline: - self.cancel_all() - raise TimeoutError("Remote call timed out (remote cancel issued)") - - time.sleep(0.01) - - if errors: - raise errors[0] - - return torch.cat(outputs, dim=0) # type: ignore[arg-type] + # ---------------- Execution unit factories ---------------- + + def _make_batch_chunker(self) -> BatchChunker: + """Build the chunk orchestration unit wired to this processor. + + ``run_chunk`` is bound at call time so monkeypatched processor + methods remain observable, matching pre-extraction behavior. + """ + return BatchChunker( + run_chunk=self._run_chunk, + chunk_concurrency=self.chunk_concurrency, + cancel_all=self.cancel_all, + ) + + def _make_job_runner(self) -> RemoteJobRunner: + """Build the remote chunk execution unit wired to this processor. + + Dependencies are injected as bound methods and deferred lambdas so + that live attribute mutation (e.g. ``max_shots_per_call``) and + monkeypatched processor methods keep working exactly as before the + extraction. + """ + return RemoteJobRunner( + create_processor=self._create_fresh_rp, + get_available_commands=lambda: self.available_commands, + effective_sample_count=self._effective_sample_count, + get_max_shots_per_call=lambda: self.max_shots_per_call, + default_shots_per_call=self.DEFAULT_SHOTS_PER_CALL, + map_results=self._process_batch_results, + register_job=self._register_job, + unregister_job=self._unregister_job, + get_microbatch_limit=( + lambda: None if self.session is not None else self.microbatch_size + ), + max_retries=self._MAX_CHUNK_RETRIES, + job_name_max=self._JOB_NAME_MAX, + ) + + def _register_job(self, job: RemoteJob) -> None: + """Track a submitted job for cancellation and history.""" + with self._lock: + self._active_jobs.add(job) + self._job_history.append(job) + + def _unregister_job(self, job: RemoteJob) -> None: + """Remove a job from active cancellation tracking.""" + with self._lock: + self._active_jobs.discard(job) def _run_chunk( self, @@ -1087,111 +1281,37 @@ def _run_chunk( config: ValidatedLayerConfig, input_chunk: torch.Tensor, nsample: int | None, - state: dict, + state: CallState, deadline: float | None, job_base_label: str | None = None, ) -> torch.Tensor: - """Submit a single chunk job with retries and return the mapped tensor.""" - from concurrent.futures import CancelledError + """Submit a single chunk job with retries and return the mapped tensor. + Local backends are dispatched to :meth:`_run_chunk_local`; remote + chunk execution is delegated to :class:`RemoteJobRunner`. + """ if self.backend_kind == "local_processor": return self._run_chunk_local( layer, config, input_chunk, nsample, state, deadline ) - batch_size = input_chunk.shape[0] - if self.session is None and batch_size > self.microbatch_size: - raise ValueError( - f"Chunk size {batch_size} exceeds microbatch {self.microbatch_size}. " - "Please report this bug." - ) - - input_param_names = self._extract_input_params(config) - input_np = input_chunk.detach().cpu().numpy() - - # Pre-compute iteration params (cheap, only done once). - iteration_params: list[dict[str, float]] = [] - for i in range(batch_size): - circuit_params = {} - for j, param_name in enumerate(input_param_names): - circuit_params[param_name] = ( - float(input_np[i, j]) if j < input_chunk.shape[1] else 0.0 - ) - iteration_params.append(circuit_params) - - def _capped_name(base: str, cmd: str) -> str: - name = f"{base}:{cmd}" - name = "".join(ch if ch.isalnum() or ch in "-_:/=." else "_" for ch in name) - if len(name) <= self._JOB_NAME_MAX: - return name - h = f"{zlib.adler32(name.encode()):08x}" - keep = self._JOB_NAME_MAX - 1 - len(h) - if keep < 1: - return h[: self._JOB_NAME_MAX] - return name[:keep] + "~" + h - - last_error: BaseException | None = None - for attempt in range(self._MAX_CHUNK_RETRIES): - if state.get("cancel_requested"): - raise CancelledError("Remote call was cancelled") - if deadline is not None and time.time() >= deadline: - raise TimeoutError("Remote call timed out (remote cancel issued)") - - # Build a fresh RemoteProcessor and Sampler on each attempt so that - # a corrupted RP doesn't poison retries. - rp = self._create_fresh_rp() - rp.set_circuit(config.circuit) - if config.input_state: - input_state = pcvl.BasicState(config.input_state) - rp.with_input(input_state) - n_photons = sum(config.input_state) - rp.min_detected_photons_filter(n_photons) - - max_shots_arg = ( - self.DEFAULT_SHOTS_PER_CALL - if self.max_shots_per_call is None - else int(self.max_shots_per_call) - ) - sampler = Sampler(rp, max_shots_per_call=max_shots_arg) - sampler.clear_iterations() - for params in iteration_params: - sampler.add_iteration(circuit_params=params) - - job = None - try: - job, is_probability = self._submit_job( - sampler, nsample, job_base_label, _capped_name - ) - with self._lock: - self._active_jobs.add(job) - self._job_history.append(job) - - return self._poll_job( - job, state, deadline, batch_size, layer, nsample, is_probability - ) - except (CancelledError, TimeoutError, KeyboardInterrupt): - raise - except Exception as exc: - last_error = exc - if job is not None: - with self._lock: - self._active_jobs.discard(job) - logger.warning( - "Chunk attempt %d/%d failed: %s", - attempt + 1, - self._MAX_CHUNK_RETRIES, - exc, - ) - if attempt < self._MAX_CHUNK_RETRIES - 1: - time.sleep(min(1.0 * (2**attempt), 5.0)) - - raise RuntimeError( - f"Chunk failed after {self._MAX_CHUNK_RETRIES} attempts" - ) from last_error + return self._make_job_runner().run_chunk( + layer, + config, + input_chunk, + nsample, + state, + deadline, + job_base_label=job_base_label, + ) def _create_fresh_local_processor(self) -> AProcessor: """Create an isolated local Perceval processor for one execution. + Delegates to :meth:`PercevalAdapter.rebuild_local_processor`, which + snapshots the experiment metadata and rebuilds the processor from a + copied experiment and fresh backend. + Returns ------- AProcessor @@ -1204,135 +1324,7 @@ def _create_fresh_local_processor(self) -> AProcessor: If the configured local processor cannot be reconstructed safely. """ assert self.processor is not None - - experiment = getattr(self.processor, "experiment", None) - backend_object = getattr(self.processor, "backend", None) - experiment_copy = getattr(experiment, "copy", None) - if ( - experiment is None - or backend_object is None - or not callable(experiment_copy) - ): - raise TypeError( - "Local execution requires a Perceval processor with copyable " - "experiment state and a reconstructable local backend." - ) - - backend_name = getattr(backend_object, "name", None) - backend: str | object - if isinstance(backend_name, str): - backend = backend_name - else: - try: - backend = type(backend_object)() - except Exception as exc: - raise TypeError( - "Local processor backend cannot be reconstructed safely." - ) from exc - - experiment_metadata = self._snapshot_local_experiment_metadata(experiment) - copied_experiment = experiment_copy() - copied_experiment.clear_input_and_circuit() - - processor = Processor(backend, copied_experiment) - processor._merlin_local_experiment_metadata = experiment_metadata - return processor - - @staticmethod - def _snapshot_local_experiment_metadata( - experiment: Any, - ) -> _LocalExperimentMetadata: - """Copy non-circuit local experiment metadata before Perceval clears it. - - Parameters - ---------- - experiment : Any - Perceval experiment owned by the caller's local processor. - - Returns - ------- - _LocalExperimentMetadata - Deep-copied metadata that is independent from the caller's processor. - """ - - in_ports = tuple( - (port, tuple(modes)) - for port, modes in copy.deepcopy(experiment._in_ports).items() - ) - out_ports = tuple( - (port, tuple(modes)) - for port, modes in copy.deepcopy(experiment._out_ports).items() - ) - return _LocalExperimentMetadata( - circuit_size=int(experiment.circuit_size), - in_ports=in_ports, - out_ports=out_ports, - detectors=tuple(copy.deepcopy(experiment.detectors)), - detectors_injected=tuple(copy.deepcopy(experiment.detectors_injected)), - in_mode_type=tuple(copy.deepcopy(experiment._in_mode_type)), - out_mode_type=tuple(copy.deepcopy(experiment._out_mode_type)), - anon_herald_num=int(experiment._anon_herald_num), - postselection=copy.copy(experiment.post_select_fn), - ) - - @staticmethod - def _restore_local_experiment_metadata( - experiment: Any, metadata: _LocalExperimentMetadata - ) -> None: - """Restore local experiment metadata after the execution circuit is set. - - Parameters - ---------- - experiment : Any - Perceval experiment owned by the fresh local execution processor. - metadata : _LocalExperimentMetadata - Metadata copied from the caller's local processor. - - Raises - ------ - ValueError - If mode-indexed metadata cannot be applied to the execution circuit - because the circuit sizes differ. - """ - - if metadata.has_mode_metadata: - circuit_size = int(experiment.circuit_size) - if circuit_size != metadata.circuit_size: - raise ValueError( - "Local processor experiment metadata is tied to circuit size " - f"{metadata.circuit_size}, but the execution circuit has size " - f"{circuit_size}." - ) - experiment._in_ports = { - port: list(modes) for port, modes in metadata.in_ports - } - experiment._out_ports = { - port: list(modes) for port, modes in metadata.out_ports - } - experiment._detectors = list(metadata.detectors) - experiment.detectors_injected = list(metadata.detectors_injected) - experiment._in_mode_type = list(metadata.in_mode_type) - experiment._out_mode_type = list(metadata.out_mode_type) - experiment._anon_herald_num = metadata.anon_herald_num - - experiment._postselect = copy.copy(metadata.postselection) - experiment._circuit_changed() - - @staticmethod - def _copy_circuit_for_execution(circuit: pcvl.ACircuit) -> pcvl.ACircuit: - """Return a circuit copy for processor execution. - - Parameters - ---------- - circuit : pcvl.ACircuit - Circuit exported by the quantum layer. - - Returns - ------- - pcvl.ACircuit - Independent circuit object used by a single backend execution. - """ - return circuit.copy() + return PercevalAdapter.rebuild_local_processor(self.processor) def _run_chunk_local( self, @@ -1340,7 +1332,7 @@ def _run_chunk_local( config: ValidatedLayerConfig, input_chunk: torch.Tensor, nsample: int | None, - state: dict, + state: CallState, deadline: float | None, ) -> torch.Tensor: """Execute a local AProcessor batch with an isolated processor. @@ -1368,9 +1360,9 @@ def _run_chunk_local( nsample : int | None Number of samples per row. ``None`` or ``<= 0`` triggers exact probability computation when the backend supports ``"probs"``. - state : dict - Shared call-state dictionary. Must contain the key - ``"cancel_requested"`` (bool). + state : CallState + Typed per-call state; its ``cancel_requested`` flag is checked + cooperatively before and after execution. deadline : float | None Absolute wall-clock deadline (``time.time()`` seconds). ``None`` means no deadline. @@ -1384,14 +1376,14 @@ def _run_chunk_local( Raises ------ concurrent.futures.CancelledError - If ``state["cancel_requested"]`` is ``True`` before or after + If ``state.cancel_requested`` is ``True`` before or after execution. TimeoutError If ``deadline`` has elapsed before or after execution. """ from concurrent.futures import CancelledError - if state.get("cancel_requested"): + if state.cancel_requested: raise CancelledError("Local call was cancelled") if deadline is not None and time.time() >= deadline: raise TimeoutError("Local call timed out") @@ -1412,41 +1404,39 @@ def _run_chunk_local( iteration_params.append(circuit_params) processor = self._create_fresh_local_processor() - processor.set_circuit(self._copy_circuit_for_execution(config.circuit)) - experiment_metadata = getattr( - processor, "_merlin_local_experiment_metadata", None + PercevalAdapter.configure_processor( + processor, PercevalAdapter.copy_circuit(config.circuit), None ) - if isinstance(experiment_metadata, _LocalExperimentMetadata): - self._restore_local_experiment_metadata( - processor.experiment, experiment_metadata + experiment_snapshot = getattr(processor, LOCAL_EXPERIMENT_SNAPSHOT_ATTR, None) + if isinstance(experiment_snapshot, LocalExperimentSnapshot): + PercevalAdapter.restore_experiment( + processor.experiment, experiment_snapshot ) - if config.input_state: - input_state = pcvl.BasicState(config.input_state) - processor.with_input(input_state) - n_photons = sum(config.input_state) - processor.min_detected_photons_filter(n_photons) + PercevalAdapter.set_input(processor, config.input_state) - sampler = Sampler(processor, max_shots_per_call=self.max_shots_per_call) - sampler.clear_iterations() - for params in iteration_params: - sampler.add_iteration(circuit_params=params) + sampler = PercevalAdapter.create_sampler( + processor, self.max_shots_per_call, iteration_params + ) is_probability = ("probs" in self.available_commands) and ( nsample is None or int(nsample) <= 0 ) if is_probability: - raw_results = sampler.probs.execute_sync() + raw_results = PercevalAdapter.execute_sync(sampler, "probs") else: use_shots = self._effective_sample_count(nsample) if "sample_count" in self.available_commands: - raw_results = sampler.sample_count.execute_sync(max_samples=use_shots) + cmd = "sample_count" elif "samples" in self.available_commands: - raw_results = sampler.samples.execute_sync(max_samples=use_shots) + cmd = "samples" else: - raw_results = sampler.sample_count.execute_sync(max_samples=use_shots) + cmd = "sample_count" + raw_results = PercevalAdapter.execute_sync( + sampler, cmd, max_samples=use_shots + ) - if state.get("cancel_requested"): + if state.cancel_requested: raise CancelledError("Local call was cancelled") if deadline is not None and time.time() >= deadline: raise TimeoutError("Local call timed out") @@ -1492,73 +1482,14 @@ def _submit_job(self, sampler, nsample, job_base_label, _capped_name): - **bool**: ``is_probability`` flag indicating execution mode: ``True`` if using exact probabilities, ``False`` if sampling. """ - is_probability = ("probs" in self.available_commands) and ( - nsample is None or int(nsample) <= 0 + return self._make_job_runner().submit_job( + sampler, nsample, job_base_label, _capped_name ) - if is_probability: - job = sampler.probs - cmd = "probs" - if job_base_label: - job.name = _capped_name(job_base_label, cmd) - self._ensure_serializable_sampler_iterator(job, sampler) - return job.execute_async(), is_probability - - use_shots = self._effective_sample_count(nsample) - - if "sample_count" in self.available_commands: - job = sampler.sample_count - cmd = "sample_count" - elif "samples" in self.available_commands: - job = sampler.samples - cmd = "samples" - else: - job = sampler.sample_count - cmd = "sample_count" - - if job_base_label: - job.name = _capped_name(job_base_label, cmd) - self._ensure_serializable_sampler_iterator(job, sampler) - return job.execute_async(max_samples=use_shots), is_probability - - @staticmethod - def _ensure_serializable_sampler_iterator(job: RemoteJob, sampler: Sampler) -> None: - """Replace Perceval 1.2 iterator objects with JSON-serializable data. - - Parameters - ---------- - job : RemoteJob - Prepared Perceval remote job whose private request payload may contain - a sampler iterator. - sampler : Sampler - Perceval sampler used to prepare the job. - - Notes - ----- - Perceval 1.1 stores sampler iterations as a plain list. Perceval 1.2 - stores them in a ``ParameterIterator`` object, but the Scaleway session - handler still serializes ``payload["payload"]`` with ``json.dumps``. - Until Perceval exposes a public serializer for that object, Merlin - normalizes the remote-job payload back to the list shape accepted by the - cloud side. - """ - iterator = getattr(sampler, "_iterator", None) - iterations = getattr(iterator, "iterations", None) - if not iterations: - return - - request_data = getattr(job, "_request_data", None) - if not isinstance(request_data, dict): - return - - payload = request_data.get("payload") - if isinstance(payload, dict) and payload.get("iterator") is iterator: - payload["iterator"] = list(iterations) - def _poll_job( self, job: RemoteJob, - state: dict, + state: CallState, deadline: float | None, batch_size: int, layer: MerlinModule, @@ -1576,8 +1507,8 @@ def _poll_job( ---------- job : RemoteJob Submitted Perceval job to poll. - state : dict - Shared state dict tracking cancellation, chunks, job IDs, etc. + state : CallState + Typed per-call state tracking cancellation, chunks, job IDs, etc. deadline : float | None Absolute time (seconds) when execution should timeout. batch_size : int @@ -1598,87 +1529,9 @@ def _poll_job( from the remote job results. Probability vs. sample interpretation is determined by ``is_probability``. """ - from concurrent.futures import CancelledError - - _MAX_NON_DICT_RETRIES = 60 # 60 * 0.1s = 6s - non_dict_retries = 0 - sleep_ms = 50 - while True: - if state.get("cancel_requested"): - cancel = getattr(job, "cancel", None) - if callable(cancel): - with suppress(Exception): - cancel() - raise CancelledError("Remote call was cancelled") - - if deadline is not None and time.time() >= deadline: - cancel = getattr(job, "cancel", None) - if callable(cancel): - with suppress(Exception): - cancel() - raise TimeoutError("Remote call timed out (remote cancel issued)") - - s = getattr(job, "status", None) - state["current_status"] = { - "state": getattr(s, "state", None) if s else None, - "progress": getattr(s, "progress", None) if s else None, - "message": getattr(s, "stop_message", None) if s else None, - } - - job_id = getattr(job, "id", None) or getattr(job, "job_id", None) - if job_id is not None and job_id not in state["job_ids"]: - state["job_ids"].append(job_id) - - if getattr(job, "is_failed", False): - msg = state["current_status"].get("message") - if msg and "Cancel requested" in str(msg): - with self._lock: - self._active_jobs.discard(job) - raise CancelledError("Remote call was cancelled") - with self._lock: - self._active_jobs.discard(job) - raise RuntimeError( - f"Remote job failed: {msg or 'unknown error'} (job_id={job_id!r})" - ) - - if getattr(job, "is_complete", False): - try: - raw = job.get_results() - except RuntimeError as ex: - msg = str(ex) - if "Results are not available" in msg: - time.sleep(0.05) - continue - if "Cancel requested" in msg: - with self._lock: - self._active_jobs.discard(job) - raise CancelledError("Remote call was cancelled") - raise - - if isinstance(raw, dict): - with self._lock: - self._active_jobs.discard(job) - return self._process_batch_results( - raw, batch_size, layer, nsample, is_probability - ) - - # The backend sometimes reports completion before the dict - # payload is actually available. Re-poll the same job for a - # bounded window before giving up to the outer retry loop. - non_dict_retries += 1 - if non_dict_retries >= _MAX_NON_DICT_RETRIES: - with self._lock: - self._active_jobs.discard(job) - raise RuntimeError( - f"Job complete but results were not a dict after " - f"{_MAX_NON_DICT_RETRIES} re-polls; " - f"job_id={job_id!r}, type={type(raw)}, value={raw!r}" - ) - time.sleep(0.1) - continue - - time.sleep(sleep_ms / 1000.0) - sleep_ms = min(sleep_ms * 2, 400) + return self._make_job_runner().poll_job( + job, state, deadline, batch_size, layer, nsample, is_probability + ) # ---------------- Per-call RP pool helpers ---------------- @@ -1715,7 +1568,7 @@ def _create_fresh_rp(self) -> RemoteProcessor: """ if self.session is not None: # Session path: create a fresh processor from the session - return self.session.build_remote_processor() + return PercevalAdapter.build_from_session(self.session) if self.remote_processor is None: raise RuntimeError( "Fresh RemoteProcessor creation is only available for remote " @@ -1730,61 +1583,21 @@ def _clone_remote_processor(self, rp: RemoteProcessor) -> RemoteProcessor: """Create a sibling RemoteProcessor with its own RPC handler (thread-safe). Forwards the token extracted at init time so that inline-token - RemoteProcessors are cloned correctly. + RemoteProcessors are cloned correctly. Delegates the Perceval + handler access to :class:`PercevalAdapter`. """ - return RemoteProcessor( - name=rp.name, - token=self._token, - url=( - rp.get_rpc_handler().url - if hasattr(rp.get_rpc_handler(), "url") - else None - ), - proxies=rp.proxies, - ) + return PercevalAdapter.clone_remote_processor(rp, self._token) @staticmethod def _extract_rp_token(rp: RemoteProcessor) -> str | None: """Extract the auth token from a RemoteProcessor. - Perceval stores the token on the RPC handler as ``handler.token`` - and also embeds it in ``handler.headers['Authorization']``. We - probe both locations so that inline-token and global-config - ``RemoteProcessor`` instances are both handled. - - As a last resort, falls back to ``RemoteConfig().get_token()``. - Returns ``None`` only if every strategy fails. + Delegates to :meth:`PercevalAdapter.extract_token`, which probes the + RPC handler token attributes, the Authorization header, and the + global ``RemoteConfig`` fallback. Returns ``None`` only if every + strategy fails. """ - try: - handler = rp.get_rpc_handler() - except Exception: - handler = None - - if handler is not None: - # Primary: handler.token (set by RPCHandler.__init__) - for attr in ("token", "_token", "auth_token"): - val = getattr(handler, attr, None) - if isinstance(val, str) and val: - return val - - # Fallback: parse 'Bearer ' from Authorization header - headers = getattr(handler, "headers", None) - if isinstance(headers, dict): - auth = headers.get("Authorization", "") - if auth.startswith("Bearer ") and len(auth) > 7: - return auth[7:] - - # Last resort: check the global config - try: - from perceval.runtime import RemoteConfig - - global_token = (RemoteConfig().get_token() or "").strip() - if global_token: - return global_token - except Exception: - logger.debug("RemoteConfig token lookup failed", exc_info=True) - - return None + return PercevalAdapter.extract_token(rp) def _iter_layers_in_order(self, module: nn.Module) -> Iterable[nn.Module]: """Yield execution leaves in deterministic order. @@ -2074,13 +1887,9 @@ def estimate_required_shots_per_input( ) config = ValidatedLayerConfig(layer.export_config()) child_rp = self._create_fresh_rp() - child_rp.set_circuit(config.circuit) - - if config.input_state: - input_state = pcvl.BasicState(config.input_state) - child_rp.with_input(input_state) - n_photons = sum(config.input_state) - child_rp.min_detected_photons_filter(n_photons) + PercevalAdapter.configure_processor( + child_rp, config.circuit, config.input_state + ) input_param_names = self._extract_input_params(config) @@ -2099,8 +1908,8 @@ def estimate_required_shots_per_input( last_ex: Exception | None = None for _attempt in range(self._MAX_ESTIMATOR_RETRIES): try: - est = child_rp.estimate_required_shots( - desired_samples_per_input, param_values=param_values + est = PercevalAdapter.estimate_required_shots( + child_rp, desired_samples_per_input, param_values ) break except requests.exceptions.ReadTimeout as ex: diff --git a/merlin/core/perceval_adapter.py b/merlin/core/perceval_adapter.py new file mode 100644 index 00000000..352904bf --- /dev/null +++ b/merlin/core/perceval_adapter.py @@ -0,0 +1,651 @@ +"""Adapter layer isolating Merlin from Perceval internal APIs (PML-306). + +MerlinProcessor and the execution units historically reached directly into +Perceval internals: RPC handler token/URL attributes, sampler command objects +(``probs`` / ``sample_count`` / ``samples``), remote job status fields, local +experiment private state, and ``RemoteConfig``. Any Perceval version bump that +renames or restructures those internals could silently break Merlin at runtime. + +:class:`PercevalAdapter` owns every such access. The rest of Merlin talks to +this facade, so a Perceval API change is localized to this module. + +The adapter is stateless (static methods) and duck-typed: it reads the same +attributes Perceval exposes today, which also makes it independently testable +with plain fakes. +""" + +from __future__ import annotations + +import copy +import logging +from contextlib import suppress +from dataclasses import dataclass +from typing import Any + +import perceval as pcvl +from perceval.algorithm import Sampler +from perceval.runtime import AProcessor, Processor, RemoteJob, RemoteProcessor +from perceval.runtime.session import ISession + +logger = logging.getLogger(__name__) + +#: Attribute name used to carry a local experiment snapshot on a rebuilt +#: processor between :meth:`PercevalAdapter.rebuild_local_processor` and +#: :meth:`PercevalAdapter.restore_experiment`. +LOCAL_EXPERIMENT_SNAPSHOT_ATTR = "_merlin_local_experiment_metadata" + + +class TokenExtractionError(ValueError): + """Raised when no auth token can be resolved for a RemoteProcessor. + + Subclasses ``ValueError`` so existing callers catching the historical + exception type keep working. + """ + + +class RemoteJobFailedError(RuntimeError): + """Raised when a remote Perceval job reports failure. + + Subclasses ``RuntimeError`` so existing callers catching the historical + exception type keep working. + """ + + +@dataclass(frozen=True) +class JobStatusSnapshot: + """Merlin-normalized view of a Perceval remote job's status. + + All ``getattr`` guards against Perceval job internals live in + :meth:`PercevalAdapter.job_snapshot`; consumers only see these fields. + """ + + job_id: str | None + state: Any + progress: Any + stop_message: Any + is_complete: bool + is_failed: bool + + +@dataclass(frozen=True) +class LocalExperimentSnapshot: + """Experiment-level state that must survive local circuit replacement. + + Captures the Perceval experiment private state (ports, detectors, mode + types, heralds, postselection) that ``clear_input_and_circuit()`` wipes. + """ + + circuit_size: int + in_ports: tuple[tuple[Any, tuple[int, ...]], ...] + out_ports: tuple[tuple[Any, tuple[int, ...]], ...] + detectors: tuple[Any | None, ...] + detectors_injected: tuple[int, ...] + in_mode_type: tuple[Any, ...] + out_mode_type: tuple[Any, ...] + anon_herald_num: int + postselection: Any + + @property + def has_mode_metadata(self) -> bool: + """Return whether metadata is tied to a concrete circuit mode layout.""" + + return ( + bool(self.in_ports) + or bool(self.out_ports) + or any(detector is not None for detector in self.detectors) + or bool(self.detectors_injected) + or self.postselection != pcvl.PostSelect() + ) + + +class PercevalAdapter: + """Stateless facade owning all direct Perceval-internal access.""" + + # ------------------------------------------------------------------ + # Token / handler / URL + # ------------------------------------------------------------------ + + @staticmethod + def extract_token(rp: RemoteProcessor) -> str | None: + """Extract the auth token from a RemoteProcessor. + + Perceval stores the token on the RPC handler as ``handler.token`` + and also embeds it in ``handler.headers['Authorization']``. We + probe both locations so that inline-token and global-config + ``RemoteProcessor`` instances are both handled. + + As a last resort, falls back to ``RemoteConfig().get_token()``. + + Parameters + ---------- + rp : perceval.runtime.RemoteProcessor + Remote processor to probe for authentication material. + + Returns + ------- + str | None + The resolved token, or ``None`` if every strategy fails. + """ + try: + handler = rp.get_rpc_handler() + except Exception: + handler = None + + if handler is not None: + # Primary: handler.token (set by RPCHandler.__init__) + for attr in ("token", "_token", "auth_token"): + val = getattr(handler, attr, None) + if isinstance(val, str) and val: + return val + + # Fallback: parse 'Bearer ' from Authorization header + headers = getattr(handler, "headers", None) + if isinstance(headers, dict): + auth = headers.get("Authorization", "") + if auth.startswith("Bearer ") and len(auth) > 7: + return auth[7:] + + # Last resort: check the global config + try: + from perceval.runtime import RemoteConfig + + global_token = (RemoteConfig().get_token() or "").strip() + if global_token: + return global_token + except Exception: + logger.debug("RemoteConfig token lookup failed", exc_info=True) + + return None + + @staticmethod + def get_url(rp: RemoteProcessor) -> str | None: + """Return the RPC handler URL of a RemoteProcessor, if exposed. + + Parameters + ---------- + rp : perceval.runtime.RemoteProcessor + Remote processor whose RPC handler is inspected. + + Returns + ------- + str | None + The handler URL, or ``None`` when the handler has no ``url`` + attribute. + """ + handler = rp.get_rpc_handler() + return handler.url if hasattr(handler, "url") else None + + # ------------------------------------------------------------------ + # Processor creation and configuration + # ------------------------------------------------------------------ + + @staticmethod + def clone_remote_processor( + rp: RemoteProcessor, token: str | None + ) -> RemoteProcessor: + """Create a sibling RemoteProcessor with its own RPC handler. + + Forwards the provided token so that inline-token RemoteProcessors + are cloned correctly. + + Parameters + ---------- + rp : perceval.runtime.RemoteProcessor + Processor whose platform name, URL, and proxies are copied. + token : str | None + Authentication token forwarded to the clone. + + Returns + ------- + perceval.runtime.RemoteProcessor + Independent processor targeting the same platform. + """ + return RemoteProcessor( + name=rp.name, + token=token, + url=PercevalAdapter.get_url(rp), + proxies=rp.proxies, + ) + + @staticmethod + def build_from_session(session: ISession) -> RemoteProcessor: + """Build a fresh RemoteProcessor from a Perceval session. + + Parameters + ---------- + session : perceval.runtime.session.ISession + Provider session (e.g. Scaleway) able to build processors. + + Returns + ------- + perceval.runtime.RemoteProcessor + Independent processor with its own handler state. + """ + return session.build_remote_processor() + + @staticmethod + def get_backend_capabilities(processor: AProcessor) -> tuple[str, tuple[str, ...]]: + """Return the backend platform name and available command snapshot. + + Parameters + ---------- + processor : perceval.runtime.AProcessor + Local or remote processor to inspect. + + Returns + ------- + tuple[str, tuple[str, ...]] + Platform name and immutable snapshot of supported commands. + """ + return processor.name, tuple(processor.available_commands) + + @staticmethod + def configure_processor( + processor: AProcessor, + circuit: pcvl.ACircuit, + input_state: Any, + ) -> None: + """Set the circuit and, when provided, the input state and photon filter. + + Parameters + ---------- + processor : AProcessor + Processor (local or remote) to configure. + circuit : pcvl.ACircuit + Circuit to install. + input_state : Any + Sequence of photon counts per mode, or falsy to skip input setup. + When set, ``min_detected_photons_filter`` is set to the total + photon count. + """ + processor.set_circuit(circuit) + PercevalAdapter.set_input(processor, input_state) + + @staticmethod + def set_input(processor: AProcessor, input_state: Any) -> None: + """Set the input state and matching photon filter, if provided. + + Split out from :meth:`configure_processor` because the local + execution path must restore experiment metadata between installing + the circuit and setting the input. + + Parameters + ---------- + processor : perceval.runtime.AProcessor + Processor to receive the input state. + input_state : Any + Sequence of photon counts per mode, or falsy to skip input setup. + """ + if input_state: + state = pcvl.BasicState(input_state) + processor.with_input(state) + n_photons = sum(input_state) + processor.min_detected_photons_filter(n_photons) + + @staticmethod + def copy_circuit(circuit: pcvl.ACircuit) -> pcvl.ACircuit: + """Return an independent copy of a circuit for one execution. + + Parameters + ---------- + circuit : pcvl.ACircuit + Circuit exported by the quantum layer. + + Returns + ------- + pcvl.ACircuit + Independent circuit object used by a single backend execution. + """ + return circuit.copy() + + @staticmethod + def estimate_required_shots( + rp: RemoteProcessor, desired_samples: int, param_values: dict[str, float] + ) -> int | None: + """Ask the remote platform estimator for the required shot count. + + Parameters + ---------- + rp : perceval.runtime.RemoteProcessor + Configured remote processor exposing the platform estimator. + desired_samples : int + Target number of usable samples. + param_values : dict[str, float] + Circuit parameter values for the input row being estimated. + + Returns + ------- + int | None + Estimated shots, or ``None`` when the platform gives no answer. + """ + return rp.estimate_required_shots(desired_samples, param_values=param_values) + + # ------------------------------------------------------------------ + # Samplers + # ------------------------------------------------------------------ + + @staticmethod + def create_sampler( + processor: AProcessor, + max_shots_per_call: int, + iterations: list[dict[str, float]], + ) -> Sampler: + """Create a Sampler on ``processor`` loaded with the given iterations. + + Parameters + ---------- + processor : perceval.runtime.AProcessor + Configured processor (circuit and input already set). + max_shots_per_call : int + Shot cap forwarded to the Perceval sampler. + iterations : list[dict[str, float]] + One circuit-parameter mapping per batch row. + + Returns + ------- + perceval.algorithm.Sampler + Sampler ready for command dispatch. + """ + sampler = Sampler(processor, max_shots_per_call=max_shots_per_call) + sampler.clear_iterations() + for params in iterations: + sampler.add_iteration(circuit_params=params) + return sampler + + @staticmethod + def submit_async( + sampler: Sampler, + command: str, + name: str | None = None, + max_samples: int | None = None, + ) -> RemoteJob: + """Submit a sampler command asynchronously and return the job handle. + + Parameters + ---------- + sampler : Sampler + Sampler prepared with circuit and iterations. + command : str + Sampler command to dispatch: ``"probs"``, ``"sample_count"``, + or ``"samples"``. + name : str | None + Remote job name to assign before submission, if any. + max_samples : int | None + Shots to request. ``None`` submits without a shot argument + (exact probabilities). + + Returns + ------- + perceval.runtime.RemoteJob + Handle of the submitted asynchronous job. + """ + job = getattr(sampler, command) + if name: + job.name = name + PercevalAdapter.ensure_serializable_sampler_iterator(job, sampler) + if max_samples is None: + return job.execute_async() + return job.execute_async(max_samples=max_samples) + + @staticmethod + def execute_sync( + sampler: Sampler, + command: str, + max_samples: int | None = None, + ) -> Any: + """Execute a sampler command synchronously and return the raw results. + + Parameters + ---------- + sampler : perceval.algorithm.Sampler + Sampler prepared with circuit and iterations. + command : str + Sampler command to dispatch: ``"probs"``, ``"sample_count"``, + or ``"samples"``. + max_samples : int | None + Shots to request. ``None`` executes without a shot argument + (exact probabilities). Default value is None. + + Returns + ------- + Any + Raw Perceval results object for the executed command. + """ + job = getattr(sampler, command) + if max_samples is None: + return job.execute_sync() + return job.execute_sync(max_samples=max_samples) + + @staticmethod + def ensure_serializable_sampler_iterator(job: RemoteJob, sampler: Sampler) -> None: + """Replace Perceval 1.2 iterator objects with JSON-serializable data. + + Parameters + ---------- + job : perceval.runtime.RemoteJob + Prepared job whose private request payload may hold an iterator. + sampler : perceval.algorithm.Sampler + Sampler used to prepare the job. + + Notes + ----- + Perceval 1.1 stores sampler iterations as a plain list. Perceval 1.2 + stores them in a ``ParameterIterator`` object, but the Scaleway session + handler still serializes ``payload["payload"]`` with ``json.dumps``. + Until Perceval exposes a public serializer for that object, Merlin + normalizes the remote-job payload back to the list shape accepted by + the cloud side. + """ + iterator = getattr(sampler, "_iterator", None) + iterations = getattr(iterator, "iterations", None) + if not iterations: + return + + request_data = getattr(job, "_request_data", None) + if not isinstance(request_data, dict): + return + + payload = request_data.get("payload") + if isinstance(payload, dict) and payload.get("iterator") is iterator: + payload["iterator"] = list(iterations) + + # ------------------------------------------------------------------ + # Jobs + # ------------------------------------------------------------------ + + @staticmethod + def job_snapshot(job: RemoteJob) -> JobStatusSnapshot: + """Read a job's status fields into a Merlin-normalized snapshot. + + Parameters + ---------- + job : perceval.runtime.RemoteJob + Job to inspect. Missing attributes map to ``None``/``False``. + + Returns + ------- + JobStatusSnapshot + Immutable view of the job's id, state, and completion flags. + """ + status = getattr(job, "status", None) + return JobStatusSnapshot( + job_id=getattr(job, "id", None) or getattr(job, "job_id", None), + state=getattr(status, "state", None) if status else None, + progress=getattr(status, "progress", None) if status else None, + stop_message=getattr(status, "stop_message", None) if status else None, + is_complete=bool(getattr(job, "is_complete", False)), + is_failed=bool(getattr(job, "is_failed", False)), + ) + + @staticmethod + def get_results(job: RemoteJob) -> Any: + """Retrieve a job's raw results, propagating Perceval errors. + + Parameters + ---------- + job : perceval.runtime.RemoteJob + Completed job to read. + + Returns + ------- + Any + Raw Perceval results object. + + Raises + ------ + RuntimeError + Propagated unchanged from Perceval (e.g. results not yet + available, cancel requested); the polling loop interprets it. + """ + return job.get_results() + + @staticmethod + def cancel_job(job: RemoteJob) -> None: + """Request best-effort cancellation of a job, swallowing errors. + + Parameters + ---------- + job : perceval.runtime.RemoteJob + Job to cancel. Objects without a callable ``cancel`` are ignored; + cancellation errors are suppressed by design (best-effort path). + """ + cancel = getattr(job, "cancel", None) + if callable(cancel): + with suppress(Exception): + cancel() + + # ------------------------------------------------------------------ + # Local processors + # ------------------------------------------------------------------ + + @staticmethod + def rebuild_local_processor(processor: AProcessor) -> AProcessor: + """Create an isolated local Perceval processor for one execution. + + The returned processor carries a :class:`LocalExperimentSnapshot` + under :data:`LOCAL_EXPERIMENT_SNAPSHOT_ATTR` so the caller can restore + experiment metadata after installing the execution circuit. + + Parameters + ---------- + processor : perceval.runtime.AProcessor + Local processor whose experiment and backend are copied. + + Returns + ------- + perceval.runtime.AProcessor + Fresh local processor with copied non-circuit experiment state + and a fresh backend instance. + + Raises + ------ + TypeError + If the configured local processor cannot be reconstructed safely. + """ + experiment = getattr(processor, "experiment", None) + backend_object = getattr(processor, "backend", None) + experiment_copy = getattr(experiment, "copy", None) + if ( + experiment is None + or backend_object is None + or not callable(experiment_copy) + ): + raise TypeError( + "Local execution requires a Perceval processor with copyable " + "experiment state and a reconstructable local backend." + ) + + backend_name = getattr(backend_object, "name", None) + backend: str | object + if isinstance(backend_name, str): + backend = backend_name + else: + try: + backend = type(backend_object)() + except Exception as exc: + raise TypeError( + "Local processor backend cannot be reconstructed safely." + ) from exc + + experiment_snapshot = PercevalAdapter.snapshot_experiment(experiment) + copied_experiment = experiment_copy() + copied_experiment.clear_input_and_circuit() + + fresh_processor = Processor(backend, copied_experiment) + setattr(fresh_processor, LOCAL_EXPERIMENT_SNAPSHOT_ATTR, experiment_snapshot) + return fresh_processor + + @staticmethod + def snapshot_experiment(experiment: Any) -> LocalExperimentSnapshot: + """Copy non-circuit local experiment metadata before Perceval clears it. + + Parameters + ---------- + experiment : Any + Perceval experiment owned by the caller's local processor. + + Returns + ------- + LocalExperimentSnapshot + Deep-copied metadata that is independent from the caller's + processor. + """ + in_ports = tuple( + (port, tuple(modes)) + for port, modes in copy.deepcopy(experiment._in_ports).items() + ) + out_ports = tuple( + (port, tuple(modes)) + for port, modes in copy.deepcopy(experiment._out_ports).items() + ) + return LocalExperimentSnapshot( + circuit_size=int(experiment.circuit_size), + in_ports=in_ports, + out_ports=out_ports, + detectors=tuple(copy.deepcopy(experiment.detectors)), + detectors_injected=tuple(copy.deepcopy(experiment.detectors_injected)), + in_mode_type=tuple(copy.deepcopy(experiment._in_mode_type)), + out_mode_type=tuple(copy.deepcopy(experiment._out_mode_type)), + anon_herald_num=int(experiment._anon_herald_num), + postselection=copy.copy(experiment.post_select_fn), + ) + + @staticmethod + def restore_experiment(experiment: Any, snapshot: LocalExperimentSnapshot) -> None: + """Restore local experiment metadata after the execution circuit is set. + + Parameters + ---------- + experiment : Any + Perceval experiment owned by the fresh local execution processor. + snapshot : LocalExperimentSnapshot + Metadata copied from the caller's local processor. + + Raises + ------ + ValueError + If mode-indexed metadata cannot be applied to the execution + circuit because the circuit sizes differ. + """ + if snapshot.has_mode_metadata: + circuit_size = int(experiment.circuit_size) + if circuit_size != snapshot.circuit_size: + raise ValueError( + "Local processor experiment metadata is tied to circuit size " + f"{snapshot.circuit_size}, but the execution circuit has size " + f"{circuit_size}." + ) + experiment._in_ports = { + port: list(modes) for port, modes in snapshot.in_ports + } + experiment._out_ports = { + port: list(modes) for port, modes in snapshot.out_ports + } + experiment._detectors = list(snapshot.detectors) + experiment.detectors_injected = list(snapshot.detectors_injected) + experiment._in_mode_type = list(snapshot.in_mode_type) + experiment._out_mode_type = list(snapshot.out_mode_type) + experiment._anon_herald_num = snapshot.anon_herald_num + + experiment._postselect = copy.copy(snapshot.postselection) + experiment._circuit_changed() diff --git a/tests/core/test_call_state.py b/tests/core/test_call_state.py new file mode 100644 index 00000000..02c02b84 --- /dev/null +++ b/tests/core/test_call_state.py @@ -0,0 +1,216 @@ +"""No-cloud unit tests for the typed CallState per-call state object (PML-303).""" + +from __future__ import annotations + +import threading +import time + +from merlin.core.merlin_processor import CallState, JobStatus + + +class TestCallStateCreation: + def test_new_assigns_short_unique_call_id(self): + """CallState.new() assigns an 8-char hex call id, unique per call.""" + first = CallState.new() + second = CallState.new() + + assert isinstance(first.call_id, str) + assert len(first.call_id) == 8 + int(first.call_id, 16) # hex-parsable + assert first.call_id != second.call_id + + def test_initial_state_is_idle_and_empty(self): + """A fresh state carries no cancellation, status, jobs, or chunks.""" + state = CallState.new() + + assert state.cancel_requested is False + assert state.current_status is None + assert state.job_ids == [] + assert state.chunks_total == 0 + assert state.chunks_done == 0 + assert state.active_chunks == 0 + + +class TestCancelPropagation: + def test_request_cancel_sets_flag(self): + """request_cancel() flips the cooperative cancellation flag.""" + state = CallState.new() + + state.request_cancel() + + assert state.cancel_requested is True + + def test_request_cancel_is_idempotent(self): + """Repeated cancellation requests keep the flag set.""" + state = CallState.new() + + state.request_cancel() + state.request_cancel() + + assert state.cancel_requested is True + + def test_cancel_is_visible_across_threads(self): + """A cancel requested on one thread is observed by another.""" + state = CallState.new() + observed = threading.Event() + + def worker(): + while not state.cancel_requested: + time.sleep(0.001) + observed.set() + + thread = threading.Thread(target=worker, daemon=True) + thread.start() + state.request_cancel() + assert observed.wait(timeout=5.0) + thread.join(timeout=5.0) + + +class TestJobIdRecording: + def test_record_job_id_appends(self): + """Recorded job ids accumulate in observation order.""" + state = CallState.new() + + state.record_job_id("job-1") + state.record_job_id("job-2") + + assert state.job_ids == ["job-1", "job-2"] + + def test_record_job_id_deduplicates(self): + """Re-observing the same job id (repeated polls) records it once.""" + state = CallState.new() + + state.record_job_id("job-123") + state.record_job_id("job-123") + state.record_job_id("job-456") + state.record_job_id("job-123") + + assert state.job_ids == ["job-123", "job-456"] + + def test_job_ids_list_identity_is_stable(self): + """The job_ids list object is stable so future.job_ids stays live.""" + state = CallState.new() + shared_reference = state.job_ids + + state.record_job_id("job-1") + + assert shared_reference is state.job_ids + assert shared_reference == ["job-1"] + + +class TestChunkCounters: + def test_ticket_example_transitions(self): + """The PML-303 minimal example holds: one started+finished chunk.""" + state = CallState.new() + state.record_job_id("job-123") + state.set_current_status(state="RUNNING", progress=0.5, message=None) + state.mark_chunk_started() + state.mark_chunk_finished() + + snapshot = state.status_snapshot() + + assert snapshot["chunks_done"] == 1 + assert snapshot["active_chunks"] == 0 + + def test_add_planned_chunks_accumulates(self): + """Planned chunk counts accumulate across quantum leaves.""" + state = CallState.new() + + state.add_planned_chunks(4) + state.add_planned_chunks(2) + + assert state.chunks_total == 6 + + def test_start_and_finish_track_in_flight_chunks(self): + """Start/finish transitions drive active and done counters.""" + state = CallState.new() + state.add_planned_chunks(2) + + state.mark_chunk_started() + state.mark_chunk_started() + assert state.active_chunks == 2 + assert state.chunks_done == 0 + + state.mark_chunk_finished() + assert state.active_chunks == 1 + assert state.chunks_done == 1 + + state.mark_chunk_finished() + assert state.active_chunks == 0 + assert state.chunks_done == 2 + + def test_finish_never_drives_active_chunks_negative(self): + """An unmatched finish clamps active_chunks at zero.""" + state = CallState.new() + + state.mark_chunk_finished() + + assert state.active_chunks == 0 + assert state.chunks_done == 1 + + +class TestStatusSnapshot: + def test_snapshot_defaults_to_idle(self): + """Before any backend status is recorded, the state is IDLE.""" + snapshot = CallState.new().status_snapshot() + + assert snapshot == { + "state": "IDLE", + "progress": 0.0, + "message": None, + "chunks_total": 0, + "chunks_done": 0, + "active_chunks": 0, + } + + def test_snapshot_reports_complete_when_future_done_without_status(self): + """A resolved future with no recorded backend status is COMPLETE.""" + snapshot = CallState.new().status_snapshot(future_done=True) + + assert snapshot["state"] == "COMPLETE" + + def test_snapshot_passes_through_recorded_status(self): + """Recorded backend status fields flow into the snapshot unchanged.""" + state = CallState.new() + state.set_current_status(state="RUNNING", progress=0.5, message="halfway") + + snapshot = state.status_snapshot() + + assert snapshot["state"] == "RUNNING" + assert snapshot["progress"] == 0.5 + assert snapshot["message"] == "halfway" + + def test_snapshot_prefers_recorded_status_over_future_done(self): + """A recorded backend status wins over the COMPLETE fallback.""" + state = CallState.new() + state.set_current_status(state="RUNNING", progress=1.0, message=None) + + snapshot = state.status_snapshot(future_done=True) + + assert snapshot["state"] == "RUNNING" + + def test_snapshot_includes_chunk_counters(self): + """Snapshots surface the chunk counters after transitions.""" + state = CallState.new() + state.add_planned_chunks(3) + state.mark_chunk_started() + state.mark_chunk_started() + state.mark_chunk_finished() + + snapshot = state.status_snapshot() + + assert snapshot["chunks_total"] == 3 + assert snapshot["chunks_done"] == 1 + assert snapshot["active_chunks"] == 1 + + def test_set_current_status_stores_immutable_job_status(self): + """Recorded status is a frozen JobStatus value object.""" + state = CallState.new() + state.set_current_status(state="RUNNING", progress=0.25, message=None) + + status = state.current_status + + assert isinstance(status, JobStatus) + assert status.state == "RUNNING" + assert status.progress == 0.25 + assert status.message is None diff --git a/tests/core/test_execution_units.py b/tests/core/test_execution_units.py new file mode 100644 index 00000000..4013d9fd --- /dev/null +++ b/tests/core/test_execution_units.py @@ -0,0 +1,638 @@ +"""No-cloud unit tests for the extracted execution units (PML-305). + +BatchChunker and RemoteJobRunner are exercised directly with fake jobs, +samplers, and processors — no MerlinProcessor instance and no cloud access. +""" + +from __future__ import annotations + +import threading +import time +from concurrent.futures import CancelledError +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +import merlin.core.execution as execution_module +import merlin.core.perceval_adapter as perceval_adapter_module +from merlin.core.execution import BatchChunker, RemoteJobRunner +from merlin.core.merlin_processor import CallState +from merlin.core.perceval_adapter import RemoteJobFailedError + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class FakeCommand: + """Sampler command that records how it was submitted.""" + + def __init__(self, job=None) -> None: + self.executed = False + self.name = None + self.execute_kwargs = None + self._job = job if job is not None else self + + def execute_async(self, **kwargs): + self.executed = True + self.execute_kwargs = kwargs + return self._job + + +class FakeSampler: + """Sampler fake exposing the commands and iteration API used by run_chunk.""" + + def __init__(self) -> None: + self.probs = FakeCommand() + self.sample_count = FakeCommand() + self.samples = FakeCommand() + self.cleared = False + self.iterations: list[dict] = [] + + def clear_iterations(self) -> None: + self.cleared = True + + def add_iteration(self, circuit_params) -> None: + self.iterations.append(circuit_params) + + +class FakeJob: + """Remote job fake with scripted polling behavior.""" + + def __init__( + self, + *, + job_id: str | None = "job-1", + is_complete: bool = True, + is_failed: bool = False, + status=None, + result_events=None, + ) -> None: + self.id = job_id + self.is_complete = is_complete + self.is_failed = is_failed + self.status = status + self.cancelled = False + self.get_results_calls = 0 + self._result_events = ( + [{"results_list": []}] if result_events is None else list(result_events) + ) + + def cancel(self) -> None: + self.cancelled = True + + def get_results(self): + self.get_results_calls += 1 + event = ( + self._result_events.pop(0) + if len(self._result_events) > 1 + else self._result_events[0] + ) + if isinstance(event, BaseException): + raise event + return event + + +def make_runner(**overrides) -> RemoteJobRunner: + """Build a RemoteJobRunner with recording no-op dependencies.""" + deps = { + "create_processor": MagicMock(name="create_processor"), + "get_available_commands": lambda: ("probs", "sample_count", "samples"), + "effective_sample_count": lambda nsample: ( + 10_000 if nsample is None else int(nsample) + ), + "get_max_shots_per_call": lambda: 100_000, + "default_shots_per_call": 10_000, + "map_results": MagicMock(return_value=torch.tensor([[1.0]])), + "register_job": MagicMock(name="register_job"), + "unregister_job": MagicMock(name="unregister_job"), + "get_microbatch_limit": lambda: 32, + } + deps.update(overrides) + return RemoteJobRunner(**deps) + + +def make_chunk_config() -> SimpleNamespace: + """Return the validated-config shape consumed by run_chunk.""" + return SimpleNamespace( + circuit=MagicMock(name="circuit"), + input_state=[1, 0], + input_param_order=["theta_0", "theta_1"], + ) + + +# --------------------------------------------------------------------------- +# BatchChunker +# --------------------------------------------------------------------------- + + +class TestSplitBatch: + def test_exact_multiple_produces_equal_chunks(self): + """An exact multiple splits into equal-size chunks.""" + assert BatchChunker.split_batch(8, 4) == [(0, 4), (4, 8)] + + def test_remainder_produces_short_final_chunk(self): + """A remainder yields a short final chunk.""" + assert BatchChunker.split_batch(10, 4) == [(0, 4), (4, 8), (8, 10)] + + def test_batch_smaller_than_microbatch_is_one_chunk(self): + """Small batches stay as a single chunk.""" + assert BatchChunker.split_batch(3, 32) == [(0, 3)] + + def test_empty_batch_produces_no_chunks(self): + """An empty batch produces no chunks.""" + assert BatchChunker.split_batch(0, 32) == [] + + +class TestBatchChunkerRunChunks: + def test_outputs_are_stitched_in_chunk_order(self): + """Chunk outputs are concatenated in submission order.""" + + def run_chunk(layer, config, chunk, nsample, state, deadline, job_base_label): + # Later chunks finish first to prove ordering is positional. + time.sleep(0.03 if chunk[0, 0] == 0 else 0.0) + return chunk.clone() + + chunker = BatchChunker( + run_chunk=run_chunk, chunk_concurrency=4, cancel_all=MagicMock() + ) + input_tensor = torch.arange(8, dtype=torch.float32).reshape(4, 2) + + output = chunker.run_chunks( + object(), + make_chunk_config(), + input_tensor, + BatchChunker.split_batch(4, 2), + None, + CallState.new(), + None, + ) + + torch.testing.assert_close(output, input_tensor) + + def test_concurrency_is_bounded(self): + """No more than chunk_concurrency chunk jobs run at once.""" + lock = threading.Lock() + active = 0 + max_active = 0 + + def run_chunk(layer, config, chunk, nsample, state, deadline, job_base_label): + nonlocal active, max_active + with lock: + active += 1 + max_active = max(max_active, active) + time.sleep(0.02) + with lock: + active -= 1 + return torch.ones(chunk.shape[0], 1) + + chunker = BatchChunker( + run_chunk=run_chunk, chunk_concurrency=2, cancel_all=MagicMock() + ) + + chunker.run_chunks( + object(), + make_chunk_config(), + torch.zeros(8, 2), + BatchChunker.split_batch(8, 1), + None, + CallState.new(), + None, + ) + + assert max_active <= 2 + + def test_chunk_labels_carry_call_id_and_position(self): + """Job base labels identify the layer, call, and chunk position.""" + labels: list[str] = [] + + def run_chunk(layer, config, chunk, nsample, state, deadline, job_base_label): + labels.append(job_base_label) + return torch.ones(chunk.shape[0], 1) + + chunker = BatchChunker( + run_chunk=run_chunk, chunk_concurrency=1, cancel_all=MagicMock() + ) + state = CallState.new() + layer = SimpleNamespace(name="qlayer") + + chunker.run_chunks( + layer, + make_chunk_config(), + torch.zeros(4, 2), + BatchChunker.split_batch(4, 2), + None, + state, + None, + ) + + assert labels == [ + f"mer:qlayer:{state.call_id}:1/2", + f"mer:qlayer:{state.call_id}:2/2", + ] + + def test_state_counters_reflect_chunk_lifecycle(self): + """Planned/done counters are updated and drain to zero active.""" + state = CallState.new() + + chunker = BatchChunker( + run_chunk=lambda *a, **k: torch.ones(1, 1), + chunk_concurrency=2, + cancel_all=MagicMock(), + ) + + chunker.run_chunks( + object(), + make_chunk_config(), + torch.zeros(3, 2), + BatchChunker.split_batch(3, 1), + None, + state, + None, + ) + + assert state.chunks_total == 3 + assert state.chunks_done == 3 + assert state.active_chunks == 0 + + def test_first_chunk_error_is_raised(self): + """A failing chunk propagates its error after the pool drains.""" + + def run_chunk(layer, config, chunk, nsample, state, deadline, job_base_label): + if chunk[0, 0] == 2: + raise RuntimeError("chunk exploded") + return torch.ones(chunk.shape[0], 1) + + chunker = BatchChunker( + run_chunk=run_chunk, chunk_concurrency=1, cancel_all=MagicMock() + ) + + with pytest.raises(RuntimeError, match="chunk exploded"): + chunker.run_chunks( + object(), + make_chunk_config(), + torch.tensor([[0.0], [2.0]]), + BatchChunker.split_batch(2, 1), + None, + CallState.new(), + None, + ) + + def test_deadline_cancels_all_and_raises_timeout(self): + """An elapsed deadline cancels in-flight jobs and raises TimeoutError.""" + cancel_all = MagicMock() + release = threading.Event() + + def run_chunk(layer, config, chunk, nsample, state, deadline, job_base_label): + release.wait(timeout=5.0) + return torch.ones(chunk.shape[0], 1) + + chunker = BatchChunker( + run_chunk=run_chunk, chunk_concurrency=1, cancel_all=cancel_all + ) + + try: + with pytest.raises(TimeoutError, match="remote cancel issued"): + chunker.run_chunks( + object(), + make_chunk_config(), + torch.zeros(1, 2), + BatchChunker.split_batch(1, 1), + None, + CallState.new(), + time.time() + 0.05, + ) + finally: + release.set() + + cancel_all.assert_called_once_with() + + def test_concurrency_floor_is_one(self): + """A non-positive concurrency setting still runs chunks serially.""" + chunker = BatchChunker( + run_chunk=lambda *a, **k: torch.ones(1, 1), + chunk_concurrency=0, + cancel_all=MagicMock(), + ) + + output = chunker.run_chunks( + object(), + make_chunk_config(), + torch.zeros(2, 2), + BatchChunker.split_batch(2, 1), + None, + CallState.new(), + None, + ) + + assert output.shape == (2, 1) + + +# --------------------------------------------------------------------------- +# RemoteJobRunner.submit_job +# --------------------------------------------------------------------------- + + +class TestSubmitJob: + def test_probs_selected_when_available_and_nsample_none(self): + """probs is selected when available and no shots requested.""" + runner = make_runner() + sampler = FakeSampler() + + job, is_probability = runner.submit_job( + sampler, None, "label", lambda base, cmd: f"{base}:{cmd}" + ) + + assert is_probability is True + assert job is sampler.probs + assert sampler.probs.executed is True + assert sampler.probs.name == "label:probs" + + def test_sampling_selected_when_nsample_positive(self): + """A positive nsample forces the sampling command path.""" + runner = make_runner() + sampler = FakeSampler() + + job, is_probability = runner.submit_job( + sampler, 500, "label", lambda base, cmd: f"{base}:{cmd}" + ) + + assert is_probability is False + assert job is sampler.sample_count + assert sampler.sample_count.execute_kwargs == {"max_samples": 500} + + def test_samples_used_when_sample_count_unavailable(self): + """samples is the fallback when sample_count is unavailable.""" + runner = make_runner(get_available_commands=lambda: ("samples",)) + sampler = FakeSampler() + + job, is_probability = runner.submit_job( + sampler, 100, None, lambda base, cmd: f"{base}:{cmd}" + ) + + assert job is sampler.samples + assert is_probability is False + + def test_sample_count_is_default_fallback_without_commands(self): + """sample_count is attempted when no commands are advertised.""" + runner = make_runner(get_available_commands=lambda: ()) + sampler = FakeSampler() + + job, _ = runner.submit_job( + sampler, None, None, lambda base, cmd: f"{base}:{cmd}" + ) + + assert job is sampler.sample_count + + def test_shot_count_flows_through_effective_sample_count(self): + """Submitted shots come from the injected effective_sample_count.""" + runner = make_runner(effective_sample_count=lambda nsample: 42) + sampler = FakeSampler() + + runner.submit_job(sampler, 999, None, lambda base, cmd: f"{base}:{cmd}") + + assert sampler.sample_count.execute_kwargs == {"max_samples": 42} + + def test_serializable_iterator_normalized_before_submit(self): + """Perceval 1.2 iterator payloads are flattened to plain lists.""" + runner = make_runner() + sampler = FakeSampler() + iterations = [{"theta_0": 0.1}, {"theta_0": 0.2}] + iterator = SimpleNamespace(iterations=iterations) + sampler._iterator = iterator + sampler.probs._request_data = {"payload": {"iterator": iterator}} + + runner.submit_job(sampler, None, None, lambda base, cmd: f"{base}:{cmd}") + + assert sampler.probs._request_data["payload"]["iterator"] == iterations + + +class TestCappedName: + def test_short_names_are_sanitized_verbatim(self): + """Short job names are sanitized but otherwise kept verbatim.""" + runner = make_runner() + + assert runner._capped_name("mer:layer 1", "probs") == "mer:layer_1:probs" + + def test_long_names_are_capped_with_hash_suffix(self): + """Long job names are capped with a stable hash suffix.""" + runner = make_runner(job_name_max=20) + + name = runner._capped_name("mer:" + "x" * 60, "sample_count") + + assert len(name) == 20 + assert "~" in name + + +# --------------------------------------------------------------------------- +# RemoteJobRunner.run_chunk +# --------------------------------------------------------------------------- + + +def run_chunk_with(runner, *, nsample=None, state=None, deadline=None, rows=1): + """Invoke run_chunk with a standard fake config and input.""" + return runner.run_chunk( + object(), + make_chunk_config(), + torch.zeros(rows, 2), + nsample, + state if state is not None else CallState.new(), + deadline, + job_base_label="label", + ) + + +class TestRunChunk: + def test_success_registers_job_and_maps_results(self): + """A successful chunk registers, polls, unregisters, and maps results.""" + raw_results = {"results_list": [{"results": {"|1,0>": 1.0}}]} + job = FakeJob(result_events=[raw_results]) + sampler = FakeSampler() + sampler.probs = FakeCommand(job=job) + mapped = torch.tensor([[0.25, 0.75]]) + layer = object() + runner = make_runner(map_results=MagicMock(return_value=mapped)) + rp = runner._create_processor.return_value + + with patch.object(perceval_adapter_module, "Sampler", return_value=sampler): + output = runner.run_chunk( + layer, + make_chunk_config(), + torch.zeros(1, 2), + None, + CallState.new(), + None, + job_base_label="label", + ) + + assert output is mapped + runner._register_job.assert_called_once_with(job) + runner._unregister_job.assert_called_once_with(job) + runner._map_results.assert_called_once_with(raw_results, 1, layer, None, True) + rp.set_circuit.assert_called_once() + rp.with_input.assert_called_once() + rp.min_detected_photons_filter.assert_called_once_with(1) + + def test_fresh_processor_and_retry_on_failure(self, monkeypatch): + """Each retry builds a fresh processor; success on a later attempt wins.""" + monkeypatch.setattr(execution_module.time, "sleep", lambda _s: None) + job = FakeJob() + good_sampler = FakeSampler() + good_sampler.probs = FakeCommand(job=job) + + attempts = [] + + def create_processor(): + attempts.append(MagicMock(name=f"rp-{len(attempts)}")) + return attempts[-1] + + submit_results = [RuntimeError("submit failed"), None] + + class FlakySampler(FakeSampler): + def __init__(self, processor, max_shots_per_call): + super().__init__() + failure = submit_results.pop(0) + if failure is None: + self.probs = FakeCommand(job=job) + else: + self.probs = MagicMock() + self.probs.execute_async.side_effect = failure + + def clear_iterations(self): + pass + + def add_iteration(self, circuit_params): + pass + + runner = make_runner(create_processor=create_processor) + + with patch.object(perceval_adapter_module, "Sampler", FlakySampler): + output = run_chunk_with(runner) + + assert output.shape == (1, 1) + assert len(attempts) == 2 # fresh processor per attempt + + def test_gives_up_after_max_retries_and_chains_last_error(self, monkeypatch): + """Persistent submission failures exhaust retries and chain the cause.""" + monkeypatch.setattr(execution_module.time, "sleep", lambda _s: None) + sampler = FakeSampler() + sampler.probs = MagicMock() + sampler.probs.execute_async.side_effect = RuntimeError("backend down") + runner = make_runner(max_retries=3) + + with ( + patch.object(perceval_adapter_module, "Sampler", return_value=sampler), + pytest.raises(RuntimeError, match="failed after 3 attempts") as excinfo, + ): + run_chunk_with(runner) + + assert isinstance(excinfo.value.__cause__, RuntimeError) + assert str(excinfo.value.__cause__) == "backend down" + assert runner._create_processor.call_count == 3 # fresh RP per attempt + + def test_cancellation_short_circuits_before_submission(self): + """A prior cancel prevents any processor construction.""" + runner = make_runner() + state = CallState.new() + state.request_cancel() + + with pytest.raises(CancelledError, match="Remote call was cancelled"): + run_chunk_with(runner, state=state) + + runner._create_processor.assert_not_called() + + def test_deadline_short_circuits_before_submission(self): + """An elapsed deadline prevents any processor construction.""" + runner = make_runner() + + with pytest.raises(TimeoutError, match="remote cancel issued"): + run_chunk_with(runner, deadline=time.time() - 1.0) + + runner._create_processor.assert_not_called() + + def test_oversized_chunk_raises_value_error(self): + """Chunks beyond the microbatch limit fail loudly.""" + runner = make_runner(get_microbatch_limit=lambda: 2) + + with pytest.raises(ValueError, match="exceeds microbatch"): + run_chunk_with(runner, rows=3) + + def test_unbounded_microbatch_limit_accepts_large_chunks(self): + """Session-style backends (limit None) skip the chunk size guard.""" + raw_results = {"results_list": []} + job = FakeJob(result_events=[raw_results]) + sampler = FakeSampler() + sampler.probs = FakeCommand(job=job) + runner = make_runner(get_microbatch_limit=lambda: None) + + with patch.object(perceval_adapter_module, "Sampler", return_value=sampler): + run_chunk_with(runner, rows=100) + + runner._map_results.assert_called_once() + + +# --------------------------------------------------------------------------- +# RemoteJobRunner.poll_job +# --------------------------------------------------------------------------- + + +class TestPollJob: + def test_success_records_job_id_and_unregisters(self): + """Successful polling records the job id and unregisters the job.""" + raw_results = {"results_list": [{"results": {"|1,0>": 1.0}}]} + job = FakeJob(job_id="job-42", result_events=[raw_results]) + state = CallState.new() + mapped = torch.tensor([[1.0]]) + runner = make_runner(map_results=MagicMock(return_value=mapped)) + + result = runner.poll_job(job, state, None, 1, object(), None) + + assert result is mapped + assert state.job_ids == ["job-42"] + runner._unregister_job.assert_called_once_with(job) + + def test_cancel_request_cancels_job_and_raises(self): + """Caller cancellation asks the backend job to cancel before raising.""" + job = FakeJob(is_complete=False) + state = CallState.new() + state.request_cancel() + runner = make_runner() + + with pytest.raises(CancelledError, match="Remote call was cancelled"): + runner.poll_job(job, state, None, 1, object(), None) + + assert job.cancelled is True + + def test_failed_job_raises_with_message_and_unregisters(self): + """Failed jobs raise the Merlin error with the backend message.""" + job = FakeJob( + is_complete=False, + is_failed=True, + status=SimpleNamespace( + state="FAILED", progress=None, stop_message="hardware rejected job" + ), + ) + runner = make_runner() + + with pytest.raises(RemoteJobFailedError, match="hardware rejected job"): + runner.poll_job(job, CallState.new(), None, 1, object(), None) + + runner._unregister_job.assert_called_once_with(job) + + def test_status_snapshot_updated_from_job_status(self): + """Polling records the backend status into the call state.""" + raw_results = {"results_list": []} + job = FakeJob( + result_events=[raw_results], + status=SimpleNamespace(state="RUNNING", progress=0.5, stop_message=None), + ) + state = CallState.new() + runner = make_runner() + + runner.poll_job(job, state, None, 1, object(), None) + + assert state.current_status is not None + assert state.current_status.state == "RUNNING" + assert state.current_status.progress == 0.5 diff --git a/tests/core/test_merlin_future.py b/tests/core/test_merlin_future.py new file mode 100644 index 00000000..6c3134f0 --- /dev/null +++ b/tests/core/test_merlin_future.py @@ -0,0 +1,196 @@ +"""No-cloud unit tests for the typed MerlinFuture async handle (PML-304). + +Covers the explicit async contract that replaced the runtime monkey-patched +future attributes: return type, status payload, job id exposure, cooperative +cancellation, and the synchronous forward() built on top of the handle. +""" + +from __future__ import annotations + +from concurrent.futures import CancelledError +from unittest.mock import MagicMock + +import perceval as pcvl +import pytest +import torch +import torch.nn as nn + +from merlin.algorithms.module import MerlinModule +from merlin.core.merlin_processor import CallState, MerlinFuture, MerlinProcessor + + +def make_future(**overrides) -> tuple[MerlinFuture, CallState, MagicMock]: + """Build a MerlinFuture around a fresh CallState and recorded cancel_all.""" + state = overrides.get("state", CallState.new()) + cancel_all = overrides.get("cancel_all", MagicMock(name="cancel_all")) + return MerlinFuture(state, cancel_all), state, cancel_all + + +def make_local_processor() -> MerlinProcessor: + """Build a MerlinProcessor around a local mocked AProcessor.""" + from perceval.runtime import AProcessor + + processor = MagicMock(spec=AProcessor) + processor.is_remote = False + processor.name = "local:slos" + processor.available_commands = ["probs"] + return MerlinProcessor(processor=processor) + + +class PassthroughLeaf(MerlinModule): + """Quantum leaf that never offloads, so pipelines run locally.""" + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(2, 2) + + def should_offload(self) -> bool: + return False + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * 2.0 + + def export_config(self): + return { + "circuit": pcvl.Circuit(m=2), + "input_state": [1, 0], + "input_param_order": ["theta_0", "theta_1"], + } + + +class TestReturnType: + def test_forward_async_returns_merlin_future(self): + """forward_async returns the typed handle, still a torch Future.""" + proc = make_local_processor() + module = PassthroughLeaf() + module.eval() + + fut = proc.forward_async(module, torch.ones(1, 2)) + + assert isinstance(fut, MerlinFuture) + assert isinstance(fut, torch.futures.Future) + fut.wait() + + def test_contract_methods_are_declared_not_patched(self): + """The async contract lives on the class, not on instances.""" + assert callable(MerlinFuture.status) + assert callable(MerlinFuture.cancel_remote) + assert isinstance(MerlinFuture.job_ids, property) + + fut, _, _ = make_future() + assert "status" not in fut.__dict__ + assert "cancel_remote" not in fut.__dict__ + assert "job_ids" not in fut.__dict__ + + +class TestStatus: + def test_status_payload_shape_when_idle(self): + """A fresh handle reports the exact IDLE status payload.""" + fut, _, _ = make_future() + + assert fut.status() == { + "state": "IDLE", + "progress": 0.0, + "message": None, + "chunks_total": 0, + "chunks_done": 0, + "active_chunks": 0, + } + + def test_status_reflects_call_state_progress(self): + """status() mirrors backend status and chunk counters live.""" + fut, state, _ = make_future() + state.add_planned_chunks(2) + state.mark_chunk_started() + state.set_current_status(state="RUNNING", progress=0.5, message="halfway") + + status = fut.status() + + assert status["state"] == "RUNNING" + assert status["progress"] == 0.5 + assert status["message"] == "halfway" + assert status["chunks_total"] == 2 + assert status["active_chunks"] == 1 + + def test_status_reports_complete_when_done_without_backend_status(self): + """A resolved handle with no backend status reports COMPLETE.""" + fut, _, _ = make_future() + fut.set_result(torch.ones(1)) + + assert fut.status()["state"] == "COMPLETE" + + +class TestJobIds: + def test_job_ids_is_live_view_of_call_state(self): + """job_ids is the same live list CallState appends to.""" + fut, state, _ = make_future() + + assert fut.job_ids == [] + state.record_job_id("job-1") + assert fut.job_ids == ["job-1"] + assert fut.job_ids is state.job_ids + + +class TestCancelRemote: + def test_cancel_remote_requests_cancel_and_cancels_jobs(self): + """cancel_remote flags the call state and cancels active jobs.""" + fut, state, cancel_all = make_future() + + fut.cancel_remote() + + assert state.cancel_requested is True + cancel_all.assert_called_once_with() + + def test_cancel_remote_resolves_future_with_cancelled_error(self): + """After cancel_remote, waiting raises CancelledError.""" + fut, _, _ = make_future() + + fut.cancel_remote() + + with pytest.raises(CancelledError, match="Remote call was cancelled"): + fut.wait() + + def test_cancel_remote_after_completion_keeps_result(self): + """Cancelling a finished call still cancels jobs but keeps the result.""" + fut, state, cancel_all = make_future() + result = torch.ones(1) + fut.set_result(result) + + fut.cancel_remote() + + assert state.cancel_requested is True + cancel_all.assert_called_once_with() + assert torch.equal(fut.wait(), result) + + def test_forward_async_cancel_propagates_cancelled_error(self): + """Cancelling a running pipeline surfaces CancelledError on wait.""" + proc = make_local_processor() + + class BlockingLeaf(PassthroughLeaf): + def forward(self, x: torch.Tensor) -> torch.Tensor: + import time + + time.sleep(0.5) + return x + + module = BlockingLeaf() + module.eval() + + fut = proc.forward_async(module, torch.ones(1, 2)) + fut.cancel_remote() + + with pytest.raises(CancelledError, match="Remote call was cancelled"): + fut.wait() + + +class TestSynchronousForward: + def test_forward_waits_on_the_typed_handle(self): + """Synchronous forward() resolves through the typed handle.""" + proc = make_local_processor() + module = PassthroughLeaf() + module.eval() + x = torch.ones(2, 2) + + output = proc.forward(module, x) + + torch.testing.assert_close(output, x * 2.0) diff --git a/tests/core/test_merlin_processor_unit.py b/tests/core/test_merlin_processor_unit.py index b640fb11..3a2e47a6 100644 --- a/tests/core/test_merlin_processor_unit.py +++ b/tests/core/test_merlin_processor_unit.py @@ -21,6 +21,7 @@ from perceval.runtime.session import ISession import merlin.core.merlin_processor as merlin_processor_module +import merlin.core.perceval_adapter as perceval_adapter_module from merlin.algorithms import QuantumLayer from merlin.algorithms.module import MerlinModule from merlin.builder.circuit_builder import CircuitBuilder @@ -28,6 +29,7 @@ from merlin.core.computation_space import ComputationSpace from merlin.core.merlin_processor import ( BackendCapabilities, + CallState, MerlinProcessor, SupportsExportConfig, ValidatedLayerConfig, @@ -237,9 +239,9 @@ def process_results(raw_results, batch_size, layer, nsample, is_probability=Fals return proc -def make_state() -> dict: - """Return the mutable polling state shape expected by _poll_job.""" - return {"cancel_requested": False, "job_ids": []} +def make_state() -> CallState: + """Return the typed per-call state expected by _poll_job.""" + return CallState.new() def make_local_chunk_config() -> SimpleNamespace: @@ -1018,7 +1020,7 @@ def fake_run_chunk_local( assert layer_arg is layer assert isinstance(config, ValidatedLayerConfig) assert nsample is None - assert state["cancel_requested"] is False + assert state.cancel_requested is False assert deadline is not None observed_chunks.append(input_chunk.clone()) return torch.ones(input_chunk.shape[0], 2) @@ -1114,7 +1116,7 @@ def test_run_chunk_local_uses_fresh_processor_per_execution(): with ( patch.object( - merlin_processor_module, "Sampler", return_value=sampler + perceval_adapter_module, "Sampler", return_value=sampler ) as sampler_cls, ): output = proc._run_chunk_local( @@ -1151,7 +1153,7 @@ def test_run_chunk_local_uses_sample_count_when_probs_unavailable(): raw_results = {"results_list": [{"results": {"|1,0>": 3}}]} sampler = FakeSyncSampler(raw_results) - with patch.object(merlin_processor_module, "Sampler", return_value=sampler): + with patch.object(perceval_adapter_module, "Sampler", return_value=sampler): proc._run_chunk_local( layer, config, @@ -1181,7 +1183,7 @@ def test_run_chunk_local_caps_default_sample_count_to_max_shots_per_call(): sampler = FakeSyncSampler(raw_results) with patch.object( - merlin_processor_module, "Sampler", return_value=sampler + perceval_adapter_module, "Sampler", return_value=sampler ) as sampler_cls: proc._run_chunk_local( layer, @@ -1209,7 +1211,7 @@ def test_run_chunk_local_uses_samples_when_sample_count_unavailable(): raw_results = {"results_list": [{"results": {"|1,0>": 3}}]} sampler = FakeSyncSampler(raw_results) - with patch.object(merlin_processor_module, "Sampler", return_value=sampler): + with patch.object(perceval_adapter_module, "Sampler", return_value=sampler): proc._run_chunk_local( layer, config, @@ -1235,7 +1237,7 @@ def test_run_chunk_local_defaults_to_sample_count_when_commands_are_empty(): raw_results = {"results_list": [{"results": {"|1,0>": 3}}]} sampler = FakeSyncSampler(raw_results) - with patch.object(merlin_processor_module, "Sampler", return_value=sampler): + with patch.object(perceval_adapter_module, "Sampler", return_value=sampler): proc._run_chunk_local( layer, config, @@ -1254,7 +1256,7 @@ def test_run_chunk_local_raises_cancelled_before_execution(): """Local chunk execution observes cancellation before starting work.""" proc = make_local_chunk_processor(["probs"]) state = make_state() - state["cancel_requested"] = True + state.request_cancel() with pytest.raises(CancelledError, match="Local call was cancelled"): proc._run_chunk_local( @@ -1293,12 +1295,10 @@ def test_run_chunk_local_raises_cancelled_after_execution(): proc = make_local_chunk_processor(["probs"]) state = make_state() raw_results = {"results_list": [{"results": {"|1,0>": 1.0}}]} - sampler = FakeSyncSampler( - raw_results, on_execute=lambda: state.__setitem__("cancel_requested", True) - ) + sampler = FakeSyncSampler(raw_results, on_execute=state.request_cancel) with ( - patch.object(merlin_processor_module, "Sampler", return_value=sampler), + patch.object(perceval_adapter_module, "Sampler", return_value=sampler), pytest.raises(CancelledError, match="Local call was cancelled"), ): proc._run_chunk_local( @@ -1322,7 +1322,7 @@ def test_run_chunk_local_raises_timeout_after_execution(monkeypatch): monkeypatch.setattr(merlin_processor_module.time, "time", lambda: next(time_values)) with ( - patch.object(merlin_processor_module, "Sampler", return_value=sampler), + patch.object(perceval_adapter_module, "Sampler", return_value=sampler), pytest.raises(TimeoutError, match="Local call timed out"), ): proc._run_chunk_local( @@ -1383,7 +1383,7 @@ def __init__(self, processor, max_shots_per_call) -> None: captured_processor["max_shots_per_call"] = max_shots_per_call super().__init__(raw_results) - with patch.object(merlin_processor_module, "Sampler", CapturingSampler): + with patch.object(perceval_adapter_module, "Sampler", CapturingSampler): output = proc._run_chunk_local( layer, config, @@ -1582,7 +1582,7 @@ def run_perceval_probabilities( processor.set_circuit(circuit.copy()) processor.with_input(pcvl.BasicState([1, 0, 0])) - sampler = merlin_processor_module.Sampler( + sampler = perceval_adapter_module.Sampler( processor, max_shots_per_call=MerlinProcessor.DEFAULT_MAX_SHOTS, ) @@ -1831,7 +1831,7 @@ def test_poll_job_success_processes_dict_payload_and_records_job_id(): ) assert torch.equal(result, output) - assert state["job_ids"] == ["job-success"] + assert state.job_ids == ["job-success"] assert proc.processed_calls == [(raw_results, 3, layer, None, False)] assert job not in proc._active_jobs @@ -1858,7 +1858,7 @@ def test_poll_job_cancel_request_cancels_remote_job(): proc = make_poll_processor() job = FakeJob(is_complete=False) state = make_state() - state["cancel_requested"] = True + state.request_cancel() with pytest.raises(CancelledError, match=r"Remote call was cancelled"): proc._poll_job(job, state, None, 1, object(), None) diff --git a/tests/core/test_perceval_adapter.py b/tests/core/test_perceval_adapter.py new file mode 100644 index 00000000..dbebbfaf --- /dev/null +++ b/tests/core/test_perceval_adapter.py @@ -0,0 +1,470 @@ +"""No-cloud unit tests for the Perceval adapter layer (PML-306). + +Every test uses plain fakes or monkeypatching — no cloud access and no real +RemoteProcessor construction. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import perceval as pcvl +import pytest + +import merlin.core.perceval_adapter as adapter_module +from merlin.core.perceval_adapter import ( + LOCAL_EXPERIMENT_SNAPSHOT_ATTR, + JobStatusSnapshot, + LocalExperimentSnapshot, + PercevalAdapter, + RemoteJobFailedError, + TokenExtractionError, +) + + +def make_rp(handler=None) -> MagicMock: + """Build a RemoteProcessor fake exposing an RPC handler.""" + rp = MagicMock(name="remote_processor") + if handler is None: + rp.get_rpc_handler.side_effect = RuntimeError("no handler") + else: + rp.get_rpc_handler.return_value = handler + return rp + + +def patch_remote_config(token: str | None): + """Patch perceval's RemoteConfig to return the given global token.""" + config = MagicMock() + config.get_token.return_value = token + return patch("perceval.runtime.RemoteConfig", return_value=config) + + +# --------------------------------------------------------------------------- +# Token extraction +# --------------------------------------------------------------------------- + + +class TestExtractToken: + @pytest.mark.parametrize("attr", ["token", "_token", "auth_token"]) + def test_handler_attribute_variants(self, attr): + """Each known handler token attribute is honored.""" + handler = SimpleNamespace(**{attr: "tok-123"}) + + assert PercevalAdapter.extract_token(make_rp(handler)) == "tok-123" + + def test_handler_attribute_priority_order(self): + """handler.token wins over the private/_alternate attributes.""" + handler = SimpleNamespace(**{"token": "primary", "_token": "secondary"}) + + assert PercevalAdapter.extract_token(make_rp(handler)) == "primary" + + def test_empty_handler_attributes_are_skipped(self): + """Empty or None token attributes fall through to later probes.""" + handler = SimpleNamespace(**{ + "token": "", + "_token": None, + "auth_token": "fallback", + }) + + assert PercevalAdapter.extract_token(make_rp(handler)) == "fallback" + + def test_bearer_header_fallback(self): + """The token is parsed from a Bearer Authorization header.""" + handler = SimpleNamespace(headers={"Authorization": "Bearer tok-from-header"}) + + assert PercevalAdapter.extract_token(make_rp(handler)) == "tok-from-header" + + def test_malformed_bearer_header_is_ignored(self): + """Non-Bearer Authorization headers are not treated as tokens.""" + handler = SimpleNamespace(headers={"Authorization": "Basic abc"}) + + with patch_remote_config(None): + assert PercevalAdapter.extract_token(make_rp(handler)) is None + + def test_remote_config_fallback(self): + """The global RemoteConfig token is used (stripped) as last resort.""" + handler = SimpleNamespace() + + with patch_remote_config(" global-tok "): + assert PercevalAdapter.extract_token(make_rp(handler)) == "global-tok" + + def test_broken_handler_falls_back_to_remote_config(self): + """A handler access error still allows the RemoteConfig fallback.""" + with patch_remote_config("global-tok"): + assert PercevalAdapter.extract_token(make_rp(None)) == "global-tok" + + def test_all_strategies_failing_returns_none(self): + """No handler token, header, or global config yields None.""" + with patch_remote_config(None): + assert PercevalAdapter.extract_token(make_rp(SimpleNamespace())) is None + + def test_remote_config_error_returns_none(self): + """A RemoteConfig failure is swallowed and resolution yields None.""" + config = MagicMock() + config.get_token.side_effect = RuntimeError("config broken") + + with patch("perceval.runtime.RemoteConfig", return_value=config): + assert PercevalAdapter.extract_token(make_rp(SimpleNamespace())) is None + + +# --------------------------------------------------------------------------- +# URL / clone / session / capabilities +# --------------------------------------------------------------------------- + + +class TestProcessorAccess: + def test_get_url_reads_handler_url(self): + """The RPC handler URL is exposed through get_url.""" + handler = SimpleNamespace(url="https://api.quandela.cloud") + + assert PercevalAdapter.get_url(make_rp(handler)) == "https://api.quandela.cloud" + + def test_get_url_returns_none_without_url_attribute(self): + """Handlers without a url attribute map to None.""" + assert PercevalAdapter.get_url(make_rp(SimpleNamespace())) is None + + def test_clone_forwards_name_token_url_and_proxies(self): + """Cloning forwards platform name, token, URL, and proxies.""" + rp = make_rp(SimpleNamespace(url="https://cloud")) + rp.name = "sim:slos" + rp.proxies = {"https": "proxy"} + clone = MagicMock(name="clone") + + with patch.object( + adapter_module, "RemoteProcessor", return_value=clone + ) as rp_cls: + result = PercevalAdapter.clone_remote_processor(rp, "tok") + + assert result is clone + rp_cls.assert_called_once() + assert rp_cls.call_args.kwargs == { + "name": "sim:slos", + "token": "tok", + "url": "https://cloud", + "proxies": {"https": "proxy"}, + } + + def test_build_from_session_delegates(self): + """Session-based construction delegates to build_remote_processor.""" + session = MagicMock() + + result = PercevalAdapter.build_from_session(session) + + assert result is session.build_remote_processor.return_value + + def test_get_backend_capabilities_snapshots_commands(self): + """Capabilities are read as (name, immutable command tuple).""" + processor = SimpleNamespace( + name="sim:slos", available_commands=["probs", "sample_count"] + ) + + name, commands = PercevalAdapter.get_backend_capabilities(processor) + + assert name == "sim:slos" + assert commands == ("probs", "sample_count") + + def test_configure_processor_sets_circuit_input_and_filter(self): + """Circuit, input state, and photon filter are installed together.""" + processor = MagicMock() + circuit = MagicMock(name="circuit") + + PercevalAdapter.configure_processor(processor, circuit, [1, 0, 1]) + + processor.set_circuit.assert_called_once_with(circuit) + processor.with_input.assert_called_once_with(pcvl.BasicState([1, 0, 1])) + processor.min_detected_photons_filter.assert_called_once_with(2) + + def test_configure_processor_skips_input_when_falsy(self): + """A falsy input state skips input and filter setup.""" + processor = MagicMock() + + PercevalAdapter.configure_processor(processor, MagicMock(), None) + + processor.set_circuit.assert_called_once() + processor.with_input.assert_not_called() + processor.min_detected_photons_filter.assert_not_called() + + def test_copy_circuit_returns_independent_copy(self): + """copy_circuit returns the circuit's own copy() result.""" + circuit = MagicMock() + + assert PercevalAdapter.copy_circuit(circuit) is circuit.copy.return_value + + def test_estimate_required_shots_delegates(self): + """Estimation forwards the target and parameter values verbatim.""" + rp = MagicMock() + rp.estimate_required_shots.return_value = 4321 + + result = PercevalAdapter.estimate_required_shots(rp, 100, {"theta": 0.5}) + + assert result == 4321 + rp.estimate_required_shots.assert_called_once_with( + 100, param_values={"theta": 0.5} + ) + + +# --------------------------------------------------------------------------- +# Samplers +# --------------------------------------------------------------------------- + + +class FakeCommand: + def __init__(self) -> None: + self.name = None + self.async_kwargs = None + self.sync_kwargs = None + + def execute_async(self, **kwargs): + self.async_kwargs = kwargs + return "async-job" + + def execute_sync(self, **kwargs): + self.sync_kwargs = kwargs + return {"results_list": []} + + +class FakeSampler: + def __init__(self) -> None: + self.probs = FakeCommand() + self.sample_count = FakeCommand() + self.samples = FakeCommand() + self.cleared = False + self.iterations: list[dict] = [] + + def clear_iterations(self) -> None: + self.cleared = True + + def add_iteration(self, circuit_params) -> None: + self.iterations.append(circuit_params) + + +class TestSamplerAccess: + def test_create_sampler_loads_iterations(self): + """Sampler creation clears then loads every iteration.""" + processor = MagicMock() + iterations = [{"theta": 0.1}, {"theta": 0.2}] + fake = FakeSampler() + + with patch.object(adapter_module, "Sampler", return_value=fake) as sampler_cls: + sampler = PercevalAdapter.create_sampler(processor, 5000, iterations) + + assert sampler is fake + sampler_cls.assert_called_once_with(processor, max_shots_per_call=5000) + assert fake.cleared is True + assert fake.iterations == iterations + + @pytest.mark.parametrize("command", ["probs", "sample_count", "samples"]) + def test_submit_async_dispatches_each_command(self, command): + """Each sampler command is dispatched by name with the job name set.""" + sampler = FakeSampler() + + job = PercevalAdapter.submit_async(sampler, command, name="job-name") + + assert job == "async-job" + assert getattr(sampler, command).name == "job-name" + assert getattr(sampler, command).async_kwargs == {} + + def test_submit_async_forwards_max_samples(self): + """Sampling submissions forward the max_samples argument.""" + sampler = FakeSampler() + + PercevalAdapter.submit_async(sampler, "sample_count", max_samples=777) + + assert sampler.sample_count.async_kwargs == {"max_samples": 777} + + def test_submit_async_without_name_leaves_command_name(self): + """Without a label the command's name attribute is untouched.""" + sampler = FakeSampler() + + PercevalAdapter.submit_async(sampler, "probs") + + assert sampler.probs.name is None + + @pytest.mark.parametrize("command", ["probs", "sample_count", "samples"]) + def test_execute_sync_dispatches_each_command(self, command): + """Each sampler command dispatches synchronously by name.""" + sampler = FakeSampler() + + result = PercevalAdapter.execute_sync(sampler, command) + + assert result == {"results_list": []} + assert getattr(sampler, command).sync_kwargs == {} + + def test_execute_sync_forwards_max_samples(self): + """Synchronous sampling forwards the max_samples argument.""" + sampler = FakeSampler() + + PercevalAdapter.execute_sync(sampler, "samples", max_samples=99) + + assert sampler.samples.sync_kwargs == {"max_samples": 99} + + def test_serializable_iterator_normalized(self): + """Perceval 1.2 iterator payloads are flattened to plain lists.""" + sampler = FakeSampler() + iterations = [{"theta": 0.1}] + iterator = SimpleNamespace(iterations=iterations) + sampler._iterator = iterator + sampler.probs._request_data = {"payload": {"iterator": iterator}} + + PercevalAdapter.submit_async(sampler, "probs") + + assert sampler.probs._request_data["payload"]["iterator"] == iterations + + def test_serializable_iterator_untouched_for_plain_lists(self): + """Payloads without a live iterator object are left alone.""" + sampler = FakeSampler() + sampler._iterator = None + sampler.probs._request_data = {"payload": {"iterator": ["kept"]}} + + PercevalAdapter.submit_async(sampler, "probs") + + assert sampler.probs._request_data["payload"]["iterator"] == ["kept"] + + +# --------------------------------------------------------------------------- +# Jobs +# --------------------------------------------------------------------------- + + +class TestJobAccess: + def test_job_snapshot_maps_all_fields(self): + """All job status fields map onto the normalized snapshot.""" + job = SimpleNamespace( + id="job-1", + status=SimpleNamespace(state="RUNNING", progress=0.4, stop_message="msg"), + is_complete=False, + is_failed=False, + ) + + snapshot = PercevalAdapter.job_snapshot(job) + + assert snapshot == JobStatusSnapshot( + job_id="job-1", + state="RUNNING", + progress=0.4, + stop_message="msg", + is_complete=False, + is_failed=False, + ) + + def test_job_snapshot_falls_back_to_job_id_attribute(self): + """job_id is read from the alternate attribute when id is absent.""" + job = SimpleNamespace(job_id="alt-id") + + assert PercevalAdapter.job_snapshot(job).job_id == "alt-id" + + def test_job_snapshot_defaults_for_bare_objects(self): + """Objects with no job attributes map to a safe default snapshot.""" + snapshot = PercevalAdapter.job_snapshot(object()) + + assert snapshot == JobStatusSnapshot( + job_id=None, + state=None, + progress=None, + stop_message=None, + is_complete=False, + is_failed=False, + ) + + def test_get_results_propagates_perceval_errors(self): + """Perceval result errors propagate unchanged to the caller.""" + job = MagicMock() + job.get_results.side_effect = RuntimeError("Results are not available") + + with pytest.raises(RuntimeError, match="Results are not available"): + PercevalAdapter.get_results(job) + + def test_cancel_job_swallows_cancel_errors(self): + """Cancellation errors are suppressed on the best-effort path.""" + job = MagicMock() + job.cancel.side_effect = RuntimeError("already finished") + + PercevalAdapter.cancel_job(job) # must not raise + + job.cancel.assert_called_once_with() + + def test_cancel_job_ignores_objects_without_cancel(self): + """Objects without a cancel method are ignored silently.""" + PercevalAdapter.cancel_job(object()) # must not raise + + +# --------------------------------------------------------------------------- +# Merlin-specific exceptions +# --------------------------------------------------------------------------- + + +class TestExceptions: + def test_token_extraction_error_is_catchable_as_value_error(self): + """Callers catching the historical ValueError still catch the new type.""" + with pytest.raises(ValueError, match="no token"): + raise TokenExtractionError("no token") + + def test_remote_job_failed_error_is_catchable_as_runtime_error(self): + """Callers catching the historical RuntimeError still catch the new type.""" + with pytest.raises(RuntimeError, match="Remote job failed"): + raise RemoteJobFailedError("Remote job failed: boom") + + +# --------------------------------------------------------------------------- +# Local processors +# --------------------------------------------------------------------------- + + +def make_local_processor() -> pcvl.Processor: + """Build a real local Perceval processor with metadata to preserve.""" + processor = pcvl.Processor("SLOS") + processor.set_circuit(pcvl.Circuit(2)) + processor.set_postselection(pcvl.PostSelect("[0] == 1")) + return processor + + +class TestLocalProcessorRebuild: + def test_rebuild_returns_isolated_processor_with_snapshot(self): + """Rebuilding yields a distinct processor carrying a metadata snapshot.""" + original = make_local_processor() + + fresh = PercevalAdapter.rebuild_local_processor(original) + + assert fresh is not original + assert fresh.experiment is not original.experiment + snapshot = getattr(fresh, LOCAL_EXPERIMENT_SNAPSHOT_ATTR) + assert isinstance(snapshot, LocalExperimentSnapshot) + + def test_rebuild_rejects_processor_without_experiment(self): + """Processors without copyable experiments are rejected clearly.""" + bare = SimpleNamespace(experiment=None, backend=MagicMock()) + + with pytest.raises(TypeError, match="copyable"): + PercevalAdapter.rebuild_local_processor(bare) + + def test_snapshot_and_restore_round_trip_preserves_postselection(self): + """Snapshot/restore keeps postselection across circuit replacement.""" + original = make_local_processor() + snapshot = PercevalAdapter.snapshot_experiment(original.experiment) + + fresh = PercevalAdapter.rebuild_local_processor(original) + fresh.set_circuit(pcvl.Circuit(2)) + PercevalAdapter.restore_experiment(fresh.experiment, snapshot) + + assert fresh.experiment.post_select_fn == original.experiment.post_select_fn + + def test_restore_rejects_mismatched_circuit_size(self): + """Restoring mode metadata onto a different-size circuit fails.""" + original = make_local_processor() + snapshot = PercevalAdapter.snapshot_experiment(original.experiment) + + fresh = PercevalAdapter.rebuild_local_processor(original) + fresh.set_circuit(pcvl.Circuit(3)) + + with pytest.raises(ValueError, match="circuit size"): + PercevalAdapter.restore_experiment(fresh.experiment, snapshot) + + def test_snapshot_is_independent_of_source_experiment(self): + """Snapshots do not alias the live experiment state.""" + original = make_local_processor() + + snapshot = PercevalAdapter.snapshot_experiment(original.experiment) + original.set_postselection(pcvl.PostSelect("[1] == 1")) + + assert snapshot.postselection != original.experiment.post_select_fn