Keynote: Building the Infrastructure that Powers the Future of AI - Vicki Cheung & Jonas Schneider

Has helped us build our infrastructure that powers the future of AI. How we might have started as an unusual use case for Kubernetes, but we've actually managed to make it work so far. So what is OpenAI? We're a non-profit AI research company, and we focus most of our efforts on both basic research and deep learning, which is an area in machine learning, as well as research in specific projects. Examples include robotics and learning to play video games. So as a non-profit company, our mission is to democratize the AI technology, and because of that, we care a lot about open source software. We rely on building infrastructure for a research lab. Running deep learning is different from infrastructure at a startup. So there are two key differences. One is the type of workload that we're running. We're mostly running large batch jobs, so the jobs can span from hours to days or even weeks. And sometimes we run a single copy, sometimes we run tens of thousands of copies. So that puts a very different strain on the cluster. And the type of jobs that we run often have very specific resource requirements. So this means running on special hardware that's configured for doing numerical computation, and some of that includes GPU.

What this means is research ideas come and go, and we don't want to invest a lot of time engineering something that might not make the cut. Yeah, so as we just mentioned, we may need to do prototyping, and of course we want to build our infrastructure to best support that and to really be optimized for this kind of workflow, which is very different from the workflow you would use when you're building, say, a web app with microservices. And so a couple of key things that are important to our researchers is a consistent platform. That is because if they run their experiments and they behave weirdly differently when they run it two times, we actually don't have any knowledge gained because we don't know what actually happened, if something happened in the experiment or something happened inside of the platform at the same time.

It's like, 'Oh yeah, let's try it, let's try some different configuration.' And so our infrastructure needs to be flexible enough to be reconfigured super easily. And so the way we're tackling this is by having a central infra team, which is Vicki and me and a couple of other people as of recently, that runs a common platform that the researchers use to run their experiments. And so the goal there is to empower users to run the experiments by themselves without needing engineering support, but at the same time be shielded from the complexity of the underlying systems, which is like the GPUs and different cloud providers that we use. At the same time, because we have such a small team, we also need low operational complexity. We want to set the thing up and in the end, up to 4,000 nodes, and it went pretty well. We had no major outages. And the key thing that really makes this work for us is that Kubernetes allows us to ship infrastructure as internal products. And what that means is if infrastructure is a product, it has to have specific product-like qualities. The most important of which is it needs to have a consistent UI that's understandable and doesn't surprise users by default. It should hide complexity, but it should also be composable so that if you need to do something that's inherently complex, you should be able to do it without ripping your hair out. And it turns out that the Kubernetes API is actually pretty nice.

The Kubernetes API is actually very nice. So Kubernetes API is our primary service layer, and what this means is we can take advantage of a lot of the abstractions that are baked into the Kubernetes API to help us compose the different components in our experiments. So typically what happens is in the life of an experiment or a new research idea, we start out as a tiny prototype that runs as a single pod in the cluster. You might have one Python script that just runs in your network, and then you see how that goes. And so as the idea turns out to be more promising, we need higher parallelization, need to think of different ways of distributing the training. And also if the idea seems promising, we also will need to run many different copies of the experiments, maybe running variants of the idea or running hyperparameter sweeps. So here's an example of a large batch of experiments that we've had to scale up very recently. The example is that we're training agents to play video games. So we're running agents against all these games on our Universe platform, and Universe is a platform of over a thousand video games.

The vanilla setup, we put a video game, let's say a Flash game, in a container, and then we also put our agent in another container, and we just run them together in a pod. And this works super fine for prototyping because you can use the same setup on your laptop and then you can ship it off, run it remotely in our cluster. But then it turns out actually that games are really complex. So you can run this for maybe solving Pong or something, and then once you start running into racing games or platformers, it actually takes our agents a lot of time to learn to solve it sufficiently. And so unless you're... Concretely, it's that you have one master process which is your TensorFlow parameter server, and then you split out your training into two parallel workers that interact with their own copy of the Flash game. And then that way you can split it up and basically scale your training. Once we start training this way, we can now tackle more difficult games. And as the complex games go, they have heavy graphics requirements and so on, and so they actually became a CPU bottleneck. And so what happened was we had to split it up. The Flash games had to run on their own in a separate host, and then the agents would interact with the Flash games over the network.

There's like how we'd network or how we schedule the different jobs on this cluster so that they can have high networking bandwidth. And oftentimes the agent might have special compute requirements. And so it's actually pretty complex to set it up, but in practice all it took was just changing the YAML file. Yeah, so as you've seen, for every experiment usually we start out very simple. There's a single person running a single job, leads to a single pod, all is well. But then eventually we'll scale it up. And so the researcher might be running on a single core for, let's say, two months, and then suddenly they realize it's a paper deadline and now they need to run all their experiments for that paper in two weeks. And now they suddenly scale up from one core to 10,000 cores, which can happen. And the thing is that this is very hard for us to predict in advance and very bursty, of course. And that actually means that we can't keep our cluster size constant because that would just mean it would either be too small, which means we can't run the big jobs, or it would be too big, which means it would be very wasteful because we have a lot of idle capacity. So it actually means we have to resize our cluster dynamically. So to do that, we wrote an autoscaler.

We need to have support from core Kubernetes necessarily. So what the autoscaler does is it gathers the state of what's happening in the cluster, like what nodes are there, what pods are waiting to be scheduled, what pods can't be currently scheduled because you don't have capacity. And that's just using the Kubernetes API. And then we actually provision new nodes exactly to fit the resource requirements that are missing from the current cluster. And this happens using the AWS provisioning APIs. And that means we have, different from how you use auto scaling groups where they just scale up after there's a lot of CPU load for a while, we actually have complete control over the types of machines that we provision. When people like, especially opt out, like 'Don't put me on the spot instance, I didn't implement checkpointing.' But basically these types of policies are something that we can very, very easily implement and very cleanly separated from the actual logic that does the scaling. And I think it's fair to say that this is a result of how the Kubernetes API is designed, and that it basically allows you to interact with the primitives that are provided using your own custom logic. So why are we using Kubernetes? So we've explained that our workloads are fundamentally different from maybe the common use case that Kubernetes was designed for, which is microservices. But Kubernetes provides good primitives and interfaces that we can build on, and it abstracts away a lot of the complexities of the underlying infrastructure, which is great for our users that may not have operational experience. So the flexibility goes both ways. It exposes a very clean API for our users, which are researchers, and also for the maintainers of our infrastructure team, so we can extend and customize the cluster as needed, because you never know when the experiment will lead us. And so we're very excited to continue to invest in Kubernetes, and we're also super excited to see the new developments coming in 1.6 and also for the next year.