Skip to content

feat: update rust to pip install code block requirements #752

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions examples/talk/6-code-json.pdl
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ text:
- def: EVAL
contribute: []
lang: python
requirements:
- textdistance
code:
|
import textdistance
Expand Down
16 changes: 13 additions & 3 deletions pdl-live-react/src-tauri/src/cli/run.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,17 @@ 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_internal_if_needed;
use crate::interpreter::pip::{pip_install_code_blocks_if_needed, pip_install_internal_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,
Expand All @@ -20,9 +27,11 @@ pub fn run_pdl_program(
);

// async the model pull and pip installs
let pull_future = pull_if_needed(&source_file_path);
let program = &from_path(&source_file_path).unwrap()[0];
let pull_future = pull_if_needed(&program);
let reqs_future = pip_install_code_blocks_if_needed(&app_handle, &program);
let bin_path_future =
pip_install_internal_if_needed(app_handle, &"interpreter/requirements.txt");
pip_install_internal_if_needed(&app_handle, &"interpreter/requirements.txt");

// wait for any model pulls to finish
block_on(pull_future).map_err(|e| match e {
Expand All @@ -33,6 +42,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,
Expand Down
2 changes: 1 addition & 1 deletion pdl-live-react/src-tauri/src/compile/beeai.rs
Original file line number Diff line number Diff line change
Expand Up @@ -291,7 +291,7 @@ fn python_source_to_json(
eprintln!("Compiling from Python source");
}
let bin_path = block_on(pip_install_internal_if_needed(
app_handle,
&app_handle,
&"interpreter/beeai-requirements.txt",
))?;

Expand Down
24 changes: 15 additions & 9 deletions pdl-live-react/src-tauri/src/interpreter/extract.rs
Original file line number Diff line number Diff line change
@@ -1,16 +1,18 @@
use yaml_rust2::Yaml;

/// Extract models referenced by the programs
pub fn extract_models(programs: Vec<Yaml>) -> Vec<String> {
extract_values(programs, "model")
/// Extract models referenced by the program
pub fn extract_models(program: &Yaml) -> Vec<String> {
extract_values(program, "model")
}

/// Extract requirements.txt referenced by the program
pub fn extract_requirements(program: &Yaml) -> Vec<String> {
extract_values(program, "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
.into_iter()
.flat_map(|p| extract_one_values(p, field))
.collect::<Vec<String>>();
pub fn extract_values(program: &Yaml, field: &str) -> Vec<String> {
let mut values = extract_one_values(program, field);

// A single program may specify the same model more than once. Dedup!
values.sort();
Expand All @@ -20,7 +22,7 @@ pub fn extract_values(programs: Vec<Yaml>, field: &str) -> Vec<String> {
}

/// Take one Yaml fragment and produce a vector of the string-valued entries of the given field
fn extract_one_values(program: Yaml, field: &str) -> Vec<String> {
fn extract_one_values(program: &Yaml, field: &str) -> Vec<String> {
let mut values: Vec<String> = Vec::new();

match program {
Expand All @@ -31,6 +33,10 @@ fn extract_one_values(program: Yaml, field: &str) -> Vec<String> {
Yaml::String(m) => {
values.push(m.to_string());
}
Yaml::Array(a) => a.into_iter().for_each(|v| match v {
Yaml::String(m) => values.push(m.to_string()),
_ => {}
}),
_ => {}
},
_ => {}
Expand Down
72 changes: 64 additions & 8 deletions pdl-live-react/src-tauri/src/interpreter/pip.rs
Original file line number Diff line number Diff line change
@@ -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("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") {
Expand All @@ -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,
program: &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(program)
.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_internal_if_needed(
app_handle: tauri::AppHandle,
app_handle: &tauri::AppHandle,
requirements: &str,
) -> Result<PathBuf, tauri::Error> {
// the interpreter requirements.txt
Expand All @@ -48,5 +104,5 @@ pub async fn pip_install_internal_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)
}
12 changes: 3 additions & 9 deletions pdl-live-react/src-tauri/src/interpreter/pull.rs
Original file line number Diff line number Diff line change
@@ -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(program: &Yaml) -> Result<(), LoadError> {
extract::extract_models(program)
.into_par_iter()
.try_for_each(|model| match model {
m if model.starts_with("ollama/") => ollama_pull_if_needed(&m[7..]),
Expand Down
11 changes: 8 additions & 3 deletions pdl-live-react/src/helpers.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,10 @@ import type {
PdlBlock,
TextBlock,
ArgsBlock,
CodeBlock,
PdlModelInput,
LocalizedExpression,
PythonCodeBlock,
NonPythonCodeBlock,
} from "./pdl_ast"

/** Re-export for convenience */
Expand Down Expand Up @@ -301,11 +302,15 @@ export function extractStructuredModelResponse({
return { resultForDisplay, lang, meta }
}

export function isArgs(block: ArgsBlock | CodeBlock): block is ArgsBlock {
export function isArgs(
block: ArgsBlock | PythonCodeBlock | NonPythonCodeBlock,
): block is ArgsBlock {
return Array.isArray((block as ArgsBlock).args)
}

export function extractCode({ code }: CodeBlock): string {
export function extractCode({
code,
}: PythonCodeBlock | NonPythonCodeBlock): string {
if (
isNonScalarPdlBlock(code) &&
hasResult(code) &&
Expand Down
Loading
Loading