mirror of
https://github.com/windmill-labs/windmill.git
synced 2026-10-19 16:03:29 +00:00
Ruben FiszelandClaude Opus 5
de98adf055
* feat: let a worker group override the dependency cache object store Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R6hBAWqNUsQMAJ7P59juug * fix: address review findings on the worker-group cache override Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R6hBAWqNUsQMAJ7P59juug * fix: close the remaining config read route and re-evaluate the override on plan change Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R6hBAWqNUsQMAJ7P59juug * fix: serialize override reloads and keep a store a failed rebuild still serves Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R6hBAWqNUsQMAJ7P59juug * fix: require enterprise for the cache override and lock its whole transition Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01R6hBAWqNUsQMAJ7P59juug --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
1039 lines
34 KiB
Rust
1039 lines
34 KiB
Rust
use std::collections::{HashMap, HashSet, VecDeque};
|
|
use std::sync::Arc;
|
|
|
|
use anyhow::bail;
|
|
use itertools::Itertools;
|
|
use process_wrap::tokio::TokioChildWrapper;
|
|
use tokio::{
|
|
fs::File,
|
|
process::Command,
|
|
sync::{broadcast::Sender, OwnedSemaphorePermit, RwLock, Semaphore},
|
|
task::JoinHandle,
|
|
};
|
|
use uuid::Uuid;
|
|
use windmill_common::{error, worker::Connection};
|
|
|
|
use crate::{common::start_child_process, is_sandboxing_enabled};
|
|
|
|
#[allow(dead_code)]
|
|
#[derive(Clone, Debug)]
|
|
pub struct RequiredDependency<T: Clone + Send + Sync> {
|
|
/// Expected directory of dependency in cache
|
|
/// For example:
|
|
/// /tmp/windmill/cache/python_311/rich==0.0.0
|
|
/// IMPORTANT!: path should not end with '/'
|
|
pub path: String,
|
|
/// Name to use for S3 tars
|
|
pub _s3_handle: String,
|
|
/// Display name
|
|
/// Name that will be used for console output and logging
|
|
pub display_name: String,
|
|
/// Custom payload can be used for passing information easily. If you wish not to use it just leave '()'
|
|
pub custom_payload: T,
|
|
}
|
|
|
|
/// Generic dependency graph that produces topologically sorted layers via Kahn's algorithm.
|
|
/// Each layer's packages only depend on packages from earlier layers, enabling parallel install
|
|
/// per layer.
|
|
#[allow(dead_code)]
|
|
pub struct DependencyGraph<T: Clone + Send + Sync> {
|
|
nodes: HashMap<String, RequiredDependency<T>>,
|
|
deps: HashMap<String, HashSet<String>>,
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
impl<T: Clone + Send + Sync> DependencyGraph<T> {
|
|
pub fn new() -> Self {
|
|
Self { nodes: HashMap::new(), deps: HashMap::new() }
|
|
}
|
|
|
|
/// Insert a dependency and the names of packages it depends on.
|
|
/// References to packages not in the graph are silently ignored during layering.
|
|
pub fn insert(
|
|
&mut self,
|
|
key: impl Into<String>,
|
|
dep: RequiredDependency<T>,
|
|
depends_on: Vec<String>,
|
|
) {
|
|
let key = key.into();
|
|
self.nodes.insert(key.clone(), dep);
|
|
self.deps.insert(key, depends_on.into_iter().collect());
|
|
}
|
|
|
|
/// Render a dependency tree string. Each package appears once, nested under the first parent
|
|
/// that pulls it in. Only includes packages present in `filter` (if provided).
|
|
pub fn print_tree(&self, filter: Option<&HashSet<String>>) -> String {
|
|
// Find roots: packages nothing else in the graph depends on
|
|
let mut depended_on: HashSet<&str> = HashSet::new();
|
|
for dep_set in self.deps.values() {
|
|
for to in dep_set {
|
|
if self.nodes.contains_key(to) {
|
|
depended_on.insert(to.as_str());
|
|
}
|
|
}
|
|
}
|
|
let roots: Vec<&String> = self
|
|
.nodes
|
|
.keys()
|
|
.filter(|k| !depended_on.contains(k.as_str()))
|
|
.filter(|k| filter.map_or(true, |f| f.contains(*k)))
|
|
.sorted()
|
|
.collect();
|
|
|
|
let mut out = String::new();
|
|
let mut seen = HashSet::new();
|
|
for root in roots {
|
|
self.print_tree_node(root, 0, &mut seen, filter, &mut out);
|
|
}
|
|
out
|
|
}
|
|
|
|
fn print_tree_node(
|
|
&self,
|
|
key: &str,
|
|
depth: usize,
|
|
seen: &mut HashSet<String>,
|
|
filter: Option<&HashSet<String>>,
|
|
out: &mut String,
|
|
) {
|
|
if !seen.insert(key.to_string()) {
|
|
return;
|
|
}
|
|
if let Some(dep) = self.nodes.get(key) {
|
|
let indent = " ".repeat(depth);
|
|
out.push_str(&format!("{}- {}\n", indent, dep.display_name));
|
|
if let Some(children) = self.deps.get(key) {
|
|
for child in children.iter().sorted() {
|
|
if self.nodes.contains_key(child)
|
|
&& filter.map_or(true, |f| f.contains(child))
|
|
&& !seen.contains(child)
|
|
{
|
|
self.print_tree_node(child, depth + 1, seen, filter, out);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Produce topologically sorted layers.
|
|
pub fn layers(self) -> Vec<Vec<RequiredDependency<T>>> {
|
|
let mut in_degree: HashMap<String, usize> =
|
|
self.nodes.keys().map(|k| (k.clone(), 0)).collect();
|
|
let mut reverse: HashMap<String, Vec<String>> = HashMap::new();
|
|
|
|
for (from, dep_set) in &self.deps {
|
|
for to in dep_set {
|
|
if self.nodes.contains_key(to) {
|
|
*in_degree.entry(from.clone()).or_default() += 1;
|
|
reverse.entry(to.clone()).or_default().push(from.clone());
|
|
}
|
|
}
|
|
}
|
|
|
|
let mut queue: VecDeque<String> = in_degree
|
|
.iter()
|
|
.filter(|(_, &d)| d == 0)
|
|
.map(|(k, _)| k.clone())
|
|
.sorted()
|
|
.collect();
|
|
|
|
let mut result = vec![];
|
|
let mut nodes = self.nodes;
|
|
|
|
while !queue.is_empty() {
|
|
let mut layer = vec![];
|
|
let mut next = VecDeque::new();
|
|
|
|
for key in queue {
|
|
if let Some(dep) = nodes.remove(&key) {
|
|
layer.push(dep);
|
|
}
|
|
if let Some(dependents) = reverse.get(&key) {
|
|
for d in dependents {
|
|
if let Some(deg) = in_degree.get_mut(d) {
|
|
*deg -= 1;
|
|
if *deg == 0 {
|
|
next.push_back(d.clone());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
if !layer.is_empty() {
|
|
result.push(layer);
|
|
}
|
|
queue = next.into_iter().sorted().collect();
|
|
}
|
|
|
|
result
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub enum InstallDeps<T: Clone + Send + Sync> {
|
|
/// Flat list of dependencies — installed in one parallel batch (existing behavior).
|
|
Flat(Vec<RequiredDependency<T>>),
|
|
/// Dependency graph — split into topological layers, each installed in parallel.
|
|
/// A `--- Layer N ---` separator is printed between layers.
|
|
Layered(DependencyGraph<T>),
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
pub enum InstallStrategy<T: Clone + Send + Sync> {
|
|
/// Will invoke callback to install single dependency
|
|
Single(Arc<dyn Fn(RequiredDependency<T>) -> Result<Command, error::Error> + Send + Sync>),
|
|
/// Will try to pull S3 first and will invoke closure to install the rest
|
|
AllAtOnce(Arc<RwLock<Vec<RequiredDependency<T>>>>),
|
|
}
|
|
|
|
/// # General
|
|
///
|
|
/// Languages that compile usually include dependencies in final executable.
|
|
/// When dynamic languages do not and runtime dependencies are provided separately.
|
|
///
|
|
/// This helper assumes that the language is dynamic.
|
|
/// Python, Ruby, Java are dynamic and they can use this helper.
|
|
///
|
|
/// # Features
|
|
///
|
|
/// This helper will install all specified dependencies in parallel and if it is EE, cache to S3
|
|
/// It has atomic success file, allowing to distinguish failed installations from succesfull.
|
|
///
|
|
/// Besides that it provides console output and does logging.
|
|
///
|
|
/// # Usage
|
|
///
|
|
/// Most important arguments in this helper are `deps` and `install_fn`
|
|
///
|
|
/// In `deps` you specify all dependencies that are needed to be on worker in order to execute script.
|
|
/// You don't know which are actually installed and which are not.
|
|
///
|
|
/// `deps` is a vector of RequiredDependency. Check [RequiredDependency] for more context.
|
|
///
|
|
/// After `deps` are provided helper will check each dependency and check if it is in cache, if not it will try to pull from S3
|
|
/// and if it does not work either, it will invoke `install_fn` closure.
|
|
/// Closure arguments has dependency name as well as its expected path in cache.
|
|
/// Closure should return Command that will install dependency to asked place.
|
|
#[allow(dead_code)]
|
|
pub async fn par_install_language_dependencies_all_at_once<
|
|
'a,
|
|
T: Clone + std::marker::Send + Sync + 'a + 'static,
|
|
>(
|
|
deps: Vec<RequiredDependency<T>>,
|
|
_language_name: &'a str,
|
|
installer_executable_name: &'a str,
|
|
_platform_agnostic: bool,
|
|
concurrent_downloads: usize,
|
|
stdout_on_err: bool,
|
|
callback: impl Fn(Vec<RequiredDependency<T>>) -> Result<Command, error::Error>
|
|
+ Send
|
|
+ Sync
|
|
+ 'static,
|
|
postinstall_cb: impl AsyncFn(Vec<RequiredDependency<T>>) -> Result<(), error::Error>,
|
|
job_id: &'a Uuid,
|
|
w_id: &'a str,
|
|
worker_name: &'a str,
|
|
jailed: bool,
|
|
conn: &'a Connection,
|
|
) -> anyhow::Result<()> {
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let is_not_pro = !matches!(
|
|
windmill_common::ee_oss::get_license_plan().await,
|
|
windmill_common::ee_oss::LicensePlan::Pro
|
|
);
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if is_not_pro {
|
|
windmill_queue::append_logs(
|
|
job_id,
|
|
w_id,
|
|
format!("\nLooking for packages on S3:\n"),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|
|
let total_time = std::time::Instant::now();
|
|
let (layers, name_max_length, total_missing) =
|
|
filter_to_missing(InstallDeps::Flat(deps), job_id, w_id, jailed, conn).await?;
|
|
let missing: Vec<RequiredDependency<T>> = layers.into_iter().flatten().collect();
|
|
if total_missing == 0 {
|
|
return Ok(());
|
|
}
|
|
let to_batch_install = Arc::new(RwLock::new(vec![]));
|
|
let handles = spawn_wrapped_installation_threads(
|
|
missing,
|
|
name_max_length,
|
|
InstallStrategy::AllAtOnce(to_batch_install.clone()),
|
|
installer_executable_name,
|
|
concurrent_downloads,
|
|
job_id,
|
|
w_id,
|
|
worker_name,
|
|
conn,
|
|
_language_name,
|
|
_platform_agnostic,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.await?;
|
|
let installation_res = process_handles(handles, w_id).await;
|
|
if !to_batch_install.read().await.is_empty() {
|
|
let for_all_at_once_copy = to_batch_install.read().await.clone();
|
|
windmill_queue::append_logs(
|
|
job_id,
|
|
w_id,
|
|
format!("\n\nFetching {} packages...\n", for_all_at_once_copy.len()),
|
|
&conn,
|
|
)
|
|
.await;
|
|
let cmd = callback(for_all_at_once_copy.clone())?;
|
|
let child = start_child_process(cmd, &installer_executable_name, false).await?;
|
|
let mut buf = "".to_owned();
|
|
let pipe_stdout = if stdout_on_err { Some(&mut buf) } else { None };
|
|
// handle_child only pings the job id it is handed, and that one is nil while output is
|
|
// buffered, so nothing would refresh this job's ping for the length of the batch install.
|
|
let _heartbeat = stdout_on_err.then(|| {
|
|
crate::worker_utils::JobPingHeartbeat::start(conn, *job_id, "dependency installation")
|
|
});
|
|
|
|
if let Err(e) = crate::handle_child::handle_child(
|
|
&(if pipe_stdout.is_some() {
|
|
Uuid::nil()
|
|
} else {
|
|
*job_id
|
|
}),
|
|
conn,
|
|
// TODO: Return mem_peak
|
|
&mut 0,
|
|
// TODO: Return canceld_by_ref
|
|
&mut None,
|
|
child,
|
|
is_sandboxing_enabled(),
|
|
&worker_name,
|
|
&w_id,
|
|
&installer_executable_name,
|
|
None,
|
|
false,
|
|
&mut None,
|
|
pipe_stdout,
|
|
None,
|
|
)
|
|
.await
|
|
{
|
|
bail!(format!(
|
|
"\nerror while installing dependencies: {e:?}\n{}",
|
|
buf
|
|
));
|
|
}
|
|
{
|
|
postinstall_cb(for_all_at_once_copy.clone()).await?;
|
|
}
|
|
for RequiredDependency { path, _s3_handle, .. } in for_all_at_once_copy.into_iter() {
|
|
mark_success(path.clone(), job_id, w_id).await;
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
{
|
|
if let Some(os) = windmill_object_store::get_cache_object_store().await {
|
|
let language_name = _language_name.to_owned();
|
|
tokio::spawn(async move {
|
|
if let Err(e) = crate::global_cache::build_tar_and_push(
|
|
os,
|
|
path,
|
|
language_name,
|
|
Some(_s3_handle),
|
|
_platform_agnostic,
|
|
)
|
|
.await
|
|
{
|
|
tracing::warn!("failed to build tar and push: {e:?}");
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
}
|
|
finish_installation(total_time, job_id, w_id, conn).await;
|
|
installation_res
|
|
}
|
|
|
|
/// # General
|
|
///
|
|
/// Languages that compile usually include dependencies in final executable.
|
|
/// When dynamic languages do not and runtime dependencies are provided separately.
|
|
///
|
|
/// This helper assumes that the language is dynamic.
|
|
/// Python, Ruby, Java are dynamic and they can use this helper.
|
|
///
|
|
/// # Features
|
|
///
|
|
/// This helper will install all specified dependencies in parallel and if it is EE, cache to S3
|
|
/// It has atomic success file, allowing to distinguish failed installations from succesfull.
|
|
///
|
|
/// Besides that it provides console output and does logging.
|
|
///
|
|
/// # Usage
|
|
///
|
|
/// Most important arguments in this helper are `deps` and `install_fn`
|
|
///
|
|
/// In `deps` you specify all dependencies that are needed to be on worker in order to execute script.
|
|
/// You don't know which are actually installed and which are not.
|
|
///
|
|
/// `deps` is a vector of RequiredDependency. Check [RequiredDependency] for more context.
|
|
///
|
|
/// After `deps` are provided helper will check each dependency and check if it is in cache, if not it will try to pull from S3
|
|
/// and if it does not work either, it will invoke `install_fn` closure.
|
|
/// Closure arguments has dependency name as well as it`s expected path in cache.
|
|
/// Closure should return Command that will install dependency to asked place.
|
|
#[allow(dead_code)]
|
|
pub async fn par_install_language_dependencies_seq<
|
|
'a,
|
|
T: Clone + std::marker::Send + Sync + 'a + 'static,
|
|
>(
|
|
install_deps: InstallDeps<T>,
|
|
_language_name: &'a str,
|
|
installer_executable_name: &'a str,
|
|
_platform_agnostic: bool,
|
|
concurrent_downloads: usize,
|
|
callback: impl Fn(RequiredDependency<T>) -> Result<Command, error::Error> + Send + Sync + 'static,
|
|
post_install: Option<
|
|
Arc<dyn Fn(&RequiredDependency<T>) -> anyhow::Result<()> + Send + Sync + 'static>,
|
|
>,
|
|
job_id: &'a Uuid,
|
|
w_id: &'a str,
|
|
worker_name: &'a str,
|
|
jailed: bool,
|
|
conn: &'a Connection,
|
|
) -> anyhow::Result<()> {
|
|
let total_time = std::time::Instant::now();
|
|
let (layers, name_max_length, total_missing) =
|
|
filter_to_missing(install_deps, job_id, w_id, jailed, conn).await?;
|
|
if total_missing == 0 {
|
|
return Ok(());
|
|
}
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let is_not_pro = !matches!(
|
|
windmill_common::ee_oss::get_license_plan().await,
|
|
windmill_common::ee_oss::LicensePlan::Pro
|
|
);
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if is_not_pro {
|
|
windmill_queue::append_logs(
|
|
job_id,
|
|
w_id,
|
|
format!("\nLooking for packages on S3:\n"),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|
|
|
|
let is_layered = layers.len() > 1;
|
|
let callback = Arc::new(callback);
|
|
let mut offset = 0usize;
|
|
|
|
windmill_queue::append_logs(
|
|
job_id,
|
|
w_id,
|
|
if jailed {
|
|
format!(
|
|
"\nStarting isolated installation... ({} tasks in parallel)\n",
|
|
concurrent_downloads
|
|
)
|
|
} else {
|
|
format!(
|
|
"\nStarting installation... ({} tasks in parallel)\n",
|
|
concurrent_downloads
|
|
)
|
|
},
|
|
conn,
|
|
)
|
|
.await;
|
|
|
|
for (i, layer_deps) in layers.into_iter().enumerate() {
|
|
if layer_deps.is_empty() {
|
|
continue;
|
|
}
|
|
|
|
if is_layered && offset > 0 {
|
|
windmill_queue::append_logs(job_id, w_id, format!("\n\n--- Layer {} ---", i + 1), conn)
|
|
.await;
|
|
}
|
|
|
|
let layer_size = layer_deps.len();
|
|
tracing::info!("Layer {}: spawning {} installs", i + 1, layer_size);
|
|
let handles = spawn_wrapped_installation_threads(
|
|
layer_deps,
|
|
name_max_length,
|
|
InstallStrategy::Single(callback.clone()),
|
|
installer_executable_name,
|
|
concurrent_downloads,
|
|
job_id,
|
|
w_id,
|
|
worker_name,
|
|
conn,
|
|
_language_name,
|
|
_platform_agnostic,
|
|
Some(offset),
|
|
Some(total_missing),
|
|
post_install.clone(),
|
|
)
|
|
.await?;
|
|
tracing::info!("Layer {}: all spawned, waiting for handles", i + 1);
|
|
|
|
process_handles(handles, w_id).await?;
|
|
tracing::info!("Layer {}: done", i + 1);
|
|
offset += layer_size;
|
|
}
|
|
|
|
finish_installation(total_time, job_id, w_id, conn).await;
|
|
Ok(())
|
|
}
|
|
|
|
type NameMaxLength = usize;
|
|
|
|
/// Returns (layers of missing deps, name_max_length, total_missing).
|
|
/// Prints the "To be installed" header once with all missing packages.
|
|
/// For `Layered`, prints a dependency tree; for `Flat`, prints a flat list.
|
|
async fn filter_to_missing<'a, T: Clone + std::marker::Send + Sync + 'a + 'static>(
|
|
install_deps: InstallDeps<T>,
|
|
job_id: &Uuid,
|
|
w_id: &str,
|
|
jailed: bool,
|
|
conn: &Connection,
|
|
) -> anyhow::Result<(Vec<Vec<RequiredDependency<T>>>, NameMaxLength, usize)> {
|
|
let (mut layers, tree_data) = match install_deps {
|
|
InstallDeps::Flat(deps) => (vec![deps], None),
|
|
InstallDeps::Layered(graph) => {
|
|
let deps_map = graph.deps.clone();
|
|
let nodes_display: HashMap<String, String> = graph
|
|
.nodes
|
|
.iter()
|
|
.map(|(k, v)| (k.clone(), v.display_name.clone()))
|
|
.collect();
|
|
let layers = graph.layers();
|
|
(layers, Some((deps_map, nodes_display)))
|
|
}
|
|
};
|
|
|
|
let mut name_ml = 0;
|
|
let mut missing_keys: HashSet<String> = HashSet::new();
|
|
let mut total_missing = 0;
|
|
|
|
for layer in layers.iter_mut() {
|
|
*layer = std::mem::take(layer)
|
|
.into_iter()
|
|
.unique_by(|rd| rd.path.clone())
|
|
.collect();
|
|
|
|
let mut missing = vec![];
|
|
for rd in std::mem::take(layer) {
|
|
if rd.path.ends_with("/") {
|
|
anyhow::bail!("Internal error: path should not end with '/'")
|
|
}
|
|
if rd.display_name.len() > name_ml {
|
|
name_ml = rd.display_name.len();
|
|
}
|
|
if tokio::fs::metadata(rd.path.clone() + ".valid.windmill")
|
|
.await
|
|
.is_err()
|
|
{
|
|
if let Some(key) = rd.path.rsplit('/').next() {
|
|
missing_keys.insert(key.to_string());
|
|
}
|
|
missing.push(rd);
|
|
}
|
|
}
|
|
total_missing += missing.len();
|
|
*layer = missing;
|
|
}
|
|
|
|
if total_missing > 0 {
|
|
windmill_queue::append_logs(
|
|
job_id,
|
|
w_id,
|
|
if jailed {
|
|
format!("\n--- ISOLATED INSTALLATION ---\n\nTo be installed:\n\n")
|
|
} else {
|
|
format!("\n--- INSTALLATION ---\n\nTo be installed:\n\n")
|
|
},
|
|
conn,
|
|
)
|
|
.await;
|
|
|
|
let to_log = if let Some((deps_map, nodes_display)) = tree_data {
|
|
let mut print_graph: DependencyGraph<()> = DependencyGraph::new();
|
|
for (key, display) in &nodes_display {
|
|
if missing_keys.contains(key) {
|
|
print_graph.insert(
|
|
key.clone(),
|
|
RequiredDependency {
|
|
path: String::new(),
|
|
_s3_handle: String::new(),
|
|
display_name: display.clone(),
|
|
custom_payload: (),
|
|
},
|
|
deps_map
|
|
.get(key)
|
|
.map(|s| s.iter().cloned().collect())
|
|
.unwrap_or_default(),
|
|
);
|
|
}
|
|
}
|
|
print_graph.print_tree(Some(&missing_keys))
|
|
} else {
|
|
layers
|
|
.iter()
|
|
.flat_map(|l| l.iter())
|
|
.map(|rd| format!("- {}", &rd.display_name))
|
|
.join("\n")
|
|
+ "\n"
|
|
};
|
|
|
|
windmill_queue::append_logs(job_id, w_id, to_log, conn).await;
|
|
}
|
|
|
|
Ok((layers, name_ml, total_missing))
|
|
}
|
|
|
|
enum Action<T: Clone + Send + Sync> {
|
|
Install(Box<dyn TokioChildWrapper + 'static>),
|
|
AddToBulk(Arc<RwLock<Vec<RequiredDependency<T>>>>),
|
|
}
|
|
struct TaskKiller(tokio::sync::broadcast::Sender<String>);
|
|
impl Drop for TaskKiller {
|
|
fn drop(&mut self) {
|
|
// We don't care if it actually stopped anything or not.
|
|
// In some cases it will fire and actually stop
|
|
// In other cases do nothing
|
|
#[allow(unused_must_use)]
|
|
self.0.send("Installation failed".to_owned());
|
|
}
|
|
}
|
|
// Callback will be invoked only if could not pull from S3
|
|
async fn spawn_wrapped_installation_threads<
|
|
'a,
|
|
T: Clone + std::marker::Send + Sync + 'a + 'static,
|
|
>(
|
|
missing: Vec<RequiredDependency<T>>,
|
|
name_ml: usize,
|
|
strategy: InstallStrategy<T>,
|
|
installer_executable_name: &str,
|
|
concurrent_downloads: usize,
|
|
job_id: &Uuid,
|
|
w_id: &str,
|
|
worker_name: &str,
|
|
conn: &Connection,
|
|
_language_name: &str,
|
|
_platform_agnostic: bool,
|
|
counter_offset: Option<usize>,
|
|
total_override: Option<usize>,
|
|
post_install: Option<
|
|
Arc<dyn Fn(&RequiredDependency<T>) -> anyhow::Result<()> + Send + Sync + 'static>,
|
|
>,
|
|
) -> anyhow::Result<(
|
|
Vec<JoinHandle<anyhow::Result<TaskKiller>>>,
|
|
tokio::sync::broadcast::Sender<()>,
|
|
)> {
|
|
// Semaphore will panic if value less then 1
|
|
let parallel_limit = concurrent_downloads.clamp(1, 30);
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
"Install parallel limit: {}, job: {}",
|
|
parallel_limit,
|
|
job_id
|
|
);
|
|
|
|
let total_to_install = total_override.unwrap_or(missing.len());
|
|
let (mut handles, semaphore, counter_arc) = (
|
|
vec![],
|
|
Arc::new(Semaphore::new(parallel_limit)),
|
|
Arc::new(tokio::sync::Mutex::new(counter_offset.unwrap_or(0))),
|
|
);
|
|
|
|
// Pretty sensitive. Single drop will fail installation
|
|
let (kill_tx, ..) = tokio::sync::broadcast::channel::<String>(10);
|
|
// ^^^^^^^ Original will go into listen_to_cancel
|
|
//
|
|
// But before that it will be used to populate t and r for every thread
|
|
let trxs: Vec<(
|
|
tokio::sync::broadcast::Sender<String>,
|
|
tokio::sync::broadcast::Receiver<String>,
|
|
)> = (0..missing.len())
|
|
.map(|_| (kill_tx.clone(), kill_tx.subscribe()))
|
|
.collect();
|
|
|
|
// Listen to cancel events and stop jobs
|
|
// NOTE:
|
|
// once closer is dropped, the listener will be closed
|
|
let closer = listen_to_cancel(job_id.clone(), w_id.to_owned(), conn.clone(), kill_tx).await;
|
|
|
|
for (dep, (kill_tx, mut kill_rx)) in missing.into_iter().zip(trxs) {
|
|
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();
|
|
|
|
let action = match strategy {
|
|
InstallStrategy::Single(ref callback) => Action::Install(
|
|
start_child_process(callback(dep.clone())?, &installer_executable_name, false)
|
|
.await?,
|
|
),
|
|
InstallStrategy::AllAtOnce(ref rw_lock) => Action::AddToBulk(Arc::clone(rw_lock)),
|
|
};
|
|
let post_install_c = post_install.clone();
|
|
let task_fut = try_install_one_detached(
|
|
dep,
|
|
installer_executable_name.to_owned(),
|
|
job_id.clone(),
|
|
w_id.to_owned(),
|
|
worker_name.to_owned(),
|
|
conn.to_owned(),
|
|
counter_arc.clone(),
|
|
total_to_install,
|
|
name_ml,
|
|
action,
|
|
_language_name.to_owned(),
|
|
_platform_agnostic,
|
|
permit,
|
|
TaskKiller(kill_tx),
|
|
post_install_c,
|
|
);
|
|
handles.push(tokio::spawn(async move {
|
|
tokio::select! {
|
|
reason = kill_rx.recv() => return Err(anyhow::anyhow!("Installation stopped. Reason: {:?}", reason)),
|
|
r = task_fut => return r
|
|
};
|
|
}));
|
|
}
|
|
|
|
Ok((handles, closer))
|
|
}
|
|
/// Do not drop the returned sender! As soon as you drop it the polling thread will stop
|
|
pub async fn listen_to_cancel(
|
|
job_id: Uuid,
|
|
w_id: String,
|
|
conn: Connection,
|
|
killpill: tokio::sync::broadcast::Sender<String>,
|
|
) -> tokio::sync::broadcast::Sender<()> {
|
|
let (close_tx, mut close_rx) = tokio::sync::broadcast::channel::<()>(1);
|
|
tokio::spawn(async move {
|
|
'outer: loop {
|
|
tokio::select! {
|
|
_ = tokio::time::sleep(tokio::time::Duration::from_secs(1)) => {
|
|
// TODO: Add memory usage tracking
|
|
if crate::worker_utils::ping_job_status(&conn, &job_id, None, None).await.map(|r| r.canceled_by.is_some()).unwrap_or(false){
|
|
tracing::info!(
|
|
// If there is listener on other side,
|
|
workspace_id = %w_id,
|
|
"cancelling installations",
|
|
);
|
|
if let Err(ref e) = killpill.send("Job was canceled".to_owned()) {
|
|
tracing::error!(
|
|
// If there is listener on other side,
|
|
workspace_id = %w_id,
|
|
"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
|
|
);
|
|
}
|
|
// Stop
|
|
break 'outer;
|
|
}
|
|
|
|
},
|
|
_ = close_rx.recv() => break 'outer,
|
|
}
|
|
}
|
|
});
|
|
|
|
// We will return the sender that can stop created thread.
|
|
// It is relatively safe, since if this handle is being dropped from the gecko the thread will stop immediately
|
|
// but it will also chainreaction all installation threads, so be carefull
|
|
close_tx
|
|
}
|
|
|
|
async fn try_install_one_detached<'a, T: Clone + std::marker::Send + Sync + 'a + 'static>(
|
|
dep: RequiredDependency<T>,
|
|
installer_executable_name: String,
|
|
job_id: Uuid,
|
|
w_id: String,
|
|
worker_name: String,
|
|
conn: Connection,
|
|
counter_arc: Arc<tokio::sync::Mutex<usize>>,
|
|
total_to_install: usize,
|
|
name_ml: usize,
|
|
action: Action<T>,
|
|
_language_name: String,
|
|
_platform_agnostic: bool,
|
|
_permit: OwnedSemaphorePermit,
|
|
// If dropped the entire installation fails and all installation threads are being stopped
|
|
// That's why we just pass it to return so it is not being dropped
|
|
kill_all_tasks: TaskKiller,
|
|
post_install: Option<
|
|
Arc<dyn Fn(&RequiredDependency<T>) -> anyhow::Result<()> + Send + Sync + 'static>,
|
|
>,
|
|
) -> anyhow::Result<TaskKiller> {
|
|
let start = std::time::Instant::now();
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let is_not_pro = !matches!(
|
|
windmill_common::ee_oss::get_license_plan().await,
|
|
windmill_common::ee_oss::LicensePlan::Pro
|
|
);
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
let s3_pull_future = if is_not_pro {
|
|
if let Some(os) = windmill_object_store::get_cache_object_store().await {
|
|
Some(crate::global_cache::pull_from_tar(
|
|
os,
|
|
dep.path.clone(),
|
|
_language_name.to_owned(),
|
|
Some(dep._s3_handle.clone()),
|
|
_platform_agnostic,
|
|
))
|
|
} else {
|
|
None
|
|
}
|
|
} else {
|
|
None
|
|
};
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
if let Some(s3_pull_future) = s3_pull_future {
|
|
if let Err(e) = s3_pull_future.await {
|
|
tracing::info!(
|
|
workspace_id = %w_id,
|
|
"No tarball was found for {:?} on S3 or different problem occurred {job_id}:\n{e}",
|
|
&dep._s3_handle.clone()
|
|
);
|
|
} else {
|
|
mark_success(dep.path, &job_id, &w_id).await;
|
|
print_success(
|
|
true,
|
|
false,
|
|
&job_id,
|
|
&w_id,
|
|
&dep.display_name,
|
|
name_ml,
|
|
counter_arc,
|
|
total_to_install,
|
|
start,
|
|
&conn,
|
|
)
|
|
.await;
|
|
return Ok(kill_all_tasks);
|
|
}
|
|
}
|
|
let child = match action {
|
|
Action::Install(child) => child,
|
|
Action::AddToBulk(rw_lock) => {
|
|
rw_lock.write().await.push(dep.clone());
|
|
return Ok(kill_all_tasks);
|
|
}
|
|
};
|
|
if let Err(e) = crate::handle_child::handle_child(
|
|
&job_id,
|
|
&conn,
|
|
// TODO: Return mem_peak
|
|
&mut 0,
|
|
// TODO: Return canceld_by_ref
|
|
&mut None,
|
|
child,
|
|
is_sandboxing_enabled(),
|
|
&worker_name,
|
|
&w_id,
|
|
&installer_executable_name,
|
|
None,
|
|
false,
|
|
&mut None,
|
|
None,
|
|
None,
|
|
)
|
|
.await
|
|
{
|
|
windmill_queue::append_logs(
|
|
&job_id,
|
|
&w_id,
|
|
format!(
|
|
"\n[x] - error while installing {}: {e:?}",
|
|
&dep.display_name
|
|
),
|
|
&conn,
|
|
)
|
|
.await;
|
|
bail!(format!(
|
|
"\n error while installing dependencies: {e:?}\n{}",
|
|
&dep.display_name
|
|
));
|
|
} else {
|
|
if let Some(ref cb) = post_install {
|
|
cb(&dep)?;
|
|
}
|
|
mark_success(dep.path.clone(), &job_id, &w_id).await;
|
|
print_success(
|
|
false,
|
|
true,
|
|
&job_id,
|
|
&w_id,
|
|
&dep.display_name,
|
|
name_ml,
|
|
counter_arc,
|
|
total_to_install,
|
|
start,
|
|
&conn,
|
|
)
|
|
.await;
|
|
|
|
#[cfg(all(feature = "enterprise", feature = "parquet"))]
|
|
{
|
|
if let Some(os) = windmill_object_store::get_cache_object_store().await {
|
|
let language_name = _language_name.to_string();
|
|
let platform_agnostic = _platform_agnostic;
|
|
let path = dep.path.clone();
|
|
let handle = dep._s3_handle.clone();
|
|
tokio::spawn(async move {
|
|
if let Err(e) = crate::global_cache::build_tar_and_push(
|
|
os,
|
|
path,
|
|
language_name,
|
|
Some(handle),
|
|
platform_agnostic,
|
|
)
|
|
.await
|
|
{
|
|
tracing::warn!("failed to build tar and push: {e:?}");
|
|
}
|
|
});
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(kill_all_tasks)
|
|
}
|
|
|
|
// Create a file to indicate that installation was successfull
|
|
async fn mark_success(path: String, job_id: &Uuid, w_id: &str) {
|
|
let valid_path = path + ".valid.windmill";
|
|
// This is atomic operation, meaning, that it either completes and dependency is valid,
|
|
// or it does not and dependency 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 jobs to function",
|
|
valid_path
|
|
);
|
|
};
|
|
}
|
|
// 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 windmill_object_store::get_cache_object_store()
|
|
.await
|
|
.is_none()
|
|
{
|
|
(s3_pull, s3_push) = (false, false);
|
|
}
|
|
|
|
let mut counter = counter_arc.lock().await;
|
|
*counter += 1;
|
|
|
|
windmill_queue::append_logs(
|
|
job_id,
|
|
w_id,
|
|
format!(
|
|
"\n{}+ {}{}{}| in {}ms",
|
|
windmill_common::worker::pad_string(&format!("[{}/{total_to_install}]", counter), 9),
|
|
// Because we want to align to max len [999/999] we take 9
|
|
// 123456789
|
|
windmill_common::worker::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
|
|
}
|
|
|
|
async fn process_handles(
|
|
(handles, closer): (Vec<JoinHandle<anyhow::Result<TaskKiller>>>, Sender<()>),
|
|
w_id: &str,
|
|
) -> anyhow::Result<()> {
|
|
let mut killers = vec![];
|
|
for handle in handles {
|
|
match handle
|
|
.await
|
|
.unwrap_or(Err(anyhow::anyhow!("Problem by joining handle")))
|
|
{
|
|
// __________< Cannot drop just yet.
|
|
Ok(kill_tasks) => killers.push(kill_tasks),
|
|
Err(e) => {
|
|
tracing::warn!(
|
|
workspace_id = %w_id,
|
|
"Env installation failed: {:?}",
|
|
e
|
|
);
|
|
// We can safely return,
|
|
// it will drop `closer`
|
|
// which will drop first kill_tx for all installation threads
|
|
// which will trigger chain reaction and kill every installation
|
|
return Err(e);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Only now we can be sure there is no installation thread running
|
|
// we can safely stop poller
|
|
drop(closer);
|
|
// Drop ONLY after we processed all handles
|
|
drop(killers);
|
|
Ok(())
|
|
}
|
|
async fn finish_installation(
|
|
total_time: std::time::Instant,
|
|
job_id: &Uuid,
|
|
w_id: &str,
|
|
conn: &Connection,
|
|
) {
|
|
let total_time = total_time.elapsed().as_millis();
|
|
windmill_queue::append_logs(
|
|
&job_id,
|
|
w_id,
|
|
format!(
|
|
"\nDone. Time spent on installation phase: {}ms\n",
|
|
total_time
|
|
),
|
|
conn,
|
|
)
|
|
.await;
|
|
}
|