Merge pull request #3008 from autonomys/cpu-plotting-cli-refactoring

CPU plotting CLI refactoring

Author: nazar-pc

crates/subspace-farmer/src/bin/subspace-farmer/commands/cluster/plotter.rs

@@ -1,5 +1,6 @@
 use crate::commands::shared::PlottingThreadPriority;
 use anyhow::anyhow;
+use async_lock::Mutex as AsyncMutex;
 use clap::Parser;
 use futures::{select, FutureExt};
 use prometheus_client::registry::Registry;
@@ -17,65 +18,73 @@ use subspace_farmer::plotter::cpu::CpuPlotter;
 use subspace_farmer::utils::{
     create_plotting_thread_pool_manager, parse_cpu_cores_sets, thread_pool_core_indices,
 };
+use subspace_farmer_components::PieceGetter;
 use subspace_proof_of_space::Table;
 use tokio::sync::Semaphore;
 use tracing::info;
 
-/// Arguments for plotter
 #[derive(Debug, Parser)]
-pub(super) struct PlotterArgs {
-    /// Piece getter concurrency.
-    ///
-    /// Increase can result in NATS communication issues if too many messages arrive via NATS, but
-    /// are not processed quickly enough for some reason and might require increasing cluster-level
-    /// `--nats-pool-size` parameter.
-    #[arg(long, default_value = "32")]
-    piece_getter_concurrency: NonZeroUsize,
+struct CpuPlottingOptions {
     /// Defines how many sectors farmer will download concurrently, allows to limit memory usage of
-    /// the plotting process, defaults to `--sector-encoding-concurrency` + 1 to download future
+    /// the plotting process, defaults to `--cpu-sector-encoding-concurrency` + 1 to download future
     /// sector ahead of time.
     ///
     /// Increase will result in higher memory usage.
     #[arg(long)]
-    sector_downloading_concurrency: Option<NonZeroUsize>,
+    cpu_sector_downloading_concurrency: Option<NonZeroUsize>,
     /// Defines how many sectors farmer will encode concurrently, defaults to 1 on UMA system and
     /// number of NUMA nodes on NUMA system or L3 cache groups on large CPUs. It is further
     /// restricted by
-    /// `--sector-downloading-concurrency` and setting this option higher than
-    /// `--sector-downloading-concurrency` will have no effect.
+    /// `--cpu-sector-downloading-concurrency` and setting this option higher than
+    /// `--cpu-sector-downloading-concurrency` will have no effect.
     ///
     /// Increase will result in higher memory usage.
     #[arg(long)]
-    sector_encoding_concurrency: Option<NonZeroUsize>,
+    cpu_sector_encoding_concurrency: Option<NonZeroUsize>,
     /// Defines how many records farmer will encode in a single sector concurrently, defaults to one
     /// record per 2 cores, but not more than 8 in total. Higher concurrency means higher memory
     /// usage and typically more efficient CPU utilization.
     #[arg(long)]
-    record_encoding_concurrency: Option<NonZeroUsize>,
+    cpu_record_encoding_concurrency: Option<NonZeroUsize>,
     /// Size of one thread pool used for plotting, defaults to number of logical CPUs available
     /// on UMA system and number of logical CPUs available in NUMA node on NUMA system or L3 cache
     /// groups on large CPUs.
     ///
-    /// Number of thread pools is defined by `--sector-encoding-concurrency` option, different
+    /// Number of thread pools is defined by `--cpu-sector-encoding-concurrency` option, different
     /// thread pools might have different number of threads if NUMA nodes do not have the same size.
     ///
     /// Threads will be pinned to corresponding CPU cores at creation.
     #[arg(long)]
-    plotting_thread_pool_size: Option<NonZeroUsize>,
+    cpu_plotting_thread_pool_size: Option<NonZeroUsize>,
     /// Specify exact CPU cores to be used for plotting bypassing any custom logic farmer might use
-    /// otherwise. It replaces both `--sector-encoding-concurrency` and
-    /// `--plotting-thread-pool-size` options if specified.
+    /// otherwise. It replaces both `--cpu-sector-encoding-concurrency` and
+    /// `--cpu-plotting-thread-pool-size` options if specified.
     ///
     /// Cores are coma-separated, with whitespace separating different thread pools/encoding
     /// instances. For example "0,1 2,3" will result in two sectors being encoded at the same time,
     /// each with a pair of CPU cores.
-    #[arg(long, conflicts_with_all = & ["sector_encoding_concurrency", "plotting_thread_pool_size"])]
-    plotting_cpu_cores: Option<String>,
+    #[arg(long, conflicts_with_all = & ["cpu_sector_encoding_concurrency", "cpu_plotting_thread_pool_size"])]
+    cpu_plotting_cores: Option<String>,
     /// Plotting thread priority, by default de-prioritizes plotting threads in order to make sure
     /// farming is successful and computer can be used comfortably for other things. Can be set to
     /// "min", "max" or "default".
     #[arg(long, default_value_t = PlottingThreadPriority::Min)]
-    plotting_thread_priority: PlottingThreadPriority,
+    cpu_plotting_thread_priority: PlottingThreadPriority,
+}
+
+/// Arguments for plotter
+#[derive(Debug, Parser)]
+pub(super) struct PlotterArgs {
+    /// Piece getter concurrency.
+    ///
+    /// Increase can result in NATS communication issues if too many messages arrive via NATS, but
+    /// are not processed quickly enough for some reason and might require increasing cluster-level
+    /// `--nats-pool-size` parameter.
+    #[arg(long, default_value = "32")]
+    piece_getter_concurrency: NonZeroUsize,
+    /// Plotting options only used by CPU plotter
+    #[clap(flatten)]
+    cpu_plotting_options: CpuPlottingOptions,
     /// Additional cluster components
     #[clap(raw = true)]
     pub(super) additional_components: Vec<String>,
@@ -92,12 +101,7 @@ where
 {
     let PlotterArgs {
         piece_getter_concurrency,
-        sector_downloading_concurrency,
-        sector_encoding_concurrency,
-        record_encoding_concurrency,
-        plotting_thread_pool_size,
-        plotting_cpu_cores,
-        plotting_thread_priority,
+        cpu_plotting_options,
         additional_components: _,
     } = plotter_args;
 
@@ -109,13 +113,62 @@ where
     .map_err(|error| anyhow!("Failed to instantiate erasure coding: {error}"))?;
     let piece_getter = ClusterPieceGetter::new(nats_client.clone(), piece_getter_concurrency);
 
+    let global_mutex = Arc::default();
+
+    let (legacy_cpu_plotter, modern_cpu_plotter) = init_cpu_plotters::<_, PosTableLegacy, PosTable>(
+        cpu_plotting_options,
+        piece_getter,
+        global_mutex,
+        kzg,
+        erasure_coding,
+        registry,
+    )?;
+    let legacy_cpu_plotter = Arc::new(legacy_cpu_plotter);
+    let modern_cpu_plotter = Arc::new(modern_cpu_plotter);
+
+    Ok(Box::pin(async move {
+        select! {
+            result = plotter_service(&nats_client, &legacy_cpu_plotter, false).fuse() => {
+                result.map_err(|error| anyhow!("Plotter service failed: {error}"))
+            }
+            result = plotter_service(&nats_client, &modern_cpu_plotter, true).fuse() => {
+                result.map_err(|error| anyhow!("Plotter service failed: {error}"))
+            }
+        }
+    }))
+}
+
+fn init_cpu_plotters<PG, PosTableLegacy, PosTable>(
+    cpu_plotting_options: CpuPlottingOptions,
+    piece_getter: PG,
+    global_mutex: Arc<AsyncMutex<()>>,
+    kzg: Kzg,
+    erasure_coding: ErasureCoding,
+    registry: &mut Registry,
+) -> anyhow::Result<(CpuPlotter<PG, PosTableLegacy>, CpuPlotter<PG, PosTable>)>
+where
+    PG: PieceGetter + Clone + Send + Sync + 'static,
+    PosTableLegacy: Table,
+    PosTable: Table,
+{
+    let CpuPlottingOptions {
+        cpu_sector_downloading_concurrency,
+        cpu_sector_encoding_concurrency,
+        cpu_record_encoding_concurrency,
+        cpu_plotting_thread_pool_size,
+        cpu_plotting_cores,
+        cpu_plotting_thread_priority,
+    } = cpu_plotting_options;
+
     let plotting_thread_pool_core_indices;
-    if let Some(plotting_cpu_cores) = plotting_cpu_cores {
-        plotting_thread_pool_core_indices = parse_cpu_cores_sets(&plotting_cpu_cores)
-            .map_err(|error| anyhow!("Failed to parse `--plotting-cpu-cores`: {error}"))?;
+    if let Some(cpu_plotting_cores) = cpu_plotting_cores {
+        plotting_thread_pool_core_indices = parse_cpu_cores_sets(&cpu_plotting_cores)
+            .map_err(|error| anyhow!("Failed to parse `--cpu-plotting-cpu-cores`: {error}"))?;
     } else {
-        plotting_thread_pool_core_indices =
-            thread_pool_core_indices(plotting_thread_pool_size, sector_encoding_concurrency);
+        plotting_thread_pool_core_indices = thread_pool_core_indices(
+            cpu_plotting_thread_pool_size,
+            cpu_sector_encoding_concurrency,
+        );
 
         if plotting_thread_pool_core_indices.len() > 1 {
             info!(
@@ -126,12 +179,12 @@ where
     }
 
     let downloading_semaphore = Arc::new(Semaphore::new(
-        sector_downloading_concurrency
-            .map(|sector_downloading_concurrency| sector_downloading_concurrency.get())
+        cpu_sector_downloading_concurrency
+            .map(|cpu_sector_downloading_concurrency| cpu_sector_downloading_concurrency.get())
             .unwrap_or(plotting_thread_pool_core_indices.len() + 1),
     ));
 
-    let record_encoding_concurrency = record_encoding_concurrency.unwrap_or_else(|| {
+    let cpu_record_encoding_concurrency = cpu_record_encoding_concurrency.unwrap_or_else(|| {
         let cpu_cores = plotting_thread_pool_core_indices
             .first()
             .expect("Guaranteed to have some CPU cores; qed");
@@ -157,39 +210,30 @@ where
         plotting_thread_pool_core_indices
             .into_iter()
             .zip(replotting_thread_pool_core_indices),
-        plotting_thread_priority.into(),
+        cpu_plotting_thread_priority.into(),
     )
     .map_err(|error| anyhow!("Failed to create thread pool manager: {error}"))?;
-    let global_mutex = Arc::default();
-    let legacy_cpu_plotter = Arc::new(CpuPlotter::<_, PosTableLegacy>::new(
+
+    let legacy_cpu_plotter = CpuPlotter::<_, PosTableLegacy>::new(
         piece_getter.clone(),
         Arc::clone(&downloading_semaphore),
         plotting_thread_pool_manager.clone(),
-        record_encoding_concurrency,
+        cpu_record_encoding_concurrency,
         Arc::clone(&global_mutex),
         kzg.clone(),
         erasure_coding.clone(),
         Some(registry),
-    ));
-    let modern_cpu_plotter = Arc::new(CpuPlotter::<_, PosTable>::new(
-        piece_getter.clone(),
+    );
+    let modern_cpu_plotter = CpuPlotter::<_, PosTable>::new(
+        piece_getter,
         downloading_semaphore,
         plotting_thread_pool_manager,
-        record_encoding_concurrency,
-        Arc::clone(&global_mutex),
-        kzg.clone(),
-        erasure_coding.clone(),
+        cpu_record_encoding_concurrency,
+        global_mutex,
+        kzg,
+        erasure_coding,
         Some(registry),
-    ));
+    );
 
-    Ok(Box::pin(async move {
-        select! {
-            result = plotter_service(&nats_client, &legacy_cpu_plotter, false).fuse() => {
-                result.map_err(|error| anyhow!("Plotter service failed: {error}"))
-            }
-            result = plotter_service(&nats_client, &modern_cpu_plotter, true).fuse() => {
-                result.map_err(|error| anyhow!("Plotter service failed: {error}"))
-            }
-        }
-    }))
+    Ok((legacy_cpu_plotter, modern_cpu_plotter))
 }