Machine Learning workflows with CPW: From complex pipelines to seamless execution

Author: Pedro Viana Gomes Pena

Reviewed by Kevin Rossell

Machine learning (ML) workflows are the backbone of data-driven innovation, but ensuring they are both scalable and adaptable can be a challenging task. That’s why we built Common Python Workflows (CPW), a versatile solution that simplifies how teams create, manage and execute machine learning pipelines.

CPW tackles common pain points, such as workflow complexity and lack of flexibility, empowering teams to focus on refining their models instead of battling inefficient processes. In this blog post, we’ll dive into the reasoning behind developing CPW, its standout features, and the key takeaways from our journey so far.

Initial challenges

Before CPW, we wrestled with several limitations with our existing tool, Sheep:

• Complexity and confusion: Over the years, Sheep became bloated with numerous functions, leading to a confusing user experience. Writing proper configuration files to govern training behavior was equally challenging, making it difficult for users to grasp its full potential..
• Lack of extensibility: Sheep was designed as a single Python library with all code and dependencies bundled together. This made it hard to extend or override certain functionalities, such as data loading. Users had little flexibility to change the pipeline: they could only override 3 pre-defined steps.
• Framework support: Supporting multiple frameworks (like Scikit-learn, Pytorch and Tensorflow) was not easy, and distributing training was cumbersome. Users had to train the model with fklearn.
• Out-of-core training: Users couldn’t perform out-of-core training efficiently.
• Scoring and evaluation: There was no abstraction for scoring, evaluating, or comparing models.
• Auto-retraining: Auto-retraining pipelines were not provided by default.
• Dependency management: Even minor changes required updating multiple dependencies in the models, which often led to dependency hell.

These challenges led us to develop CPW—a flexible, future-proof solution to streamline ML workflows and address these pain points.

Check our job opportunies

What is CPW?

CPW is a library designed to streamline the creation and management of ML workflows. It introduces several key features:

• Custom training workflows: Tailor workflows to meet specific needs.
• Inheriting and adapting training flows: Reuse workflows with minimal changes.
• Sharing flows across teams: Foster collaboration by enabling easy workflow sharing.
• Serialization/deserialization handling: Simplify the saving and loading of models and data.
• Auto retraining implementation: Implement auto retraining mechanisms easily.

To execute model pipelines, we use Kubeflow Pipelines—an open-source machine learning platform that makes it easy to deploy and manage end-to-end ML workflows. Besides orchestrating ML pipelines on Kubernetes, it provides tools for managing and monitoring pipeline runs.

CPW integrates seamlessly with Kubeflow through Argo Workflows. Once a Flow is authored, it can be compiled to an Argo Workflow and submitted for execution on Kubeflow, which provides features like:

• Run tasks in parallel, each in a different container.
• Hyperparameter Tuning with Katib (Kubeflow’s component for hyperparameter search).
• Allocating different machine resources for each step.

CPW is based on Dagster, leveraging mainly the concepts of software-defined assets and IOManagers.

Why We Chose Dagster

When designing CPW, we deliberately avoided re-implementing existing solutions for Direct Acyclic Graphs (DAG) and orchestration. There were several open-source frameworks that did that well. We wanted to focus on the pipeline authoring experience.

After evaluating open-source frameworks such as Metaflow and Dagster, we ultimately chose Dagster for its:

• Modular design: Dagster’s modular design allows for clear separation of concerns. Each step in a workflow can be defined independently, and dependencies are inferred based on the artifacts produced by each step. This contrasts with Metaflow, which required users to bundle together what a step would do with the description of which step(s) would follow.
• Software-defined assets: Dagster’s concept of software-defined assets aligns well with our goal of creating reusable and flexible workflows. This allows us to manage the components and artifacts of our workflows in a modular and reusable way.
• IOManagers: Dagster’s IOManagers provide a robust way to handle input and output operations, ensuring that data is managed efficiently and consistently across different steps of the workflow. And freeing our end user from that burden.
• Extensibility: Dagster’s design allows for greater and focused extensibility, enabling users to customize just the component they need to meet their specific needs, such as an IOManager for for pictures, which can be further reused and shared with other teams.

Workflow Class Design

To facilitate the creation of reusable workflows, we introduced a FlowSpec class. This class allows teams to create workflows that can be inherited and mutated by other teams. For example, a team can inherit a workflow (a DAG) and replace a step (a node) by overriding the appropriate attribute of the inherited workflow.

This approach has additional benefits:

• Integration with infrastructure: It creates a single entry point for our team to maintain and evolve. Once something changes in our infrastructure, we can update the base class and be sure that all workflows will adapt to the change.
• Ease of development: Developing new functionality became easier because we could rely on the modular concepts we had for Assets, IOManagers, and Flows (repositories, in Dagster).

As an example, since we introduced CPW, we have enabled the following functionalities for users:

• Different backend for execution: by changing a keyword in the asset definition, users can execute that step in Spark clusters in Databricks or in Sagemaker.
• Profiling: by passing a keyword in the command line, cProfiling is enabled for all steps in the Flow.

This allows for minimally intrusive changes that require little to no modification of the user’s code to be used.

This design also facilitated collaboration, with teams developing base training Flows to be used for multiple models, allowing improvements to be applied once across all models.

Core Concepts

CPW builds on several core concepts from Dagster, mainly software-defined assets and IOManagers. These concepts allow you to define and manage the components of your workflows in a modular and reusable way.

Here is an example of a Flow definition in CPW:

In this example, we define a simple workflow with three assets:

• params: A dictionary containing static definitions needed throughout the workflow, such as hyperparameters.
• my_dataset: A dataset that is split based on the parameters defined in params.
• train_model: A model trained using the dataset and parameters.

Based on the asset (function) names and arguments, CPW automatically establishes the dependencies between the functions. The resulting DAG (direct acyclic graph) will look like:

Artifacts and Persistence

CPW ensures that each asset’s artifacts are persisted in a shared S3 (or local) directory. For example:

• params: Stored as a JSON file.
• my_dataset: Saved as Parquet.
• train_model: Serialized and saved as a Pickle object.

CPW handles the serialization and deserialization of these artifacts, ensuring that they are stored and retrieved efficiently. This automatic persistence allows users to focus on defining their workflow without worrying about the underlying data management.

By leveraging Dagster’s software-defined assets and IOManagers, CPW provides a robust and flexible framework for building and managing ML workflows. This modular approach makes it easy to define, share, and adapt workflows, enabling teams to work more efficiently and effectively, by just focusing on what matters:

1. The function their Flow must perform, and;
2. The artifacts their Flow must produce.

Case Studies

CPW has enabled multiple teams to significantly streamline their workflows:

Sharing a same base flow for a multitude of NLP models

One of the most significant success stories comes from a team working on natural language processing (NLP) models. By sharing a common base Flow for various NLP models, the team was able to reduce the effort to update the training workflow of the models by 80%. 

This was achieved by reusing the same workflow structure and only modifying specific parameters for each model. The modular design of CPW allowed the team to efficiently manage and adapt their workflows, leading to substantial time savings and increased productivity.

Distributing improvements across teams

Another success story involves a data scientist who developed evaluation and chart assets for his model. Using CPW, she was able to immediately deploy these improvements to all other models within his team. This saved weeks of engineering effort and re-work, as the new assets could be seamlessly integrated into existing workflows. The ability to share and reuse workflow components across different models and teams highlights the power of CPW in promoting collaboration and efficiency.

These case studies demonstrate the tangible benefits of using CPW to manage ML workflows. By leveraging the modular and reusable nature of CPW, teams can significantly improve their workflow efficiency and reduce the time and effort required to develop and maintain ML models.

Lessons learned

Through the development and use of CPW, we identified several key lessons:

1. Balancing flexibility and simplicity: While CPW enabled complex use cases, it also introduced a new domain specific language (DSL), with a steeper learning curve. Providing simpler abstractions for basic workflows is critical for wider adoption. With straightforward entry points, we can make CPW more accessible to new users while still offering advanced capabilities for complex workflows.

2. Providing examples and templates: Real-world examples and comprehensive templates are essential to demonstrate CPW’s capabilities and ease adoption.

3. Identifying “champions” in each team: Identifying and supporting “champions” within teams fostered deeper adoption and reduced the need for constant guidance from our core team.

These lessons have informed our future roadmap, ensuring CPW continues to evolve and meet the needs of our users while maintaining its flexibility and power.