Authors: Dong Wang, Arunkumar Eli

Motivation

Graph neural networks (GNNs) have emerged as one of the leading solutions for ML applications. Most real-world data can be represented as graphs - see this blog for a comprehensive overview of what use cases are best solved with GNNs and their key advantages. With tools like PinSage and the proliferation of Open Source Software like PyG, GNNs have had a considerable impact across many applications and thus gained significant popularity.

While graph learning is very powerful, there are significant challenges data teams run into when building a production-grade ML platform using GNNs. Running a SaaS service at scale introduces significant operational complexity and overhead due to the requirements of running the models in the cloud. In addition to the intensive storage and compute requirements, a platform needs to deliver fast response times to many customers in parallel while also minimizing the cost of expensive hardware usage.

At Kumo, we built an enterprise ML platform powered by GNNs - this platform makes it possible for users to immediately tackle many prediction problems by setting up a single data connection over a data warehouse and to query the future using a simple and intuitive interface. Kumo leverages Kubernetes and other tools in the ecosystem to build a scalable, robust, and cost-efficient platform.

This article will discuss the work that Kumo has done so that customers can leverage an enterprise-grade solution out of the box. In subsequent blog posts, our team will talk about modeling innovations that enable Kumo to provide strong modeling accuracy for different use cases.

Operationalizing Large-Scale Machine Learning in the Public Cloud

Design Principles

Here are the core design principles of our platform that enable us to best support our customers.

Ability to scale to support thousands of training models
Ability to run many inference calls in parallel while supporting customers at the individual level to optimize their models with thousands of hyper-parameters
Ability to quickly and dynamically provision resources as needed by customers, while minimizing idle utilization of the hardware
Performing the above capabilities in a cost-effective manner

Key Challenges

Operational complexity

For a SaaS company like ours, the model training and inference workload can vary significantly based on the two factors.

Data size and computation cost: Our customers' data vary in size as they come from different domains. It ranges from 32GB to 10TB. While smaller data can be fit into the memory, the large graph requires distributed storage to store. Applying machine learning models to customer data in a cost-efficient way becomes crucial to our customers.
Training and prediction schedule: The workloads from customers can vary significantly from day to day, creating high variability and unpredictability in demanding resources. A customer can submit many training jobs during a short period or can schedule predictions daily or hourly. The optimal customer experience requires a highly elastic solution that can handle the arrival of many jobs at once, spinning up resources as needed rather than queueing.

Because these factors are highly variable, the proper sizing of resources cannot be determined beforehand and preconfigured. To serve hundreds or thousands of customers, we need to ensure fast model training and inference in our design.

Cost complexity

Training GNNs requires large memory and compute resources. In addition, training time is significantly longer because it scales directly with the size of the graph.

GPUs are the standard in training GNNs due to their ability to provide highly-parallel computations. While GPUs offer great efficiency for training, the instance that host GPUs are expensive to use in the public cloud.

On average, training a workload on a public cloud will require anywhere from 5 to 100 GPU hours at ~$5/hr, running through hundreds of hyper-parameter settings. Running the machines as a SaaS platform introduces additional operational complexity - in order to support hundreds of customers, you need a truly elastic platform so that the machines are not running idle 24/7.

Kumo’s Solution

Here is the solution we implemented at Kumo.ai to overcome the challenges of delivering an efficient and highly performant AI solution as-a-service.

Figure 1: Training Infrastructure Diagram. A depiction of how the queries from N customers are handled by JIT nodes in the Kubernetes cluster

There are two major components of our solution design - the Kubernetes orchestration engine and the Smart Node Provisioner.

(1) Kubernetes Orchestration Engine

Kubernetes forms the backbone of the infrastructure at Kumo, orchestrating both training and prediction jobs.

The primary benefits Kubernetes brings to our platform are

Ability to package the training code and its dependencies as docker containers
Ability to use Kubernetes Jobs for the creation of transient trainer pods that perform the training
Ability to use node scaling mechanism to avoid warm virtual machines idling

(2) Smart Node Provisioner

One of the major design considerations of an enterprise-scale product is its ability to provision compute resources quickly and upon request by the user. We considered two different tools in building our solution - the Cluster Autoscaler and Karpenter.

Cluster Autoscaler

Cluster Autoscaler is an open-source tool built for Kubernetes hardware capacity management that is supported by major cloud providers. It increases or decreases the size of a Kubernetes cluster (by adding or removing nodes) based on the presence of pending pods and node utilization metrics. It automatically adjusts the size of the Kubernetes cluster when one of the following conditions is true:

Some pods failed to run in the cluster due to insufficient resources
There are nodes in the cluster that have been underutilized for an extended period of time, and their pods can be placed on other existing nodes

We tested the Cluster Autoscaler for our use case and concluded that the scaling-up capability is not immediate. A pod will be in a pending state for a few minutes until a new worker is added. This introduces lag and uncertainty which does not meet the third criterion of our design. In order to provide predictive analytics at scale, we looked to Karpenter as the tool of choice.

Karpenter

We leverage Karpenter to dynamically allocate the proper compute resources in AWS and to automatically terminate unused compute. Karpenter is an open-source node provisioning project built for Kubernetes.

Karpenter has two control loops that maximize the availability and efficiency of the cluster.

Allocator – fast-acting controller ensuring that pods are scheduled as quickly as possible
Reallocator – slow-acting controller replaces nodes as pods capacity shifts over time

The ability to launch Just-in-time ephemeral nodes has helped us scale. We have The rest of the blog talks about implementing Karpenter and our future plans.

Implementation on Amazon EKS

We built a scalable solution using Amazon EKS - the primary Kubernetes engine for AWS. To provide an overview for how this could look in a basic production environment, we outline a few steps to set up a working environment and deploy a job.

As a prerequisite, install the following tools in your environment:

Once the prerequisite tools are installed,, you can deploy an EKS cluster and configure KarpenterNode IAM Roles in your environment, and run the following steps:

Install NVIDIA device plugin. Kubernetes provides a device plugin framework that we can use to advertise system hardware resources to Kubelet. The NVIDIA device plugin for Kubernetes is a Daemonset that allows you to automatically:
- Expose the number of GPUs on each node of your cluster
- Keep track of the health of your GPUs
- Run GPU-enabled containers in your Kubernetes cluster.
Install Karpenter Helm Chart and
A sample Karpenter GPU node Provisioner is here. This provisioner allows a few well-tested GPU instances to be allowed to launch. With this, when a customer launches a prediction job, our ML platform decides the memory requirements on the fly and creates a Kubernetes job request. Karpenter selects the node that satisfies the capacity request and launches just in time.
A sample Job deployment is here. This sample job will be launched on a g4dn.16xlarge node as it is requesting a minimum of 250Gi memory.

This is a simplified recipe for a basic environment. For a platform like Kumo’s, these steps would be significantly more involved given the nature of the scale and complexity of our platform. Kumo provisions many nodes for each training job that serve different purposes - for example, a feature node will be deployed entirely dedicated to storing features used in training - and would leverage multiple Karpenter provisioners for managing the different nodes.

Future Work - Expanding to Bigger Scale

We currently have a robust solution that can handle hundreds of concurrent training jobs, scaling quickly to meet demand. However, when looking to expand to thousands of concurrent jobs, the cloud provider alone cannot guarantee the availability of computing resources, especially GPUs, in a single region or zone. We started our development with Karpenter in AWS, so we have more work to do in expanding to support other areas of deployment.

Multi-cloud and Multi-cluster Pattern - While there is a community push to make Karpenter support other cloud providers, we need a flexible job dispatcher that works across Multi-regions and Multi-cloud. This creates an additional layer of complexity when running a job because Karpenter now needs to determine which GPU and cluster to run it on.
Proactive Autoscale - we plan to introduce a very different and improved customer experience with our autoscaling system. We will redesign to scale well in advance before a machine learning workload is triggered - guaranteeing the compute resource will be ready to use in time while maintaining a low cost of utilization. We will also provide a visual demonstration of training progress in real-time to provide visibility into the system.
Consistent and Manageable Platform - We are expected to deliver consistent and manageable service for all of our customers as we continually add new use cases, models, and users. A job controller and job monitor will be thoughtfully designed to enable our platform to execute machine learning jobs on the predictable hardware between iterations. Frequent checkpoints and a failure-over mechanism will be implemented for improved reliability.

Acknowledgment

Everything accomplished at Kumo is a result of collaboration across multiple teams. Many thanks to Ivaylo Bahtchevanov and Hema Raghavan for reviewing this blog.

As Kumo solves many business problems for its customers using GNNs, stay tuned for exciting articles from our fantastic team. If you are interested in working on such exciting problems in this space, see the job opening for the Kumo Cloud Platform team

Cookies settings