perf(tagger): chunked, yielding inference so tagging stops freezing the UI
The tagging worker claimed a batch from the DB but ran the model one image at a time, so `tagger_batch_size` had no effect on inference. Batching the whole claim into a single DirectML forward pass fixed that but introduced a worse problem: on a shared GPU each wide dispatch locks the device (and the WebView2 compositor with it) for 1-3.7s, freezing the entire app while tagging runs — worst of all on first launch with a cold graph compile. The WD model is compute-bound here (~50-230ms/image of actual GPU work), so a wide batch buys almost no throughput; it only lengthens each uninterruptible GPU lock. So decouple DB claim size from GPU granularity: claim `tagger_batch_size` for DB efficiency, but feed the GPU in TAGGER_INFER_CHUNK (4) images per forward pass with a brief yield between chunks. Each dispatch now lasts ~200-900ms, the UI gets windows to paint, peak decode memory is bounded, and the cold compile is for a smaller shape. As a bonus the wide-batch throughput spikes disappear — steady ~1.6-1.8s/16 vs the old 0.8-3.7s swings. - run_batch: parallel (rayon) decode + single forward pass per chunk, with per-image fallback and decode failures kept attached to their input slot. - Worker iterates source_paths.chunks(TAGGER_INFER_CHUNK), yielding 40ms between chunks; outputs zip back to jobs 1:1, write tx unchanged. - Remove now-dead WdTagger::run() (fallback uses infer_one directly).
This commit is contained in:
+43
-10
@@ -1248,13 +1248,16 @@ fn process_tagging_batch(
|
||||
|
||||
// Exclude actively-indexing folders for the same reason as the other
|
||||
// workers: don't compete with a running scan.
|
||||
let batch_started_at = Instant::now();
|
||||
let mut excluded_folders = paused_folder_ids("tagging");
|
||||
excluded_folders.extend(active_indexing_folders());
|
||||
let batch_size = crate::tagger::tagger_batch_size(app_data_dir);
|
||||
let claim_started_at = Instant::now();
|
||||
let jobs = with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
db::claim_tagging_jobs(&mut conn, &excluded_folders, batch_size)
|
||||
})?;
|
||||
let claim_elapsed = claim_started_at.elapsed();
|
||||
|
||||
if jobs.is_empty() {
|
||||
return Ok(false);
|
||||
@@ -1288,21 +1291,41 @@ fn process_tagging_batch(
|
||||
.as_mut()
|
||||
.expect("tagger should be initialized before tagging batch processing");
|
||||
|
||||
let tag_results = jobs
|
||||
// Resolve each job's source image (AVIF can't be decoded directly, so it
|
||||
// falls back to its thumbnail), then tag the batch in small chunks (below).
|
||||
let source_paths: Vec<PathBuf> = jobs
|
||||
.iter()
|
||||
.map(|job| {
|
||||
let source_path = if is_avif_path(Path::new(&job.path)) {
|
||||
job.thumbnail_path.as_deref().unwrap_or(&job.path)
|
||||
if is_avif_path(Path::new(&job.path)) {
|
||||
PathBuf::from(job.thumbnail_path.as_deref().unwrap_or(&job.path))
|
||||
} else {
|
||||
&job.path
|
||||
};
|
||||
(
|
||||
job.clone(),
|
||||
tagger_ref.run(Path::new(source_path), tagger::DEFAULT_MAX_TAGS),
|
||||
)
|
||||
PathBuf::from(&job.path)
|
||||
}
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
.collect();
|
||||
|
||||
// Tag in small micro-batches instead of one wide forward pass. On a shared
|
||||
// GPU every DirectML dispatch blocks the WebView2 compositor for its whole
|
||||
// duration, so a 16-wide batch freezes the UI for seconds; small chunks keep
|
||||
// each GPU lock short (and bound peak decode memory) while a brief yield
|
||||
// between them lets the UI and other workers grab the GPU/CPU. The model is
|
||||
// compute-bound here, so this costs almost no throughput.
|
||||
let infer_started_at = Instant::now();
|
||||
let mut outputs = Vec::with_capacity(source_paths.len());
|
||||
let mut chunks = source_paths.chunks(tagger::TAGGER_INFER_CHUNK).peekable();
|
||||
while let Some(chunk) = chunks.next() {
|
||||
outputs.extend(tagger_ref.run_batch(chunk, tagger::DEFAULT_MAX_TAGS));
|
||||
if chunks.peek().is_some() {
|
||||
std::thread::sleep(std::time::Duration::from_millis(
|
||||
tagger::TAGGER_INFER_YIELD_MS,
|
||||
));
|
||||
}
|
||||
}
|
||||
let infer_elapsed = infer_started_at.elapsed();
|
||||
|
||||
let tag_results = jobs.iter().cloned().zip(outputs).collect::<Vec<_>>();
|
||||
|
||||
let write_started_at = Instant::now();
|
||||
let updated_images = with_db_write_lock(|| {
|
||||
let mut conn = pool.get()?;
|
||||
let tx = conn.transaction()?;
|
||||
@@ -1354,6 +1377,7 @@ fn process_tagging_batch(
|
||||
db::requeue_tagging_jobs(&conn, &image_ids)
|
||||
});
|
||||
})?;
|
||||
let write_elapsed = write_started_at.elapsed();
|
||||
|
||||
if !updated_images.is_empty() {
|
||||
let folder_ids = updated_images
|
||||
@@ -1369,6 +1393,15 @@ fn process_tagging_batch(
|
||||
emit_folder_job_progress(app, pool, &folder_ids.into_iter().collect::<Vec<_>>(), true);
|
||||
}
|
||||
|
||||
log::info!(
|
||||
"Tagging batch timing: {} items, claim {:?}, tag {:?}, write {:?}, total {:?}",
|
||||
jobs.len(),
|
||||
claim_elapsed,
|
||||
infer_elapsed,
|
||||
write_elapsed,
|
||||
batch_started_at.elapsed()
|
||||
);
|
||||
|
||||
Ok(true)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user