Holistic Optimization: Implementing "Pipeline-as-a-Trial" HPO with Ray and Cloud Infra

Speaker 1 [00:03]

Now we'll start the talk, like now we start the talk. Just for base small admin thing, if anyone has a question, please let's write the question.

Speaker 2 [00:15]

and in the end we can answer all the questions now let's give a big round of applause to Abdullah and for his talk thank you very much so I want to start this talk with this scene from Rick and Morty but just before starting can I see a show of hands who are familiar with the show or not okay yeah so just for you does know it's not mandatory to know the show. So on the left here we have Morty. Morty is in love with Jessica. Jessica doesn't care about Morty. So Morty goes to his uncle Rick, who is this genius scientist who can solve everything with science. And so Rick comes up with this love potion. He gives it to Morty. Morty goes to a party, makes Jessica drink it, and it works. Jessica is in love with Morty, but the fallback here is that Jessica had the flu, and then it spread across the party and now all the party and the whole city is just crazy about Morty Morty goes back to his uncle and then Rick come up with a new solution and they spread across the city it doesn't work and now all the humans just turns into monsters and the whole city is just chaotic and chaos and they cannot reverse the changes and so they just have to ditch the universe they're living in and go to another parallel universe and start living there and so Morty is traumatized yeah it reminds me of our last Friday production hotfix so I want you to take a couple of seconds just to look at the look at Morty's face and see if it's familiar to you like have you seen it anywhere so basically it's it's it's this look I mean I saw this look many many times working in tech it's it's look of a data scientist when they're working on their model for weeks, they put it into an A-B test, it fails the A-B test. It's the look of a product manager when their feature or improvement just causes the revenue to fall. And it's the same problem that Rick and Morty have fallen into. It's this problem of when you're trying to solve something, but without considering the whole system, right? So let's take a step back here and see what happens so Rick did come up with a solution which is this love potion but he only considered the input at hand he couldn't consider the city the flu the whole system and normally in real life like you would have let's consider this this diagram on the left let's say this is a working system it can be anything we want to we notice that C is not is not doing very well so we want to optimize C and let's consider that C is a machine learning model just for the sake of stock. And so we got the data scientist. The data scientist takes the C, do some machine learning magic, come up with a Cv2. It might be improved C, like an improved model or another model. We take C, we put it in the system. Now E is green, so everything is good. But we broke D. Why? Because D also takes input from F, which we didn't know about. And now G is also broken. And this problem, once you start to see it, you see it everywhere. Like the problem when you're trying to optimize is the problem of local optimization. You optimize something, when you put it in a bigger system, it fails. It's everywhere. It's in backend, it's in product, it's in machine learning, and it's even in your real life. Like you're doing some recipe, you're following some recipe, you taste the stock, it needs some salt, you add salt. It's good. you need to add soy sauce. And now everything is just too salty. Sorry for the long introduction. My name is Abdullah Taha. I work as an MLOps engineer at Zalando. And today I'm going to talk to you about this exact problem of how to optimise holistically, how to optimise a pipeline. And we're going to focus today on the problem of HBO, of hyperparameter optimisation, how we at Zalando implemented a scalable solution with Ray and cloud infrastructure. So just to give you a little bit of content, I'm going to start with some context and what we do at my team at Zalando, and how are we dealing with hundreds of machine learning models in production, just to give you a background of where are we starting from. we're going to narrow down a little bit on the problem of how to optimize a pipeline. Then I'm going to present you with the solution design that we did, how to implement, like we tried, a couple of POCs, a couple of implementations using Ray and cloud infrastructure. And then I'm going to show you some comparisons and some insights that you can take home. Okay, let's start. So I work at a forecasting team at Zalando, and our team is responsible for predicting couple of business KPIs, like number of sold items, cancellation rate, etc, etc. Predicting those targets helped Zalando to plan some logistics, to plan budget, etc, etc. And we do this for 24 countries roughly. And like each of those boxes is not a single machine learning model, but it looks something like this. So we might have like model one is a short horizon model to predict like five weeks of data we have like model two that would predict like 15 weeks of data along horizon model then we might have like an assembler that take the input from the models and do some assembling techniques or so then we would have like a post processor and the kpi consolidator that take the input from another kpi and try to make result consistent so yeah it's kind of pipeline and then we get the final target so if you try to visualize this for each of those kpis multiply by 24 you would understand that we're dealing with kind of hundreds of models so how are we doing this in in production we are following a config based approach kind of so you see here this blue library is like we have a library that holds every machine learning code so it's training evaluation splitting post processing whatever and and this library is is used by configs so normally a data scientist just try different configs to experiment with different machine learning models and then whenever he get the best model or the best config he would just give the config to the mlops which is me i put it into production and i take my salary um even that like i could i give it to cloud now but yeah uh so um and and and for the dependency here you see that we have some kind of dag like here we have dependencies like it's a it's a kind of diastolic graph. For the dependency management of the DAG building, we use Airflow. We also follow a config-based approach there. If you watch the talk by Akif, he talked that this is a good design. So, we also have, like, we define the dependency there in kind of config, and then we have a dynamically generated DAG in Airflow. So, this part on the the right is our production but we're going to focus today about experimentation part so normally when it scientists want to experiment with machine learning model he would use this part this process called hyper parameter tuning or hyper parameter optimization i'm sure most of you are familiar with the topic uh but just an overview wow this resolution is bad okay uh so you start with ml model or you start the machine learning model you define a set of parameters that you want to optimize and then the search space is created so this search space is basically like combination of different parameters that you want to optimize against and then the hyper parameter tuning happens i will talk a little bit more on that and then you get the best configs of your model and the hyper parameter tuning looks something like this so there's a couple of approaches to this but yeah just that's an example we're talking here about the bayesian optimization So you would have an optimizer, the optimizer creates the search space, and then the optimizer decides on an initial set of hyperparameters, and then you would trigger a couple of trials. On each trial, you would train the model, evaluate the model, and then the result would be reported back to the optimizer. After a couple of iterations, you get your best results. So, looping back on our problem. if you recall this picture from before so basically in the normal circumstances you would optimize you know model one for the output model one which we don't want what we really want is that we want to optimize model one for the output of the pipeline and not only that we want to optimize each step of the pipeline like we want to optimize the parameters at the ensemble the parameters at the post processor etc etc and so if you if you'd like this diagram that I showed before about the hyper parameter optimization what we really want here for each trial that we try we don't want to try training only the machine learning model what we want to do is this basically this is the core of this presentation this is what we want to implement we want for each trial evaluate the whole pipeline. Okay, so how did we do this? Like a very good, like every good engineering team, we go and consult our senior principal data engineer. No, just joking. So how did we do this? Well, we used Ray, and for those of you who don't know, Ray is a very cool library, a machine learning library that that allows you to run machine learning in a distributed way. So you would write the same code. It would run on either machine or run on like 10 machine cluster, and it's the same code. So we used Ray as the main implementation library. And for the infrastructure, for where the code would be executed in a scalable way, we tried a couple of POCs, one on SageMaker, one on Databricks. And we also tried our own custom solution, running array cluster on EC2s. So from Ray, we used the HyperOpt search model, which is a component under the Ray tune. So Ray has a couple of libraries, one for training, one for data processing, et cetera, et cetera. So this one is under the Ray tune library. And this diagram is from the documentation of Ray. And as we talked before, as a hyperparameter tuner, you have a search space. And here, you trigger the trainable. So the trainable here is the objective function that you want to tune against. So normally, this is a machine learning model. But what we will do is that we will use this customizable component to feed a pipeline to this. And so this is just an example from our code base. So here, when we define the trainables, we take, like, we tried a couple of POCs. So this is just a dictionary of those. you see here when we define tuner we pass the trainable so here like each of those functions are responsible of building a pipeline and I will talk a little bit more on that so just to give you a rough idea this is the config that a data scientist might write so without going a lot to details you can see that we have steps and instead under steps we have three steps basically here the naive seasonal which is a model we define the name of the model the class that it will bind to and then we would define the parameters that you will optimize and like here we have another model which is Auto-ATS and here we have an assembler that take as input the two models so here we have the structure of how the pipeline which is a very simple pipeline here but just to give you an idea now in terms of solution design here we need two components both please clean solution one which is the entry point so here where the data scientist would interact with this whole system so the data scientist would write those configs and then what you need to put them somewhere and so this entry point is where data scientists will just define the configs define the class that we wrote and then trigger whatever method is it's working and the what is triggered basically so the infrastructure should be able to build this concept of a DAG or this concept of pipeline, it should run multiple trials in a scalable way, and it should report back to the entry point. So this is just an overview of the solution. So we tried a couple of POCs. As I mentioned before, I will talk a little bit on details on those. We tried on SageMaker, on Databricks, and on EC2 machines. There are two worth mentioning things that we didn't try but if I have time I will talk a bit more. So what are the considerations when we build such POCs? There are a couple. So there are performance, so each infrastructure here has its own boot time, you know running performance, maintainability, how hard does it, do you need engineers to maintain your environment or is it self maintainable like you can imagine that building custom thing would require more maintainability. Usability, the data scientists are able to use this easily. Traceability, where can you see logs, where can you see failed step, failed pipeline, cost, limits of the infra, etc. Okay let's start. So we start with the SageMaker POC. So in SageMaker we use this concept of SageMaker pipelines. If you're not familiar with it, so you can define multiple training jobs and you combine them together and it's called SageMaker Pipeline. So as our entry point, we used SageMaker Studio notebook, so the data scientist would log in to SageMaker Studio, open a Jupyter notebook, define the configs, trigger the system, define the class and run the cell, and then in the background it would use the config, use SageMaker pipelines to build a DAG and run each trial on a single SageMaker pipeline, basically. So you would run this in a scalable way. And this is some code. So here, as you see, for step-in-step configs, we're iterating over each step, and in each step we're doing defined job training. So we're defining a SageMaker training job, so this is one step of it. So we're defining multiple jobs, and then in the second for loop, we're looping over dependencies between the steps so we can bind them together so we can so when you use this basically you create this DAG and then so yeah and then you return the pipeline and this ray would use to do to evaluate against and this is a screenshot from SageMaker as you can see there are two steps and they are bound into this insular step it has nice UI you can see the running steps you can see the queued steps and and if here in SageMaker if you want to see logs they are all go to CloudWatch if you are familiar with the AWS ecosystem tech system and if you want to also look for each instance CPU and memory usage you would also go to CloudWatch so another similar one which is Databricks Databricks also have this concept of pipeline which called workflows. Again, same system here for the entry point. I'm going to skip that part because it's not very important. It has a similar, just to keep the timekeeping, it has similar concepts and it has like this nice visual UI. And in Databricks I think it's very nice that they have this air flow kind of look so you can see a couple of like running pipelines you can you can see the logs directly from there you can restart each step like if you have failed step look at local we could try every try which isn't there in SageMaker and so yeah now the third POC that we tried one thing to note about the first POC is that here we didn't use Ray cluster. We basically used Ray library and we did the scalability on ourselves. Like we tricked Ray to triggering each SageMaker pipeline or each database workflow so we provide scalability in this way. But in this one we use basically Ray cluster. So as an entry point we use again SageMaker notebook but it can be your local, it can be anywhere. Here we spinned up a Ray cluster so we said okay I want a pipe machine ec2 cluster on spin it up for me and after that we're reading the context and submitting the job to Dre and then Ray would use those machines to run this in parallel and if you recall we needed something to define this tag and here we used metaphor because we need if you're familiar with meta flow you can build Python classes and you can build attack from the class basically so yeah this is a bit more complex but it's yeah it's without using a platform we needed something to measure against to just worth mentioning things I'm not going to spend a lot of time so Databricks also have this concept of rate cluster inside Databricks so you can run one notebook run array cluster behind those notebook it's it's we didn't try this but basically like one downside here is that you cannot see the pipelines running behind the notebook because you have one notebook and a lot of things running in the background. Another thing is this platform which is called AnyScale. So this is a platform that's built by the people who built Ray. First they built the library, they said okay we built the library, now we built the infra. So they basically request to connect to your AWS system and they would handle like like spinning up the machine, terminating the machine, et cetera, et cetera. Yeah. And they give you this nice notebook kind of thing, and you can run things, and it will just scale by itself. So some comparisons. So in terms of spin time and warm pools, we'll start with that. So each of SageMaker and Databricks have this concept of warm pools that you can, you know, don't need to wait for machines to boot up. On ray on EC2, this can be configured, we can just trace, pin up ray cluster beforehand and so we can configure that. In terms of traceability, monitoring, debugging, so where to see the logs, how to trace the pipelines, I think Databricks have the nicer UI in terms of this. On Rayon EC2, it's none. It's basically a terminal, but it can be configured, again. So just a quick insight, like here, Rayon EC2 is the most configurable one, but it requires more maintainability. You need an engineering team to maintain this. The other one is also better. You pay more for a platform, but you get less maintainability. Okay, final thoughts, so what's the winner POC between those? None, obviously. I mean, we ended up using SageMaker just because our data scientists were just using SageMaker from day to day. And it really depends on your requirement. If you're required to more configurable thing, you want to publish logs to some other platform or so, you can build your own stuff. If you don't have engineers, like one engineer and like 50 data scientists, you probably don't want to build something custom. And again, like the second insight is your ML infra is only as good as the ability of data science to navigate it. You can build the most amazing system, engineered system, if the data scientist cannot use it. It's useless. The third one is go look around. I mean, take the local optimization problem in your system, try to figure out the solution. Don't be a rake. I mean the insight here is that we talked about this concept in the context of hyper parameter optimization, but this problem is everywhere. It's it's it's it's in on all products in back-end engineering. You're trying. Yeah, just go around try looking for it The fourth insight. Well, yeah, I need to fill out the template and I didn't have the time so Okay Want to finish with This quote from a great author in system design and the performance of a system is not the sum of performance of its part It is the product of their interactions Thank you very much. And if you're interested in the talk, if you like the talk, if you want to work on these similar problems, please join us at Zalando. You can look at our open jobs at jobsteadzalando.com. If you want to reach out to me through LinkedIn, you can use the link on the right.

Speaker 1 [21:58]

Thank you so much First question is when testing the POCs Did you use enterprise level or were the test using the open source levels of different POCs?

Speaker 2 [22:11]

I mean, in Zalando, we all have both Databricks and SageMaker are using enterprise level.

Speaker 1 [22:29]

Any more questions anyone you have anything

Speaker 2 [22:37]

given him a big round of applause.