Develop a Custom Artifact Store
How to develop a custom artifact store
This is an older version of the ZenML documentation. To read and view the latest version please visit this up-to-date URL.
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 it is the base class for a specific type of StackComponent, it inherits from the StackComponent class. This sets the
TYPE
variable to a StackComponentType. TheFLAVOR
class variable needs to be set in the specific subclass.As ZenML only supports filesystem-based artifact stores, it features an instance configuration parameter called
path
, which will indicate the root path of the artifact store. When creating an instance of any flavor of anArtifactStore
, users will have to define this parameter.Moreover, there is an empty class variable called
SUPPORTED_SCHEMES
that needs to be defined in every flavor implementation. It indicates the supported filepath 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
abstractmethod
s:open
,copyfile
,exists
,glob
,isdir
,listdir
,makedirs
,mkdir
,remove
,rename
,rmtree
,stat
,walk
. In the implementation of everyArtifactStore
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:
This is a slimmed-down version of the base implementation which aims to highlight the abstraction layer. In order to see the full implementation and get the complete docstrings, please check the API docs.
The effect on the zenml.io.fileio
If you created an instance of an artifact store, added it to your stack and activated the stack, it will create a filesystem each time you run a ZenML pipeline and make it available to the zenml.io.fileio
module.
This means that when you utilize a method such as fileio.open(...)
with a filepath which starts with one of the SUPPORTED_SCHEMES
within your steps or materializers, it will be able to use the open(...)
method that you defined within your artifact store.
Build your own custom artifact store
If you want to implement your own custom Artifact Store, you can follow the following steps:
Create a class which inherits from the
BaseArtifactStore
base class.Define the
FLAVOR
andSUPPORTED_SCHEMES
class variables.Implement the
abstractmethod
s based on your desired filesystem.
Once you are done with the implementation, you can register it through the CLI as:
ZenML includes a range of Artifact Store implementations, some built-in and other provided by specific integration modules. You can use them as examples of how you can extend the base Artifact Store abstraction to implement your own custom Artifact Store:
Last updated