use crate::Result;
use std::collections::HashMap;
use std::fs::File;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::SyncSender;
use std::sync::{Arc, Mutex, OnceLock, Weak};
use super::GgmlDType;
static ENABLED: AtomicBool = AtomicBool::new(false);
pub fn set_enabled(on: bool) {
ENABLED.store(on, Ordering::Relaxed);
}
pub fn enabled() -> bool {
ENABLED.load(Ordering::Relaxed)
}
const PAGE_CACHE_RESERVE: u64 = 2_500_000_000;
const ACTIVATION_RESERVE: u64 = 1_200_000_000;
const BUDGET_FRACTION: f64 = 0.88;
const PIN_FRACTION: f64 = 0.25;
const REPIN_MAX_SWAPS: usize = 4;
pub struct ExpertStreamBank {
name: String,
file: File,
base_offset: u64,
expert_bytes: usize,
n_experts: usize,
dtype: GgmlDType,
n: usize,
k: usize,
inner: Mutex<BankState>,
}
struct BankState {
cap: usize,
clock: u64,
lru: HashMap<u32, LruEntry>,
pinned: HashMap<u32, Arc<[u8]>>,
usage: Vec<u64>,
heat: Vec<u32>,
hits: u64,
misses: u64,
}
struct LruEntry {
bytes: Arc<[u8]>,
used: u64,
}
impl ExpertStreamBank {
#[allow(clippy::too_many_arguments)]
pub fn open(
name: String,
path: &std::path::Path,
base_offset: u64,
expert_bytes: usize,
n_experts: usize,
dtype: GgmlDType,
n: usize,
k: usize,
) -> Result<Arc<Self>> {
let file = File::open(path)?;
let bank = Arc::new(Self {
name,
file,
base_offset,
expert_bytes,
n_experts,
dtype,
n,
k,
inner: Mutex::new(BankState {
cap: 1,
clock: 0,
lru: HashMap::new(),
pinned: HashMap::new(),
usage: vec![0; n_experts],
heat: vec![0; n_experts],
hits: 0,
misses: 0,
}),
});
registry().lock().expect("registry lock").register(&bank);
Ok(bank)
}
pub fn dtype(&self) -> GgmlDType {
self.dtype
}
pub fn n_experts(&self) -> usize {
self.n_experts
}
pub fn expert_dims(&self) -> (usize, usize) {
(self.n, self.k)
}
pub fn expert_bytes(&self) -> usize {
self.expert_bytes
}
pub fn logical_bytes(&self) -> usize {
self.expert_bytes * self.n_experts
}
pub fn fetch(&self, eid: u32) -> Result<Arc<[u8]>> {
let mut guard = self.inner.lock().expect("expert bank lock poisoned");
let st = &mut *guard;
if (eid as usize) < st.usage.len() {
st.usage[eid as usize] = st.usage[eid as usize].saturating_add(1);
st.heat[eid as usize] = st.heat[eid as usize].saturating_add(1);
}
if let Some(bytes) = st.pinned.get(&eid) {
st.hits += 1;
return Ok(bytes.clone());
}
if let Some(entry) = st.lru.get_mut(&eid) {
st.hits += 1;
st.clock += 1;
entry.used = st.clock;
return Ok(entry.bytes.clone());
}
st.misses += 1;
let bytes = self.read_expert(eid)?;
self.admit(st, eid, bytes.clone());
Ok(bytes)
}
fn admit(&self, st: &mut BankState, eid: u32, bytes: Arc<[u8]>) {
if st.lru.len() >= st.cap {
if let Some(&victim) = st.lru.iter().min_by_key(|(_, e)| e.used).map(|(k, _)| k) {
st.lru.remove(&victim);
}
}
st.clock += 1;
let used = st.clock;
st.lru.insert(eid, LruEntry { bytes, used });
}
fn read_expert(&self, eid: u32) -> Result<Arc<[u8]>> {
let off = self.base_offset + eid as u64 * self.expert_bytes as u64;
let mut buf = vec![0u8; self.expert_bytes];
read_exact_at(&self.file, &mut buf, off)?;
fadvise_dontneed(&self.file, off, self.expert_bytes as u64);
Ok(Arc::from(buf.into_boxed_slice()))
}
pub fn pin(&self, eid: u32) -> Result<()> {
let mut st = self.inner.lock().expect("expert bank lock poisoned");
if st.pinned.contains_key(&eid) {
return Ok(());
}
let bytes = if let Some(entry) = st.lru.remove(&eid) {
entry.bytes
} else {
drop(st);
let b = self.read_expert(eid)?;
st = self.inner.lock().expect("expert bank lock poisoned");
b
};
st.pinned.insert(eid, bytes);
Ok(())
}
pub fn repin(&self, max_swaps: usize) -> Result<usize> {
let mut guard = self.inner.lock().expect("expert bank lock poisoned");
let st = &mut *guard;
tier_decay(&mut st.heat);
let mut swaps = 0usize;
while swaps < max_swaps {
let pinned: Vec<u32> = st.pinned.keys().copied().collect();
let Some((slot, hot, _gain)) = tier_pick_swap(&st.heat, &pinned) else {
break;
};
let cold = pinned[slot];
if let Some(bytes) = st.pinned.remove(&cold) {
self.admit(st, cold, bytes);
}
let bytes = match st.lru.remove(&hot) {
Some(entry) => entry.bytes,
None => self.read_expert(hot)?,
};
st.pinned.insert(hot, bytes);
swaps += 1;
}
Ok(swaps)
}
fn prefetch_resident(&self, eid: u32) -> Result<()> {
{
let st = self.inner.lock().expect("expert bank lock poisoned");
if st.pinned.contains_key(&eid) || st.lru.contains_key(&eid) {
return Ok(());
}
}
let bytes = self.read_expert(eid)?;
let mut guard = self.inner.lock().expect("expert bank lock poisoned");
let st = &mut *guard;
if st.pinned.contains_key(&eid) || st.lru.contains_key(&eid) {
return Ok(());
}
self.admit(st, eid, bytes);
Ok(())
}
pub fn prefetch(self: &Arc<Self>, eid: u32) {
if !prefetch_enabled() || eid as usize >= self.n_experts {
return;
}
let _ = prefetcher().try_send(PrefetchJob {
bank: Arc::downgrade(self),
eid,
});
}
pub fn set_cap(&self, cap: usize) {
let mut st = self.inner.lock().expect("expert bank lock poisoned");
st.cap = cap.max(1);
}
pub fn stats(&self) -> (usize, usize, usize, u64, u64) {
let st = self.inner.lock().expect("expert bank lock poisoned");
(st.pinned.len(), st.lru.len(), st.cap, st.hits, st.misses)
}
pub fn resident_bytes(&self) -> usize {
let st = self.inner.lock().expect("expert bank lock poisoned");
(st.pinned.len() + st.lru.len()) * self.expert_bytes
}
fn usage_snapshot(&self) -> Vec<(u32, u64)> {
let st = self.inner.lock().expect("expert bank lock poisoned");
st.usage
.iter()
.enumerate()
.filter(|(_, &c)| c > 0)
.map(|(e, &c)| (e as u32, c))
.collect()
}
}
struct Registry {
banks: Vec<Weak<ExpertStreamBank>>,
sidecar: Option<PathBuf>,
}
fn registry() -> &'static Mutex<Registry> {
static REG: OnceLock<Mutex<Registry>> = OnceLock::new();
REG.get_or_init(|| {
Mutex::new(Registry {
banks: Vec::new(),
sidecar: None,
})
})
}
impl Registry {
fn register(&mut self, bank: &Arc<ExpertStreamBank>) {
self.banks.push(Arc::downgrade(bank));
}
fn live(&self) -> Vec<Arc<ExpertStreamBank>> {
self.banks.iter().filter_map(Weak::upgrade).collect()
}
}
pub fn set_usage_sidecar(path: PathBuf) {
registry().lock().expect("registry lock").sidecar = Some(path);
}
pub fn total_resident_bytes() -> usize {
registry()
.lock()
.expect("registry lock")
.live()
.iter()
.map(|b| b.resident_bytes())
.sum()
}
fn tier_pick_swap(heat: &[u32], pinned: &[u32]) -> Option<(usize, u32, i64)> {
if heat.is_empty() || pinned.is_empty() {
return None;
}
let cold = pinned
.iter()
.enumerate()
.min_by_key(|&(_, &p)| heat[p as usize])
.map(|(z, _)| z)
.expect("pinned is non-empty");
let mut hot: Option<usize> = None;
let mut hot_heat = 0u32;
for (e, &h) in heat.iter().enumerate() {
let resident = pinned.iter().any(|&p| p as usize == e);
if !resident && h > hot_heat {
hot_heat = h;
hot = Some(e);
}
}
let hot = hot?;
let cold_heat = heat[pinned[cold] as usize];
if hot_heat <= cold_heat + (cold_heat >> 2) + 4 {
return None;
}
Some((cold, hot as u32, hot_heat as i64 - cold_heat as i64))
}
fn tier_decay(heat: &mut [u32]) {
for h in heat.iter_mut() {
*h >>= 1;
}
}
pub fn repin_interval() -> usize {
static N: OnceLock<usize> = OnceLock::new();
*N.get_or_init(|| {
std::env::var("STREAM_EXPERTS_REPIN")
.ok()
.and_then(|v| v.parse::<usize>().ok())
.unwrap_or(0)
})
}
pub fn repin_all() -> usize {
if repin_interval() == 0 {
return 0;
}
let banks = registry().lock().expect("registry lock").live();
banks
.iter()
.map(|b| b.repin(REPIN_MAX_SWAPS).unwrap_or(0))
.sum()
}
pub fn prefetch_enabled() -> bool {
static ON: OnceLock<bool> = OnceLock::new();
*ON.get_or_init(|| {
std::env::var("STREAM_EXPERTS_PREFETCH")
.map(|v| !v.is_empty() && v != "0")
.unwrap_or(false)
})
}
struct PrefetchJob {
bank: Weak<ExpertStreamBank>,
eid: u32,
}
fn prefetcher() -> &'static SyncSender<PrefetchJob> {
static P: OnceLock<SyncSender<PrefetchJob>> = OnceLock::new();
P.get_or_init(|| {
let (tx, rx) = std::sync::mpsc::sync_channel::<PrefetchJob>(256);
std::thread::Builder::new()
.name("stream-experts-prefetch".into())
.spawn(move || {
while let Ok(job) = rx.recv() {
if let Some(bank) = job.bank.upgrade() {
let _ = bank.prefetch_resident(job.eid);
}
}
})
.expect("spawn stream-experts prefetch worker");
tx
})
}
pub fn finalize() {
let banks = registry().lock().expect("registry lock").live();
if banks.is_empty() {
return;
}
let expert_bytes = banks.iter().map(|b| b.expert_bytes).max().unwrap_or(1).max(1);
let n_banks = banks.len() as u64;
let forced_gb = std::env::var("STREAM_EXPERTS_RAM_GB")
.ok()
.and_then(|v| v.parse::<f64>().ok())
.filter(|g| *g > 0.0);
let avail = match forced_gb {
Some(gb) => (gb * 1e9) as u64,
None => mem_available_bytes(),
};
let budget = (avail as f64 * BUDGET_FRACTION) as u64;
let slack = PAGE_CACHE_RESERVE + ACTIVATION_RESERVE;
let for_cache = budget.saturating_sub(slack);
let mut cap = (for_cache / (n_banks * expert_bytes as u64)) as usize;
let max_experts = banks.iter().map(|b| b.n_experts).max().unwrap_or(1);
cap = cap.clamp(1, max_experts);
for b in &banks {
b.set_cap(cap);
}
let pinned = load_and_pin(&banks, cap);
register_atexit_save();
spawn_usage_saver();
eprintln!(
"[stream-experts] {} banks x {:.1} MB/expert; budget {:.1} GB -> cap {}/bank \
(cache {:.1} GB), pinned {}",
banks.len(),
expert_bytes as f64 / 1e6,
avail as f64 / 1e9,
cap,
(cap as u64 * n_banks * expert_bytes as u64) as f64 / 1e9,
pinned,
);
}
fn load_and_pin(banks: &[Arc<ExpertStreamBank>], cap: usize) -> usize {
let sidecar = match registry().lock().expect("registry lock").sidecar.clone() {
Some(p) if p.exists() => p,
_ => return 0,
};
let text = match std::fs::read_to_string(&sidecar) {
Ok(t) => t,
Err(_) => return 0,
};
let mut by_name: HashMap<&str, Vec<(u32, u64)>> = HashMap::new();
for line in text.lines() {
let mut it = line.split_whitespace();
let (Some(name), Some(eid), Some(cnt)) = (it.next(), it.next(), it.next()) else {
continue;
};
if let (Ok(eid), Ok(cnt)) = (eid.parse::<u32>(), cnt.parse::<u64>()) {
by_name.entry(name).or_default().push((eid, cnt));
}
}
let pin_budget = ((cap as f64 * PIN_FRACTION) as usize).max(1);
let mut pinned = 0usize;
for b in banks {
let Some(rows) = by_name.get(b.name.as_str()) else {
continue;
};
let mut rows = rows.clone();
rows.sort_by(|a, c| c.1.cmp(&a.1));
for (eid, _) in rows.into_iter().take(pin_budget) {
if eid as usize >= b.n_experts {
continue;
}
if b.pin(eid).is_ok() {
pinned += 1;
}
}
}
pinned
}
pub fn save_usage() -> Result<()> {
let (banks, sidecar) = {
let reg = registry().lock().expect("registry lock");
(reg.live(), reg.sidecar.clone())
};
let Some(path) = sidecar else {
return Ok(());
};
let mut merged: HashMap<(String, u32), u64> = HashMap::new();
if let Ok(text) = std::fs::read_to_string(&path) {
for line in text.lines() {
let mut it = line.split_whitespace();
if let (Some(name), Some(eid), Some(cnt)) = (it.next(), it.next(), it.next()) {
if let (Ok(eid), Ok(cnt)) = (eid.parse::<u32>(), cnt.parse::<u64>()) {
*merged.entry((name.to_string(), eid)).or_default() += cnt;
}
}
}
}
for b in &banks {
for (eid, cnt) in b.usage_snapshot() {
*merged.entry((b.name.clone(), eid)).or_default() += cnt;
}
}
let mut out = String::new();
for ((name, eid), cnt) in merged {
out.push_str(&format!("{name} {eid} {cnt}\n"));
}
std::fs::write(&path, out)?;
Ok(())
}
fn register_atexit_save() {
#[cfg(unix)]
{
static ONCE: std::sync::Once = std::sync::Once::new();
ONCE.call_once(|| {
extern "C" fn on_exit() {
let _ = save_usage();
}
unsafe {
libc::atexit(on_exit);
}
});
}
}
const USAGE_SAVE_INTERVAL: std::time::Duration = std::time::Duration::from_secs(60);
fn spawn_usage_saver() {
static ONCE: std::sync::Once = std::sync::Once::new();
ONCE.call_once(|| {
let _ = std::thread::Builder::new()
.name("expert-usage-saver".into())
.spawn(|| loop {
std::thread::sleep(USAGE_SAVE_INTERVAL);
let _ = save_usage();
});
});
}
#[cfg(unix)]
fn read_exact_at(file: &File, buf: &mut [u8], offset: u64) -> Result<()> {
use std::os::unix::fs::FileExt;
file.read_exact_at(buf, offset)?;
Ok(())
}
#[cfg(not(unix))]
fn read_exact_at(_file: &File, _buf: &mut [u8], _offset: u64) -> Result<()> {
crate::bail!("streaming experts require a unix positional-read (pread) platform")
}
#[cfg(target_os = "linux")]
fn fadvise_dontneed(file: &File, offset: u64, len: u64) {
use std::os::unix::io::AsRawFd;
unsafe {
libc::posix_fadvise(
file.as_raw_fd(),
offset as libc::off_t,
len as libc::off_t,
libc::POSIX_FADV_DONTNEED,
);
}
}
#[cfg(not(target_os = "linux"))]
fn fadvise_dontneed(_file: &File, _offset: u64, _len: u64) {}
fn mem_available_bytes() -> u64 {
#[cfg(target_os = "linux")]
{
if let Ok(text) = std::fs::read_to_string("/proc/meminfo") {
for line in text.lines() {
if let Some(rest) = line.strip_prefix("MemAvailable:") {
if let Some(kb) = rest.split_whitespace().next().and_then(|v| v.parse::<u64>().ok())
{
return kb.saturating_mul(1024);
}
}
}
}
}
static WARNED: AtomicU64 = AtomicU64::new(0);
if WARNED.swap(1, Ordering::Relaxed) == 0 {
eprintln!("[stream-experts] MemAvailable unreadable; assuming 8 GB free");
}
8_000_000_000
}
#[cfg(test)]
mod tests {
use super::{tier_decay, tier_pick_swap, REPIN_MAX_SWAPS};
fn simulate(heat: &[u32], pinned: &mut Vec<u32>, max_swaps: usize) -> usize {
let mut swaps = 0;
while swaps < max_swaps {
match tier_pick_swap(heat, pinned) {
Some((slot, hot, _gain)) => {
pinned[slot] = hot;
swaps += 1;
}
None => break,
}
}
swaps
}
#[test]
fn swaps_coldest_pinned_for_hottest_streamed() {
let heat = [100, 5, 200, 50];
let pinned = [0, 1];
let (slot, hot, gain) = tier_pick_swap(&heat, &pinned).expect("beneficial swap");
assert_eq!(slot, 1, "coldest pinned is at slot 1 (eid 1)");
assert_eq!(hot, 2, "hottest streamed is eid 2");
assert_eq!(gain, 195, "gain is hot_heat - cold_heat = 200 - 5");
}
#[test]
fn hysteresis_blocks_near_equal_experts() {
let heat = [100, 100, 110];
let pinned = [0, 1];
assert!(tier_pick_swap(&heat, &pinned).is_none());
}
#[test]
fn hysteresis_margin_is_exclusive_at_the_boundary() {
assert!(
tier_pick_swap(&[100, 100, 129], &[0, 1]).is_none(),
"equal to the threshold must not swap"
);
let (slot, hot, gain) =
tier_pick_swap(&[100, 100, 130], &[0, 1]).expect("one over the threshold swaps");
assert_eq!((slot, hot, gain), (0, 2, 30));
}
#[test]
fn no_swap_when_every_expert_is_already_pinned() {
assert!(tier_pick_swap(&[10, 20], &[0, 1]).is_none());
}
#[test]
fn empty_inputs_are_safe() {
assert!(tier_pick_swap(&[], &[]).is_none());
assert!(tier_pick_swap(&[1, 2, 3], &[]).is_none());
assert!(tier_pick_swap(&[], &[0, 1]).is_none());
}
#[test]
fn picks_the_coldest_slot_and_hottest_candidate_among_many() {
let heat = [50, 10, 30, 200, 80, 5];
let pinned = [0, 2, 4];
let (slot, hot, gain) = tier_pick_swap(&heat, &pinned).expect("swap");
assert_eq!(slot, 1, "eid 2 (heat 30) is the coldest pinned, at slot 1");
assert_eq!(hot, 3, "eid 3 (heat 200) is the hottest streamed");
assert_eq!(gain, 170);
}
#[test]
fn decay_halves_every_expert() {
let mut heat = [10, 3, 0, 255, 1];
tier_decay(&mut heat);
assert_eq!(heat, [5, 1, 0, 127, 0]);
}
#[test]
fn repeated_swaps_converge_then_stop_under_hysteresis() {
let heat = [100, 90, 5, 4];
let mut pinned = vec![2, 3];
let swaps = simulate(&heat, &mut pinned, REPIN_MAX_SWAPS);
assert_eq!(swaps, 2, "exactly the two hot experts get pinned");
pinned.sort_unstable();
assert_eq!(pinned, vec![0, 1], "hot-set converged to the two hottest experts");
assert_eq!(simulate(&heat, &mut pinned, REPIN_MAX_SWAPS), 0);
}
#[test]
fn one_dominant_expert_settles_after_a_single_swap() {
let heat = [50, 10, 30, 200, 80, 5];
let mut pinned = vec![0, 2, 4];
assert_eq!(simulate(&heat, &mut pinned, REPIN_MAX_SWAPS), 1);
pinned.sort_unstable();
assert_eq!(pinned, vec![0, 3, 4]);
}
}