liboxen 0.51.2

Oxen is a fast, unstructured data version control, to help version large machine learning datasets written in Rust.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

use crate::constants;

/// Returns the number of threads to use for a given number of items
/// Can be overridden by setting the environment variable OXEN_NUM_THREADS
/// Defaults to constants::DEFAULT_NUM_WORKERS or the number of CPUs we have if we have less than that
pub fn num_threads_for_items(num_items: usize) -> usize {
    let num_workers = parse_num_threads_env();

    // Cap the worker count at the item count, but never return 0: callers feed this straight into
    // `buffer_unordered` / `for_each_concurrent`, where a concurrency limit of 0 stalls the stream
    // forever instead of completing on an empty input.
    num_workers.min(num_items).max(1)
}

/// Returns the default number of threads to use
/// Can be overridden by setting the environment variable OXEN_NUM_THREADS
/// Defaults to constants::DEFAULT_NUM_WORKERS or the number of CPUs we have if we have less than that
pub fn default_num_threads() -> usize {
    parse_num_threads_env()
}

fn parse_num_threads_env() -> usize {
    if let Ok(num_threads) = std::env::var("OXEN_NUM_THREADS") {
        if let Ok(num_threads) = num_threads.parse::<usize>() {
            num_threads
        } else {
            get_default_num_workers()
        }
    } else {
        get_default_num_workers()
    }
}

fn get_default_num_workers() -> usize {
    // Check how many CPUs we have
    let num_cpus = num_cpus::get();

    // Default to constants::DEFAULT_NUM_WORKERS, but if we have less cpus than that, use that instead
    if constants::DEFAULT_NUM_WORKERS > num_cpus {
        num_cpus
    } else {
        constants::DEFAULT_NUM_WORKERS
    }
}

#[cfg(test)]
mod tests {
    use super::num_threads_for_items;

    #[test]
    fn test_num_threads_for_items_never_zero() {
        // A worker count of 0 fed into `buffer_unordered` / `for_each_concurrent` stalls the
        // stream forever, so an empty work set must still yield at least one worker.
        assert_eq!(num_threads_for_items(0), 1);
    }
}