SBTB23:Oleg Avdeëv& Riley Hun, Lesson learned from orchestrating large-scale GenAI, ML & data on k8s
[Music] uh my name is OLG I'm kaar at arond I'm going to talk today a little bit about the lessons we learn from orchestrating large scale ml workloads including gen uh and data on kubernetes where helping our customers our buns and then Riley from Autodesk will talk about their more practical side of things how they do it there um so previous speaker actually did pretty good job kind of introducing all the concepts so I'm not going to spend too much time on this but um still it's kind of useful to zoom out a little bit and talk about the kind of the full picture and uh if you think about um before you even think about ML and AI uh it's useful to think about what's what's been the difference between normal traditional software uh and the m and AI so with traditional software you usually have those yellowad people Engineers writing code and that code gets put in this box uh and it's some sort of valuable system that produces uh in this case money so produces some value to your customers in the end because we're uh talking about uh still applied ml Andi here we're not uh academic uh conference uh so what changes if you introduce ml not even talking about ji yet um I mean changed there's still this big box there is still a lot of software there's a lot of traditional code that you have to write uh but now you have those couple new things so you have um this data box on the left uh that somehow you using some magic uh this yellow head or purple head guy data scientist ml engineer uh makes an ml model out of it and now that thing gets plugged into the your traditional software still hopefully to somehow produce the value for your customers or you know power your product or some part of this product um and this thing is a little different because in the previous slide you had this very deterministic uh situation where you just write code you can write some unit tests and things like this uh now you have this more fuzzy Cloud that is based on you know probabilistic principles buil built on based on statistics trained on data uh so it's not as predictable or it's predictable maybe not as easily testable but it still has to run correctly it's still has to plug somewhere we just we're not building those models just for the sake of building models uh and it has to work with the rest of your software uh so looking at this slightly differently just putting those boxes in the Triangle um you have um data which is a new part you have models that you train based on this data uh using some compute platform and then again very important to remember again they have to somehow plug into your production systems uh become featur or something in your product uh so you're not they don't just exist in vacuum for you to write a paper about um and you have this uh little data scientist or machine learning engineer uh who has to work on those models and has to be productive um and they usually have a slightly different set of skills than uh your traditional software Engineers uh they know those math and statistics really well they're also usually actually pretty good at software engineering but they're not uh necessarily spending all day you know designing some beautiful architectures of you know message buses and things like this so we want to make them productive and we want them to be able to rapidly prototype and build those models iteratively and deploy them production ideally with like as little help as possible uh so that's um probably like the biggest goal if you're building any kind of ml infrastructure or AI or geni even before you get this um and another way uh would like to put it is there are those two conjoin triangle of success uh of a data scientists on the left uh who care mostly about top levels in this stack so they care about modeling and maybe deploying their models and versioning and then you have those um yellow head um infrastructure Engineers that mostly work on the lower layers uh they build this data and compute platform and build some pieces for orchestrating uh those small workflows so really the data scientists really care about the modeling part and and the infrastructure Parts kind of underneath this uh stack so now closer to geni uh does anything actually change with geni is it even different when you throw llms and gen and foundational models in the mix um and here this is kind of just random rough architecture of you know maybe your data infrastructure M infrastructure if you're starting to get into geni and llms and um I wouldn't say a lot of things change there's emphasis changes a little bit on certain things uh so if you talk about compute specifically um one big difference is that uh you now have those big GPU boxes to train from scratch or fine-tune your llm models uh and this is different probably from the at least philosophy that uh ml systems and Data Systems used to have like five seven years ago uh where you would be talking about you know spark and all those kind of things where you just have those very elastically scalable clusters of compute you just kind of throw a few more CS into the mix and Things become faster now you have um you actually have to have gpus you have to use you know certain gpus types for certain models uh it becomes I mean it doesn't change completely but you have a lot more heterogenous clusters in terms of uh instance types and Hardware types so that's one thing uh another thing that is always uh there for a Mel is that you have a lot of uh problems that uh maybe don't require a giant amount of uh gpus or don't require a giant GPU box with like a aunds or Hy wage 100s uh but you're running uh some sort of what they call embarrassingly parallel workloads uh let's say you want to run inference uh in a batch mode maybe PR compute some recommendations or something embeddings and your computer platform has to support that as well uh which is pretty different so this is now this is more closer to uh spark and things like this where you have those super paralyzable uh tasks um and last but at least these days for AI a lot of times you will be relying on vendors and something like open AI AP or something like this it's not as exciting as you know running giant clusters of gpus on your own own uh but you have to remember that this is also has to be part of your platform it has to play nicely with other things in your uh infrastructure uh stack so and those tasks they don't require as much compute resources but they still have to finish start on time um and the results maybe get plugged into uh your other models and uh thinking about like quickly from orchestration perspective is um few things here uh that are different so one is it's becomes really important to think about failure modes because again those large models are really expensive to train uh they cost a lot of money if something fails uh you need to make sure that your data scientists or ml Engineers uh can understand what went wrong and recover and ideally not to lose too much compute uh budget and another thing is it has to react to new data uh very few people here probably would ever train uh foundational model from scratch uh that means you're taking someone else's and you f tuning this uh and that becomes an interesting problem in versioning and organizing those pipelines because uh even though people were talking about fine-tuning models and maybe some sort of online training for many years I would say if you asked me like five years ago uh it's not really something that 99% % of use cases would need because you could a lot of times just train from scratch just train on daily data uh retrain the model and you would be probably better off it maybe not most computationally uh efficient way to do things but it would be a lot easier from infrastructure perspective now that actually makes sense that you take the model from somewhere you fine tune it uh on your data on some Cadence um and that poses a set of new challenges for uh managing this dat data and results and uh models um so since the talk is titled kubernetes uh I'll be talking mostly about the compute part of this um and maybe getting closer to uh practical learnings that we got from this uh in helping our customers to run those models on our platform uh first thing is again kind of hard to easy to overlook is really the user experience for uh dat scientist or ml engineer uh talking about kubernetes specifically it's awesome for people like me infrastructure Engineers uh I love how like they have this Universal declarative configuration management language and all those things uh lots of yaml uh and it doesn't lend well for use for by data scientists directly for a few reasons so one is I would say there's really philosophical difference between uh um kubernetes that was built at least initially to run more like Services microservices you know your usual web apps uh so it's very natural for someone managing the cluster or configuring the cluster uh to use this awesome uh declarative language where you describe what you you want end state to be and it kind of happens for you so you could say you know I want the service to be running at this port with this host name and things got to happen and you don't really care how it gets there uh when you start thinking about ML and batch workloads it's actually not really the case because it's not really about the end result so end result is all your training jobs are finished so nothing is running so it's really about the kind of natural way to describe it as you all do probably day day today is to describe the steps to get there and it's very imperative it's not really a declarative thing you kind of you is saying hey run the this job take the results you know listen for this event trigger a fine-tuning job when this data arrives and of course the other thing is there's a lot of yamel and it's very technical and very uh lowlevel uh language if you're doing gaml customize or whatever uh so that's where metaflow comes in real quickly because we're company built around metaflow it's an open source project uh Open Source by Netflix a few years back uh and I think the biggest value is really it provides this very natural user interface pythonic for data scientists where they can write their code in Python not caring much about kubernetes and all those things um and metaflow kind of takes care of scheduling all those uh workloads on a kubernetes cluster you can you know version your artifacts you can uh download data really fast from cloud storage you can use parallel training and so on okay there is no video so but it looks like this but it actually runs uh if you if you you can basically write your code run it locally and then with pretty much no changes run out kubernetes in the distributed fashion okay the video is missing uh and here it runs on the cluster uh another learning uh or I guess the cirk once you uh start trying those workloads kubernetes uh machine learning related is the workload diversity is really huge so if you just again like you take your vanilla normal kubernetes cluster uh with some Services running in it uh it's kind of one thing you can have all those pods running come up you know Auto scale all that stuff collect metrics now you start throwing batch workloads it's pretty different um because they require a lot of Auto scaling up and down when you know fan out on 100 tasks then you want to those tasks finish you have some join step that sort of shuts up everything down it's very different requirements for the autoscaler and for scheduling and uh maybe using cluster upgrades for for an as an example uh if you want to upgrade the comp version or you want to replace the underlying nodes if you just run the services as usual you can usually just add a few more replicas shut down some old containers uh start new ones on new notes and kind of off you go with high availability and all that uh when you start having patch workloads they actually usually not really interruptable so you cannot just spin another replica of bad job that trains you know a foundational model for two days or three days or whatever uh so you kind of have to wait for it to finish before you can do anything with the underly node so kind of the good news that they actually they don't run forever every every bad job some at some point stops uh but if you really want to have High availability for your data scientist and it's important because it's so expensive to use those gpus uh your cluster upgrades could easily take you know days or weeks if you want to re really do the rolling uh gradual uh replacement of those notes without interrupting anyone uh and in addition to this there is also something we provide is uh Cloud workstations where uh data scientists uh kind of workspace is not on their laptop but actually in remote node in a cluster which is great because they get access to all those fency gpus directly uh but it also adds a lot of requirements on scheduling uh again like you cannot interrupt those things very easily um next thing again it's not super fancy I would say it's just uh something that you run into almost immediately is uh container image sizes and the managing those sizes because when you start dealing with the m workflows I did a small plot here of uh by torch uh Transformers uh image container image from dockerhub from hiking phase so you can see a few short years it went from two gigs to 8 gigs it's actually not unusual for us to see images in like 10 gigs and even 20 gigs uh range um especially when people start to try to build some sort of universal image for their data scientist they add P torch and tensor flow and iterate on those images uh so you can imagine you start a pretty small uh parallel training job uh using this new image and it's just five nodes now you have to suddenly pull 100 gigs before anything useful happens at all so that becomes a problem both from uh performance and uh actually uh it makes user experience really bad if they have to wait every time for 5 10 minutes to launch uh whatever they're working on uh that kind of hurts productivity a lot and in terms of solutions there's some Solutions in open source like open source project I think it's cncf hardore uh that provides in cluster image caching but um it's it's easy on the surface but it gets really complicated when you get into details so as a kind of a example of that complexity is if you run your cluster on let's say AWS uh in a lot of cases you probably want to run it across availability zones and uh it's not super obvious from the pricing on the Amazon's page but it's actually you pay pretty much as much uh cost for networking across availability zones as much as you pay for across region data transfer and they really charge you through the nose for those and especially if something like a n Gateway is in the middle uh it it gets expensive and you're kind of pulling those gigabytes and tens of gigabytes and terabytes of data just for container images so you have to have your if you do in cluster image caching you have to make sure it's U availability Zone aware um and now to the kind of the font task is uh running multinode mgpu workloads it's also not Supernatural thing that comes out of the box with most orchestration platforms including kubernetes uh because you do want this uh gank scheduling semantics which means uh if you have a highly parallel job you don't want to just schedule 100 tasks you want them to run concurrently and communicate with each other to share parameters or something and uh it's crucial especially now with like those large models that sometimes don't fit on a single node and uh there are a few open source Solutions like q and jobset and volcano I think CU flow has a few operators but it's still early days when you start using this in Anger uh you run into a lot of Cur work and it's it makes things more complicated for the Schuler uh autoscaler uh whatever you're using cluster autoscaler or maybe something more modern uh it's not always aware of this and when you add this and also your normal workloads are in the same cluster and normal batch workloads and those workstations that again becomes pretty tricky um and lastly uh one big problem with everyone runs into if they want to train those foundational models is just scarcity of GPU there's just not enough gpus uh on major public Cloud providers uh they cost a lot uh there's some alternative providers coming up like corv for example that are specific specifically focused on gpus and provide more competitive pricing uh but now you're in the situation where maybe most of your uh infrastructure runs on your traditional public cloud and you kind of have to burst out out for certain capacity for certain workloads into this alternative Cloud uh and how do you marry those proes together uh it's I mean there's no like Silver Bullet here there are a few projects like Submariner uh some multicluster kubernetes Management Solutions which help with this uh they don't completely there's a lot of uh interesting complications in terms of how those clusters communicate and especially where the data sits if you start pulling from uh again like taking Amazon as an example they will charge you through the nose for egress um so you have to think about those things and if you can make your infrastructure uh multicluster aware from the GetGo so it's not necessarily has to be managed on the kuus layer but kind of an application layer uh that makes things a lot easier all right and then I will pass it to Riley who will talk about more specifical specific things that do it out of with llms great um so my name is Riley and I work as a ml platform engineer at Autodesk um and I'll be talking a little bit about how we've been able to leverage a lot of the little goodies that um all L described associated with metaflow and our ml platform so a little bit of context about our organization um so for those who are unfamiliar Autodesk is a multinational uh company that specializes in creating software that basically empowers our customers to design and make across um various uh sectors including manufacturing construction engineering and media entertainment um among our most uh popular software are AutoCAD in Maya where Maya is being extensively used um in visual effects for feature films uh Autodesk software serves Architects Builders Engineers designers manufacturers um 3D artists and production teams Glo globally and our Innovative um design and creation Technologies cater to a diverse array of different Industries um aiming to support innovators worldwide at scale so Ai and ml are active in expanding areas of focus in both research and product engineering at our uh company uh numerous product teams at Autodesk aim to innovate by harnessing Ai and ml um basically to introduce um intelligent capabilities for our customers the Autodesk AI lab was established in 2018 and it's a Research Unit dedicated to um applied research um often collaborating with academic institutions like MIT the lab has contributed over 60 peer-reviewed research papers uh with P past uh research efforts revolving around um problems in the 2D and 3D geometric uh design space and the team produces some fascinating work on using language to communicate the semantics of geometry uh to make design work more intuitive so I'm so I'm part of uh Autodesk machine learning um platform team focused on building a unified uh ML and AI platform our goal is to offer comprehensive tooling across the entire ml endtoend ml workflow um enabling data scientists and ml Engineers to accelerate model development um um from research to production so earlier this year our team assessed various ml Frameworks uh and basically we selected metaflow as a foundational part of our infrastructure due to its human Centric design and it aligns um with some of like the preferences of many data scientists and and ml engineers at audit desk who are more comfortable with tooling like Jupiter notebook and python um given the diverse uh the diverse problem space at our our organization um we basically recognize the necessity for a platform capable of evolving alongside um these new use cases we like that meta flows extensibility and interoperability make it an ideal framework that can grow and adapt to our platform's evolving needs um enabling us to introduce new capabilities on top of metaflow so we really like that metaflow could be used for a variety of different applications like data compute orchestration um and versioning and all of this suited our needs metaflow is particularly advantageous uh for constructing machine learning pipelines by chaining tasks throughout the entire ml workflow um all the way from data preparation to model training and evaluation and we've been able to leverage meta's convenient wrapper for utilizing adop batch for scaling workflows um particularly in the context of training models it is also helpful to us that metaflow also excels in um versioning in ensuring reproducibility by versioning flow runs data snapshots um and artifacts and it can also serve as a useful experiment tracker aside from that metaflow has proven to be really versatile for us while we intended to use metaflow specifically for mod train training some of our users have actually been able to uh use metaflow for building and orchestrating data extraction pipelines as well this is a high level overview of how we've integrated metaflow into our managed IDE using sagemaker Studio we call it uh am stud am studio and it basically acts as the productivity Suite um for developing and training models in order to access studio uh users interact with a UI which spins up their um personal Studio instance uh once they are logged into their studio notebooks they're they have access to metaflow which um they then leverage to create and orchestrate their jobs at scale and they'll be able to monitor their workflows through the metaflow UI additionally users um are able to bring in their own data and metaflow and Studio have access to import this data we also successfully integrated metaflow into our giops workflow uh enabling data scientists to deploy their jobs all the way into production one of our users have adopted this pattern to deploy their data extraction pipeline using metaflow on step functions where metaflow basically compiles and Maps their flow to a workflow orchestrator for scheduling um and our users can conveniently conduct local testing of their flows including local runs um they can perform unit testing effortlessly before pushing their flows all the way to production So currently we are using step functions as our orchestrator but we are in the process of transitioning to Argo workflows as our workflow orchestrator of choice our AI lab research team is heavily invested in using Ray to train their models so we needed to figure out a way to integrate Ray with our manage training infrastructure powered by metaflow to ease their transition to adopting metaflow Ray plays a critical uh piece in transform forming our ml platform into an AI platform and as such we work with outer bounds to evolve metaflow compute layer such that we could orchestrate the create the creation of a transient Ray cluster using aw spatch multi-node parallel jobs um through metaflow and in this way batch handles the infrastructure setting up the driver and workers while Ray optimizes the hardware usage So currently we're we're creating the ray clusters on top of AWS batch but our team is also in the midst of working with outer bounds to add compatibility with kubernetes as well we used the existing abstractions from the parallel decorator and then teamed up with outer bounds to introduce the metaflow ray decorator um which basically is like an extension or Plugin and this metap Ray decorator can essentially decorate a step and subsequently sets up the necessary Hardware using aw AWS batch multi-node parallel jobs users then insert their Ray code in the in in the steps and then during execution metaflow starts up the transient Ray cluster runs the ray application and then subsequently shuts down the cluster upon task execution and thank you very much um happy to in uh happy to entertain any questions and answers as well we've got a couple of minutes if anybody has any questions oh uh sorry for the microphone I have a small question I only present an Amazon Cloud get the mic yes uh so metlow by itself uh was originally built with AWS as a kind of primary target but it's actually flexible enough that we have Azure and gcp Sport as well uh you can technically use it on the on brem cluster as well so it's not super strongly tied to anything specific kind of an obnoxious question but meant in a friendly way I given it seems like shoehorning AIML into kubernetes is a bunch of challenges so why use kubernetes I'm just sort of why not just set aside a bunch of a cluster of G and just run it that way well so the KU is really flexible for those things so it's really uh I think for me personally it's mostly the configuration management platform if anything else it's not necessarily like about containers it's just how do you manage applications and like those distributed fashion and um I mean I've been using AWS batch for for years before this I've been using you know just ad hoc container R platforms we actually built one with one of the previous companies and uh kind of scheduling containers on servers is the usually not the hardest part it's actually everything around it when you start thinking about you know networking and different types of scheduling and autoscaling that's where it gets complicated and kuus happens to have a lot of momentum behind it uh and there's a lot of open source project that help with those things got it open AI uses them so it can't uses it can't be that bad it was just more yeah exactly I guess this is a question for the second speaker probably but um I was just wondering uh you you went into kind of how Ray is in your platform and you can spin it up spin it down I was wondering um do you also use Ray for production inferencing or can you speak to yeah so right now we're focusing on Ray for training because um our aab research team is right now working on the training stage but we are definitely um looking to do right use Ray for inferencing in the next phase probably the next fiscal year um right now our research team is kind of experimenting I I can't really divulge too many details because it is kind of proprietary but they're basically trying out different um models um using different deep learning techniques and um they're not at the stage right now to uh move that into production but we definitely are looking into you leveraging Ray for infuencing for sure are some of the models trained in Ray are they uh able to be deployed in the production inferencing environment you have right now is that using stag maker is that correct um I I didn't catch oh I was just wondering whether the models that you train with Ray whether you have a path for them to be deployed outside of Ray like in Sag maker yeah that's actually exactly what we're using um right now for um our more basic models where we're using P maker end points to deploy but we we definitely need Ray for more complex use cases which we anticipate in the next fiscal [Music] year