Scaling Python: An End-to-End ML Pipeline for ISS Anomaly Detection with Kubeflow

Speaker 1 [00:07]

Good morning, everybody. I'm really glad so many of you showed up here. And I'm Henrik. This is my colleague Christian. And we will, as already introduced, tell you something about how we scale Python using Kubeflow for our use case, which is anomaly detection for ISS data. But before we jump into the content, let me introduce myself real quick. You know my name already. I would consider myself as a data scientist, even though I now slowly turned into a rather an MLOps engineer maybe and I'm doing research in ML for cyber physical systems and my use case is DISS obviously and I used to be a consultant for five years or so I have a background in engineering and during my research I really got to know Kubeflow a lot and I'd say I'm a research Kubeflow enthusiast and with that Christian what What about you?

Speaker 2 [01:05]

Thank you. So my name is Christian.

Speaker 1 [01:10]

Back one slide.

Speaker 2 [01:10]

One slide. I work with data. I used to be a physicist and consider myself a long-time Unix nerd. Also, I love Python. So according to this t-shirt I found at home, my first PyCon was in 2011, which was really a long time ago. And nowadays I'm building a Kubernetes-based MLOps platform, which is what this use case is running on. But us two here are not the only authors of the work that was done here, but all those people together wrote a paper last year where we showcased more or less also what we are showing here. Today more with a focus on the Python part of it, in that paper Python isn't that much in focus. But the work is not only done by those seven people, but by the whole KISS project team. And so let's quickly talk a bit about what is KISS. It's Künstliche Intelligenz, so artificial intelligence for the International Space Station. And there we are doing anomaly detection, diagnostics and reconfiguration for the European part of the International Space Station, which is a so-called Columbus module. Unfortunately, I can't tell you where that is on that image there, but it's somewhere on there probably. And you can see also there are quite a lot of project partners involved in that. And before we really start, I'll quickly tell you what are we going to talk about today. So we're going to give you a very high-level overview of our use case, but nothing very specific to this use case in general. And then we're going to tell you which parts of Kubeflow and also the wider open source ecosystem we're going to use for which parts, how it all fits together. And then we're going to share some challenges we had and then the solutions and also share some tips and tricks. Unfortunately, there are also things we are not going to talk about today. And so if you came here because the ISS thing is like really appealing to you, we're very sorry to disappoint you. and now is really your time to leave. We won't really judge you. We're also not going to talk a lot about the scientific methods in detail, but as I said, it's more going to be Python-focused. If you're interested in those scientific methods, at the end of the talk we have some references. Check those out and especially also the references in the references and probably some more levels of recursion. Also, what we're not going to talk about, but which seems to be the main topic of the conference, is large language models, model context protocol or any parts of agents. So yes, I saw the jump. Please, this is also your chance to leave if you want to learn something about that. In general also, we're not going to show the exact code we use for the ISS and also not the data. We are not allowed to show the data and the code part, which also already for that paper, we just changed a bit to make it easier to understand. it doesn't really change anything like in general but it's just like an easier starting point and you also can download that code from github there's also a link at the end and with that let's quickly talk about the use case a bit so as mentioned before it's anomaly detection on time series data so we really are trying to figure out where does something go wrong on the ISS so is it like in an anomalic situation and then if we found some situation like that we're going to do root cause analysis on it so for example we find out it's getting warmer and warmer for example than it should be what is the root cause of that might be like a fan that's not spinning correctly and after we've done that we're trying to find out how can we reconfigure the ISS so that it all gets back to normal which in the case of the fan might be for example spinning another fan faster. And finally there's also some supervision part that makes sure that we don't bring the ISS or this Columbus module into a dangerous situation. There's like a dashboard on it which the people who operate the Columbus module use. This is also not like the real dashboard which we are also not allowed to show, but it's just basically some simple web dashboard and then if they think like yeah this makes sense they take those values that were suggested by the reconfiguration and transfer them to another system and then reconfigure everything. The first part here, the anomaly detection, uses a neural network model, so really like proper machine learning, while everything else uses old-school symbolic AI so we are also not going to talk in much detail about that today because there's like no training involved or anything. It uses some components mainly model serving but nothing else so today we are mostly gonna talk about that and just for your information everything else is cobbled together using Kafka. So challenges. We want to do some distributed pre- processing and we have really massive amount of data so we are not trying to do that only on one machine but in many machines and as we love Python we want to do that in Python and don't want to use Hadoop or something like that. Then we want to do hyper parameter search directly in the beginning not to optimize the model in the end but to find optimal starting parameters to start training the model. Then we want to do model training while it's rather easy nowadays to use like multi GPUs on one machine what we want to do is use multi GPUs on multiple machines and that also is not as easy as doing it on just one machine. Then we want to deploy the model so also run it in production also not just like one version but multiple just for resiliency and also because we might need to serve really lots and lots of requests and finally we want to tie all of that together in one workflow so first do the distributed pre-processing then hyperparameter search, train the model and even bring it into at least semi-production or like a test deployment like completely with one workflow. But before we go into what tools we use for that let's very quickly talk about Kubernetes perhaps quick show of hands who of you knows what Kubernetes is? That seems to be everyone we can skip this slide and who of you is actually using Kubernetes already? Well that's also quite a lot of people. So we skipped that slide. So we mostly use Kubeflow which is a set of open source tools which are some parts are developed under that Kubeflow name some are also contributed or like just taken from other open source projects, but what all of them have together is they're like running on Kubernetes, Kubernetes native, so most things then run in a pod and it makes distributing things a lot easier. What isn't really covered by Kubeflow, at least wasn't to our satisfaction, are like some things that we got covered by MLflow and Dask, which we also integrated. And this is then also the point where we're going to talk about more in detail how we We solved each step and Henrik is going to take over here.

Speaker 1 [08:52]

Thank you. All right, we know the challenges we faced, and now I'd like to talk in detail about each single one. Let's start with distributed preprocessing. And Christian already showed on the last slide we use Dask to actually distribute our preprocessing. Quick show of hands again. Who uses or knows Dask? Well, that's about half of you. Good. For the others, it's a Python-based distributed computation framework. And it's not actually part officially of the Kubeflow ecosystem, but good news is it works with Kubernetes in the back end, which means it's relatively easy to integrate it in a Kubeflow-based platform that we used. Now, once an admin set it up, you actually only need Python code to use it. Let me show you what we did. So, finally, we have Python code displayed here. This is a little part of our pre-processing code, and you see here this class says Dask preprocessor, and it has two Dask-specific imports here. The first is kubecluster, which is imported from Dask Kubernetes, so it's Kubernetes-specific, and the other is the Dask client, which is not. It's imported from Dask distributed. And somewhere inside our class, we have a method called createDaskCluster, and it does That's exactly what it promises. First, we create a cluster with a self.cluster, which is an object of type cube cluster. And we give it a name, that our cluster has a nice name. We also provide an image. That's the container image that actually runs inside the worker pods of the Dask cluster. Not to confuse with the Kubernetes cluster. So it's a couple of pods that run our Dask cluster on top of Kubernetes. And we also specify the namespace just to make sure that we have the according rights to create the cluster inside that. In the next line, we say cluster.scale to a specific number of workers. In our case, we picked three because our basically our cluster had three nodes, but it can also easily be 100. And then once the cluster is there, we actually use this self.client to connect to the cluster we just created. And once that line is executed, everything we do in Dask actually is computed on the cluster. So every Dask.read Parquet file or Dask compute to get something will be executed on the worker nodes that have been created in the lines on top of that. When we go ahead and create a cluster, the SDK will actually talk to the Dask operator, which will have to be installed by the cluster admin. But if we say we now want a cluster, that operator goes ahead and creates a scheduler pod, which is this one here. And that's in charge of creating this fancy computational graph that Dask is known for. So that scheduler pod decides which part of the processing will be done by which worker, and he will then send the task down to all those workers, which, as I said, might be plenty and this thing also serves the dash desk UI which I really love to use because it really helped me a lot in debugging things and in our use case we managed to create this like textbooks MapReduce computational graph so what you see here is all the colored dots are process tasks and from left to right the computation goes and the further we go to the right the higher the level of aggregation and we can also see nicely here that this is truly out of memory even for the whole cluster because as soon as a process task is completed it releases and and communicated to the next one it releases its memory so even if data is larger than all the memory we have in the whole cluster we can still compute our metrics and this is basically what we do for pre-processing the large data chunks we had for the time series so So great, we now have a data set that we can actually use for training. As Christian already mentioned, the next thing to do, or the next challenge, was find a good starting point in terms of hyperparameters. And for that, we use a Kubeflow-specific tool, or at least that's integrated into the Kubeflow ecosystem, which is called Katib. Who knows or uses Katib? One of you. Great. So it's new to all of us. to do to abstract away the complexity of hyperparameter tuning while using Kubernetes in the backend. And let's see how well that goes. So the way we use it, you need three things. First, this might be a big surprise, you need a training script. Second, we need a containerized version of that because we are on Kubernetes or a Docker file. And third, we need a description of the hyperparameter experiment. So let's have a look at each of those. First, training script. So no surprises here. We have our function somewhere in the main.py file called runTraining, and we decorate it using click decorators, such that we can use the CLI to call the function. So number of GPUs might be interesting. RunAsPy.drop will be interesting in a minute. But here for hyperparameter tuning, especially this learning rate is interesting. So we have to specify every hyperparameter as a CLI argument here. More emphasis on that. We use PyTorch Lightning. So this PL is imported from PyTorch Lightning to initialize the trainer. And we use MLflow to lock the validation loss and so on, and also the model. And then finally we say model.trainer.fit. trainer.fit. I assume most of you know how this stuff works. Next, we need the image, also pretty standard here. We have a two-stage build just to make sure that we don't very often install the packages again, which might take 15 minutes or so because Torch is pretty large. But we use this new and fancy package manager called UV, which so far we really love. I can only recommend it. I also loved poetry at the beginning, which I changed my mind later, so it might happen for you too. So first we install the dependencies and then second we copy over our source code files and then also install all the packages we wrote ourselves in order to have them available in the image, which will be important later. Now the third step is to actually describe the hyperparameter tuning experiment. And for that we need a YAML file. YAML is very common in Kubernetes. You will see it more often in this talk. But most of you know Kubernetes, so you are used to these kind of files. And here we need to describe, for example, how many trials. And the trial is a single training we want to run in parallel. We picked three here, but it might be also 10 or 100 again. How many maximum trials we want to run, and so on. And don't worry, if you are not the YAML type of guy, you could also use the UI to do that. will give you the initial version for the YAML file, if you're more the clicky type of person, which I assume you're not, because this is nicely trackable in Git, for example. And the next step is to specify what exactly gets optimized. And in this case, we are minimizing the validation loss, and we're using a very fancy algorithm called random. But Katib also has things like Bayesian optimization and the most state-of-the-art algorithm you you might expect of such a tool. You will find, again, the learning rate here again. And we give it a couple of options. And there are more parameters. I only didn't want to list all of them here. And finally, we describe what exactly happens in a single trial. And in this case, it's a Kubernetes job, which means we have our image again, which is the image we built before. And that holds the training script We saw even before that. We tell Kubernetes that the trial is allowed to use a GPU. And then we have the command. And that's what we constructed using the CLI decorators on top of our function. So here we say, Python, run the main.py file. Run this run training command. Use one GPU. We also say, run as py.drop equals false. We'll later say true. But here we don't need that. And here you see the learning rate. And most importantly, that's managed by CarTip itself. So the value will be set by the hyperparameter tuning tool here. And accordingly, it will do it for all the other parameters. Again, more emphasis on the learning rate. So what happens when we write something like kubectl apply this file, or in the UI we say launch or submit? We will actually talk to the experiment controller. He will notice, oh, there's a new experiment, and spin up a trial controller. The trial controller runs, in our case, a random algorithm or the Bayesian optimization and comes up with an initial set of hyperparameters and tells each worker, in this case three again, which hyperparameters to use. And those three workers will then run the training in parallel. Once they are done, they will report back their metrics, validation loss, or whatever we are optimizing. And the trial controller will think of the next best option for the next set of hyperparameters this worker should be using. As I noted, the Kubernetes job is only the most simple version of a CUTTIP trial. You could also go for a PyTorch job or even a whole Argo workflow. In that case, it really gets complicated because we are parallelizing already using CUTTIP trials. But the Kubernetes job itself could be using lots of pods. So if you really want, you can scale massively here. This is how it looks like in the UI if a couple of trials succeeded. So you might know this spot from other hyperparameter tuning tools. Every single line is a trial. And you see the combination of hyperparameters and how they affected the metric we are interested in. You could also then drill deeper into single trials and see how the validation loss evolves. Things like that. So great. This worked. And now we have a reasonable set of initial hyperparameters. The next thing we'd like to do is train the model really seriously on multiple GPUs, and our GPUs are distributed on multiple machines. So let's say we have three nodes, three GPUs each. Now we want to train the model on nine GPUs at the same time in parallel. And for that, we use a tool called Kubeflow Trainer. Those who used Kubeflow, did you also use the trainer or PyTorch job? One or two? So, this might be new to you. To run a PyTorch job, we're going to need three things, a training script, surprise, the image, also no surprise because we are on Kubernetes, and the PyTorch job YAML again. Good news is we already have the training script, and we can use exactly that training script to run the PyTorch job. Also, for the image, we don't have to change anything. We only need to take care of the YAML file, so let's have a look at that. It looks similar, but this time it's not an experiment, but a PyTorch job. And we have to specify the master node. We just need one of those. And again, we just run the image and the command we already know. It's the image we built, and the command looks very similar, only now we specify all the values and nobody is managing our values here. And importantly, we say run as PyTorch job equals true. Now we'll see that for the worker, which we also specify in the same file, in this case we have two worker nodes, we do exactly the same thing, same image, same command. And since we say run as PyTorch job equals true, the PyTorch job in Kubeflow will actually take care of managing environment variables such that the container itself will know which rank it has. So the master has rank zero, and then the workers have one and two. And in the background, this will happen. So we submit the file we just saw. The training operator will see it, spin up worker nodes, and in case of a PyTorch job, we have PyTorch just using data distributed parallel. So it will split up the training data set, send different batches to different workers. So they will all do their forward and backward paths and share the gradients in a very efficient way in terms of communication with each other, and the master will then eventually only one of them will only update the weights and share them back to the workers for the next batch. And the algorithm they use is ring or reduce. And importantly, only the master node then is in charge of logging results. I messed that up first, or we messed it up first, and had suddenly multiple experiments in our MSO server. But it's easily taken care of in the script. So in this case here, we see the logs of one PyTorch job, and luckily, both train and validation loss both decreased. So all good here. So great. We have a model. It's trained. The metrics are fine. What do we do next? Obviously, we want to deploy it. And for me personally, this used to be the most challenging part, because I'm rather a data science guy than a hardcore IT nerd. And Kubeflow promises to make this easier and more accessible to all the data scientists. Let's have a look at KSERV. Who knows KSERV? Again same people. Yeah. Great. So to use KSERV, we have the trained model, which we now discussed how to get. And again, we need the YAML file. In this case, the KSERV YAML. In the case of the model we have, let's have a look at the YAML file. We saw experiment and PyTorch job. Now this is an inference service. And we can give it a name and point it to the model. Now where do we get that string from? Well, we have MLflow, and that can manage this artifact for us, and the UI, it looks like this. So we can just hit copy and paste the value here. And also, again, we can tell Kubernetes that also for serving we can use a GPU because it's much quicker. Now if we submit this file, it will take minutes, and then we will get a fully fledged API, a REST API that does the serving for us. So here we can send requests, too, including metrics and logs and everything that you might expect from a served model. So this is really handy. Here you see, for example, a couple of metrics that get logs, so number of requests per second, the rate, and how many memory and CPU it uses and things that you would expect. All only given the small YAML file. So this is an end. One of the most important features is that it scales up and down as you wish. So we can define whenever the response rate drops, Because I want another pod that serves the model in order to really get the user certified. Great. So basically we covered everything. The only thing we didn't is how do we now orchestrate everything together. And I think the tool that Kubeflow is most famous for is its pipelines engines. So just because we can, we decided let's automate this whole process in a holistic pipeline That first does the task stuff, then uses Catip to train the hyperparameters and all the things until we finally deploy our model in an automated pipeline. And this is how it looks like. I'm aware you can't read the steps, but you see already it's a kind of directed graph. Let me first scroll down and then zoom in. You see here, for example, that we have this run Cartib experiment task in the graph Also run PyTorch job and evaluate model, things like that And every box with the green icon is a task in Kubeflow pipelines Which means it's basically a container image And a command and inputs and outputs And the yellow ones are artifacts So Kubeflow pipeline takes care of storing the artifacts Somewhere where we can track them later Even though we need an image and so on, the framework makes it easy to create those tasks. You need a component that's kind of the template for a task, and here's an example. I see really quick, so you define a function. You can give it a base image if you like, but it can also just be Python, so you don't have to specify this. You can install packages before your actual code gets executed. And this is an example of how I.O. works. So the function has an input called testdf, which is of type input dataset, and that's Kubeflow pipeline-specific. And later in this code, I do something with this object, and I use the .path attribute, which actually holds the string to the location in the container where the file lives that I want to read. And it was injected by Kubeflow pipelines. I don't have to take care of where it lives and so on. And we return a dictionary. And if we have a look at the UI again here, in the UI I just clicked on the evaluate task, and I can see the output dictionary here in the UI pretty nicely and also other results, artifacts, and where they live in MinIO or S3. So that's it. That's all the things we wanted to integrate in a holistic framework. And which brings me to the end of our talk. If you want to remember anything from what we've told you, we'd like you to remember first that we are still big fans of PowerPoint's morph feature. And the open source community has a lot of tools that can help you in the ML lifecycle. You don't have to go to expensive solutions here. Also, specifically, Kubeflow Pipelines lets you integrate This lets you integrate everything in a holistic end-to-end framework, and data scientists like myself can actually use tools like Kubeflow to massively scale ML workloads. And please check out our code repository, what we showed you today. We updated it. It's in the branch PyCon. We will probably merge it later. We wrote a paper about this, as Christian already said, and we gave a talk specifically about how this integration and pipeline works for all these fancy tools on last year's Kubeflow Summit. And you will probably find one of us at the ProKube booth downstairs, or at least someone that knows where we are. And we love to discuss details. Thanks for your attention.

Speaker 2 [28:41]

so we have probably time for just one question from Slido and so that question is is distributing the pre-processing on three nodes really faster than using a larger single node and using something like DuckDB

Speaker 1 [29:00]

I didn't run benchmarks, so I can't answer. We didn't have such a large node available, but we had pretty powerful nodes in our cluster, so it wasn't really an option to use a larger one for us. Sorry, I can't really answer. Probably not. Distributing is always complicated and adds complexity.

Speaker 2 [29:19]

Awesome. If you did happen to ask questions on Slido.com, you guys can continue that conversation on Discord and continue liking the questions that are posted. So until then, let's give our speaker another round of applause.