ecp_core/
utils.rs

1//! The default memory limit. Also re-exports the `metric` and `dtype` items.
2
3pub use crate::dtype::{EmbeddingDtype, dtype_of_array, read_subset_as_f32};
4pub use crate::metric::{Metric, calculate_distances, negative_squared_distances};
5
6/// Share of system RAM used as the memory limit when the caller doesn't set one.
7const DEFAULT_MEMORY_LIMIT_RAM_FRACTION: f64 = 0.8;
8
9/// Returns the default memory limit in bytes, 80% of RAM rounded down to a
10/// whole GiB. Inside a Linux container it uses the container's memory cap,
11/// since the host's RAM would be too high.
12pub fn default_memory_limit_bytes() -> usize {
13    let system = sysinfo::System::new_with_specifics(
14        sysinfo::RefreshKind::nothing()
15            .with_memory(sysinfo::MemoryRefreshKind::nothing().with_ram()),
16    );
17    let total_ram_bytes = match system.cgroup_limits() {
18        Some(limits) => limits.total_memory,
19        None => {
20            // No cgroup limits is the normal, every-call result outside
21            // Linux, not a failure worth warning about there.
22            if cfg!(target_os = "linux") {
23                log::warn!(
24                    "failed to read this Linux host's cgroup memory limit, \
25                     falling back to its total RAM"
26                );
27            }
28            system.total_memory()
29        }
30    };
31    default_memory_limit_bytes_for(total_ram_bytes)
32}
33
34/// Computes the default memory limit for `total_ram_bytes` of RAM.
35/// Note: This function is only split out from `default_memory_limit_bytes` for unit tests.
36fn default_memory_limit_bytes_for(total_ram_bytes: u64) -> usize {
37    const GIB: f64 = (1024 * 1024 * 1024) as f64;
38    let total_gib = total_ram_bytes as f64 / GIB;
39    let default_gib = (total_gib * DEFAULT_MEMORY_LIMIT_RAM_FRACTION).floor();
40    (default_gib * GIB) as usize
41}
42
43#[cfg(test)]
44#[path = "utests/utils.rs"]
45mod tests;
ecp_core/utils.rs

ecp_core/
utils.rs
