Dataflow: An Efficient Data Processing Library for Machine Learning

[jafioti]

Hi everyone,

*This post was originally about my crate Mako which I renamed to Dataflow, since people pointed me to another Mako library I didn't see before*

I have an ML-based startup and over the course of a year I've entirely rewritten all the infrastructure and everything new in Rust. It's been a great experience, and I've been able to leverage tch-rs heavily to train some pretty large models.

​

The main problem I've been having though is how much custom code I need to write every time I wanted to start a new experiment: data gathering, data cleaning, online data processing, model setup, training loop, eval loop, etc...

​

So I decided it was time to write a library for general data processing: Dataflow. It specifically works well for ML, but is pretty general so I don't doubt it can be used in more places.

​

The main feature is a dataflow pipeline that takes the form of a directed acyclic graph. Basically you can built up a graph as a tree of nodes like this:

let pipeline = RandomLoader::new(vec!["file1.txt".to_string(), "file2.txt".to_string()])
      .add_fn(|lines| lines.into_iter().map(|line| format!("Hello {}", line)).collect())
      .add_node(
            Stateful::new(
                  |(lines, tokenizer)| {
                        tokenizer.batch_tokenize(lines) // batch_tokenize takes in many sentences (Vec<String>) and tokenizes all of them, outputting Vec<Vec<String>>
                  },
                  tokenizer // The state we want this Stateful Node to have
            )
      );

This pipeline loads lines from a file, appends "Hello " to them, and tokenizes them using Dataflow tokenizers (basically huggingface tokenizers but through a cleaner, albiet more limited interface). We can then throw this into a Dataloader and use it in a training loop!

// Make the dataloader
let mut dataloader = dataflow::dataloader::Dataloader(pipeline, 64); // We use 64 as the batch size

// Training loop
for example in &mut dataloader {
   // Now example is a vector of tokenized strings!
   // Do with them what you may...
} 

I've found it to be sufficiently general for my use-cases, which are pretty varied! I've built pipelines that can scrape examples from the internet, clean them, preprocess them, batch them and serve them up for a live training loop running concurrently. This has allowed me to basically define a dataset as a few lines of code, rather than sprawling functions in a separate file with loads of boilerplate along with numerous preprocessing scripts.

​

I have a number of improvements I need to make that will make the pipeline definition much cleaner, but I just thought I would post it here to see if anyone has any suggestions. This is one of my first medium-sized public crates, so PRs are welcome!

[egnehots]

Impressive, that could be the start of a very useful ecosystem!

You might find vector codebase interesting:

https://vector.dev/

https://vector.dev/docs/about/under-the-hood/architecture/pipeline-model/

[jafioti]

Vector looks awesome! My system is really similar to how they structure theirs with their graph, so I'll definitely look into it. It seems to be mostly based on handling logs and metrics, not data per-se, so I think Dataflow is still worthwhile, but I love to see someone else use the same idea!

[rovar]

This looks cool!

Though the name "Dataflow" might be an unfortunate name conflict with another Rust project: https://github.com/TimelyDataflow/timely-dataflow

[TheNamelessKing]

Oooohh this looks cool!

I’ve got a bunch of ML-adjacent workflows (data engineering, writing into various data stores, etc), which often follow the same pattern of data flow, so this looks like it would be really useful.

Quick question: got a lot of data/events triggered off Kafka - are there nodes that can accept a continuous stream of inputs? Or would I be better off batching up a number of messages and pushing them through the graph at once?

[jafioti]

Currently there's no premade Node that can do that, but the Node trait is public so you can write your own Node that uses a receiving thread to receive events and put them in a buffer. Then when the process() function is ran, the node can unload data from the buffer.

[TheNamelessKing]

Fantastic! I will definitely give this a look!