Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Abstractions for artifact stores #474

Merged
merged 66 commits into from
Apr 6, 2022

Conversation

bcdurak
Copy link
Contributor

@bcdurak bcdurak commented Mar 23, 2022

Hello everyone!

Even though this PR will be mainly focused on the abstraction of artifact stores, there is a lot to unpack here. That's why I will try to separate it into different sections.

Abstraction for Artifact Stores

First off, let's start with artifact stores. Previously, our artifact stores looked exactly as:

class BaseArtifactStore(StackComponent, ABC):
    """Base class for all ZenML artifact stores.
    Attributes:
        path: The root path of the artifact store.
    """

    path: str

    @property
    def type(self) -> StackComponentType:
        """The component type."""
        return StackComponentType.ARTIFACT_STORE

    @property
    @abstractmethod
    def flavor(self) -> ArtifactStoreFlavor:
        """The artifact store flavor."""

You can see at a first glance, the only thing that our base implementation featured was a path attribute and that was the entire interface for our artifact stores. As we are strictly using filesystem-dependant artifact stores, we have decided to merge the concept of Filesystems and ArtifactStores together and create a new interface as a combination of both:

class BaseArtifactStore(StackComponent, ABC):
    """Base class for all ZenML artifact stores.
    Attributes:
        path: The root path of the artifact store.
    """

    ...
    
    @staticmethod
    def open(name: PathType, mode: str = "r") -> Any:
        """Open a file at the given path."""
        raise NotImplementedError()

    @staticmethod
    def copyfile(src: PathType, dst: PathType, overwrite: bool = False) -> None:
        """Copy a file from the source to the destination."""
        raise NotImplementedError()

    @staticmethod
    def exists(path: PathType) -> bool:
        """Returns `True` if the given path exists."""
        raise NotImplementedError()

    ...

We still needed to register a FileSystem, that's why the base class now has a _register method, which creates a proper FileSystem and registers it to tfx.dsl.io.filesystem_registry.DEFAULT_FILESYSTEM_REGISTRY upon the __init__ of the instance.

def _register(self, priority: int = 5) -> None:
        """Create and register a filesystem within the TFX registry"""
        from tfx.dsl.io.filesystem import Filesystem
        from tfx.dsl.io.filesystem_registry import DEFAULT_FILESYSTEM_REGISTRY

        filesystem_class = type(
            self.__class__.__name__,
            (Filesystem,),
            {
                "SUPPORTED_SCHEMES": self.SUPPORTED_SCHEMES,
                "open": staticmethod(_catch_not_found_error(self.open)),
                "copy": staticmethod(_catch_not_found_error(self.copyfile)),
                ...
            },
        )

        DEFAULT_FILESYSTEM_REGISTRY.register(
            filesystem_class, priority=priority
        )

Due to the internal functionalities of TFX, there was a need to catch the FileNotFoundErrors and
raise a from tfx.dsl.io.fileio.NotFoundErrors instead. Due to this issue, the implementations looked like this:

    @staticmethod
    def rmtree(path: PathType) -> None:
        """Remove the given directory."""
        try:
            fs.delete(path=path, recursive=True)
        except FileNotFoundError as e:
            raise NotFoundError() from e

    @staticmethod
    def stat(path: PathType) -> Dict[str, Any]:
        """Return stat info for the given path."""
        try:
            return fs.stat(path=path)  # type: ignore[no-any-return]
        except FileNotFoundError as e:
            raise NotFoundError() from e

This did not only create a lot of clutter in code but also put the responsibility of handling this on the person who is extending it. As a solution, I have created a decorator _catch_not_found_error which is now used as a wrapper around these methods when you create the corresponding Filesystem. (special thanks to @schustmi and @jwwwb for the idea)

fileio changes

For the sake of consistency and reduced complications, our implementations for filesystem and filesystem_registry have been removed. The remaining zenml.io.fileio and zenml.io.utils modules are now using the file systems which are registered in the tfx.dsl.io.filesystem_registry.DEFAULT_FILESYSTEM_REGISTRY (which also holds the Filesystems generated by our ArtifactStores)

Previously Implemented Artifact Stores

We had already built a few ArtifactStore implementations either as a part of the core or as a part of an integration. However, due to the recent changes in the base implementation, I needed to update the current implementations of the artifact stores. Since every ArtifactStore-Filesystem pair got merged, the plugins were removed and as a direct result, the activation of corresponding integrations got updated.

Also, as the concepts are now merged, the path attribute is now validated through a pydantic validator within the base implementation using the SUPPORTED_SCHEMES, thus ensure_XX_path validators have been removed.

Changes to the StackComponent

Working on the previously implemented artifact stores proved to be an excellent opportunity as I was able to experience the process of how a user can extend one of the base implementations of a stack component. With this in mind, I have made a few fundamental changes to the StackComponent:

  • Anyone who was extending any StackComponent had to define a supports_local_execution variable and a supports_remote_execution variable. Since they were not used anywhere, both of these variables have been removed to make it cleaner.

  • The type and flavor of the StackComponents have also been reworked as follows:

    @property
    @abstractmethod
    def type(self) -> StackComponentType:
        """The component type."""

    @property
    @abstractmethod
    def flavor(self) -> StackComponentFlavor:
        """The component flavor."""

to

    TYPE: ClassVar[StackComponentType]
    FLAVOR: ClassVar[str]
  • As a direct outcome of these changes, we can now treat these attributes as class variables and access them when we register the stack component.

Old version of extension:

@register_stack_component_class(
    component_type=StackComponentType.ARTIFACT_STORE, # previously we need to put these params here since they are properties, not class vars
    component_flavor=ArtifactStoreFlavor.GCP,
)
class GCPArtifactStore(BaseArtifactStore):
    """Artifact Store for Google Cloud Storage based artifacts."""

    supports_local_execution = True
    supports_remote_execution = True

    @property
    def flavor(self) -> ArtifactStoreFlavor:
        """The artifact store flavor."""
        return ArtifactStoreFlavor.GCP

New version of extension:

@register_stack_component_class
class GCPArtifactStore(BaseArtifactStore):
    """Artifact Store for Google Cloud Storage based artifacts."""

    # Class Configuration
    FLAVOR: ClassVar[str] = "gcp"
  • In order to have the same effect that the @abstractmethod has on the first implementation, the StackComponent now features a @root_validator which checks whether the TYPE and FLAVOR is defined when creating a new StackComponent.

  • For the sake of extensibility, the flavor enums have been removed and replaced with string values.

Open Questions

There are still a few questions:

  • When I implement a new artifact store which is not a part of an integration, how do I make the CLI aware of this new implementation when I want to register an instance of it?
  • How can I use my runtime configuration to parameterize a stack component?

TODOs

  • Updating the docs
  • More beautiful error handling
  • Add more tests

New Changes

  • S3ArtifactStore has been moved to a separate integration.

Pre-requisites

Please ensure you have done the following:

  • I have read the CONTRIBUTING.md document.
  • If my change requires a change to docs, I have updated the documentation accordingly.
  • If I have added an integration, I have updated the integrations table.
  • I have added tests to cover my changes.

Types of changes

  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Other (add details above)

Copy link
Contributor

@strickvl strickvl left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Some small changes / nits in here. Otherwise, really painstaking work you did! Looking forward to getting this merged in as soon as possible as I think it will affect some other branches currently under construction...

Copy link
Contributor

@schustmi schustmi left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Awesome, I especially love the nice error messages when defining an incorrect stack component/artifact store! Only got some very minor nits

@@ -66,7 +66,7 @@ def evaluator(
return test_acc


@pipeline(required_integrations=[SKLEARN])
@pipeline(required_integrations=[SKLEARN, S3])
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't this the general step operator example, in which case S3 could also be any other artifact store? In any case as long as a StackComponent (e.g. the artifact store here) is part of your active stack, it should be recognized automatically and there is no need to specify it manually

@@ -136,19 +123,19 @@ def _activate_integrations():
# getting the local orchestrator should not require activating integrations
StackComponentClassRegistry.get_class(
component_type=StackComponentType.ORCHESTRATOR,
component_flavor=OrchestratorFlavor.LOCAL,
component_flavor="local",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm wondering if we should at least have string constants somewhere for all the ZenML internal flavors, what do you think?

bcdurak and others added 5 commits April 5, 2022 16:33
Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com>
Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com>
Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com>
Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com>
Co-authored-by: Alex Strick van Linschoten <strickvl@users.noreply.github.com>
Copy link
Contributor

@jwwwb jwwwb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I love it! Always great to see user experience improved by simplifying the interface. I just left a few small suggestions, and a question or two for my own understanding. Excited to see this merged soon

Copy link
Contributor

@stefannica stefannica left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work redesigning the artifact store ! This was a real pleasure to review.

Copy link
Contributor

@AlexejPenner AlexejPenner left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Wow what a giant of a PR - great job Baris !!!

Copy link
Contributor

@jwwwb jwwwb left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Love what you've done with the ClassVars

@bcdurak bcdurak merged commit 191e2bb into develop Apr 6, 2022
@bcdurak bcdurak deleted the feature/ENG-616-artifact-store-abstractions branch April 6, 2022 09:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

7 participants