diff --git a/crates/engine/src/analysis/models.rs b/crates/engine/src/analysis/models.rs index 7aa50b6b9..cc4e8506b 100644 --- a/crates/engine/src/analysis/models.rs +++ b/crates/engine/src/analysis/models.rs @@ -254,9 +254,9 @@ impl Session { } /// Only installing the chosen model is a retry of that model's failed load. - fn retry_installed_model(&mut self, task: Task, id: &str, wait: bool) -> bool { + fn retry_installed_model(&mut self, task: Task, id: &str, wait: bool) -> std::result::Result { if self.chosen_model(task) != id { - return false; + return Ok(false); } set(self.models.lane(task), |s| { s.error = None; @@ -264,8 +264,22 @@ impl Session { // A load that read the old file must not report failure after these weights replaced it. s.generation = s.generation.wrapping_add(1); }); + if wait { + self.wait_model_worker(task)?; + } self.apply_model(task, wait); - true + Ok(true) + } + + fn wait_model_worker(&self, task: Task) -> std::result::Result<(), String> { + let started = Instant::now(); + while self.models.status(task).busy.is_some() { + if started.elapsed() >= ROTO_MODEL_WAIT { + return Err(format!("{} is still loading after {} minutes", self.chosen_model(task), ROTO_MODEL_WAIT.as_secs() / 60)); + } + std::thread::sleep(Duration::from_millis(5)); + } + Ok(()) } /// The Roto Brush model Settings chose, when it is one (not the classic engine). @@ -330,14 +344,8 @@ impl Session { let Some(info) = self.roto_model_info() else { return Ok(()) }; // A download or load under way (from Settings, or `select` without `wait`) finishes first. // Nothing is in flight in the browser build, so this never sleeps there. - if wait && self.models.status(Task::Mask).busy.is_some() { - let started = Instant::now(); - while self.models.status(Task::Mask).busy.is_some() { - if started.elapsed() >= ROTO_MODEL_WAIT { - return Err(format!("{} is still loading after {} minutes", info.name, ROTO_MODEL_WAIT.as_secs() / 60)); - } - std::thread::sleep(Duration::from_millis(5)); - } + if wait { + self.wait_model_worker(Task::Mask)?; } if self.model_path(info).is_none() { if self.models.status(Task::Mask).busy.is_some() { @@ -587,7 +595,7 @@ fn install(s: &mut Session, p: &Value, task: Task) -> Result { }; let dest = install_bytes(&dir, m, &bytes).map_err(|e| bad(&cmd, e))?; let wait = b_p(p, "wait").unwrap_or(false); - if s.retry_installed_model(task, m.id, wait) && wait { + if s.retry_installed_model(task, m.id, wait).map_err(|error| bad(&cmd, error))? && wait { waited(s, task, &cmd)?; } Ok(json!({"id": m.id, "path": dest.display().to_string()})) @@ -828,16 +836,50 @@ mod availability_tests { }); let before = session.models.status(Task::Mask); for wait in [false, true] { - assert!(!session.retry_installed_model(Task::Mask, "sam2.1-hiera-tiny", wait)); + assert!(!session.retry_installed_model(Task::Mask, "sam2.1-hiera-tiny", wait).unwrap()); assert_eq!(session.models.status(Task::Mask), before); } // Installing the selected weights is an explicit retry, even if the installed path // has since gone away (no load is started from this session's missing model folder). - assert!(session.retry_installed_model(Task::Mask, "mobilesam", false)); + assert!(session.retry_installed_model(Task::Mask, "mobilesam", false).unwrap()); assert!(session.models.status(Task::Mask).load_error.is_none()); assert!(session.models.status(Task::Mask).error.is_none()); } + #[cfg(not(target_arch = "wasm32"))] + #[test] + fn installing_selected_weights_with_wait_finishes_the_replacement_load() { + crate::tests::roto::hold_roto_cache_test_lock(); + let dir = std::env::temp_dir().join(format!("ec-model-install-running-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let info = seg::info("mobilesam").unwrap(); + std::fs::File::create(dir.join(info.file_name)).unwrap().set_len(info.size).unwrap(); + let mut session = Session { models_dir: Some(dir.clone()), ..Default::default() }; + session.choose_model(Task::Mask, info.id); + let generation = session.models.status(Task::Mask).generation; + let lane = session.models.lane(Task::Mask).clone(); + spawn(&lane, "Roto Brush", "Loading the previous file".into(), false, Some(generation), move |status| { + let started = Instant::now(); + while status.lock().unwrap().generation == generation { + if started.elapsed() > Duration::from_secs(3) { + return Err("test did not start the replacement".into()); + } + std::thread::sleep(Duration::from_millis(1)); + } + Err("previous file failed".into()) + }); + assert!(session.retry_installed_model(Task::Mask, info.id, true).unwrap()); + let status = session.models.status(Task::Mask); + assert!(status.busy.is_none()); + assert!( + status.load_error.as_ref().is_some_and(|error| error.contains("SHA-256")), + "wait must finish loading the replacement weights and report their actual validation error" + ); + assert_ne!(status.error.as_deref(), Some("previous file failed")); + crate::effects::roto::set_model(None); + std::fs::remove_dir_all(dir).unwrap(); + } + #[test] fn browser_models_exclude_native_trackers_and_keep_portable_models() { let browser: Vec<_> = models_for_platform(Task::Mask, true).map(|m| m.id).collect();