✨ Add configurable RL pass timeouts

CHANGELOG.md

@@ -12,6 +12,8 @@ releases may include breaking changes.
 
 ### Added
 
+- ✨ Add opt-in per-pass timeouts for RL training and inference ([#789])
+  ([**@flowerthrower**])
 - ✨ Add Qiskit's `TrivialLayout`, `ElidePermutations`, `SabreSwap`,
   `BasicSwap`, `LookaheadSwap`, `RemoveIdentityEquivalent`, and
   `Optimize1qGatesSimpleCommutation` passes to the RL actions ([#785])
@@ -98,6 +100,7 @@ for previous changelogs._
 
 <!-- PR links -->
 
+[#789]: https://github.com/munich-quantum-toolkit/predictor/pull/789
 [#773]: https://github.com/munich-quantum-toolkit/predictor/pull/771
 [#769]: https://github.com/munich-quantum-toolkit/predictor/pull/769
 [#784]: https://github.com/munich-quantum-toolkit/predictor/pull/784

docs/setup.md

@@ -45,9 +45,14 @@ from mqt.bench.targets import get_device
 
 device = get_device("ibm_falcon_27")
 rl_pred = RL_Predictor(device=device, figure_of_merit="expected_fidelity")
-rl_pred.train_model(timesteps=100000)
+rl_pred.train_model(timesteps=100000, pass_timeout=600)
 ```
 
+`pass_timeout` is optional and limits each compilation pass during training. If
+it is omitted, pass execution is unbounded. Pass timeouts require POSIX signals
+and execution on the main thread; unsupported environments emit a warning and
+continue without a timeout.
+
 Currently, the following figures of merit are supported:
 
 ```{code-cell} ipython3
@@ -121,9 +126,16 @@ from mqt.predictor import qcompile
 from mqt.bench import get_benchmark, BenchmarkLevel
 
 uncompiled_qc = get_benchmark("ghz", level=BenchmarkLevel.ALG, circuit_size=5)
-compiled_qc, compilation_info, selected_device = qcompile(uncompiled_qc, figure_of_merit="expected_fidelity")
+compiled_qc, compilation_info, selected_device = qcompile(
+    uncompiled_qc,
+    figure_of_merit="expected_fidelity",
+    pass_timeout=60,
+)
 ```
 
+Inference has its own optional `pass_timeout`, so it can use a different limit
+than training.
+
 This returns:
 
 - the compiled quantum circuit,

src/mqt/predictor/qcompile.py

@@ -27,19 +27,29 @@ def qcompile(
     qc: QuantumCircuit,
     figure_of_merit: figure_of_merit = "expected_fidelity",
     tracer_output_path: str | Path | None = None,
+    pass_timeout: float | None = None,
 ) -> tuple[QuantumCircuit, list[str], str]:
     """Compiles a given quantum circuit to a device with the highest predicted figure of merit.
 
     Arguments:
         qc: The quantum circuit to be compiled.
         figure_of_merit: The figure of merit to be used for compilation. Defaults to "expected_fidelity".
         tracer_output_path: If provided, enables compiler tracing and exports the JSON log to this path/directory.
+        pass_timeout: Maximum duration in seconds for one compilation pass.
+            Defaults to None, which disables pass timeouts.
 
     Returns:
         A tuple containing the compiled quantum circuit, the compilation information, and the name of the device used for compilation.
+
+    Raises:
+        ValueError: If ``pass_timeout`` is not positive.
     """
     predicted_device = predict_device_for_figure_of_merit(qc, figure_of_merit=figure_of_merit)
     res = rl_compile(
-        qc, device=predicted_device, figure_of_merit=figure_of_merit, tracer_output_path=tracer_output_path
+        qc,
+        device=predicted_device,
+        figure_of_merit=figure_of_merit,
+        tracer_output_path=tracer_output_path,
+        pass_timeout=pass_timeout,
     )
     return *res, predicted_device

src/mqt/predictor/rl/predictor.py

@@ -72,26 +72,32 @@ def compile_as_predicted(
         self,
         qc: QuantumCircuit | str,
         tracer_output_path: str | Path | None = None,
+        pass_timeout: float | None = None,
     ) -> tuple[QuantumCircuit, list[str]]:
         """Compiles a given quantum circuit such that the given figure of merit is maximized by using the respectively trained optimized compiler.
 
         Arguments:
             qc: The quantum circuit to be compiled or the path to a qasm file containing the quantum circuit.
             tracer_output_path: Optional temporary path to export the compilation trace for this specific run.
+            pass_timeout: Maximum duration in seconds for one compilation pass.
+                Defaults to None, which disables pass timeouts.
 
         Returns:
             A tuple containing the compiled quantum circuit and the compilation information. If compilation fails, False is returned.
 
         Raises:
             RuntimeError: If an error occurs during compilation.
+            ValueError: If ``pass_timeout`` is not positive.
         """
         original_tracer_output_path = self.env.tracer_output_path
-
-        # Temporarily override singleton if a new path is explicitly provided
-        if tracer_output_path is not None:
-            self.env.tracer_output_path = tracer_output_path
+        original_pass_timeout = self.env.pass_timeout
 
         try:
+            # Temporarily override singleton settings for this compilation.
+            if tracer_output_path is not None:
+                self.env.tracer_output_path = tracer_output_path
+            self.env.pass_timeout = pass_timeout
+
             trained_rl_model = load_model(self.model_name)
 
             obs, _ = self.env.reset(qc, seed=0)
@@ -114,15 +120,16 @@ def compile_as_predicted(
             raise RuntimeError(msg)
 
         finally:
-            # Restore original singleton path
             self.env.tracer_output_path = original_tracer_output_path
+            self.env.pass_timeout = original_pass_timeout
 
     def train_model(
         self,
         timesteps: int = 1000,
         verbose: int = 2,
         test: bool = False,
         seed: int | None = None,
+        pass_timeout: float | None = None,
     ) -> None:
         """Trains all models for the given reward functions and device.
 
@@ -132,6 +139,11 @@ def train_model(
             test: Whether to train the model for testing purposes. Defaults to False.
             seed: The random seed to use for reproducible training. Set to None to use true randomness.
                 Defaults to None.
+            pass_timeout: Maximum duration in seconds for one compilation pass.
+                Defaults to None, which disables pass timeouts.
+
+        Raises:
+            ValueError: If ``pass_timeout`` is not positive.
         """
         if seed is not None:
             set_random_seed(seed)
@@ -148,22 +160,27 @@ def train_model(
             batch_size = 64
             progress_bar = True
 
-        logger.debug("Start training for: " + self.figure_of_merit + " on " + self.device_name)
-        model = MaskablePPO(
-            MaskableMultiInputActorCriticPolicy,
-            self.env,
-            verbose=verbose,
-            tensorboard_log=f"./{self.model_name}",
-            gamma=0.98,
-            n_steps=n_steps,
-            batch_size=batch_size,
-            n_epochs=n_epochs,
-            seed=seed,
-        )
-        # Training Loop: In each iteration, the agent collects n_steps steps (rollout),
-        # updates the policy for n_epochs, and then repeats the process until total_timesteps steps have been taken.
-        model.learn(total_timesteps=timesteps, progress_bar=progress_bar)
-        model.save(get_path_trained_model() / self.model_name)
+        original_pass_timeout = self.env.pass_timeout
+        self.env.pass_timeout = pass_timeout
+        try:
+            logger.debug("Start training for: " + self.figure_of_merit + " on " + self.device_name)
+            model = MaskablePPO(
+                MaskableMultiInputActorCriticPolicy,
+                self.env,
+                verbose=verbose,
+                tensorboard_log=f"./{self.model_name}",
+                gamma=0.98,
+                n_steps=n_steps,
+                batch_size=batch_size,
+                n_epochs=n_epochs,
+                seed=seed,
+            )
+            # Training Loop: In each iteration, the agent collects n_steps steps (rollout),
+            # updates the policy for n_epochs, and then repeats the process until total_timesteps steps have been taken.
+            model.learn(total_timesteps=timesteps, progress_bar=progress_bar)
+            model.save(get_path_trained_model() / self.model_name)
+        finally:
+            self.env.pass_timeout = original_pass_timeout
 
 
 def load_model(model_name: str) -> MaskablePPO:
@@ -194,6 +211,7 @@ def rl_compile(
     predictor_singleton: Predictor | None = None,
     tracer_output_path: str | Path | None = None,
     mdp: MDPPolicy = "v3",
+    pass_timeout: float | None = None,
 ) -> tuple[QuantumCircuit, list[str]]:
     """Compiles a given quantum circuit to a device optimizing for the given figure of merit.
 
@@ -206,12 +224,15 @@ def rl_compile(
         mdp: The MDP transition policy used when constructing a predictor. ``v2``
             is the original strategy and ``v3`` is the default. When
             ``predictor_singleton`` is provided, its configured policy is used instead.
+        pass_timeout: Maximum duration in seconds for one compilation pass.
+            Defaults to None, which disables pass timeouts.
 
     Returns:
         A tuple containing the compiled quantum circuit and the compilation information. If compilation fails, False is returned.
 
     Raises:
-        ValueError: If figure_of_merit or device is None and predictor_singleton is also None.
+        ValueError: If figure_of_merit or device is None and predictor_singleton is also None,
+            or if ``pass_timeout`` is not positive.
     """
     if predictor_singleton is None:
         if figure_of_merit is None:
@@ -226,6 +247,8 @@ def rl_compile(
             tracer_output_path=tracer_output_path,
             mdp=mdp,
         )
-        return predictor.compile_as_predicted(qc)
+        return predictor.compile_as_predicted(qc, pass_timeout=pass_timeout)
 
-    return predictor_singleton.compile_as_predicted(qc, tracer_output_path=tracer_output_path)
+    return predictor_singleton.compile_as_predicted(
+        qc, tracer_output_path=tracer_output_path, pass_timeout=pass_timeout
+    )

src/mqt/predictor/rl/predictorenv.py

@@ -10,14 +10,20 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import re
+import signal
+import threading
 import time
 import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal, get_args
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from types import FrameType
+
     from gymnasium.spaces import Space
     from qiskit.transpiler import Layout, Target
 
@@ -60,6 +66,43 @@
 MDP_POLICIES: frozenset[str] = frozenset(get_args(MDPPolicy))
 
 
+@contextlib.contextmanager
+def _enforce_pass_timeout(pass_timeout: float | None) -> Iterator[None]:
+    if pass_timeout is None:
+        yield
+        return
+
+    if not all(hasattr(signal, attribute) for attribute in ("SIGALRM", "ITIMER_REAL", "getitimer", "setitimer")):
+        warnings.warn("Pass timeouts are not supported on this platform.", RuntimeWarning, stacklevel=2)
+        yield
+        return
+    if threading.current_thread() is not threading.main_thread():
+        warnings.warn("Pass timeouts are only supported on the main thread.", RuntimeWarning, stacklevel=2)
+        yield
+        return
+
+    def timeout_handler(_signum: int, _frame: FrameType | None) -> None:
+        msg = f"Compilation pass exceeded the timeout of {pass_timeout:g} seconds."
+        raise TimeoutError(msg)
+
+    previous_delay, previous_interval = signal.getitimer(signal.ITIMER_REAL)
+    if 0 < previous_delay <= pass_timeout:
+        yield
+        return
+
+    previous_handler = signal.signal(signal.SIGALRM, timeout_handler)
+    start_time = time.monotonic()
+    try:
+        signal.setitimer(signal.ITIMER_REAL, pass_timeout)
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, 0)
+        signal.signal(signal.SIGALRM, previous_handler)
+        if previous_delay > 0:
+            remaining_delay = max(previous_delay - (time.monotonic() - start_time), 1e-6)
+            signal.setitimer(signal.ITIMER_REAL, remaining_delay, previous_interval)
+
+
 class PredictorEnv(Env):
     """Predictor environment for reinforcement learning."""
 
@@ -71,6 +114,7 @@ def __init__(
         max_steps: int | None = None,
         tracer_output_path: str | Path | None = None,
         mdp: MDPPolicy = "v3",
+        pass_timeout: float | None = None,
     ) -> None:
         """Initializes the PredictorEnv object.
 
@@ -83,9 +127,12 @@ def __init__(
             mdp: The MDP transition policy. ``v2`` is the original MQT Predictor
                 strategy. ``v3`` is the default and is permissive before layout while
                 preserving the established compilation structure afterwards.
+            pass_timeout: Maximum duration in seconds for one compilation pass.
+                Defaults to None, which disables pass timeouts.
 
         Raises:
-            ValueError: If ``mdp`` is unsupported, if the reward function is
+            ValueError: If ``mdp`` is unsupported, ``pass_timeout`` is not positive,
+                if the reward function is
                 "estimated_success_probability" and no calibration data is available
                 for the device, or if the reward function is
                 "estimated_hellinger_distance" and no trained model is available for
@@ -100,6 +147,7 @@ def __init__(
         self.path_training_circuits = path_training_circuits or get_path_training_circuits()
         self.max_steps = max_steps
         self.mdp = mdp
+        self.pass_timeout = pass_timeout
 
         self.action_set = {}
         self.actions_synthesis_indices = []
@@ -201,6 +249,18 @@ def __init__(
         self.observation_space = Dict(spaces)
         self.filename = ""
 
+    @property
+    def pass_timeout(self) -> float | None:
+        """The current per-pass timeout in seconds."""
+        return self._pass_timeout
+
+    @pass_timeout.setter
+    def pass_timeout(self, pass_timeout: float | None) -> None:
+        if pass_timeout is not None and pass_timeout <= 0:
+            msg = "pass_timeout must be positive."
+            raise ValueError(msg)
+        self._pass_timeout = pass_timeout
+
     def _collect_tracer_data(
         self,
         step_index: int,
@@ -290,7 +350,8 @@ def step(self, action: int) -> tuple[dict[str, Any], float, bool, bool, dict[Any
         start_time = time.perf_counter()
         try:
             self.used_actions.append(action_name)
-            altered_qc = self.apply_action(action)
+            with _enforce_pass_timeout(self.pass_timeout):
+                altered_qc = self.apply_action(action)
             action_duration = time.perf_counter() - start_time
         except Exception as exc:  # ruff:ignore[blind-except]
             action_duration = time.perf_counter() - start_time