Develop a custom artifact store

Learning how to develop a custom artifact store.

Before diving into the specifics of this component type, it is beneficial to familiarize yourself with our general guide to writing custom component flavors in ZenML. This guide provides an essential understanding of ZenML's component flavor concepts.

ZenML comes equipped with Artifact Store implementations that you can use to store artifacts on a local filesystem or in the managed AWS, GCP, or Azure cloud object storage services. However, if you need to use a different type of object storage service as a backend for your ZenML Artifact Store, you can extend ZenML to provide your own custom Artifact Store implementation.

Base Abstraction

The Artifact Store establishes one of the main components in every ZenML stack. Now, let us take a deeper dive into the fundamentals behind its abstraction, namely the BaseArtifactStore class:

As ZenML only supports filesystem-based artifact stores, it features a configuration parameter called path, which will indicate the root path of the artifact store. When registering an artifact store, users will have to define this parameter.

Moreover, there is another variable in the config class called SUPPORTED_SCHEMES. This is a class variable that needs to be defined in every subclass of the base artifact store configuration. It indicates the supported file path schemes for the corresponding implementation. For instance, for the Azure artifact store, this set will be defined as {"abfs://", "az://"}.

Lastly, the base class features a set of abstractmethods: open, copyfile,exists,glob,isdir,listdir ,makedirs,mkdir,remove, rename,rmtree,stat,walk. In the implementation of every ArtifactStore flavor, it is required to define these methods with respect to the flavor at hand.

Putting all these considerations together, we end up with the following implementation:

from zenml.enums import StackComponentType