Ask or search…

Write Custom Stack Component Flavors

How to write a custom stack component flavor
This is an older version of the ZenML documentation. To read and view the latest version please visit this up-to-date URL.
When building sophisticated ML workflows, you will often need to come up with custom-tailed solutions. Sometimes, this might even require you to use custom components for your infrastructure or tooling.
That is exactly why the stack components in ZenML were designed to be modular and straightforward to extend. Using ZenML's base abstractions, you can create your own stack component flavors and use custom solutions for any stack component.

Base Abstractions

Before we get into how custom stack component flavors can be defined, let us briefly discuss how ZenML's abstraction for stack components are designed.

Abstract Stack Component Base Abstraction

All stack components in ZenML inherit from a common abstract base class StackComponent, parts of which you can see below:
from abc import ABC
from pydantic import BaseModel, Field
from typing import ClassVar
from uuid import UUID, uuid4
from zenml.enums import StackComponentType
class StackComponent(BaseModel, ABC):
"""Abstract class for all components of a ZenML stack."""
# Instance configuration
name: str
uuid: UUID = Field(default_factory=uuid4)
# Class parameters
TYPE: ClassVar[StackComponentType]
FLAVOR: ClassVar[str]
There are a few things to unpack here. Let's talk about Pydantic first. Pydantic is a library for data validation and settings management. Using their BaseModel is helping us to configure and serialize these components while allowing us to add a validation layer to each stack component instance/implementation.
You can already see how that comes into play here within the base StackComponent implementation. As you can see, each instance of a StackComponent needs to include a name and an auto-generated uuid. These variables will be tracked when we serialize the stack component object. (You can exclude an instance configuration parameter from the serialization by giving it a name which starts with _.)
Moreover, you can use class variables by denoting them with the ClassVar[..], which are also excluded from the serialization. Each StackComponent implementation features two important class variables called the TYPE and the FLAVOR. The TYPE is utilized when we set up the base implementation for a specific type of stack component whereas the FLAVOR parameter is used to denote different flavors (which we will cover in the next section).

Component-Specific Base Abstraction

For each stack component, there then exists another component-specific base abstraction which all flavors should inherit from. These component-specific are themselves subclasses of StackComponent.
As an example, let us take a look at the BaseArtifactStore:
from typing import ClassVar, Set
from zenml.enums import StackComponentType
from zenml.stack import StackComponent
class BaseArtifactStore(StackComponent):
"""Abstract class for all ZenML artifact stores."""
# Instance configuration
path: str
# Class parameters
TYPE: ClassVar[StackComponentType] = StackComponentType.ARTIFACT_STORE
As you can see, the BaseArtifactStore sets the correct TYPE, while introducing a new instance variable called path and class variable called SUPPORTED_SCHEMES, which will be used by all the subclasses of this base implementation.

Building Custom Stack Component Flavors

In order to build a new flavor for a stack component, you need to create a custom class that inherits from the respective component-specific base abstraction, defines the FLAVOR class variable, and implements any abstract or flavor-specific methods and properties.
As an example, this is how you could define a custom artifact store:
from typing import ClassVar, Set
from zenml.artifact_stores import BaseArtifactStore
class MyCustomArtifactStore(BaseArtifactStore):
"""Custom artifact store implementation."""
# Class configuration
FLAVOR: ClassVar[str] = "custom" # the name you want your flavor to have
SUPPORTED_SCHEMES: ClassVar[Set[str]] = {"custom://"} # implement this
... # custom functionality
As you can see from the example above, MyCustomArtifactStore inherits from the corresponding base abstraction BaseArtifactStore and implements a custom flavor.
You could now register this custom artifact store via zenml artifact-store flavor register:
zenml artifact-store flavor register <path.to.MyCustomArtifacStore>
Afterwards, you should see the new custom artifact store in the list of available artifact store flavors:
zenml artifact-store flavor list
And that's it, you now have defined a custom stack component flavor that you can use in any of your stacks just like any other flavor you used before, e.g.:
zenml artifact-store register <ARTIFACT_STORE_NAME> \
--flavor=custom \
zenml stack register <STACK_NAME> \
--artifact-store <ARTIFACT_STORE_NAME> \
If your custom stack component flavor requires special setup before it can be used, check out the Managing Stack Component States section for more details.

Extending Specific Stack Components

If you would like to learn more about how to build a custom stack component flavor for a specific stack component, please check the links below:
Type of Stack Component
Orchestrating the runs of your pipeline
Storage for the artifacts created by your pipelines
Tracking the execution of your pipelines/steps
Store for your containers
Centralized location for the storage of your secrets
Execution of individual steps in specialized runtime environments
Services/platforms responsible for online model serving
Management of your data/features
Tracking your ML experiments
Sending alerts through specified channels