Skip to main content

max / audiofiles

7.9 KB · 210 lines History Blame Raw
1 //! Background classifier worker: runs heavy library-wide classifier operations off the
2 //! GUI thread so the window stays responsive (training, library auto-tagging, clustering,
3 //! and `.afcl` export can each take seconds-to-minutes on a large library).
4 //!
5 //! Mirrors `cleanup.rs`: a one-shot thread with its own `Database` connection (WAL +
6 //! `busy_timeout`, so it runs concurrently with the GUI connection), reporting back over a
7 //! channel the GUI polls each frame. One job runs at a time, the UI disables the trigger
8 //! buttons while a job is in flight, and dispatching a new job drops (joins) the prior one.
9
10 use std::path::PathBuf;
11
12 use tracing::{error, instrument};
13
14 use audiofiles_core::analysis::{afcl, cluster, exemplar, trained_head};
15 use audiofiles_core::db::Database;
16 use audiofiles_core::worker_runtime::{JobHandle, spawn_job};
17
18 use crate::backend::{ClassifierJob, ClassifierJobResult, TrainedHeadInfo};
19
20 /// Event sent from the classifier worker back to the GUI thread.
21 pub enum ClassifierWorkerEvent {
22 Done(ClassifierJobResult),
23 Failed(String),
24 }
25
26 /// Handle for a single in-flight classifier job. Thin wrapper over the shared
27 /// one-shot [`JobHandle`], which owns the result channel, the `catch_unwind`
28 /// panic-isolation, and the join-on-drop.
29 pub struct ClassifierHandle(JobHandle<ClassifierWorkerEvent>);
30
31 impl ClassifierHandle {
32 /// Poll for the job's result without blocking.
33 pub fn try_recv(&self) -> Option<ClassifierWorkerEvent> {
34 self.0.try_recv()
35 }
36 }
37
38 // Join-on-drop is owned by the inner JobHandle (so a replaced job finishes cleanly).
39
40 /// Spawn a one-shot worker thread to run `job` against the database at `db_path`.
41 ///
42 /// The shared [`spawn_job`] runtime owns the panic isolation: a panic in the job
43 /// (or `Database::open`) becomes a `Failed` event rather than a silently-dead
44 /// thread, so the GUI's `try_recv` always resolves and the trigger buttons
45 /// re-enable.
46 #[instrument(skip_all)]
47 pub fn spawn_classifier_job(
48 db_path: PathBuf,
49 job: ClassifierJob,
50 ) -> std::io::Result<ClassifierHandle> {
51 let handle = spawn_job(
52 "classifier-job",
53 || ClassifierWorkerEvent::Failed("classifier job panicked (internal error)".to_string()),
54 move || match Database::open(&db_path) {
55 Ok(db) => run_job(&db, job),
56 Err(e) => {
57 error!("Classifier worker failed to open database: {e}");
58 ClassifierWorkerEvent::Failed(format!("could not open database: {e}"))
59 }
60 },
61 )?;
62 Ok(ClassifierHandle(handle))
63 }
64
65 fn run_job(db: &Database, job: ClassifierJob) -> ClassifierWorkerEvent {
66 let result = match job {
67 ClassifierJob::TrainHead => trained_head::train_and_save(db)
68 .map(|h| ClassifierJobResult::Trained(h.map(|head| head_info(&head)))),
69 ClassifierJob::AutoApply { k } => auto_apply(db, k).map(ClassifierJobResult::AutoApplied),
70 ClassifierJob::Cluster { k } => {
71 cluster::cluster_library(db, k, 50).map(|r| ClassifierJobResult::Clustered(r.clusters))
72 }
73 ClassifierJob::Export { path, opts } => {
74 afcl::export_to_path(db, &path, &opts).map(|()| ClassifierJobResult::Exported(path))
75 }
76 ClassifierJob::SuggestSample { hash, k } => {
77 // Build the O(library) exemplar index and score this one sample off the
78 // GUI thread (was done under the DB lock on the render thread).
79 let index = match exemplar::build_index(db) {
80 Ok(idx) => idx,
81 Err(e) => return ClassifierWorkerEvent::Failed(e.to_string()),
82 };
83 exemplar::preview_sample(db, &hash, &index, k)
84 .map(|outcome| ClassifierJobResult::Suggested { hash, outcome })
85 }
86 };
87 match result {
88 Ok(r) => ClassifierWorkerEvent::Done(r),
89 Err(e) => ClassifierWorkerEvent::Failed(e.to_string()),
90 }
91 }
92
93 /// Library auto-tagging: route through the trained head when one is current (O(tags) per
94 /// sample), else brute-force k-NN. Mirrors `DirectBackend::ml_auto_apply_all`.
95 fn auto_apply(db: &Database, k: usize) -> audiofiles_core::error::Result<usize> {
96 if let Some(head) = trained_head::load(db)? {
97 trained_head::auto_apply_library(db, &head)
98 } else {
99 exemplar::auto_apply_library(db, k)
100 }
101 }
102
103 fn head_info(h: &trained_head::TrainedHead) -> TrainedHeadInfo {
104 TrainedHeadInfo {
105 class_count: h.classes.len(),
106 exemplar_count: h.exemplar_count,
107 trained_at: h.trained_at,
108 }
109 }
110
111 #[cfg(test)]
112 mod tests {
113 use super::*;
114 use audiofiles_core::analysis::classify::{FEATURE_VERSION, NUM_FEATURES};
115
116 fn seed(db_path: &std::path::Path) {
117 let db = Database::open(db_path).unwrap();
118 for i in 0..6 {
119 let hash = format!("k{i}");
120 db.conn()
121 .execute(
122 "INSERT INTO samples (hash, original_name, file_extension, file_size, import_date, last_modified) \
123 VALUES (?1, ?1, 'wav', 1, 0, 0)",
124 [&hash],
125 )
126 .unwrap();
127 let v = serde_json::to_string(&vec![0.01 * i as f64; NUM_FEATURES]).unwrap();
128 db.conn()
129 .execute(
130 "INSERT INTO sample_features (hash, feat_version, vector, computed_at) VALUES (?1, ?2, ?3, 0)",
131 rusqlite::params![hash, FEATURE_VERSION, v],
132 )
133 .unwrap();
134 audiofiles_core::tags::add_tag(&db, &hash, "instrument.drum.kick").unwrap();
135 }
136 }
137
138 fn run_to_completion(db_path: PathBuf, job: ClassifierJob) -> ClassifierWorkerEvent {
139 let handle = spawn_classifier_job(db_path, job).unwrap();
140 loop {
141 if let Some(ev) = handle.try_recv() {
142 return ev;
143 }
144 std::thread::sleep(std::time::Duration::from_millis(5));
145 }
146 }
147
148 #[test]
149 fn train_job_completes_on_worker() {
150 let dir = tempfile::TempDir::new().unwrap();
151 let db_path = dir.path().join("audiofiles.db");
152 seed(&db_path);
153 match run_to_completion(db_path, ClassifierJob::TrainHead) {
154 ClassifierWorkerEvent::Done(ClassifierJobResult::Trained(Some(info))) => {
155 assert_eq!(info.class_count, 1);
156 }
157 _ => panic!("expected a trained head"),
158 }
159 }
160
161 #[test]
162 fn cluster_job_completes_on_worker() {
163 let dir = tempfile::TempDir::new().unwrap();
164 let db_path = dir.path().join("audiofiles.db");
165 seed(&db_path);
166 match run_to_completion(db_path, ClassifierJob::Cluster { k: 2 }) {
167 ClassifierWorkerEvent::Done(ClassifierJobResult::Clustered(clusters)) => {
168 assert!(!clusters.is_empty());
169 }
170 _ => panic!("expected clusters"),
171 }
172 }
173
174 #[test]
175 fn suggest_sample_job_completes_on_worker() {
176 let dir = tempfile::TempDir::new().unwrap();
177 let db_path = dir.path().join("audiofiles.db");
178 seed(&db_path);
179 let job = ClassifierJob::SuggestSample {
180 hash: "k0".to_string(),
181 k: 5,
182 };
183 match run_to_completion(db_path, job) {
184 ClassifierWorkerEvent::Done(ClassifierJobResult::Suggested { hash, .. }) => {
185 assert_eq!(hash, "k0");
186 }
187 _ => panic!("expected suggestions"),
188 }
189 }
190
191 #[test]
192 fn export_job_writes_file_on_worker() {
193 let dir = tempfile::TempDir::new().unwrap();
194 let db_path = dir.path().join("audiofiles.db");
195 seed(&db_path);
196 let out = dir.path().join("out.afcl");
197 let job = ClassifierJob::Export {
198 path: out.clone(),
199 opts: afcl::ExportOptions::default(),
200 };
201 match run_to_completion(db_path, job) {
202 ClassifierWorkerEvent::Done(ClassifierJobResult::Exported(p)) => {
203 assert_eq!(p, out);
204 assert!(out.exists());
205 }
206 _ => panic!("expected export"),
207 }
208 }
209 }
210