kalugin66/windmill
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Handle pip install by uv (#4517)

Browse Source 
* feat: Handle `pip install` by `uv`

Dirty and untested, but already something working

* Integrate with NSJAIL and prepare fallbacks

* Refactor fallback
no_uv disable compile and install
where no_uv_install and no_uv_compile are a bit more specific

* Remove `--disable-pip-version-check`
Reason:
   warning: pip's `--disable-pip-version-check` has no effect

* Fix backend compilation error

* Pip fallback overwrite UV's cache

* Initially refactor cache (No S3)

* Support S3

* Remove unused import

* Handle flags for NSJAIL

* Return deleted flag

* Update Dockerfile

* Update docker-image.yml

* Update docker-image.yml

* Add --link-mode=copy and remove -v

* Fix NSJAIL INDEX_URL

* Fix flock and windows

* Update python_executor.rs

* Remove line from Dockerfile

We dont need it and to trigger build

* fixing for windows

* Dont pin python to specific version

* Change TMP for windows

* Revert docker-image.yml

* Disable UV for ansible

Will be enabled later.
Needs proper testing and its better to split onto 2 PRs with first modifying python and second ansible

---------

Co-authored-by: Ruben Fiszel <ruben@windmill.dev>
Co-authored-by: Alexander Petric <petric.al@gmail.com>
This commit is contained in:
pyranota
2024-11-14 23:12:56 +00:00
committed by GitHub
parent fd510a5ed1
commit 04e2ef5809
10 changed files with 285 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
backend/windmill-worker/src/python_executor.rs
View File

@@ -42,6 +42,10 @@ lazy_static::lazy_static! {
static ref USE_PIP_COMPILE: bool = std::env::var("USE_PIP_COMPILE")
.ok().map(|flag| flag == "true").unwrap_or(false);
/// Use pip install
static ref USE_PIP_INSTALL: bool = std::env::var("USE_PIP_INSTALL")
.ok().map(|flag| flag == "true").unwrap_or(false);
static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(((u|f)\.)|\.)"#).unwrap();
@@ -50,6 +54,8 @@ lazy_static::lazy_static! {
}
const NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT: &str = include_str!("../nsjail/download.py.config.proto");
const NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT_FALLBACK: &str =
include_str!("../nsjail/download.py.pip.config.proto");
const NSJAIL_CONFIG_RUN_PYTHON3_CONTENT: &str = include_str!("../nsjail/run.python3.config.proto");
const RELATIVE_PYTHON_LOADER: &str = include_str!("../loader.py");
@@ -66,8 +72,8 @@ use crate::{
},
handle_child::handle_child,
AuthedClientBackgroundTask, DISABLE_NSJAIL, DISABLE_NUSER, HOME_ENV, LOCK_CACHE_DIR,
NSJAIL_PATH, PATH_ENV, PIP_CACHE_DIR, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, PROXY_ENVS, TZ_ENV,
UV_CACHE_DIR,
NSJAIL_PATH, PATH_ENV, PIP_CACHE_DIR, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, PROXY_ENVS,
PY311_CACHE_DIR, TZ_ENV, UV_CACHE_DIR,
};
#[cfg(windows)]
@@ -901,10 +907,10 @@ async fn handle_python_deps(
.unwrap_or_else(|| vec![])
.clone();
let annotations = windmill_common::worker::PythonAnnotations::parse(inner_content);
let requirements = match requirements_o {
Some(r) => r,
None => {
let annotation = windmill_common::worker::PythonAnnotations::parse(inner_content);
let mut already_visited = vec![];
let requirements = windmill_parser_py_imports::parse_python_imports(
@@ -929,8 +935,8 @@ async fn handle_python_deps(
worker_name,
w_id,
occupancy_metrics,
annotation.no_uv,
annotation.no_cache,
annotations.no_uv || annotations.no_uv_compile,
annotations.no_cache,
)
.await
.map_err(|e| {
@@ -955,6 +961,8 @@ async fn handle_python_deps(
job_dir,
worker_dir,
occupancy_metrics,
annotations.no_uv || annotations.no_uv_install,
false,
)
.await?;
additional_python_paths.append(&mut venv_path);
@@ -966,6 +974,7 @@ lazy_static::lazy_static! {
static ref PIP_SECRET_VARIABLE: Regex = Regex::new(r"\$\{PIP_SECRET:([^\s\}]+)\}").unwrap();
}
/// pip install, include cached or pull from S3
pub async fn handle_python_reqs(
requirements: Vec<&str>,
job_id: &Uuid,
@@ -977,12 +986,22 @@ pub async fn handle_python_reqs(
job_dir: &str,
worker_dir: &str,
occupancy_metrics: &mut Option<&mut OccupancyMetrics>,
// TODO: Remove (Deprecated)
mut no_uv_install: bool,
is_ansible: bool,
) -> error::Result<Vec<String>> {
let mut req_paths: Vec<String> = vec![];
let mut vars = vec![("PATH", PATH_ENV.as_str())];
let pip_extra_index_url;
let pip_index_url;
no_uv_install |= *USE_PIP_INSTALL;
if no_uv_install && !is_ansible {
append_logs(&job_id, w_id, "\nFallback to pip (Deprecated!)\n", db).await;
tracing::warn!("Fallback to pip");
}
if !*DISABLE_NSJAIL {
pip_extra_index_url = PIP_EXTRA_INDEX_URL
.read()
@@ -1013,10 +1032,21 @@ pub async fn handle_python_reqs(
let _ = write_file(
job_dir,
"download.config.proto",
&NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT
.replace("{WORKER_DIR}", &worker_dir)
.replace("{CACHE_DIR}", PIP_CACHE_DIR)
.replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()),
&(if no_uv_install {
NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT_FALLBACK
} else {
NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT
})
.replace("{WORKER_DIR}", &worker_dir)
.replace(
"{CACHE_DIR}",
if no_uv_install {
PIP_CACHE_DIR
} else {
PY311_CACHE_DIR
},
)
.replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string()),
)?;
};
@@ -1026,8 +1056,14 @@ pub async fn handle_python_reqs(
if req.starts_with('#') {
continue;
}
let py_prefix = if no_uv_install {
PIP_CACHE_DIR
} else {
PY311_CACHE_DIR
};
let venv_p = format!(
"{PIP_CACHE_DIR}/{}",
"{py_prefix}/{}",
req.replace(' ', "").replace('/', "").replace(':', "")
);
if metadata(&venv_p).await.is_ok() {
@@ -1073,7 +1109,10 @@ pub async fn handle_python_reqs(
.map(|(req, venv_p)| {
let os = os.clone();
async move {
if pull_from_tar(os, venv_p.clone()).await.is_ok() {
if pull_from_tar(os, venv_p.clone(), no_uv_install)
.await
.is_ok()
{
PullFromTar::Pulled(venv_p.to_string())
} else {
PullFromTar::NotPulled(req.to_string(), venv_p.to_string())
@@ -1114,7 +1153,7 @@ pub async fn handle_python_reqs(
for (req, venv_p) in req_with_penv {
let mut logs1 = String::new();
logs1.push_str("\n\n--- PIP INSTALL ---\n");
logs1.push_str("\n\n--- UV PIP INSTALL ---\n");
logs1.push_str(&format!("\n{req} is being installed for the first time.\n It will be cached for all ulterior uses."));
append_logs(&job_id, w_id, logs1, db).await;
@@ -1150,21 +1189,47 @@ pub async fn handle_python_reqs(
#[cfg(windows)]
let req = format!("{}", req);
let mut command_args = vec![
PYTHON_PATH.as_str(),
"-m",
"pip",
"install",
&req,
"-I",
"--no-deps",
"--no-color",
"--isolated",
"--no-warn-conflicts",
"--disable-pip-version-check",
"-t",
venv_p.as_str(),
];
let mut command_args = if no_uv_install {
vec![
PYTHON_PATH.as_str(),
"-m",
"pip",
"install",
&req,
"-I",
"--no-deps",
"--no-color",
"--isolated",
"--no-warn-conflicts",
"--disable-pip-version-check",
"-t",
venv_p.as_str(),
]
} else {
vec![
UV_PATH.as_str(),
"pip",
"install",
&req,
"--no-deps",
"--no-color",
// "-p",
// "3.11",
// Prevent uv from discovering configuration files.
"--no-config",
"--link-mode=copy",
// TODO: Doublecheck it
"--system",
// Prefer main index over extra
// https://docs.astral.sh/uv/pip/compatibility/#packages-that-exist-on-multiple-indexes
// TODO: Use env variable that can be toggled from UI
"--index-strategy",
"unsafe-best-match",
"--target",
venv_p.as_str(),
"--no-cache",
]
};
let pip_extra_index_url = PIP_EXTRA_INDEX_URL
.read()
.await
@@ -1196,7 +1261,7 @@ pub async fn handle_python_reqs(
envs.push(("HOME", HOME_ENV.as_str()));
tracing::debug!("pip install command: {:?}", command_args);
tracing::debug!("uv pip install command: {:?}", command_args);
#[cfg(unix)]
{
@@ -1207,7 +1272,12 @@ pub async fn handle_python_reqs(
.envs(envs)
.args([
"-x",
&format!("{}/pip-{}.lock", LOCK_CACHE_DIR, fssafe_req),
&format!(
"{}/{}-{}.lock",
LOCK_CACHE_DIR,
if no_uv_install { "pip" } else { "py311" },
fssafe_req
),
"--command",
&command_args.join(" "),
])
@@ -1218,16 +1288,20 @@ pub async fn handle_python_reqs(
#[cfg(windows)]
{
let mut pip_cmd = Command::new(PYTHON_PATH.as_str());
pip_cmd
.env_clear()
let installer_path = if no_uv_install { command_args[0] } else { "uv" };
let mut cmd: Command = Command::new(&installer_path);
cmd.env_clear()
.envs(envs)
.envs(PROXY_ENVS.clone())
.env("SystemRoot", SYSTEM_ROOT.as_str())
.env(
"TMP",
std::env::var("TMP").unwrap_or_else(|_| String::from("/tmp")),
)
.args(&command_args[1..])
.stdout(Stdio::piped())
.stderr(Stdio::piped());
start_child_process(pip_cmd, PYTHON_PATH.as_str()).await?
start_child_process(cmd, installer_path).await?
}
};
@@ -1240,7 +1314,7 @@ pub async fn handle_python_reqs(
false,
worker_name,
&w_id,
&format!("pip install {req}"),
&format!("uv pip install {req}"),
None,
false,
occupancy_metrics,
@@ -1260,7 +1334,7 @@ pub async fn handle_python_reqs(
tracing::warn!("S3 cache not available in the pro plan");
} else {
let venv_p = venv_p.clone();
tokio::spawn(build_tar_and_push(os, venv_p));
tokio::spawn(build_tar_and_push(os, venv_p, no_uv_install));
}
}
req_paths.push(venv_p);