From fe7d40e4a176375c22fae5e52d42aeaf9ab38d33 Mon Sep 17 00:00:00 2001 From: Luca Mondada <72734770+lmondada@users.noreply.github.com> Date: Thu, 25 Jul 2024 12:04:18 +0100 Subject: [PATCH] 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`. --- badger-optimiser/src/main.rs | 6 +++--- tket2-py/src/optimiser.rs | 6 +++--- tket2-py/src/passes.rs | 16 +++++++++++++--- tket2-py/tket2/_tket2/passes.pyi | 14 +++++++++++--- tket2-py/tket2/passes.py | 7 +++++-- tket2/src/optimiser/badger.rs | 12 ++++++------ 6 files changed, 41 insertions(+), 20 deletions(-) diff --git a/badger-optimiser/src/main.rs b/badger-optimiser/src/main.rs index 399a5d7f..33a755d6 100644 --- a/badger-optimiser/src/main.rs +++ b/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, + max_circuit_count: Option, /// Number of threads (default=1) #[arg( short = 'j', @@ -176,7 +176,7 @@ fn main() -> Result<(), Box> { 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, }, ); diff --git a/tket2-py/src/optimiser.rs b/tket2-py/src/optimiser.rs index e5bb6cde..ac04c667 100644 --- a/tket2-py/src/optimiser.rs +++ b/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, progress_timeout: Option, - max_circuit_cnt: Option, + max_circuit_count: Option, n_threads: Option, split_circ: Option, queue_size: Option, @@ -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), diff --git a/tket2-py/src/passes.rs b/tket2-py/src/passes.rs index a747738a..eb463be1 100644 --- a/tket2-py/src/passes.rs +++ b/tket2-py/src/passes.rs @@ -102,18 +102,27 @@ fn lower_to_pytket<'py>(circ: &Bound<'py, PyAny>) -> PyResult> /// 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, timeout: Option, progress_timeout: Option, + max_circuit_count: Option, log_dir: Option, rebase: Option, ) -> PyResult> { @@ -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); diff --git a/tket2-py/tket2/_tket2/passes.pyi b/tket2-py/tket2/_tket2/passes.pyi index e0d19775..6d85cfc0 100644 --- a/tket2-py/tket2/_tket2/passes.pyi +++ b/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. """ diff --git a/tket2-py/tket2/passes.py b/tket2-py/tket2/passes.py index e63a5087..c10333dc 100644 --- a/tket2-py/tket2/passes.py +++ b/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 `_ 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, ) diff --git a/tket2/src/optimiser/badger.rs b/tket2/src/optimiser/badger.rs index 75741928..0df2505d 100644 --- a/tket2/src/optimiser/badger.rs +++ b/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, + pub max_circuit_count: Option, /// 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();