Authors: Han Li, Sahil Khanna, Jocelyn De La Rosa, Moping Dou, Sharad Gupta, Chenyang Yu and Rajpal Paryani

Background

About a year ago, we introduced the first version of Griffin, Instacart’s first ML Platform, detailing its development and support for end-to-end ML in a previous post. In our last week announcement, we unveiled Griffin 2.0, our advanced ML Platform, shedding light on the transition from Griffin 1.0 to 2.0 for a more cohesive user experience, refined model management, distributed computation capabilities, and enhanced deployment automation. The Griffin 2.0 post provided insights into all key components of the system, including the Griffin UI, Feature Marketplace, ML Training Platform, and ML Serving Platform.

In this article, we focus on the technical details of constructing the ML Training Platform (MLTP) within Griffin 2.0. MLTP introduces innovative functionalities such as a centralized web interface, a distributed computation framework, standard ML builds, orchestration services, and a scalable metadata store, which collectively contribute to the comprehensive creation and management of training workloads at Instacart.

Crafting Strategies

In the development of Griffin 2.0 MLTP, our goal was to establish a unified and centralized platform where all MLEs can seamlessly create, track, and manage their training workloads. We aimed to support distributed machine learning, preparing the platform for tasks such as distributed training, batch inference, and Large Language Model (LLM) fine-tuning. Simultaneously, we wanted to address the limitations identified in Griffin 1.0, as detailed in our previous post. These considerations led us to make the following strategic design decisions:

• A Singular Interface: Griffin 1.0 required MLEs to navigate through multiple systems for comprehensive training workload information. Griffin 2.0 brings all tools into a unified web interface, greatly simplifying the user experience and streamlining model training development.
• Centralized Unity: Griffin 1.0 incorporated various training backend platforms, resulting in increased maintenance overhead. Griffin 2.0 consolidates all training workloads onto a unified Kubernetes platform.
• Standard ML Runtime: Griffin 1.0 lacked a standardized approach for modeling frameworks and versions, leading to maintenance challenges. Griffin 2.0 introduces standard runtimes for various ML frameworks with consistent building blocks and package versions to address this concern.
• Horizontal Scalability: Griffin 1.0 relied on vertical scaling for accommodating increased training data and evolving model architectures. In Griffin 2.0, we adopted Ray to support horizontally scalable distributed workloads without imposing excessive complexity on our MLEs.
• Metadata Store for All: Griffin 1.0 lacked effective model lineage information, making it challenging for users to understand and manage the complete model lifecycle. Griffin 2.0 implements a centralized metadata store to ensure thorough model lineage management throughout the lifecycle.

How It’s Built

System Architecture

We’ve constructed the MLTP with the following major building blocks (Figure 1) to make it a centralized service and platform with distributed computation capabilities:

Metadata Store

We’ve established a series of data models to manage all aspects of the model lifecycle, including:

• Model Store: Specifies the raw, untrained model architecture used as input for training jobs.
• Offline Feature Store: Contains metadata for features in our offline storage to be utilized for training.
• Workflow Run: It represents a specific training job request submitted to our workflow services.
• Model Registry: Maintains information about models generated by training jobs. This registry is used in post-training processes like evaluation, batch inference, and real-time inference.

API Endpoints

We’ve developed an API server with accessible REST endpoints for users to interact with the Metadata Store to achieve the following:

• /api/training/models/: Create and retrieve model architecture in the Model Store
• /api/registry/models/: Generate, save, and retrieve the trained model from Model Registry during post-training processes
• /api/features/: Index and retrieve offline features
• /api/training/dataset/: Generate and retrieve training dataset
• /api/workflows: Create, retrieve and terminate training jobs in Workflow Run

Workflow Orchestrator

The workflow orchestrator consists of two major components, as illustrated in Figure 2:

• MLTP API service: The workflow APIs empower users to customize their training jobs with specifications like worker count, CPU/GPU units, memory, SSD attachment, and runtime choice, all while shielding them from the complexities of Kubernetes operations. Figure 3 shows an example of workflow creation and resource specification UI to create a Ray Cluster. When a user initiates a workflow, the MLTP API service communicates with the ISC backend worker to generate Kubernetes resources.

• ISC (Instacart Software Center) + Kubernetes: We chose Kubernetes as the orchestration platform for our training workloads. This choice allowed us to centralize the methods for launching workflows to different training backends within a single, unified platform. We collaborated closely with Instacart’s core infrastructure team who built and managed the Instacart Software Center (ISC). ISC is a suite of tools for creating, validating, and deploying software across the company. We integrated Ray and Kubernetes in ISC to enable company-wide adoption and leverage existing build & deploy features. Upon receiving workflow requests from the MLTP service, the ISC worker establishes a unique Kubernetes namespace for each workflow definition to ensure resource isolation and share identical environment variables and authorization settings. Subsequently, the ISC worker generates RayCluster and RayJob Custom Resource Definitions (CRDs), initializes them, provides an endpoint for the Ray dashboard URL, monitors the status of the launched containers, and manages cluster lifespan.

Training Application

Here’s the typical process on MLTP for creating a training job:

1. Users begin by customizing their input, which includes tasks like organizing features and trying out different model designs. They then select a specific model from the Model Store, pick input data from the Training Dataset, and set training configurations to start a formal training job.
2. When they’re ready to initiate a training workload, they have the option to either access the Griffin UI or directly employ Python SDKs to send a request to the workflow services.
3. The workflow services create Kubernetes resources based on the user’s input. This can be a simple single-container job or a more complex multi-node Ray cluster.
4. After training, the results, like metrics from MLFlow and logs from Datadog, are shown visually. The model weights and other training items are then stored in the Model Registry for easy access in tasks like evaluation and inference.

Our design considerations for streamlined and standardized training job creation are as follows:

Unified Interface

The centralized service and APIs offer a unified interface for managing the model development lifecycle:

1. Users can prototype by scheduling training to a Ray cluster from their laptop or Jupyter server to leverage distributed computation.
2. Once the model is ready for production, users can utilize the Griffin UI to create a “production” workflow definition.
3. Alternatively, if there’s an existing production pipeline in Airflow, they can rely on Griffin task operator and sensor operator, which integrate with Griffin workflow APIs to submit and query workflow runs.

Throughout all stages of the model lifecycle, users interact with MLTP using the same API interface.

Standard ML Runtimes

When we surveyed training applications implemented by our users, we found the majority of code isn’t focused on model development. Those applications shared many similarities, for instance, parallel executors for feature transformations, batch inference, and distributed GPU workers for parallelizing training batches. Optimizing these components was time-consuming for MLEs, requiring them to learn various frameworks and APIs, not to mention the debugging complexity.

To simplify model development and encourage more users to adopt MLTP, we’ve developed standardized ML runtimes with structured and configurable input parameters. This solution offers predefined approaches that cover the majority of use cases for building different models at Instacart, be it decision trees or neural networks. Figure 5 shows the high-level structure of training configurations.

Underneath, we rely on Ray APIs to implement various components of the training pipeline as shown in Figure 6. Most of these building blocks are shared across different standard runtimes for various ML frameworks. Moreover, these standard runtimes are easily scalable, allowing users to initially test in Ray local mode and then effortlessly scale the workflow in remote distributed environments.

Lineage

Thanks to the centralized data storage layer, the model lineage is now much better managed. We can effortlessly trace the following lineage pairs, as illustrated in Figure 7:

• Model architecture — Workflow run
• Training dataset — Workflow run
• Workflow run — Model Registry
• Workflow run — Output Metadata URLs (Datadog, MLFlow, Ray Dashboard, etc.)

This simplifies the training-to-post-training process for users, helps them be more productive in both experimentation and deployment and enables them to reproduce results for sharing and debugging purposes.

Lessons Learned

We learned a few lessons during the journey of building the next generation of ML training infrastructure:

• Unified Solutions: To scale ourselves effectively with a growing number of ML use cases and a limited ML infra team, we opted to unify our ML training solutions. As a result, we can now offer a more consistent training job abstraction. For instance, transitioning to Kubernetes as our only orchestration platform requires a one-time migration for existing training jobs, but it enables us to provide a consistent experience between tasks and offers users additional benefits like distributed computation and enhanced metadata management.
• Balancing Flexibility and Standardization: We aim for our platform to be highly flexible and capable of accommodating a wide range of ML applications for model training. Simultaneously, we recognize the importance of providing standardization to cater to the majority of use cases and accelerate development velocity.
• Consider the Bigger Picture: In our redesign of next-gen MLTP, we didn’t just focus on training alone. We also took into account model serving and feature engineering. We collaborated closely with stakeholders from other Griffin 2.0 projects to co-design the data models of training jobs, ensuring seamless integration into the end-to-end process for easier deployment and metadata management.

Acknowledgments

The MLTP project received support and funding from various teams such as Cloud Foundation, Build & Deploy, Developers Productivity, Ads ML, Fulfillment ML and many more at Instacart. Additionally, we are immensely thankful for the guidance offered by members of our ML consulting group: Jin Zhang, Chuanwei Ruan, Aditya Subramanian, Liang Chen, Trace Levinson, and Peng Qi.

Unveiling the Core of Instacart’s Griffin 2.0: was originally published in tech-at-instacart on Medium, where people are continuing the conversation by highlighting and responding to this story.

Author: Han Li

At Instacart, we take pride in offering a diverse range of machine learning (ML) products that empower every aspect of our marketplace, including customers, shoppers, retailers, and brands. Along with the growth of business and ML applications, we have encountered an increasing number of use cases that require distributed ML techniques to effectively scale our ML products.

To support the emerging requests, we have designed the distributed ML system with the following key design considerations:

• Scalability: We must be able to scale a large number of distributed ML workloads on both CPU and GPU instances with a robust system that maintains high performance and reliability.
• Resource Efficiency: We aim to fully utilize distributed computation resources to maximize system throughput and achieve the fastest execution at optimal cost.
• Diversity: The selected computation framework should be as extensible as possible to support diverse distributed ML paradigms. In addition, we should be able to handle diverse environments that are specific to different ML workloads.

Our team is currently working on a variety of ML applications that utilize the power of distributed computing to efficiently solve complex problems. Some examples include:

• Parallelizing thousands of small to mid-sized training models across distributed hosts with efficient resource utilization and reasonable queuing time.
• Using data parallelism to enable deep learning models to work with large datasets more effectively.
• Developing scalable batch inference capabilities with a large number of parallel jobs processing over a vast amount of data.

To meet these requirements, we have chosen Ray as the foundational computation framework for distributed ML. In this article, we will walk through the system architecture, analyze a case study, compare it with legacy systems, and demonstrate the value of Ray-based distributed ML.

System Architecture

We recognize that a system with a simple user experience is essential for increasing developer productivity and maximizing the impact of our distributed ML applications. Therefore, in addition to the design principles of scalability, resource efficiency, and diversity, we prioritize the following objectives to simplify the user experience:

• Development Simplicity: We aim to simplify the development process for ML engineers by providing them with tools to test their distributed code seamlessly in a local environment and distributed environments.
• Maintenance Simplicity: Our objective is to offer a serverless development model that empowers our engineers to develop and run workloads without the burden of managing any services. We believe that by eliminating the need to handle service management, we can effectively improve MLE’s productivity.
• Seamless Integration: Our goal is to seamlessly integrate distributed ML workloads with Instacart ML Platform Griffin, without creating a redundant set of MLOps tools or platforms. Rather than reinventing the wheel, we aim to enhance our existing infrastructure.

We have extended our workflow orchestration tools on Griffin to provide a unified experience for running both prototyping and production workloads on Kubernetes-based distributed ML environments powered by Ray. Fig 1 illustrates the end-to-end user experience.

Here are breakdowns of each component interacting with Ray distributed environment:

• During the development stage, users can package prototyping code snippets on their development environment and launch them on remote AWS EKS hosts through an internal launcher API that abstracts Ray Cluster access through Ray Job API (see Fig 1a).

• When users are ready to automate their code in production pipelines (i.e. Airflow), they can use the same set of APIs to launch their containerized application to production Ray Clusters (see Fig 1b).

• Each Ray Cluster is independently configured with Python installations and environment variables. This effectively separates the workspace environments between clusters, avoiding overhead of maintaining a monolithic python environment, which will be discussed in detail in the next section (see Fig 1c).

• At the application level, users need to build ML applications using Ray APIs to achieve a wide variety of distributed computation patterns, such as Ray Core APIs, Ray AIR and Ray Serve.
• On the controller side, we use KubeRay as the controller of provisioning, scaling, and deleting Ray Cluster resources on Kubernetes.

Next, we will dive deep into a case study of Instacart Fulfillment ML, the very first use case of Ray-based distributed ML at Instacart, in which thousands of models need to be trained in parallel efficiently.

Case Study: Parallel Fulfillment ML Jobs

At Instacart, ML is applied throughout the lifecycle of fulfilling every order placed on the Instacart App (Fig 2). This includes batching multiple orders together to improve efficiency, routing to determine optimal delivery routes, ETA prediction on order arrival time, and supply/demand forecasting to prevent lost deliveries when there are not enough available shoppers, and many more. A significant number of such models are trained with data collected in a specific neighborhood, and some of those models are also sensitive to the data collected within a particular time frame, therefore it’s common practice to partition the entire national dataset based on one or multiple space/time attributes, and then launch one training job per data partition, preferably in parallel.

The most common “attribute” is the geographical location of orders, called “zones.” At Instacart, we manage thousands of unique zones to partition our national dataset. For each Fulfillment ML application, it’s typical to launch as many as thousands of parallel training jobs to cover all zones in a single model experiment.

Previous Solutions & Limitations

Our legacy solution implemented a distributed task queue service using Celery; Fig 3 illustrates the architectural diagram. For each unique model application (represented as Model in the diagram) that requires parallel training across zone 1~n, training jobs for all zones are published as tasks to the same task queue on Message Broker. Subsequently, a group of Celery workers subscribes to each task queue and executes the tasks inside it. The tasks are executed asynchronously, and whenever a Celery worker is free, the task at the front of the queue is removed and assigned to the free worker. Once the worker finishes, it updates the task status in the Message Backend and fetches the next task in the queue to execute.

As the number of ML applications requiring zone level parallel training increased over time, this type of monolithic service started to get hard to further scale due to the following reasons:

• Inefficient resource usage: Many ML applications shared the same task queue and Celery distributed worker services. It turned out that distributed workers couldn’t get fully utilized very efficiently:
  - Worker node had to be over-provisioned with enough resource headroom to fit all types of models. However, in Fulfillment ML, most models are quite lightweight, such as regression, classification, etc. Therefore, when a small size model was running, the worker was poorly utilized. For example, in the left-hand side graph of Fig 4, a lightweight model was running on a 16-CPU instance, with CPU utilization only around 10% to 15%.
  - The long-running service could leave the system idle if the load was unbalanced throughout the day. For example, in the right-hand side of Fig 4, the queue was not running any tasks for over 60% of the day, leading to inefficient resource usage.

• Long queuing time: It’s very hard to achieve faster execution time without upscaling the task queue services:
  - Low Celery worker concurrency is configured per host of task queue, because it’s not feasible to simply increase concurrency to fit all models on a monolithic service.
  - When a task queue gets assigned with too many tasks, such as Fig 5 in which 300 to 1k+ tasks are in queue all day round, the queuing time of each task to be executed will significantly increase, slowing down the end-to-end execution time of each ML application.
  - Therefore, the only way to speed up a busy queue is to do upscaling to add more physical hosts, which continues adding poorly utilized resources to the system.

• Difficult dependency management: It’s difficult to manage Python dependencies for all model applications running on the same queue. Upgrading a specific python package version is challenging.
• Maintenance burden: This system is complex to automate, and does not provide a simple approach to replicate the system in a local environment for testing purposes.

Improvements by New System

When we started adopting Fulfillment ML workloads for our new distributed ML systems, the most significant change was the migration from long-running monolithic services to Ray-based serverless applications. Instead of all models running on the same environment, each Fulfillment ML application is launched as an independent Ray job associated with a dedicated Ray Cluster to handle all its zone level training jobs, as illustrated in Fig 6. This new design offers several advantages compared to our legacy solution:

• With our new system, each model now has its own isolated workspace environment, enabling independent management of Python dependencies and resource usage on a per-model basis (Fig 7).

• Long-running task queue services are now deprecated. This reduces costs on computation resources, and most importantly, it makes the infrastructure much easier to scale to host upcoming new applications.
• Serverless applications offer users the ability to provision resources per Ray worker more accurately, resulting in much more efficient resource utilization. Fig 8 provides a real-world example of this improvement by comparing the before/after CPU utilization of a Fulfillment production model running 1.5k unique training tasks on the same group of hosts:
  - In the legacy system described earlier, only a few Celery workers were configured to run on each host, resulting in CPU utilization of around 10–15%.
  - With our new Ray-based architecture, we can more accurately customize resource requirements. Through benchmarking, we found that allocating just 2 CPUs per zonal training job is sufficient. As a result, we are able to create more concurrent Ray workers on the same instance, significantly increasing CPU utilization per host to up to 80%.

• Improved resource utilization contributed to faster end-to-end execution. With the same production model mentioned above running on 10 16-CPU instances, there were only 10 Celery concurrent workers in total on the legacy service while a Ray based application can launch as many as 70+ Ray workers. As a result, we are now able to greatly speed up end-to-end completion time–previously ~4 hours– to 20 minutes!

We have also found that converting existing training functions into Ray distributed function calls is a relatively straightforward process. As an example, we have refactored an existing project implemented in the legacy solution to use Ray (see Fig 9). The conversion required only a few changes to existing functions or class objects to make them Ray remote executable. Additionally, we needed only a few lines of code to orchestrate the remote calls asynchronously with rate limit handling, if necessary.

Learnings & Future Work

We’ve gained valuable insights from our experience building distributed ML solutions and migrating early use cases away from legacy systems. We’ve discovered that hosting a monolithic service as the computation backend for all distributed ML applications has limitations in scalability, efficiency, and diversity, given the rapidly evolving and highly diversified nature of ML applications. For instance, by transitioning from monolithic services to standalone serverless applications for Fulfillment zone level training, we were able to significantly boost execution time, resource utilization, and development simplicity.

Although the topology of each application can be highly customized by the application owner in distributed ML, it’s essential to have unified build and launch tools on the ML platform layer. This approach enables us to support more types of distributed ML workloads in the future while ensuring our platform remains extensible.

Looking ahead, we aim to foster broader adoption of distributed ML applications across Instacart on multiple ML product lines. Our goal for 2023 is to further mature our Ray-based distributed ML solution as the unified solution running as part of Griffin. As we embark on this journey, we’re excited to see what distributed ML can help us achieve next at Instacart.

Distributed Machine Learning at Instacart was originally published in tech-at-instacart on Medium, where people are continuing the conversation by highlighting and responding to this story.