Setting the file. One moment.
Skill 29 · Knowledge Base For Startups
Subchapter 29.112
references/learn/genai/how-snorkel-ai-achieved-over-40-percent-cost-savings-by-scaling-machine-learning-workloads-using-amazon-eks.mdMarkdown19 KBView on GitHub
Machine learning (ML) startups are often heavy compute users, because they train large models using high-end GPUs and deploy them at scale for inference. AWS Startups (opens in a new tab) partners with startups from inception to IPO, and has helped thousands of founders and artificial intelligence (AI) innovators build their businesses on Amazon Elastic Kubernetes Service (Amazon EKS) (opens in a new tab). Amazon EKS is a popular choice to build and host ML models because it provides the flexibility of Kubernetes with the security and resiliency of being an AWS managed service that is optimized for building highly available containerized workloads.
Snorkel AI (opens in a new tab) is one such company that benefits from Amazon EKS. Snorkel AI equips Fortune 500 companies, federal agencies, and AI innovators to build, adapt, and distill foundation models (opens in a new tab) (FMs) and large language models (opens in a new tab) (LLMs) to perform with high accuracy on domain-specific datasets. Using Snorkel’s data-centric approach toward AI development, organizations have built production-ready AI services for use cases, including insurance claims processing, financial spreading, clinical trial analytics, and accelerating proactive well management for offshore drilling.
Over the past several months, the Snorkel team has been hard at work addressing the unique challenges of designing efficient infrastructure to support ML development workloads without increasing infrastructure bills, lowering developer velocity, or impairing user experience. Their ultimate objective was to reduce the cluster compute expenditure for Snorkel Flow (opens in a new tab), their end-to-end ML platform, by more than 40%.
Snorkel’s Snorkel Flow AI data development platform enables data teams to rapidly build AI applications by employing an iterative loop of programmatic labeling, quick model training, and error analysis. Each project starts when users create a small number of labeling functions.
Labeling functions employ simple heuristics, external databases, legacy models, or even calls to large language models to apply labels to swaths of unlabeled data based on encoded expert intuition. The platform’s weak supervision algorithm combines these rule-based functions to determine the most likely label for each record. Users then train a simple model based on these probabilistic data points and assess the impact of each labeling function. In the analysis phase, users investigate slices of the data where the model underperforms. Then they build or modify labeling functions, train another quick model, and continue the loop. When users are satisfied with the quality of their labels, they build a final model on an architecture from the model zoo—ranging from logistic regression to FMs—and export it for deployment.
Due to the nature of this workflow, Snorkel Flow’s infrastructure experiences varied periods of high compute usage. Operating costs naturally increased as the customer base and ML product capabilities of Snorkel Flow scaled. To achieve efficient growth, Snorkel aimed to understand how to enhance margins while operating state-of-the-art ML software. Snorkel implemented the following practices to achieve more than 40% reduction in cluster compute costs.
Software-as-a-service (SaaS) startups often have opportunities to optimize their cloud spending. It is essential to understand the unique factors that drive these costs.
For Snorkel, there were two significant factors:
Snorkel’s team is keen on creating systems that allow efficient scaling without a linear increase in infrastructure costs. Consequently, Snorkel developed a comprehensive autoscaling solution tailored for their ML workloads on Amazon EKS to address cloud expense concerns. This solution not only expedited workloads that require burst compute but also achieved their cost reduction goals.
In addition to the autoscaling solution, key strategies that contributed to the over 40% reduction in cloud expenses include:
In this post, Snorkel shares the process for addressing these scaling challenges to help facilitate the design of better infrastructure for ML systems. If you’re new to Kubernetes, read Snorkel’s Introduction to Kubernetes (opens in a new tab) post to learn more about the basics and their Machine learning on Kubernetes: wisdom learned at Snorkel (opens in a new tab) post to learn more about their journey with Kubernetes thus far.
In practice, Snorkel Flow’s interaction with AWS follows the sequence outlined as follows. As the Snorkel Flow platform relies heavily on containers, the migration to AWS was almost seamless.
This architecture has yielded a stable and snappy experience for Snorkel Flow users.
Prior to the architecture described in Figure 1, early iterations of Snorkel’s infrastructure used fixed resources. Snorkel’s users shared that these bursty workloads could take too long to complete and therefore negatively impacted their experience.
Manual scaling of compute resources proved unscalable and error-prone, leading to cloud costs that stayed elevated even during periods of low usage. It was the worst of both worlds: low cloud cost efficiency and slower-than-needed performance.
To address these challenges, Snorkel implemented autoscaling at multiple levels in their infrastructure, as discussed in the following sections.
The Kubernetes distribution of Snorkel Flow involves a set of deployments (opens in a new tab) running in an EKS cluster that contains pods (opens in a new tab) that run various components of the platform.
As shown in Figure 2, to address the unique challenges of working with bursty compute workloads, Snorkel’s team introduced a new concept for Kubernetes pods: semantically categorizing them as either “fixed” or “flexible.”
This fixed/flexible framework gives Snorkel a domain-specific means to enable automated cluster downscaling, which allows them to turn on the cluster autoscaler (opens in a new tab) on Amazon EKS without their finance department messaging them every hour.
Snorkel’s initial approach was to deploy podDisruptionBudgets (opens in a new tab) on the EKS cluster to prevent the cluster autoscaler from moving flexible pods during the day and from moving fixed pods at all. While effective, this approach left the Snorkel team unsatisfied because it downscaled far fewer nodes than what was theoretically optimal.
To address this, Snorkel layered on a pod scheduling optimization that isolated fixed pods to a small fixed group of nodes. It scheduled flexible and worker pods (which are considered to be fixed pods but are ephemeral due to worker node autoscaling) in the remaining flexible group of nodes.
These changes allowed Snorkel to efficiently downscale the flexible nodes at night, when it became safe to move around flexible pods and scale down the vast majority of worker pods.
Enabling efficient downscaling of the vast majority of the cluster’s nodes (i.e., the flexible nodes) allowed Snorkel to meet their target of reducing the cloud costs for hosting Snorkel Flow by over 40%.
Snorkel divides the implementation of the solution described in the previous section into three sequential efforts:
The Snorkel Flow platform abstracts compute into a paradigm where jobs wait in Redis queues and workers run as processes in worker pods.
Snorkel implemented a worker autoscaling solution (Figure 3) for worker pods by running a recurring function in Snorkel Flow’s backend API. Every few seconds, this function checks the Kubernetes cluster and Redis for both upscaling and downscaling eligibility.
If there are jobs waiting in one or more relevant Redis based queues, the function will ask the Kubernetes API to provision additional worker pods to process these jobs. If the Redis queue is empty and there are no running jobs in the job registry, it will ask the Kubernetes API to destroy the worker pods to free up reserved CPU and RAM resources.
As shown in Figure 4, with this worker autoscaling implementation rolled out, Snorkel Flow’s worker pods became ephemeral, appearing in the cluster only when jobs needed to be processed.
The PodDisruptionBudget (opens in a new tab) resource protects certain pods against disruption (for example, voluntary restarts) by allowing for the specification of the maximum number of pod replicas that can be unavailable at any given time. As shown in Figure 5, setting this value explicitly to 0 for a deployment ensures that the cluster autoscaler will not downscale nodes running the deployment’s pods.
Implementing this resource on hosted Snorkel Flow instances safely allowed the cluster autoscaler to downscale underutilized nodes. However, the cost savings Snorkel realized were marginal—they were still unable to downscale the majority of their nodes because all Snorkel Flow pods were protected by an associated podDisruptionBudget.
Upon closer examination, Snorkel’s team realized that this protection does not need to exist all of the time. Workloads are bursty, and most user interaction with Snorkel Flow occurs during a customer’s business day, meaning it is safe to relax this protection outside of business hours. Similar to worker autoscaling, Snorkel implemented a recurring function that toggled podDisruptionBudgets “off” overnight for an instance’s flexible pods by setting the maximum number of unavailable pod replicas to 1, up from 0 (Figure 6). The previous worker autoscaling solution combined with the ClusterAutoscaler and PodDisruptionBudget feature was able to downscale many more underutilized worker nodes than before. Customers deploying Snorkel Flow in their cloud can configure this as needed.
Even with these improvements, Snorkel saw that a majority of underutilized nodes were not being downscaled at all.
Upon further investigation, Snorkel realized that the issue stemmed from fixed and flexible pods occupying the same node. This was problematic because a fixed pod, pseudo-randomly assigned to a node containing flexible pods, would “pin” that node and prevent it from being downscaled, even when it is underutilized. This lack of control over the scheduling of fixed pods led to periods where the vast majority of the cluster’s nodes could not be downscaled, even though they represented far more compute power than needed at the time.
Snorkel leveraged the podAffinities (opens in a new tab) Kubernetes resource to address this, which enabled them to constrain which nodes a pod is eligible to run on based on the labels (opens in a new tab) of other pods already running on any given node. They added labels to the pods to differentiate between fixed versus flexible pods, and added a podAntiAffinity (opens in a new tab) stanza to their deployments configuration to ensure that fixed pods are not scheduled on nodes running flexible pods, and vice versa.
This implementation of podAffinities allowed Snorkel AI to split nodes into two functional groups: the fixed group of nodes containing fixed pods, which can never be safely moved between nodes (for example, Redis due to cache), and the flexible group of nodes containing “flexible” pods that are either ephemeral (such as worker pods) or safe to move outside of business hours (such as overnight).
Although it is possible with manual intervention during platform maintenance, Snorkel cannot automatically downscale the fixed nodes. This solution, however, allows them to automatically downscale the flexible nodes because they have now isolated the unmovable pods into the fixed nodes.
The Snorkel and AWS Startups teams hope that sharing this thought process and these solutions helps other startups to build better infrastructure for ML workloads, which are rapidly becoming more important as ML, large language models (opens in a new tab), and other FMs (opens in a new tab) make their way into production for organizations all over the world.
Many thanks to David Hao, Edmond Liu, and Alec Xiang for helping to make this technical vision a reality for Snorkel. Special thanks to the aforementioned as well as Matt Casey, Henry Ehrenberg, Anthony Bishopric, and the entire Snorkel infrastructure engineering team for their thoughtful feedback on this article.
Ganapathi Krishnamoorthi is a Senior ML Solutions Architect at AWS. Ganapathi provides prescriptive guidance to startup and enterprise customers helping them to design and deploy cloud applications at scale. He is specialized in machine learning and is focused on helping customers leverage AI/ML for their business outcomes. When not at work, he enjoys exploring outdoors and listening to music.