Add WIP priority scheduler

Author: Guilucand

crates/api/src/lib.rs

@@ -16,6 +16,7 @@ use parallel_processor::enable_counters_logging;
 use parallel_processor::memory_data_size::MemoryDataSize;
 use parallel_processor::memory_fs::MemoryFs;
 use parallel_processor::phase_times_monitor::PHASES_TIMES_MONITOR;
+use parallel_processor::scheduler::PriorityScheduler;
 use parking_lot::Mutex;
 use std::cmp::max;
 use std::fs::create_dir_all;
@@ -233,7 +234,7 @@ impl GGCATInstance {
 
         disk_optimization_level: u32,
     ) -> anyhow::Result<PathBuf> {
-        // PriorityScheduler::set_max_threads_count(threads_count);
+        PriorityScheduler::set_max_threads_count(threads_count);
 
         let merging_hash_dispatch = utils::get_hash_static_id(
             debug::DEBUG_HASH_TYPE.lock().clone(),
@@ -334,7 +335,7 @@ impl GGCATInstance {
         // Query output format
         color_output_format: ColoredQueryOutputFormat,
     ) -> anyhow::Result<PathBuf> {
-        // PriorityScheduler::set_max_threads_count(threads_count);
+        PriorityScheduler::set_max_threads_count(threads_count);
 
         let merging_hash_dispatch = utils::get_hash_static_id(
             debug::DEBUG_HASH_TYPE.lock().clone(),
@@ -436,7 +437,7 @@ impl GGCATInstance {
         single_thread_output_function: bool,
         output_function: impl Fn(&[u8], &[ColorIndexType], bool) + Send + Sync,
     ) -> anyhow::Result<()> {
-        // PriorityScheduler::set_max_threads_count(threads_count);
+        PriorityScheduler::set_max_threads_count(threads_count);
 
         let temp_dir = create_tempdir(self.0.temp_dir.clone());
 

crates/assembler/src/lib.rs

@@ -14,7 +14,7 @@ use colors::colors_manager::ColorsMergeManager;
 use config::{
     get_compression_level_info, get_memory_mode, SwapPriority, DEFAULT_PER_CPU_BUFFER_SIZE,
     INTERMEDIATE_COMPRESSION_LEVEL_FAST, INTERMEDIATE_COMPRESSION_LEVEL_SLOW, KEEP_FILES,
-    MAXIMUM_SECOND_BUCKETS_LOG, MINIMUM_LOG_DELTA_TIME,
+    MAXIMUM_SECOND_BUCKETS_LOG, MINIMUM_LOG_DELTA_TIME, PRIORITY_SCHEDULING_BASE,
 };
 use hashes::HashFunctionFactory;
 use io::concurrent::structured_sequences::binary::StructSeqBinaryWriter;
@@ -33,6 +33,7 @@ use parallel_processor::buckets::MultiThreadBuckets;
 use parallel_processor::memory_data_size::MemoryDataSize;
 use parallel_processor::memory_fs::{MemoryFs, RemoveFileMode};
 use parallel_processor::phase_times_monitor::PHASES_TIMES_MONITOR;
+use parallel_processor::scheduler::PriorityScheduler;
 use parallel_processor::utils::scoped_thread_local::ScopedThreadLocal;
 use pipeline::eulertigs::build_eulertigs;
 use std::fs::remove_file;
@@ -187,13 +188,15 @@ pub fn run_assembler<
         buckets.par_iter().enumerate().for_each(|(index, bucket)| {
             ggcat_logging::info!("Stats for bucket index: {}", index);
             for chunk in &bucket.chunks {
+                let thread_handle = PriorityScheduler::declare_thread(PRIORITY_SCHEDULING_BASE);
                 kmers_transform::debug_bucket_stats::compute_stats_for_bucket::<MergingHash>(
                     chunk.clone(),
                     index,
                     buckets.len(),
                     MAXIMUM_SECOND_BUCKETS_LOG,
                     k,
                     m,
+                    &thread_handle,
                 );
             }
         });

crates/cmdline/Cargo.toml

@@ -65,7 +65,7 @@ ggcat-logging = { version = "1.1.1", path = "../logging" }
 [features]
 mem-analysis = ["parallel-processor/track-usage"]
 no-stats = ["parallel-processor/no-stats"]
-process-stats = ["parallel-processor/process-stats"]
+process-stats = ["parallel-processor/process-stats", "ggcat-logging/stats"]
 tracing = ["instrumenter/enabled"]
 devel-build = ["assembler/devel-build", "querier/devel-build"]
 kmer-counters = ["assembler/support_kmer_counters"]

crates/config/src/lib.rs

@@ -76,6 +76,10 @@ pub const QUERIES_COUNT_MIN_BATCH: u64 = 1000;
 pub const DEFAULT_COMPACTION_MAP_SUBBUCKET_ELEMENTS: usize = 512;
 pub const MAX_COMPACTION_MAP_SUBBUCKET_ELEMENTS: usize = 1024 * 4;
 
+pub const PRIORITY_SCHEDULING_HIGH: usize = 0;
+pub const PRIORITY_SCHEDULING_BASE: usize = 1;
+pub const PRIORITY_SCHEDULING_LOW: usize = 2;
+
 pub struct SwapPriority {}
 #[allow(non_upper_case_globals)]
 impl SwapPriority {

crates/io/src/concurrent/temp_reads/creads_utils.rs

@@ -276,7 +276,8 @@ pub mod helpers {
                 $allowed_passtrough:expr,
                 |$passtrough_info:ident| $p:expr,
                 |$checkpoint_data:ident| $c:expr,
-                |$data: ident, $extra_buffer: ident| $f:expr
+                |$data: ident, $extra_buffer: ident| $f:expr,
+                $thread_handle:ident
             );
 
         ) => {
@@ -297,7 +298,7 @@ pub mod helpers {
                         $FlagsCount,
                         $BucketMode,
                         WithMultiplicity,
-                    >>(read_thread, Vec::new(), <$E>::new_temp_buffer(), $allowed_passtrough);
+                    >>(read_thread, Vec::new(), <$E>::new_temp_buffer(), $allowed_passtrough, &$thread_handle);
                 while let Some(checkpoint) = items.get_next_checkpoint_extended() {
                     match checkpoint {
                         AsyncBinaryReaderIteratorData::Stream(items, $checkpoint_data) => {
@@ -324,7 +325,7 @@ pub mod helpers {
                         $FlagsCount,
                         $BucketMode,
                         NoMultiplicity,
-                    >>(read_thread, Vec::new(), <$E>::new_temp_buffer(), $allowed_passtrough);
+                    >>(read_thread, Vec::new(), <$E>::new_temp_buffer(), $allowed_passtrough, &$thread_handle);
                 while let Some(checkpoint) = items.get_next_checkpoint_extended() {
                     match checkpoint {
                         AsyncBinaryReaderIteratorData::Stream(items, $checkpoint_data) => {

crates/kmers_transform/src/debug_bucket_stats.rs

@@ -15,6 +15,7 @@ use parallel_processor::buckets::readers::async_binary_reader::{
     AllowedCheckpointStrategy, AsyncBinaryReader, AsyncReaderThread,
 };
 use parallel_processor::memory_fs::RemoveFileMode;
+use parallel_processor::scheduler::ThreadPriorityHandle;
 use std::collections::HashSet;
 use std::path::PathBuf;
 
@@ -50,6 +51,7 @@ pub fn compute_stats_for_bucket<MH: HashFunctionFactory>(
     second_buckets_log_max: usize,
     k: usize,
     m: usize,
+    thread_handle: &ThreadPriorityHandle,
 ) {
     let reader = AsyncBinaryReader::new(
         &bucket,
@@ -78,6 +80,7 @@ pub fn compute_stats_for_bucket<MH: HashFunctionFactory>(
         Vec::new(),
         (),
         AllowedCheckpointStrategy::DecompressOnly,
+        thread_handle,
     );
 
     let mut total_counters = vec![0; second_buckets_max];

crates/kmers_transform/src/processor.rs

@@ -3,13 +3,14 @@ use crate::{
     KmersTransformContext, KmersTransformExecutorFactory, KmersTransformFinalExecutor,
     KmersTransformMapProcessor,
 };
-use config::WORKERS_PRIORITY_BASE;
+use config::{PRIORITY_SCHEDULING_HIGH, WORKERS_PRIORITY_BASE};
 use parallel_processor::execution_manager::executor::{AsyncExecutor, ExecutorReceiver};
 use parallel_processor::execution_manager::memory_tracker::MemoryTracker;
 use parallel_processor::execution_manager::objects_pool::PoolObjectTrait;
 use parallel_processor::execution_manager::packet::{Packet, PacketTrait};
 use parallel_processor::mt_debug_counters::counter::{AtomicCounter, SumMode};
 use parallel_processor::mt_debug_counters::declare_counter_i64;
+use parallel_processor::scheduler::PriorityScheduler;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::path::PathBuf;
@@ -57,9 +58,11 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformProcessor
                 <F::MapProcessorType as KmersTransformMapProcessor<F>>::MapStruct::allocate_new(&()),
             );
 
+            let thread_handle = PriorityScheduler::declare_thread(PRIORITY_SCHEDULING_HIGH);
+
             while let Ok((address, proc_info)) = track!(
                 receiver
-                    .obtain_address_with_priority(WORKERS_PRIORITY_BASE)
+                    .obtain_address_with_priority(WORKERS_PRIORITY_BASE, &thread_handle)
                     .await,
                 ADDR_WAITING_COUNTER
             ) {
@@ -69,9 +72,10 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformProcessor
                 let mut total_kmers = 0;
                 let mut unique_kmers = 0;
 
-                while let Some(input_packet) =
-                    track!(address.receive_packet().await, PACKET_WAITING_COUNTER)
-                {
+                while let Some(input_packet) = track!(
+                    address.receive_packet(&thread_handle).await,
+                    PACKET_WAITING_COUNTER
+                ) {
                     real_size += input_packet.reads.len() as usize;
                     let stats = map_processor.process_group_batch_sequences(
                         &global_context.global_extra_data,

crates/kmers_transform/src/reader.rs

@@ -10,7 +10,7 @@ use config::{
     DEFAULT_PER_CPU_BUFFER_SIZE, DEFAULT_PREFETCH_AMOUNT, KEEP_FILES,
     MAXIMUM_JIT_PROCESSED_BUCKETS, MAX_INTERMEDIATE_MAP_SIZE, MIN_BUCKET_CHUNKS_FOR_READING_THREAD,
     PACKETS_PRIORITY_DEFAULT, PACKETS_PRIORITY_REWRITTEN, PARTIAL_VECS_CHECKPOINT_SIZE,
-    USE_SECOND_BUCKET, WORKERS_PRIORITY_BASE,
+    PRIORITY_SCHEDULING_BASE, PRIORITY_SCHEDULING_LOW, USE_SECOND_BUCKET, WORKERS_PRIORITY_BASE,
 };
 use instrumenter::local_setup_instrumenter;
 use io::compressed_read::CompressedReadIndipendent;
@@ -39,6 +39,7 @@ use parallel_processor::execution_manager::packet::{Packet, PacketTrait, Packets
 use parallel_processor::memory_fs::RemoveFileMode;
 use parallel_processor::mt_debug_counters::counter::{AtomicCounter, SumMode};
 use parallel_processor::mt_debug_counters::declare_counter_i64;
+use parallel_processor::scheduler::{PriorityScheduler, ThreadPriorityHandle};
 use parallel_processor::utils::replace_with_async::replace_with_async;
 use std::cmp::{max, min, Reverse};
 use std::collections::{BinaryHeap, VecDeque};
@@ -58,6 +59,7 @@ pub(crate) struct KmersTransformReader<F: KmersTransformExecutorFactory> {
 pub struct InputBucketDesc {
     pub(crate) paths: Vec<PathBuf>,
     pub(crate) sub_bucket_counters: Vec<BucketCounter>,
+    #[allow(dead_code)]
     pub(crate) compaction_delta: i64,
     pub(crate) out_data_format: MinimizerBucketMode,
     pub(crate) resplitted: bool,
@@ -168,11 +170,6 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
             })
             .collect();
 
-        let compaction_ratio = sequences_count
-            .saturating_add_signed(file.compaction_delta)
-            .max(16) as f64
-            / sequences_count as f64;
-
         let total_file_size = readers.iter().map(|r| r.get_file_size()).sum();
 
         bucket_sizes.make_contiguous().sort();
@@ -202,7 +199,7 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
 
             let is_outlier = !file.resplitted
                 && (total_sequences > 0)
-                && (biggest_sub_bucket.0.count as f64 * unique_estimator_factor * compaction_ratio
+                && (biggest_sub_bucket.0.count as f64 * unique_estimator_factor
                     >= (MAX_INTERMEDIATE_MAP_SIZE / F::MapProcessorType::MAP_SIZE as u64) as f64);
 
             // if is_outlier {
@@ -398,6 +395,7 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
         bucket_info: &BucketsInfo,
         async_reader_thread: Arc<AsyncReaderThread>,
         packets_pool: Arc<PoolObject<PacketsPool<ReadsBuffer<F::AssociatedExtraData>>>>,
+        thread_handle: &ThreadPriorityHandle,
     ) {
         if bucket_info.readers.iter().all(|r| r.is_finished()) {
             return;
@@ -494,9 +492,9 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
                         if buffers[bucket].reads.len() == buffers[bucket].reads.capacity() {
                             match &bucket_info.addresses[bucket] {
                                 AddressMode::Send(address) => {
-                                    replace_with_async(&mut buffers[bucket], |mut buffer| async move {
+                                    replace_with_async(&mut buffers[bucket], |mut buffer| async {
                                         buffer.sub_bucket = bucket;
-                                        ops.packet_send(address.clone(), buffer);
+                                        ops.packet_send(address.clone(), buffer, thread_handle);
                                         track!(packets_pool.alloc_packet().await, PACKET_ALLOC_COUNTER)
                                     })
                                     .await;
@@ -512,7 +510,8 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
                             }
                         }
                         F::AssociatedExtraData::clear_temp_buffer(extra_buffer);
-                    }
+                    },
+                    thread_handle
                 );
             }
         }
@@ -526,7 +525,7 @@ impl<F: KmersTransformExecutorFactory> KmersTransformReader<F> {
                 packet.sub_bucket = bucket;
                 match address {
                     AddressMode::Send(address) => {
-                        ops.packet_send(address.clone(), packet);
+                        ops.packet_send(address.clone(), packet, &thread_handle);
                     }
                     AddressMode::Rewrite(writer, seq_count, _) => {
                         Self::flush_rewrite_bucket::<MultiplicityModeFromBoolean<WITH_MULTIPLICITY>>(
@@ -563,14 +562,16 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformReader<F>
         async move {
             let mut async_threads = Vec::new();
 
+            let thread_handle = PriorityScheduler::declare_thread(PRIORITY_SCHEDULING_BASE);
+
             while let Ok((address, _)) = track!(
                 receiver
-                    .obtain_address_with_priority(WORKERS_PRIORITY_BASE)
+                    .obtain_address_with_priority(WORKERS_PRIORITY_BASE, &thread_handle)
                     .await,
                 ADDR_WAITING_COUNTER
             ) {
                 let file = track!(
-                    address.receive_packet().await.unwrap(),
+                    address.receive_packet(&thread_handle).await.unwrap(),
                     PACKET_WAITING_COUNTER
                 );
                 let is_main_bucket = !file.resplitted && !file.rewritten;
@@ -604,13 +605,18 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformReader<F>
                     let address = &address;
                     let buckets_info = &buckets_info;
                     let packets_pool = address
-                        .pool_alloc_await(max(
-                            global_context.max_buckets / 2,
-                            2 * buckets_info.addresses.len(),
-                        ))
+                        .pool_alloc_await(
+                            max(
+                                global_context.max_buckets / 2,
+                                2 * buckets_info.addresses.len(),
+                            ),
+                            &thread_handle,
+                        )
                         .await;
 
                     spawner.spawn_executor(async move {
+                        let thread_handle =
+                            PriorityScheduler::declare_thread(PRIORITY_SCHEDULING_LOW);
                         match buckets_info.data_format {
                             MinimizerBucketMode::Single => {
                                 Self::read_bucket::<false>(
@@ -619,6 +625,7 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformReader<F>
                                     buckets_info,
                                     async_thread,
                                     packets_pool,
+                                    &thread_handle,
                                 )
                                 .await;
                             }
@@ -629,6 +636,7 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformReader<F>
                                     buckets_info,
                                     async_thread,
                                     packets_pool,
+                                    &thread_handle,
                                 )
                                 .await;
                             }
@@ -695,7 +703,7 @@ impl<F: KmersTransformExecutorFactory> AsyncExecutor for KmersTransformReader<F>
                 }
 
                 assert!(track!(
-                    address.receive_packet().await.is_none(),
+                    address.receive_packet(&thread_handle).await.is_none(),
                     PACKET_WAITING_COUNTER
                 ));
             }