Testing Frameworks for Simulators

@pytest.fixture(scope="session")
def baseline_run():
    """One 300-request run, built once and shared by every test that asks for it."""
    wl = poisson_workload(4.0, 300, LengthDist(2048, 0.5), LengthDist(128, 0.5), seed=1)
    return simulate(SimConfig(), wl)


@pytest.fixture
def small_workload():
    """A fresh workload per test: simulate() mutates requests, so never share one."""
    return poisson_workload(2.0, 50, LengthDist(512, 0.3), LengthDist(32, 0.3), seed=7)

@pytest.mark.parametrize("cfg", [
    SimConfig(),
    SimConfig(mode="colocated", n_colocated=2),
    SimConfig(n_prefill=2, n_decode=2),
], ids=["1P1D", "colocated", "2P2D"])
def test_kv_is_released(cfg, small_workload):
    res = simulate(cfg, small_workload)
    assert all(i.kv_used == 0 for i in res.instances)


@pytest.mark.slow
def test_long_run_is_stable():
    from disagg_sim.workload import LengthDist, poisson_workload
    wl = poisson_workload(3.0, 3000, LengthDist(2048, 0.5), LengthDist(128, 0.5), seed=2)
    m = summarise(simulate(SimConfig(fast_forward=True), wl))
    assert m["throughput"]["slo_attainment"] > 0.9


@pytest.mark.xfail(strict=True, reason="known optimism: tensor-parallel all-reduce is not modelled")

@settings(max_examples=40, deadline=None, derandomize=True)
@given(rate=st.floats(0.5, 20.0))
def test_ttft_p99_under_one_second_at_any_rate(rate):
    wl = poisson_workload(rate, 200, LengthDist(2048, 0.5), LengthDist(64, 0.5), seed=0)
    m = summarise(simulate(SimConfig(fast_forward=True), wl))
    assert m["latency_s"]["ttft"]["p99"] < 1.0

class EventQueueMachine(RuleBasedStateMachine):
    def __init__(self):
        super().__init__()
        self.q = EventQueue(fifo_ties=FIFO)
        self.model = []          # (time, priority, insertion order, name): sorted() is the spec
        self.n = 0

    @rule(dt=st.integers(0, 3), urgent=st.booleans())
    def schedule(self, dt, urgent):
        prio = URGENT if urgent else NORMAL
        name = f"e{self.n}"
        self.q.schedule(self.q.now + dt, name, prio)
        self.model.append((self.q.now + dt, prio, self.n, name))
        self.n += 1

    @precondition(lambda self: self.model)
    @rule()
    def pop(self):
        want = min(self.model)
        self.model.remove(want)
        assert self.q.pop() == (want[0], want[3])

    @invariant()
    def sizes_agree(self):
        assert len(self.q) == len(self.model)


TestEventQueue = EventQueueMachine.TestCase

state = EventQueueMachine()
state.schedule(dt=0, urgent=False)
state.schedule(dt=0, urgent=False)
state.pop()
AssertionError: assert (0.0, 'e1') == (0.0, 'e0')

@pytest.fixture
def small_workload():
    """A fresh workload per test: simulate() mutates requests, so never share one."""
    return poisson_workload(2.0, 50, LengthDist(512, 0.3), LengthDist(32, 0.3), seed=7)


@pytest.fixture
def golden(request, pytestconfig):
    """Compare a result with tests/golden/<test name>.json, or rewrite it with --bless."""
    path = request.path.parent / "golden" / f"{request.node.name}.json"

    def check(result: dict):
        if pytestconfig.getoption("bless") or not path.exists():
            path.parent.mkdir(exist_ok=True)
            path.write_text(json.dumps(result, indent=1, sort_keys=True))
            pytest.skip(f"blessed {path.name}")
        assert result == json.loads(path.read_text())

    return check

#[test]
fn matches_the_python_simulator_exactly() {
    let cases: Vec<Value> = serde_json::from_str(include_str!("fixtures/golden.json")).unwrap();
    assert!(cases.len() >= 5);
    for case in &cases {
        let name = case["name"].as_str().unwrap();
        let spec: ConfigSpec = serde_json::from_value(case["config"].clone()).unwrap();
        let rows = case["rows"].as_array().unwrap();
        let wl = rows
            .iter()
            .enumerate()
            .map(|(i, r)| {
                Request::new(
                    i,
                    r[0].as_f64().unwrap(),
                    r[1].as_i64().unwrap(),
                    r[2].as_i64().unwrap(),
                )
            })
            .collect();
        let res = simulate(spec.build().unwrap(), wl).unwrap();
        for (r, want) in res.requests.iter().zip(case["stamps"].as_array().unwrap()) {
            let got = [
                r.prefill_start,
                r.first_token,
                r.kv_start,
                r.kv_ready,
                r.decode_start,
                r.finish,
            ];
            let want: Vec<Option<f64>> =
                want.as_array().unwrap().iter().map(Value::as_f64).collect();
            assert_eq!(got.to_vec(), want, "{name}: request {}", r.rid);
        }
        assert_eq!(summarise(&res), case["summary"], "{name}: summary differs");
    }
}

def step_time(flops, nbytes, peak_flops, bandwidth, overhead=0.5e-3):
    """Seconds for one step, and which roof bounds it."""
    tc = flops / peak_flops
    tm = nbytes / bandwidth
    if tc >= tm:
        return tc + overhead, "compute"
    return tm + overhead, "memory"

-def step_time(flops, nbytes, peak_flops, bandwidth, overhead=0.5e-3):
+def step_time(flops, nbytes, peak_flops, bandwidth, overhead=1.0005):

-    if tc >= tm:
+    if tc > tm:

-        return tc + overhead, "compute"
+        return tc - overhead, "compute"

-    return tm + overhead, "memory"
+    return tm - overhead, "memory"

// A fixture: SetUp() runs before every TEST_F that names it.
class TiedQueue : public ::testing::Test {
protected:
    void SetUp() override {
        q.schedule(1.0, "n1");
        q.schedule(1.0, "n2");
        q.schedule(1.0, "urgent", Priority::Urgent);
    }
    EventQueue q;
};

TEST_F(TiedQueue, UrgentBeatsNormalAtTheSameTime) {
    EXPECT_EQ(q.pop().name, "urgent");
}

TEST_F(TiedQueue, EqualKeysAreFirstInFirstOut) {
    q.pop();
    ASSERT_EQ(q.size(), 2u);              // ASSERT stops this test if it fails; EXPECT carries on
    EXPECT_EQ(q.pop().name, "n1");
    EXPECT_EQ(q.pop().name, "n2");
}

// A value-parameterised test: one body, many inputs.
class ClockNeverGoesBack : public ::testing::TestWithParam<std::tuple<double, double, double>> {};

TEST_P(ClockNeverGoesBack, AcrossInsertionOrders) {
    auto [a, b, c] = GetParam();
    EventQueue q;
    for (double t : {a, b, c}) q.schedule(t, "e");
    double last = -1.0;
    while (q.size()) {
        double t = q.pop().time;
        EXPECT_GE(t, last);
        last = t;
    }
}

INSTANTIATE_TEST_SUITE_P(Orders, ClockNeverGoesBack,

fn spec() -> impl Strategy<Value = ConfigSpec> {
    (
        prop::bool::ANY,
        1usize..4,
        1usize..4,
        prop::sample::select(vec!["ib-ndr", "eth-25g", "nvlink4", "pcie5", "eth-100g"]),
        1usize..4,
        prop::option::of(200.0f64..700.0),
        prop::bool::ANY,
        prop::sample::select(vec![8usize, 64, 256]),
    )
        .prop_map(|(colo, a, b, link, ch, cap, dvfs, maxb)| ConfigSpec {
            mode: if colo { "colocated" } else { "disagg" }.into(),
            n_prefill: a,
            n_decode: b,
            n_colocated: a,
            link: link.into(),
            link_channels: ch,
            power_cap_w: cap,
            dvfs,
            max_decode_batch: maxb,
            ..Default::default()
        })
}

// Every request either finishes or is rejected, never both.
let finished = res.requests.iter().filter(|r| r.finish.is_some()).count();
prop_assert_eq!(finished + res.rejected.len(), res.requests.len());

for r in res.requests.iter().filter(|r| r.finish.is_some()) {
    // Timestamps never go backwards along a request's path.
    let path: Vec<f64> = [Some(r.arrival), r.prefill_start, r.first_token, r.kv_start,
                          r.kv_ready, r.decode_start, r.finish].into_iter().flatten().collect();
    prop_assert!(path.windows(2).all(|w| w[0] <= w[1]), "{:?}", path);
    // One inter-token latency per token after the first, all positive.
    prop_assert_eq!(r.itls.len() as i64, (r.output_len - 1).max(0));
    prop_assert!(r.itls.iter().all(|&x| x > 0.0));
    // Stages partition end-to-end latency.
    let sum: f64 = r.stages().iter().sum();
    prop_assert!((sum - r.e2e()).abs() <= 1e-9 * r.e2e().max(1.0));

def golden(a: int, b: int) -> int:
    """The reference model: what the hardware must compute."""
    return (a + b) % Q


def stimulus(rng: random.Random, n: int):
    """Constrained random: mostly uniform, plus the corners a uniform draw rarely hits."""
    corners = [(0, 0), (Q - 1, 0), (Q - 1, 1), (Q - 1, Q - 1), (Q // 2, Q - Q // 2)]
    yield from corners
    for _ in range(n):
        yield rng.randrange(Q), rng.randrange(Q)


@cocotb.test()
async def matches_golden_model_with_full_coverage(dut):
    cocotb.start_soon(Clock(dut.clk, 10, unit="ns").start())
    dut.rst_n.value, dut.in_valid.value = 0, 0
    await RisingEdge(dut.clk)
    dut.rst_n.value = 1

    bins = {"no wrap": 0, "wrap": 0, "sum == Q (result 0)": 0, "max operands": 0}
    for a, b in stimulus(random.Random(2026), 2000):
        await FallingEdge(dut.clk)                      # drive away from the active edge
        dut.a.value, dut.b.value, dut.in_valid.value = a, b, 1
        await RisingEdge(dut.clk)
        await ReadOnly()                                # sample after the register updates
        assert dut.out_valid.value == 1
        assert int(dut.r.value) == golden(a, b), f"{a} + {b}: got {int(dut.r.value)}"
        bins["wrap" if a + b >= Q else "no wrap"] += 1
        bins["sum == Q (result 0)"] += a + b == Q
        bins["max operands"] += a == b == Q - 1

    dut._log.info("coverage: %s", bins)
    assert all(bins.values()), f"coverage hole: {bins}"

stage('Rust tests') {
    steps {
        sh(env.WITH_CARGO + 'cargo nextest run --release --profile ci')
    }
}

stage('Python differential tests') {
    steps {
        sh(env.WITH_CARGO + '''
            python3 -m venv .venv
            .venv/bin/pip install -q maturin
            .venv/bin/maturin develop --release -q -E dev
            # The venv outlives builds, and pip keeps an installed git dependency whose version
            # number has not changed, so fetch the Python reference's current commit every time.
            .venv/bin/pip install -q --force-reinstall --no-deps "disagg-sim @ git+https://github.com/BrendanJamesLynskey/Disaggregated_Inference_Sim"
            .venv/bin/pytest pytests --junitxml=pytest-junit.xml
        ''')
    }
}

stage('Coverage') {
    steps {
        sh(env.WITH_CARGO + 'cargo llvm-cov --release --cobertura --output-path coverage.xml')
        // cargo-llvm-cov lists each generic instantiation as its own method, and the
        // Coverage plugin's Cobertura parser rejects duplicate names: skip them.
        recordCoverage(tools: [[parser: 'COBERTURA', pattern: 'coverage.xml']],