mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-09-11 16:09:39 +00:00
hugocasaandClaude Opus 5
8aa8b7ee6c
* fix(python): stop uv pip compile emitting lockfile annotations `uv pip compile` annotates each resolved package with the indented `# via <pkg>` comments that name what pulled it in. Windmill installs a lockfile one entry at a time as a `uv pip install` argument, so those lines are unparseable package names to the reader, and the compile step was stripping them back out of its own output. Pass `--no-annotate` so they are never written in the first place: the stored lockfile is annotation-free at the source rather than by virtue of the reader filtering them. Lockfiles that arrive already annotated (deployed from outside Windmill) are handled by `requirement_from_lockfile_line` (#11035); this only covers the ones Windmill generates itself. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * fix: state what the annotation filter below cannot catch The comment justified the flag against the previous design instead of recording the constraint that keeps it there. uv's annotation style is configurable, and `[pip] annotation-style = "line"` in the worker HOME's uv.toml emits annotations inline, which the whole-line `#` filter on the output does not touch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
3854 lines
140 KiB
Rust
3854 lines
140 KiB
Rust
use std::{
|
|
collections::{HashMap, HashSet},
|
|
fs,
|
|
path::Path,
|
|
process::Stdio,
|
|
str::FromStr,
|
|
sync::Arc,
|
|
};
|
|
|
|
#[cfg(unix)]
|
|
use std::os::unix::process::ExitStatusExt;
|
|
|
|
use anyhow::anyhow;
|
|
use itertools::Itertools;
|
|
use regex::Regex;
|
|
use serde_json::value::RawValue;
|
|
use tokio::{
|
|
fs::{metadata, DirBuilder, File},
|
|
io::AsyncReadExt,
|
|
process::Command,
|
|
sync::Semaphore,
|
|
task,
|
|
};
|
|
use windmill_queue::MiniPulledJob;
|
|
|
|
use uuid::Uuid;
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
use windmill_common::ee_oss::{get_license_plan, LicensePlan};
|
|
|
|
use windmill_common::{
|
|
error::{
|
|
self,
|
|
Error::{self},
|
|
},
|
|
jobs::JobKind,
|
|
scripts::ScriptLang,
|
|
utils::calculate_hash,
|
|
worker::{
|
|
copy_dir_recursively, is_allowed_file_location, lockfile_line_has_continuation, pad_string,
|
|
requirement_from_lockfile_line, split_python_requirements, write_file, Connection,
|
|
PyVAlias, PythonAnnotations, WORKER_CONFIG,
|
|
},
|
|
};
|
|
|
|
#[cfg(feature = "enterprise")]
|
|
use windmill_common::variables::get_secret_value_as_admin;
|
|
|
|
use std::env::var;
|
|
use windmill_queue::{append_logs, CanceledBy, PrecomputedAgentInfo};
|
|
|
|
use process_wrap::tokio::TokioChildWrapper;
|
|
|
|
lazy_static::lazy_static! {
|
|
pub(crate) static ref PYTHON_PATH: Option<String> = var("PYTHON_PATH").ok().map(|v| {
|
|
tracing::warn!("PYTHON_PATH is set to {} and thus python will not be managed by uv and stay static regardless of annotation and instance settings. NOT RECOMMENDED", v);
|
|
v
|
|
});
|
|
|
|
pub(crate) static ref UV_PATH: String =
|
|
var("UV_PATH").unwrap_or_else(|_| "/usr/local/bin/uv".to_string());
|
|
|
|
static ref PY_CONCURRENT_DOWNLOADS: usize =
|
|
var("PY_CONCURRENT_DOWNLOADS").ok().map(|flag| flag.parse().unwrap_or(20)).unwrap_or(20);
|
|
|
|
static ref NON_ALPHANUM_CHAR: Regex = regex::Regex::new(r"[^0-9A-Za-z=.-]").unwrap();
|
|
|
|
static ref RELATIVE_IMPORT_REGEX: Regex = Regex::new(r#"(import|from)\s(((u|f)\.)|\.)"#).unwrap();
|
|
|
|
static ref EPHEMERAL_TOKEN_CMD: Option<String> = var("EPHEMERAL_TOKEN_CMD").ok();
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
lazy_static::lazy_static! {
|
|
static ref PIPTAR_UPLOAD_CHANNEL: tokio::sync::mpsc::UnboundedSender<PiptarUploadTask> = {
|
|
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
|
|
|
|
// Spawn background task to handle uploads sequentially
|
|
tokio::spawn(handle_piptar_uploads(rx));
|
|
|
|
tx
|
|
};
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
#[derive(Debug)]
|
|
struct PiptarUploadTask {
|
|
venv_path: String,
|
|
cache_dir: String,
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
async fn handle_piptar_uploads(mut rx: tokio::sync::mpsc::UnboundedReceiver<PiptarUploadTask>) {
|
|
use crate::global_cache::build_tar_and_push;
|
|
use windmill_object_store::get_cache_object_store;
|
|
|
|
while let Some(task) = rx.recv().await {
|
|
if let Some(os) = get_cache_object_store().await {
|
|
match build_tar_and_push(os, task.venv_path.clone(), task.cache_dir, None, false).await
|
|
{
|
|
Ok(()) => {
|
|
tracing::info!("Successfully uploaded piptar for {}", task.venv_path);
|
|
}
|
|
Err(e) => {
|
|
tracing::error!("Failed to upload piptar for {}: {}", task.venv_path, e);
|
|
}
|
|
}
|
|
} else {
|
|
tracing::warn!(
|
|
"S3 object store not available for piptar upload: {}",
|
|
task.venv_path
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
const NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT: &str = include_str!("../nsjail/download.py.config.proto");
|
|
const NSJAIL_CONFIG_RUN_PYTHON3_CONTENT: &str = include_str!("../nsjail/run.python3.config.proto");
|
|
pub const RELATIVE_PYTHON_LOADER: &str = include_str!("../loader.py");
|
|
|
|
/// Every file exchanged with a job is UTF-8 by construction, so the interpreter
|
|
/// must agree. A job env carries no locale: Linux then picks UTF-8 on its own
|
|
/// (PEP 540), Windows picks the ANSI code page. Applied after the whitelisted
|
|
/// envs so the protocol is not the user's to opt out of.
|
|
pub const PYTHON_UTF8_ENVS: [(&str, &str); 1] = [("PYTHONUTF8", "1")];
|
|
|
|
/// Render loader.py with the TEMP_SCRIPT_REFS placeholder substituted by a
|
|
/// Python dict literal. Preview jobs pass a path -> temp-hash map so relative
|
|
/// imports resolve from not-yet-deployed local content; deployed runs pass
|
|
/// `None` which renders an empty dict (deployed resolution unchanged).
|
|
fn render_relative_python_loader(temp_script_refs: &Option<HashMap<String, String>>) -> String {
|
|
let temp_refs_py = temp_script_refs
|
|
.as_ref()
|
|
.and_then(|m| serde_json::to_string(m).ok())
|
|
.unwrap_or_else(|| "{}".to_string());
|
|
RELATIVE_PYTHON_LOADER.replace("TEMP_SCRIPT_REFS_PLACEHOLDER", &temp_refs_py)
|
|
}
|
|
|
|
#[cfg(any(feature = "private", test))]
|
|
pub fn has_relative_imports(content: &str) -> bool {
|
|
RELATIVE_IMPORT_REGEX.is_match(content)
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
use crate::global_cache::pull_from_tar;
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
use windmill_object_store::get_cache_object_store;
|
|
|
|
use crate::{
|
|
common::{
|
|
build_command_with_isolation, create_args_and_out_file, get_reserved_variables, read_file,
|
|
read_result, render_nsjail_rlimit_as, resolve_nsjail_timeout,
|
|
resolve_nsjail_tmp_mount_block, start_child_process, OccupancyMetrics, StreamNotifier,
|
|
DEV_CONF_NSJAIL,
|
|
},
|
|
get_proxy_envs_for_lang,
|
|
handle_child::handle_child,
|
|
is_sandboxing_enabled, read_ee_registry_with_workspace_override,
|
|
worker_utils::ping_job_status,
|
|
PyV, DISABLE_NUSER, HOME_ENV, INDEX_CERT, NATIVE_CERT, NSJAIL_AVAILABLE, NSJAIL_PATH,
|
|
NSJAIL_PY_RLIMIT_AS_MB, PATH_ENV, PIP_EXTRA_INDEX_URL, PIP_INDEX_URL, PROXY_ENVS,
|
|
PY_INSTALL_DIR, TRACING_PROXY_CA_CERT_PATH, TRUSTED_HOST, TZ_ENV, UV_CACHE_DIR,
|
|
UV_EXCLUDE_NEWER, UV_HTTP_TIMEOUT, UV_INDEX_STRATEGY, UV_PYTHON_INSTALL_MIRROR,
|
|
};
|
|
use windmill_common::client::AuthedClient;
|
|
|
|
#[cfg(windows)]
|
|
use crate::SYSTEM_ROOT;
|
|
|
|
pub async fn create_dependencies_dir(job_dir: &str) {
|
|
DirBuilder::new()
|
|
.recursive(true)
|
|
.create(&format!("{job_dir}/dependencies"))
|
|
.await
|
|
.expect("could not create dependencies dir");
|
|
}
|
|
|
|
#[inline(always)]
|
|
pub fn handle_ephemeral_token(x: String) -> String {
|
|
#[cfg(feature = "enterprise")]
|
|
{
|
|
if let Some(full_cmd) = EPHEMERAL_TOKEN_CMD.as_ref() {
|
|
let mut splitted = full_cmd.split(" ");
|
|
let cmd = splitted.next().unwrap();
|
|
let args = splitted.collect::<Vec<&str>>();
|
|
let output = std::process::Command::new(cmd)
|
|
.args(args)
|
|
.output()
|
|
.map(|x| String::from_utf8(x.stdout).unwrap())
|
|
.unwrap_or_else(|e| panic!("failed to execute replace_ephemeral command: {}", e));
|
|
let r = x.replace("EPHEMERAL_TOKEN", &output.trim());
|
|
tracing::debug!("replaced ephemeral token: '{}'", r);
|
|
return r;
|
|
}
|
|
}
|
|
x
|
|
}
|
|
|
|
/// Removes lockfile/requirements entries matching the worker's `pip_local_dependencies`
|
|
/// regexes. Those packages are already provided locally (e.g. via `additional_python_paths`),
|
|
/// so installing them again duplicates files and triggers expensive `postinstall` copies on
|
|
/// every job. `#`-prefixed comment lines (e.g. the `# py:` lockfile header) are always kept.
|
|
/// Returns `(kept_lines, ignored_lines)`.
|
|
fn filter_pip_local_dependencies(lines: Vec<String>) -> (Vec<String>, Vec<String>) {
|
|
let Some(pip_local_dependencies) = WORKER_CONFIG.load().pip_local_dependencies.clone() else {
|
|
return (lines, vec![]);
|
|
};
|
|
|
|
let compiled_deps = pip_local_dependencies
|
|
.iter()
|
|
.filter_map(|dep| match Regex::new(dep) {
|
|
Ok(r) => Some(r),
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
"regex compilation failed for Python local dependency: '{}' - it will be ignored",
|
|
e
|
|
);
|
|
None
|
|
}
|
|
})
|
|
.collect::<Vec<Regex>>();
|
|
|
|
filter_lines_by_deps(lines, &compiled_deps)
|
|
}
|
|
|
|
/// Pure core of [`filter_pip_local_dependencies`]: partitions `lines` into
|
|
/// `(kept, ignored)`. A line is ignored when it is not a `#` comment and matches any of
|
|
/// `compiled_deps`. Kept separate from config/regex loading so it can be unit-tested.
|
|
fn filter_lines_by_deps(lines: Vec<String>, compiled_deps: &[Regex]) -> (Vec<String>, Vec<String>) {
|
|
let (ignored, kept): (Vec<String>, Vec<String>) = lines.into_iter().partition(|s| {
|
|
!s.trim_start().starts_with('#') && compiled_deps.iter().any(|dep| dep.is_match(s))
|
|
});
|
|
|
|
(kept, ignored)
|
|
}
|
|
|
|
// This function only invoked during deployment of script or test run.
|
|
// And never for already deployed scripts, these have their lockfiles in PostgreSQL
|
|
// thus this function call is skipped.
|
|
/// Returns lockfile and python version
|
|
pub async fn uv_pip_compile(
|
|
job_id: &Uuid,
|
|
requirements: &str,
|
|
mem_peak: &mut i32,
|
|
canceled_by: &mut Option<CanceledBy>,
|
|
job_dir: &str,
|
|
conn: &Connection,
|
|
worker_name: &str,
|
|
w_id: &str,
|
|
occupancy_metrics: &mut Option<&mut OccupancyMetrics>,
|
|
py_version: PyV,
|
|
// Debug-only flag
|
|
no_cache: bool,
|
|
) -> error::Result<String> {
|
|
let mut logs = String::new();
|
|
logs.push_str(&format!("\nresolving dependencies..."));
|
|
logs.push_str(&format!("\ncontent of requirements:\n{}\n", requirements));
|
|
|
|
let requirements = {
|
|
let (kept, ignored) =
|
|
filter_pip_local_dependencies(requirements.lines().map(str::to_owned).collect());
|
|
for line in ignored {
|
|
logs.push_str(&format!("\nignoring local dependency: {}", line));
|
|
}
|
|
kept.join("\n")
|
|
};
|
|
|
|
let uv_index_strategy = UV_INDEX_STRATEGY.read().await.clone();
|
|
let uv_index_strategy = uv_index_strategy.as_deref().unwrap_or("unsafe-best-match");
|
|
let uv_exclude_newer = (*UV_EXCLUDE_NEWER.read().await).map(|secs| format!("{}s", secs));
|
|
|
|
let py_version_str = py_version.clone().to_string();
|
|
// Include python version to requirements.in
|
|
// We need it because same hash based on requirements.in can get calculated even for different python versions
|
|
// To prevent from overwriting same requirements.in but with different python versions, we include version to hash
|
|
let requirements = format!("# py: {}\n{}", py_version.to_string(), requirements);
|
|
|
|
#[cfg(feature = "enterprise")]
|
|
let requirements = replace_pip_secret(conn, w_id, &requirements, worker_name, job_id).await?;
|
|
|
|
let ws_suffix = crate::workspace_registry_cache_suffix(w_id).await;
|
|
let exclude_newer_suffix = uv_exclude_newer
|
|
.as_deref()
|
|
.map(|v| format!("-en{}", v))
|
|
.unwrap_or_default();
|
|
let req_hash = format!(
|
|
"py-{}-{uv_index_strategy}{exclude_newer_suffix}{ws_suffix}",
|
|
calculate_hash(&requirements)
|
|
);
|
|
|
|
if !no_cache {
|
|
if let Some(db) = conn.as_sql() {
|
|
if let Some(cached) = sqlx::query_scalar!(
|
|
"SELECT lockfile FROM pip_resolution_cache WHERE hash = $1",
|
|
// Python version is included in hash,
|
|
// hash will be the different for every python version
|
|
req_hash
|
|
)
|
|
.fetch_optional(db)
|
|
.await?
|
|
{
|
|
logs.push_str(&format!(
|
|
"\nFound cached resolution: {req_hash}, on python version: {}",
|
|
&py_version_str
|
|
));
|
|
return Ok(cached);
|
|
}
|
|
}
|
|
}
|
|
|
|
let file = "requirements.in";
|
|
|
|
write_file(job_dir, file, &requirements)?;
|
|
|
|
{
|
|
// Make sure we have python runtime installed
|
|
py_version
|
|
.try_get_python(job_id, mem_peak, conn, worker_name, w_id, occupancy_metrics)
|
|
.await?;
|
|
|
|
let mut args = vec![
|
|
"pip",
|
|
"compile",
|
|
"-q",
|
|
"--no-header",
|
|
// The `#`-line filter applied to the output below only catches whole-line
|
|
// annotations, and uv's annotation style is configurable: `[pip]
|
|
// annotation-style = "line"` in the worker HOME's uv.toml emits them inline
|
|
// ("anyio==4.15.1 # via httpx"), which that filter keeps.
|
|
"--no-annotate",
|
|
file,
|
|
"--strip-extras",
|
|
"-o",
|
|
"requirements.txt",
|
|
// Target to /tmp/windmill/cache/uv
|
|
"--cache-dir",
|
|
&*UV_CACHE_DIR,
|
|
];
|
|
|
|
args.extend(["-p", &py_version_str, "--python-preference", "only-managed"]);
|
|
|
|
if no_cache {
|
|
args.extend(["--no-cache"]);
|
|
}
|
|
let pip_extra_index_url = read_ee_registry_with_workspace_override(
|
|
PIP_EXTRA_INDEX_URL.read().await.clone(),
|
|
"pip_extra_index_url",
|
|
"pip extra index url",
|
|
job_id,
|
|
w_id,
|
|
conn,
|
|
)
|
|
.await
|
|
.map(handle_ephemeral_token);
|
|
if let Some(url) = pip_extra_index_url.as_ref() {
|
|
url.split(",").for_each(|url| {
|
|
args.extend(["--extra-index-url", url]);
|
|
});
|
|
}
|
|
let pip_index_url = read_ee_registry_with_workspace_override(
|
|
PIP_INDEX_URL.read().await.clone(),
|
|
"pip_index_url",
|
|
"pip index url",
|
|
job_id,
|
|
w_id,
|
|
conn,
|
|
)
|
|
.await
|
|
.map(handle_ephemeral_token);
|
|
if let Some(url) = pip_index_url.as_ref() {
|
|
args.extend(["--index-url", url]);
|
|
}
|
|
if let Some(host) = TRUSTED_HOST.as_ref() {
|
|
host.split_whitespace().for_each(|h| {
|
|
args.extend(["--trusted-host", h]);
|
|
});
|
|
}
|
|
if let Some(cert_path) = INDEX_CERT.as_ref() {
|
|
args.extend(["--cert", cert_path]);
|
|
}
|
|
if *NATIVE_CERT {
|
|
args.extend(["--native-tls"]);
|
|
}
|
|
if let Some(exclude_newer) = uv_exclude_newer.as_deref() {
|
|
args.extend(["--exclude-newer", exclude_newer]);
|
|
}
|
|
tracing::debug!("uv args: {:?}", args);
|
|
|
|
#[cfg(windows)]
|
|
let uv_cmd = "uv";
|
|
|
|
#[cfg(unix)]
|
|
let uv_cmd = UV_PATH.as_str();
|
|
|
|
let mut child_cmd = Command::new(uv_cmd);
|
|
child_cmd
|
|
.current_dir(job_dir)
|
|
.env_clear()
|
|
.env("HOME", HOME_ENV.to_string())
|
|
.env("PATH", PATH_ENV.to_string())
|
|
.env("UV_PYTHON_INSTALL_DIR", PY_INSTALL_DIR.to_string())
|
|
.env("UV_INDEX_STRATEGY", uv_index_strategy)
|
|
.envs(PROXY_ENVS.clone())
|
|
.args(&args)
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
|
|
if let Some(mirror) = UV_PYTHON_INSTALL_MIRROR.read().await.as_ref() {
|
|
child_cmd.env("UV_PYTHON_INSTALL_MIRROR", mirror);
|
|
}
|
|
|
|
#[cfg(windows)]
|
|
{
|
|
child_cmd
|
|
.env("SystemRoot", SYSTEM_ROOT.as_str())
|
|
.env("USERPROFILE", crate::USERPROFILE_ENV.as_str())
|
|
.env("HOME", crate::USERPROFILE_ENV.as_str())
|
|
.env(
|
|
"LOCALAPPDATA",
|
|
std::env::var("LOCALAPPDATA")
|
|
.unwrap_or_else(|_| format!("{}\\AppData\\Local", HOME_ENV.as_str())),
|
|
)
|
|
.env(
|
|
"TMP",
|
|
std::env::var("TMP").unwrap_or_else(|_| String::from("/tmp")),
|
|
)
|
|
.env(
|
|
"APPDATA",
|
|
std::env::var("APPDATA").unwrap_or_else(|_| {
|
|
format!("{}\\AppData\\Roaming", crate::USERPROFILE_ENV.as_str())
|
|
}),
|
|
)
|
|
.env(
|
|
"ComSpec",
|
|
std::env::var("ComSpec")
|
|
.unwrap_or_else(|_| String::from("C:\\Windows\\System32\\cmd.exe")),
|
|
)
|
|
.env(
|
|
"PATHEXT",
|
|
std::env::var("PATHEXT").unwrap_or_else(|_| {
|
|
String::from(".COM;.EXE;.BAT;.CMD;.VBS;.VBE;.JS;.JSE;.WSF;.WSH;.MSC;.CPL")
|
|
}),
|
|
)
|
|
.env(
|
|
"ProgramData",
|
|
std::env::var("ProgramData")
|
|
.unwrap_or_else(|_| String::from("C:\\ProgramData")),
|
|
)
|
|
.env(
|
|
"ProgramFiles",
|
|
std::env::var("ProgramFiles")
|
|
.unwrap_or_else(|_| String::from("C:\\Program Files")),
|
|
);
|
|
}
|
|
|
|
let child_process = start_child_process(child_cmd, uv_cmd, false).await?;
|
|
append_logs(&job_id, &w_id, logs, conn).await;
|
|
handle_child(
|
|
job_id,
|
|
conn,
|
|
mem_peak,
|
|
canceled_by,
|
|
child_process,
|
|
false,
|
|
worker_name,
|
|
&w_id,
|
|
// TODO: Rename to uv-pip-compile?
|
|
"uv",
|
|
None,
|
|
false,
|
|
occupancy_metrics,
|
|
None,
|
|
None,
|
|
)
|
|
.await
|
|
.map_err(|e| {
|
|
Error::ExecutionErr(format!(
|
|
"Lock file generation failed.\n\ncommand: {uv_cmd} {}\n\n{e:?}",
|
|
args.join(" ")
|
|
))
|
|
})?;
|
|
}
|
|
|
|
let path_lock = format!("{job_dir}/requirements.txt");
|
|
let mut file = File::open(path_lock).await?;
|
|
let mut req_content = "".to_string();
|
|
file.read_to_string(&mut req_content).await?;
|
|
let lockfile = format!(
|
|
"# py: {}\n{}",
|
|
py_version.to_string(),
|
|
req_content
|
|
.lines()
|
|
.filter(|x| !x.trim_start().starts_with('#'))
|
|
.map(|x| x.to_string())
|
|
.collect::<Vec<String>>()
|
|
.join("\n")
|
|
);
|
|
if let Some(db) = conn.as_sql() {
|
|
sqlx::query!(
|
|
"INSERT INTO pip_resolution_cache (hash, lockfile, expiration) VALUES ($1, $2, now() + ('5 mins')::interval) ON CONFLICT (hash) DO UPDATE SET lockfile = EXCLUDED.lockfile",
|
|
req_hash,
|
|
lockfile
|
|
).fetch_optional(db).await?;
|
|
}
|
|
|
|
Ok(lockfile)
|
|
}
|
|
|
|
/**
|
|
Iterate over all python paths and if same folder has same name multiple times,
|
|
then merge the content and put to <job_dir>/site-packages
|
|
|
|
Solves problem with imports for some dependencies.
|
|
|
|
Default layout (/windmill/cache/):
|
|
|
|
dep==x.y.z
|
|
└── X
|
|
└── A
|
|
dep-ext==x.y.z
|
|
└── X
|
|
└── B
|
|
|
|
In this case python would be confused with finding B module.
|
|
|
|
This function will convert it to (/<job_dir>):
|
|
|
|
site-packages
|
|
└── X
|
|
├── A
|
|
└── B
|
|
|
|
This way python has no problems with finding correct module
|
|
*/
|
|
#[tracing::instrument(level = "trace", skip_all)]
|
|
async fn postinstall(
|
|
additional_python_paths: &mut Vec<String>,
|
|
job_dir: &str,
|
|
job: &MiniPulledJob,
|
|
conn: &Connection,
|
|
) -> windmill_common::error::Result<()> {
|
|
// It is guranteed that additional_python_paths only contains paths within windmill/cache/
|
|
// All other paths you would usually expect in PYTHONPATH are NOT included. These are added in downstream
|
|
//
|
|
// <PackageName, Vec<GlobalPath>>
|
|
let mut lookup_table: HashMap<String, Vec<String>> = HashMap::new();
|
|
// e.g.: <"requests", ["/tmp/windmill/cache/python_311/requests==1.0.0"]>
|
|
for path in additional_python_paths.iter() {
|
|
for entry in fs::read_dir(&path)? {
|
|
let entry = entry?;
|
|
// Ignore all files, we only need directories.
|
|
// We cannot merge files.
|
|
if entry.file_type()?.is_dir() {
|
|
// Short name, e.g.: requests
|
|
let name = entry
|
|
.file_name()
|
|
.to_str()
|
|
.ok_or(anyhow::anyhow!("Cannot convert OsString to String"))?
|
|
.to_owned();
|
|
|
|
if name == "bin" || name == "__pycache__" || name.contains("dist-info") {
|
|
continue;
|
|
}
|
|
|
|
if let Some(existing_paths) = lookup_table.get_mut(&name) {
|
|
tracing::debug!(
|
|
"Found existing package name: {:?} in {}",
|
|
entry.file_name(),
|
|
path
|
|
);
|
|
existing_paths.push(path.to_owned())
|
|
} else {
|
|
lookup_table.insert(name, vec![path.to_owned()]);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
let mut paths_to_remove: HashSet<String> = HashSet::new();
|
|
// Copy to shared dir
|
|
for existing_paths in lookup_table.values() {
|
|
if existing_paths.len() == 1 {
|
|
// There is only single path for given name
|
|
// So we skip it
|
|
continue;
|
|
}
|
|
|
|
for path in existing_paths {
|
|
copy_dir_recursively(
|
|
Path::new(path),
|
|
&std::path::PathBuf::from(job_dir).join("site-packages"),
|
|
)?;
|
|
paths_to_remove.insert(path.to_owned());
|
|
}
|
|
}
|
|
|
|
if !paths_to_remove.is_empty() {
|
|
append_logs(
|
|
&job.id,
|
|
&job.workspace_id,
|
|
"\n\nCopying some packages from cache to job_dir...\n".to_string(),
|
|
conn,
|
|
)
|
|
.await;
|
|
// Remove PATHs we just moved
|
|
additional_python_paths.retain(|e| !paths_to_remove.contains(e));
|
|
// Instead add shared path
|
|
additional_python_paths.insert(0, format!("{job_dir}/site-packages"));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Python hard keywords cannot be used as a bare name in `import <name>` /
|
|
/// `from <pkg> import <name>`. A flow inline step whose id (or a folder on its
|
|
/// path) is such a keyword — e.g. a step id `in` — otherwise generates
|
|
/// `from pkg import in as inner_script`, a SyntaxError. Prefix these with `_`,
|
|
/// mirroring the existing digit-leading guard.
|
|
fn is_python_keyword(s: &str) -> bool {
|
|
matches!(
|
|
s,
|
|
"False"
|
|
| "None"
|
|
| "True"
|
|
| "and"
|
|
| "as"
|
|
| "assert"
|
|
| "async"
|
|
| "await"
|
|
| "break"
|
|
| "class"
|
|
| "continue"
|
|
| "def"
|
|
| "del"
|
|
| "elif"
|
|
| "else"
|
|
| "except"
|
|
| "finally"
|
|
| "for"
|
|
| "from"
|
|
| "global"
|
|
| "if"
|
|
| "import"
|
|
| "in"
|
|
| "is"
|
|
| "lambda"
|
|
| "nonlocal"
|
|
| "not"
|
|
| "or"
|
|
| "pass"
|
|
| "raise"
|
|
| "return"
|
|
| "try"
|
|
| "while"
|
|
| "with"
|
|
| "yield"
|
|
)
|
|
}
|
|
|
|
/// Compute the directory (relative to job_dir) where Python writes the main script.
|
|
/// Module files must be placed in this same directory for relative imports to work.
|
|
pub fn compute_python_module_dir(script_path: &str) -> String {
|
|
let script_path_splitted = script_path.split("/").map(|x| {
|
|
if x.starts_with(|x: char| x.is_ascii_digit()) || is_python_keyword(x) {
|
|
format!("_{}", x)
|
|
} else {
|
|
x.to_string()
|
|
}
|
|
});
|
|
let dirs_full = script_path_splitted
|
|
.clone()
|
|
.take(script_path_splitted.clone().count() - 1)
|
|
.join("/")
|
|
.replace("-", "_")
|
|
.replace("@", ".");
|
|
if dirs_full.len() > 0 {
|
|
let dirs = dirs_full.strip_prefix("/").unwrap_or(&dirs_full);
|
|
// This directory is appended to job_dir and written to. Neutralize any
|
|
// `.`/`..` segment so the result stays a relative path inside job_dir: a
|
|
// Preview path is request-supplied and skips the DB `proper_id` CHECK that
|
|
// deployed runnables get, and the `@`->`.` rewrite above can also turn a
|
|
// segment like `@.` into `..`.
|
|
dirs.split('/')
|
|
.map(|seg| if seg == "." || seg == ".." { "_" } else { seg })
|
|
.collect::<Vec<_>>()
|
|
.join("/")
|
|
} else {
|
|
"tmp".to_string()
|
|
}
|
|
}
|
|
|
|
#[tracing::instrument(level = "trace", skip_all)]
|
|
pub async fn handle_python_job(
|
|
requirements_o: Option<&String>,
|
|
job_dir: &str,
|
|
worker_dir: &str,
|
|
worker_name: &str,
|
|
job: &MiniPulledJob,
|
|
mem_peak: &mut i32,
|
|
canceled_by: &mut Option<CanceledBy>,
|
|
conn: &Connection,
|
|
client: &AuthedClient,
|
|
parent_runnable_path: Option<String>,
|
|
inner_content: &String,
|
|
shared_mount: &str,
|
|
base_internal_url: &str,
|
|
envs: HashMap<String, String>,
|
|
new_args: &mut Option<HashMap<String, Box<RawValue>>>,
|
|
occupancy_metrics: &mut OccupancyMetrics,
|
|
precomputed_agent_info: Option<PrecomputedAgentInfo>,
|
|
has_stream: &mut bool,
|
|
modules: &Option<std::collections::HashMap<String, windmill_common::scripts::ScriptModule>>,
|
|
) -> windmill_common::error::Result<Box<RawValue>> {
|
|
let script_path = crate::common::use_flow_root_path(job.runnable_path());
|
|
|
|
let annotations = PythonAnnotations::parse(inner_content);
|
|
|
|
let is_wac_v2 = job.script_entrypoint_override.is_none()
|
|
&& crate::wac_executor::is_wac_v2_py(inner_content);
|
|
|
|
if annotations.sandbox && NSJAIL_AVAILABLE.is_none() {
|
|
return Err(Error::ExecutionErr(
|
|
"Script has #sandbox annotation but nsjail is not available on this worker. \
|
|
Please ensure nsjail is installed or remove the #sandbox annotation."
|
|
.to_string(),
|
|
));
|
|
}
|
|
|
|
// Preview jobs may carry _TEMP_SCRIPT_REFS so relative imports resolve from
|
|
// not-yet-deployed local content uploaded to raw_script_temp. Gated on
|
|
// JobKind::Preview because job.args includes caller-controlled request
|
|
// args; honoring this key on deployed runs would let a caller swap import
|
|
// resolution targets in deployed code.
|
|
let temp_script_refs: Option<HashMap<String, String>> = if matches!(job.kind, JobKind::Preview)
|
|
{
|
|
job.args
|
|
.as_ref()
|
|
.and_then(|x| x.get("_TEMP_SCRIPT_REFS"))
|
|
.and_then(|v| serde_json::from_str(v.get()).ok())
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let (py_version, mut additional_python_paths) = handle_python_deps(
|
|
job_dir,
|
|
requirements_o,
|
|
inner_content,
|
|
&job.workspace_id,
|
|
&script_path,
|
|
&job.id,
|
|
conn,
|
|
worker_name,
|
|
worker_dir,
|
|
mem_peak,
|
|
canceled_by,
|
|
&mut Some(occupancy_metrics),
|
|
precomputed_agent_info,
|
|
annotations.clone(),
|
|
&temp_script_refs,
|
|
)
|
|
.await?;
|
|
|
|
tracing::debug!("Finished handling python dependencies");
|
|
let python_path = py_version
|
|
.get_python(
|
|
worker_name,
|
|
&job.id,
|
|
&job.workspace_id,
|
|
mem_peak,
|
|
conn,
|
|
&mut Some(occupancy_metrics),
|
|
)
|
|
.await?;
|
|
|
|
if !annotations.no_postinstall {
|
|
if let Err(e) = postinstall(&mut additional_python_paths, job_dir, job, conn).await {
|
|
tracing::error!("Postinstall stage has failed. Reason: {e}");
|
|
}
|
|
tracing::debug!("Finished deps postinstall stage");
|
|
}
|
|
|
|
{
|
|
let mut logs = format!(
|
|
"\n\n--- PYTHON ({}) CODE EXECUTION ---\n",
|
|
py_version.clone().to_string()
|
|
);
|
|
if annotations.sandbox {
|
|
logs.push_str("sandbox mode (nsjail)\n");
|
|
}
|
|
append_logs(&job.id, &job.workspace_id, logs, conn).await;
|
|
}
|
|
let (
|
|
import_loader,
|
|
import_base64,
|
|
import_datetime,
|
|
module_dir_dot,
|
|
dirs,
|
|
last,
|
|
transforms,
|
|
spread,
|
|
main_name,
|
|
pre_spread,
|
|
wac_pre_spread,
|
|
) = prepare_wrapper(
|
|
job_dir,
|
|
job.flow_step_id.as_deref(),
|
|
job.preprocessed,
|
|
job.script_entrypoint_override.as_deref(),
|
|
inner_content,
|
|
&script_path,
|
|
&temp_script_refs,
|
|
)
|
|
.await?;
|
|
|
|
tracing::debug!("Finished preparing wrapper");
|
|
|
|
let apply_preprocessor = pre_spread.is_some();
|
|
|
|
create_args_and_out_file(&client, job, job_dir, conn).await?;
|
|
tracing::debug!("Finished preparing wrapper");
|
|
|
|
let preprocessor = if let Some(pre_spread) = pre_spread {
|
|
format!(
|
|
r#"if inner_script.preprocessor is None or not callable(inner_script.preprocessor):
|
|
raise ValueError("preprocessor function is missing")
|
|
else:
|
|
pre_args = {{}}
|
|
{pre_spread}
|
|
for k, v in list(pre_args.items()):
|
|
if v == '<function call>':
|
|
del pre_args[k]
|
|
kwargs = inner_script.preprocessor(**pre_args)
|
|
kwrags_json = res_to_json(kwargs, type(kwargs))
|
|
with open("args.json", 'w', encoding="utf-8") as f:
|
|
f.write(kwrags_json)"#
|
|
)
|
|
} else {
|
|
"".to_string()
|
|
};
|
|
|
|
let postprocessor = get_result_postprocessor(annotations.skip_result_postprocessing);
|
|
|
|
let os_main_override = if let Some(main_override) = main_name.as_ref() {
|
|
format!("os.environ[\"MAIN_OVERRIDE\"] = \"{main_override}\"\n")
|
|
} else {
|
|
String::new()
|
|
};
|
|
|
|
let main_override = main_name.unwrap_or_else(|| "main".to_string());
|
|
let res_to_json_body = python_res_to_json_body(postprocessor);
|
|
// Indented at 4 spaces so it can be inlined inside the wrapper's
|
|
// `try:` block — preprocessor failures route through the same error
|
|
// serializer as workflow failures.
|
|
let wac_preprocessor = if let Some(pre_spread) = wac_pre_spread.as_ref() {
|
|
format!(
|
|
r#"if not hasattr(inner_script, 'preprocessor') or not callable(inner_script.preprocessor):
|
|
raise ValueError("preprocessor function is missing")
|
|
pre_args = {{}}
|
|
{pre_spread}
|
|
for k, v in list(pre_args.items()):
|
|
if v == '<function call>':
|
|
del pre_args[k]
|
|
_pre_result = inner_script.preprocessor(**pre_args)
|
|
if hasattr(_pre_result, '__await__'):
|
|
import asyncio
|
|
_pre_result = asyncio.run(_pre_result)
|
|
kwargs = _pre_result if _pre_result is not None else {{}}
|
|
_pre_json = json.dumps(kwargs, separators=(',', ':'), default=str)
|
|
with open("args.json", 'w', encoding="utf-8") as f:
|
|
f.write(_pre_json)
|
|
sys.stdout.write("wm_res[preprocessed_args]:" + _pre_json + "\n")
|
|
sys.stdout.flush()"#
|
|
)
|
|
} else {
|
|
"".to_string()
|
|
};
|
|
let wrapper_content: String = if is_wac_v2 {
|
|
format!(
|
|
r#"
|
|
import os
|
|
import json
|
|
{import_loader}
|
|
{import_base64}
|
|
{import_datetime}
|
|
import traceback
|
|
import sys
|
|
from {module_dir_dot} import {last} as inner_script
|
|
from wmill.client import _run_workflow
|
|
|
|
with open("args.json", encoding="utf-8") as f:
|
|
kwargs = json.load(f, strict=False)
|
|
{transforms}
|
|
|
|
with open("checkpoint.json", encoding="utf-8") as f:
|
|
checkpoint = json.load(f, strict=False)
|
|
|
|
result_json = os.path.join(os.path.abspath(os.path.dirname(__file__)), "result.json")
|
|
|
|
try:
|
|
{wac_preprocessor}
|
|
args = kwargs
|
|
for k, v in list(args.items()):
|
|
if v == '<function call>':
|
|
del args[k]
|
|
|
|
workflow_fn = None
|
|
for name in dir(inner_script):
|
|
obj = getattr(inner_script, name)
|
|
if callable(obj) and getattr(obj, '_is_workflow', False):
|
|
workflow_fn = obj
|
|
break
|
|
if workflow_fn is None:
|
|
raise ValueError("No @workflow function found in script")
|
|
|
|
output = _run_workflow(workflow_fn, checkpoint, args)
|
|
if isinstance(output, dict) and output.get("type") == "complete":
|
|
print("")
|
|
print("--- WAC: complete ---")
|
|
output_json = json.dumps(output, separators=(',', ':'), default=str)
|
|
with open(result_json, 'w', encoding="utf-8") as f:
|
|
f.write(output_json)
|
|
except BaseException as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
tb = traceback.format_tb(exc_traceback)
|
|
with open(result_json, 'w', encoding="utf-8") as f:
|
|
err = {{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}
|
|
extra = e.__dict__
|
|
if extra and len(extra) > 0:
|
|
err['extra'] = extra
|
|
flow_node_id = os.environ.get('WM_FLOW_STEP_ID')
|
|
if flow_node_id:
|
|
err['step_id'] = flow_node_id
|
|
err_json = json.dumps(err, separators=(',', ':'), default=str).replace('\n', '')
|
|
f.write(err_json)
|
|
sys.exit(1)
|
|
"#,
|
|
)
|
|
} else {
|
|
format!(
|
|
r#"
|
|
import os
|
|
import json
|
|
{import_loader}
|
|
{import_base64}
|
|
{import_datetime}
|
|
import traceback
|
|
import sys
|
|
{os_main_override}
|
|
from {module_dir_dot} import {last} as inner_script
|
|
import re
|
|
|
|
with open("args.json", encoding="utf-8") as f:
|
|
kwargs = json.load(f, strict=False)
|
|
args = {{}}
|
|
{transforms}
|
|
|
|
def to_b_64(v: bytes):
|
|
import base64
|
|
b64 = base64.b64encode(v)
|
|
return b64.decode('ascii')
|
|
|
|
_u=re.compile(r'\\\\|\\u0000')
|
|
_us=lambda m:' null ' if m.group(0)[1]=='u' else m.group(0)
|
|
_r=lambda m,s='':(_u.sub(_us,s) if '\\u0000' in s else s) if (s:=m.group(0))[0]=='"' else ' null '
|
|
replace_invalid_fields=re.compile(r'"(?:\\.|[^"\\])*"|\bNaN\b|-?Infinity')
|
|
_fix=lambda s:s if 'Infinity' not in s and 'NaN' not in s and '\\u0000' not in s else re.sub(replace_invalid_fields,_r,s)
|
|
|
|
result_json = os.path.join(os.path.abspath(os.path.dirname(__file__)), "result.json")
|
|
|
|
def res_to_json(res, typ):
|
|
{res_to_json_body}
|
|
|
|
try:
|
|
{preprocessor}
|
|
{spread}
|
|
for k, v in list(args.items()):
|
|
if v == '<function call>':
|
|
del args[k]
|
|
if inner_script.{main_override} is None or not callable(inner_script.{main_override}):
|
|
raise ValueError("{main_override} function is missing")
|
|
res = inner_script.{main_override}(**args)
|
|
typ = type(res)
|
|
if hasattr(res, '__iter__') and not isinstance(res, (str, dict, list, bytes, tuple, set, frozenset, range, memoryview, bytearray)) and typ.__name__ != 'DataFrame':
|
|
for chunk in res:
|
|
print("WM_STREAM: " + chunk.replace('\n', '\\n'))
|
|
res = None
|
|
res_json = res_to_json(res, typ)
|
|
with open(result_json, 'w', encoding="utf-8") as f:
|
|
f.write(res_json)
|
|
except BaseException as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
tb = traceback.format_tb(exc_traceback)
|
|
with open(result_json, 'w', encoding="utf-8") as f:
|
|
err = {{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}
|
|
extra = e.__dict__
|
|
if extra and len(extra) > 0:
|
|
err['extra'] = extra
|
|
flow_node_id = os.environ.get('WM_FLOW_STEP_ID')
|
|
if flow_node_id:
|
|
err['step_id'] = flow_node_id
|
|
err_json = json.dumps(err, separators=(',', ':'), default=str).replace('\n', '')
|
|
f.write(err_json)
|
|
sys.exit(1)
|
|
"#,
|
|
)
|
|
};
|
|
write_file(job_dir, "wrapper.py", &wrapper_content)?;
|
|
|
|
// For WAC v2, write checkpoint.json before python runs.
|
|
if is_wac_v2 {
|
|
if let Connection::Sql(db) = conn {
|
|
let checkpoint = crate::wac_executor::load_checkpoint(db, &job.id).await?;
|
|
let checkpoint =
|
|
crate::wac_executor::prepare_checkpoint_for_resume(db, &job.id, checkpoint).await?;
|
|
|
|
let checkpoint_json = serde_json::to_string(&checkpoint).map_err(|e| {
|
|
error::Error::internal_err(format!("Failed to serialize checkpoint: {e}"))
|
|
})?;
|
|
write_file(job_dir, "checkpoint.json", &checkpoint_json)?;
|
|
} else {
|
|
write_file(job_dir, "checkpoint.json", r#"{"completed_steps":{}}"#)?;
|
|
}
|
|
}
|
|
|
|
tracing::debug!("Finished writing wrapper");
|
|
|
|
let mut reserved_variables =
|
|
get_reserved_variables(job, &client.token, conn, parent_runnable_path).await?;
|
|
|
|
// Add /tmp/windmill/cache/python_x_y_z/global-site-packages to PYTHONPATH.
|
|
// Usefull if certain wheels needs to be preinstalled before execution.
|
|
let global_site_packages_path = py_version.to_cache_dir(true) + "/global-site-packages";
|
|
let additional_python_paths_folders = {
|
|
let mut paths = additional_python_paths.clone();
|
|
if std::fs::metadata(&global_site_packages_path).is_ok() {
|
|
// We want global_site_packages_path to be included in additonal_python_paths_folders, but
|
|
// we don't want it to be included in global_site_packages_path.
|
|
// The reason for this is that additional_python_paths_folders is used to fill PYTHONPATH env variable for jailed script
|
|
// When global_site_packages_path used to place mount point of wheels to the jail config.
|
|
// Since we handle mount of global_site_packages on our own, we don't want it to be mounted automatically.
|
|
// We do this because existence of every wheel in cache is mandatory and if it is not there and nsjail expects it, it is a bug.
|
|
// On the other side global_site_packages is purely optional.
|
|
// NOTE: This behaviour can be changed in future, so verification of wheels can be delegated from nsjail to windmill
|
|
paths.insert(0, global_site_packages_path.clone());
|
|
// ^^^^^^ ^
|
|
// We also want this be priorotized, that's why we insert it to the beginning
|
|
}
|
|
#[cfg(windows)]
|
|
{
|
|
paths.iter().join(";")
|
|
}
|
|
#[cfg(not(windows))]
|
|
{
|
|
paths.iter().join(":")
|
|
}
|
|
};
|
|
|
|
#[cfg(windows)]
|
|
let additional_python_paths_folders = additional_python_paths_folders.replace(":", ";");
|
|
|
|
if is_sandboxing_enabled() || annotations.sandbox {
|
|
let shared_deps = additional_python_paths
|
|
.into_iter()
|
|
.map(|pp| {
|
|
format!(
|
|
r#"
|
|
mount {{
|
|
src: "{pp}"
|
|
dst: "{pp}"
|
|
is_bind: true
|
|
rw: false
|
|
}}
|
|
"#
|
|
)
|
|
})
|
|
.join("\n");
|
|
let nsjail_timeout =
|
|
resolve_nsjail_timeout(conn, &job.workspace_id, job.id, job.timeout).await;
|
|
let _ = write_file(
|
|
job_dir,
|
|
"run.config.proto",
|
|
&NSJAIL_CONFIG_RUN_PYTHON3_CONTENT
|
|
.replace(
|
|
"{RLIMIT_AS}",
|
|
&render_nsjail_rlimit_as(NSJAIL_PY_RLIMIT_AS_MB.as_deref(), 4096),
|
|
)
|
|
.replace("{JOB_DIR}", job_dir)
|
|
.replace("{PY_INSTALL_DIR}", &*PY_INSTALL_DIR)
|
|
.replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string())
|
|
.replace("{SHARED_MOUNT}", shared_mount)
|
|
.replace("{SHARED_DEPENDENCIES}", shared_deps.as_str())
|
|
.replace("{GLOBAL_SITE_PACKAGES}", &global_site_packages_path)
|
|
.replace("{MAIN}", format!("{dirs}/{last}").as_str())
|
|
.replace(
|
|
"{ADDITIONAL_PYTHON_PATHS}",
|
|
additional_python_paths_folders.as_str(),
|
|
)
|
|
.replace("{TRACING_PROXY_CA_CERT_PATH}", &*TRACING_PROXY_CA_CERT_PATH)
|
|
.replace("#{DEV}", DEV_CONF_NSJAIL)
|
|
.replace(
|
|
"{TMP_MOUNT_BLOCK}",
|
|
&resolve_nsjail_tmp_mount_block(job_dir).await,
|
|
)
|
|
.replace("{TIMEOUT}", &nsjail_timeout),
|
|
)?;
|
|
} else {
|
|
reserved_variables.insert("PYTHONPATH".to_string(), additional_python_paths_folders);
|
|
}
|
|
|
|
tracing::info!(
|
|
workspace_id = %job.workspace_id,
|
|
"started python code execution {}",
|
|
job.id
|
|
);
|
|
|
|
let child = if is_sandboxing_enabled() || annotations.sandbox {
|
|
let mut nsjail_cmd = Command::new(NSJAIL_PATH.as_str());
|
|
nsjail_cmd
|
|
.current_dir(job_dir)
|
|
.env_clear()
|
|
// inject PYTHONPATH here - for some reason I had to do it in nsjail conf
|
|
.envs(reserved_variables)
|
|
.envs(
|
|
get_proxy_envs_for_lang(
|
|
&ScriptLang::Python3,
|
|
job.kind,
|
|
&job.id,
|
|
&job.workspace_id,
|
|
conn,
|
|
)
|
|
.await?,
|
|
)
|
|
.envs(PYTHON_UTF8_ENVS)
|
|
.env("PATH", PATH_ENV.as_str())
|
|
.env("TZ", TZ_ENV.as_str())
|
|
.env("BASE_INTERNAL_URL", base_internal_url)
|
|
.env("BASE_URL", base_internal_url)
|
|
.args(vec![
|
|
"--config",
|
|
"run.config.proto",
|
|
"--",
|
|
&python_path,
|
|
"-u",
|
|
"-m",
|
|
"wrapper",
|
|
])
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
start_child_process(nsjail_cmd, NSJAIL_PATH.as_str(), false).await?
|
|
} else {
|
|
let args = vec!["-u", "-m", "wrapper"];
|
|
|
|
let mut python_cmd = build_command_with_isolation(&python_path, &args);
|
|
python_cmd
|
|
.current_dir(job_dir)
|
|
.env_clear()
|
|
.envs(envs)
|
|
.envs(reserved_variables)
|
|
.envs(
|
|
get_proxy_envs_for_lang(
|
|
&ScriptLang::Python3,
|
|
job.kind,
|
|
&job.id,
|
|
&job.workspace_id,
|
|
conn,
|
|
)
|
|
.await?,
|
|
)
|
|
.envs(PYTHON_UTF8_ENVS)
|
|
.env("PATH", PATH_ENV.as_str())
|
|
.env("TZ", TZ_ENV.as_str())
|
|
.env("BASE_INTERNAL_URL", base_internal_url)
|
|
.env("HOME", HOME_ENV.as_str())
|
|
.stdin(Stdio::null())
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
|
|
#[cfg(windows)]
|
|
{
|
|
python_cmd.env("SystemRoot", SYSTEM_ROOT.as_str());
|
|
python_cmd.env("USERPROFILE", crate::USERPROFILE_ENV.as_str());
|
|
python_cmd.env("windir", SYSTEM_ROOT.as_str());
|
|
python_cmd.env(
|
|
"LOCALAPPDATA",
|
|
std::env::var("LOCALAPPDATA")
|
|
.unwrap_or_else(|_| format!("{}\\AppData\\Local", HOME_ENV.as_str())),
|
|
);
|
|
}
|
|
|
|
start_child_process(python_cmd, &python_path, false).await?
|
|
};
|
|
|
|
let stream_notifier = StreamNotifier::new(conn, job);
|
|
|
|
let handle_result = handle_child(
|
|
&job.id,
|
|
conn,
|
|
mem_peak,
|
|
canceled_by,
|
|
child,
|
|
is_sandboxing_enabled(),
|
|
worker_name,
|
|
&job.workspace_id,
|
|
"python run",
|
|
job.timeout,
|
|
false,
|
|
&mut Some(occupancy_metrics),
|
|
None,
|
|
stream_notifier,
|
|
)
|
|
.await?;
|
|
|
|
*has_stream = handle_result.result_stream.is_some();
|
|
|
|
if apply_preprocessor {
|
|
let args = read_file(&format!("{job_dir}/args.json"))
|
|
.await
|
|
.map_err(|e| {
|
|
error::Error::internal_err(format!(
|
|
"error while reading args from preprocessing: {e:#}"
|
|
))
|
|
})?;
|
|
let args: HashMap<String, Box<RawValue>> =
|
|
serde_json::from_str(args.get()).map_err(|e| {
|
|
error::Error::internal_err(format!(
|
|
"error while deserializing args from preprocessing: {e:#}"
|
|
))
|
|
})?;
|
|
*new_args = Some(args.clone());
|
|
}
|
|
|
|
let result = read_result(job_dir, handle_result.result_stream).await?;
|
|
|
|
// WAC v2 post-execution: parse output and handle dispatch/suspend.
|
|
// Box::pin to avoid bloating handle_python_job's async state machine (stack overflow).
|
|
if is_wac_v2 {
|
|
return Box::pin(crate::bun_executor::handle_wac_v2_output(
|
|
result,
|
|
job,
|
|
conn,
|
|
canceled_by,
|
|
modules,
|
|
new_args.as_ref(),
|
|
))
|
|
.await;
|
|
}
|
|
|
|
Ok(result)
|
|
}
|
|
|
|
/// Generate Python code to spread preprocessor args from `kwargs` into `pre_args`.
|
|
/// The `indent` parameter controls the join separator for multi-arg spreads.
|
|
fn python_preprocessor_spread(sig: &windmill_parser::MainArgSignature, indent: &str) -> String {
|
|
if sig.star_kwargs {
|
|
"pre_args = kwargs".to_string()
|
|
} else {
|
|
sig.args
|
|
.iter()
|
|
.map(|x| {
|
|
let name = &x.name;
|
|
if x.default.is_none() {
|
|
format!("pre_args[\"{name}\"] = kwargs.get(\"{name}\")")
|
|
} else {
|
|
format!(
|
|
r#"pre_args["{name}"] = kwargs.get("{name}")
|
|
{indent}if pre_args["{name}"] is None:
|
|
{indent} del pre_args["{name}"]"#
|
|
)
|
|
}
|
|
})
|
|
.join(&format!("\n{indent}"))
|
|
}
|
|
}
|
|
|
|
/// Pre-computed codegen data for a Python script.
|
|
/// Computed in Rust from the parsed signature, then baked into the wrapper template.
|
|
#[cfg(any(feature = "private", test))]
|
|
pub struct PyScriptCodegen {
|
|
/// Python module directory dot notation (e.g., "f.my")
|
|
pub module_dir_dot: String,
|
|
/// Python module name / last path component (e.g., "script")
|
|
pub module_name: String,
|
|
/// Directory path for the module (e.g., "f/my")
|
|
pub dirs: String,
|
|
/// Inline Python code for type transforms (dates, bytes, etc.)
|
|
pub transforms: String,
|
|
/// Inline Python code for arg spread / filtering
|
|
pub spread: String,
|
|
/// Inline Python code for preprocessor arg spread (if applicable)
|
|
pub pre_spread: Option<String>,
|
|
}
|
|
|
|
/// Parse a Python script and compute the codegen data.
|
|
/// This reuses the same logic that was used on main in `prepare_wrapper`.
|
|
#[cfg(any(feature = "private", test))]
|
|
pub fn compute_py_codegen(content: &str, script_path: &str) -> PyScriptCodegen {
|
|
let dirs = compute_python_module_dir(script_path);
|
|
let last = script_path
|
|
.split("/")
|
|
.map(|x| {
|
|
if x.starts_with(|x: char| x.is_ascii_digit()) {
|
|
format!("_{}", x)
|
|
} else {
|
|
x.to_string()
|
|
}
|
|
})
|
|
.last()
|
|
.unwrap()
|
|
.replace("-", "_")
|
|
.replace(" ", "_")
|
|
.to_lowercase();
|
|
// `last` is lowercased above, so this catches a keyword id in any case.
|
|
let last = if is_python_keyword(&last) {
|
|
format!("_{last}")
|
|
} else {
|
|
last
|
|
};
|
|
|
|
let sig = windmill_parser_py::parse_python_signature(content, None, false).unwrap_or_default();
|
|
let pre_sig = windmill_parser_py::parse_python_signature(
|
|
content,
|
|
Some("preprocessor".to_string()),
|
|
false,
|
|
)
|
|
.ok()
|
|
.filter(|s| !s.args.is_empty());
|
|
|
|
let init_sig = pre_sig.as_ref().unwrap_or(&sig);
|
|
|
|
let transforms = init_sig
|
|
.args
|
|
.iter()
|
|
.map(|x| match x.typ {
|
|
windmill_parser::Typ::Bytes => {
|
|
let name = &x.name;
|
|
format!(
|
|
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
|
kwargs[\"{name}\"] = base64.b64decode(kwargs[\"{name}\"])\n",
|
|
)
|
|
}
|
|
windmill_parser::Typ::Datetime => {
|
|
let name = &x.name;
|
|
format!(
|
|
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
|
kwargs[\"{name}\"] = datetime.fromisoformat(kwargs[\"{name}\"])\n",
|
|
)
|
|
}
|
|
windmill_parser::Typ::Date => {
|
|
let name = &x.name;
|
|
format!(
|
|
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
|
try:\n \
|
|
kwargs[\"{name}\"] = date.fromisoformat(kwargs[\"{name}\"])\n \
|
|
except ValueError:\n \
|
|
for _fmt in (\"%d-%m-%Y\", \"%m/%d/%Y\", \"%d/%m/%Y\", \"%Y/%m/%d\"):\n \
|
|
try:\n \
|
|
kwargs[\"{name}\"] = datetime.strptime(kwargs[\"{name}\"], _fmt).date()\n \
|
|
break\n \
|
|
except ValueError:\n \
|
|
continue\n",
|
|
)
|
|
}
|
|
_ => "".to_string(),
|
|
})
|
|
.collect::<Vec<String>>()
|
|
.join("");
|
|
|
|
let spread = if sig.star_kwargs {
|
|
"args = kwargs".to_string()
|
|
} else {
|
|
sig.args
|
|
.into_iter()
|
|
.map(|x| {
|
|
let name = &x.name;
|
|
if x.default.is_none() {
|
|
format!("args[\"{name}\"] = kwargs.get(\"{name}\")")
|
|
} else {
|
|
format!(
|
|
r#"args["{name}"] = kwargs.get("{name}")
|
|
if args["{name}"] is None:
|
|
del args["{name}"]"#
|
|
)
|
|
}
|
|
})
|
|
.join("\n ")
|
|
};
|
|
|
|
let pre_spread = pre_sig.map(|sig| python_preprocessor_spread(&sig, " "));
|
|
|
|
let module_dir_dot = dirs.replace("/", ".").replace("-", "_");
|
|
|
|
PyScriptCodegen { module_dir_dot, module_name: last, dirs, transforms, spread, pre_spread }
|
|
}
|
|
|
|
/// Script entry for the Python unified wrapper generator.
|
|
#[cfg(any(feature = "private", test))]
|
|
pub struct PyScriptEntry<'a> {
|
|
pub original_path: &'a str,
|
|
pub codegen: &'a PyScriptCodegen,
|
|
}
|
|
|
|
/// Generate a wrapper for Python dedicated workers and runner groups.
|
|
/// All scripts are baked in at codegen time with proper Python imports and inline arg handling.
|
|
/// Protocol:
|
|
/// execd:<json_args> -> execute the single registered script (non-runner-group)
|
|
/// exec:<path>:<json_args> -> execute script by path (runner groups)
|
|
/// end -> exit
|
|
#[cfg(any(feature = "private", test))]
|
|
pub fn generate_multi_script_wrapper(
|
|
scripts: &[PyScriptEntry<'_>],
|
|
skip_result_postprocessing: bool,
|
|
any_relative_imports: bool,
|
|
) -> String {
|
|
let postprocessor = get_result_postprocessor(skip_result_postprocessing);
|
|
let res_to_json_body = python_res_to_json_body(postprocessor);
|
|
|
|
let imports: String = scripts
|
|
.iter()
|
|
.enumerate()
|
|
.map(|(i, e)| {
|
|
format!(
|
|
"from {module_dir_dot} import {module_name} as _s{i}",
|
|
module_dir_dot = e.codegen.module_dir_dot,
|
|
module_name = e.codegen.module_name,
|
|
)
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
|
|
let mut functions = String::new();
|
|
let mut registrations = String::new();
|
|
|
|
for (i, entry) in scripts.iter().enumerate() {
|
|
let cg = entry.codegen;
|
|
let indented_transforms = cg
|
|
.transforms
|
|
.split('\n')
|
|
.map(|line| {
|
|
if line.is_empty() {
|
|
String::new()
|
|
} else {
|
|
format!(" {}", line)
|
|
}
|
|
})
|
|
.collect::<Vec<_>>()
|
|
.join("\n");
|
|
|
|
functions.push_str(&format!(
|
|
r#"
|
|
def transform_{i}(kwargs):
|
|
{indented_transforms}
|
|
args = dict()
|
|
{spread}
|
|
for k, v in list(args.items()):
|
|
if v == '<function call>':
|
|
del args[k]
|
|
return args
|
|
"#,
|
|
spread = cg.spread
|
|
));
|
|
|
|
let pre_fn = if let Some(ref pre_spread) = cg.pre_spread {
|
|
functions.push_str(&format!(
|
|
r#"
|
|
def pre_transform_{i}(kwargs):
|
|
pre_args = dict()
|
|
{pre_spread}
|
|
for k, v in list(pre_args.items()):
|
|
if v == '<function call>':
|
|
del pre_args[k]
|
|
return pre_args
|
|
"#,
|
|
));
|
|
format!("pre_transform_{i}")
|
|
} else {
|
|
"None".to_string()
|
|
};
|
|
|
|
registrations.push_str(&format!(
|
|
"scripts[\"{path}\"] = {{ 'mod': _s{i}, 'transform': transform_{i}, 'pre_transform': {pre_fn} }}\n",
|
|
path = entry.original_path,
|
|
));
|
|
}
|
|
|
|
let import_loader = if any_relative_imports {
|
|
"import loader"
|
|
} else {
|
|
""
|
|
};
|
|
|
|
format!(
|
|
r#"
|
|
import json
|
|
import sys
|
|
import traceback
|
|
import re
|
|
import base64
|
|
from datetime import datetime, date
|
|
{import_loader}
|
|
|
|
{imports}
|
|
|
|
scripts = {{}}
|
|
|
|
def to_b_64(v: bytes):
|
|
b64 = base64.b64encode(v)
|
|
return b64.decode('ascii')
|
|
|
|
_u=re.compile(r'\\\\|\\u0000')
|
|
_us=lambda m:' null ' if m.group(0)[1]=='u' else m.group(0)
|
|
_r=lambda m,s='':(_u.sub(_us,s) if '\\u0000' in s else s) if (s:=m.group(0))[0]=='"' else ' null '
|
|
replace_invalid_fields=re.compile(r'"(?:\\.|[^"\\])*"|\bNaN\b|-?Infinity')
|
|
_fix=lambda s:s if 'Infinity' not in s and 'NaN' not in s and '\\u0000' not in s else re.sub(replace_invalid_fields,_r,s)
|
|
|
|
def res_to_json(res, typ):
|
|
{res_to_json_body}
|
|
{functions}
|
|
{registrations}
|
|
|
|
sys.stdout.write('start\n')
|
|
sys.stdout.flush()
|
|
|
|
for line in sys.stdin:
|
|
line = line.strip()
|
|
if line == 'end':
|
|
break
|
|
|
|
if line.startswith('exec_preprocess:'):
|
|
try:
|
|
rest = line[len('exec_preprocess:'):]
|
|
colon_idx = rest.index(':')
|
|
script_path = rest[:colon_idx]
|
|
args_json = rest[colon_idx + 1:]
|
|
|
|
entry = scripts.get(script_path)
|
|
if not entry:
|
|
err_json = json.dumps({{ "message": "Script not found: " + script_path, "name": "Error" }}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
|
|
mod = entry['mod']
|
|
if not hasattr(mod, 'preprocessor') or not callable(mod.preprocessor):
|
|
err_json = json.dumps({{"message": "preprocessor function is missing", "name": "Error"}}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
kwargs = json.loads(args_json, strict=False)
|
|
pre_args = entry['pre_transform'](kwargs)
|
|
preprocessed = mod.preprocessor(**pre_args)
|
|
preprocessed_json = json.dumps(preprocessed, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[preprocessed_args]:" + preprocessed_json + "\n")
|
|
main_args = entry['transform'](preprocessed if preprocessed else {{}})
|
|
res = mod.main(**main_args)
|
|
typ = type(res)
|
|
res_json = res_to_json(res, typ)
|
|
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
|
except BaseException as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
tb = traceback.format_tb(exc_traceback)
|
|
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
|
|
if line.startswith('execd_preprocess:'):
|
|
try:
|
|
args_json = line[len('execd_preprocess:'):]
|
|
entry = next(iter(scripts.values()))
|
|
mod = entry['mod']
|
|
if not hasattr(mod, 'preprocessor') or not callable(mod.preprocessor):
|
|
err_json = json.dumps({{"message": "preprocessor function is missing", "name": "Error"}}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
kwargs = json.loads(args_json, strict=False)
|
|
pre_args = entry['pre_transform'](kwargs)
|
|
preprocessed = mod.preprocessor(**pre_args)
|
|
preprocessed_json = json.dumps(preprocessed, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[preprocessed_args]:" + preprocessed_json + "\n")
|
|
main_args = entry['transform'](preprocessed if preprocessed else {{}})
|
|
res = mod.main(**main_args)
|
|
typ = type(res)
|
|
res_json = res_to_json(res, typ)
|
|
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
|
except BaseException as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
tb = traceback.format_tb(exc_traceback)
|
|
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
|
|
if line.startswith('execd:'):
|
|
try:
|
|
args_json = line[len('execd:'):]
|
|
entry = next(iter(scripts.values()))
|
|
kwargs = json.loads(args_json, strict=False)
|
|
args = entry['transform'](kwargs)
|
|
res = entry['mod'].main(**args)
|
|
typ = type(res)
|
|
res_json = res_to_json(res, typ)
|
|
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
|
except BaseException as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
tb = traceback.format_tb(exc_traceback)
|
|
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
|
|
if line.startswith('exec:'):
|
|
try:
|
|
rest = line[len('exec:'):]
|
|
colon_idx = rest.index(':')
|
|
script_path = rest[:colon_idx]
|
|
args_json = rest[colon_idx + 1:]
|
|
|
|
entry = scripts.get(script_path)
|
|
if not entry:
|
|
err_json = json.dumps({{ "message": "Script not found: " + script_path, "name": "Error" }}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
|
|
kwargs = json.loads(args_json, strict=False)
|
|
args = entry['transform'](kwargs)
|
|
res = entry['mod'].main(**args)
|
|
typ = type(res)
|
|
res_json = res_to_json(res, typ)
|
|
sys.stdout.write("wm_res[success]:" + res_json + "\n")
|
|
except BaseException as e:
|
|
exc_type, exc_value, exc_traceback = sys.exc_info()
|
|
tb = traceback.format_tb(exc_traceback)
|
|
err_json = json.dumps({{ "message": str(e), "name": e.__class__.__name__, "stack": '\n'.join(tb[1:]) }}, separators=(',', ':'), default=str).replace('\n', '')
|
|
sys.stdout.write("wm_res[error]:" + err_json + "\n")
|
|
sys.stdout.flush()
|
|
continue
|
|
|
|
sys.stderr.write("Unknown command: " + line + "\n")
|
|
"#
|
|
)
|
|
}
|
|
|
|
async fn prepare_wrapper(
|
|
job_dir: &str,
|
|
job_flow_step_id: Option<&str>,
|
|
job_preprocessed: Option<bool>,
|
|
job_script_entrypoint_override: Option<&str>,
|
|
inner_content: &str,
|
|
script_path: &str,
|
|
temp_script_refs: &Option<HashMap<String, String>>,
|
|
) -> error::Result<(
|
|
&'static str,
|
|
&'static str,
|
|
&'static str,
|
|
String,
|
|
String,
|
|
String,
|
|
String,
|
|
String,
|
|
Option<String>,
|
|
Option<String>,
|
|
Option<String>,
|
|
)> {
|
|
let main_override = job_script_entrypoint_override.as_deref();
|
|
let apply_preprocessor =
|
|
job_flow_step_id != Some("preprocessor") && job_preprocessed == Some(false);
|
|
|
|
let relative_imports = RELATIVE_IMPORT_REGEX.is_match(&inner_content);
|
|
|
|
let dirs = compute_python_module_dir(script_path);
|
|
let script_path_splitted = script_path.split("/").map(|x| {
|
|
if x.starts_with(|x: char| x.is_ascii_digit()) {
|
|
format!("_{}", x)
|
|
} else {
|
|
x.to_string()
|
|
}
|
|
});
|
|
let last = script_path_splitted
|
|
.clone()
|
|
.last()
|
|
.unwrap()
|
|
.replace("-", "_")
|
|
.replace(" ", "_")
|
|
.to_lowercase();
|
|
// `last` is lowercased above, so this catches a keyword id in any case.
|
|
let last = if is_python_keyword(&last) {
|
|
format!("_{last}")
|
|
} else {
|
|
last
|
|
};
|
|
let module_dir = format!("{}/{}", job_dir, dirs);
|
|
// Defense-in-depth: `dirs`/`last` derive from the (request-supplied for
|
|
// previews) script path. compute_python_module_dir already neutralizes `..`,
|
|
// but assert containment here too so the write can never escape job_dir.
|
|
is_allowed_file_location(job_dir, &format!("{dirs}/{last}.py"))?;
|
|
tokio::fs::create_dir_all(format!("{module_dir}/")).await?;
|
|
|
|
let _ = write_file(&module_dir, &format!("{last}.py"), inner_content)?;
|
|
if relative_imports {
|
|
let _ = write_file(
|
|
job_dir,
|
|
"loader.py",
|
|
&render_relative_python_loader(temp_script_refs),
|
|
)?;
|
|
}
|
|
|
|
let sig = windmill_parser_py::parse_python_signature(
|
|
inner_content,
|
|
main_override.map(ToString::to_string),
|
|
false,
|
|
)?;
|
|
|
|
let pre_sig = if apply_preprocessor {
|
|
Some(windmill_parser_py::parse_python_signature(
|
|
inner_content,
|
|
Some("preprocessor".to_string()),
|
|
false,
|
|
)?)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
// transforms should be applied based on the signature of the first function called
|
|
let init_sig = pre_sig.as_ref().unwrap_or(&sig);
|
|
|
|
let transforms = init_sig
|
|
.args
|
|
.iter()
|
|
.map(|x| match x.typ {
|
|
windmill_parser::Typ::Bytes => {
|
|
let name = &x.name;
|
|
format!(
|
|
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
|
kwargs[\"{name}\"] = base64.b64decode(kwargs[\"{name}\"])\n",
|
|
)
|
|
}
|
|
windmill_parser::Typ::Datetime => {
|
|
let name = &x.name;
|
|
format!(
|
|
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
|
kwargs[\"{name}\"] = datetime.fromisoformat(kwargs[\"{name}\"])\n",
|
|
)
|
|
}
|
|
windmill_parser::Typ::Date => {
|
|
let name = &x.name;
|
|
format!(
|
|
"if \"{name}\" in kwargs and kwargs[\"{name}\"] is not None:\n \
|
|
try:\n \
|
|
kwargs[\"{name}\"] = date.fromisoformat(kwargs[\"{name}\"])\n \
|
|
except ValueError:\n \
|
|
for _fmt in (\"%d-%m-%Y\", \"%m/%d/%Y\", \"%d/%m/%Y\", \"%Y/%m/%d\"):\n \
|
|
try:\n \
|
|
kwargs[\"{name}\"] = datetime.strptime(kwargs[\"{name}\"], _fmt).date()\n \
|
|
break\n \
|
|
except ValueError:\n \
|
|
continue\n",
|
|
)
|
|
}
|
|
_ => "".to_string(),
|
|
})
|
|
.collect::<Vec<String>>()
|
|
.join("");
|
|
|
|
let import_loader = if relative_imports {
|
|
"import loader"
|
|
} else {
|
|
""
|
|
};
|
|
let import_base64 = if init_sig
|
|
.args
|
|
.iter()
|
|
.any(|x| x.typ == windmill_parser::Typ::Bytes)
|
|
{
|
|
"import base64"
|
|
} else {
|
|
""
|
|
};
|
|
let has_datetime = init_sig
|
|
.args
|
|
.iter()
|
|
.any(|x| x.typ == windmill_parser::Typ::Datetime);
|
|
let has_date = init_sig
|
|
.args
|
|
.iter()
|
|
.any(|x| x.typ == windmill_parser::Typ::Date);
|
|
let import_datetime = match (has_datetime, has_date) {
|
|
(true, true) => "from datetime import datetime, date",
|
|
(true, false) => "from datetime import datetime",
|
|
(false, true) => "from datetime import datetime, date",
|
|
(false, false) => "",
|
|
};
|
|
let spread = if sig.star_kwargs {
|
|
"args = kwargs".to_string()
|
|
} else {
|
|
sig.args
|
|
.into_iter()
|
|
.map(|x| {
|
|
let name = &x.name;
|
|
if x.default.is_none() {
|
|
format!("args[\"{name}\"] = kwargs.get(\"{name}\")")
|
|
} else {
|
|
format!(
|
|
r#"args["{name}"] = kwargs.get("{name}")
|
|
if args["{name}"] is None:
|
|
del args["{name}"]"#
|
|
)
|
|
}
|
|
})
|
|
.join("\n ")
|
|
};
|
|
|
|
let pre_spread = pre_sig
|
|
.as_ref()
|
|
.map(|sig| python_preprocessor_spread(sig, " "));
|
|
// 4-space indent for the WAC wrapper, where the preprocessor block is
|
|
// injected inside the `try:` body so failures get serialized to result.json.
|
|
let wac_pre_spread = pre_sig
|
|
.as_ref()
|
|
.map(|sig| python_preprocessor_spread(sig, " "));
|
|
|
|
let module_dir_dot = dirs.replace("/", ".").replace("-", "_");
|
|
|
|
Ok((
|
|
import_loader,
|
|
import_base64,
|
|
import_datetime,
|
|
module_dir_dot,
|
|
dirs,
|
|
last,
|
|
transforms,
|
|
spread,
|
|
main_override.map(ToString::to_string),
|
|
pre_spread,
|
|
wac_pre_spread,
|
|
))
|
|
}
|
|
|
|
#[cfg(feature = "enterprise")]
|
|
async fn replace_pip_secret(
|
|
conn: &Connection,
|
|
w_id: &str,
|
|
req: &str,
|
|
worker_name: &str,
|
|
job_id: &Uuid,
|
|
) -> error::Result<String> {
|
|
if let Some(db) = conn.as_sql() {
|
|
if PIP_SECRET_VARIABLE.is_match(req) {
|
|
let mut joined = "".to_string();
|
|
for req in req.lines() {
|
|
let nreq = if PIP_SECRET_VARIABLE.is_match(req) {
|
|
let capture = PIP_SECRET_VARIABLE.captures(req);
|
|
let variable = capture.unwrap().get(1).unwrap().as_str();
|
|
if !variable.contains("/PIP_SECRET_") {
|
|
return Err(error::Error::internal_err(format!(
|
|
"invalid secret variable in pip requirements, (last part of path ma): {}",
|
|
req
|
|
)));
|
|
}
|
|
let secret = get_secret_value_as_admin(db, w_id, variable).await?;
|
|
tracing::info!(
|
|
worker = %worker_name,
|
|
job_id = %job_id,
|
|
workspace_id = %w_id,
|
|
"found secret variable in pip requirements: {}",
|
|
req
|
|
);
|
|
PIP_SECRET_VARIABLE
|
|
.replace(req, secret.as_str())
|
|
.to_string()
|
|
} else {
|
|
req.to_string()
|
|
};
|
|
joined.push_str(&nreq);
|
|
joined.push_str("\n");
|
|
}
|
|
|
|
Ok(joined)
|
|
} else {
|
|
Ok(req.to_string())
|
|
}
|
|
} else {
|
|
Ok(req.to_string())
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn handle_python_deps(
|
|
job_dir: &str,
|
|
requirements_o: Option<&String>,
|
|
inner_content: &str,
|
|
w_id: &str,
|
|
script_path: &str,
|
|
job_id: &Uuid,
|
|
conn: &Connection,
|
|
worker_name: &str,
|
|
worker_dir: &str,
|
|
mem_peak: &mut i32,
|
|
canceled_by: &mut Option<CanceledBy>,
|
|
occupancy_metrics: &mut Option<&mut OccupancyMetrics>,
|
|
precomputed_agent_info: Option<PrecomputedAgentInfo>,
|
|
annotations: PythonAnnotations,
|
|
temp_script_refs: &Option<HashMap<String, String>>,
|
|
) -> error::Result<(PyV, Vec<String>)> {
|
|
create_dependencies_dir(job_dir).await;
|
|
|
|
let mut additional_python_paths: Vec<String> = WORKER_CONFIG
|
|
.load()
|
|
.additional_python_paths
|
|
.clone()
|
|
.unwrap_or_else(|| vec![]);
|
|
|
|
let (pyv, resolved_lines) = match requirements_o {
|
|
// Deployed
|
|
Some(r) => {
|
|
let rl = split_python_requirements(r);
|
|
(PyV::parse_from_requirements(&rl), rl)
|
|
}
|
|
// Preview
|
|
None => {
|
|
let (v, requirements_lines, error_hint) = match conn {
|
|
Connection::Sql(db) => {
|
|
let (mut version_specifiers, mut locked_v) = (vec![], None);
|
|
let (r, h) = Box::pin(windmill_parser_py_imports::parse_python_imports(
|
|
inner_content,
|
|
w_id,
|
|
script_path,
|
|
db,
|
|
&mut version_specifiers,
|
|
&mut locked_v,
|
|
&None,
|
|
temp_script_refs,
|
|
))
|
|
.await?;
|
|
|
|
let v = if let Some(v) = locked_v {
|
|
v.into()
|
|
} else {
|
|
PyV::resolve(
|
|
version_specifiers,
|
|
job_id,
|
|
w_id,
|
|
annotations.py_select_latest,
|
|
Some(conn.clone()),
|
|
None,
|
|
None,
|
|
)
|
|
.await?
|
|
};
|
|
|
|
(v, r, h)
|
|
}
|
|
|
|
Connection::Http(_) => match precomputed_agent_info {
|
|
Some(PrecomputedAgentInfo::Python {
|
|
requirements,
|
|
py_version,
|
|
py_version_v2,
|
|
}) => {
|
|
let v = {
|
|
let v_v2 = py_version_v2
|
|
.clone()
|
|
.and_then(|s| pep440_rs::Version::from_str(&s).ok().map(PyV::from));
|
|
let v_v1 = py_version.and_then(PyVAlias::try_from_v1).map(PyV::from);
|
|
|
|
match v_v2.or(v_v1) {
|
|
Some(v) => v,
|
|
None => {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"
|
|
Failed to get precomputed python version from server. Fallback to Default ({})
|
|
Returned from server: py_version - {:?}, py_version_v2 - {:?}
|
|
",
|
|
*PyV::default(),
|
|
py_version,
|
|
py_version_v2
|
|
);
|
|
Default::default()
|
|
}
|
|
}
|
|
};
|
|
|
|
let r = split_python_requirements(requirements.unwrap_or_default());
|
|
let h = None;
|
|
|
|
(v, r, h)
|
|
}
|
|
_ => Default::default(),
|
|
},
|
|
};
|
|
|
|
(
|
|
v.clone(),
|
|
if !requirements_lines.is_empty() {
|
|
uv_pip_compile(
|
|
job_id,
|
|
&requirements_lines.join("\n"),
|
|
mem_peak,
|
|
canceled_by,
|
|
job_dir,
|
|
conn,
|
|
worker_name,
|
|
w_id,
|
|
occupancy_metrics,
|
|
// annotated_pyv.unwrap_or(instance_pyv),
|
|
v,
|
|
annotations.no_cache,
|
|
)
|
|
.await
|
|
.map_err(|e| {
|
|
Error::ExecutionErr(format!(
|
|
"pip compile failed: {}{}",
|
|
e.to_string(),
|
|
error_hint.unwrap_or_default()
|
|
))
|
|
})?
|
|
.lines()
|
|
.map(|s| s.to_owned())
|
|
.collect_vec()
|
|
} else {
|
|
vec![]
|
|
},
|
|
)
|
|
}
|
|
};
|
|
|
|
// Filter out packages matched by pip_local_dependencies. For preview runs this is also
|
|
// handled inside uv_pip_compile, but deployed scripts skip uv_pip_compile entirely and
|
|
// would otherwise pass every lockfile entry to handle_python_reqs — causing duplicate
|
|
// installs alongside additional_python_paths and triggering expensive postinstall copies.
|
|
let resolved_lines = {
|
|
let (kept, ignored) = filter_pip_local_dependencies(resolved_lines);
|
|
if !ignored.is_empty() {
|
|
append_logs(
|
|
job_id,
|
|
w_id,
|
|
format!("\nignoring local dependencies:\n{}\n", ignored.join("\n")),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|
|
kept
|
|
};
|
|
|
|
if !resolved_lines.is_empty() {
|
|
let mut venv_path = handle_python_reqs(
|
|
resolved_lines,
|
|
job_id,
|
|
w_id,
|
|
mem_peak,
|
|
canceled_by,
|
|
conn,
|
|
worker_name,
|
|
job_dir,
|
|
worker_dir,
|
|
occupancy_metrics,
|
|
pyv.clone(),
|
|
None,
|
|
)
|
|
.await?;
|
|
additional_python_paths.append(&mut venv_path);
|
|
}
|
|
|
|
Ok((pyv, additional_python_paths))
|
|
}
|
|
|
|
lazy_static::lazy_static! {
|
|
static ref PIP_SECRET_VARIABLE: Regex = Regex::new(r"\$\{PIP_SECRET:([^\s\}]+)\}").unwrap();
|
|
|
|
/// venv paths whose wheel RECORD this process has already verified against
|
|
/// disk. A cache entry is only damaged out-of-band (disk-pressure eviction,
|
|
/// an interrupted extraction on a shared cache volume, or a corrupt entry
|
|
/// that predates this worker), never spontaneously while we keep running, so
|
|
/// re-verifying it once per process is enough — every later reuse trusts the
|
|
/// in-memory marker and pays only the original single stat.
|
|
static ref VERIFIED_VENVS: tokio::sync::Mutex<HashSet<String>> =
|
|
tokio::sync::Mutex::new(HashSet::new());
|
|
|
|
// In-process locks serializing concurrent installs into the same shared
|
|
// `venv_p` cache dir; `uv --reinstall` removes a package's .dist-info/RECORD
|
|
// before rewriting it, so a sibling install/verify racing it corrupts the
|
|
// dir. Keyed by venv_p so distinct deps still install in parallel.
|
|
static ref PY_INSTALL_LOCKS: tokio::sync::Mutex<std::collections::HashMap<String, Arc<tokio::sync::Mutex<()>>>> =
|
|
tokio::sync::Mutex::new(std::collections::HashMap::new());
|
|
}
|
|
|
|
/// Returns the in-process install lock for a given target cache dir, creating it
|
|
/// on first use. Idle entries (only the map holds a reference) are pruned each
|
|
/// call so the map stays bounded by the number of in-flight installs.
|
|
async fn get_venv_install_lock(venv_p: &str) -> Arc<tokio::sync::Mutex<()>> {
|
|
let mut map = PY_INSTALL_LOCKS.lock().await;
|
|
map.retain(|_, v| Arc::strong_count(v) > 1);
|
|
map.entry(venv_p.to_string())
|
|
.or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
|
|
.clone()
|
|
}
|
|
|
|
/// Spawn process of uv install
|
|
/// Can be wrapped by nsjail depending on configuration
|
|
#[inline]
|
|
async fn spawn_uv_install(
|
|
w_id: &str,
|
|
req: &str,
|
|
venv_p: &str,
|
|
job_dir: &str,
|
|
(pip_extra_index_url, pip_index_url): (Option<String>, Option<String>),
|
|
// If none, it is system python
|
|
py_path: Option<String>,
|
|
worker_dir: &str,
|
|
) -> Result<Box<dyn TokioChildWrapper>, Error> {
|
|
let uv_index_strategy_guard = UV_INDEX_STRATEGY.read().await.clone();
|
|
let uv_index_strategy = uv_index_strategy_guard
|
|
.as_deref()
|
|
.unwrap_or("unsafe-best-match");
|
|
let uv_exclude_newer = (*UV_EXCLUDE_NEWER.read().await).map(|secs| format!("{}s", secs));
|
|
let uv_exclude_newer = uv_exclude_newer.as_deref();
|
|
let uv_python_install_mirror = UV_PYTHON_INSTALL_MIRROR.read().await.clone();
|
|
|
|
if is_sandboxing_enabled() {
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
"starting nsjail"
|
|
);
|
|
|
|
let mut vars = vec![("PATH", PATH_ENV.as_str())];
|
|
if let Some(url) = pip_extra_index_url.as_ref() {
|
|
vars.push(("EXTRA_INDEX_URL", url));
|
|
}
|
|
if let Some(url) = pip_index_url.as_ref() {
|
|
vars.push(("INDEX_URL", url));
|
|
}
|
|
if let Some(cert_path) = INDEX_CERT.as_ref() {
|
|
vars.push(("SSL_CERT_FILE", cert_path));
|
|
}
|
|
if let Some(host) = TRUSTED_HOST.as_ref() {
|
|
vars.push(("TRUSTED_HOST", host));
|
|
}
|
|
if *NATIVE_CERT {
|
|
vars.push(("UV_NATIVE_TLS", "true"));
|
|
}
|
|
if let Some(timeout) = UV_HTTP_TIMEOUT.as_ref() {
|
|
vars.push(("UV_HTTP_TIMEOUT", timeout.as_str()));
|
|
}
|
|
|
|
let _owner;
|
|
if let Some(py_path) = py_path.as_ref() {
|
|
_owner = format!(
|
|
"-p {} --python-preference only-managed",
|
|
py_path.as_str() //
|
|
);
|
|
vars.push(("PY_PATH", &_owner));
|
|
}
|
|
vars.push(("REQ", &req));
|
|
vars.push(("TARGET", venv_p));
|
|
vars.push(("UV_INDEX_STRATEGY", uv_index_strategy));
|
|
if let Some(v) = uv_exclude_newer {
|
|
vars.push(("UV_EXCLUDE_NEWER", v));
|
|
}
|
|
if let Some(mirror) = uv_python_install_mirror.as_ref() {
|
|
vars.push(("UV_PYTHON_INSTALL_MIRROR", mirror));
|
|
}
|
|
|
|
std::fs::create_dir_all(venv_p)?;
|
|
let nsjail_proto = format!("{req}.config.proto");
|
|
// Prepare NSJAIL
|
|
let _ = write_file(
|
|
job_dir,
|
|
&nsjail_proto,
|
|
NSJAIL_CONFIG_DOWNLOAD_PY_CONTENT
|
|
.replace("{WORKER_DIR}", worker_dir)
|
|
.replace("{PY_INSTALL_DIR}", &*PY_INSTALL_DIR)
|
|
.replace("{TARGET_DIR}", &venv_p)
|
|
.replace("{CLONE_NEWUSER}", &(!*DISABLE_NUSER).to_string())
|
|
.replace("{TRACING_PROXY_CA_CERT_PATH}", &*TRACING_PROXY_CA_CERT_PATH)
|
|
.replace("#{DEV}", DEV_CONF_NSJAIL)
|
|
.replace(
|
|
"{TMP_MOUNT_BLOCK}",
|
|
&resolve_nsjail_tmp_mount_block(job_dir).await,
|
|
)
|
|
.as_str(),
|
|
)?;
|
|
|
|
let mut nsjail_cmd = Command::new(NSJAIL_PATH.as_str());
|
|
nsjail_cmd
|
|
.current_dir(job_dir)
|
|
.env_clear()
|
|
.envs(vars)
|
|
.envs(PROXY_ENVS.clone())
|
|
.args(vec!["--config", &nsjail_proto])
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
start_child_process(nsjail_cmd, NSJAIL_PATH.as_str(), false).await
|
|
} else {
|
|
#[cfg(unix)]
|
|
let req = req.to_owned();
|
|
|
|
#[cfg(windows)]
|
|
let req = format!("{}", req);
|
|
|
|
let mut command_args = vec![
|
|
UV_PATH.as_str(),
|
|
"pip",
|
|
"install",
|
|
&req,
|
|
"--no-deps",
|
|
"--no-color",
|
|
// Prevent uv from discovering configuration files.
|
|
"--no-config",
|
|
"--link-mode=copy",
|
|
"--system",
|
|
"--target",
|
|
venv_p,
|
|
"--no-cache",
|
|
// If we invoke uv pip install, then we want to overwrite existing data
|
|
"--reinstall",
|
|
// Compile .py to .pyc at install time so imports are fast even
|
|
// through read-only nsjail mounts (no in-memory compilation per job).
|
|
"--compile-bytecode",
|
|
];
|
|
|
|
if let Some(py_path) = py_path.as_ref() {
|
|
command_args.extend([
|
|
"-p",
|
|
py_path.as_str(),
|
|
"--python-preference",
|
|
"only-managed", //
|
|
]);
|
|
} else {
|
|
command_args.extend([
|
|
"--python-preference",
|
|
"only-system", //
|
|
]);
|
|
}
|
|
|
|
if let Some(url) = pip_extra_index_url.as_ref() {
|
|
url.split(",").for_each(|url| {
|
|
command_args.extend(["--extra-index-url", url]);
|
|
});
|
|
}
|
|
if let Some(v) = uv_exclude_newer {
|
|
command_args.extend(["--exclude-newer", v]);
|
|
}
|
|
|
|
let mut envs = vec![("PATH", PATH_ENV.as_str())];
|
|
envs.push(("HOME", HOME_ENV.as_str()));
|
|
envs.push(("UV_INDEX_STRATEGY", uv_index_strategy));
|
|
if let Some(timeout) = UV_HTTP_TIMEOUT.as_ref() {
|
|
envs.push(("UV_HTTP_TIMEOUT", timeout.as_str()));
|
|
}
|
|
if let Some(mirror) = uv_python_install_mirror.as_ref() {
|
|
envs.push(("UV_PYTHON_INSTALL_MIRROR", mirror));
|
|
}
|
|
|
|
if let Some(url) = pip_index_url.as_ref() {
|
|
command_args.extend(["--index-url", url]);
|
|
}
|
|
if let Some(host) = TRUSTED_HOST.as_ref() {
|
|
host.split_whitespace().for_each(|h| {
|
|
command_args.extend(["--trusted-host", h]);
|
|
});
|
|
}
|
|
if *NATIVE_CERT {
|
|
command_args.extend(["--native-tls"]);
|
|
}
|
|
// TODO:
|
|
// Track https://github.com/astral-sh/uv/issues/6715
|
|
if let Some(cert_path) = INDEX_CERT.as_ref() {
|
|
// Once merged --cert can be used instead
|
|
//
|
|
// command_args.extend(["--cert", cert_path]);
|
|
envs.push(("SSL_CERT_FILE", cert_path));
|
|
}
|
|
|
|
tracing::debug!("uv pip install command: {:?}", command_args);
|
|
|
|
#[cfg(unix)]
|
|
{
|
|
let mut cmd = Command::new(command_args[0]);
|
|
cmd.env_clear()
|
|
.envs(PROXY_ENVS.clone())
|
|
.envs(envs)
|
|
.args(&command_args[1..])
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
start_child_process(cmd, UV_PATH.as_str(), false).await
|
|
}
|
|
|
|
#[cfg(windows)]
|
|
{
|
|
let mut cmd: Command = Command::new("uv");
|
|
cmd.env_clear()
|
|
.envs(envs)
|
|
.envs(PROXY_ENVS.clone())
|
|
.env("SystemRoot", SYSTEM_ROOT.as_str())
|
|
.env("USERPROFILE", crate::USERPROFILE_ENV.as_str())
|
|
.env("HOME", HOME_ENV.as_str())
|
|
.env(
|
|
"TMP",
|
|
std::env::var("TMP").unwrap_or_else(|_| String::from("/tmp")),
|
|
)
|
|
.env(
|
|
"LOCALAPPDATA",
|
|
std::env::var("LOCALAPPDATA")
|
|
.unwrap_or_else(|_| format!("{}\\AppData\\Local", HOME_ENV.as_str())),
|
|
)
|
|
.env(
|
|
"APPDATA",
|
|
std::env::var("APPDATA").unwrap_or_else(|_| {
|
|
format!("{}\\AppData\\Roaming", crate::USERPROFILE_ENV.as_str())
|
|
}),
|
|
)
|
|
.env(
|
|
"ComSpec",
|
|
std::env::var("ComSpec")
|
|
.unwrap_or_else(|_| String::from("C:\\Windows\\System32\\cmd.exe")),
|
|
)
|
|
.env(
|
|
"PATHEXT",
|
|
std::env::var("PATHEXT").unwrap_or_else(|_| {
|
|
String::from(".COM;.EXE;.BAT;.CMD;.VBS;.VBE;.JS;.JSE;.WSF;.WSH;.MSC;.CPL")
|
|
}),
|
|
)
|
|
.env(
|
|
"ProgramData",
|
|
std::env::var("ProgramData")
|
|
.unwrap_or_else(|_| String::from("C:\\ProgramData")),
|
|
)
|
|
.env(
|
|
"ProgramFiles",
|
|
std::env::var("ProgramFiles")
|
|
.unwrap_or_else(|_| String::from("C:\\Program Files")),
|
|
)
|
|
.args(&command_args[1..])
|
|
.stdout(Stdio::piped())
|
|
.stderr(Stdio::piped());
|
|
start_child_process(cmd, "uv", false).await
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Verify that every file listed in the wheel's RECORD exists on disk under
|
|
/// `venv_p`. Used as a structural integrity check after both a successful
|
|
/// `pull_from_tar` (object-store cache hit) and a successful local
|
|
/// `uv pip install`, so a truncated tar or a dropped wheel entry can never
|
|
/// become an authoritative cache entry. Returns Err with a short description
|
|
/// on the first integrity issue (no .dist-info, no RECORD, or any listed
|
|
/// path missing on disk).
|
|
async fn verify_wheel_record(venv_p: &str) -> Result<(), String> {
|
|
let mut entries = tokio::fs::read_dir(venv_p)
|
|
.await
|
|
.map_err(|e| format!("read_dir({venv_p}): {e}"))?;
|
|
|
|
let mut dist_info: Option<String> = None;
|
|
loop {
|
|
match entries.next_entry().await {
|
|
Ok(Some(entry)) => {
|
|
let name = entry.file_name();
|
|
let name_s = name.to_string_lossy();
|
|
if name_s.ends_with(".dist-info") {
|
|
if let Ok(ft) = entry.file_type().await {
|
|
if ft.is_dir() {
|
|
dist_info = Some(name_s.into_owned());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Ok(None) => break,
|
|
Err(e) => return Err(format!("read_dir entry in {venv_p}: {e}")),
|
|
}
|
|
}
|
|
|
|
let dist_info = match dist_info {
|
|
Some(d) => d,
|
|
None => return Err(format!("no .dist-info directory in {venv_p}")),
|
|
};
|
|
|
|
let record_path = format!("{venv_p}/{dist_info}/RECORD");
|
|
let record_content = tokio::fs::read_to_string(&record_path)
|
|
.await
|
|
.map_err(|e| format!("read RECORD at {record_path}: {e}"))?;
|
|
|
|
let mut missing: Vec<String> = Vec::new();
|
|
for line in record_content.lines() {
|
|
let trimmed = line.trim();
|
|
if trimmed.is_empty() {
|
|
continue;
|
|
}
|
|
let rel_path = match trimmed.split(',').next() {
|
|
Some(p) if !p.is_empty() => p,
|
|
_ => continue,
|
|
};
|
|
// Defensive: skip absolute paths or escaping entries — we only
|
|
// validate package-relative files.
|
|
if rel_path.starts_with('/') || rel_path.contains("..") {
|
|
continue;
|
|
}
|
|
let full = format!("{venv_p}/{rel_path}");
|
|
if tokio::fs::metadata(&full).await.is_err() {
|
|
missing.push(rel_path.to_string());
|
|
// Bound error size in pathological cases (e.g. wholly empty dir).
|
|
if missing.len() >= 10 {
|
|
missing.push("...".to_string());
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
if missing.is_empty() {
|
|
Ok(())
|
|
} else {
|
|
Err(format!(
|
|
"wheel RECORD lists files missing on disk: {}",
|
|
missing.join(", ")
|
|
))
|
|
}
|
|
}
|
|
|
|
/// uv pip install, include cached or pull from S3
|
|
pub async fn handle_python_reqs(
|
|
requirements: Vec<String>,
|
|
job_id: &Uuid,
|
|
w_id: &str,
|
|
mem_peak: &mut i32,
|
|
_canceled_by: &mut Option<CanceledBy>,
|
|
conn: &Connection,
|
|
_worker_name: &str,
|
|
job_dir: &str,
|
|
worker_dir: &str,
|
|
_occupancy_metrics: &mut Option<&mut OccupancyMetrics>,
|
|
py_version: PyV,
|
|
reduced_concurrent_downloads: Option<usize>,
|
|
) -> error::Result<Vec<String>> {
|
|
let worker_dir = worker_dir.to_string();
|
|
|
|
let counter_arc = Arc::new(tokio::sync::Mutex::new(0));
|
|
// Append logs with line like this:
|
|
// [9/21] + requests==2.32.3 << (S3) | in 57ms
|
|
#[allow(unused_assignments)]
|
|
async fn print_success(
|
|
mut s3_pull: bool,
|
|
mut s3_push: bool,
|
|
job_id: &Uuid,
|
|
w_id: &str,
|
|
req: &str,
|
|
req_tl: usize,
|
|
counter_arc: Arc<tokio::sync::Mutex<usize>>,
|
|
total_to_install: usize,
|
|
instant: std::time::Instant,
|
|
conn: &Connection,
|
|
) {
|
|
#[cfg(not(all(feature = "enterprise", feature = "parquet")))]
|
|
{
|
|
(s3_pull, s3_push) = (false, false);
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if get_cache_object_store().await.is_none() {
|
|
(s3_pull, s3_push) = (false, false);
|
|
}
|
|
|
|
let mut counter = counter_arc.lock().await;
|
|
*counter += 1;
|
|
|
|
append_logs(
|
|
job_id,
|
|
w_id,
|
|
format!(
|
|
"\n{}+ {}{}{}| in {}ms",
|
|
pad_string(&format!("[{}/{total_to_install}]", counter), 9),
|
|
// Because we want to align to max len [999/999] we take ^
|
|
// 123456789
|
|
pad_string(&req, req_tl + 1),
|
|
// Margin to the right ^
|
|
if s3_pull { "<< (S3) " } else { "" },
|
|
if s3_push { " > (S3) " } else { "" },
|
|
instant.elapsed().as_millis(),
|
|
),
|
|
conn,
|
|
)
|
|
.await;
|
|
// Drop lock, so next print success can fire
|
|
}
|
|
|
|
// Parallelism level (N)
|
|
let parallel_limit = reduced_concurrent_downloads
|
|
.unwrap_or(*PY_CONCURRENT_DOWNLOADS)
|
|
.clamp(1, 30);
|
|
// Semaphore will panic if value less then 1
|
|
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
// is_ok = out,
|
|
"Parallel limit: {}, job: {}",
|
|
parallel_limit,
|
|
job_id
|
|
);
|
|
|
|
let pip_indexes = (
|
|
read_ee_registry_with_workspace_override(
|
|
PIP_EXTRA_INDEX_URL.read().await.clone(),
|
|
"pip_extra_index_url",
|
|
"pip extra index url",
|
|
job_id,
|
|
w_id,
|
|
conn,
|
|
)
|
|
.await
|
|
.map(handle_ephemeral_token),
|
|
read_ee_registry_with_workspace_override(
|
|
PIP_INDEX_URL.read().await.clone(),
|
|
"pip_index_url",
|
|
"pip index url",
|
|
job_id,
|
|
w_id,
|
|
conn,
|
|
)
|
|
.await
|
|
.map(handle_ephemeral_token),
|
|
);
|
|
|
|
// Cached paths
|
|
let mut req_with_penv: Vec<(String, String)> = vec![];
|
|
// Requirements to pull (not cached)
|
|
let mut req_paths: Vec<String> = vec![];
|
|
// Find out if there is already cached dependencies
|
|
// If so, skip them
|
|
let mut in_cache = vec![];
|
|
if requirements
|
|
.iter()
|
|
.any(|r| lockfile_line_has_continuation(r))
|
|
{
|
|
tracing::warn!(workspace_id = %w_id, job_id = %job_id, "lockfile continues entries across lines; the continued lines are dropped");
|
|
append_logs(
|
|
job_id,
|
|
w_id,
|
|
"\n[!] lockfile continues entries across lines and the continued lines are dropped: `--hash=` pins, extras and markers written that way do not apply\n".to_string(),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|
|
for req in &requirements {
|
|
let Some(req) = requirement_from_lockfile_line(req) else {
|
|
continue;
|
|
};
|
|
let py_prefix = &py_version.to_cache_dir(false);
|
|
|
|
let venv_p = format!(
|
|
"{py_prefix}/{}",
|
|
req.replace(' ', "").replace('/', "").replace(':', "")
|
|
);
|
|
if metadata(venv_p.clone() + "/.valid.windmill").await.is_ok() {
|
|
// The .valid.windmill marker is written once at creation time, after
|
|
// verify_wheel_record passes on the install/pull paths. It is an empty
|
|
// file with no binding to the directory contents, so a file dropped
|
|
// out-of-band afterwards (disk-pressure eviction, interrupted tar
|
|
// extraction on a shared cache volume, or a corrupt entry that
|
|
// predates this worker) leaves the marker intact while the wheel is
|
|
// incomplete. Re-verify the RECORD once per process so such an entry
|
|
// is repaired rather than trusted; VERIFIED_VENVS makes every later
|
|
// reuse skip the scan and pay only the single stat above.
|
|
let already_verified = VERIFIED_VENVS.lock().await.contains(&venv_p);
|
|
let verify_res = if already_verified {
|
|
Ok(())
|
|
} else {
|
|
verify_wheel_record(&venv_p).await
|
|
};
|
|
match verify_res {
|
|
Ok(()) => {
|
|
if !already_verified {
|
|
VERIFIED_VENVS.lock().await.insert(venv_p.clone());
|
|
}
|
|
req_paths.push(venv_p);
|
|
in_cache.push(req.to_string());
|
|
}
|
|
Err(verify_err) => {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
"Local cache for {venv_p} failed wheel RECORD verification, will reinstall: {verify_err}"
|
|
);
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!(
|
|
"\n[!] cached wheel for {req} failed integrity check, reinstalling: {verify_err}\n"
|
|
),
|
|
conn,
|
|
)
|
|
.await;
|
|
if let Err(rm_err) = tokio::fs::remove_dir_all(&venv_p).await {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"could not remove broken cache dir {venv_p}: {rm_err}"
|
|
);
|
|
}
|
|
req_with_penv.push((req.to_string(), venv_p));
|
|
}
|
|
}
|
|
} else {
|
|
// There is no valid or no wheel at all. Regardless of if there is content or not, we will overwrite it with --reinstall flag
|
|
req_with_penv.push((req.to_string(), venv_p));
|
|
}
|
|
}
|
|
if in_cache.len() > 0 {
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!("\nenv deps from local cache: {}\n", in_cache.join(", ")),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
let (kill_tx, ..) = tokio::sync::broadcast::channel::<()>(1);
|
|
let kill_rxs: Vec<tokio::sync::broadcast::Receiver<()>> = (0..req_with_penv.len())
|
|
.map(|_| kill_tx.subscribe())
|
|
.collect();
|
|
|
|
// _______ Read comments at the end of the function to get more context
|
|
let (done_tx, mut done_rx) = tokio::sync::mpsc::channel::<()>(1);
|
|
|
|
let job_id_2 = job_id.clone();
|
|
let conn_2 = conn.clone();
|
|
let w_id_2 = w_id.to_string();
|
|
|
|
// Wheels to install
|
|
let total_to_install = req_with_penv.len();
|
|
let pids = Arc::new(tokio::sync::Mutex::new(vec![None; total_to_install]));
|
|
let mem_peak_thread_safe = Arc::new(tokio::sync::Mutex::new(0));
|
|
{
|
|
let pids = pids.clone();
|
|
let mem_peak_thread_safe = mem_peak_thread_safe.clone();
|
|
tokio::spawn(async move {
|
|
loop {
|
|
tokio::select! {
|
|
_ = tokio::time::sleep(tokio::time::Duration::from_secs(1)) => {
|
|
let mut local_mem_peak = 0;
|
|
for pid_o in pids.lock().await.iter() {
|
|
if pid_o.is_some(){
|
|
let mem = crate::handle_child::get_mem_peak(*pid_o, is_sandboxing_enabled()).await;
|
|
if mem < 0 {
|
|
tracing::warn!(
|
|
workspace_id = %w_id_2,
|
|
"Cannot get memory peak for pid: {:?}, job_id: {:?}, exit code: {mem}",
|
|
pid_o,
|
|
job_id_2
|
|
);
|
|
} else {
|
|
local_mem_peak += mem;
|
|
}
|
|
}
|
|
}
|
|
|
|
let mem_peak_actual = {
|
|
let mut mem_peak_lock = mem_peak_thread_safe.lock().await;
|
|
|
|
if local_mem_peak > *mem_peak_lock{
|
|
*mem_peak_lock = local_mem_peak;
|
|
} else {
|
|
tracing::debug!(
|
|
workspace_id = %w_id_2,
|
|
"Local mem_peak {:?}mb is smaller then global one {:?}mb, ignoring. job_id: {:?}",
|
|
local_mem_peak / 1000,
|
|
*mem_peak_lock / 1000,
|
|
job_id_2
|
|
);
|
|
|
|
}
|
|
// Get the copy of value and drop lock itself, to release it as fast as possible
|
|
*mem_peak_lock
|
|
};
|
|
|
|
|
|
// Notify server that we are still alive
|
|
// Detect if job has been canceled
|
|
let canceled = match conn_2 {
|
|
Connection::Sql(ref db) => {
|
|
sqlx::query_scalar!(
|
|
"UPDATE v2_job_runtime r SET
|
|
memory_peak = $1,
|
|
ping = now()
|
|
FROM v2_job_queue q
|
|
WHERE r.id = $2 AND q.id = r.id
|
|
RETURNING canceled_by IS NOT NULL AS \"canceled!\"",
|
|
mem_peak_actual,
|
|
job_id_2
|
|
)
|
|
.fetch_optional(db)
|
|
.await
|
|
.unwrap_or_else(|e| {
|
|
tracing::error!(%e, "error updating job {job_id_2}: {e:#}");
|
|
Some(false)
|
|
})
|
|
.unwrap_or_else(|| {
|
|
// if the job is not in queue, it can only be in the completed_job so it is already complete
|
|
false
|
|
})
|
|
}
|
|
Connection::Http(_) => {
|
|
if let Err(e) = ping_job_status(&conn_2, &job_id_2, Some(mem_peak_actual), None).await {
|
|
tracing::error!(%e, "error pinging job {job_id_2}: {e:#}");
|
|
}
|
|
false
|
|
}
|
|
};
|
|
|
|
if canceled {
|
|
tracing::info!(
|
|
// If there is listener on other side,
|
|
workspace_id = %w_id_2,
|
|
"cancelling installations",
|
|
);
|
|
|
|
if let Err(ref e) = kill_tx.send(()){
|
|
tracing::error!(
|
|
// If there is listener on other side,
|
|
workspace_id = %w_id_2,
|
|
"failed to send done: Probably receiving end closed too early or have not opened yet\n{}",
|
|
// If there is no listener, it will be dropped safely
|
|
e
|
|
);
|
|
}
|
|
}
|
|
}
|
|
// Once done_tx is dropped, this will be fired
|
|
_ = done_rx.recv() => break
|
|
}
|
|
}
|
|
});
|
|
}
|
|
|
|
// tl = total_length
|
|
// "small".len == 5
|
|
// "middle".len == 6
|
|
// "largest".len == 7
|
|
// ==> req_tl = 7
|
|
let mut req_tl = 0;
|
|
if total_to_install > 0 {
|
|
let mut logs = String::new();
|
|
logs.push_str("\n\n--- UV PIP INSTALL ---\n");
|
|
logs.push_str("\nTo be installed: \n\n");
|
|
for (req, _) in &req_with_penv {
|
|
if req.len() > req_tl {
|
|
req_tl = req.len();
|
|
}
|
|
logs.push_str(&format!("{} \n", &req));
|
|
}
|
|
|
|
// Do we use Nsjail?
|
|
if is_sandboxing_enabled() {
|
|
logs.push_str(&format!(
|
|
"\nStarting isolated installation... ({} tasks in parallel) \n",
|
|
parallel_limit
|
|
));
|
|
} else {
|
|
logs.push_str(&format!(
|
|
"\nStarting installation... ({} tasks in parallel) \n",
|
|
parallel_limit
|
|
));
|
|
}
|
|
append_logs(&job_id, w_id, logs, conn).await;
|
|
}
|
|
|
|
let semaphore = Arc::new(Semaphore::new(parallel_limit));
|
|
let mut handles = Vec::with_capacity(total_to_install);
|
|
// let mem_peak_thread_safe = Arc::new(tokio::sync::Mutex::new(0));
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let is_not_pro = !matches!(get_license_plan().await, LicensePlan::Pro);
|
|
|
|
let total_time = std::time::Instant::now();
|
|
let py_path = py_version
|
|
.try_get_python(
|
|
job_id,
|
|
mem_peak,
|
|
conn,
|
|
_worker_name,
|
|
w_id,
|
|
_occupancy_metrics,
|
|
)
|
|
.await?;
|
|
|
|
let has_work = req_with_penv.len() > 0;
|
|
for ((i, (req, venv_p)), mut kill_rx) in
|
|
req_with_penv.iter().enumerate().zip(kill_rxs.into_iter())
|
|
{
|
|
let permit = semaphore.clone().acquire_owned().await; // Acquire a permit
|
|
|
|
if let Err(_) = permit {
|
|
tracing::error!(
|
|
workspace_id = %w_id,
|
|
"Cannot acquire permit on semaphore, that can only mean that semaphore has been closed."
|
|
);
|
|
break;
|
|
}
|
|
|
|
let permit = permit.unwrap();
|
|
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
"started setup python dependencies"
|
|
);
|
|
|
|
let conn = conn.clone();
|
|
let job_id = job_id.clone();
|
|
let job_dir = job_dir.to_owned();
|
|
let w_id = w_id.to_owned();
|
|
let req = req.clone();
|
|
let venv_p = venv_p.clone();
|
|
let counter_arc = counter_arc.clone();
|
|
let pip_indexes = pip_indexes.clone();
|
|
let py_path = py_path.clone();
|
|
let pids = pids.clone();
|
|
let worker_dir = worker_dir.clone();
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let py_version = py_version.clone();
|
|
|
|
handles.push(task::spawn(async move {
|
|
// permit will be dropped anyway if this thread exits at any point
|
|
// so we dont have to drop it manually
|
|
// but we need to move permit into scope to take ownership
|
|
let _permit = permit;
|
|
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
// is_ok = out,
|
|
"started thread to install wheel {}",
|
|
venv_p
|
|
);
|
|
|
|
let start = std::time::Instant::now();
|
|
|
|
// Lock the shared target dir (see PY_INSTALL_LOCKS). In-process lock
|
|
// first; only one task per dir then contends the cross-process file
|
|
// lock below. Both guards drop on every return path.
|
|
let venv_lock = get_venv_install_lock(&venv_p).await;
|
|
let _venv_guard = tokio::select! {
|
|
_ = kill_rx.recv() => {
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::from(anyhow::anyhow!(
|
|
"install of {venv_p} canceled while waiting for venv lock"
|
|
)));
|
|
}
|
|
guard = venv_lock.lock_owned() => guard,
|
|
};
|
|
|
|
// Cross-process advisory lock. Best-effort: if the filesystem doesn't
|
|
// support locking we log and proceed — verify_wheel_record + job retry
|
|
// still guard correctness, just without the dedup. Cross-platform
|
|
// (flock on unix, LockFileEx on windows) so agents sharing a wheel-cache
|
|
// dir on a Windows host serialize just as they do on unix.
|
|
let _venv_file_lock: Option<std::fs::File> = {
|
|
use fs4::fs_std::FileExt;
|
|
let lock_path = format!("{venv_p}.lock");
|
|
if let Some(parent) = std::path::Path::new(&lock_path).parent() {
|
|
let _ = std::fs::create_dir_all(parent);
|
|
}
|
|
match std::fs::OpenOptions::new().create(true).write(true).open(&lock_path) {
|
|
Ok(f) => {
|
|
// Bounded wait: a holder that crashes releases the lock (the
|
|
// OS drops it on handle close), but a live-but-stuck holder
|
|
// (e.g. uv wedged on a hung mount) would otherwise block us
|
|
// forever. After the cap, proceed degraded rather than hang —
|
|
// verify_wheel_record + retry still guard correctness.
|
|
const MAX_WAIT: std::time::Duration = std::time::Duration::from_secs(300);
|
|
let waited_since = std::time::Instant::now();
|
|
loop {
|
|
match f.try_lock_exclusive() {
|
|
Ok(true) => break Some(f),
|
|
// Another holder has the lock.
|
|
Ok(false) => {
|
|
if waited_since.elapsed() >= MAX_WAIT {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"venv install lock {lock_path} still held after {}s, proceeding without cross-process install lock",
|
|
MAX_WAIT.as_secs()
|
|
);
|
|
break Some(f);
|
|
}
|
|
tokio::select! {
|
|
_ = kill_rx.recv() => {
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::from(anyhow::anyhow!(
|
|
"install of {venv_p} canceled while waiting for venv file lock"
|
|
)));
|
|
}
|
|
_ = tokio::time::sleep(std::time::Duration::from_millis(200)) => {}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"could not lock {lock_path}, proceeding without cross-process install lock: {e}"
|
|
);
|
|
break Some(f);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"could not open install lock file {lock_path}, proceeding without cross-process install lock: {e}"
|
|
);
|
|
None
|
|
}
|
|
}
|
|
};
|
|
|
|
// Double-checked: another job (this process or another sharing the
|
|
// mount) may have installed this exact dep while we waited on the
|
|
// locks. Reuse it instead of reinstalling.
|
|
if metadata(format!("{venv_p}/.valid.windmill")).await.is_ok() {
|
|
print_success(
|
|
false, false, &job_id, &w_id, &req, req_tl, counter_arc,
|
|
total_to_install, start, &conn,
|
|
)
|
|
.await;
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Ok(());
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if is_not_pro {
|
|
if let Some(os) = windmill_object_store::get_cache_object_store().await {
|
|
tokio::select! {
|
|
// Cancel was called on the job
|
|
_ = kill_rx.recv() => return Err(Error::from(anyhow::anyhow!("S3 pull was canceled"))),
|
|
pull = pull_from_tar(os, venv_p.clone(), py_version.to_cache_dir_top_level(false), None, false) => {
|
|
if let Err(e) = pull {
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
"No tarball was found for {venv_p} on S3 or different problem occurred {job_id}:\n{e}",
|
|
);
|
|
} else if let Err(verify_err) = verify_wheel_record(&venv_p).await {
|
|
// The object-store tar extracted cleanly but the resulting
|
|
// directory is missing files referenced by the wheel RECORD.
|
|
// Wipe the broken cache entry and fall through to a fresh
|
|
// local install rather than treating it as authoritative.
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
"Object-store cache for {venv_p} failed wheel RECORD verification, will reinstall locally: {verify_err}"
|
|
);
|
|
if let Err(rm_err) = tokio::fs::remove_dir_all(&venv_p).await {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"could not remove broken cache dir {venv_p}: {rm_err}"
|
|
);
|
|
}
|
|
append_logs(
|
|
&job_id,
|
|
&w_id,
|
|
format!(
|
|
"\n[!] cached wheel for {req} from object store failed integrity check, reinstalling: {verify_err}\n"
|
|
),
|
|
&conn,
|
|
).await;
|
|
} else {
|
|
print_success(
|
|
true,
|
|
false,
|
|
&job_id,
|
|
&w_id,
|
|
&req,
|
|
req_tl,
|
|
counter_arc,
|
|
total_to_install,
|
|
start,
|
|
&conn
|
|
).await;
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
|
|
// Create a file to indicate that installation was successfull
|
|
let valid_path = venv_p.clone() + "/.valid.windmill";
|
|
// This is atomic operation, meaning, that it either completes and wheel is valid,
|
|
// or it does not and wheel is invalid and will be reinstalled next run
|
|
if let Err(e) = File::create(&valid_path).await{
|
|
tracing::error!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
"Failed to create {}!\n{e}\n
|
|
This file needed for python jobs to function", valid_path)
|
|
};
|
|
return Ok(());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut uv_install_proccess = match spawn_uv_install(
|
|
&w_id,
|
|
&req,
|
|
&venv_p,
|
|
&job_dir,
|
|
pip_indexes,
|
|
py_path,
|
|
&worker_dir
|
|
).await {
|
|
Ok(r) => r,
|
|
Err(e) => {
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!(
|
|
"\nError while spawning proccess:\n{e}",
|
|
),
|
|
&conn,
|
|
)
|
|
.await;
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::from(e));
|
|
}
|
|
};
|
|
|
|
let (mut stderr_buf, mut stdout_buf) = Default::default();
|
|
let (mut stderr_pipe, mut stdout_pipe) = (
|
|
uv_install_proccess
|
|
.stderr()
|
|
.take()
|
|
.ok_or(anyhow!("Cannot take stderr from uv_install_proccess"))?,
|
|
uv_install_proccess
|
|
.stdout()
|
|
.take()
|
|
.ok_or(anyhow!("Cannot take stdout from uv_install_proccess"))?
|
|
);
|
|
let (stderr_future, stdout_future) = (
|
|
stderr_pipe.read_to_string(&mut stderr_buf),
|
|
stdout_pipe.read_to_string(&mut stdout_buf)
|
|
);
|
|
|
|
if let Some(pid) = pids.lock().await.get_mut(i) {
|
|
*pid = uv_install_proccess.id();
|
|
#[cfg(unix)]
|
|
if let Err(e) = uv_install_proccess
|
|
.id()
|
|
.ok_or(Error::InternalErr(format!(
|
|
"failed to get PID for python installation process: {}",
|
|
&req
|
|
)))
|
|
.and_then(|pid| write_file(
|
|
&format!("/proc/{pid}"),
|
|
"oom_score_adj",
|
|
&windmill_common::worker::JOB_OOM_SCORE_ADJ.to_string(),
|
|
))
|
|
{
|
|
tracing::error!(
|
|
req = %req,
|
|
"Failed to create oom_score_adj for python dependency installation process: {e}"
|
|
);
|
|
}
|
|
} else {
|
|
tracing::error!(
|
|
workspace_id = %w_id,
|
|
"Index out of range for uv pids",
|
|
);
|
|
}
|
|
|
|
tokio::select! {
|
|
// Canceled
|
|
_ = kill_rx.recv() => {
|
|
Box::into_pin(uv_install_proccess.kill()).await?;
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::from(anyhow::anyhow!("uv pip install was canceled")));
|
|
},
|
|
(_, _, exitstatus) = async {
|
|
// See tokio::process::Child::wait_with_output() for more context
|
|
// Sometimes uv_install_proccess.wait() is not exiting if stderr is not awaited first
|
|
(stderr_future.await, stdout_future.await, Box::into_pin(uv_install_proccess.wait()).await)
|
|
} => match exitstatus {
|
|
Ok(status) => if !status.success() {
|
|
#[cfg(unix)]
|
|
let code = status.signal();
|
|
#[cfg(not(unix))]
|
|
let code = status.code();
|
|
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"uv install {} did not succeed, exit status: {:?}",
|
|
&req,
|
|
code
|
|
);
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!(
|
|
"\nError while installing {}: \nStderr:\n{stderr_buf}\nStdout:\n{stdout_buf}\nExit status: {:?}",
|
|
&req,
|
|
code
|
|
),
|
|
&conn,
|
|
)
|
|
.await;
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::ExitStatus(stderr_buf, code.unwrap_or(1)));
|
|
},
|
|
Err(e) => {
|
|
tracing::error!(
|
|
workspace_id = %w_id,
|
|
"Cannot wait for uv_install_proccess, ExitStatus is Err: {e:?}",
|
|
);
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::from(e));
|
|
}
|
|
}
|
|
};
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let s3_push = is_not_pro;
|
|
|
|
#[cfg(not(all(feature = "enterprise", feature = "parquet")))]
|
|
let s3_push = false;
|
|
|
|
if is_sandboxing_enabled() {
|
|
let _ = std::fs::remove_file(format!("{job_dir}/{req}.config.proto"));
|
|
}
|
|
|
|
// Verify the install before declaring success: if uv exited 0 but
|
|
// the on-disk directory is missing files the wheel RECORD says
|
|
// should exist, do NOT write .valid.windmill, do NOT queue the
|
|
// piptar upload, and fail the job. This prevents a broken tar
|
|
// from ever being pushed to the object store and propagated to
|
|
// every other replica.
|
|
if let Err(verify_err) = verify_wheel_record(&venv_p).await {
|
|
tracing::error!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
"uv pip install of {req} into {venv_p} failed wheel RECORD verification: {verify_err}"
|
|
);
|
|
append_logs(
|
|
&job_id,
|
|
&w_id,
|
|
format!(
|
|
"\nWheel RECORD verification failed after install of {req}: {verify_err}. \
|
|
Aborting to avoid publishing a corrupt cache entry."
|
|
),
|
|
&conn,
|
|
).await;
|
|
if let Err(rm_err) = tokio::fs::remove_dir_all(&venv_p).await {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"could not remove broken install dir {venv_p}: {rm_err}"
|
|
);
|
|
}
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
return Err(Error::from(anyhow!(
|
|
"wheel RECORD verification failed after install of {req}"
|
|
)));
|
|
}
|
|
|
|
print_success(
|
|
false,
|
|
s3_push,
|
|
&job_id,
|
|
&w_id,
|
|
&req,
|
|
req_tl,
|
|
counter_arc,
|
|
total_to_install,
|
|
start,
|
|
&conn, //
|
|
)
|
|
.await;
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if s3_push {
|
|
// Send to upload channel for sequential processing
|
|
let upload_task = PiptarUploadTask {
|
|
venv_path: venv_p.clone(),
|
|
cache_dir: py_version.to_cache_dir_top_level(false),
|
|
};
|
|
|
|
if let Err(e) = PIPTAR_UPLOAD_CHANNEL.send(upload_task) {
|
|
tracing::warn!("Failed to queue piptar upload for {venv_p}: {e}");
|
|
} else {
|
|
tracing::info!("Queued piptar upload for {venv_p}");
|
|
}
|
|
}
|
|
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
// is_ok = out,
|
|
"finished setting up python dependency {}",
|
|
venv_p
|
|
);
|
|
|
|
pids.lock().await.get_mut(i).and_then(|e| e.take());
|
|
// Create a file to indicate that installation was successfull
|
|
let valid_path = venv_p.clone() + "/.valid.windmill";
|
|
// This is atomic operation, meaning, that it either completes and wheel is valid,
|
|
// or it does not and wheel is invalid and will be reinstalled next run
|
|
if let Err(e) = File::create(&valid_path).await{
|
|
tracing::error!(
|
|
workspace_id = %w_id,
|
|
job_id = %job_id,
|
|
"Failed to create {}!\n{e}\n
|
|
This file needed for python jobs to function", valid_path)
|
|
};
|
|
Ok(())
|
|
}));
|
|
}
|
|
|
|
let (mut failed, mut oom_killed) = (false, false);
|
|
for (handle, (_, venv_p)) in handles.into_iter().zip(req_with_penv.into_iter()) {
|
|
if let Err(e) = handle
|
|
.await
|
|
.unwrap_or(Err(Error::from(anyhow!("Problem by joining handle"))))
|
|
{
|
|
failed = true;
|
|
|
|
// OOM code is 9 or 137
|
|
if matches!(e, Error::ExitStatus(_, 9 | 137)) {
|
|
oom_killed = true;
|
|
}
|
|
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!("\nEnv installation failed: {:?}", e),
|
|
conn,
|
|
)
|
|
.await;
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"Env installation failed: {:?}",
|
|
e
|
|
);
|
|
} else {
|
|
req_paths.push(venv_p);
|
|
}
|
|
}
|
|
|
|
if has_work {
|
|
let total_time = total_time.elapsed().as_millis();
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!("\nenv set in {}ms", total_time),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
*mem_peak = *mem_peak_thread_safe.lock().await;
|
|
|
|
// Usually done_tx will drop after this return
|
|
// If there is listener on other side,
|
|
// it will be triggered
|
|
// If there is no listener, it will be dropped safely
|
|
return if failed {
|
|
if cfg!(unix) && oom_killed && parallel_limit > 1 {
|
|
// We want to drop it and stop monitor
|
|
// new invocation will create another one
|
|
drop(done_tx);
|
|
|
|
let reduced_limit = parallel_limit / 2;
|
|
|
|
append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!(
|
|
"\n
|
|
======================
|
|
===== IMPORTANT! =====
|
|
======================
|
|
|
|
Some of installations have been killed by OOM,
|
|
restarting with reduced concurrency: {parallel_limit} -> {reduced_limit}
|
|
|
|
This is not normal behavior, please make sure all workers have enough memory.\n
|
|
"
|
|
),
|
|
conn,
|
|
)
|
|
.await;
|
|
|
|
// restart with half of concurrency
|
|
Box::pin(handle_python_reqs(
|
|
requirements,
|
|
job_id,
|
|
w_id,
|
|
mem_peak,
|
|
_canceled_by,
|
|
conn,
|
|
_worker_name,
|
|
job_dir,
|
|
&worker_dir,
|
|
_occupancy_metrics,
|
|
py_version,
|
|
Some(reduced_limit),
|
|
))
|
|
.await
|
|
} else {
|
|
Err(anyhow!("Env installation did not succeed, check logs").into())
|
|
}
|
|
} else {
|
|
Ok(req_paths)
|
|
};
|
|
}
|
|
|
|
/// Python function body for `res_to_json(res, typ)`.
|
|
/// Handles DataFrame, bytes, dict coercion + JSON serialization.
|
|
fn python_res_to_json_body(postprocessor: &str) -> String {
|
|
format!(
|
|
r#" if typ.__name__ == 'DataFrame':
|
|
if typ.__module__ == 'pandas.core.frame':
|
|
res = res.values.tolist()
|
|
elif typ.__module__ == 'polars.dataframe.frame':
|
|
res = res.rows()
|
|
elif typ.__name__ == 'bytes':
|
|
res = to_b_64(res)
|
|
elif typ.__name__ == 'dict':
|
|
for k, v in res.items():
|
|
if type(v).__name__ == 'bytes':
|
|
res[k] = to_b_64(v)
|
|
unprocessed = json.dumps(res, separators=(',', ':'), default=str).replace('\n', '')
|
|
return {postprocessor}"#
|
|
)
|
|
}
|
|
|
|
// Returns code snippet that needs to be injected into wrapper to post-process results or leave unprocessed
|
|
fn get_result_postprocessor<'a>(skip: bool) -> &'a str {
|
|
if skip {
|
|
"unprocessed"
|
|
} else {
|
|
"_fix(unprocessed)"
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "private")]
|
|
use crate::JobCompletedSender;
|
|
#[cfg(feature = "private")]
|
|
use crate::{common::build_envs_map, dedicated_worker_oss::handle_dedicated_process};
|
|
#[cfg(feature = "private")]
|
|
use windmill_common::variables;
|
|
#[cfg(feature = "private")]
|
|
use windmill_queue::DedicatedWorkerJob;
|
|
|
|
#[cfg(feature = "private")]
|
|
pub async fn start_worker(
|
|
requirements_o: Option<&String>,
|
|
db: &sqlx::Pool<sqlx::Postgres>,
|
|
inner_content: &str,
|
|
base_internal_url: &str,
|
|
job_dir: &str,
|
|
worker_name: &str,
|
|
envs: HashMap<String, String>,
|
|
w_id: &str,
|
|
script_path: &str,
|
|
token: &str,
|
|
job_completed_tx: JobCompletedSender,
|
|
jobs_rx: tokio::sync::mpsc::Receiver<DedicatedWorkerJob>,
|
|
killpill_rx: tokio::sync::broadcast::Receiver<()>,
|
|
client: windmill_common::client::AuthedClient,
|
|
concurrency_semaphore: Option<std::sync::Arc<tokio::sync::Semaphore>>,
|
|
) -> error::Result<()> {
|
|
use crate::PyV;
|
|
tracing::info!("script path: {}", script_path);
|
|
|
|
let codegen = compute_py_codegen(inner_content, script_path);
|
|
|
|
// Write script to proper module path (e.g., f/my/script.py)
|
|
let module_dir = format!("{}/{}", job_dir, codegen.dirs);
|
|
tokio::fs::create_dir_all(&module_dir).await?;
|
|
write_file(
|
|
&module_dir,
|
|
&format!("{}.py", codegen.module_name),
|
|
inner_content,
|
|
)?;
|
|
|
|
let any_relative_imports = RELATIVE_IMPORT_REGEX.is_match(inner_content);
|
|
if any_relative_imports {
|
|
// Dedicated worker runs deployed scripts only — no temp refs.
|
|
let _ = write_file(job_dir, "loader.py", &render_relative_python_loader(&None))?;
|
|
}
|
|
|
|
let mut mem_peak: i32 = 0;
|
|
let mut canceled_by: Option<CanceledBy> = None;
|
|
let context = variables::get_reserved_variables(
|
|
&Connection::Sql(db.clone()),
|
|
w_id,
|
|
&token,
|
|
"dedicated_worker@windmill.dev",
|
|
"dedicated_worker",
|
|
"NOT_AVAILABLE",
|
|
"dedicated_worker",
|
|
Some(script_path.to_string()),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.await
|
|
.to_vec();
|
|
|
|
let annotations = PythonAnnotations::parse(inner_content);
|
|
let context_envs = build_envs_map(context).await;
|
|
let (_, additional_python_paths) = handle_python_deps(
|
|
job_dir,
|
|
requirements_o,
|
|
inner_content,
|
|
w_id,
|
|
script_path,
|
|
&Uuid::nil(),
|
|
&Connection::Sql(db.clone()),
|
|
worker_name,
|
|
job_dir,
|
|
&mut mem_peak,
|
|
&mut canceled_by,
|
|
&mut None,
|
|
None,
|
|
annotations,
|
|
&None, // dedicated worker runs deployed scripts only
|
|
)
|
|
.await?;
|
|
|
|
{
|
|
let scripts = [PyScriptEntry { original_path: script_path, codegen: &codegen }];
|
|
let wrapper_content = generate_multi_script_wrapper(
|
|
&scripts,
|
|
annotations.skip_result_postprocessing,
|
|
any_relative_imports,
|
|
);
|
|
write_file(job_dir, "wrapper.py", &wrapper_content)?;
|
|
}
|
|
|
|
let reserved_variables = windmill_common::variables::get_reserved_variables(
|
|
&Connection::Sql(db.clone()),
|
|
w_id,
|
|
token,
|
|
"dedicated_worker",
|
|
"dedicated_worker",
|
|
Uuid::nil().to_string().as_str(),
|
|
"dedicated_worker",
|
|
Some(script_path.to_string()),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.await;
|
|
|
|
let mut proc_envs: HashMap<String, String> = PYTHON_UTF8_ENVS
|
|
.iter()
|
|
.map(|(k, v)| (k.to_string(), v.to_string()))
|
|
.collect();
|
|
let additional_python_paths_folders = additional_python_paths.iter().join(":");
|
|
proc_envs.insert("PYTHONPATH".to_string(), additional_python_paths_folders);
|
|
proc_envs.insert("PATH".to_string(), PATH_ENV.to_string());
|
|
proc_envs.insert("TZ".to_string(), TZ_ENV.to_string());
|
|
proc_envs.insert(
|
|
"BASE_INTERNAL_URL".to_string(),
|
|
base_internal_url.to_string(),
|
|
);
|
|
proc_envs.insert("BASE_URL".to_string(), base_internal_url.to_string());
|
|
|
|
let py_version = if let Some(requirements) = requirements_o {
|
|
PyV::parse_from_requirements(&split_python_requirements(requirements.as_str()))
|
|
} else {
|
|
tracing::warn!(workspace_id = %w_id, "lockfile is empty for dedicated worker, thus python version cannot be inferred. Fallback to 3.11");
|
|
PyVAlias::default().into()
|
|
};
|
|
|
|
let python_path = py_version
|
|
.get_python(
|
|
worker_name,
|
|
&Uuid::nil(),
|
|
w_id,
|
|
&mut mem_peak,
|
|
&Connection::Sql(db.clone()),
|
|
&mut None,
|
|
)
|
|
.await?;
|
|
|
|
handle_dedicated_process(
|
|
&python_path,
|
|
job_dir,
|
|
context_envs,
|
|
envs,
|
|
reserved_variables,
|
|
proc_envs,
|
|
["-u", "-m", "wrapper"].to_vec(),
|
|
killpill_rx,
|
|
job_completed_tx,
|
|
token,
|
|
jobs_rx,
|
|
worker_name,
|
|
db,
|
|
script_path,
|
|
"python",
|
|
client,
|
|
false,
|
|
concurrency_semaphore,
|
|
)
|
|
.await
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_nested_path() {
|
|
assert_eq!(
|
|
compute_python_module_dir("f/my_folder/my_script"),
|
|
"f/my_folder"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_deep_path() {
|
|
assert_eq!(compute_python_module_dir("f/a/b/c/script"), "f/a/b/c");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_root_level() {
|
|
// Root-level script (no parent dirs) should fall back to "tmp"
|
|
assert_eq!(compute_python_module_dir("my_script"), "tmp");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_single_folder() {
|
|
assert_eq!(compute_python_module_dir("f/script"), "f");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_digit_prefix() {
|
|
// Dirs starting with digits get underscore-prefixed
|
|
assert_eq!(
|
|
compute_python_module_dir("1st_folder/script"),
|
|
"_1st_folder"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_hyphens_replaced() {
|
|
// Hyphens are replaced with underscores
|
|
assert_eq!(
|
|
compute_python_module_dir("my-folder/sub-dir/script"),
|
|
"my_folder/sub_dir"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_at_replaced() {
|
|
// @ is replaced with .
|
|
assert_eq!(compute_python_module_dir("u/@admin/script"), "u/.admin");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_keyword_segment() {
|
|
// A folder whose name is a Python keyword would otherwise produce an
|
|
// invalid `from f.in.x import ...`; it is underscore-prefixed.
|
|
assert_eq!(compute_python_module_dir("f/in/script"), "f/_in");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_python_module_dir_neutralizes_traversal() {
|
|
// A Preview path skips the DB `proper_id` CHECK, so it can carry `..`.
|
|
// `..`/`.` segments must be neutralized so the dir stays inside job_dir.
|
|
let dirs = compute_python_module_dir("u/x/../../../../tmp/evil/payload");
|
|
assert!(!dirs.split('/').any(|s| s == ".." || s == "."));
|
|
assert_eq!(dirs, "u/x/_/_/_/_/tmp/evil");
|
|
// The `@`->`.` rewrite must not be able to synthesize a `..` segment.
|
|
assert_eq!(compute_python_module_dir("u/@./script"), "u/_");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_py_codegen_basic_args() {
|
|
let code = "def main(x: str, y: int):\n return x\n";
|
|
let cg = compute_py_codegen(code, "f/test/script");
|
|
assert!(cg.spread.contains("args[\"x\"]"));
|
|
assert!(cg.spread.contains("args[\"y\"]"));
|
|
assert!(cg.transforms.is_empty());
|
|
assert!(cg.pre_spread.is_none());
|
|
assert_eq!(cg.module_name, "script");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_py_codegen_keyword_step_id() {
|
|
// Regression for a flow inline step whose auto-assigned id is a Python
|
|
// keyword (e.g. `in`): the generated wrapper must not emit
|
|
// `from pkg import in as inner_script` (SyntaxError). The module name is
|
|
// underscore-prefixed, matching the digit-leading guard.
|
|
let code = "def main():\n return 1\n";
|
|
let cg = compute_py_codegen(code, "u/admin/myflow/in");
|
|
assert_eq!(cg.module_name, "_in");
|
|
assert_eq!(cg.module_dir_dot, "u.admin.myflow");
|
|
|
|
// Non-keyword ids are unaffected.
|
|
let cg2 = compute_py_codegen(code, "u/admin/myflow/step");
|
|
assert_eq!(cg2.module_name, "step");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_py_codegen_with_datetime_and_bytes() {
|
|
let code = "import datetime\n\ndef main(name: str, created_at: datetime.datetime, file: bytes):\n return name\n";
|
|
let cg = compute_py_codegen(code, "f/my/handler");
|
|
assert!(cg.transforms.contains("datetime.fromisoformat"));
|
|
assert!(cg.transforms.contains("base64.b64decode"));
|
|
assert!(cg.spread.contains("args[\"name\"]"));
|
|
assert_eq!(cg.module_dir_dot, "f.my");
|
|
assert_eq!(cg.module_name, "handler");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_py_codegen_star_kwargs() {
|
|
let code = "def main(**kwargs):\n return kwargs\n";
|
|
let cg = compute_py_codegen(code, "f/test/star");
|
|
assert_eq!(cg.spread, "args = kwargs");
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_py_codegen_with_preprocessor() {
|
|
let code = "import datetime\n\ndef main(x: str, ts: datetime.datetime):\n return x\n\ndef preprocessor(input: str, when: datetime.datetime):\n return {\"x\": input, \"ts\": when}\n";
|
|
let cg = compute_py_codegen(code, "f/test/pre");
|
|
assert!(cg.spread.contains("args[\"x\"]"));
|
|
assert!(cg.pre_spread.is_some());
|
|
let pre = cg.pre_spread.as_ref().unwrap();
|
|
assert!(pre.contains("pre_args[\"input\"]"));
|
|
}
|
|
|
|
fn lines(v: &[&str]) -> Vec<String> {
|
|
v.iter().map(|s| s.to_string()).collect()
|
|
}
|
|
|
|
#[test]
|
|
fn test_filter_lines_by_deps_no_deps_keeps_everything() {
|
|
let (kept, ignored) = filter_lines_by_deps(lines(&["requests==2.0", "numpy==1.0"]), &[]);
|
|
assert_eq!(kept, lines(&["requests==2.0", "numpy==1.0"]));
|
|
assert!(ignored.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_filter_lines_by_deps_matches_and_partitions() {
|
|
let deps = vec![Regex::new("^my-local-pkg").unwrap()];
|
|
let (kept, ignored) = filter_lines_by_deps(
|
|
lines(&["requests==2.0", "my-local-pkg==1.2.3", "numpy==1.0"]),
|
|
&deps,
|
|
);
|
|
assert_eq!(kept, lines(&["requests==2.0", "numpy==1.0"]));
|
|
assert_eq!(ignored, lines(&["my-local-pkg==1.2.3"]));
|
|
}
|
|
|
|
#[test]
|
|
fn test_filter_lines_by_deps_preserves_comment_lines() {
|
|
// `#` lines (e.g. the `# py: 3.11` lockfile header) must survive even when a
|
|
// dependency regex would otherwise match them.
|
|
let deps = vec![Regex::new("py").unwrap()];
|
|
let (kept, ignored) = filter_lines_by_deps(
|
|
lines(&["# py: 3.11", "pyyaml==6.0", "requests==2.0"]),
|
|
&deps,
|
|
);
|
|
assert_eq!(kept, lines(&["# py: 3.11", "requests==2.0"]));
|
|
assert_eq!(ignored, lines(&["pyyaml==6.0"]));
|
|
}
|
|
|
|
/// Materialize a fake installed wheel: every file in `files` is created, and
|
|
/// `record_entries` is written verbatim as the RECORD (so a test can list a
|
|
/// path in RECORD without creating it, to simulate out-of-band loss).
|
|
fn write_fake_wheel(root: &std::path::Path, files: &[&str], record_entries: &[&str]) {
|
|
for f in files {
|
|
let full = root.join(f);
|
|
std::fs::create_dir_all(full.parent().unwrap()).unwrap();
|
|
std::fs::write(full, b"x").unwrap();
|
|
}
|
|
let dist_info = root.join("pkg-1.0.0.dist-info");
|
|
std::fs::create_dir_all(&dist_info).unwrap();
|
|
std::fs::write(dist_info.join("RECORD"), record_entries.join("\n") + "\n").unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_verify_wheel_record_ok_when_all_present() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
write_fake_wheel(
|
|
dir.path(),
|
|
&["pkg/__init__.py", "pkg/mod.py"],
|
|
&[
|
|
"pkg/__init__.py,sha256=aaa,1",
|
|
"pkg/mod.py,sha256=bbb,1",
|
|
"pkg-1.0.0.dist-info/RECORD,,",
|
|
],
|
|
);
|
|
assert!(verify_wheel_record(dir.path().to_str().unwrap())
|
|
.await
|
|
.is_ok());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_verify_wheel_record_err_when_file_missing() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
// RECORD lists pkg/mod.py but we never create it: the exact failure mode
|
|
// the customer hit (wmill/s3_reader.py present in RECORD, gone on disk).
|
|
write_fake_wheel(
|
|
dir.path(),
|
|
&["pkg/__init__.py"],
|
|
&[
|
|
"pkg/__init__.py,sha256=aaa,1",
|
|
"pkg/mod.py,sha256=bbb,1",
|
|
"pkg-1.0.0.dist-info/RECORD,,",
|
|
],
|
|
);
|
|
let err = verify_wheel_record(dir.path().to_str().unwrap())
|
|
.await
|
|
.unwrap_err();
|
|
assert!(err.contains("pkg/mod.py"), "unexpected error: {err}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_verify_wheel_record_err_when_no_dist_info() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
std::fs::write(dir.path().join("loose.py"), b"x").unwrap();
|
|
assert!(verify_wheel_record(dir.path().to_str().unwrap())
|
|
.await
|
|
.is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_verify_wheel_record_skips_absolute_and_escaping_entries() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
// Absolute and `..` RECORD entries are not package-relative and must be
|
|
// skipped rather than reported missing.
|
|
write_fake_wheel(
|
|
dir.path(),
|
|
&["pkg/__init__.py"],
|
|
&[
|
|
"pkg/__init__.py,sha256=aaa,1",
|
|
"/etc/passwd,sha256=ccc,1",
|
|
"../outside.py,sha256=ddd,1",
|
|
"pkg-1.0.0.dist-info/RECORD,,",
|
|
],
|
|
);
|
|
assert!(verify_wheel_record(dir.path().to_str().unwrap())
|
|
.await
|
|
.is_ok());
|
|
}
|
|
|
|
// Regression tests for the concurrent-install guard. Two jobs installing the
|
|
// same uncached dep into the shared `venv_p` used to race uv's `--reinstall`,
|
|
// corrupting the on-disk wheel and failing with "Env installation did not
|
|
// succeed". The guard serializes those installs.
|
|
|
|
#[tokio::test]
|
|
async fn test_venv_install_lock_serializes_same_path() {
|
|
use std::sync::atomic::{AtomicUsize, Ordering};
|
|
// Same target path => one shared lock => no two tasks install at once.
|
|
let active = Arc::new(AtomicUsize::new(0));
|
|
let max_seen = Arc::new(AtomicUsize::new(0));
|
|
let mut handles = vec![];
|
|
for _ in 0..8 {
|
|
let active = active.clone();
|
|
let max_seen = max_seen.clone();
|
|
handles.push(tokio::spawn(async move {
|
|
let lock = get_venv_install_lock("/cache/py/3.11/samedep==1.0").await;
|
|
let _g = lock.lock_owned().await;
|
|
let cur = active.fetch_add(1, Ordering::SeqCst) + 1;
|
|
max_seen.fetch_max(cur, Ordering::SeqCst);
|
|
// Yield so any concurrency would be observed by another task.
|
|
tokio::time::sleep(std::time::Duration::from_millis(5)).await;
|
|
active.fetch_sub(1, Ordering::SeqCst);
|
|
}));
|
|
}
|
|
for h in handles {
|
|
h.await.unwrap();
|
|
}
|
|
assert_eq!(
|
|
max_seen.load(Ordering::SeqCst),
|
|
1,
|
|
"installs into the same target dir must be serialized"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_venv_install_lock_distinct_paths_are_independent() {
|
|
// Different target paths get different locks and never block each other.
|
|
let a = get_venv_install_lock("/cache/py/3.11/depA==1.0").await;
|
|
let b = get_venv_install_lock("/cache/py/3.11/depB==1.0").await;
|
|
let _ga = a.lock_owned().await;
|
|
// Holding depA's lock must not prevent acquiring depB's.
|
|
assert!(
|
|
b.try_lock().is_ok(),
|
|
"distinct deps must install in parallel"
|
|
);
|
|
// Same path returns the same underlying lock.
|
|
let a2 = get_venv_install_lock("/cache/py/3.11/depA==1.0").await;
|
|
assert!(
|
|
a2.try_lock().is_err(),
|
|
"same target dir must map to the same lock"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_venv_file_lock_excludes_across_descriptions() {
|
|
// The cross-process layer: an advisory lock on a sibling `.lock` excludes a
|
|
// second independent open file handle (i.e. another worker process) while
|
|
// held, and frees it on close. Mirrors the loop in handle_python_reqs and
|
|
// must hold on every platform (flock on unix, LockFileEx on windows) — a
|
|
// Windows host running several agents against one wheel cache relies on it.
|
|
use fs4::fs_std::FileExt;
|
|
|
|
let dir = std::env::temp_dir().join("wm_venv_lock_test");
|
|
std::fs::create_dir_all(&dir).unwrap();
|
|
let lock_path = dir.join("dep==1.0.lock");
|
|
|
|
let f1 = std::fs::OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.open(&lock_path)
|
|
.unwrap();
|
|
assert!(
|
|
f1.try_lock_exclusive().unwrap(),
|
|
"first holder must acquire the lock"
|
|
);
|
|
|
|
// A second handle (stand-in for another process) cannot take it.
|
|
let f2 = std::fs::OpenOptions::new()
|
|
.create(true)
|
|
.write(true)
|
|
.open(&lock_path)
|
|
.unwrap();
|
|
assert!(
|
|
!f2.try_lock_exclusive().unwrap(),
|
|
"a second holder must be blocked while the lock is held"
|
|
);
|
|
|
|
// Releasing the first lets the second acquire it.
|
|
drop(f1);
|
|
assert!(
|
|
f2.try_lock_exclusive().unwrap(),
|
|
"lock must be acquirable once the holder releases it"
|
|
);
|
|
|
|
drop(f2);
|
|
let _ = std::fs::remove_file(&lock_path);
|
|
}
|
|
}
|