Building a Machine Learning Training Pipeline on Kubeflow
By: Pierre H. Arnoux
Tags: Tutorials, Scaling, Kubeflow, AI, Machine Learning

Setting up your Kubeflow work environment

Defining a convention to organize a codebase can greatly simplify a development process. It is even more true for orchestration repositories as they deal with a lot of moving parts.

With Kubeflow pipelines, we are handling two main categories of assets: pipelines and components. Pipelines are holding the orchestration logic while components are the containerized actions executed by the pipeline.

To emphasize this, and following some examples in the community, we organized our repository in two directories: pipelines and components.

Under components, one directory per component, with the following structure:

- Dockerfile - src

- main

- test - resources

We found this structure very convenient as it makes component self contained. Yet, it makes it harder to share assets between components.

Under pipelines, one python file per pipeline and several utils and config modules:

my_kubeflow_pipeline.py

- utils

- container_operations.py

- kubernetes_utils.py - config

- config_holder.py

- site_config.json

...

container_operations.py holds all the component classes for our project. We found it necessary to separate the class definitions from the pipelines to ensure readability. It also makes the components more reusable as it provides an interface outside the context of a pipeline.

kubernetes_utils.py holds all the common Kubernetes related functionalities (mounting a volume, setting a resource pool, etc.)

config_holder.py holds the config loading and validation mechanism. Having a site_config file saved us a lot of time as we could define configurations parameters in a single place.

Designing for scale

Example of a Kubeflow Pipeline following a Map-Reduce implementation schema

When dealing with large datasets, scaling is often achieved though parallelism. A pattern commonly used in the Big Data community is Map-Reduce (see Hadoop or Spark). While those frameworks offer great performances, they require the code to be written using their libraries. As explained above, we are dealing with code in many shapes in forms, so we cannot always leverage those technologies. Still Kubelow pipelines can be designed to follow the same general architecture.

In a Machine Learning context, Map-Reduce is best suited for pre-modeling steps. For example converting a dataset from a format to another, or generating features. At this stage, most operations are executed at a data point level, so they are easy to spread across multiple executors.

To allow data sharing across executors, we designed our components to execute all operation on Persistent Volumes Claims (PVC) shared by all executors. This also has the benefit of making all training assets persistent and recoverable by design.

Finally, we designed all our components to take as parameters a set of input directories, treated as read only, and a single output directory. With this design, it is very easy to chain one component after another by passing along the output directory as input to the next step. It also allows each component to be able to work interchangeably at a single chunk level or a dataset level. Lastly, a nice side effect is that the content generated by each step is preserved for debugging purposes.

The benefits of Test Driven Development

Test Driven Methodology

When developing components, we found Test Driven Development to be critical and a MAJOR time saver.

Indeed, as you implement the orchestration levels of your pipeline, the abstraction level make it harder to detect non breaking bugs. An indexing error may not break your code but it can completely corrupt your data.

Also, having a large scheduling system in your dev cycle makes feedback loops very long. Being able to get instant feedback from Unit Tests considerably speeds up development.

Yet, building unit tests for a multi stage pipeline with a variety of tools, input — output formats etc. can be challenging.

To improve our efficiency, we adopted few conventions:

For complex tasks, our components wrap codebases maintained in separate repositories. This separates functionality testing from orchestration testing. Testing components becomes a matter of making sure the code executes without error.

Our Unit Tests leverage the SetUp feature offered by most frameworks to consistently generate synthetic environments (directory structure, environment variables, synthetic data etc.). Sometimes, we simply zip the output of the Unit Tests from the previous step and use it as a starter environment for the next one.

Finally, a big time saver was to use multi staged docker builds to execute our unit tests. This makes the execution of unit tests consistent across all components. It also makes sure the tests are executed in the production environment (no more “I don’t understand, it worked on my local machine…”). And it makes production images cleaner as you can build them without adding the testing suite.

Here is an example of our Dockerfile: