Two Sigma: Push computing research boundaries with Google Cloud

Key Takeaways

[Music] all right welcome everyone uh to Ops 210 2 Sigma push Computing research boundaries with Google Cloud I'm Alex I'm a software engineer at to Sigma hi I'm Dax I'm a hedge fund customer engineer hey and my name is Ma kraski I'm a product manager on Google kubera engine great all right thank you so much for joining us again I'm Alex Hayes I'm an engineer on the team that builds and maintains uh two sigma's research compute platform I'm fortunate to have the opportunity uh this afternoon to share my team's story and it's the story of providing gpus given just unprecedented demand and it's a problem I'm sure many of you have uh experienced over the last few years then Dax and Magic will uh introduce you to a new Google Cloud product Dynamic workload scheduler or DWS that helps solve this problem and then I'll cover how we integrated with DWS to acquire gpus for our researchers uh this is just a quick disclaimer saying this isn't financial advice let's get started all right so what does compute look like at to Sigma I think the best way to understand our compute platform is to get an idea for what it enables our researchers to do and then this will be good context for you to understand our scale in gke and how we run millions of jobs a day uh then I want to tell you about our challenges with gpus in today's industrywide capacity crunch so two Sigma is a quantitative hedge fund that develops And Trades uh large data strategies these strategies are packaged up in some something we call a model and so here I'm highlighting the highle pipeline for developing a model at 2 Sigma so where do we start well two Sigma researchers are scientists So based on an observation an idea researchers form a hypothesis about how the markets might move and then they'll explore this idea using techniques and data and generally quite a bit of data uh but they have to test and verify that their idea with this data makes sense so they run these things they run simulations and these simulations are compute intensive and with lots of researchers it's incredibly compute intensive and that's where a research platform plays its biggest role and so after final review successful models are then traded in the markets okay so I've covered what we enable let's dive into what compute intensive really means and so over the past few years we've grown this platform to to Super Compu scale our compute platform in gke regularly exceeds 10 pedabytes of concurrent memory and 1 and a half million cores and that's a lot of compute and we do this with dozens of very large gke clusters spanning multiple Google regions and it offers users a uh wide range of machine shapes including all types of gpus in Google cloud and then we tune this environment for batch workloads so these clusters aren't running web services they're highly optimized to run a wide range of different batch compute workloads U including simulations and so given the scale the dozens of clusters and the diverse machine shapes we have to make this easy to use for our users and we do this with cook cook is two sigma's own internal batch scheduler and so cook was created in 2015 so it's had a really long and important history at the firm it originally started as a Mesa Schuler but today it's evolved to orchestrate all of this compute in kubernetes so users submit workloads directly to cook cook will cue them rank them and then launch them in our clusters and cook is capable of doing this uh millions of times a day so millions of PODS a day um it's capable of running a couple hundred 100 pods a second or launching them and it supports queuing hundreds of thousands of pending workloads and so cook is really the single integration point for researchers and developers at two Sigma to run their workloads at scale Cooks API saves every user from having to create and then manage their own infrastructure it's centrally managed and it's run by an amazing team and so it integrates with many open source distributed compute Frameworks as well so this includes spark dask and Ray but it's also integrated with the rest of two sigma's platform uh storage and to sigma's own internal compute Frameworks and so that's a really quick overview of our platform I want to switch Focus uh to something that's probably that probably isn't news to a lot of you um so over the last year the demand for gpus has just increased dramatically and we've all been challenged by this industrywide capacity shortage and this is true across many models but it's especially cute uh with the latest ones like nvidia's a100 and h100 and so two Sigma is no different from other companies that want to access and use these gpus and so platform teams are in a real bind right where you're trying to satisfy user demand at a reasonable cost or sometimes at all and it's doing that in a highly constrained environment so naturally we Face questions from users like why can't I get gpus we use several different strategies to tackle uh a problem like this one of them is just by shifting to Demand right so you can suggest to users well why don't you try another model like v00 or T4 they're relatively more obtainable and this often works but sometimes users just need more they need more vram they need their results faster you know there are a number of reasons and so some users can get away with the other models but some users clearly need or benefit from these latest gpus and so I so helping these users in particular is really difficult because h100s and A1 100s are at this just entirely different level of demand and stock out so I'm sure you've seen reports of companies uh stockpiling h100s and you can solve this a number of ways we consider a lot of them and one of them is by buying GPU reservations in Google but reservations are a long-term solution and the GPU ecosystem is just rapidly changing and then this reservation uh solution only gets worse when users tell you well I only need these Gus for a limited period of time in this situation we would uh severely underutilize that reservation and so you can see how platform teams are in a tough spot so given unprecedented demand this industrywide capacity shortage and understandable needs from users what do you do if you're facing these same issues I think you're in the right place we have some good news for you so you're going to hear from uh our partners at Google about two really exciting new technologies that we use to solve this problem uh T Sigma is one of the first customers uh to use these new technologies and I'll dive into that integration uh at the end of the session but first i'm going to pass it over to Dax who will do a deep dive into Dynamic workload scheduler thanks Alex before uh before we dive into Dynamic workload schedule I'd like to share sh something I hear from uh customers all the time as a as a customer engineer it goes something like this Dax we have a great idea that we need 200 h100s for to run for a few days and we need them as soon as possible and we can't start our job until we have them all can you help and my answer before DWS was well it might be tricky it might be tricky to get 200 machines or 200 cards immediately right now now we're probably going to have to grab them what we can grab what's available and then over time build up to the 200 and that can take time that can take that can take days or or even weeks and during that time those machines are underutilized and it's not a it's not a great way to to to optimize um your your Cloud spend because you have them sitting around waiting for all the machines to come with DWS uh We've make this problem a little bit better we can I have a great answer to the customer I could say you can use flex start mode um and everyone will be happy so what is DWS and why is the obtainability of scarse GPU resources like a100 and h100s so much better compared to On Demand well there's a few reasons for starters DWS has access to additional capacity than is available via on demand when you ask us for a machine via DWS you're going to let us know the region you want it in the machine type and count of machines and the runtime duration that last one is really important knowing that we will get the machines released back back to us at a certain time unlocks additional capacity as we can Target DWS dedicated pools and expand across limited duration pools for unparalleled obtainability to use DWS you can leverage several Integrations across gcp and we continue to add to that list um first and foremost you can use GCE directly via manage instance groups you can leverage it with our Google Cloud batch uh managed offering you can use gke which we're going to hear a little bit more about um in a bit and vertex AI training um there's two main modes to Dynamic workload scheduler we have Flex start mode and calendar mode drill into Flex start mode for a little bit it basically offers just in time capacity for defined duration jobs so in the example I mentioned earlier when the you know the requests for all those machines um just came to us and wasn't planned and we wanted them as soon as possible Flex start mode is the way to go so basically you're going to say hey Google I need to run a workload for two days using four a100 machines with h gpus each in this region we'll put it in a queue and as soon as we have them ready we'll deliver to them to you all at once and you can run your workload you don't have to worry about paying for idal resources while you build up to to all those machines um and flex start mode the max runtime duration is limited to 7 days so if you need to run longer than that you'll need to be able to checkpoint and save your data somewhere and um start it up uh again with another with another request um it supports all gpus in all regions and you only pay for what you use so in this example if you requested these machines for seven days but you only use them for two and then you turn them off you're only going to be used build for two now on to calendar mode calendar mode offers a simple way to buy a short-term reservation for optimized capacity this option is great for users who need a guaranteed start time so in this case you're going to know the machines and regions region that you want to run in but you're going to have a little bit of an advanced notice and you'll you'll make the request to us we'll we'll approve it and then you'll get them at at that start time calendar mode is currently in private preview um it supports h100s and you can get up to 128 gpus um and you can get them for up to 14 days which is a bit longer than um Flex start mode um in calendar mode you do pay for the entire duration of of the reservation so this is great for that use case where you know you have something you need to train maybe once a quarter um at at a schedule and you and you don't need to purchase a reservation you know 24/7 you can just you can just use this to to save some money with that I'll uh I'll hand it over to matrick to talk about how gke can integrate with DWS thanks Doc so uh dyamic cor CL scheduler is a foundational API or a capability that is built by our colleagues in the core infrastructure compute layers and uh you still when you run your workloads in kubernetes cluster or other uh service you need other apis that allow you to get access to it and orchestrate the entire interaction with cute provisioning of machines so in kubernetes in Google kubernetes engine we recommend two surfaces two API surfaces that you should use to leverage and the dynamic workload scheduler the one that probably will fit the needs of most of you uh is q a that's a open source job scheduler um built for kubernetes um that orchestrates the entire interaction of the dynamic or scheduler or as an alternative to to Q you can also use an extension to the cluster autoscaler with a new provisioning request API which is a more basic and foundational level API that on one hand will be most powerful and will give you the most the biggest level of customizability when you when it comes to interactions and building the entire workflow of leveraging cute VMS and cute provisioning but at the same time it also will require the most work out of you so that you can uh you can um you can integrate it in your more maybe best Po and customized solution that you might have in your company so let's take a closer look on how to leverage these apis and how the interactions would look like um so H wrong button okay so uh first maybe what is cure uh so a couple of years ago I think two or three we've recognized that uh kubernetes really lacked a good um job regestration system that is cloud friendly and together with the open source community at cloud cloud native Foundation we've started the work on a project called Q um with a uh that is essentially a add-on to kubernetes that's performing the task of scheduling full jobs so the core kubernetes scheduler is capable of scheduling individual PS and there are use cases where you need to make decisions whether something will be scheduled or not and and will will have access to the infrastructure on full workload level the most common use cases are essentially various research platforms batch processing uh platforms which where typically the end users the researchers they will create a very long backlog of jobs way bigger than whatever the platform is capable of handling and then you need a system which will through policies and the defined configuration decide okay these jobs can run these need to wait maybe these need to be preempted this is capacity gets divided between teams and the full orchestration on workload level is performed and that's the role of Q uh it's uh if you want to learn more about it uh it has a dedicated website qsh and here we're going to talk how do you really leverage it so that you can orchestrate the dynamic work CL schedular um provisioning uh again wrong button uh so uh first what do you need to do to perform the configuration in here the platform admin first needs to model the capacity in the cluster uh using cluster cues so what you essentially do with the sample um custom resource. configur CU that I'm showing here on the screen is that you need to say uh what's the level of what's the amount of capacity and what types of machines uh your and type of consumption model that your users will be allowed to use and um then Q looks at the backlock of jobs that want to get access to the cluster wants to execute and it verifies what resources they need versus the available quot and then decides okay this job can be admitted this one needs to wait perhaps this one needs to be preempted because a more important workload is coming in in the backlog so we need to make space for that one and uh on the workload level so your researcher or end user on the end user system uh that is interacting with this needs to make very simple two very simple levels of configuration one in a label it needs to indicate that this work should be orchestrated by que using a certain cluster Quee to provide the name in the first label that you see here on the screen as an example second important uh element of the workflow that uh it needs to conform to is that uh you need to create it in suspended mode uh suspended mode is something like for the job API that we've added U I think two years ago also in anticipation for for the need for of capabilities like that where you create the job resource just the job uh definition in the cluster but actually no pods get created by the job controller there are no containers running in the cluster just yet there's only the information about the intent to run this job and then Q with a K will be responsible for changing that flag and the the logic behind making that decision is going to be driven by available capacity existing machines and those various factors that you code in into Q's configuration so really you just need these two elements in the world workload so on workload level the integration with Q is very simple and uh I'm showing the example of uh a job and how do you do it with a job but you can also use various custom workloads or even be pods all they need to do is they need to have an API that allows you to get the semantic of suspension so that you create you you express the intent that you want to run this workload but pods will get created only when Q allows them to be created uh the whole the whole tool and the whole system has quite Advanced capabilities that allow you to build a very sophisticated and complicated batch processing platform where you can have multiple teams these teams can have their own dedicated slices of the cluster you may Define policies where perhaps team a and Team B have their own dedicated capacity but they can also borrow from one another if the other team is not not utilizing fully the capacity you can model a variety of resource flavors like types of machines shapes of those machines maybe consumption models maybe you want to prioritize running on your reservation but when you're out of that reservation then you'd like a low priority work close to spill over to spot machines all that logic can be coded into q but for the dynamic work scheduler what you think what you can think about in here is that the core and the essence of what Q does is it takes a list of jobs that want to be executed and decides which need to be run when there is available capacity so if you think about also the model that Dax explained with how Dynamic Ro CL schedule Works where it simply looks at the uh incoming in Flex mode it looks at the incoming request to get VMS so that something can be executed on those VMS and like if these two layers these two concepts can be connected for a simple API they could actually work jointly together so that Q recognizes what's the backlock of workloads that are coming in and the dynamic workload scheduler can act on that and prepare capacity for those workloads and the API to build that integration is a provision request so provision request is an extension to the open source cluster autoscaler so it's an API that will work on all clouds um but um as you probably know um uh cluster autoscalers typically have custom implementations on every cloud so GK has its own cluster autoscaler that is well integrated and with lots of features that are that are baked in into into its logic of interactions with with our infrastructure while at the same time it still conforms to the open source apis and provisioning request is an extension to Cluster autoscaler which allows you you through a custom research to essentially Express hey cluster of scaler I have this workload that I want to run with this spec uh can you make sure you're going to have resources for it and actually when you have them can you send me a thumbs up so that I can trigger the workload and there are a variety of use cases for it but uh uh for the dynamic workload scheduler uh we've built a special mode of operations for it so current provisioning request uh has three modes how you can enable it the first and the simplest mode is you ask the a provision request through um a custom resource like the one that you see here on the screen you ask it a question hey do you have capacity for this job and it tells you yes or no it's actually a very valuable mode of operations for on premises use cases where capacity is fixed you just you run your workload when you see that actually there's space for it then the second mode is meant more for cloud use cases where cluster autoscaler will check does it have the capacity and if it doesn't it will resize the cluster so that you actually do have the capacity for it and again that's a very generic scale up of the cluster and then last but not least there is a special Mode called cute provisioning where um the cluster autoscaler will go to Dynamic workload schedule and tell that hey I have the workload coming in can you get VM plan VMS for me and when you find them let me know and I will tell uh my users to to start their workloads and that's the mode that can be used that is used by Q to get machines you can also leverage provision quest yourself so that you build the integration yourself U and uh and uh and integrate it also in a more complicated or best book platform um maybe as it may look complicated and with all of the code Snippets that I've shown let me show you on a simple example of a conversation how these old various components how they interact together and how they provide capacity for you so first you have your user right they want to this is a researcher they want to run a workload maybe they use a notebook maybe they use a shell script um create a yaml file everybody loves them with a job and um and so they interact with the cluster and tell it okay please execute this workload for me so Q recognizes that workload and accepts it if everything is okay and um and sends information to Cluster autoscaler through the provision request API hey can you give me the necessary machines that are needed to run this workload so cluster autoscaler picks it up and uh translates it into a request to the dynamic workload scheduler uh to dispatch it to so that it can start looking for machines in our infrastructure that will be used to fulfill that request and the run that workload after a little bit of time uh the the um uh the uh the machines get created so Dynamic work scheduler U gets the machine the VMS uh and running a cluster autoscaler recognizes that so it attaches those machines to the cluster uh based on that it is able to send the information to Q telling it that the job can now be executed so Q changes that suspended flag that you've seen earlier on the screen that the workload can be triggered and then job controller takes over or whatever is the controller like Ray job and it starts to create pods and containers that will actually execute the necessary tasks so after some time um whatever is necessary to actually fulfill the job up to seven days as Dax mentioned uh the workload gets completed so uh on one hand q q or whole kubernetes system interacts with the user confirming that the workload is completed they get the necessary data that they wanted or whatever was the output of the processing workload while cluster autoscaler kicks in and takes and standard autoscaling mechanisms start to take over so um unless you want to run another job on that uh on these VMS uh cluster autoscaler will just remove them resizing the cluster down so that you don't pay necessary costs associated with using uh these expensive vmms now this whole flow that I have shown you may look simple uh but but keep in mind that this is the happy p this is this is what happens when everything goes well and we all know Happy paths are easy to implement things start to get funky when something goes sideways maybe the user canels their request maybe something failed how do I inform all of the components and Cascade the information about the error everywhere so that then later it's debuggable and and the system is in a healthy State still after the failure uh all of that is covered for you in the implementation of Q and it handles all of those interactions uh greatly simplifying the entire effort and the experience necessary to integrate with dyamic work CL scheduler as I mentioned uh in the examples you've seen that this works with the with job controller but we also support custom Cate classes of jobs um we went on and worked together with various open source communities on implementing in the job set in flux in Cube flow Ray uh it also supports bare pods so uh which are typically used by workflow schedulers or orchestrators like airflow so um you can use a variety of tools to orchestrate these jobs and essentially when you configure queue you add those labels to indicate which classer que should be used that's it you don't need to do many changes in airflow or in Ray job uh so that you can leverage WS it will just work for your users out of the box and uh if you're interested in those more advanced features uh of queue then a lot more is also coming we very heavily innovate on the capability add think like a new there is a new release coming out every two months things like Network topology over scheduling multicluster um queuing uh more advanced algorithm for sharing resources in hierarchies in cohorts to model more complicated Enterprises and how they manage budgets plus ux improvements all of that is coming in here uh of course if you don't need that all you need is a wrapper that simplifies Dynamic CL scheduler that is absolutely an option and we give you all of the code templates and necessary to get started the stuff that you will find on our website allows you to um get going with your first initial POC within like 10 minutes of work effectively on this and with that I'll pass to Alex how the integrated with this great thanks guys uh for the great overview of DWS and Q so these two technologies are really exciting um uh especially I think the open source collaboration so I want to share uh with you now how to Sigma integrated um with DWS and Q to access A1 100s and h100s so the integration primarily consists of two things the first one is DWS node pools and then the second is uh using Q to access the node pools and ultimately the gpus through cook if you're not familiar with gke um I'll just give you a quick primer on node pools so each gke cluster can have one or more node pools and node pools are how you define a specific machine shape is uh available for your workloads to use and so for example you can define a node pool with machines that have 64 cores 500 gigs of memory and 4 a100 attached and so the first step to using DWS is to uh create node pools with DWS enabled and you can do that pretty easily with just a single flag so here in the CLI enabled C provisioning or you can even do it through terraform like we do and that's really simple this is really all you need to do on the gke side the second part to accessing gpus is that you have provisioning request enabled so you have a couple options as magic was uh just hinting at one of those is by doing it the hard way and that's using provisioning requests directly so like you said provisioning requests are very powerful but it's a very low-level API integrating directly with provisioning requests means you have to essentially build your own controller you have to submit the provisioning requests with pods specs that outline the shape of the nodes you want then you have to build a custom watch on those provisioning request continually check their status if the request didn't fail if it was actually provisioned and you actually have that node then and only then do you submit well you can submit pods whenever but you would then um submit pods to consume those provisioning request and then they would schedule on the notes and then finally you should do um just some cleanup when you're done with the process and so really that's a lot of coordination and it's not an unreasonable task for platform teams uh to do something like this it just takes a bit of time and resources to implement when we were doing this initial evaluation we estimated that it would take us about two months to do it at our scale but I want to show you an easier way and that's by using Q so using Q accessing the gpus through DWS is as simple as setting uh a label on your pods right so Q handles the complexities of creating and managing the provisioning request and so here I've drawn a simplified uh diagram of how we integrated Q into cook so once you've installed Q uh and configured it to work with provisioning requests the only active part of this is to submit the pods uh with that label right so so to start the process cook will create pods and it will set the label to the name of the cluster queue that you want the pods to go to and they're configured to request gpus we've configured Q does this work no so we've configured Q um to monitor the single cluster Q for pods it's looking at the specs and it's deciding okay so you're requesting this this amount of gpus I need to create a provisioning request that actually will go out and scale scale that and that's when the cluster autoscaler will pick up the provisioning request and it will go to those DWS node pools you uh we had created and will actually request nodes with those gpus once the Q controller which is watching these provisioning request sees that the nodes have been created um in your cluster through DWS it will then admit the work workloads um so for jobs it will um unsuspend for be pods like us it will remove uh scheduling gate and that's when the cube scheduler will actually see the pods essentially for the first time and it will schedule those on the provisioned GPU nodes and that's it your workloads are running on gpus so uh this integration with DWS really had an amazing impact on GPU obtainability 2 Sigma and I'm just showing a very short snapshot here um but we were experiencing some pretty severe stockouts for a100 gpus and when we integrated with DWS andq in December we uh improved obtainability for A1 100's on demand to 80% and when h100s became available in January we also enabled those through DWS and Q and so I would confidently say that this integration with DWS and Q is what really unlocked these uh scarce gpus for customers in our platform and I think it does it at a pretty um it does it in an economical way and it does it at really great scale So yeah thank you that's uh that's it we appreciate your feedback in the mobile app and I think we have some time for [Applause] questions we'd be happy to we'd be happy to take your questions and also if uh you would like to uh spend some time with folks from our product team or user ux researchers uh we would really love to get your feedback if you scan this code it will ask you to fill in a short form that will allow you to uh schedule a session uh together with product or ux researchers um working on [Music] JK