open-standards/eoc/eoc-rs/crates/eoc-cascade/tests/neural_span_process_energy.rs

//! What a resolve the neural stage answers is charged, against the
//! hardware: the cascade on this process's own CPU energy counter, its span
//! held against an oracle that reads the same count its own way.
//!
//! The cascade reads its meter only once the three deterministic tiers have
//! missed, and again after the neural stage. So the figure is the neural
//! stage's, and work done before the walk reached it is not in it. Here the
//! KV stage burns 8 ms of this thread's CPU on its way to missing, every
//! resolve, straight before the cascade opens its span.
//!
//! One test binary with one test: the counter's domain is every thread of
//! the process, and the oracle's check is the process's context-switch
//! count, so nothing else may run here. No joule figure and no timing is
//! asserted: what a millisecond of CPU costs follows the core that ran it,
//! and a loaded host stretches everything's wall time.

#![cfg(target_os = "macos")]

mod oracle;

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use eoc_cache::LruCache;
use eoc_cascade::Cascade;
use eoc_core::{JouleCost, JouleSource, Query, Response, Stage};
use eoc_graph::GraphStage;
use eoc_kv::{KvBackend, KvStage};
use eoc_meter::{CounterStatus, JouleCounter, MacosProcessEnergyCounter};
use eoc_neural::{NeuralBackend, NeuralStage};

use oracle::{Counted, Span, Watched, burn, energy_nj_off_core};

/// What the KV stage burns before it misses, and what the neural stage
/// then burns or sleeps.
const BEFORE_THE_SPAN: Duration = Duration::from_millis(8);
const NEURAL: Duration = Duration::from_millis(2);

/// A KV backend that knows nothing, and burns this thread's CPU finding
/// that out. It counts its lookups, so the test can see the walk ran.
struct BurnsThenMisses(Arc<AtomicU64>);

impl KvBackend for BurnsThenMisses {
    fn get(&self, _key: &str) -> Option<Vec<u8>> {
        burn(BEFORE_THE_SPAN);
        self.0.fetch_add(1, Ordering::SeqCst);
        None
    }
    fn put(&self, _key: &str, _value: Vec<u8>) {}
    fn embeddings(&self) -> Vec<(Vec<f32>, String)> {
        Vec::new()
    }
    fn put_embedding(&self, _embedding: Vec<f32>, _key: String) {}
}

/// A neural stage that works, or only sleeps, and claims no figure for it:
/// the resolve's cost is then the cascade's span alone, measured.
enum Neural {
    Burns,
    Sleeps,
}

#[async_trait::async_trait]
impl NeuralBackend for Neural {
    async fn infer(&self, q: &Query) -> Response {
        match self {
            Self::Burns => burn(NEURAL),
            Self::Sleeps => std::thread::sleep(NEURAL),
        }
        Response::new(q.id, "inferred".to_string(), Stage::Neural, JouleCost::zero())
    }
}

/// A cascade whose KV stage burns and misses, around `neural`, metered by
/// `meter`. Nothing is memoized, so every resolve walks to the neural stage.
fn cascade(meter: Arc<dyn JouleCounter>, neural: Neural, lookups: &Arc<AtomicU64>) -> Cascade {
    Cascade::new(
        Arc::new(LruCache::new(8)),
        Arc::new(KvStage::new(Box::new(BurnsThenMisses(lookups.clone())))),
        Arc::new(GraphStage::new()),
        Arc::new(NeuralStage::new(Box::new(neural))),
    )
    .with_meter(meter)
    .without_memoization()
}

#[tokio::test]
async fn a_neural_resolve_reads_the_neural_stage_and_nothing_before_it() {
    const TRIES: usize = 40;
    let status = eoc_meter::process_energy_status();
    eprintln!("neural span: status={status:?}");
    assert_eq!(status, CounterStatus::Live, "this process's CPU energy counter is not live here");
    let q = Query::new("nothing knows this");
    let lookups = Arc::new(AtomicU64::new(0));

    // 1. A neural stage that burns 2 ms, through a counter the oracle reads
    //    on either side of. The resolve takes exactly two reads, both after
    //    the KV stage's burn. Its figure lies between the oracle's window
    //    inside the span, which holds the neural stage, and the oracle's
    //    window around the span and both reads, which opens after the burn:
    //    that much holds for any count that only grows. What the windows
    //    leave open is their own distance apart, the two readers' reads and
    //    wake-ups: the figure is pinned within a third when the window
    //    inside is at least two thirds of the one around.
    let counter = MacosProcessEnergyCounter::try_new().expect("the process counter is live here");
    let watched = Arc::new(Watched::beside(counter, lookups.clone()));
    let through_oracle = cascade(watched.clone(), Neural::Burns, &lookups);
    let (mut tight, mut loose, mut unread, mut burned) = (Vec::new(), Vec::new(), 0usize, Vec::new());
    for attempt in 0..TRIES {
        watched.take();
        let (walked, before) = (lookups.load(Ordering::SeqCst), energy_nj_off_core());
        let r = through_oracle.resolve(q.clone()).await;
        let reads = watched.take();
        assert_eq!(r.stage, Stage::Neural);
        assert_eq!(lookups.load(Ordering::SeqCst), walked + 1, "the walk went through the KV stage");
        assert_eq!(reads.len(), 2, "#{attempt}: a neural resolve reads the meter exactly twice: {reads:?}");
        assert_eq!(reads[0].count, walked + 1, "#{attempt}: the span opened before the KV stage had missed");
        let span = Span::between(&reads[0], &reads[1]);
        // The burn is in the oracle's count before the cascade's first read.
        burned.push(reads[0].before_nanojoules - before);
        let Some(read) = span.read else {
            // A read failed: the stage's own cost stands, an estimated zero.
            assert_eq!(r.joule_cost, JouleCost::zero(), "#{attempt}");
            unread += 1;
            continue;
        };
        assert_eq!(r.joule_cost, JouleCost::measured(read / 1_000), "#{attempt}: the span, measured");
        assert!(
            span.holds(read),
            "#{attempt}: {} µJ measured outside the oracle's windows, {} nJ inside the span and {} nJ around it",
            read / 1_000,
            span.inside,
            span.around
        );
        if span.tightness() * 3 < 2_000 {
            loose.push((attempt, read, span.inside, span.around));
        } else {
            tight.push(span.tightness());
        }
    }
    burned.sort_unstable();
    eprintln!(
        "neural span: a 2 ms neural stage x {TRIES}: {} figures within the oracle's two windows with the inner at \
         least {}/1000 of the outer, {} looser {loose:?} (try, measured nJ, inside nJ, around nJ), {unread} not read; \
         the 8 ms before each span cost {} to {} µJ and none of it is in a figure",
        tight.len(),
        tight.iter().min().copied().unwrap_or(0),
        loose.len(),
        burned.first().copied().unwrap_or(0) / 1_000,
        burned.last().copied().unwrap_or(0) / 1_000,
    );

    // 2. A neural stage that only sleeps 2 ms, on the bare counter, as the
    //    cascade is used: nothing takes this thread off its core between the
    //    KV stage's 8 ms and the cascade's first read. The oracle's window
    //    is the whole resolve, burn and sleep. A span opened before the
    //    walk would read the burn; opened after it, the figure for the
    //    sleep is under a quarter of that window, or no reading at all.
    let counted = Arc::new(Counted::new(MacosProcessEnergyCounter::try_new().expect("the process counter is live here")));
    let bare = cascade(counted.clone(), Neural::Sleeps, &lookups);
    let (mut slept, mut slept_unread, mut leaked) = (Vec::new(), 0usize, Vec::new());
    for attempt in 0..TRIES {
        let (walked, reads, before) = (lookups.load(Ordering::SeqCst), counted.reads(), energy_nj_off_core());
        let r = bare.resolve(q.clone()).await;
        let whole = energy_nj_off_core() - before;
        assert_eq!(lookups.load(Ordering::SeqCst), walked + 1, "the walk went through the KV stage");
        assert_eq!(counted.reads(), reads + 2, "#{attempt}: a neural resolve reads the meter exactly twice");
        match r.joule_cost.source {
            // No span: the stage's own cost stands, an estimated zero.
            JouleSource::Estimated => {
                assert_eq!(r.joule_cost, JouleCost::zero(), "#{attempt}");
                slept_unread += 1;
            }
            JouleSource::Measured if r.joule_cost.microjoules * 1_000 * 4 >= whole => {
                leaked.push((attempt, r.joule_cost.microjoules, whole));
            }
            JouleSource::Measured => slept.push(r.joule_cost.microjoules * 1_000_000 / whole.max(1)),
        }
    }
    eprintln!(
        "neural span: a 2 ms sleep after the KV stage's 8 ms burn x {TRIES}: {} read under a quarter of the oracle's \
         window over the whole resolve (at most {}/1000 of it), {slept_unread} not read, {} read the burn before the \
         span {leaked:?} (try, µJ, oracle nJ)",
        slept.len(),
        slept.iter().max().copied().unwrap_or(0),
        leaked.len()
    );

    // A tenth may miss: on a loaded host a reader's wake-up can land on a
    // costlier core than the work did, and a read whose thread was not seen
    // to leave its core is no reading by design. A span opened before the
    // walk misses every time.
    assert!(
        (loose.len() + unread) * 10 <= TRIES,
        "{} of {TRIES} neural spans were not pinned within a third by the oracle {loose:?} and {unread} were not read",
        loose.len()
    );
    assert!(leaked.len() * 10 <= TRIES, "{} of {TRIES} sleeps read the burn before them: {leaked:?}", leaked.len());
    assert!(slept_unread * 10 <= TRIES, "{slept_unread} of {TRIES} sleeping neural stages were not read");
}