Implement a custom integration

Creating an external integration and contributing to ZenML

One of the main goals of ZenML is to find some semblance of order in the ever-growing MLOps landscape. ZenML already provides numerous integrations into many popular tools, and allows you to come up with ways to implement your own stack component flavors in order to fill in any gaps that are remaining.

However, what if you want to make your extension of ZenML part of the main codebase, to share it with others? If you are such a person, e.g., a tooling provider in the ML/MLOps space, or just want to contribute a tooling integration to ZenML, this guide is intended for you.

Step 1: Plan out your integration

In the previous page, we looked at the categories and abstractions that core ZenML defines. In order to create a new integration into ZenML, you would need to first find the categories that your integration belongs to. The list of categories can be found here as well.

Note that one integration may belong to different categories: For example, the cloud integrations (AWS/GCP/Azure) contain container registries, artifact stores etc.

Step 2: Create individual stack component flavors

Each category selected above would correspond to a stack component type. You can now start developing individual stack component flavors for this type by following the detailed instructions on the respective pages.

Before you package your new components into an integration, you may want to use/test them as a regular custom flavor. For instance, if you are developing a custom orchestrator and your flavor class MyOrchestratorFlavor is defined in flavors/my_flavor.py, you can register it by using:

zenml orchestrator flavor register flavors.my_flavor.MyOrchestratorFlavor