62d4f31b1e
Every heap ref carries its heap tag, generation checks always on, rooted async values. Splash isolates get a domain table and curated minimal preludes; system modules under `mod` expand lazily on first touch. No wall-clock script caps: isolation and force-stop instead. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
196 lines
10 KiB
Rust
196 lines
10 KiB
Rust
//! Admission and counted work for untrusted origins, in ops (the units a
|
|
//! kernel counts while it runs): the worst case of a job is checked before
|
|
//! it runs (the element bound, the untrusted ceiling, the device ledger),
|
|
//! a running job stops once its counted ops pass its budget (within one
|
|
//! slice), whatever the machine, backend or thread count, and a cancel
|
|
//! lands at the next element.
|
|
|
|
use makepad_script_compute_sched::kernel::KernelExt;
|
|
use makepad_script_compute_sched::admission::{DeviceLimits, JobBudget, Ledger, Origin, Refused};
|
|
use makepad_script_compute::cost::{element_ops};
|
|
use makepad_script_compute::kernel::{compile_with, Kernel, KernelRef};
|
|
use makepad_script_compute_sched::kernel::{KernelError};
|
|
use makepad_script_compute::Backend;
|
|
|
|
|
|
|
|
fn kernel(body: &str, backend: Backend) -> KernelRef {
|
|
let src = format!("let src = input(f32)\nlet o = output(f32)\nfn element(i) {{\n {}\n }}", body);
|
|
compile_with(&src, &[], backend).unwrap_or_else(|e| panic!("{:?}", e))
|
|
}
|
|
|
|
fn budget(ops: u64) -> JobBudget {
|
|
JobBudget { work: ops }
|
|
}
|
|
|
|
/// Runs `k` over `n` elements under `limit` counted ops: the ops it counted,
|
|
/// or the error.
|
|
fn counted(k: &Kernel, n: usize, limit: u64, threads: usize, interp: bool, simd: bool) -> Result<u64, KernelError> {
|
|
let src = vec![0.0f32; n];
|
|
let mut out = vec![0.0f32; n];
|
|
let mut c = k.call();
|
|
c.input("src", &src).unwrap();
|
|
c.output("o", &mut out).unwrap();
|
|
c.set_work_limit(limit);
|
|
c.set_simd(simd);
|
|
let r = if interp { c.run_interp(n) } else if threads > 1 { c.run_parallel(n, threads) } else { c.run(n) };
|
|
r.map(|stats| stats.work)
|
|
}
|
|
|
|
#[test]
|
|
fn counted_work_is_the_same_on_every_backend_and_thread_count() {
|
|
for body in [
|
|
"let s = float(i)\n for a in 0..1000 { s = sin(s) }\n o[i] = s",
|
|
// A data-dependent exit: the count follows what ran, not the cap.
|
|
"let s = float(i)\n for a in 0..1000 { s = s + 1.0\n if s > float(i % 97) + 40.0 { break } }\n o[i] = s",
|
|
] {
|
|
let k = kernel(body, Backend::Native);
|
|
let ki = kernel(body, Backend::Interp);
|
|
let n = 10_001;
|
|
let base = counted(&k, n, u64::MAX, 1, false, false).unwrap();
|
|
assert!(base > 0 && base <= element_ops(&k) * n as u64, "{base} within the worst case");
|
|
for (threads, interp, simd) in [(1, false, true), (8, false, true), (8, false, false), (1, true, false)] {
|
|
let kk = if interp { &ki } else { &k };
|
|
assert_eq!(counted(kk, n, u64::MAX, threads, interp, simd).unwrap(), base, "{body}: threads {threads} interp {interp} simd {simd}");
|
|
}
|
|
// Exactly its count runs; one op less stops it, on every backend.
|
|
for (threads, interp) in [(1, false), (8, false), (1, true)] {
|
|
let kk = if interp { &ki } else { &k };
|
|
assert_eq!(counted(kk, n, base, threads, interp, true).unwrap(), base);
|
|
assert!(matches!(counted(kk, n, base - 1, threads, interp, true), Err(KernelError::OverBudget { limit, .. }) if limit == base - 1));
|
|
}
|
|
}
|
|
// The early exit counts far less than its static worst case.
|
|
let early = kernel("let s = 0.0\n for a in 0..1000 { s = s + 1.0\n if s > 3.0 { break } }\n o[i] = s", Backend::Native);
|
|
let full = kernel("let s = 0.0\n for a in 0..1000 { s = s + 1.0 }\n o[i] = s", Backend::Native);
|
|
let (e, f) = (counted(&early, 100, u64::MAX, 1, false, true).unwrap(), counted(&full, 100, u64::MAX, 1, false, true).unwrap());
|
|
assert!(e * 20 < f, "early {e} vs full {f}");
|
|
}
|
|
|
|
#[test]
|
|
fn untrusted_elements_are_bounded_host_elements_are_not() {
|
|
let ledger = Ledger::new(DeviceLimits { element_work: 100_000, ..DeviceLimits::default() });
|
|
let heavy = kernel("let s = i\n for a in 0..20000 { s = int(src[s]) }\n o[i] = float(s)", Backend::Native);
|
|
assert!(element_ops(&heavy) > 100_000);
|
|
for origin in [Origin::Ai, Origin::Store, Origin::Lan, Origin::Livecode] {
|
|
match ledger.admit(&heavy, 10, origin, &budget(1 << 30)) {
|
|
Err(Refused::ElementTooHeavy { worst_ops, limit_ops }) => assert!(worst_ops > limit_ops),
|
|
other => panic!("{:?}: expected ElementTooHeavy, got {:?}", origin, other.map(|t| t.estimate())),
|
|
}
|
|
}
|
|
assert!(ledger.admit(&heavy, 10, Origin::Host, &budget(1 << 30)).is_ok());
|
|
assert_eq!(ledger.stats().jobs, 0, "tickets return their share on drop");
|
|
}
|
|
|
|
#[test]
|
|
fn job_budgets_and_untrusted_ceilings() {
|
|
let limits = DeviceLimits { untrusted_job_work: 1_000_000, untrusted_worst_case: 1 << 30, ..DeviceLimits::default() };
|
|
let ledger = Ledger::new(limits);
|
|
let k = kernel("let s = float(i)\n for a in 0..1000 { s = sin(s) }\n o[i] = s", Backend::Native);
|
|
let per = element_ops(&k);
|
|
// An untrusted job's budget is capped; a trusted one's is its own.
|
|
let t = ledger.admit(&k, 10, Origin::Ai, &budget(1 << 40)).unwrap();
|
|
assert_eq!(t.work_limit(), 1_000_000);
|
|
assert_eq!(t.estimate().charge_ops, (per * 10).min(1_000_000));
|
|
let t = ledger.admit(&k, 10, Origin::Host, &budget(1 << 40)).unwrap();
|
|
assert_eq!(t.work_limit(), 1 << 40);
|
|
// The worst case only refuses past the ceiling.
|
|
let fits = ((1u64 << 30) / per) as usize;
|
|
assert!(ledger.admit(&k, fits, Origin::Ai, &budget(1000)).is_ok());
|
|
assert!(matches!(ledger.admit(&k, fits + 1, Origin::Ai, &budget(1000)), Err(Refused::OverCeiling { .. })));
|
|
assert!(ledger.admit(&k, fits + 1, Origin::Host, &budget(1000)).is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn the_device_ledger_sums_every_live_job() {
|
|
let limits = DeviceLimits { max_in_flight: 100_000, max_untrusted_in_flight: 50_000, max_jobs: 8, max_untrusted_jobs: 3, ..DeviceLimits::default() };
|
|
let ledger = Ledger::new(limits);
|
|
let k = kernel("let s = float(i)\n for a in 0..1000 { s = sin(s) }\n o[i] = s", Backend::Native);
|
|
// Each job is charged what it may run: here its budget.
|
|
let a = ledger.admit(&k, 1000, Origin::Ai, &budget(20_000)).unwrap();
|
|
let b = ledger.admit(&k, 1000, Origin::Store, &budget(20_000)).unwrap();
|
|
assert!(matches!(ledger.admit(&k, 1000, Origin::Lan, &budget(20_000)), Err(Refused::DeviceBusy { .. })));
|
|
// Host work uses the rest of the device, not the untrusted share.
|
|
let c = ledger.admit(&k, 1000, Origin::Host, &budget(20_000)).unwrap();
|
|
let d = ledger.admit(&k, 1000, Origin::Host, &budget(20_000)).unwrap();
|
|
assert!(matches!(ledger.admit(&k, 1000, Origin::Host, &budget(40_000)), Err(Refused::DeviceBusy { .. })));
|
|
assert_eq!(ledger.stats().jobs, 4);
|
|
drop(a);
|
|
assert!(ledger.admit(&k, 1000, Origin::Lan, &budget(20_000)).is_ok());
|
|
drop((b, c, d));
|
|
assert_eq!(ledger.stats(), Default::default());
|
|
// Job-count limits.
|
|
let small: Vec<_> = (0..3).map(|_| ledger.admit(&k, 1, Origin::Ai, &budget(100)).unwrap()).collect();
|
|
assert!(matches!(ledger.admit(&k, 1, Origin::Ai, &budget(100)), Err(Refused::TooManyJobs { .. })));
|
|
drop(small);
|
|
}
|
|
|
|
/// The element a test's cancel lands in, and the call to cancel.
|
|
static CANCEL_AT: std::sync::Mutex<Option<(u32, makepad_script_compute_sched::kernel::CancelToken)>> = std::sync::Mutex::new(None);
|
|
|
|
/// `test.cancel_at(i)`: cancels the running call when its element `i` is
|
|
/// the chosen one; returns `i + 1`.
|
|
unsafe extern "C-unwind" fn cancel_at(context: *mut (), frame: &mut makepad_script_compute::host::HostCallFrame, _errors: makepad_script_compute::host::ErrorSink<'_>) -> bool {
|
|
let (args, _s, rets) = frame.parts();
|
|
let state = &*context.cast::<std::sync::Mutex<Option<(u32, makepad_script_compute_sched::kernel::CancelToken)>>>();
|
|
if let Some((at, token)) = state.lock().unwrap().as_ref() {
|
|
if args[0] == *at { token.cancel(); }
|
|
}
|
|
rets[0] = args[0] + 1;
|
|
true
|
|
}
|
|
unsafe extern "C-unwind" fn cancel_cost(_: *mut (), _: *const u32, _: usize) -> u64 { 1 }
|
|
|
|
/// A cancel lands at the next element: what runs after it is at most the
|
|
/// element it came in, whose worst case admission bounds in ops (see
|
|
/// `untrusted_elements_are_bounded_host_elements_are_not`). The same on
|
|
/// every machine, at any load: the element that cancels is the last one
|
|
/// written, native and interpreted, whatever the element's own work.
|
|
#[test]
|
|
fn a_cancel_stops_the_call_at_the_next_element() {
|
|
use makepad_script_compute::host::{self, HostFn, Tier};
|
|
use makepad_script_compute::ir::Ty;
|
|
let f = host::find("test.cancel_at").unwrap_or_else(|| {
|
|
host::register(HostFn { name: "test.cancel_at", params: &[Ty::I32], slices: &[], rets: &[Ty::I32], callbacks: unsafe { host::HostCallbacks::new((&CANCEL_AT as *const _ as *mut ()).cast(), cancel_cost, cancel_at) }, misses: 0, tier: Tier::D, doc: "" }).into_result().unwrap()
|
|
});
|
|
let _ = f;
|
|
let src = "let o = output(i32)\nfn element(i) { let s = float(i)\n for a in 0..2000 { s = sin(s) + 0.5 }\n o[i] = test.cancel_at(i) + int(s) * 0 }";
|
|
for backend in [Backend::Native, Backend::Interp] {
|
|
let k = compile_with(src, &[], backend).unwrap_or_else(|e| panic!("{:?}", e));
|
|
for at in [0u32, 7, 4095, 4096, 9000] {
|
|
let n = 20_000;
|
|
let mut out = vec![0u32; n];
|
|
let mut c = k.call();
|
|
c.output_u32("o", &mut out).unwrap();
|
|
*CANCEL_AT.lock().unwrap() = Some((at, c.cancel_token()));
|
|
let r = if backend == Backend::Interp { c.run_interp(n) } else { c.run(n) };
|
|
*CANCEL_AT.lock().unwrap() = None;
|
|
drop(c);
|
|
assert!(matches!(r, Err(KernelError::Cancelled)), "{backend:?} at {at}: {r:?}");
|
|
let last = out.iter().rposition(|w| *w != 0).map(|k| k as u32);
|
|
assert_eq!(last, Some(at), "{backend:?}: the cancel came in element {at}");
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A call stopped by its count stops close to its limit (in ops, the
|
|
/// same on every machine): at most four elements a worker past it.
|
|
#[test]
|
|
fn a_count_stops_the_call_within_four_elements_a_worker_of_its_limit() {
|
|
let k = kernel("let s = float(i)\n for a in 0..2000 { s = sin(s) + 0.5 }\n o[i] = s", Backend::Native);
|
|
let per = element_ops(&k);
|
|
let ran = counted(&k, 1000, u64::MAX, 1, false, true).unwrap();
|
|
assert!(ran <= per * 1000, "an element counts {} of its worst case {per}", ran / 1000);
|
|
for threads in [1usize, 8] {
|
|
for limit in [per * 10, per * 5000 + 17, per * 40_000] {
|
|
match counted(&k, 100_000, limit, threads, false, true) {
|
|
Err(KernelError::OverBudget { work, limit: l }) => {
|
|
assert_eq!(l, limit);
|
|
let most = 4 * threads as u64;
|
|
assert!(work > limit && work - limit <= most * per, "{threads} threads: {work} counted past {limit}");
|
|
}
|
|
other => panic!("{other:?}"),
|
|
}
|
|
}
|
|
}
|
|
}
|