feat: update rust to pip install code block requirements

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

Author: starpit

examples/talk/6-code-json.pdl

@@ -26,6 +26,7 @@ text:
 - def: EVAL
   contribute: []
   lang: python
+  requirements: textdistance
   code:
     |
     import textdistance

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

@@ -2,10 +2,19 @@ use ::std::path::Path;
 use duct::cmd;
 use futures::executor::block_on;
 use yaml_rust2::yaml::LoadError;
+use yaml_rust2::{ScanError, Yaml, YamlLoader};
 
-use crate::interpreter::pip::pip_install_interpreter_if_needed;
+use crate::interpreter::pip::{
+    pip_install_code_blocks_if_needed, pip_install_interpreter_if_needed,
+};
 use crate::interpreter::pull::pull_if_needed;
 
+/// Read the given filesystem path and produce a potentially multi-document Yaml
+fn from_path(path: &String) -> Result<Vec<Yaml>, ScanError> {
+    let content = std::fs::read_to_string(path).unwrap();
+    YamlLoader::load_from_str(&content)
+}
+
 #[cfg(desktop)]
 pub fn run_pdl_program(
     source_file_path: String,
@@ -20,8 +29,10 @@ pub fn run_pdl_program(
     );
 
     // async the model pull and pip installs
-    let pull_future = pull_if_needed(&source_file_path);
-    let bin_path_future = pip_install_interpreter_if_needed(app_handle);
+    let programs = from_path(&source_file_path).unwrap();
+    let pull_future = pull_if_needed(programs.clone());
+    let reqs_future = pip_install_code_blocks_if_needed(&app_handle, programs.clone());
+    let bin_path_future = pip_install_interpreter_if_needed(&app_handle);
 
     // wait for any model pulls to finish
     block_on(pull_future).map_err(|e| match e {
@@ -32,6 +43,7 @@ pub fn run_pdl_program(
 
     // wait for any pip installs to finish
     let bin_path = block_on(bin_path_future)?;
+    block_on(reqs_future)?;
 
     let mut args = vec![
         source_file_path,

pdl-live-react/src-tauri/src/interpreter/extract.rs

@@ -5,6 +5,11 @@ pub fn extract_models(programs: Vec<Yaml>) -> Vec<String> {
     extract_values(programs, "model")
 }
 
+/// Extract requirements.txt referenced by the programs
+pub fn extract_requirements(programs: Vec<Yaml>) -> Vec<String> {
+    extract_values(programs, "requirements")
+}
+
 /// Take a list of Yaml fragments and produce a vector of the string-valued entries of the given field
 pub fn extract_values(programs: Vec<Yaml>, field: &str) -> Vec<String> {
     let mut values = programs

pdl-live-react/src-tauri/src/interpreter/pip.rs

@@ -1,23 +1,32 @@
-use ::std::fs::{copy, create_dir_all};
+use ::std::fs::{copy, create_dir_all, write};
 use ::std::path::{Path, PathBuf};
 
 use duct::cmd;
+use rayon::prelude::*;
 use tauri::path::BaseDirectory;
 use tauri::Manager;
+use tempfile::Builder;
+use yaml_rust2::Yaml;
 
+use crate::interpreter::extract;
 use crate::interpreter::shasum;
 
 #[cfg(desktop)]
-pub async fn pip_install_if_needed(
+fn pip_install_if_needed_with_hash(
     cache_path: &Path,
     requirements_path: &Path,
+    hash: String,
+    force: bool,
 ) -> Result<PathBuf, tauri::Error> {
     create_dir_all(&cache_path)?;
 
-    let hash = shasum::sha256sum(&requirements_path)?;
-    let venv_path = cache_path.join(hash);
+    let venv_path = cache_path.join("venvs").join(hash);
     let bin_path = venv_path.join(if cfg!(windows) { "Scripts" } else { "bin" });
 
+    // re: force, this is part of the short-term hack to install all
+    // code block dependencies in the main interpreter venv. Once we
+    // figure out how to support a separate venv for each code block
+    // (that needs it), we can undo this hack.
     if !venv_path.exists() {
         println!("Creating virtual environment...");
         let python = if cfg!(target_os = "macos") {
@@ -27,18 +36,65 @@ pub async fn pip_install_if_needed(
         };
         cmd!(python, "-mvenv", &venv_path).run()?;
 
-        cmd!(bin_path.join("pip"), "install", "-r", &requirements_path,).run()?;
+        if !force {
+            cmd!(bin_path.join("pip"), "install", "-r", &requirements_path).run()?;
+
+            let cached_requirements_path = venv_path.join("requirements.txt");
+            copy(requirements_path, cached_requirements_path)?;
+        }
+    }
 
-        let cached_requirements_path = venv_path.join("requirements.txt");
-        copy(requirements_path, cached_requirements_path)?;
+    if force {
+        cmd!(bin_path.join("pip"), "install", "-r", &requirements_path).run()?;
     }
 
     Ok(bin_path.to_path_buf())
 }
 
+#[cfg(desktop)]
+fn pip_install_if_needed(
+    cache_path: &Path,
+    requirements_path: &Path,
+) -> Result<PathBuf, tauri::Error> {
+    let hash = shasum::sha256sum(&requirements_path)?;
+    pip_install_if_needed_with_hash(cache_path, requirements_path, hash, false)
+}
+
+#[cfg(desktop)]
+pub async fn pip_install_code_blocks_if_needed(
+    app_handle: &tauri::AppHandle,
+    programs: Vec<Yaml>,
+) -> Result<(), tauri::Error> {
+    let cache_path = app_handle.path().cache_dir()?.join("pdl");
+
+    // for now, install the requirements in the main interpreter venv
+    let requirements_path = app_handle
+        .path()
+        .resolve("interpreter/requirements.txt", BaseDirectory::Resource)?;
+
+    extract::extract_requirements(programs)
+        .into_par_iter()
+        .try_for_each(|req| -> Result<(), tauri::Error> {
+            let req_path = Builder::new()
+                .prefix("pdl-requirements-")
+                .suffix(".txt")
+                .tempfile()?;
+            // This is part of the "force" hack described above, where
+            // we force the code block dependencies to be installed in
+            // the main interpreter venv.
+            let hash = shasum::sha256sum(&requirements_path)?;
+            write(&req_path, req)?;
+            pip_install_if_needed_with_hash(&cache_path, &req_path.path(), hash, true)?;
+            Ok(())
+        })
+        .expect("code block requirements installed");
+
+    Ok(())
+}
+
 #[cfg(desktop)]
 pub async fn pip_install_interpreter_if_needed(
-    app_handle: tauri::AppHandle,
+    app_handle: &tauri::AppHandle,
 ) -> Result<PathBuf, tauri::Error> {
     // the interpreter requirements.txt
     let requirements_path = app_handle
@@ -47,5 +103,5 @@ pub async fn pip_install_interpreter_if_needed(
 
     let cache_path = app_handle.path().cache_dir()?.join("pdl");
 
-    pip_install_if_needed(&cache_path, &requirements_path).await
+    pip_install_if_needed(&cache_path, &requirements_path)
 }

pdl-live-react/src-tauri/src/interpreter/pull.rs

@@ -1,19 +1,13 @@
 use duct::cmd;
 use rayon::prelude::*;
 use yaml_rust2::yaml::LoadError;
-use yaml_rust2::{ScanError, Yaml, YamlLoader};
+use yaml_rust2::Yaml;
 
 use crate::interpreter::extract;
 
-/// Read the given filesystem path and produce a potentially multi-document Yaml
-fn from_path(path: &String) -> Result<Vec<Yaml>, ScanError> {
-    let content = std::fs::read_to_string(path).unwrap();
-    YamlLoader::load_from_str(&content)
-}
-
 /// Pull models (in parallel) from the PDL program in the given filepath.
-pub async fn pull_if_needed(path: &String) -> Result<(), LoadError> {
-    extract::extract_models(from_path(path).unwrap())
+pub async fn pull_if_needed(programs: Vec<Yaml>) -> Result<(), LoadError> {
+    extract::extract_models(programs)
         .into_par_iter()
         .try_for_each(|model| match model {
             m if model.starts_with("ollama/") => ollama_pull_if_needed(&m[7..]),

pdl-live-react/src-tauri/src/interpreter/shasum.rs

@@ -2,7 +2,7 @@ use ::std::fs::File;
 use ::std::io::{copy, Result};
 use ::std::path::Path;
 
-use base64ct::{Base64, Encoding};
+use base64ct::{Base64Url, Encoding};
 use sha2::{Digest, Sha256};
 
 pub fn sha256sum(path: &Path) -> Result<String> {
@@ -12,5 +12,5 @@ pub fn sha256sum(path: &Path) -> Result<String> {
     copy(&mut file, &mut hasher)?;
     let hash_bytes = hasher.finalize();
 
-    Ok(Base64::encode_string(&hash_bytes))
+    Ok(Base64Url::encode_string(&hash_bytes))
 }

pdl-live-react/src/pdl_ast.d.ts

@@ -1919,6 +1919,11 @@ export type Kind15 = "code"
  *
  */
 export type Lang = "python" | "command" | "jinja" | "pdl"
+/**
+ * Pip requirements.txt
+ *
+ */
+export type Requirements = string | null
 /**
  * Code to execute.
  *
@@ -2837,6 +2842,7 @@ export interface CodeBlock {
   pdl__is_leaf?: PdlIsLeaf15
   kind?: Kind15
   lang: Lang
+  requirements: Requirements
   code: Code
 }
 /**

src/pdl/pdl-schema.json

@@ -1306,6 +1306,18 @@
           "title": "Lang",
           "type": "string"
         },
+        "requirements": {
+          "anyOf": [
+            {
+              "type": "string"
+            },
+            {
+              "type": "null"
+            }
+          ],
+          "description": "Pip requirements.txt\n    ",
+          "title": "Requirements"
+        },
         "code": {
           "anyOf": [
             {
@@ -1390,6 +1402,7 @@
       },
       "required": [
         "lang",
+        "requirements",
         "code"
       ],
       "title": "CodeBlock",

src/pdl/pdl_ast.py

@@ -462,6 +462,9 @@ class CodeBlock(LeafBlock):
     ]
     """Programming language of the code.
     """
+    requirements: Optional[str]
+    """Pip requirements.txt
+    """
     code: "BlockType"
     """Code to execute.
     """