Browse Source

Adding run comand within python API (#875)

This PR makes it possible to run dataflow from python using:

```python
from dora import run

# Make sure to build it first with the CLI.
# Build step is on the todo list.

run("qwen2.5.yaml", uv=True)
```

Making it easier to run dataflow based on our needs
tags/v0.3.11-rc1
Haixuan Xavier Tao GitHub 10 months ago
parent
commit
066d47007c
No known key found for this signature in database GPG Key ID: B5690EEEBB952194
5 changed files with 57 additions and 2 deletions
  1. +3
    -0
      Cargo.lock
  2. +3
    -0
      apis/python/node/Cargo.toml
  3. +43
    -0
      apis/python/node/src/lib.rs
  4. +2
    -2
      tests/llm/qwen2.5.yaml
  5. +6
    -0
      tests/llm/run_script.py

+ 3
- 0
Cargo.lock View File

@@ -2874,6 +2874,8 @@ name = "dora-node-api-python"
version = "0.3.10"
dependencies = [
"arrow 54.2.1",
"dora-daemon",
"dora-download",
"dora-node-api",
"dora-operator-api-python",
"dora-ros2-bridge-python",
@@ -2884,6 +2886,7 @@ dependencies = [
"pyo3",
"pythonize",
"serde_yaml 0.8.26",
"tokio",
]

[[package]]


+ 3
- 0
apis/python/node/Cargo.toml View File

@@ -22,11 +22,14 @@ eyre = "0.6"
serde_yaml = "0.8.23"
flume = "0.10.14"
dora-runtime = { workspace = true, features = ["tracing", "metrics", "python"] }
dora-daemon = { workspace = true }
dora-download = { workspace = true }
arrow = { workspace = true, features = ["pyarrow"] }
pythonize = { workspace = true }
futures = "0.3.28"
dora-ros2-bridge-python = { workspace = true }
# pyo3_special_method_derive = "0.4.2"
tokio = { version = "1.24.2", features = ["rt"] }

[lib]
name = "dora"


+ 43
- 0
apis/python/node/src/lib.rs View File

@@ -1,10 +1,15 @@
#![allow(clippy::borrow_deref_ref)] // clippy warns about code generated by #[pymethods]

use std::env::current_dir;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use arrow::pyarrow::{FromPyArrow, ToPyArrow};
use dora_daemon::Daemon;
use dora_download::download_file;
use dora_node_api::dora_core::config::NodeId;
use dora_node_api::dora_core::descriptor::source_is_url;
use dora_node_api::merged::{MergeExternalSend, MergedEvent};
use dora_node_api::{DataflowId, DoraNode, EventStream};
use dora_operator_api_python::{pydict_to_metadata, DelayedCleanup, NodeCleanupHandle, PyEvent};
@@ -310,11 +315,49 @@ pub fn start_runtime() -> eyre::Result<()> {
dora_runtime::main().wrap_err("Dora Runtime raised an error.")
}

pub fn resolve_dataflow(dataflow: String) -> eyre::Result<PathBuf> {
let dataflow = if source_is_url(&dataflow) {
// try to download the shared library
let target_path = current_dir().context("Could not access the current dir")?;
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.context("tokio runtime failed")?;
rt.block_on(async { download_file(&dataflow, &target_path).await })
.wrap_err("failed to download dataflow yaml file")?
} else {
PathBuf::from(dataflow)
};
Ok(dataflow)
}

/// Run a Dataflow
///
/// :rtype: None
#[pyfunction]
#[pyo3(signature = (dataflow_path, uv=None))]
pub fn run(dataflow_path: String, uv: Option<bool>) -> eyre::Result<()> {
let dataflow_path = resolve_dataflow(dataflow_path).context("could not resolve dataflow")?;
let rt = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.context("tokio runtime failed")?;
let result = rt.block_on(Daemon::run_dataflow(&dataflow_path, uv.unwrap_or_default()))?;
match result.is_ok() {
true => Ok(()),
false => Err(eyre::eyre!(
"Dataflow failed to run with error: {:?}",
result.node_results
)),
}
}

#[pymodule]
fn dora(_py: Python, m: Bound<'_, PyModule>) -> PyResult<()> {
dora_ros2_bridge_python::create_dora_ros2_bridge_module(&m)?;

m.add_function(wrap_pyfunction!(start_runtime, &m)?)?;
m.add_function(wrap_pyfunction!(run, &m)?)?;
m.add_class::<Node>()?;
m.setattr("__version__", env!("CARGO_PKG_VERSION"))?;
m.setattr("__author__", "Dora-rs Authors")?;


+ 2
- 2
tests/llm/qwen2.5.yaml View File

@@ -8,8 +8,8 @@ nodes:
DATA: "Please only output: This is a test"

- id: dora-qwen2.5
build: pip install -e ../../node-hub/dora-qwen2.5
path: dora-qwen2-5
build: pip install -e ../../node-hub/dora-qwen
path: dora-qwen
inputs:
text: pyarrow-sender/data
outputs:


+ 6
- 0
tests/llm/run_script.py View File

@@ -0,0 +1,6 @@
from dora import run

# Make sure to build it first with the CLI.
# Build step is on the todo list.

run("qwen2.5.yaml", uv=True)

Loading…
Cancel
Save