Parallelise select_best_assets

Cargo.toml

@@ -44,6 +44,7 @@ dirs = "6.0.0"
 edit = "0.1.5"
 erased-serde = "0.4.10"
 context_manager = "0.1.3"
+rayon = "1.12.0"
 
 [dev-dependencies]
 assert_cmd = "2.2.2"

benches/assets.rs

@@ -16,8 +16,9 @@ use muse2::simulation::market::{
 };
 use muse2::simulation::optimisation::DispatchRun;
 use muse2::simulation::prices::{Prices, calculate_prices};
+use rayon::ThreadPoolBuilder;
 use std::hint::black_box;
-use std::rc::Rc;
+use std::sync::Arc;
 use std::time::Duration;
 use tempfile::TempDir;
 
@@ -43,8 +44,8 @@ fn load_bench_model() -> (Model, DataWriter, TempDir, TempDir) {
     let example_dir = TempDir::new().expect("Failed to create temp dir for example");
     let model_path = example_dir.path().join(EXAMPLE_NAME);
     Example::from_name(EXAMPLE_NAME)
-        .expect("Invalid example name")
-        .extract(&model_path)
+        .and_then(|example| example.extract(&model_path))
+        .expect("Failed to extract example");
     let model = load_model(&model_path).expect("Failed to load model");
 
     let output_dir = TempDir::new().expect("Failed to create temp dir for output");
@@ -114,11 +115,11 @@ fn calculate_seed_prices(
 ///
 /// Each synthetic process is a copy of one of `templates`, but given a unique ID so that it is
 /// treated as a distinct candidate technology.
-fn build_synthetic_processes(templates: &[Rc<Process>], n: usize) -> Vec<Rc<Process>> {
+fn build_synthetic_processes(templates: &[Arc<Process>], n: usize) -> Vec<Arc<Process>> {
     (0..n)
         .map(|i| {
             let template = &templates[i % templates.len()];
-            Rc::new(Process {
+            Arc::new(Process {
                 id: ProcessID::new(&format!("{}#{i}", template.id)),
                 ..(**template).clone()
             })
@@ -128,6 +129,11 @@ fn build_synthetic_processes(templates: &[Rc<Process>], n: usize) -> Vec<Rc<Proc
 
 /// Benchmark [`select_best_assets`] using inputs derived from the `two_outputs` example model,
 /// scaling the number of competing candidate technologies over `N_TECHNOLOGIES_RANGE`.
+///
+/// Two groups are run:
+/// - `parallel`: uses the default Rayon global thread pool (all available cores).
+/// - `sequential`: installs a single-thread Rayon pool so that `par_iter` degenerates to serial
+///   execution, giving a fair like-for-like comparison with the overhead of parallelism removed.
 fn criterion_benchmark(c: &mut Criterion) {
     let (model, mut writer, _example_dir, _output_dir) = load_bench_model();
     let (base_year_assets, existing_assets, candidates) = build_assets_for_investment_year(&model);
@@ -157,67 +163,84 @@ fn criterion_benchmark(c: &mut Criterion) {
 
     // Real candidate technologies for this market, used as templates to build up to
     // `N_TECHNOLOGIES_RANGE.end()` synthetic competing technologies
-    let templates: Vec<Rc<Process>> = agent
+    let templates: Vec<Arc<Process>> = agent
         .iter_search_space(region_id, &commodity.id, YEAR)
         .cloned()
         .collect();
 
-    let mut group = c.benchmark_group("select_best_assets");
-    group
-        .noise_threshold(0.05)
-        .sample_size(20)
-        .measurement_time(Duration::from_secs(3));
-
-    for n in N_TECHNOLOGIES_RANGE {
-        // Give the agent a synthetic search space of `n` competing technologies for this market
-        let mut agent = agent.clone();
-        agent.search_space.insert(
-            (commodity.id.clone(), region_id.clone(), YEAR),
-            Rc::new(build_synthetic_processes(&templates, n)),
-        );
-
-        let opt_assets: Vec<AssetRef> = get_asset_options(
-            &existing_assets,
-            &demand,
-            &agent,
-            commodity,
-            region_id,
-            YEAR,
-            model.parameters.capacity_limit_factor,
-        )
-        .collect();
-        let investment_limits =
-            collect_investment_limits_for_candidates(&opt_assets, commodity_portion);
-
-        group.bench_with_input(BenchmarkId::from_parameter(n), &n, |b, _| {
-            b.iter_batched(
-                || {
-                    (
-                        opt_assets.clone(),
-                        investment_limits.clone(),
-                        demand.clone(),
-                    )
-                },
-                |(opt_assets, investment_limits, demand)| {
-                    select_best_assets(
-                        black_box(&model),
-                        opt_assets,
-                        investment_limits,
-                        black_box(commodity),
-                        black_box(&agent),
-                        black_box(region_id),
-                        black_box(&prices),
-                        demand,
-                        black_box(YEAR),
-                        &mut writer,
-                    )
-                    .expect("select_best_assets failed")
-                },
-                BatchSize::SmallInput,
+    // Single-thread pool used to run `select_best_assets` sequentially for comparison.
+    // Because `select_best_assets` uses `par_iter` internally, installing this pool makes it
+    // degenerate to serial execution while keeping all other code paths identical.
+    let sequential_pool = ThreadPoolBuilder::new()
+        .num_threads(1)
+        .build()
+        .expect("Failed to build sequential thread pool");
+
+    for (group_name, use_parallel) in &[("parallel", true), ("sequential", false)] {
+        let mut group = c.benchmark_group(format!("select_best_assets/{group_name}"));
+        group
+            .noise_threshold(0.05)
+            .sample_size(20)
+            .measurement_time(Duration::from_secs(3));
+
+        for n in N_TECHNOLOGIES_RANGE {
+            // Give the agent a synthetic search space of `n` competing technologies
+            let mut agent = agent.clone();
+            agent.search_space.insert(
+                (commodity.id.clone(), region_id.clone(), YEAR),
+                Arc::new(build_synthetic_processes(&templates, n)),
             );
-        });
+
+            let opt_assets: Vec<AssetRef> = get_asset_options(
+                &existing_assets,
+                &demand,
+                &agent,
+                commodity,
+                region_id,
+                YEAR,
+                model.parameters.capacity_limit_factor,
+            )
+            .collect();
+            let investment_limits =
+                collect_investment_limits_for_candidates(&opt_assets, commodity_portion);
+
+            group.bench_with_input(BenchmarkId::from_parameter(n), &n, |b, _| {
+                b.iter_batched(
+                    || {
+                        (
+                            opt_assets.clone(),
+                            investment_limits.clone(),
+                            demand.clone(),
+                        )
+                    },
+                    |(opt_assets, investment_limits, demand)| {
+                        let run = || {
+                            select_best_assets(
+                                black_box(&model),
+                                opt_assets,
+                                investment_limits,
+                                black_box(commodity),
+                                black_box(&agent),
+                                black_box(region_id),
+                                black_box(&prices),
+                                demand,
+                                black_box(YEAR),
+                                &mut writer,
+                            )
+                            .expect("select_best_assets failed")
+                        };
+                        if *use_parallel {
+                            run()
+                        } else {
+                            sequential_pool.install(run)
+                        }
+                    },
+                    BatchSize::SmallInput,
+                );
+            });
+        }
+        group.finish();
     }
-    group.finish();
 }
 
 criterion_group!(benches, criterion_benchmark);

src/agent.rs

@@ -8,7 +8,7 @@ use crate::units::Dimensionless;
 use indexmap::{IndexMap, IndexSet};
 use serde::Deserialize;
 use std::collections::HashMap;
-use std::rc::Rc;
+use std::sync::Arc;
 
 define_id_type! {AgentID, "agent ID"}
 
@@ -19,7 +19,7 @@ pub type AgentMap = IndexMap<AgentID, Agent>;
 pub type AgentCommodityPortionsMap = HashMap<(CommodityID, u32), Dimensionless>;
 
 /// A map for the agent's search space, keyed by commodity, region, and year
-pub type AgentSearchSpaceMap = HashMap<(CommodityID, RegionID, u32), Rc<Vec<Rc<Process>>>>;
+pub type AgentSearchSpaceMap = HashMap<(CommodityID, RegionID, u32), Arc<Vec<Arc<Process>>>>;
 
 /// A map of objectives for an agent, keyed by year.
 ///
@@ -59,7 +59,7 @@ impl Agent {
         region_id: &RegionID,
         commodity_id: &CommodityID,
         year: u32,
-    ) -> impl Iterator<Item = &Rc<Process>> {
+    ) -> impl Iterator<Item = &Arc<Process>> {
         self.search_space[&(commodity_id.clone(), region_id.clone(), year)].iter()
     }
 }