feat: Expose advanced Badger timeout options to tket2-py (#506)

#496 introduced a new timeout criterion for Badger based on the number
of circuits seen during optimisation. This was not consistently exposed
to the Python API.

Fly-by: The `progress_timeout` timeout criterion that was added in #259
was also not exposed properly. This PR also fixes this.

The timeout option `max_circuit_cnt` is now called `max_circuit_count`.

Author: lmondada

badger-optimiser/src/main.rs

@@ -85,10 +85,10 @@ struct CmdLineArgs {
     #[arg(
         short = 'c',
         long,
-        value_name = "MAX_CIRCUIT_CNT",
+        value_name = "MAX_CIRCUIT_COUNT",
         help = "Maximum number of circuits to process (default=None)."
     )]
-    max_circuit_cnt: Option<usize>,
+    max_circuit_count: Option<usize>,
     /// Number of threads (default=1)
     #[arg(
         short = 'j',
@@ -176,7 +176,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             n_threads,
             split_circuit: opts.split_circ,
             queue_size: opts.queue_size,
-            max_circuit_cnt: opts.max_circuit_cnt,
+            max_circuit_count: opts.max_circuit_count,
         },
     );
 

tket2-py/src/optimiser.rs

@@ -61,7 +61,7 @@ impl PyBadgerOptimiser {
     ///     If `None` the optimiser will run indefinitely, or until `timeout` is
     ///     reached.
     ///
-    /// * `max_circuit_cnt`: The maximum number of circuits to process before
+    /// * `max_circuit_count`: The maximum number of circuits to process before
     ///     stopping the optimisation.
     ///
     ///
@@ -94,7 +94,7 @@ impl PyBadgerOptimiser {
         circ: &Bound<'py, PyAny>,
         timeout: Option<u64>,
         progress_timeout: Option<u64>,
-        max_circuit_cnt: Option<usize>,
+        max_circuit_count: Option<usize>,
         n_threads: Option<NonZeroUsize>,
         split_circ: Option<bool>,
         queue_size: Option<usize>,
@@ -103,7 +103,7 @@ impl PyBadgerOptimiser {
         let options = BadgerOptions {
             timeout,
             progress_timeout,
-            max_circuit_cnt,
+            max_circuit_count,
             n_threads: n_threads.unwrap_or(NonZeroUsize::new(1).unwrap()),
             split_circuit: split_circ.unwrap_or(false),
             queue_size: queue_size.unwrap_or(100),

tket2-py/src/passes.rs

@@ -102,18 +102,27 @@ fn lower_to_pytket<'py>(circ: &Bound<'py, PyAny>) -> PyResult<Bound<'py, PyAny>>
 /// optimising. This can be deactivated by setting `rebase` to `false`, in which
 /// case the circuit is expected to be in the Nam gate set.
 ///
-/// Will use at most `max_threads` threads (plus a constant) and take at most
-/// `timeout` seconds (plus a constant). Default to the number of cpus and
-/// 15min respectively.
+/// Will use at most `max_threads` threads (plus a constant). Defaults to the
+/// number of CPUs available.
+///
+/// The optimisation will terminate at the first of the following timeout
+/// criteria, if set:
+/// - `timeout` seconds (default: 15min) have elapsed since the start of the
+///    optimisation
+/// - `progress_timeout` (default: None) seconds have elapsed since progress
+///    in the cost function was last made
+/// - `max_circuit_count` (default: None) circuits have been explored.
 ///
 /// Log files will be written to the directory `log_dir` if specified.
 #[pyfunction]
+#[allow(clippy::too_many_arguments)]
 fn badger_optimise<'py>(
     circ: &Bound<'py, PyAny>,
     optimiser: &PyBadgerOptimiser,
     max_threads: Option<NonZeroUsize>,
     timeout: Option<u64>,
     progress_timeout: Option<u64>,
+    max_circuit_count: Option<usize>,
     log_dir: Option<PathBuf>,
     rebase: Option<bool>,
 ) -> PyResult<Bound<'py, PyAny>> {
@@ -165,6 +174,7 @@ fn badger_optimise<'py>(
                 progress_timeout,
                 n_threads: n_threads.try_into().unwrap(),
                 split_circuit: true,
+                max_circuit_count,
                 ..Default::default()
             };
             circ = optimiser.optimise(circ, log_file, options);

tket2-py/tket2/_tket2/passes.pyi

@@ -35,6 +35,7 @@ def badger_optimise(
     max_threads: int | None = None,
     timeout: int | None = None,
     progress_timeout: int | None = None,
+    max_circuit_count: int | None = None,
     log_dir: Path | None = None,
     rebase: bool | None = False,
 ) -> CircuitClass:
@@ -47,9 +48,16 @@ def badger_optimise(
     optimising. This can be deactivated by setting `rebase` to `false`, in which
     case the circuit is expected to be in the Nam gate set.
 
-    Will use at most `max_threads` threads (plus a constant) and take at most
-    `timeout` seconds (plus a constant). Default to the number of cpus and
-    15min respectively.
+    Will use at most `max_threads` threads (plus a constant). Defaults to the
+    number of CPUs available.
+
+    The optimisation will terminate at the first of the following timeout
+    criteria, if set:
+    - `timeout` seconds (default: 15min) have elapsed since the start of the
+      optimisation
+    - `progress_timeout` (default: None) seconds have elapsed since progress
+      in the cost function was last made
+    - `max_circuit_count` (default: None) circuits have been explored.
 
     Log files will be written to the directory `log_dir` if specified.
     """

tket2-py/tket2/passes.py

@@ -35,6 +35,7 @@ def badger_pass(
     max_threads: Optional[int] = None,
     timeout: Optional[int] = None,
     progress_timeout: Optional[int] = None,
+    max_circuit_count: Optional[int] = None,
     log_dir: Optional[Path] = None,
     rebase: bool = False,
 ) -> BasePass:
@@ -44,8 +45,9 @@ def badger_pass(
     `compile-rewriter <https://github.com/CQCL/tket2/tree/main/badger-optimiser>`_
     utility. If `rewriter` is not specified, a default one will be used.
 
-    The arguments `max_threads`, `timeout`, `log_dir` and `rebase` are optional
-    and will be passed on to the Badger optimiser if provided."""
+    The arguments `max_threads`, `timeout`, `progress_timeout`, `max_circuit_count`,
+    `log_dir` and `rebase` are optional and will be passed on to the Badger
+    optimiser if provided."""
     if rewriter is None:
         with resources.as_file(
             resources.files("tket2").joinpath("data/nam_6_3.rwr")
@@ -61,6 +63,7 @@ def apply(circuit: Circuit) -> Circuit:
             max_threads=max_threads,
             timeout=timeout,
             progress_timeout=progress_timeout,
+            max_circuit_count=max_circuit_count,
             log_dir=log_dir,
             rebase=rebase,
         )

tket2/src/optimiser/badger.rs

@@ -57,7 +57,7 @@ pub struct BadgerOptions {
     /// per-thread basis, otherwise applies globally.
     ///
     /// Defaults to `None`, which means no limit.
-    pub max_circuit_cnt: Option<usize>,
+    pub max_circuit_count: Option<usize>,
     /// The number of threads to use.
     ///
     /// Defaults to `1`.
@@ -86,7 +86,7 @@ impl Default for BadgerOptions {
             n_threads: NonZeroUsize::new(1).unwrap(),
             split_circuit: Default::default(),
             queue_size: 20,
-            max_circuit_cnt: None,
+            max_circuit_count: None,
         }
     }
 }
@@ -251,8 +251,8 @@ where
                     break;
                 }
             }
-            if let Some(max_circuit_cnt) = opt.max_circuit_cnt {
-                if seen_hashes.len() >= max_circuit_cnt {
+            if let Some(max_circuit_count) = opt.max_circuit_count {
+                if seen_hashes.len() >= max_circuit_count {
                     timeout_flag = true;
                     break;
                 }
@@ -347,8 +347,8 @@ where
                         Ok(PriorityChannelLog::CircuitCount{processed_count: proc, seen_count: seen, queue_length}) => {
                             processed_count = proc;
                             seen_count = seen;
-                            if let Some(max_circuit_cnt) = opt.max_circuit_cnt {
-                                if seen_count > max_circuit_cnt {
+                            if let Some(max_circuit_count) = opt.max_circuit_count {
+                                if seen_count > max_circuit_count {
                                     timeout_flag = true;
                                     // Signal the workers to stop.
                                     let _ = pq.close();