bge_m3_embedding_server/

sysinfo.rs

// Copyright (c) 2026 J. Patrick Fulton
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Memory detection for auto-budget computation.
//!
//! Production target is Linux (Fargate/ECS). On Linux we walk the cgroup
//! hierarchy to find the container memory limit, then fall back to host RAM
//! reported by `/proc/meminfo`. On macOS we read host RAM via `sysctl`;
//! cgroup support requires unsafe FFI so it is deferred.
//!
//! ## cgroup-v2 detection on ECS Managed Instances (Bottlerocket)
//!
//! ECS Managed Instances launch containers **without** `--cgroupns=private`,
//! so `/sys/fs/cgroup/memory.max` resolves to the unified-hierarchy root,
//! which reads `"max"` (no limit). The actual container memory limit is
//! set at a deeper path whose last component is recorded in
//! `/proc/self/cgroup` (unified-hierarchy format: a single line
//! `0::<path>`, e.g. `0::/ecs.slice/ecs-…-task.scope/<id>`).
//!
//! `cgroup_memory()` reads `/proc/self/cgroup`, extracts that path, then
//! reads `memory.max` at each ancestor (deepest first) until it finds a
//! numeric limit or exhausts the tree. Falls through to `host_ram` only
//! when the entire walk yields `"max"` (truly unconstrained host).
//!
//! RSS tracking (`read_process_rss_bytes`) is Linux-only (parses
//! `/proc/self/statm`). On macOS it returns `None`; the auto-budget logic
//! treats `None` as "cannot measure model footprint" and uses conservative
//! defaults.
use tracing::warn;

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

/// Where the available-memory reading came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[allow(dead_code)] // CgroupV2/CgroupV1 are constructed only on Linux; macOS sees them as unused.
pub enum MemorySource {
    /// `BGE_M3_AVAILABLE_MEMORY_BYTES` env override.
    Override,
    /// Linux cgroup v2 `memory.max`.
    CgroupV2,
    /// Linux cgroup v1 `memory.limit_in_bytes`.
    CgroupV1,
    /// `/proc/meminfo` `MemAvailable` (Linux) or `sysctl hw.memsize` (macOS).
    HostRam,
}

impl std::fmt::Display for MemorySource {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Override => f.write_str("override"),
            Self::CgroupV2 => f.write_str("cgroup_v2"),
            Self::CgroupV1 => f.write_str("cgroup_v1"),
            Self::HostRam => f.write_str("host_ram"),
        }
    }
}

/// A memory reading with its provenance.
#[derive(Debug, Clone, Copy)]
pub struct MemoryReading {
    /// Total available memory bytes detected from the source.
    pub available_bytes: usize,
    /// Detection method that produced this reading.
    pub source: MemorySource,
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------

/// Detects available memory for the process.
///
/// Detection chain (first success wins):
/// 1. `BGE_M3_AVAILABLE_MEMORY_BYTES` env override.
/// 2. Linux cgroup v2: `/sys/fs/cgroup/memory.max`.
/// 3. Linux cgroup v1: `/sys/fs/cgroup/memory/memory.limit_in_bytes`.
/// 4. Linux: `/proc/meminfo` `MemAvailable`.
/// 5. macOS: `sysctl hw.memsize` (total host RAM; no cgroup support).
/// 6. Fallback: 4 GiB constant with a warning log.
pub(crate) fn detect_available_memory() -> MemoryReading {
    // --- step 1: explicit override ---
    if let Some(bytes) = env_override() {
        return MemoryReading {
            available_bytes: bytes,
            source: MemorySource::Override,
        };
    }

    // --- step 2 / 3: Linux cgroup ---
    #[cfg(target_os = "linux")]
    if let Some(r) = cgroup_memory() {
        return r;
    }

    // --- step 4 / 5: OS-level RAM ---
    if let Some(bytes) = host_ram() {
        return MemoryReading {
            available_bytes: bytes,
            source: MemorySource::HostRam,
        };
    }

    // --- fallback ---
    let fallback: usize = 4 * 1024 * 1024 * 1024; // 4 GiB
    warn!(
        available_bytes = fallback,
        "Memory detection failed on all paths; using 4 GiB fallback. \
         Set BGE_M3_AVAILABLE_MEMORY_BYTES to override."
    );
    MemoryReading {
        available_bytes: fallback,
        source: MemorySource::HostRam,
    }
}

/// Detects the number of NVIDIA GPU devices available on this instance.
///
/// Detection order (first success wins):
/// 1. `override_val` — the parsed value of `BGE_M3_GPU_COUNT` env var, if set.
/// 2. Linux `/proc/driver/nvidia/gpus/` directory entry count (compiled out on
///    non-Linux targets).
/// 3. Default: `1` (macOS `CoreML` is always single-device; Linux fallback when
///    the NVIDIA driver proc path is absent or empty).
///
/// Returns at least `1` regardless of the detection path. Logs the chosen
/// count and source at `INFO`.
pub(crate) fn detect_gpu_count(override_val: Option<usize>) -> usize {
    if let Some(n) = override_val {
        let n = n.max(1);
        tracing::info!(
            gpu_count = n,
            source = "env_override",
            "detected GPU(s) on this instance"
        );
        return n;
    }

    #[cfg(target_os = "linux")]
    if let Some(n) = count_nvidia_gpus_from_proc() {
        tracing::info!(
            gpu_count = n,
            source = "proc_driver",
            "detected GPU(s) on this instance"
        );
        return n;
    }

    tracing::info!(
        gpu_count = 1_usize,
        source = "default",
        "detected GPU(s) on this instance"
    );
    1
}

/// Counts NVIDIA GPU entries in `/proc/driver/nvidia/gpus/`.
///
/// Returns `None` when the directory does not exist, is unreadable, or
/// contains no entries (indicating no NVIDIA driver is loaded). Returns
/// `Some(n)` with `n ≥ 1` when entries are found.
#[cfg(target_os = "linux")]
fn count_nvidia_gpus_from_proc() -> Option<usize> {
    let count = std::fs::read_dir("/proc/driver/nvidia/gpus")
        .ok()?
        .filter_map(std::result::Result::ok)
        .count();
    if count > 0 { Some(count) } else { None }
}

/// Returns the current process's RSS (Resident Set Size) in bytes, or `None`
/// if measurement is not supported on this platform.
///
/// Linux: parses `/proc/self/statm`. Field 1 (index 1) is RSS in pages;
/// multiplied by the system page size (typically 4096).
///
/// macOS: returns `None` — requires `task_info` FFI which conflicts with
/// `unsafe_code = "forbid"`. A future release can add it via the `mach2`
/// crate.
pub(crate) fn read_process_rss_bytes() -> Option<usize> {
    #[cfg(target_os = "linux")]
    return linux_rss();

    #[cfg(not(target_os = "linux"))]
    None
}

// ---------------------------------------------------------------------------
// Private helpers
// ---------------------------------------------------------------------------

fn env_override() -> Option<usize> {
    std::env::var("BGE_M3_AVAILABLE_MEMORY_BYTES")
        .ok()
        .and_then(|v| {
            v.parse::<usize>().ok().or_else(|| {
                warn!(
                    value = %v,
                    "BGE_M3_AVAILABLE_MEMORY_BYTES is not a valid usize; ignoring"
                );
                None
            })
        })
}

#[cfg(target_os = "linux")]
fn cgroup_memory() -> Option<MemoryReading> {
    // Sentinel threshold: the cgroup v1 kernel uses a near-i64::MAX value when
    // no limit is configured. Treat any value ≥ 1 TiB as "unlimited".
    const ONE_TIB: usize = 1024 * 1024 * 1024 * 1024;

    // --- cgroup v2: path-walk from /proc/self/cgroup ---
    //
    // ECS Managed Instances (Bottlerocket) do NOT set --cgroupns=private, so
    // /sys/fs/cgroup/memory.max resolves to the host root where value is "max".
    // The container's actual limit lives at a deeper path recorded in
    // /proc/self/cgroup (unified v2 format: `0::<path>`).
    //
    // Walk ancestors deepest-first until a numeric limit < 1 TiB is found.
    // If the entire walk yields "max", fall through to cgroup v1 then host_ram.
    if let Some(reading) = cgroup_v2_walk("/sys/fs/cgroup", "/proc/self/cgroup") {
        return Some(reading);
    }

    // --- cgroup v1: /sys/fs/cgroup/memory/memory.limit_in_bytes ---
    if let Ok(raw) = std::fs::read_to_string("/sys/fs/cgroup/memory/memory.limit_in_bytes") {
        let trimmed = raw.trim();
        if let Ok(bytes) = trimmed.parse::<usize>()
            && bytes < ONE_TIB
        {
            tracing::debug!(bytes, source = "cgroup_v1", "Detected memory limit");
            return Some(MemoryReading {
                available_bytes: bytes,
                source: MemorySource::CgroupV1,
            });
        }
    }

    None
}

/// Reads the cgroup v2 memory limit by walking ancestors of the container's
/// cgroup path.
///
/// # Arguments
///
/// - `cgroup_fs_root`: the mountpoint of the cgroup v2 filesystem (normally
///   `/sys/fs/cgroup`; injectable for unit tests).
/// - `proc_self_cgroup`: path to the per-process cgroup file (normally
///   `/proc/self/cgroup`; injectable for unit tests).
///
/// Parses the unified-hierarchy line (`0::<path>`), then iterates from the
/// deepest ancestor up to the root, reading `memory.max` at each level.
/// Returns the first numeric limit found that is below 1 TiB, or `None`
/// when the entire walk yields `"max"` or the file is unreadable.
#[cfg(target_os = "linux")]
pub(crate) fn cgroup_v2_walk(
    cgroup_fs_root: &str,
    proc_self_cgroup: &str,
) -> Option<MemoryReading> {
    const ONE_TIB: usize = 1024 * 1024 * 1024 * 1024;

    let cgroup_content = std::fs::read_to_string(proc_self_cgroup).ok()?;

    // Unified hierarchy: exactly one line, format `0::<path>` (e.g. `0::/ecs.slice/…`)
    // Legacy v1 has multiple lines, each with `<hierarchy_id>:<controllers>:<path>`.
    // We only attempt v2 if we find the unified `0::` prefix.
    let cgroup_rel_path = cgroup_content
        .lines()
        .find_map(|line| line.strip_prefix("0::"))?;

    // Build the absolute cgroup directory path.
    let cgroup_dir = std::path::PathBuf::from(cgroup_fs_root).join(
        // Strip the leading '/' so PathBuf::join doesn't replace the root.
        cgroup_rel_path.trim_start_matches('/'),
    );

    // Walk ancestors from deepest to shallowest (inclusive of the container
    // cgroup itself, exclusive of the root mountpoint).
    let mut current = cgroup_dir.as_path();
    let fs_root = std::path::Path::new(cgroup_fs_root);

    loop {
        let memory_max = current.join("memory.max");
        if let Ok(raw) = std::fs::read_to_string(&memory_max) {
            let trimmed = raw.trim();
            if trimmed != "max"
                && let Ok(bytes) = trimmed.parse::<usize>()
                && bytes < ONE_TIB
            {
                tracing::debug!(
                    bytes,
                    source = "cgroup_v2",
                    path = %memory_max.display(),
                    "Detected memory limit"
                );
                return Some(MemoryReading {
                    available_bytes: bytes,
                    source: MemorySource::CgroupV2,
                });
            }
        }

        // Stop at the cgroup filesystem root — don't walk above it.
        if current == fs_root {
            break;
        }

        match current.parent() {
            Some(parent) => current = parent,
            None => break,
        }
    }

    None
}

/// Linux: parse `MemAvailable` from `/proc/meminfo` (kB → bytes).
/// macOS: read total host RAM via `sysctl hw.memsize`.
fn host_ram() -> Option<usize> {
    #[cfg(target_os = "linux")]
    return linux_meminfo_available();

    #[cfg(target_os = "macos")]
    return macos_host_ram();

    #[cfg(not(any(target_os = "linux", target_os = "macos")))]
    None
}

#[cfg(target_os = "linux")]
fn linux_meminfo_available() -> Option<usize> {
    let content = std::fs::read_to_string("/proc/meminfo").ok()?;
    for line in content.lines() {
        if line.starts_with("MemAvailable:") {
            let kb: usize = line.split_whitespace().nth(1)?.parse().ok()?;
            return Some(kb * 1024);
        }
    }
    None
}

#[cfg(target_os = "macos")]
fn macos_host_ram() -> Option<usize> {
    // `sysctl -n hw.memsize` returns an integer in bytes printed to stdout.
    let output = std::process::Command::new("sysctl")
        .args(["-n", "hw.memsize"])
        .output()
        .ok()?;
    let stdout = std::str::from_utf8(&output.stdout).ok()?.trim();
    stdout.parse::<usize>().ok()
}

#[cfg(target_os = "linux")]
fn linux_rss() -> Option<usize> {
    // /proc/self/statm: all values in pages.
    // Fields: size, rss, shared, text, lib, data, dt
    let raw = std::fs::read_to_string("/proc/self/statm").ok()?;
    let rss_pages: usize = raw.split_whitespace().nth(1)?.parse().ok()?;
    // SAFETY: page_size is a compile-time constant on Linux (4096 on x86_64/arm64).
    // We use sysconf(SC_PAGESIZE) via libc-free approach: fallback to 4096.
    let page_size = page_size_bytes();
    Some(rss_pages * page_size)
}

#[cfg(target_os = "linux")]
fn page_size_bytes() -> usize {
    // Read from /proc/self/auxv would be ideal but requires parsing ELF aux
    // vectors. Parsing /proc/$pid/smaps is too heavy. sysconf(SC_PAGESIZE)
    // requires libc. The practical answer on Linux/x86_64 and Linux/aarch64
    // is always 4096; we hard-code that to avoid any unsafe.
    4096
}

#[cfg(test)]
mod tests;