From 28af216c8f11cfd590e32121456378ceb1217b45 Mon Sep 17 00:00:00 2001 From: haixuantao Date: Wed, 29 Jan 2025 15:18:04 +0100 Subject: [PATCH] Add uv flag to enable build and run through uv run and uv pip --- binaries/cli/pyproject.toml | 2 +- binaries/cli/src/build.rs | 37 +++++++++++-------- binaries/cli/src/lib.rs | 16 +++++--- binaries/coordinator/src/run/mod.rs | 1 + binaries/daemon/src/lib.rs | 8 +++- binaries/daemon/src/spawn.rs | 36 ++++++++++++++++-- libraries/core/src/descriptor/mod.rs | 12 ++++++ .../message/src/coordinator_to_daemon.rs | 1 + 8 files changed, 85 insertions(+), 28 deletions(-) diff --git a/binaries/cli/pyproject.toml b/binaries/cli/pyproject.toml index f8bb2914..b8a74928 100644 --- a/binaries/cli/pyproject.toml +++ b/binaries/cli/pyproject.toml @@ -8,7 +8,7 @@ dynamic = ["version"] scripts = { "dora" = "dora_cli:py_main" } license = { text = "MIT" } requires-python = ">=3.8" -dependencies = ['dora-rs'] +dependencies = ['dora-rs', 'uv'] [tool.maturin] features = ["python", "pyo3/extension-module"] diff --git a/binaries/cli/src/build.rs b/binaries/cli/src/build.rs index ee88bc65..9c18f74b 100644 --- a/binaries/cli/src/build.rs +++ b/binaries/cli/src/build.rs @@ -7,7 +7,7 @@ use std::{path::Path, process::Command}; use crate::resolve_dataflow; -pub fn build(dataflow: String) -> eyre::Result<()> { +pub fn build(dataflow: String, uv: bool) -> eyre::Result<()> { let dataflow = resolve_dataflow(dataflow).context("could not resolve dataflow")?; let descriptor = Descriptor::blocking_read(&dataflow)?; let dataflow_absolute = if dataflow.is_relative() { @@ -22,29 +22,28 @@ pub fn build(dataflow: String) -> eyre::Result<()> { for node in descriptor.nodes { match node.kind()? { dora_core::descriptor::NodeKind::Standard(_) => { - run_build_command(node.build.as_deref(), working_dir).with_context(|| { + run_build_command(node.build.as_deref(), working_dir, uv).with_context(|| { format!("build command failed for standard node `{}`", node.id) })? } dora_core::descriptor::NodeKind::Runtime(runtime_node) => { for operator in &runtime_node.operators { - run_build_command(operator.config.build.as_deref(), working_dir).with_context( - || { + run_build_command(operator.config.build.as_deref(), working_dir, uv) + .with_context(|| { format!( "build command failed for operator `{}/{}`", node.id, operator.id ) - }, - )?; + })?; } } dora_core::descriptor::NodeKind::Custom(custom_node) => { - run_build_command(custom_node.build.as_deref(), working_dir).with_context(|| { - format!("build command failed for custom node `{}`", node.id) - })? + run_build_command(custom_node.build.as_deref(), working_dir, uv).with_context( + || format!("build command failed for custom node `{}`", node.id), + )? } dora_core::descriptor::NodeKind::Operator(operator) => { - run_build_command(operator.config.build.as_deref(), working_dir).with_context( + run_build_command(operator.config.build.as_deref(), working_dir, uv).with_context( || { format!( "build command failed for operator `{}/{}`", @@ -60,16 +59,22 @@ pub fn build(dataflow: String) -> eyre::Result<()> { Ok(()) } -fn run_build_command(build: Option<&str>, working_dir: &Path) -> eyre::Result<()> { +fn run_build_command(build: Option<&str>, working_dir: &Path, uv: bool) -> eyre::Result<()> { if let Some(build) = build { let lines = build.lines().collect::>(); for build_line in lines { let mut split = build_line.split_whitespace(); - let mut cmd = Command::new( - split - .next() - .ok_or_else(|| eyre!("build command is empty"))?, - ); + + let program = split + .next() + .ok_or_else(|| eyre!("build command is empty"))?; + let mut cmd = if uv && (program == "pip" || program == "pip3") { + let mut cmd = Command::new("uv"); + cmd.arg("pip"); + cmd + } else { + Command::new(program) + }; cmd.args(split); cmd.current_dir(working_dir); let exit_status = cmd diff --git a/binaries/cli/src/lib.rs b/binaries/cli/src/lib.rs index a59e49cd..1854ce4a 100644 --- a/binaries/cli/src/lib.rs +++ b/binaries/cli/src/lib.rs @@ -83,6 +83,9 @@ enum Command { /// Path to the dataflow descriptor file #[clap(value_name = "PATH")] dataflow: String, + // Use UV to build nodes. + #[clap(long, action)] + uv: bool, }, /// Generate a new project or node. Choose the language between Rust, Python, C or C++. New { @@ -99,6 +102,9 @@ enum Command { /// Path to the dataflow descriptor file #[clap(value_name = "PATH")] dataflow: String, + // Use UV to run nodes. + #[clap(long, action)] + uv: bool, }, /// Spawn coordinator and daemon in local mode (with default config) Up { @@ -345,20 +351,20 @@ fn run(args: Args) -> eyre::Result<()> { } => { graph::create(dataflow, mermaid, open)?; } - Command::Build { dataflow } => { - build::build(dataflow)?; + Command::Build { dataflow, uv } => { + build::build(dataflow, uv)?; } Command::New { args, internal_create_with_path_dependencies, } => template::create(args, internal_create_with_path_dependencies)?, - Command::Run { dataflow } => { + Command::Run { dataflow, uv } => { let dataflow_path = resolve_dataflow(dataflow).context("could not resolve dataflow")?; let rt = Builder::new_multi_thread() .enable_all() .build() .context("tokio runtime failed")?; - let result = rt.block_on(Daemon::run_dataflow(&dataflow_path))?; + let result = rt.block_on(Daemon::run_dataflow(&dataflow_path, uv))?; handle_dataflow_result(result, None)? } Command::Up { config } => { @@ -519,7 +525,7 @@ fn run(args: Args) -> eyre::Result<()> { ); } - let result = Daemon::run_dataflow(&dataflow_path).await?; + let result = Daemon::run_dataflow(&dataflow_path, false).await?; handle_dataflow_result(result, None) } None => { diff --git a/binaries/coordinator/src/run/mod.rs b/binaries/coordinator/src/run/mod.rs index 1ba6b74d..d4ac6270 100644 --- a/binaries/coordinator/src/run/mod.rs +++ b/binaries/coordinator/src/run/mod.rs @@ -55,6 +55,7 @@ pub(super) async fn spawn_dataflow( nodes: nodes.clone(), machine_listen_ports, dataflow_descriptor: dataflow, + uv: false, }; let message = serde_json::to_vec(&Timestamped { inner: DaemonCoordinatorEvent::Spawn(spawn_command), diff --git a/binaries/daemon/src/lib.rs b/binaries/daemon/src/lib.rs index e016f0c6..682d1c5a 100644 --- a/binaries/daemon/src/lib.rs +++ b/binaries/daemon/src/lib.rs @@ -133,7 +133,7 @@ impl Daemon { .map(|_| ()) } - pub async fn run_dataflow(dataflow_path: &Path) -> eyre::Result { + pub async fn run_dataflow(dataflow_path: &Path, uv: bool) -> eyre::Result { let working_dir = dataflow_path .canonicalize() .context("failed to canonicalize dataflow path")? @@ -152,6 +152,7 @@ impl Daemon { nodes, machine_listen_ports: BTreeMap::new(), dataflow_descriptor: descriptor, + uv, }; let clock = Arc::new(HLC::default()); @@ -384,6 +385,7 @@ impl Daemon { nodes, machine_listen_ports, dataflow_descriptor, + uv, }) => { match dataflow_descriptor.communication.remote { dora_core::config::RemoteCommunicationConfig::Tcp => {} @@ -409,7 +411,7 @@ impl Daemon { }; let result = self - .spawn_dataflow(dataflow_id, working_dir, nodes, dataflow_descriptor) + .spawn_dataflow(dataflow_id, working_dir, nodes, dataflow_descriptor, uv) .await; if let Err(err) = &result { tracing::error!("{err:?}"); @@ -625,6 +627,7 @@ impl Daemon { working_dir: PathBuf, nodes: Vec, dataflow_descriptor: Descriptor, + uv: bool, ) -> eyre::Result<()> { let dataflow = RunningDataflow::new(dataflow_id, self.machine_id.clone()); let dataflow = match self.running.entry(dataflow_id) { @@ -696,6 +699,7 @@ impl Daemon { dataflow_descriptor.clone(), self.clock.clone(), node_stderr_most_recent, + uv, ) .await .wrap_err_with(|| format!("failed to spawn node `{node_id}`")) diff --git a/binaries/daemon/src/spawn.rs b/binaries/daemon/src/spawn.rs index a75259a1..887b82b1 100644 --- a/binaries/daemon/src/spawn.rs +++ b/binaries/daemon/src/spawn.rs @@ -47,6 +47,7 @@ pub async fn spawn_node( dataflow_descriptor: Descriptor, clock: Arc, node_stderr_most_recent: Arc>, + uv: bool, ) -> eyre::Result { let node_id = node.id.clone(); tracing::debug!("Spawning node `{dataflow_id}/{node_id}`"); @@ -113,15 +114,42 @@ pub async fn spawn_node( // If extension is .py, use python to run the script let mut cmd = match resolved_path.extension().map(|ext| ext.to_str()) { Some(Some("py")) => { - let python = get_python_path().context("Could not get python path")?; - tracing::info!("spawning: {:?} {}", &python, resolved_path.display()); - let mut cmd = tokio::process::Command::new(&python); + let mut cmd = if uv { + let mut cmd = tokio::process::Command::new("uv"); + cmd.arg("run"); + cmd.arg("python"); + tracing::info!( + "spawning: uv run python -u {}", + resolved_path.display() + ); + cmd + } else { + let python = get_python_path().wrap_err( + "Could not find python path when spawning custom node", + )?; + tracing::info!( + "spawning: {:?} -u {}", + &python, + resolved_path.display() + ); + let cmd = tokio::process::Command::new(python); + cmd + }; + // Force python to always flush stdout/stderr buffer + cmd.arg("-u"); cmd.arg(&resolved_path); cmd } _ => { tracing::info!("spawning: {}", resolved_path.display()); - tokio::process::Command::new(&resolved_path) + if uv { + let mut cmd = tokio::process::Command::new(&"uv"); + cmd.arg("run"); + cmd.arg(&resolved_path); + cmd + } else { + tokio::process::Command::new(&resolved_path) + } } }; diff --git a/libraries/core/src/descriptor/mod.rs b/libraries/core/src/descriptor/mod.rs index 021734d7..07b74fc6 100644 --- a/libraries/core/src/descriptor/mod.rs +++ b/libraries/core/src/descriptor/mod.rs @@ -7,7 +7,9 @@ use std::{ collections::{BTreeMap, HashMap}, env::consts::EXE_EXTENSION, path::{Path, PathBuf}, + process::Stdio, }; +use tokio::process::Command; // reexport for compatibility pub use dora_message::descriptor::{ @@ -210,6 +212,16 @@ pub fn resolve_path(source: &str, working_dir: &Path) -> Result { // Search path within $PATH } else if let Ok(abs_path) = which::which(&path) { Ok(abs_path) + } else if which::which("uv").is_ok() { + // spawn: uv run which + let _output = Command::new("uv") + .arg("run") + .arg("which") + .arg(&path) + .stdout(Stdio::null()) + .spawn() + .context("Could not find binary within uv")?; + Ok(path) } else { bail!("Could not find source path {}", path.display()) } diff --git a/libraries/message/src/coordinator_to_daemon.rs b/libraries/message/src/coordinator_to_daemon.rs index 463a857e..da98af36 100644 --- a/libraries/message/src/coordinator_to_daemon.rs +++ b/libraries/message/src/coordinator_to_daemon.rs @@ -54,4 +54,5 @@ pub struct SpawnDataflowNodes { pub nodes: Vec, pub machine_listen_ports: BTreeMap, pub dataflow_descriptor: Descriptor, + pub uv: bool, }