337 lines
16 KiB
Rust
337 lines
16 KiB
Rust
//! The `kernel` domain end to end: a loopback node, its HTTP routes and the
|
|
//! blocking client. Tiles from the node must be the bits the same kernel
|
|
//! gives in-process, on the native backend and the interpreter.
|
|
|
|
use makepad_ai_hub::download::Downloader;
|
|
use makepad_ai_hub::kernel_client::KernelNode;
|
|
use makepad_ai_hub::kernel_compute::{compile, PreparedJob};
|
|
use makepad_ai_hub::kernel_wire::kernels::{fractal_job, fractal_tile, FractalParams};
|
|
use makepad_ai_hub::kernel_wire::*;
|
|
use makepad_ai_hub::registry::Registry;
|
|
use makepad_ai_hub::server::{start_service, ServiceConfig};
|
|
use makepad_script_compute::Backend;
|
|
use std::sync::OnceLock;
|
|
use std::time::{Duration, Instant};
|
|
|
|
/// One node for the whole test binary (the service holds a process lock).
|
|
fn node() -> KernelNode {
|
|
static URL: OnceLock<String> = OnceLock::new();
|
|
let url = URL.get_or_init(|| {
|
|
let dir = std::env::temp_dir().join(format!("makepad-ai-hub-kernel-{}", std::process::id()));
|
|
let _ = std::fs::create_dir_all(&dir);
|
|
let handle = start_service(ServiceConfig {
|
|
host: "127.0.0.1".to_string(),
|
|
port: 0,
|
|
cache_dir: dir,
|
|
registry: Registry::embedded().unwrap(),
|
|
downloader: Downloader::new("http://127.0.0.1:1", None).unwrap(),
|
|
peer: makepad_ai_hub::peer_serve::PeerOptions { serve: Some(false), sources: Some(Vec::new()), ..Default::default() },
|
|
fleet: makepad_ai_hub::discovery::DEFAULT_FLEET.to_string(),
|
|
})
|
|
.unwrap();
|
|
let url = format!("http://{}", handle.addr);
|
|
std::mem::forget(handle);
|
|
url
|
|
});
|
|
KernelNode::new(url)
|
|
}
|
|
|
|
fn local(job: &KernelJobJson, backend: Backend) -> Vec<TileResult> {
|
|
let kernel = compile(&job.source, backend).unwrap();
|
|
let prepared = PreparedJob::new(kernel, job).unwrap();
|
|
(0..prepared.tile_count()).map(|k| prepared.run_tile(k, backend)).collect()
|
|
}
|
|
|
|
fn by_id(mut tiles: Vec<TileResult>) -> Vec<TileResult> {
|
|
tiles.sort_by_key(|t| t.id);
|
|
tiles
|
|
}
|
|
|
|
fn same_bits(a: &[TileResult], b: &[TileResult]) {
|
|
assert_eq!(a.len(), b.len());
|
|
for (x, y) in a.iter().zip(b) {
|
|
assert_eq!(x.id, y.id);
|
|
assert_eq!(x.status, STATUS_OK, "{}", x.message);
|
|
assert_eq!(y.status, STATUS_OK, "{}", y.message);
|
|
assert_eq!(x.work, y.work, "tile {}: counted work", x.id);
|
|
assert_eq!(x.outputs, y.outputs, "tile {}", x.id);
|
|
assert_eq!(x.reduced.iter().map(|r| r.to_bits()).collect::<Vec<_>>(), y.reduced.iter().map(|r| r.to_bits()).collect::<Vec<_>>());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn node_fractal_tiles_are_the_local_bits_for_every_kind() {
|
|
let node = node();
|
|
let info = node.info().unwrap();
|
|
assert!(info.cores >= 1);
|
|
for (kind, de, trap) in (0..5).flat_map(|k| [(k, false, false), (k, true, false), (k, false, true), (k, true, true)]).chain([(5, false, false)]) {
|
|
let p = FractalParams { kind, max_iter: 300, de, trap, ..FractalParams::default() };
|
|
let tiles = (0..6u64).map(|k| fractal_tile(k, -2.0 + 0.6 * k as f64, -0.4, 0.8, 32, &p)).collect();
|
|
let job = fractal_job(&p, tiles);
|
|
let remote = by_id(node.run(&job, 32 * 32).unwrap());
|
|
let native = local(&job, Backend::Native);
|
|
let interp = local(&job, Backend::Interp);
|
|
same_bits(&remote, &native);
|
|
same_bits(&native, &interp);
|
|
// a texel inside the set and one outside
|
|
if kind == 5 {
|
|
// near 1 the cubic converges to root 0 (class 1); some texel reaches each root
|
|
let t = remote.iter().find(|t| t.id == 3).unwrap().output("out").unwrap();
|
|
let classes: std::collections::HashSet<u32> = t.iter().map(|w| w >> 24).collect();
|
|
assert!((1..=3).all(|c| classes.contains(&c)), "{classes:?}");
|
|
}
|
|
if kind == 0 && de {
|
|
// escaped texels carry distance codes (top byte 1..254)
|
|
let t = remote.iter().find(|t| t.id == 3).unwrap().output("out").unwrap();
|
|
assert!(t.iter().any(|w| (1..255).contains(&(w >> 24))));
|
|
}
|
|
if kind == 0 && !de && !trap {
|
|
let t0 = remote.iter().find(|t| t.id == 3).unwrap().output("out").unwrap().to_vec();
|
|
assert!(t0.iter().any(|w| w >> 24 == 255) && t0.iter().any(|w| w >> 24 == 0));
|
|
}
|
|
}
|
|
}
|
|
|
|
/// A second, unrelated kernel: a map over an f32 input with a per-tile
|
|
/// param and the seed, and a reduce. Nothing in the node knows fractals.
|
|
#[test]
|
|
fn any_kernel_runs_maps_and_reduces() {
|
|
let node = node();
|
|
let src: Vec<f32> = (0..1000).map(|i| i as f32 * 0.5).collect();
|
|
let map = KernelJobJson {
|
|
source: "let src = input(f32)\nlet out = output(vec2)\nlet scale = param(1.0, -100, 100)\nfn element(i) { out[i] = vec2(src[i] * scale, float(i + seed)) }".into(),
|
|
params: vec![KernelParamJson { name: "scale".into(), value: 2.0 }],
|
|
inputs: vec![KernelBufferJson::from_f32("src", &src)],
|
|
tiles: vec![
|
|
KernelTileJson { id: 7, count: 1000, params: vec![], inputs: vec![] },
|
|
KernelTileJson { id: 8, count: 10, params: vec![KernelParamJson { name: "scale".into(), value: -1.0 }], inputs: vec![] },
|
|
],
|
|
encoding: "raw".into(),
|
|
seed: 5,
|
|
passes: None,
|
|
feedback: None,
|
|
};
|
|
let remote = by_id(node.run(&map, 2000).unwrap());
|
|
same_bits(&remote, &local(&map, Backend::Native));
|
|
let a = remote[0].output("out").unwrap();
|
|
assert_eq!(a.len(), 2000);
|
|
assert_eq!((f32::from_bits(a[2 * 10]), f32::from_bits(a[2 * 10 + 1])), (10.0, 15.0));
|
|
let b = remote[1].output("out").unwrap();
|
|
assert_eq!(b.len(), 20);
|
|
assert_eq!(f32::from_bits(b[2 * 3]), -1.5);
|
|
|
|
let reduce = KernelJobJson {
|
|
source: "let v = input(f32)\nfn reduce_sum(i) { v[i] * v[i] }".into(),
|
|
params: vec![],
|
|
inputs: vec![],
|
|
tiles: (0..3u64)
|
|
.map(|k| KernelTileJson { id: k, count: 100, params: vec![], inputs: vec![KernelBufferJson::from_f32("v", &vec![k as f32 + 1.0; 100])] })
|
|
.collect(),
|
|
encoding: "delta-deflate".into(),
|
|
seed: 0,
|
|
passes: None,
|
|
feedback: None,
|
|
};
|
|
let remote = by_id(node.run(&reduce, 0).unwrap());
|
|
same_bits(&remote, &local(&reduce, Backend::Native));
|
|
assert_eq!(remote.iter().map(|t| t.reduced[0]).collect::<Vec<_>>(), vec![100.0, 400.0, 900.0]);
|
|
}
|
|
|
|
#[test]
|
|
fn cancel_drops_queued_tiles_and_stops_running_ones() {
|
|
let node = node();
|
|
let cores = node.info().unwrap().cores as usize;
|
|
// every texel inside the set at the cap: the slowest tiles there are
|
|
let p = FractalParams { kind: 0, max_iter: 16384, ..FractalParams::default() };
|
|
let n = (cores * 4).max(8);
|
|
let tiles = (0..n as u64).map(|k| fractal_tile(k, -1.37, -0.02, 0.04, 128, &p)).collect();
|
|
let job = fractal_job(&p, tiles);
|
|
let submitted = node.submit(&job).unwrap();
|
|
std::thread::sleep(Duration::from_millis(100));
|
|
let t0 = Instant::now();
|
|
node.cancel(submitted.job_id).unwrap();
|
|
let mut got = Vec::new();
|
|
loop {
|
|
let batch = node.poll(submitted.job_id, Duration::from_millis(2000), 128 * 128).unwrap();
|
|
got.extend(batch.tiles);
|
|
if batch.done {
|
|
break;
|
|
}
|
|
assert!(t0.elapsed() < Duration::from_secs(5), "a cancelled job must end promptly");
|
|
}
|
|
let stopped = t0.elapsed();
|
|
let ok = got.iter().filter(|t| t.status == STATUS_OK).count();
|
|
let cancelled = got.iter().filter(|t| t.status == STATUS_CANCELLED).count();
|
|
eprintln!("cancel: {n} tiles, {ok} finished, {cancelled} stopped while running, {} never ran; ended {stopped:?} after the cancel", n - got.len());
|
|
assert!(got.len() < n, "queued tiles must not run after a cancel");
|
|
assert!(stopped < Duration::from_secs(2));
|
|
// the job is gone
|
|
assert!(node.poll(submitted.job_id, Duration::ZERO, 16).is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn bad_jobs_are_refused_with_the_reason() {
|
|
let node = node();
|
|
let mut job = fractal_job(&FractalParams::default(), vec![fractal_tile(0, 0.0, 0.0, 1.0, 8, &FractalParams::default())]);
|
|
job.source = "let out = output(u32)\nfn element(i) { out[i] = nope }".into();
|
|
let e = node.submit(&job).unwrap_err().to_string();
|
|
assert!(e.contains("does not compile") && e.contains("2:"), "{e}");
|
|
|
|
let mut job = fractal_job(&FractalParams::default(), vec![fractal_tile(0, 0.0, 0.0, 1.0, 8, &FractalParams::default())]);
|
|
job.params.push(KernelParamJson { name: "nope".into(), value: 1.0 });
|
|
assert!(node.submit(&job).unwrap_err().to_string().contains("no param `nope`"));
|
|
|
|
let mut job = fractal_job(&FractalParams::default(), vec![fractal_tile(0, 0.0, 0.0, 1.0, 8, &FractalParams::default())]);
|
|
job.tiles[0].inputs.clear();
|
|
assert!(node.submit(&job).unwrap_err().to_string().contains("input `view` is not bound"));
|
|
|
|
let emit = KernelJobJson {
|
|
source: "let seg = emit_buffer(2, 4)\nfn primitive(i) { emit(seg, vec2(1.0, 2.0)) }".into(),
|
|
params: vec![],
|
|
inputs: vec![],
|
|
tiles: vec![KernelTileJson { id: 0, count: 4, params: vec![], inputs: vec![] }],
|
|
encoding: "raw".into(),
|
|
seed: 0,
|
|
passes: None,
|
|
feedback: None,
|
|
};
|
|
let e = node.submit(&emit).unwrap_err().to_string();
|
|
assert!(e.contains("not served") || e.contains("does not compile"), "{e}");
|
|
|
|
let mut job = fractal_job(&FractalParams::default(), vec![fractal_tile(0, 0.0, 0.0, 1.0, 8, &FractalParams::default())]);
|
|
job.tiles[0].count = (MAX_TILE_ELEMENTS + 1) as u32;
|
|
assert!(node.submit(&job).unwrap_err().to_string().contains("elements"));
|
|
}
|
|
|
|
/// The stream (one websocket per node): the same tiles as by request,
|
|
/// pushed as they finish, and a cancel up the same stream.
|
|
#[test]
|
|
fn the_stream_pushes_the_same_tiles_and_cancels() {
|
|
use makepad_ai_hub::kernel_client::KernelStream;
|
|
let node = node();
|
|
let stream = KernelStream::open(&node).unwrap();
|
|
let p = FractalParams { kind: 2, max_iter: 400, ..FractalParams::default() };
|
|
let job = fractal_job(&p, (0..8u64).map(|k| fractal_tile(k, -2.0 + 0.4 * k as f64, -1.0, 0.5, 48, &p)).collect());
|
|
let submitted = stream.submit(&job).unwrap();
|
|
let mut got = Vec::new();
|
|
loop {
|
|
let batch = stream.next_batch(submitted.job_id, Duration::from_secs(5)).unwrap();
|
|
got.extend(batch.tiles);
|
|
if batch.done {
|
|
break;
|
|
}
|
|
}
|
|
same_bits(&by_id(got), &by_id(node.run(&job, 48 * 48).unwrap()));
|
|
|
|
let heavy = FractalParams { kind: 0, max_iter: 16384, ..FractalParams::default() };
|
|
let job = fractal_job(&heavy, (0..32u64).map(|k| fractal_tile(k, -1.37, -0.02, 0.04, 128, &heavy)).collect());
|
|
let submitted = stream.submit(&job).unwrap();
|
|
std::thread::sleep(Duration::from_millis(50));
|
|
let t0 = Instant::now();
|
|
stream.cancel(submitted.job_id);
|
|
let mut n = 0;
|
|
loop {
|
|
let batch = stream.next_batch(submitted.job_id, Duration::from_secs(5)).unwrap();
|
|
n += batch.tiles.len();
|
|
if batch.done {
|
|
break;
|
|
}
|
|
assert!(t0.elapsed() < Duration::from_secs(3), "a cancelled stream job must end promptly");
|
|
}
|
|
assert!(n < 32);
|
|
// a bad job is refused through the stream with the reason
|
|
let mut bad = job.clone();
|
|
bad.source = "nope".into();
|
|
assert!(stream.submit(&bad).unwrap_err().to_string().contains("does not compile"));
|
|
}
|
|
|
|
/// The fleet over two "nodes" (the loopback node twice): every tag comes
|
|
/// back exactly once, with the local bits, even when the frame's tail is
|
|
/// copied to the idle node.
|
|
#[test]
|
|
fn the_fleet_answers_every_tile_once_with_the_local_bits() {
|
|
use makepad_ai_hub::kernel_client::{Batch, KernelFleet};
|
|
use std::collections::{HashMap, VecDeque};
|
|
use std::sync::{Arc, Mutex};
|
|
let node = node();
|
|
let info = node.info().unwrap();
|
|
let p = FractalParams { kind: 0, max_iter: 3000, ..FractalParams::default() };
|
|
// cheap tiles and a few dear ones (inside the set at the cap)
|
|
let tiles: Vec<KernelTileJson> = (0..48u64)
|
|
.map(|k| if k % 12 == 11 { fractal_tile(k, -0.3, -0.05, 0.1, 64, &p) } else { fractal_tile(k, 0.3 + 0.01 * k as f64, 0.5, 0.05, 64, &p) })
|
|
.collect();
|
|
let want: HashMap<u64, Vec<u32>> = local(&fractal_job(&p, tiles.clone()), Backend::Native).into_iter().map(|t| (t.id, t.output("out").unwrap().to_vec())).collect();
|
|
let queue = Arc::new(Mutex::new(tiles.into_iter().collect::<VecDeque<_>>()));
|
|
let got: Arc<Mutex<Vec<(u64, TileResult)>>> = Arc::new(Mutex::new(Vec::new()));
|
|
let source = {
|
|
let queue = queue.clone();
|
|
Arc::new(move |max: usize| {
|
|
let mut q = queue.lock().unwrap();
|
|
let take: Vec<KernelTileJson> = (0..max.min(q.len())).filter_map(|_| q.pop_front()).collect();
|
|
if take.is_empty() {
|
|
return None;
|
|
}
|
|
let tags = take.iter().map(|t| t.id).collect();
|
|
Some(Batch { job: fractal_job(&p, take), tags })
|
|
})
|
|
};
|
|
let sink = {
|
|
let got = got.clone();
|
|
Arc::new(move |_node: usize, tag: u64, r: TileResult| got.lock().unwrap().push((tag, r)))
|
|
};
|
|
let fleet = KernelFleet::start(vec![(node.clone(), info.clone()), (node.clone(), info)], source, sink);
|
|
let t0 = Instant::now();
|
|
while got.lock().unwrap().len() < 48 {
|
|
assert!(t0.elapsed() < Duration::from_secs(60), "the fleet must answer every tile");
|
|
std::thread::sleep(Duration::from_millis(5));
|
|
}
|
|
std::thread::sleep(Duration::from_millis(300));
|
|
drop(fleet);
|
|
let got = got.lock().unwrap();
|
|
assert_eq!(got.len(), 48, "exactly once");
|
|
for (tag, r) in got.iter() {
|
|
assert_eq!(r.status, STATUS_OK, "{}", r.message);
|
|
assert_eq!(r.output("out").unwrap(), want[tag].as_slice(), "tile {tag}");
|
|
}
|
|
}
|
|
|
|
/// Long iteration counts run in passes with the orbit state fed back on the
|
|
/// node: a 2-pass long render equals the one-call render at the same
|
|
/// iteration count, and counts past one call's bound render (and take longer).
|
|
#[test]
|
|
fn long_iteration_counts_run_in_passes_with_the_same_bits() {
|
|
use makepad_ai_hub::kernel_wire::kernels::FRACTAL_LONG;
|
|
let node = node();
|
|
for (kind, de, trap) in [(0, false, false), (0, true, false), (1, false, true), (2, true, true), (3, false, false)] {
|
|
let p = FractalParams { kind, max_iter: 3000, de, trap, ..FractalParams::default() };
|
|
let tiles: Vec<KernelTileJson> = (0..4u64).map(|k| fractal_tile(k, -1.9 + 0.5 * k as f64, -0.3, 0.6, 32, &p)).collect();
|
|
let one = fractal_job(&p, tiles.clone());
|
|
// the same count through the long kernel: 2 passes
|
|
let mut long = fractal_job(&FractalParams { max_iter: 20000, ..p }, tiles);
|
|
assert_eq!(long.source, FRACTAL_LONG);
|
|
long.params.retain(|q| q.name != "max_iter" && q.name != "passes");
|
|
long.params.push(KernelParamJson { name: "max_iter".into(), value: 3000.0 });
|
|
long.params.push(KernelParamJson { name: "passes".into(), value: 2.0 });
|
|
long.passes = Some(2);
|
|
let a = by_id(node.run(&one, 32 * 32).unwrap());
|
|
let b = by_id(node.run(&long, 32 * 32).unwrap());
|
|
same_bits_outputs(&a, &b);
|
|
same_bits(&b, &local(&long, Backend::Native));
|
|
assert!(b.iter().all(|t| t.outputs.len() == 1), "fed-back state stays on the node");
|
|
}
|
|
// past the one-call cap: 40000 iterations inside the set
|
|
let p = FractalParams { kind: 0, max_iter: 40000, ..FractalParams::default() };
|
|
let job = fractal_job(&p, vec![fractal_tile(0, -0.3, -0.05, 0.1, 16, &p)]);
|
|
assert_eq!(job.passes, Some(20));
|
|
let t = node.run(&job, 16 * 16).unwrap();
|
|
assert_eq!(t[0].status, STATUS_OK, "{}", t[0].message);
|
|
assert!(t[0].output("out").unwrap().iter().all(|w| w >> 24 == 255));
|
|
}
|
|
|
|
fn same_bits_outputs(a: &[TileResult], b: &[TileResult]) {
|
|
assert_eq!(a.len(), b.len());
|
|
for (x, y) in a.iter().zip(b) {
|
|
assert_eq!((x.status, y.status), (STATUS_OK, STATUS_OK), "{} {}", x.message, y.message);
|
|
assert_eq!(x.output("out"), y.output("out"), "tile {}", x.id);
|
|
}
|
|
}
|