#![cfg(target_os = "macos")]
#![allow(clippy::expect_used, clippy::indexing_slicing)]
use std::os::unix::ffi::OsStringExt;
use std::path::{Path, PathBuf};
use std::process::Command;
use tempfile::{TempDir, tempdir};
use zeppelin_embed_text::runtime::coreml::{ComputeUnits, CoreMlComputePlan, CoreMlRuntime};
use zeppelin_embed_text::runtime::{ModelRuntime, RuntimeError};
use zeppelin_embed_text::tower::TokenBatch;
mod common;
use common::compute_plan_available;
#[test]
fn coreml_open_does_not_load_compute_plan() {
if !compute_plan_available() {
return;
}
let (_root, model) = compile_fixture();
let (_harness_root, executable) = common::compile_coreml_plan_contract();
let output = Command::new(executable)
.arg("open")
.arg(model)
.output()
.expect("native model open without plan");
assert!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
}
#[test]
fn coreml_lazy_plan_retries_failure_and_caches_success() {
if !compute_plan_available() {
return;
}
let (_root, model) = compile_fixture();
let hidden = model.with_extension("hidden");
let plan = CoreMlComputePlan::new(&model, ComputeUnits::Cpu);
std::fs::rename(&model, &hidden).expect("hide compiled model");
assert!(matches!(
plan.observed_compute_units(),
Err(RuntimeError::Mlx(_))
));
std::fs::rename(&hidden, &model).expect("restore model so retry can succeed");
let observed = plan
.observed_compute_units()
.expect("retry failed request")
.expect("available compute plan");
assert_eq!(observed.cpu_ops, 1);
std::fs::rename(&model, &hidden).expect("make reloading fail");
assert_eq!(
plan.observed_compute_units().expect("cached success"),
Some(observed)
);
}
#[test]
fn coreml_unknown_preferred_device_counts_as_other() {
if !compute_plan_available() {
return;
}
run_plan_contract("unknown");
}
#[test]
fn coreml_plan_timeout_and_load_errors_fail_loudly() {
if !compute_plan_available() {
return;
}
run_plan_contract("timeout");
run_plan_contract("failure");
}
fn run_plan_contract(mode: &str) {
let (_root, executable) = common::compile_coreml_plan_contract();
let output = Command::new(executable)
.arg(mode)
.output()
.expect("native shim contract");
assert!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
}
fn compile_fixture() -> (TempDir, PathBuf) {
let root = tempdir().expect("CoreML compiler directory");
let source = root.path().join("tiny.mlmodel");
std::fs::write(&source, include_bytes!("fixtures/ze115/tiny.mlmodel"))
.expect("write portable model");
let output = Command::new("xcrun")
.args(["coremlcompiler", "compile"])
.arg(&source)
.arg(root.path())
.output()
.expect("run platform CoreML compiler");
assert!(
output.status.success(),
"CoreML compilation failed: {}",
String::from_utf8_lossy(&output.stderr)
);
let model = root.path().join("tiny.mlmodelc");
assert!(model.is_dir());
(root, model)
}
#[test]
fn coreml_compute_plan_reports_one_fixture_op_under_every_requested_policy() {
let (_root, model) = compile_fixture();
for policy in [
ComputeUnits::Cpu,
ComputeUnits::CpuAndGpu,
ComputeUnits::CpuAndNeuralEngine,
ComputeUnits::All,
] {
let plan = CoreMlComputePlan::new(&model, policy);
let observed = plan.observed_compute_units().expect("lazy compute plan");
if compute_plan_available() {
let observed = observed.expect("compute plan on macOS 14.4+");
assert_eq!(
observed.cpu_ops
+ observed.gpu_ops
+ observed.neural_engine_ops
+ observed.other_ops,
1
);
if policy == ComputeUnits::Cpu {
assert_eq!(observed.cpu_ops, 1);
assert_eq!(observed.gpu_ops, 0);
assert_eq!(observed.neural_engine_ops, 0);
assert_eq!(observed.other_ops, 0);
}
} else {
assert_eq!(observed, None);
}
}
}
#[test]
fn coreml_runtime_evaluates_exact_rows_and_masks_under_every_requested_policy() {
let (_root, model) = compile_fixture();
let tokens = TokenBatch::new(
vec![7, 11, 3, 9, -2, 5],
vec![1., 0., 0.01, -1., f32::NAN, f32::INFINITY],
3,
2,
)
.expect("three mask cases");
for (policy, name, gpu) in [
(ComputeUnits::Cpu, "coreml-cpu", false),
(ComputeUnits::CpuAndGpu, "coreml-cpu-gpu", true),
(ComputeUnits::CpuAndNeuralEngine, "coreml-cpu-ane", false),
(ComputeUnits::All, "coreml-all", true),
] {
let mut runtime = CoreMlRuntime::load(&model, 2, policy).expect("actual compiled model");
assert_eq!(runtime.sequence(), 2);
assert_eq!(runtime.identity().name, name);
assert_eq!(runtime.identity().gpu, gpu);
runtime.warm().expect("real one-row warmup");
let output = runtime
.embed_batch(&tokens)
.expect("actual CoreML prediction");
assert_eq!((output.rows(), output.dims()), (3, 2));
assert_eq!(output.values(), [8., 11., 4., 9., -2., 6.]);
assert_eq!(
runtime
.embed_batch(&tokens)
.expect("repeat after warm and prediction")
.values(),
output.values()
);
}
}
#[test]
fn coreml_reports_bad_artifacts_paths_shapes_and_prediction_width_without_partial_results() {
let (root, model) = compile_fixture();
assert!(
matches!(CoreMlRuntime::load(&model, 0, ComputeUnits::Cpu), Err(RuntimeError::Shape(detail)) if detail == "CoreML dims must be non-zero")
);
let invalid_utf8 = PathBuf::from(std::ffi::OsString::from_vec(vec![0xff]));
assert!(
matches!(CoreMlRuntime::load(&invalid_utf8, 2, ComputeUnits::Cpu), Err(RuntimeError::Mlx(detail)) if detail == "CoreML model path is not valid UTF-8")
);
assert!(
matches!(CoreMlRuntime::load(Path::new("invalid\0path"), 2, ComputeUnits::Cpu), Err(RuntimeError::Mlx(detail)) if detail == "CoreML model path contains NUL")
);
assert!(
matches!(CoreMlRuntime::load(&root.path().join("absent.mlmodelc"), 2, ComputeUnits::Cpu), Err(RuntimeError::Mlx(detail)) if !detail.is_empty())
);
let mut runtime = CoreMlRuntime::load(&model, 2, ComputeUnits::Cpu).expect("clean runtime");
let wrong_width = TokenBatch::new(vec![1, 2, 3], vec![1.; 3], 1, 3).expect("different width");
assert!(
matches!(runtime.embed_batch(&wrong_width), Err(RuntimeError::Shape(detail)) if detail == "CoreML query has 3 tokens, expected 1..=2")
);
let overflow = TokenBatch {
token_ids: vec![],
attention_mask: vec![],
rows: usize::MAX,
tokens_per_row: 2,
};
assert!(
matches!(runtime.embed_batch(&overflow), Err(RuntimeError::Shape(detail)) if detail == "CoreML batch size overflows")
);
let mut invalid = TokenBatch::new(vec![7, 11, 3, 9], vec![1.; 4], 2, 2).expect("two rows");
invalid.token_ids.pop();
assert!(
matches!(runtime.embed_batch(&invalid), Err(RuntimeError::Shape(detail)) if detail == "token id row is out of range")
);
invalid.token_ids.push(9);
invalid.attention_mask.pop();
assert!(
matches!(runtime.embed_batch(&invalid), Err(RuntimeError::Shape(detail)) if detail == "attention mask row is out of range")
);
let valid = TokenBatch::new(vec![7, 11], vec![1., 0.], 1, 2).expect("valid row");
assert_eq!(
runtime
.embed_batch(&valid)
.expect("errors leave model usable")
.values(),
[8., 11.]
);
let mut wrong_dims = CoreMlRuntime::load(&model, 3, ComputeUnits::Cpu)
.expect("declaration checked at prediction");
assert!(
matches!(wrong_dims.embed_batch(&valid), Err(RuntimeError::Mlx(detail)) if detail == "output embedding has an unexpected width")
);
let short = TokenBatch::new(vec![7], vec![1.], 1, 1).expect("short row");
assert_eq!(
runtime
.embed_batch(&short)
.expect("automatic padding")
.values(),
[8., 0.]
);
}
#[test]
fn discovered_coreml_query_model_is_reported_and_used_by_the_store() {
let (_root, model) = compile_fixture();
let output = Command::new(std::env::current_exe().expect("test executable"))
.args(["--exact", "coreml_store_child", "--nocapture"])
.env("ZE115_COREML_FIXTURE", &model)
.env("ZE_QUERY_COREML", &model)
.output()
.expect("isolated environment for model discovery");
assert!(
output.status.success(),
"CoreML store child failed:\n{}\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
assert!(
String::from_utf8_lossy(&output.stdout).contains("ZE115 real CoreML store query completed")
);
}
#[test]
fn coreml_store_open_refuses_sequence_below_bundle_query_max_tokens() {
let (_root, model) = compile_fixture();
let output = Command::new(std::env::current_exe().expect("test executable"))
.args(["--exact", "coreml_store_mismatch_child", "--nocapture"])
.env("ZE516_COREML_MISMATCH", "1")
.env("ZE_QUERY_COREML", &model)
.output()
.expect("isolated model discovery");
assert!(
output.status.success(),
"{}\n{}",
String::from_utf8_lossy(&output.stdout),
String::from_utf8_lossy(&output.stderr)
);
assert!(
String::from_utf8_lossy(&output.stdout).contains("ZE516 mismatched pair refused at open")
);
}
#[test]
fn coreml_store_mismatch_child() {
if std::env::var_os("ZE516_COREML_MISMATCH").is_none() {
return;
}
let root = tempdir().expect("store root");
let bundle = root.path().join("document.zem");
common::write_symmetric_fixture_bundle(&bundle);
let result = zeppelin_embed_text::TextStore::open(
root.path().join("store"),
&bundle,
Default::default(),
);
assert!(
matches!(result, Err(zeppelin_embed_text::TextError::Runtime(RuntimeError::Shape(detail)))
if detail == "CoreML query sequence length 2 is below bundle query max_tokens 16")
);
println!("ZE516 mismatched pair refused at open");
}
#[test]
fn coreml_store_child() {
if std::env::var_os("ZE115_COREML_FIXTURE").is_none() {
return;
}
use zeppelin_embed_text::{IngestOptions, Legs, QueryOptions, TextDocument, TextStore};
let root = tempdir().expect("store root");
let model = root.path().join("document.zem");
let mut query = common::FixtureTower::query("fixture-query");
query.max_tokens = 2;
query.prefix = "";
common::write_fixture_bundle(
&model,
&[common::FixtureTower::document("fixture-document"), query],
&[],
);
let store = TextStore::open(root.path().join("store"), &model, Default::default())
.expect("actual mixed runtime store");
let backend = store.query_backend();
assert_eq!(backend.runtime.name, "coreml-cpu-ane");
assert_eq!(
backend.requested_compute_units,
zeppelin_embed::epoch::ComputeUnits::CpuAndNeuralEngine
);
assert_eq!(backend.observed_compute_units, None);
assert_eq!(backend.sequence_length, Some(2));
let report = store
.ingest_text(
&[TextDocument::new(7, 1, "bronze")],
IngestOptions::default(),
)
.expect("MLX document ingest");
assert_eq!(report.chunks, 1);
assert!(report.all_threads_joined);
let result = store
.query_text_with_diagnostics("", QueryOptions::new(1).with_legs(Legs::Dense))
.expect("CoreML query");
assert_eq!(result.hits.len(), 1);
assert_eq!(result.hits[0].doc_id, 7);
assert_eq!(result.query_tokens, 2);
assert_eq!(result.embedding_calls, 1);
assert_eq!(result.backend, Some(backend), "queries never load a plan");
let mut diagnostic_backend = backend;
diagnostic_backend.observed_compute_units =
store.observed_compute_units().expect("explicit plan load");
if compute_plan_available() {
let counts = diagnostic_backend
.observed_compute_units
.expect("observed placement");
assert_eq!(
counts.cpu_ops + counts.gpu_ops + counts.neural_engine_ops + counts.other_ops,
1
);
}
let result = store
.query_text_with_diagnostics("", QueryOptions::new(1).with_legs(Legs::Dense))
.expect("query with cached placement");
assert_eq!(result.backend, Some(diagnostic_backend));
assert_eq!(
store.query_backend(),
backend,
"cheap snapshot never loads or returns the plan"
);
store.close().expect("join MLX and CoreML owner");
println!("ZE115 real CoreML store query completed");
}