10 Training pipelines

Now, with the company following trends and the AI hype train full steam ahead, it’s no surprise that it hires Jane, a young and promising ML engineer whose interest in ML is undeniable. After onboarding, the CTO delegates her first problem to solve: build an AI-powered assistant that will help Pizzaiolos perform basic visual assessments like listing and calculating the number of components on the pizza base for each order.

The software development lifecycle in this pizza company hasn’t included ML systems so far. Thus, engineering manager Alex asks Jane to prepare a model and a small code snippet showing how to run it; the internal systems team will handle the rest.

Fast-forward several months later: Jane gathered a small dataset and trained a model, and it all looked fine during the initial testing, so Alex’s team managed to wrap it in a service. But right before deployment, the product manager brought multiple new recipes that were added to the menu and said the model should be able to support those too. This was nothing complicated and involved just adding some more data, changing the labels map, and retraining the model—sounds like something that shouldn’t affect the deployment schedule much. However, after a discussion, Jane and Alex realized it would take several more months, even assuming the new dataset is readily available. What went wrong here? Jane performed all the required steps to train a model—manually validating datasets, applying numerous data processing and cleaning steps in the Jupyter Notebook environment, training the model with multiple interruptions, validating the result with the chef and customer happiness team on an ad hoc basis, uploading the trained model to the company’s shared storage, and sending the link back to Alex.

To make things more complicated, the model itself is usually a pipeline of another kind. For example, a simple logistic regression classifier is often enhanced with a feature scaling step, resulting in a pipeline of at least two steps. Usually, there’s also basic feature engineering (e.g., one-hot encoding for categorical variables), so even the simplest model has the properties of a pipeline. Other modalities like images, text, audio, etc., require additional preprocessing steps. For instance, a typical image classification model is a pipeline of image reading, normalization, resizing, and the model itself. If we switch to natural language processing, the pipeline almost always starts with text tokenization, etc. All in all, there’s a lot of space for ambiguity and confusion surrounding the term “pipeline” itself. To make things clearer, we will use the following terminology.

Training pipeline refers to a pipeline used to train a model. It’s a DAG of steps that takes the full dataset and, optionally, additional metadata as input and produces a trained model as output. It is a higher-level abstraction than the model itself (see figure 10.1).

Data fetching is the first step in the pipeline. It is responsible for downloading the data from the sources and making it available for the subsequent steps. As mentioned in chapter 6, we do not consider ourselves data engineering experts, so we will not discuss data fetching and storage in detail.

Preprocessing is usually the second step in the pipeline. It is a very generic term that can have a different meaning for a respective task. In general, preprocessing is a set of actions performed to prepare the data for model training. While we separate the training and inference pipelines for the sake of the book’s structure, the distinction can be somewhat blurry in practice. For example, you can fully preprocess the raw dataset before training the model or, alternatively, make it part of a single model inference. In this context, we are discussing training-specific preprocessing. Feature selection is one example of such preprocessing: we only do it before training and freeze the selected features for the subsequent steps. Model training is the core of the training pipeline. It is a step that takes preprocessed data and produces a trained model; it is usually the longest and most complex (especially in deep learning-based systems) step in the pipeline.

After the model has been trained, it can be evaluated and tested. These are different aspects aiming to answer the same question: how good is the model? Evaluation is a step of computing metrics, while tests are a set of checks performed to ensure that the model is working as expected.

Postprocessing is a step that is performed after evaluation and testing. It is a set of actions performed to prepare the model for deployment. Here, we can convert the model to a format supported by the target platforms, apply posttraining quantization or other optimizations if applicable, prepare tasks for human evaluation, and so on. It is worth noting that postprocessing and evaluation can be swapped. For example, we can evaluate the model before converting it to the target format, or alternatively, we can convert the model to the target format and then evaluate it using the deployment format.

Artifact packaging is the last step in the pipeline. It is responsible for packaging the model and other artifacts (e.g., config files with preprocessor parameters) into a format that can be easily deployed to production. The goal here is to simplify further deployment and separate the training and deployment pipelines. Ideally, the output should be as agnostic as possible to the training pipeline. For example, the model is exported to a universal format like ONNX for backend serving or CoreML for iOS serving, all config files are exported to a universal format like JSON, and any changes in the training pipeline should affect deployment as little as possible. Otherwise, the deployment pipeline will be tightly coupled with the training pipeline and will require many changes after each training pipeline update, becoming an obstacle for rapid model development and related experiments.

Reports are a special case of artifacts. It is a generic term that can be related to many things, including a basic table containing validation/test metrics, various types of error analysis (recall chapter 9), additional visualizations, and other auxiliary information. While these artifacts are not directly used for deployment, they’re crucial to consider the training successful; no responsible engineer can release a newly trained model without having at least a short look at proper reports. The only exception we have seen is a variation of an AutoML scenario when many new models are trained automatically per user requests. In this case, manual validation is not always possible; thus, engineers can only review suspicious outliers. We will touch on the topic of the release cycle in chapter 13.

Some of these artifacts are related to experiment tracking and reproducibility, and those are crucial for projects with multiple contributors or involved parties. When a researcher works alone on their own problem, they can track all their ideas and experiments using a simple notebook or a text file. However, when a team of researchers works on the same problem, they need a more structured way to track, compare, and reproduce their experiments. A centralized repository for all experiments is one tool that can help achieve this.

In this book, we don’t want to highlight any specific platform or tool, so we will not discuss them in detail. Given the pace of changes in the MLOps landscape, it’s very likely that our current understanding of the tools and platforms will be outdated by the time the book is published. Instead, we will focus on the principles and practices that are common to all platforms and tools. In the simplest case, you can implement a full training pipeline using generic non-ML tools—by creating a series of Python scripts connected with shell scripts, for example. However, this is rarely the case in practice: usually, there are some ML-specific tools that introduce abstractions and simplify the pipeline implementation. Most MLOps tools are “opinionated,” which means that using them strongly suggests a particular way of structuring your code. In the long run, this improves the training pipeline’s code consistency and makes it easier to maintain in the future (see chapter 16). Typical features required from the training pipeline platform are

It’s important to emphasize here that features don’t have to be covered by the same platform. For example, you can use a generic platform for running the pipeline and a custom tool for experiment tracking because market solutions don’t satisfy custom requirements. Sometimes it’s just a matter of cost optimization. Arseny worked in a company that used a hybrid of two tools just because one provided many useful features and an overall nice developer experience, while the other was integrated with a cloud provider with the cheapest GPUs for training. In this situation, it was reasonable to spend some time on integration and save a lot of money on training costs.

Choosing proper tools is determined by the problem scale and the infrastructure of the company. FAANG-level companies usually have their own ML platforms and tools that can operate on a proper scale, while smaller companies typically prefer a set of open-source tools and cloud services. Every tool has its own adoption cost, so it’s important to choose the proper tool for the right problem. To the best of our knowledge, there is no one-size-fits-all solution, unlike with many other more mature software engineering problems (see figure 10.3).

The current size of a dataset should not be confused with the size of a dataset expected to be used in the future. We may face a cold-start problem (see chapter 6), and having even thousands of samples may be a significant advantage for the initial phase of system development. However, in the future, it can grow by several orders of magnitude, and if you want to use all the data, you need to be able to handle it.

While there is no silver bullet for training models on huge datasets, there are two classic software engineering approaches to scaling. Those are vertical and horizontal scaling (see figure 10.4).

Horizontal scaling involves splitting the load between multiple machines. The first level of horizontal scaling is using multiple GPU machines; in this case, an ML engineer often needs to introduce small changes to the pipeline code, as most of the heavy lifting is already done by the training framework. However, this isn’t true horizontal scaling, as we’re still talking about a single machine. Genuine horizontal scaling involves using multiple machines and distributing the load among them. Nowadays, this type of scaling is often provided by frameworks as well, but it is more complex and usually requires more engineering effort during implementation. DeepSpeed by Microsoft, Accelerate by Hugging Face, and Horovod, originated in Uber, are some examples of such frameworks.

One ML-specific way of scaling is subsampling: if your dataset is overly huge, it could be reasonable to subsample it and reduce the required compute resources. The most straightforward way of subsampling is applicable to most problems and involves removing duplicated samples. However, there are more aggressive methods: downsampling near-duplicates (samples that are very similar based on a simple distance function; e.g., Levenshtein distance for strings) and downsampling based on an internal ID (e.g., not more than X samples per user) or an artificial ID (e.g., clusterizing the full dataset with a simple method and keeping not more than X samples per cluster).

In ML pipelines, scaling often requires changing other pipeline parameters: for example, batch size is limited by GPU memory, larger batches lead to faster convergence, and the learning rate schedule depends on both batch size and expected convergence schedule. So the ability to alter the parameters is important.

Underconfiguration means that the pipeline is not configurable enough, making it difficult to change the model’s architecture, dataset, preprocessing steps, etc. Things are hardcoded here and there in a convoluted way, and it’s hard to understand how the pipeline works or alter even the simplest aspects. This is a typical problem in the early stages of ML development. When the pipeline is small and simple, it’s easy to understand and change. Therefore, researchers without a software engineering background may find it unnecessary to introduce proper software abstractions, leading to the addition of more and more code without proper structure. This antipattern often occurs among such researchers.

Overconfiguration is not ideal either. A typical ML pipeline has many hyperparameters related to dataset processing, model architecture, feature engineering, and the training process. In reality, it’s hard to predict all the possible use cases and parameters that can be changed, and inexperienced developers may try to cover all the possible cases and introduce as many abstractions as possible. At some point, these additional abstraction layers only increase complexity. Note that in this section, we use “parameters of the training pipeline” and “hyperparameters of the model” interchangeably. Just as a reminder, hyperparameters are those parameters of the model that are not learned during the training process but are set by the user.

In the example of overconfigured code in the following listing, we see how a multilevel hierarchy of subconfigs may complicate things.

After preliminarily deciding which hyperparameters are tunable, it’s important to determine the tuning strategy. When computational resources are limited, handcrafted experiments are preferable. When resources are abundant, it makes sense to apply an automated hyperparameter tuning method, such as a straightforward random search or a more advanced Bayesian optimization. Tools for hyperparameter tuning (e.g., Hyperopt, Optuna, and scikit-optimize) can be part of the ML platform and may dictate how the configuration files should look.

From our experience, extensive hyperparameter tuning is more applicable for small datasets, where it’s possible to run numerous experiments in a reasonable time. When a single experiment takes weeks, it’s more practical to rely on the intuition of the ML engineer and run a few experiments manually. It is worth noting that experiments with smaller datasets may help build this intuition, although not every conclusion may generalize to a large training run.

It is important to find a proper way to config the training pipeline (see figure 10.5).

Proper tests serve three purposes:

A smoke test should be as fast as possible so you can run it on a small subset of the dataset, a small number of epochs, and maybe with a reduced version of the model. It should check that the pipeline runs without errors and produces reasonable output—for example, it ensures that the loss is decreasing on this toy dataset. The following listing shows a simplified example of a smoke test for a training pipeline.

In the context of an ML project, property-based testing can be used to ensure that the final trained model behaves as expected and satisfies certain properties. The following are some examples of properties that can be tested in an ML project:

An exact list of tests is not usually included in the design document; however, we recommend thinking about it in advance and mentioning it in the document. The design document is often used as a reference for implementation, so it’s useful to mention the tests in it.