Rust AI & Data Meetup: Everett Kleven & Srinivas Lade, Daft running Python coroutines in Tokio
Um I'm Everett. >> I'm Srinu. >> Uh we work at Eventual, uh which stewards the open-source project called Daft. It's a distributed data engine, uh commonly known for its data frame API. Um that's our repo. And >> Pick up. >> that's our repo. >> Cool
>> Uh basically, um we're the kings of IO. Uh there's a lot of really great data engines out there. Um there's some really excellent technologies that have, you know, come and talked tonight. Um we uh participate in that ecosystem as well. And uh but one of the big things that we're known for is uh multimodal. Um we've got our this is our main repository. We've got a bunch of examples. Um lately we've been on a really big push with uh physically ID assets and uh supporting different formats there
Um and we have a bunch of different projects that have built on top of Daft. Uh and so we've got a broad ecosystem that's uh building on top of us. Um what we're going to talk about today is we're going to dive into um a talk that was given at the Tokyo Conf um earlier this year. And uh we're just going to talk about the the the Rust layer that really underlies Daft, um even though you're maybe interacting with Python a lot more when you're uh using Daft. So we already did our our intros. Already showed you the repo. Um showed you the connectors. Um so yeah, I already mentioned this, but you know, DuckDB, DataFusion, Spark, they they there are a lot of really fast um excellent technologies out there
Um, and on several analytical queries and specific workloads, we're very competitive. Um, but one of the things that we're particularly good at um is multimodal. And that uh comes from the fact that we have a lot of technology built around increasing our bandwidth at the IO layer um between storage uh and compute. And that really bottlenecks the amount of uh throughput that you can really push uh particularly on top of S3. So, when it comes to running models on multimodal data, that's it's it's pretty hard to compete with us. Um, so like, you know, why not Pandas? Uh you've got a lot of traditional data science built on top of the um maybe initial data frame API. A lot of us know that data frame API and and that's maybe the first way that we got started with data analysis. Um, you know, the the rows are the primitives there and you can just really quickly get started
Uh all of that's in memory and you're processing each uh row uh serially with each operation that you compose. Daft is lazily evaluated and so we uh by doing that, we have a lot like a specific materialization boundary uh when we either collect or write to storage. And that gives us an opportunity to do a lot of optimizations and planning under the hood. Um, so uh Pandas is just uh built for a different job. It also only runs on a single machine. We can work on uh clusters as well as your own laptop, which um and you know, Pandas is synchronous. We support uh async and sync workloads. Um Why not PySpark? Uh you know, it scales, but the execution model uh certainly fights uh our multimodal workloads, and particularly running like something like a stateful user-defined function where you have, let's say, a Transformers model that you want to load once for a worker and then parallelize its inference calls
Um that helps you amortize that model loading cost. So, you know, it just it doesn't have the async UDFs, the generators that we have. And I think the the biggest thing that I love about Daft is like how expressive the API is. Um and uh it's still being very actively developed. And there are there um It just gives you a lot of uh flexibility. So, uh let's just like dive into like a common example of uh what makes Daft powerful, um captioning video frames. So, if you're familiar with how uh like an MP4 file is actually structured, how you read it, um you have to uh you have all of these different frames, and uh those keyframes end up needing to be decoded um because the compression that MP4 comes with uh will, you know, explode the memory once you actually read it from disk into your compute. Handling all of that memory is non-trivial
The same thing goes for your images once you go from JPEG into your bitmap, and you actually want to do inference on that. Uh you need to be able to handle the memory uh and the the row throughput there. Uh so, let's just quickly walk through a simple example of like what that looks like. We got some videos on an object store. You've got some frames that you're going to do decode. Uh maybe you're going to go resize it because you want to uh run that through a VLA um on and run a benchmark on like Libero or something. And then uh you want to annotate that. Uh that annotation can come in a variety of forms
Maybe you're using like a vision language uh uh model to be able to uh review the video, identify different uh features, uh and then curate data sets for downstream training. Um so then once you have that this kind of we're going to help you Daft is really great at that like curation uh step. And generally when we're talking about Python, this is what uh you're going to be interacting with. Um which hides a lot of the complexity um and optimizations that we're running in Rust under the hood. All right. Yeah, so Daft has first-class IO for images, audio, video, uh recently adding HDF5 uh and MCAP. Um we can work locally and distributed um as well as a lot more expressive user-defined functions that are keeping things Pythonic. Uh but at the end of the day um we're still getting all of the performance benefits of Rust
And we're going to kind of dive into like how that actually happens. Uh you know, let's switch to you. >> Okay. excuse me if I am a little bit awkward about going through these slides. I remember vibe coding or vibe designing these at like 2:00 a.m. 3 months ago and got pulled into this like 2 hours ago. So, let's see how well this goes. So, what's going on under the hood? When we're given a query, a construction of a uh piece of like code that we want to run
In this case kind of representing that vision example we had before. Mhm. Is this better? Yes. Okay. On the left-hand side you have kind of like what the user code looks like. Maybe you define like okay, you're going to read a set of videos. You're going to transform the individual frames of those videos, perform some sort of captioning process, and then writing out. All right
>> Or maybe use the admin You're good. No, thank you. This is another 2 seconds. That's when you are magically on. Yeah, okay. Yeah. [clears throat] >> Much better. >> Thanks
Cool. Thanks, guys. So, yeah. Behind the scenes, we take that piece of code and transform it into what we call a plan, a logical that represents the step of actions that occur in some sort of like graph flow or graph format, a DAG that we like to call. We do with some set of optimizations on top of that DAG based on things that we understand that we can optimize at this level. And at some point, we transform that into a set of steps that we want to execute it. And all the execution happens kind of within our engine itself. That's all we're interested in
Uh since this was a talk from a Tokyo talk, we are focusing on how we use Tokyo as a lot of the operations, as you'll see in the next slides, end up being very IO-heavy rather than more traditional computation. And actually to get to that point, if you had a situation where maybe you were reading from local disk and you were writing to local disk, then a lot of the operations that we talked about before would be mostly computation. And all we really would need is some sort of thread pool, some sort of way to just utilize all of your cores. To execute, you know, your set of uh um tasks and kind of see them as quickly as possible. Pretty straightforward. But when you're working at large scale, petabytes of data or something like that, then you're starting to deal with a lot of IO, a lot of transfer of data from like an external source. And just waiting on that in a traditional system that's just distributing all of the medical would end up leading to a lot of stalls. And so you'll see in this case for example that we're waiting a lot for just like getting the video blocks on this three
Similarly, if you're doing captioning through an API like you're going to open AI's SDK and generate captions from GPT or maybe you're hosting a model locally that's on GPU that you're communicating with in order to generate those captions. In a way that's also wasting time because you're communicating with some external service and GPUs are effectively external services from the perspective of the program. To do this do some sort of action that you're waiting for to come back to you. And so you're also not actually making progress on your project. And this is the problem that we saw when using a traditional data engine that was meant for analytics like Spark. A lot of the time would be wasted on this kind of stuff and you wouldn't get to utilize the resources on your machine to 100%. And so that's why we ended up building this custom back end in our application instead of adopting an existing one. The main way we did that is by dividing the work between a computation pool and an IO pool
The IO pool could focus on doing a lot of asynchronous operations using Tokio which is like pretty freaking amazing at this kind of stuff while the computation can be done on a separate login runtime. And so you can see in this example where overlapping with a lot of these like three-step operations like the read write the hashing and then the resizing even across the you know these were the capital And then you can actually kind of like see in this diagram you can see how the process is being made across each subset of each batch of data. And so yeah, how does Tokio kind of help with this? I mean, said provides a way to treat a lot of async situations in a very coherent manner to really make our stuff make sense. And a lot of the work that we have scheduling, which is whether it's like the query or the compute or resizing HP operations to be distributed, we can treat that effectively like a async function and generalize that across our entire backing system. So, Torque was useful because it's work stealing. It can already handle stuff like multithreading and and this work stealing and um uh I mean it can back in so if something takes too long, it can kind of be more reasonable if necessary and while other things are picked up by other threads. A lot of the overhead that's kind of involved in this process is going to kind of be optimized by the team to make a lot of these things happen at the second or third level. And so, we can kind of be focused on like the data operations themselves and not worry about scheduling
And lastly, despite the structure, a lot of the blocking operations because of these isolated to a CPU core end up being taken up by CPU threads while a lot of the IO can have the time it needs to make progress. And yeah, this is just kind of like the representation of the code that we talked about before. Because it's all written in this style style API, we have an understanding of okay, read operations end up flowing towards IO. A lot of the expressions and computations end up flowing to the CPU. A lot of like user defined functions that find a specific thing can also be IO. And in specific cases where we see that okay, we're using a specific library or tool, we can also treat that as IO because it's being interacted with the CPU. Uh let's take a look at this stuff because it's more like Okay. And yeah, so it ends up being that by kind of taking this approach where we're treating code as these same mechanical existences, it allows us to not have to implement a lot of these features that that are useful in multimodal AI workloads
Instead relying on external libraries, which is what really common in the scene of programming languages. And so that's things like if you're going to be doing video coding, you can use libav/ffmpeg. If you're doing image operations, you can use Pillow, very popular image IO operation. PyTorch, they can do shot and one sort of output or a lot of like model handling as well as associated things like Transformers to go along and so on. >> Yeah, pretty much correct. >> Yeah. >> Okay. >> Yeah
>> Um thank you for coming. Thank you for listening to this convoluted topic. Um if you have any questions, feel free to ask. >> Yeah. >> We'll leave the time for one question because it turns out that the 8:30 end time is a hard stop at AWS. We will not. So, we'll have one more question to ask and then we'll call it a day. Who wants to ask the most interesting question? >> So, I saw you batch the computer and IO out or did the computer you have CPUs and you have GPUs which are inherently different uh um sort of uh the CPU is more like scheduling threads and all versus the GPU you wanted to batch inferences sort of
Does it take into account that part? >> Sorry, I didn't miss uh say it. The GPU is treated as part of the IO rather than the compute because in the GPU it's doing the computation. On the perspective of the program, it's like you're communicating with an external entity. And so, you can treat it like I like you're sending a piece of work to something and then later when you're listening you can see that it's done. >> Cool. Thanks, guys. >> [applause] >> I guess