feat: update auto-pull logic to run in parallel with pip install (#744)

This also fixes a bug in the model puller to avoid pulling the same model more than once.

Signed-off-by: Nick Mitchell <[email protected]>

Author: starpit

pdl-live-react/src-tauri/Cargo.lock

@@ -31,6 +31,7 @@ file_diff = "1.0.0"
 duct = "0.13.7"
 rayon = "1.10.0"
 yaml-rust2 = "0.10.0"
+futures = "0.3.31"
 
 [target.'cfg(not(any(target_os = "android", target_os = "ios")))'.dependencies]
 tauri-plugin-cli = "2"

pdl-live-react/src-tauri/Cargo.toml

@@ -2,14 +2,16 @@ use ::file_diff::diff;
 use ::std::fs::{copy, create_dir_all};
 use ::std::path::{Path, PathBuf};
 use duct::cmd;
+use futures::executor::block_on;
+use yaml_rust2::yaml::LoadError;
 
 use tauri::path::BaseDirectory;
 use tauri::Manager;
 
 use crate::interpreter::load;
 
 #[cfg(desktop)]
-fn pip_install_if_needed(app_handle: tauri::AppHandle) -> Result<PathBuf, tauri::Error> {
+async fn pip_install_if_needed(app_handle: tauri::AppHandle) -> Result<PathBuf, tauri::Error> {
     let cache_path = app_handle.path().cache_dir()?.join("pdl");
 
     create_dir_all(&cache_path)?;
@@ -82,9 +84,9 @@ pub fn run_pdl_program(
         Path::new(&source_file_path).file_name().unwrap()
     );
 
-    let _ = load::pull_if_needed(&source_file_path);
+    let pull_future = load::pull_if_needed(&source_file_path);
+    let bin_path_future = pip_install_if_needed(app_handle);
 
-    let bin_path = pip_install_if_needed(app_handle)?;
     let trace_arg = if let Some(arg) = trace_file {
         if let serde_json::Value::String(f) = &arg.value {
             "--trace=".to_owned() + f
@@ -115,6 +117,15 @@ pub fn run_pdl_program(
         "".to_owned()
     };
 
+    // wait for any model pulls to finish
+    block_on(pull_future).map_err(|e| match e {
+        LoadError::IO(ee) => tauri::Error::Io(ee),
+        LoadError::Scan(ee) => tauri::Error::Anyhow(ee.into()),
+        _ => tauri::Error::FailedToReceiveMessage,
+    })?;
+
+    let bin_path = block_on(bin_path_future)?;
+
     let mut args = vec![
         source_file_path.as_str(),
         trace_arg.as_str(),

pdl-live-react/src-tauri/src/cli/run.rs

@@ -47,12 +47,18 @@ fn extract_models(program: Yaml) -> Vec<String> {
 }
 
 /// Pull models (in parallel) from the PDL program in the given filepath.
-pub fn pull_if_needed(path: &String) -> Result<(), LoadError> {
-    from_path(path)
+pub async fn pull_if_needed(path: &String) -> Result<(), LoadError> {
+    let mut models = from_path(path)
         .unwrap()
         .into_iter()
         .flat_map(extract_models)
-        .collect::<Vec<String>>()
+        .collect::<Vec<String>>();
+
+    // A single program may specify the same model more than once. Dedup!
+    models.sort();
+    models.dedup();
+
+    models
         .into_par_iter()
         .try_for_each(|model| match model {
             m if model.starts_with("ollama/") => ollama_pull(&m[7..]),