Add a uv flag to make it possible to automatically replace `pip` with `uv pip` and prepend run command with `uv run` (#765)

This uv flag makes it possible to run dora within a `uv venv` envirement
instead of a venv.

This has the advantage, to:
- not depend on a shell to run command within the virtual env.
- be faster than normal venv
- be easier to maintain a clear separation of environment.

## Example


https://github.com/user-attachments/assets/fbb58afb-76df-4447-85db-87806367846b

## Usage

```bash
uv venv
dora build dataflow.yml --uv
dora run dataflow.yml --uv
```

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"]

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::<Vec<_>>();
         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

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 => {

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),

binaries/daemon/src/lib.rs

@@ -133,7 +133,7 @@ impl Daemon {
         .map(|_| ())
     }
 
-    pub async fn run_dataflow(dataflow_path: &Path) -> eyre::Result<DataflowResult> {
+    pub async fn run_dataflow(dataflow_path: &Path, uv: bool) -> eyre::Result<DataflowResult> {
         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<ResolvedNode>,
         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}`"))

binaries/daemon/src/spawn.rs

@@ -47,6 +47,7 @@ pub async fn spawn_node(
     dataflow_descriptor: Descriptor,
     clock: Arc<HLC>,
     node_stderr_most_recent: Arc<ArrayQueue<String>>,
+    uv: bool,
 ) -> eyre::Result<RunningNode> {
     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)
+                            }
                         }
                     };
 

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<PathBuf> {
     // 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 <path>
+        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())
     }

libraries/message/src/coordinator_to_daemon.rs

@@ -54,4 +54,5 @@ pub struct SpawnDataflowNodes {
     pub nodes: Vec<ResolvedNode>,
     pub machine_listen_ports: BTreeMap<String, SocketAddr>,
     pub dataflow_descriptor: Descriptor,
+    pub uv: bool,
 }