Implement alias syntax for runtime nodes with single python operator

apis/rust/node/src/config.rs

@@ -6,7 +6,7 @@ use std::{
     str::FromStr,
 };
 
-#[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(deny_unknown_fields)]
 pub struct NodeRunConfig {
     #[serde(default)]

binaries/coordinator/examples/mini-dataflow.yml

@@ -54,5 +54,14 @@ nodes:
         python: ../runtime/examples/python-operator/op.py
         inputs:
           time: timer/time
+          test: python-operator/counter
         outputs:
           - counter
+
+  - id: python-operator
+    python:
+      path: ../runtime/examples/python-operator/op.py
+      inputs:
+        time: timer/time
+      outputs:
+        - counter

binaries/coordinator/src/main.rs

@@ -1,4 +1,4 @@
-use dora_core::descriptor::{self, Descriptor, NodeKind};
+use dora_core::descriptor::{self, CoreNodeKind, Descriptor};
 use dora_node_api::config::NodeId;
 use eyre::{bail, eyre, WrapErr};
 use futures::{stream::FuturesUnordered, StreamExt};
@@ -49,17 +49,21 @@ async fn main() -> eyre::Result<()> {
 }
 
 async fn run_dataflow(dataflow_path: PathBuf, runtime: &Path) -> eyre::Result<()> {
-    let Descriptor {
-        mut communication,
-        nodes,
-    } = read_descriptor(&dataflow_path).await.wrap_err_with(|| {
+    let descriptor = read_descriptor(&dataflow_path).await.wrap_err_with(|| {
         format!(
             "failed to read dataflow descriptor at {}",
             dataflow_path.display()
         )
     })?;
 
-    if nodes.iter().any(|n| matches!(n.kind, NodeKind::Runtime(_))) && !runtime.is_file() {
+    let nodes = descriptor.resolve_aliases();
+    let mut communication = descriptor.communication;
+
+    if nodes
+        .iter()
+        .any(|n| matches!(n.kind, CoreNodeKind::Runtime(_)))
+        && !runtime.is_file()
+    {
         bail!(
             "There is no runtime at {}, or it is not a file",
             runtime.display()
@@ -74,12 +78,12 @@ async fn run_dataflow(dataflow_path: PathBuf, runtime: &Path) -> eyre::Result<()
         let node_id = node.id.clone();
 
         match node.kind {
-            descriptor::NodeKind::Custom(node) => {
+            descriptor::CoreNodeKind::Custom(node) => {
                 let result = spawn_custom_node(node_id.clone(), &node, &communication)
                     .wrap_err_with(|| format!("failed to spawn custom node {node_id}"))?;
                 tasks.push(result);
             }
-            descriptor::NodeKind::Runtime(node) => {
+            descriptor::CoreNodeKind::Runtime(node) => {
                 if !node.operators.is_empty() {
                     let result =
                         spawn_runtime_node(runtime, node_id.clone(), &node, &communication)

libraries/core/src/descriptor/mod.rs

@@ -2,6 +2,7 @@ use dora_node_api::config::{
     CommunicationConfig, DataId, InputMapping, NodeId, NodeRunConfig, OperatorId,
 };
 use serde::{Deserialize, Serialize};
+use std::collections::HashSet;
 use std::fmt;
 use std::{
     collections::{BTreeMap, BTreeSet},
@@ -18,14 +19,69 @@ pub struct Descriptor {
 }
 
 impl Descriptor {
+    pub fn resolve_aliases(&self) -> Vec<ResolvedNode> {
+        const PYTHON_OP_NAME: &str = "op";
+
+        let python_nodes: HashSet<_> = self
+            .nodes
+            .iter()
+            .filter(|n| matches!(n.kind, NodeKind::Python(_)))
+            .map(|n| &n.id)
+            .collect();
+
+        let mut resolved = vec![];
+        for mut node in self.nodes.clone() {
+            // adjust input mappings
+            let input_mappings: Vec<_> = match &mut node.kind {
+                NodeKind::Runtime(node) => node
+                    .operators
+                    .iter_mut()
+                    .flat_map(|op| op.inputs.values_mut())
+                    .collect(),
+                NodeKind::Custom(node) => node.run_config.inputs.values_mut().collect(),
+                NodeKind::Python(node) => node.inputs.values_mut().collect(),
+            };
+            for mapping in input_mappings {
+                if python_nodes.contains(&mapping.source) {
+                    assert_eq!(mapping.operator, None);
+                    mapping.operator = Some(OperatorId::from(PYTHON_OP_NAME.to_string()));
+                }
+            }
+
+            // resolve nodes
+            let kind = match node.kind {
+                NodeKind::Custom(node) => CoreNodeKind::Custom(node),
+                NodeKind::Runtime(node) => CoreNodeKind::Runtime(node),
+                NodeKind::Python(node) => CoreNodeKind::Runtime(RuntimeNode {
+                    operators: vec![OperatorConfig {
+                        id: OperatorId::from(PYTHON_OP_NAME.to_string()),
+                        name: None,
+                        description: None,
+                        inputs: node.inputs,
+                        outputs: node.outputs,
+                        source: OperatorSource::Python(node.path),
+                    }],
+                }),
+            };
+            resolved.push(ResolvedNode {
+                id: node.id,
+                name: node.name,
+                description: node.description,
+                kind,
+            });
+        }
+        resolved
+    }
+
     pub fn visualize_as_mermaid(&self) -> eyre::Result<String> {
-        let flowchart = visualize::visualize_nodes(&self.nodes);
+        let resolved = self.resolve_aliases();
+        let flowchart = visualize::visualize_nodes(&resolved);
 
         Ok(flowchart)
     }
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct Node {
     pub id: NodeId,
     pub name: Option<String>,
@@ -35,16 +91,36 @@ pub struct Node {
     pub kind: NodeKind,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "lowercase")]
 pub enum NodeKind {
     /// Dora runtime node
     #[serde(rename = "operators")]
     Runtime(RuntimeNode),
     Custom(CustomNode),
+    Python(PythonOperatorConfig),
 }
 
 #[derive(Debug, Serialize, Deserialize)]
+pub struct ResolvedNode {
+    pub id: NodeId,
+    pub name: Option<String>,
+    pub description: Option<String>,
+
+    #[serde(flatten)]
+    pub kind: CoreNodeKind,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum CoreNodeKind {
+    /// Dora runtime node
+    #[serde(rename = "operators")]
+    Runtime(RuntimeNode),
+    Custom(CustomNode),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(transparent)]
 pub struct RuntimeNode {
     pub operators: Vec<OperatorConfig>,
@@ -73,7 +149,16 @@ pub enum OperatorSource {
     Wasm(PathBuf),
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct PythonOperatorConfig {
+    pub path: PathBuf,
+    #[serde(default)]
+    pub inputs: BTreeMap<DataId, InputMapping>,
+    #[serde(default)]
+    pub outputs: BTreeSet<DataId>,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct CustomNode {
     pub run: String,
     pub env: Option<BTreeMap<String, EnvValue>>,
@@ -83,7 +168,7 @@ pub struct CustomNode {
     pub run_config: NodeRunConfig,
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(untagged)]
 pub enum EnvValue {
     Bool(bool),

libraries/core/src/descriptor/visualize.rs

@@ -1,11 +1,11 @@
-use super::{CustomNode, Node, NodeKind, OperatorConfig, RuntimeNode};
+use super::{CoreNodeKind, CustomNode, OperatorConfig, ResolvedNode, RuntimeNode};
 use dora_node_api::config::{DataId, InputMapping, NodeId};
 use std::{
     collections::{BTreeMap, HashMap},
     fmt::Write as _,
 };
 
-pub fn visualize_nodes(nodes: &[Node]) -> String {
+pub fn visualize_nodes(nodes: &[ResolvedNode]) -> String {
     let mut flowchart = "flowchart TB\n".to_owned();
     let mut all_nodes = HashMap::new();
 
@@ -21,11 +21,11 @@ pub fn visualize_nodes(nodes: &[Node]) -> String {
     flowchart
 }
 
-fn visualize_node(node: &Node, flowchart: &mut String) {
+fn visualize_node(node: &ResolvedNode, flowchart: &mut String) {
     let node_id = &node.id;
     match &node.kind {
-        NodeKind::Custom(node) => visualize_custom_node(node_id, node, flowchart),
-        NodeKind::Runtime(RuntimeNode { operators }) => {
+        CoreNodeKind::Custom(node) => visualize_custom_node(node_id, node, flowchart),
+        CoreNodeKind::Runtime(RuntimeNode { operators }) => {
             visualize_runtime_node(node_id, operators, flowchart)
         }
     }
@@ -63,16 +63,20 @@ fn visualize_runtime_node(node_id: &NodeId, operators: &[OperatorConfig], flowch
     flowchart.push_str("end\n");
 }
 
-fn visualize_node_inputs(node: &Node, flowchart: &mut String, nodes: &HashMap<&NodeId, &Node>) {
+fn visualize_node_inputs(
+    node: &ResolvedNode,
+    flowchart: &mut String,
+    nodes: &HashMap<&NodeId, &ResolvedNode>,
+) {
     let node_id = &node.id;
     match &node.kind {
-        NodeKind::Custom(node) => visualize_inputs(
+        CoreNodeKind::Custom(node) => visualize_inputs(
             &node_id.to_string(),
             &node.run_config.inputs,
             flowchart,
             nodes,
         ),
-        NodeKind::Runtime(RuntimeNode { operators }) => {
+        CoreNodeKind::Runtime(RuntimeNode { operators }) => {
             for operator in operators {
                 visualize_inputs(
                     &format!("{node_id}/{}", operator.id),
@@ -89,7 +93,7 @@ fn visualize_inputs(
     target: &str,
     inputs: &BTreeMap<DataId, InputMapping>,
     flowchart: &mut String,
-    nodes: &HashMap<&NodeId, &Node>,
+    nodes: &HashMap<&NodeId, &ResolvedNode>,
 ) {
     for (input_id, mapping) in inputs {
         let InputMapping {
@@ -101,7 +105,7 @@ fn visualize_inputs(
         let mut source_found = false;
         if let Some(source_node) = nodes.get(source) {
             match (&source_node.kind, operator) {
-                (NodeKind::Custom(custom_node), None) => {
+                (CoreNodeKind::Custom(custom_node), None) => {
                     if custom_node.run_config.outputs.contains(output) {
                         let data = if output == input_id {
                             format!("{output}")
@@ -112,7 +116,7 @@ fn visualize_inputs(
                         source_found = true;
                     }
                 }
-                (NodeKind::Runtime(RuntimeNode { operators }), Some(operator_id)) => {
+                (CoreNodeKind::Runtime(RuntimeNode { operators }), Some(operator_id)) => {
                     if let Some(operator) = operators.iter().find(|o| &o.id == operator_id) {
                         if operator.outputs.contains(output) {
                             let data = if output == input_id {
@@ -126,7 +130,7 @@ fn visualize_inputs(
                         }
                     }
                 }
-                (NodeKind::Custom(_), Some(_)) | (NodeKind::Runtime(_), None) => {}
+                (CoreNodeKind::Custom(_), Some(_)) | (CoreNodeKind::Runtime(_), None) => {}
             }
         }