static PROCESS_REGISTRY: OnceLock<Arc<Mutex<Vec<u32>>>> = OnceLock::new();
fn get_registry() -> Arc<Mutex<Vec<u32>>> {
PROCESS_REGISTRY
.get_or_init(|| Arc::new(Mutex::new(Vec::new())))
.clone()
}
fn register_process(pid: u32) {
if let Ok(mut registry) = get_registry().lock() {
registry.push(pid);
}
}
fn unregister_process(pid: u32) {
if let Ok(mut registry) = get_registry().lock() {
registry.retain(|&p| p != pid);
}
}
fn kill_all_registered() -> usize {
let registry = get_registry();
let pids = match registry.lock() {
Ok(guard) => guard.clone(),
Err(_) => return 0,
};
let count = pids.len();
for pid in pids {
#[cfg(unix)]
{
let _ = Command::new("kill").args(["-9", &pid.to_string()]).output();
}
#[cfg(windows)]
{
let _ = Command::new("taskkill")
.args(["/F", "/PID", &pid.to_string()])
.output();
}
}
count
}
struct ProcessGuard {
child: Option<Child>,
pid: u32,
}
impl ProcessGuard {
fn new(child: Child) -> Self {
let pid = child.id();
register_process(pid);
Self {
child: Some(child),
pid,
}
}
fn child_mut(&mut self) -> Option<&mut Child> {
self.child.as_mut()
}
fn kill_and_wait(&mut self) {
if let Some(ref mut child) = self.child {
let _ = child.kill();
let _ = child.wait();
}
unregister_process(self.pid);
self.child = None;
}
}
impl Drop for ProcessGuard {
fn drop(&mut self) {
if let Some(ref mut child) = self.child {
let _ = child.kill();
let _ = child.wait();
unregister_process(self.pid);
}
}
}
fn setup_signal_handler() {
if let Err(e) = ctrlc::set_handler(move || {
let count = kill_all_registered();
eprintln!(
"\n{}[JIDOKA] SIGINT received. Reaping {} active child process(es)...{}",
"\x1b[1;33m", count, "\x1b[0m"
);
std::process::exit(130); }) {
eprintln!(
"{}Warning: Could not set SIGINT handler: {}{}",
"\x1b[33m", e, "\x1b[0m"
);
}
}
const RED: &str = "\x1b[0;31m";
const GREEN: &str = "\x1b[0;32m";
const YELLOW: &str = "\x1b[0;33m";
const BLUE: &str = "\x1b[0;34m";
const CYAN: &str = "\x1b[0;36m";
const MAGENTA: &str = "\x1b[0;35m";
const NC: &str = "\x1b[0m";
const BOLD: &str = "\x1b[1m";
const DEFAULT_TIMEOUT: Duration = Duration::from_secs(60);
#[derive(Debug, Clone, Copy, PartialEq)]
enum Modality {
Run,
Chat,
Serve,
}
impl Modality {
fn display_name(&self) -> &'static str {
match self {
Modality::Run => "apr run",
Modality::Chat => "apr chat",
Modality::Serve => "apr serve",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
enum Backend {
Cpu,
Gpu,
}
impl Backend {
fn as_str(&self) -> &'static str {
match self {
Backend::Cpu => "CPU",
Backend::Gpu => "GPU",
}
}
fn flag(&self) -> Option<&'static str> {
match self {
Backend::Cpu => Some("--no-gpu"),
Backend::Gpu => None, }
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum Format {
Gguf,
SafeTensors,
Apr,
}
impl Format {
fn as_str(&self) -> &'static str {
match self {
Format::Gguf => "GGUF",
Format::SafeTensors => "SafeTensors",
Format::Apr => "APR",
}
}
}
#[derive(Debug, Clone, Copy)]
enum TraceLevel {
None,
Brick,
Step,
Layer,
Profile,
}
#[derive(Debug, Clone, Copy, PartialEq)]
enum TestClass {
Quantized,
FullPrecision,
All,
}
impl TestClass {
fn includes_quantized(&self) -> bool {
matches!(self, TestClass::Quantized | TestClass::All)
}
fn includes_full_precision(&self) -> bool {
matches!(self, TestClass::FullPrecision | TestClass::All)
}
}
#[derive(Debug, Clone)]
struct MatrixCell {
id: String,
modality: Modality,
backend: Backend,
format: Format,
model_uri: String, with_trace: bool, }
impl MatrixCell {
fn new(id: &str, backend: Backend, format: Format, model_uri: String) -> Self {
Self {
id: id.to_string(),
modality: Modality::Run, backend,
format,
model_uri,
with_trace: false,
}
}
fn with_modality(mut self, modality: Modality) -> Self {
self.modality = modality;
self
}
fn with_trace(mut self, trace: bool) -> Self {
self.with_trace = trace;
self
}
fn label(&self) -> String {
let trace_suffix = if self.with_trace { " +trace" } else { "" };
format!(
"{} × {} × {}{}",
self.modality.display_name(),
self.backend.as_str(),
self.format.as_str(),
trace_suffix
)
}
}
#[allow(dead_code)] struct ModelFixture {
uri: String,
resolved_path: Option<String>,
format: Format,
verified: bool,
error: Option<String>,
}
impl ModelFixture {
fn new(uri: &str, format: Format) -> Self {
Self {
uri: uri.to_string(),
resolved_path: None,
format,
verified: false,
error: None,
}
}
fn verify(&mut self, config: &Config) -> bool {
if !self.uri.starts_with("hf://") {
let path = std::path::Path::new(&self.uri);
if path.exists() {
self.resolved_path = Some(self.uri.clone());
self.verified = true;
return true;
}
self.error = Some(format!("Local path not found: {}", self.uri));
self.verified = false;
return false;
}
let args = vec!["inspect", &self.uri, "--quiet"];
match run_apr(config, &args) {
Ok(output) => {
if output.contains("Path:") || output.contains("/") {
self.resolved_path = Some(self.uri.clone()); self.verified = true;
true
} else {
self.resolved_path = Some(self.uri.clone());
self.verified = true;
true
}
}
Err(e) => {
self.error = Some(format!("Failed to verify model: {}", e));
self.verified = false;
false
}
}
}
}
#[allow(dead_code)] fn verify_model_fixtures(config: &Config, fixtures: &mut [ModelFixture]) -> (usize, Vec<String>) {
let mut successes = 0;
let mut failures = Vec::new();
for fixture in fixtures.iter_mut() {
if fixture.verify(config) {
successes += 1;
} else {
failures.push(format!(
"{} ({:?}): {}",
fixture.uri,
fixture.format,
fixture.error.as_deref().unwrap_or("Unknown error")
));
}
}
(successes, failures)
}