diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 099634301e89..9a75a427684e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -6,6 +6,12 @@ target_compile_features(test-server-prefill PRIVATE cxx_std_17) add_test(NAME test-server-prefill COMMAND test-server-prefill) set_tests_properties(test-server-prefill PROPERTIES LABELS "main") +# Model-free learning-budget validator: the ledger every learning step debits. +add_executable(test-server-learn-ledger test-server-learn-ledger.cpp) +target_compile_features(test-server-learn-ledger PRIVATE cxx_std_17) +add_test(NAME test-server-learn-ledger COMMAND test-server-learn-ledger) +set_tests_properties(test-server-learn-ledger PROPERTIES LABELS "main") + function(llama_build source) set(TEST_SOURCES ${source} ${ARGN}) diff --git a/tests/test-server-learn-ledger.cpp b/tests/test-server-learn-ledger.cpp new file mode 100644 index 000000000000..333da6405701 --- /dev/null +++ b/tests/test-server-learn-ledger.cpp @@ -0,0 +1,133 @@ +// what this catches: the learning ledger's contract (server-learn-ledger.h) without a model, in +// milliseconds: authority-only registration, immutable live grants, no replenishment by retry, +// free reservations as live custody, cancel only before start, checked ranges, bounded accounts. +#include "../tools/server/server-learn-ledger.h" + +#include +#include +#include + +static void require(bool condition, const char * message) { + if (!condition) { + std::cerr << "FAILED: " << message << '\n'; + std::exit(1); + } +} + +using R = server_learn_refusal; + +static int64_t used(const server_learn_ledger & l, const char * grant, int64_t now) { + int64_t out = -1; + require(l.used_us(grant, now, out) == R::none, "a valid observation"); + return out; +} + +int main() { + server_learn_ledger ledger("engine-a"); + const int64_t day_us = 86400000000; + // 10% of a 1000 ms period: 100 ms of contended learning per period, for a day + const server_learn_grant g{ "grant-1", "engine-a", 1000, 100000, day_us }; + + // registration is the authority's, with checked ranges + require(ledger.register_grant({ "", "engine-a", 1000, 100000, day_us }, 0) == R::invalid, "empty id"); + require(ledger.register_grant({ "x", "engine-a", 0, 100000, day_us }, 0) == R::invalid, "zero period"); + require(ledger.register_grant({ "x", "engine-a", server_learn_ledger::kMaxPeriodMs + 1, 100000, day_us }, 0) == R::invalid, + "a period past the checked range"); + require(ledger.register_grant({ "x", "engine-a", 1000, 0, day_us }, 0) == R::invalid, "zero share is no learning"); + require(ledger.register_grant({ "x", "engine-a", 1000, 1000001, day_us }, 0) == R::invalid, "share over 1e6"); + require(ledger.register_grant({ "x", "engine-a", 1000, 100000, 0 }, 0) == R::invalid, "already expired"); + require(ledger.register_grant({ "x", "engine-b", 1000, 100000, day_us }, 0) == R::stale_engine, "another incarnation"); + require(ledger.n_accounts() == 0, "refused grants register nothing"); + require(ledger.admit("never-registered", 10, 0, true).refusal == R::unknown_grant, "a step cannot mint a grant"); + require(ledger.admit("grant-1", server_learn_ledger::kMaxCostUs + 1, 0, true).refusal == R::invalid, "a cost past the range"); + + require(ledger.register_grant(g, 0) == R::none, "the authority registers"); + require(ledger.register_grant(g, 0) == R::none, "re-confirming the same terms is a no-op"); + require(ledger.register_grant({ "grant-1", "engine-a", 1000, 900000, day_us }, 0) == R::grant_mismatch, "a live grant is immutable"); + + // reserve 60 ms, then a second 60 ms does not fit beside it + const auto a1 = ledger.admit("grant-1", 60000, 0, true); + require(a1.refusal == R::none && a1.reservation != 0, "first step admitted"); + const auto a2 = ledger.admit("grant-1", 60000, 1000, true); + require(a2.refusal == R::budget_exhausted && a2.retry_after_ms == 999, "reserved time counts against the share"); + + // started: it settles with what it measured, never cancels + require(ledger.finish(a1.reservation, 30000, 1, 900) == R::not_started, "finish before start is refused"); + require(ledger.start(a1.reservation, 1000) == R::none, "start"); + require(ledger.start(a1.reservation, 1001) == R::already_started, "one reservation authorizes one execution"); + require(ledger.cancel(a1.reservation, 30000) == R::started, "a started step cannot be cancelled"); + require(ledger.finish(a1.reservation, 30000, 1, 40000) == R::none, "it settles after quiescence"); + require(used(ledger, "grant-1", 40000) == 30000 && ledger.committed_steps("grant-1") == 1, "debit and steps measured"); + require(ledger.finish(a1.reservation, 30000, 1, 40000) == R::unknown_reservation, "never settled twice"); + + // 70 ms remain: a retry of the same request spends the same account + const auto a3 = ledger.admit("grant-1", 70000, 41000, true); + require(a3.refusal == R::none, "the remaining 70 ms fit"); + require(ledger.admit("grant-1", 1, 42000, true).refusal == R::budget_exhausted, "a duplicate never replenishes"); + require(ledger.finish(a3.reservation, 70000, 1, 42500) == R::not_started, "a never-started step cannot settle committed work"); + require(ledger.cancel(a3.reservation, 43000) == R::none, "never started: cancelled (custody kept for it)"); + require(used(ledger, "grant-1", 43000) == 30000, "cancel releases the whole reservation"); + + // a step larger than the whole period's share can never run: said so, not retried + require(ledger.admit("grant-1", 100001, 44000, true).refusal == R::never_fits, "never fits"); + + // a free (uncontended) reservation is live custody: revoking the grant does not erase its + // account while it is outstanding, and its later settlement still lands + const auto free_step = ledger.admit("grant-1", 500000, 45000, false); + require(free_step.refusal == R::none, "uncontended step admitted"); + ledger.revoke("grant-1", 46000); + require(ledger.n_accounts() == 1, "a live free reservation keeps its account through revocation"); + require(ledger.admit("grant-1", 1, 47000, true).refusal == R::unknown_grant, "a revoked grant admits nothing new"); + require(ledger.start(free_step.reservation, 46500) == R::unknown_grant, "a revoked grant's reservation does not start"); + require(ledger.cancel(free_step.reservation, 3 * 1000000) == R::none, "it never ran: cancelled, custody released"); + require(ledger.n_accounts() == 0, "then the revoked account goes"); + + // admitted, then the grant expires before the step starts: refused, cancellable + require(ledger.register_grant({ "brief", "engine-a", 1000, 100000, 3500000 }, 3000000) == R::none, "a brief grant"); + const auto late = ledger.admit("brief", 1000, 3100000, true); + require(late.refusal == R::none, "admitted before expiry"); + require(ledger.start(late.reservation, 3600000) == R::unknown_grant, "expired before start: refused"); + require(ledger.cancel(late.reservation, 3600000) == R::none, "and cancelled, never erased silently"); + + // a started step still settles after its grant ends + require(ledger.register_grant({ "ending", "engine-a", 1000, 100000, 3900000 }, 3700000) == R::none, "an ending grant"); + const auto running = ledger.admit("ending", 1000, 3700000, true); + require(ledger.start(running.reservation, 3700000) == R::none, "started while live"); + require(ledger.finish(running.reservation, 1000, 1, 4000000) == R::none, "settles after the grant expired"); + + // the clock domain and the per-settle step bound are explicit, checked refusals + require(ledger.admit("grant-x", 1, -1, true).refusal == R::invalid, "a negative clock"); + require(ledger.admit("grant-x", 1, server_learn_ledger::kMaxClockUs + 1, true).refusal == R::invalid, "a clock past the domain"); + require(ledger.register_grant({ "far", "engine-a", 1000, 100000, server_learn_ledger::kMaxClockUs + 1 }, 4000000) == R::invalid, + "an expiry past the domain"); + require(ledger.register_grant({ "steps", "engine-a", 1000, 100000, day_us }, 4000000) == R::none, "steps grant"); + const auto big = ledger.admit("steps", 1, 4000000, true); + require(ledger.start(big.reservation, 4000000) == R::none, "start"); + require(ledger.finish(big.reservation, 1, INT64_MAX, 4000001) == R::invalid, "INT64_MAX steps in one settle is refused, not wrapped"); + require(ledger.finish(big.reservation, 1, 1, 4000001) == R::none, "a valid settle still lands"); + + // bounded without idleness eviction: an expired grant with nothing outstanding goes at the + // next registration, and a live idle one stays (its terms cannot be reset) + server_learn_ledger fresh("engine-a"); + require(fresh.register_grant({ "short", "engine-a", 1000, 100000, 2000000 }, 1000000) == R::none, "short grant"); + require(fresh.register_grant({ "idle", "engine-a", 1000, 100000, day_us }, 1000000) == R::none, "idle grant"); + require(fresh.register_grant({ "next", "engine-a", 1000, 100000, day_us }, 5000000) == R::none, "a later registration"); + require(fresh.n_accounts() == 2, "the expired grant is gone, the idle live one stays"); + require(fresh.register_grant({ "idle", "engine-a", 1000, 500000, day_us }, 9000000) == R::grant_mismatch, + "an idle live grant keeps its terms"); + + // an observation outside the clock domain is refused and changes nothing + { + server_learn_ledger obs("engine-a"); + require(obs.register_grant({ "o", "engine-a", 1000, 100000, day_us }, 0) == R::none, "obs grant"); + const auto o1 = obs.admit("o", 50000, 0, true); + require(obs.start(o1.reservation, 0) == R::none && obs.finish(o1.reservation, 50000, 1, 10) == R::none, "debit 50 ms"); + int64_t out = 0; + require(obs.used_us("o", INT64_MAX, out) == R::invalid, "a clock past the domain is refused"); + require(used(obs, "o", 20) == 50000, "and the period it observed is unchanged"); + require(obs.admit("o", 60000, 30, true).refusal == R::budget_exhausted, "admission still sees the spent period"); + } + + std::cout << "test-server-learn-ledger: OK\n"; + return 0; +} diff --git a/tools/server/server-learn-ledger.h b/tools/server/server-learn-ledger.h new file mode 100644 index 000000000000..fd46c1b61bee --- /dev/null +++ b/tools/server/server-learn-ledger.h @@ -0,0 +1,288 @@ +#pragma once + +// ONE learning budget for the whole server. The core's resource authority REGISTERS a grant +// { id, engine incarnation, period, share, expiry }: the share of engine time, per period, the +// learning it covers may take while another slot is serving. Every learning step (a POST /learn, +// a completion's own learn, a /train window) references a registered grant and debits its +// account. Receiving grant fields on a step authorizes nothing: an unregistered id is refused. +// +// What it guarantees (each one a check in test-server-learn-ledger): +// - registration is the authority's: an unknown grant id is refused, never minted from a step; +// - a grant from another engine incarnation (issued before a restart) is refused at registration; +// - a grant is immutable while it lives: re-registering its id with other terms is refused, with +// the same terms is a no-op; it ends only by revocation or expiry, never by idleness, and a +// changed budget is a new grant the authority issues; +// - a retry or a duplicate never adds budget: every step naming the id debits the same account; +// - every reservation, a free (uncontended) one included, is live custody until settled; +// - a reservation is cancelled (its whole cost released) only before its step started; a +// started step settles with its measured time and committed steps, after quiescence; it +// starts once, and only a started step settles; +// - a step that could never fit its grant's whole period is refused as such; +// - a step whose grant was revoked or expired after admission does not start (typed refusal); +// - arithmetic is checked: clocks live in an explicit domain, periods and costs are bounded, and +// accumulated time and steps are checked before they change: a typed refusal, never a wrap. + +#include +#include +#include +#include +#include + +struct server_learn_grant { + std::string id; + std::string engine; + int64_t period_ms = 0; + int64_t share_ppm = 0; + int64_t expires_at_us = 0; // the authority's expiry, on the engine's clock +}; + +enum class server_learn_refusal { + none, + invalid, // a field out of range (see kMaxPeriodMs, share in [1, 1e6], an empty id, a negative cost) + stale_engine, // registered for another engine incarnation + grant_mismatch, // this live id is registered with other terms + unknown_grant, // no live grant has this id (never registered, revoked or expired) + never_fits, // the step costs more than the grant's whole period allows + budget_exhausted, // the period's share is spent: retry_after_ms says when it renews + started, // cancel after the step started: it must settle with its measured time + not_started, // finish before start: a step that never started has nothing measured to settle + already_started, // start again: one reservation authorizes one execution + unknown_reservation, +}; + +struct server_learn_admission { + server_learn_refusal refusal = server_learn_refusal::none; + int64_t retry_after_ms = 0; + uint64_t reservation = 0; // nonzero when admitted +}; + +class server_learn_ledger { +public: + // The longest period a grant may name, and the largest step cost: the bounds under which every + // product and sum below stays in range (1 day of ms x 1e6 ppm, and 1 day of us). + static constexpr int64_t kMaxPeriodMs = 86400000; + static constexpr int64_t kMaxCostUs = 86400000000; + // The accepted engine clock domain (us since the engine's epoch): with periods <= 1 day, every + // sum of a clock value and a period stays below INT64_MAX. The core translates its own expiry + // into this clock; a wall-clock time interpreted as engine time is its bug, not this ledger's. + static constexpr int64_t kMaxClockUs = int64_t(1) << 62; + // One settlement commits at most this many optimizer steps. + static constexpr int64_t kMaxStepsPerSettle = int64_t(1) << 20; + + explicit server_learn_ledger(std::string engine) : engine(std::move(engine)) {} + + // The authority registers (or re-confirms) a grant. + server_learn_refusal register_grant(const server_learn_grant & g, int64_t now_us) { + if (!in_clock(now_us) || g.id.empty() || g.period_ms <= 0 || g.period_ms > kMaxPeriodMs || g.share_ppm < 1 || + g.share_ppm > 1000000 || g.expires_at_us <= now_us || !in_clock(g.expires_at_us)) { + return server_learn_refusal::invalid; + } + if (g.engine != engine) { + return server_learn_refusal::stale_engine; + } + // bounded: grants the authority let expire, whose last reservation settled, go here + for (auto a = accounts.begin(); a != accounts.end();) { + const bool done = now_us >= a->second.grant.expires_at_us && a->second.outstanding == 0; + a = done ? accounts.erase(a) : std::next(a); + } + const auto it = accounts.find(g.id); + if (it != accounts.end()) { + const auto & r = it->second.grant; + const bool same = r.period_ms == g.period_ms && r.share_ppm == g.share_ppm && r.expires_at_us == g.expires_at_us; + return same ? server_learn_refusal::none : server_learn_refusal::grant_mismatch; + } + accounts.emplace(g.id, account{ g, now_us, 0, 0, 0, false }); + return server_learn_refusal::none; + } + + // The authority ends a grant. Its account goes once its last reservation settles. + void revoke(const std::string & id, int64_t now_us) { + if (!in_clock(now_us)) { + return; + } + const auto it = accounts.find(id); + if (it != accounts.end()) { + it->second.revoked = true; + retire_if_done(it, now_us); + } + } + + // est_us: the step's expected cost (its last measured time). contended: another slot is + // serving at admission, so the step's time is debited; an uncontended step is admitted + // without a debit but is still a live reservation. + server_learn_admission admit(const std::string & grant, int64_t est_us, int64_t now_us, bool contended) { + server_learn_admission out; + if (est_us < 0 || est_us > kMaxCostUs || !in_clock(now_us)) { + out.refusal = server_learn_refusal::invalid; + return out; + } + const auto it = accounts.find(grant); + if (it == accounts.end() || it->second.revoked || now_us >= it->second.grant.expires_at_us) { + out.refusal = server_learn_refusal::unknown_grant; + return out; + } + account & a = it->second; + roll(a, now_us); + const int64_t cost = contended ? est_us : 0; + const int64_t budget = budget_us(a.grant); + if (cost > budget) { + out.refusal = server_learn_refusal::never_fits; + return out; + } + if (a.used_us + a.reserved_us + cost > budget) { + out.refusal = server_learn_refusal::budget_exhausted; + out.retry_after_ms = (a.period_start_us + a.grant.period_ms * 1000 - now_us + 999) / 1000; + return out; + } + a.reserved_us += cost; + a.outstanding += 1; + out.reservation = ++next_reservation; + reservations.emplace(out.reservation, reservation{ grant, cost, contended, false }); + return out; + } + + // The step begins executing: from here it settles only with its measured time. Refused when its + // grant was revoked or expired after admission (unknown_grant): the reservation stays held, so + // the caller cancels it (it never ran) and nothing is erased silently. + server_learn_refusal start(uint64_t id, int64_t now_us) { + if (!in_clock(now_us)) { + return server_learn_refusal::invalid; + } + const auto r = reservations.find(id); + if (r == reservations.end()) { + return server_learn_refusal::unknown_reservation; + } + if (r->second.started) { + return server_learn_refusal::already_started; + } + const auto a = accounts.find(r->second.grant); + if (a == accounts.end() || a->second.revoked || now_us >= a->second.grant.expires_at_us) { + return server_learn_refusal::unknown_grant; + } + r->second.started = true; + return server_learn_refusal::none; + } + + // The step quiesced after measured_us with committed_steps optimizer steps: debit the measured + // time (when it was admitted contended) and release the rest of its reservation. + server_learn_refusal finish(uint64_t id, int64_t measured_us, int64_t committed_steps, int64_t now_us) { + if (measured_us < 0 || measured_us > kMaxCostUs || committed_steps < 0 || committed_steps > kMaxStepsPerSettle || + !in_clock(now_us)) { + return server_learn_refusal::invalid; + } + return settle(id, measured_us, committed_steps, now_us, /*cancel =*/ false); + } + + // The step never ran: its whole reservation returns. Refused once it started. + server_learn_refusal cancel(uint64_t id, int64_t now_us) { + if (!in_clock(now_us)) { + return server_learn_refusal::invalid; + } + return settle(id, 0, 0, now_us, /*cancel =*/ true); + } + + // Debited time of a grant's current period, observed without changing anything: a clock + // outside the domain is refused, and an unknown grant reads 0. + server_learn_refusal used_us(const std::string & grant, int64_t now_us, int64_t & out) const { + if (!in_clock(now_us)) { + return server_learn_refusal::invalid; + } + const auto it = accounts.find(grant); + if (it == accounts.end()) { + out = 0; + return server_learn_refusal::none; + } + const account & a = it->second; + out = now_us >= a.period_start_us + a.grant.period_ms * 1000 ? 0 : a.used_us; + return server_learn_refusal::none; + } + + int64_t committed_steps(const std::string & grant) const { + const auto it = accounts.find(grant); + return it == accounts.end() ? 0 : it->second.steps; + } + + size_t n_accounts() const { + return accounts.size(); + } + +private: + struct account { + server_learn_grant grant; + int64_t period_start_us; + int64_t used_us; + int64_t reserved_us; + int64_t outstanding; // live reservations, free ones included + bool revoked; + int64_t steps = 0; + }; + struct reservation { + std::string grant; + int64_t cost_us; + bool contended; + bool started; + }; + + static bool in_clock(int64_t t) { + return t >= 0 && t <= kMaxClockUs; + } + + // share_ppm of period_ms, in us: in range by kMaxPeriodMs x 1e6 (8.64e13) + static int64_t budget_us(const server_learn_grant & g) { + return g.period_ms * g.share_ppm / 1000; + } + + // A new period starts where the old one ended; whole idle periods are skipped, not banked. + static void roll(account & a, int64_t now_us) { + const int64_t period_us = a.grant.period_ms * 1000; + if (now_us >= a.period_start_us + period_us) { + a.period_start_us += (now_us - a.period_start_us) / period_us * period_us; + a.used_us = 0; + } + } + + server_learn_refusal settle(uint64_t id, int64_t measured_us, int64_t committed, int64_t now_us, bool cancel) { + const auto r = reservations.find(id); + if (r == reservations.end()) { + return server_learn_refusal::unknown_reservation; + } + if (cancel && r->second.started) { + return server_learn_refusal::started; + } + if (!cancel && !r->second.started) { + return server_learn_refusal::not_started; + } + const auto a = accounts.find(r->second.grant); + if (a != accounts.end()) { + roll(a->second, now_us); + const int64_t debit = r->second.contended ? measured_us : 0; + // checked before anything changes: a sum past INT64_MAX is refused, never wrapped + if (a->second.used_us > INT64_MAX - debit || a->second.steps > INT64_MAX - committed) { + return server_learn_refusal::invalid; + } + a->second.reserved_us -= r->second.cost_us; + a->second.outstanding -= 1; + a->second.used_us += debit; + a->second.steps += committed; + } + reservations.erase(r); + if (a != accounts.end()) { + retire_if_done(a, now_us); + } + return server_learn_refusal::none; + } + + // An account ends only when the authority ended its grant (revoked or expired) and its last + // reservation settled: never by idleness, so a live grant's terms cannot be reset. + void retire_if_done(std::map::iterator it, int64_t now_us) { + const account & a = it->second; + if ((a.revoked || now_us >= a.grant.expires_at_us) && a.outstanding == 0) { + accounts.erase(it); + } + } + + std::string engine; + std::map accounts; + std::map reservations; + uint64_t next_reservation = 0; +};