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:
2026-06-29 01:09:47 +01:00
parent cebd709391
commit 992417710f
3 changed files with 200 additions and 49 deletions
+43 -10
View File
@@ -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)
}