use serde::{Deserialize, Serialize};
use crate::{
config::{ProviderPreference, RuntimeConfig},
error::{RapidOcrError, Result},
};
pub(crate) fn auto_tuned_thread_budget() -> usize {
let physical_cores = num_cpus::get_physical().max(1);
let available = std::thread::available_parallelism()
.ok()
.map(|value| value.get())
.unwrap_or(1);
available.clamp(1, physical_cores)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ThreadSource {
Auto,
Explicit,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ThreadPlan {
pub source: ThreadSource,
pub budget: usize,
pub ort_intra: Option<usize>,
pub ort_inter: Option<usize>,
pub rayon: Option<usize>,
pub sessions: usize,
}
#[derive(Debug, Clone, PartialEq)]
pub struct RuntimeProfile {
pub provider: ProviderPreference,
pub cpu_mem_arena: bool,
pub fail_if_provider_unavailable: bool,
pub formula_batch: usize,
pub threads: ThreadPlan,
}
impl RuntimeProfile {
pub fn resolve(runtime: &RuntimeConfig, classifier_enabled: bool) -> Result<Self> {
let mut profile = Self::plan(runtime, classifier_enabled);
profile.apply_rayon_global_pool()?;
Ok(profile)
}
pub fn plan(runtime: &RuntimeConfig, classifier_enabled: bool) -> Self {
let budget = auto_tuned_thread_budget().max(1);
let explicit_rayon = runtime.rayon_threads.filter(|value| *value > 0);
let (ort_intra, ort_inter) = runtime.effective_session_threads();
let source = if runtime.has_explicit_thread_request() {
ThreadSource::Explicit
} else {
ThreadSource::Auto
};
let auto = runtime.auto_tune_threads;
let rayon = explicit_rayon.or_else(|| auto.then(|| (budget / 4).clamp(1, 8)));
Self {
provider: runtime.provider_preference,
cpu_mem_arena: runtime.enable_cpu_mem_arena,
fail_if_provider_unavailable: runtime.fail_if_provider_unavailable,
formula_batch: runtime.formula_batch,
threads: ThreadPlan {
source,
budget,
ort_intra,
ort_inter,
rayon,
sessions: if classifier_enabled { 3 } else { 2 },
},
}
}
pub fn apply_rayon_global_pool(&mut self) -> Result<()> {
let Some(requested) = self.threads.rayon else {
return Ok(());
};
if rayon::ThreadPoolBuilder::new()
.num_threads(requested)
.build_global()
.is_ok()
{
return Ok(());
}
if let Some(actual) = reconcile_existing_rayon_pool(
self.threads.source,
requested,
rayon::current_num_threads(),
)? {
self.threads.rayon = Some(actual);
}
Ok(())
}
pub fn session_runtime(&self) -> RuntimeConfig {
RuntimeConfig {
intra_threads: self.threads.ort_intra,
inter_threads: self.threads.ort_inter,
auto_tune_threads: false,
rayon_threads: self.threads.rayon,
enable_cpu_mem_arena: self.cpu_mem_arena,
fail_if_provider_unavailable: self.fail_if_provider_unavailable,
provider_preference: self.provider,
formula_batch: self.formula_batch,
}
}
}
fn reconcile_existing_rayon_pool(
source: ThreadSource,
requested: usize,
actual: usize,
) -> Result<Option<usize>> {
if actual == requested {
return Ok(None);
}
if source == ThreadSource::Explicit {
return Err(RapidOcrError::Config(format!(
"the Rayon global thread pool is already initialised with {actual} threads, \
so the requested rayon_threads = {requested} cannot be applied; Rayon allows \
only one global pool per process"
)));
}
Ok(Some(actual))
}
#[cfg(test)]
mod tests {
use super::{RuntimeProfile, ThreadPlan, ThreadSource, auto_tuned_thread_budget};
use crate::{
config::{ProviderPreference, RuntimeConfig},
error::RapidOcrError,
};
#[test]
fn explicit_values_win_and_are_marked_explicit() {
let runtime = RuntimeConfig {
intra_threads: Some(6),
inter_threads: Some(2),
auto_tune_threads: false,
rayon_threads: Some(3),
enable_cpu_mem_arena: false,
fail_if_provider_unavailable: true,
provider_preference: ProviderPreference::Cpu,
formula_batch: 5,
};
let profile = RuntimeProfile::plan(&runtime, false);
assert_eq!(profile.threads.source, ThreadSource::Explicit);
assert_eq!(profile.threads.ort_intra, Some(6));
assert_eq!(profile.threads.ort_inter, Some(2));
assert_eq!(profile.threads.rayon, Some(3));
assert_eq!(profile.threads.sessions, 2);
assert_eq!(profile.formula_batch, 5);
assert!(!profile.cpu_mem_arena);
assert!(profile.fail_if_provider_unavailable);
assert_eq!(
RuntimeProfile::plan(&runtime, true).threads.sessions,
3,
"enabling the classifier adds its session to the shared budget"
);
}
#[test]
fn default_config_derives_the_same_session_threads_in_both_paths() {
let budget = auto_tuned_thread_budget();
let runtime = RuntimeConfig::default();
assert!(runtime.auto_tune_threads, "the default must auto-tune");
assert_eq!(
runtime.effective_session_threads(),
(Some(budget), Some(1)),
"a default config derives intra = budget and inter = 1"
);
let plan = RuntimeProfile::plan(&runtime, false);
assert_eq!(plan.threads.budget, budget.max(1));
assert_eq!(
(plan.threads.ort_intra, plan.threads.ort_inter),
runtime.effective_session_threads(),
"the plan must not re-derive the ORT thread counts on its own"
);
let session = plan.session_runtime();
assert_eq!(session.intra_threads, Some(budget));
assert_eq!(session.inter_threads, Some(1));
}
#[test]
fn auto_tune_disabled_derives_no_session_threads() {
let runtime = RuntimeConfig {
auto_tune_threads: false,
..RuntimeConfig::default()
};
assert_eq!(runtime.effective_session_threads(), (None, None));
let plan = RuntimeProfile::plan(&runtime, false);
assert_eq!(
(plan.threads.ort_intra, plan.threads.ort_inter),
runtime.effective_session_threads()
);
assert_eq!(
(plan.threads.ort_intra, plan.threads.ort_inter),
(None, None)
);
}
#[test]
fn explicit_threads_win_in_both_auto_tune_modes() {
let both_explicit = RuntimeConfig {
intra_threads: Some(6),
inter_threads: Some(2),
auto_tune_threads: false,
..RuntimeConfig::default()
};
assert_eq!(
both_explicit.effective_session_threads(),
(Some(6), Some(2))
);
let plan = RuntimeProfile::plan(&both_explicit, false);
assert_eq!(
(plan.threads.ort_intra, plan.threads.ort_inter),
(Some(6), Some(2))
);
assert_eq!(plan.threads.source, ThreadSource::Explicit);
let mixed_on = RuntimeConfig {
intra_threads: Some(6),
inter_threads: None,
auto_tune_threads: true,
..RuntimeConfig::default()
};
assert_eq!(mixed_on.effective_session_threads(), (Some(6), Some(1)));
let plan = RuntimeProfile::plan(&mixed_on, false);
assert_eq!(
(plan.threads.ort_intra, plan.threads.ort_inter),
(Some(6), Some(1))
);
let mixed_off = RuntimeConfig {
intra_threads: None,
inter_threads: Some(2),
auto_tune_threads: false,
..RuntimeConfig::default()
};
assert_eq!(mixed_off.effective_session_threads(), (None, Some(2)));
let plan = RuntimeProfile::plan(&mixed_off, false);
assert_eq!(
(plan.threads.ort_intra, plan.threads.ort_inter),
(None, Some(2))
);
}
#[test]
fn auto_values_follow_the_documented_formula() {
let budget = auto_tuned_thread_budget().max(1);
let profile = RuntimeProfile::plan(&RuntimeConfig::default(), false);
assert_eq!(profile.threads.source, ThreadSource::Auto);
assert_eq!(profile.threads.budget, budget);
assert_eq!(profile.threads.ort_intra, Some(budget));
assert_eq!(profile.threads.ort_inter, Some(1));
assert_eq!(profile.threads.rayon, Some((budget / 4).clamp(1, 8)));
assert!(
(1..=8).contains(&profile.threads.rayon.expect("auto plan sets rayon")),
"the auto Rayon share must stay small and bounded"
);
}
#[test]
fn auto_tune_disabled_configures_nothing() {
let runtime = RuntimeConfig {
auto_tune_threads: false,
..RuntimeConfig::default()
};
let profile = RuntimeProfile::plan(&runtime, false);
assert_eq!(profile.threads.source, ThreadSource::Auto);
assert_eq!(profile.threads.ort_intra, None);
assert_eq!(profile.threads.ort_inter, None);
assert_eq!(profile.threads.rayon, None);
let session = profile.session_runtime();
assert_eq!(session.intra_threads, None);
assert_eq!(session.inter_threads, None);
assert_eq!(session.rayon_threads, None);
assert!(!session.auto_tune_threads);
}
#[test]
fn auto_tune_disabled_keeps_explicit_intra_only() {
let runtime = RuntimeConfig {
intra_threads: Some(6),
auto_tune_threads: false,
..RuntimeConfig::default()
};
let profile = RuntimeProfile::plan(&runtime, false);
assert_eq!(profile.threads.source, ThreadSource::Explicit);
assert_eq!(profile.threads.ort_intra, Some(6));
assert_eq!(profile.threads.ort_inter, None);
assert_eq!(profile.threads.rayon, None);
let session = profile.session_runtime();
assert_eq!(session.intra_threads, Some(6));
assert_eq!(session.inter_threads, None);
assert_eq!(session.rayon_threads, None);
}
#[test]
fn unconfigured_rayon_is_left_alone_even_with_an_existing_pool() {
let _ = rayon::ThreadPoolBuilder::new().build_global();
let before = rayon::current_num_threads();
let runtime = RuntimeConfig {
auto_tune_threads: false,
..RuntimeConfig::default()
};
let profile = RuntimeProfile::resolve(&runtime, false)
.expect("not configuring Rayon must never fail, even if a pool already exists");
assert_eq!(profile.threads.rayon, None);
assert_eq!(
rayon::current_num_threads(),
before,
"an unconfigured plan must not resize the live pool"
);
}
#[test]
fn zero_is_treated_as_unset() {
let runtime = RuntimeConfig {
intra_threads: Some(0),
inter_threads: Some(0),
rayon_threads: Some(0),
auto_tune_threads: true,
..RuntimeConfig::default()
};
let profile = RuntimeProfile::plan(&runtime, false);
assert_eq!(profile.threads.source, ThreadSource::Auto);
assert!(profile.threads.ort_intra.is_some_and(|value| value >= 1));
assert_eq!(profile.threads.ort_inter, Some(1));
assert!(profile.threads.rayon.is_some_and(|value| value >= 1));
}
#[test]
fn session_runtime_reflects_the_plan() {
let runtime = RuntimeConfig {
intra_threads: Some(5),
inter_threads: None,
auto_tune_threads: true,
rayon_threads: None,
enable_cpu_mem_arena: false,
fail_if_provider_unavailable: true,
provider_preference: ProviderPreference::Cpu,
formula_batch: 9,
};
let profile = RuntimeProfile::plan(&runtime, false);
let session = profile.session_runtime();
assert_eq!(session.intra_threads, Some(5));
assert_eq!(session.inter_threads, Some(1));
assert_eq!(session.rayon_threads, profile.threads.rayon);
assert!(
!session.auto_tune_threads,
"ORT must not re-derive threads after the plan was resolved"
);
assert!(!session.enable_cpu_mem_arena);
assert!(session.fail_if_provider_unavailable);
assert_eq!(session.provider_preference, ProviderPreference::Cpu);
assert_eq!(session.formula_batch, 9);
}
#[test]
fn thread_plan_serialises_for_reports() {
let plan = ThreadPlan {
source: ThreadSource::Auto,
budget: 14,
ort_intra: Some(14),
ort_inter: Some(1),
rayon: Some(3),
sessions: 2,
};
let value = serde_json::to_value(&plan).expect("thread plan must serialise");
assert_eq!(value["source"], serde_json::json!("auto"));
assert_eq!(value["budget"], serde_json::json!(14));
assert_eq!(value["rayon"], serde_json::json!(3));
let round_trip: ThreadPlan =
serde_json::from_value(value).expect("thread plan must deserialise");
assert_eq!(round_trip, plan);
let unconfigured = ThreadPlan {
source: ThreadSource::Auto,
budget: 14,
ort_intra: None,
ort_inter: None,
rayon: None,
sessions: 2,
};
let value = serde_json::to_value(&unconfigured).expect("thread plan must serialise");
assert_eq!(value["ort_intra"], serde_json::Value::Null);
assert_eq!(value["ort_inter"], serde_json::Value::Null);
assert_eq!(value["rayon"], serde_json::Value::Null);
let explicit = serde_json::to_value(ThreadSource::Explicit).expect("source serialises");
assert_eq!(explicit, serde_json::json!("explicit"));
}
#[test]
fn explicit_rayon_request_conflicting_with_an_existing_pool_is_reported() {
let _ = rayon::ThreadPoolBuilder::new().build_global();
let actual = rayon::current_num_threads();
assert!(actual > 0, "rayon must report a live global pool");
let requested = actual + 1;
let runtime = RuntimeConfig {
intra_threads: Some(requested),
inter_threads: Some(1),
auto_tune_threads: false,
rayon_threads: Some(requested),
..RuntimeConfig::default()
};
let error = RuntimeProfile::resolve(&runtime, false)
.expect_err("an explicit Rayon request that cannot be applied must fail");
assert!(
matches!(error, RapidOcrError::Config(_)),
"unexpected error: {error}"
);
let message = error.to_string();
assert!(
message.contains(&actual.to_string()) && message.contains(&requested.to_string()),
"the error must name both the actual and the requested thread counts: {message}"
);
}
#[test]
fn auto_plan_adopts_and_reports_an_existing_pool() {
let _ = rayon::ThreadPoolBuilder::new().build_global();
let profile = RuntimeProfile::resolve(&RuntimeConfig::default(), false)
.expect("the auto plan must tolerate an existing global pool");
assert_eq!(
profile.threads.rayon,
Some(rayon::current_num_threads()),
"the reported Rayon value must be the effective one"
);
assert_eq!(
profile.session_runtime().rayon_threads,
profile.threads.rayon
);
}
#[test]
fn explicit_request_matching_the_existing_pool_is_accepted() {
let _ = rayon::ThreadPoolBuilder::new().build_global();
let actual = rayon::current_num_threads();
let runtime = RuntimeConfig {
intra_threads: Some(actual),
inter_threads: Some(1),
auto_tune_threads: false,
rayon_threads: Some(actual),
..RuntimeConfig::default()
};
let profile = RuntimeProfile::resolve(&runtime, false)
.expect("an explicit request equal to the live pool size must be accepted");
assert_eq!(profile.threads.source, ThreadSource::Explicit);
assert_eq!(profile.threads.rayon, Some(actual));
assert_eq!(profile.session_runtime().rayon_threads, Some(actual));
}
#[test]
fn rayon_only_explicit_request_conflicts_with_an_existing_pool() {
let runtime = RuntimeConfig {
intra_threads: None,
inter_threads: None,
auto_tune_threads: false,
rayon_threads: Some(4),
..RuntimeConfig::default()
};
assert!(runtime.has_explicit_thread_request());
let profile = RuntimeProfile::plan(&runtime, false);
assert_eq!(
profile.threads.source,
ThreadSource::Explicit,
"an explicit `rayon_threads` request must be marked explicit even when intra/inter \
are unset"
);
assert_eq!(profile.threads.ort_intra, None);
assert_eq!(profile.threads.ort_inter, None);
let error = super::reconcile_existing_rayon_pool(ThreadSource::Explicit, 4, 7)
.expect_err("an explicit Rayon request that cannot be applied must fail");
assert!(
matches!(error, RapidOcrError::Config(_)),
"unexpected error: {error}"
);
let message = error.to_string();
assert!(
message.contains("with 7 threads"),
"the error must name the actual pool size: {message}"
);
assert!(
message.contains("rayon_threads = 4"),
"the error must name the requested size: {message}"
);
assert_eq!(
super::reconcile_existing_rayon_pool(ThreadSource::Auto, 4, 7)
.expect("the auto path must still tolerate an existing pool"),
Some(7),
"the auto path adopts the existing pool and reports the effective value"
);
assert_eq!(
super::reconcile_existing_rayon_pool(ThreadSource::Explicit, 7, 7)
.expect("an explicit request equal to the live pool size is not a conflict"),
None
);
let _ = rayon::ThreadPoolBuilder::new().build_global();
let actual = rayon::current_num_threads();
let requested = actual + 1;
let runtime = RuntimeConfig {
rayon_threads: Some(requested),
..runtime
};
let error = RuntimeProfile::resolve(&runtime, false)
.expect_err("a rayon-only explicit request that cannot be applied must fail");
assert!(
matches!(error, RapidOcrError::Config(_)),
"unexpected error: {error}"
);
let message = error.to_string();
assert!(
message.contains(&format!("with {actual} threads"))
&& message.contains(&format!("rayon_threads = {requested}")),
"the error must name both the actual ({actual}) and the requested ({requested}) \
thread counts: {message}"
);
}
}