Introduction
At Krea, research and production run on the same GPU clusters. We make no distinction between them. That works until you think about the economics of the GPUs, at least in our case: is a GPU more valuable running research or production? For us, the answer has almost always been research.
Part of the reason is how we buy. When we buy GPU clusters for research, we optimize for the quality of the GPUs and the InfiniBand fabric, while keeping the price reasonable. In our experience, it is worth paying a bit more for good GPUs than settling for cheap, problematic ones.
So these GPUs are shared across every kind of workload, and the demands compete. Researchers want the whole cluster for training runs. Users on the website want to generate content. It is a constant tug of war, and when research needs the GPUs, research gets them. That is the whole reason we paid for them.
This raises the obvious question: if all the GPUs are busy training, how does production run? Doesn’t that hurt the user experience? Yes, it can. But if you build the right systems around it, both sides get what they need: researchers use the entire cluster, production keeps running, and the user never notices.
Building on Virtual Kubelet
We built our solution on top of the Virtual Kubelet project (VK from now on). To explain why it fits, here is a brief Kubernetes 101. Every node in a Kubernetes cluster runs a background process called the Kubelet. One of its jobs is to register the node with the apiserver. When the scheduler assigns a pod to the node, the Kubelet makes sure the pod’s containers are actually running there. It then reports status back, keeping the node and the Kubernetes state reconciled.
There is nothing special about this background process. That is the whole idea behind VK: if you emulate a Kubernetes node and fully control every step of the pod lifecycle, you can deploy your pods and containers anywhere. The funny part is that VK itself runs as a Kubernetes pod. You end up with a pod that registers itself as a node. And since you fully control this node, you can report whatever you want back to Kubernetes. We use this fact to report that our node has effectively infinite GPU capacity.
Many projects have been built with VK as a base, and at first we tried one of them: interLink. It did not fit our needs. It was too constrained for what we wanted to do, and it had some bugs. So quite early we decided to build on top of VK itself.
That turned out to be surprisingly simple and pleasant, because the API VK asks you to implement is quite nice:
type PodLifecycleHandler interface {
// CreatePod takes a Kubernetes Pod and deploys it within the provider.
CreatePod(ctx context.Context, pod *corev1.Pod) error
// UpdatePod takes a Kubernetes Pod and updates it within the provider.
UpdatePod(ctx context.Context, pod *corev1.Pod) error
// DeletePod takes a Kubernetes Pod and deletes it from the provider.
DeletePod(ctx context.Context, pod *corev1.Pod) error
// GetPod retrieves a pod by name from the provider (can be cached).
GetPod(ctx context.Context, namespace, name string) (*corev1.Pod, error)
// GetPodStatus retrieves the status of a pod by name from the provider.
GetPodStatus(ctx context.Context, namespace, name string) (*corev1.PodStatus, error)
// GetPods retrieves a list of all pods running on the provider (can be cached).
GetPods(context.Context) ([]*corev1.Pod, error)
}External Providers
Of course a lot more was built, but in theory that is all you need, plus some setup code. Our objective is to scale anywhere. So we abstracted a lot of the business logic and common routines on top of VK’s APIs into a new set of APIs, for what we call External Providers. This gives us a clean separation of concerns. To add a new provider, we implement only the basics of how to talk to it. Everything underneath handles the rest, including:
- reconciliation, at startup and at runtime
- garbage collection
- correct translation to Kubernetes status
Inside the VK node: what runs there, and the steps it performs with the API server.
One nice fact: this system is 100% stateless. The whole state can be re-derived from what is running on your cluster and in the External Provider. So there is no concern about keeping state. If the VK pod restarts, it picks up where it left off and keeps working. You do have to build this, but it comes very naturally.
Here is how pod deletion looks when implemented in an External Provider.
func (p *MyProvider) DeletePod(ctx context.Context, pod *v1.Pod) error {
podUID := string(pod.UID)
if podUID == "" {
return fmt.Errorf("pod metadata must include a UUID")
}
// Create unique name for that pod
name := p.makeUniqueName(podUID, pod.Namespace, pod.Name)
// Delete the pod on the provider
// ... Call the api or cli using `name`
if err != nil {
log.G(ctx).WithError(err).Errorf("Deletion failed")
if strings.Contains(err.Error(), "not found") {
return errdefs.NotFoundf("deletion failed: pod %s not found (%w)", name, err)
}
return fmt.Errorf("deletion failed: %w", err)
}
log.G(ctx).WithFields(map[string]interface{}{
"namespace": pod.Namespace,
"name": pod.Name,
"uid": pod.UID,
}).Info("Deletion successful")
return nil
}The implementation is short. The one constraint it exposes is that every pod in the cluster must be individually addressable on the External Provider side. Per-pod operations depend on this: it is what lets you run kubectl delete pod <pod running in vk>, and what lets an HPA kill pods one at a time.
Other steps of the lifecycle can and will be more complex, especially because you are adapting something that was never part of the Kubernetes world to fit inside it. Creating a pod is a good example. Most of the logic there is translation: turning the Kubernetes pod spec into something the provider can understand. Sometimes the spec alone does not carry enough metadata to do that. Our solution is to add special annotations to the pods that enrich the spec with the missing information.
Self-healing and architecture
Another property we like about this system is that it is pretty much self-healing. Suppose you scale a deployment with an HPA and one of the replicas running in the External Provider fails or gets evicted, for whatever reason. As long as your reconciliation loop reports the correct status back to Kubernetes, the HPA detects the failure and creates a new pod to replace the problematic one. You get a lot of the systems Kubernetes has already built, for free.
There is also an architecture decision to make: one VK pod per provider, or one big VK that manages all providers and decides internally how to dispatch between them. With the first approach you depend on the Kubernetes scheduler, which might not be the optimal solution to your problem. We chose the second: a single VK that manages all providers.
Kubernetes schedules a pod onto the VK node like any other node in the cluster.
What’s next
In the next blog post we will talk about scheduling: how we guarantee high GPU usage in the cluster, and how the offload system works when we run out of GPUs.
.png)
