zenml-nightly 0.58.2.dev20240625__py3-none-any.whl → 0.60.0.dev20240627__py3-none-any.whl

README.md

@@ -306,7 +306,7 @@ the Apache License Version 2.0.
     <a href="https://github.com/zenml-io/zenml-projects">Projects Showcase</a>
     <br />
     <br />
-    🎉 Version 0.58.2 is out. Check out the release notes
+    🎉 Version 0.60.0 is out. Check out the release notes
     <a href="https://github.com/zenml-io/zenml/releases">here</a>.
     <br />
     🖥️ Download our VS Code Extension <a href="https://marketplace.visualstudio.com/items?itemName=ZenML.zenml-vscode">here</a>.

RELEASE_NOTES.md

@@ -1,5 +1,211 @@
 <!-- markdown-link-check-disable -->
 
+# 0.60.0
+
+ZenML now uses Pydantic v2. 🥳
+
+This upgrade comes with a set of critical updates. While your user experience
+mostly remains unaffected, you might see unexpected behavior due to the
+changes in our dependencies. Moreover, since Pydantic v2 provides a slightly
+stricter validation process, you might end up bumping into some validation
+errors which was not caught before, but it is all for the better 🙂 If
+you run into any other errors, please let us know either on
+[GitHub](https://github.com/zenml-io/zenml) or on
+our [Slack](https://zenml.io/slack-invite).
+
+## Changes in some of the critical dependencies
+
+- SQLModel is one of the core dependencies of ZenML and prior to this upgrade,
+  we were utilizing version `0.0.8`. However, this version is relatively
+  outdated and incompatible with Pydantic v2. Within the scope of this upgrade,
+  we upgraded it to `0.0.18`.
+- Due to the change in the SQLModel version, we also had to upgrade our
+  SQLAlchemy dependency from V1 to v2. While this does not affect the way
+  that you are using ZenML, if you are using SQLAlchemy in your environment,
+  you might have to migrate your code as well. For a detailed list of changes,
+  feel free to
+  check [their migration guide](https://docs.sqlalchemy.org/en/20/changelog/migration_20.html).
+
+## Changes in `pydantic`
+
+Pydantic v2 brings a lot of new and exciting changes to the table. The core
+logic now uses Rust, and it is much faster and more efficient in terms of
+performance. On top of it, the main concepts like model design, configuration,
+validation, or serialization now include a lot of new cool features. If you are
+using `pydantic` in your workflow and are interested in the new changes, you can
+check [the brilliant migration guide](https://docs.pydantic.dev/2.7/migration/)
+provided by the `pydantic` team to see the full list of changes.
+
+## Changes in our integrations changes
+
+Much like ZenML, `pydantic` is an important dependency in many other Python
+packages. That’s why conducting this upgrade helped us unlock a new version for
+several ZenML integration dependencies. Additionally, in some instances, we had
+to adapt the functionality of the integration to keep it compatible
+with `pydantic`. So, if you are using any of these integrations, please go
+through the changes.
+
+### Airflow
+
+As mentioned above upgrading our `pydantic` dependency meant we had to upgrade
+our `sqlmodel` dependency. Upgrading our `sqlmodel` dependency meant we had to
+upgrade our `sqlalchemy` dependency as well. Unfortunately, `apache-airflow`
+is still using `sqlalchemy` v1 and is incompatible with pydantic v2. As a
+solution, we have removed the dependencies of the `airflow` integration. Now,
+you can use ZenML to create your Airflow pipelines and use a separate
+environment to run them with Airflow. You can check the updated docs
+[right here](https://docs.zenml.io/stack-components/orchestrators/airflow).
+
+### AWS
+
+Some of our integrations now require `protobuf` 4. Since our
+previous `sagemaker` version (`2.117.0`) did not support `protobof` 4, we could
+not pair it with these new integrations. Thankfully `sagemaker` started
+supporting `protobuf` 4 with version `2.172.0` and relaxing its dependency
+solved the compatibility issue.
+
+### Evidently
+
+The old version of our `evidently` integration was not compatible with Pydantic
+v2. They started supporting it starting from version `0.4.16`. As their latest
+version is `0.4.22`, the new dependency of the integration is limited between
+these two versions.
+
+### Feast
+
+Our previous implementation of the `feast` integration was not compatible with
+Pydantic v2 due to the extra `redis` dependency we were using. This extra
+dependency is now removed and the `feast` integration is working as intended.
+
+### GCP
+
+The previous version of the Kubeflow dependency (`kfp==1.8.22`) in our GCP
+integration required Pydantic V1 to be installed. While we were upgrading our
+Pydantic dependency, we saw this as an opportunity and wanted to use this chance
+to upgrade the `kfp` dependency to v2 (which has no dependencies on the Pydantic
+library). This is why you may see some functional changes in the vertex step
+operator and orchestrator. If you would like to go through the changes in
+the `kfp` library, you can
+find [the migration guide here](https://www.kubeflow.org/docs/components/pipelines/v2/migration/).
+
+### Great Expectations
+
+Great Expectations started supporting Pydantic v2 starting from
+version `0.17.15` and they are closing in on their `1.0` release. Since this
+release might include a lot of big changes, we adjusted the dependency in our
+integration to `great-expectations>=0.17.15,<1.0`. We will try to keep it
+updated in the future once they release the `1.0` version
+
+### Kubeflow
+
+Similar to the GCP integration, the previous version of the kubeflow
+dependency (`kfp==1.8.22`) in our `kubeflow` integration required Pydantic V1 to
+be installed. While we were upgrading our Pydantic dependency, we saw this as an
+opportunity and wanted to use this chance to upgrade the `kfp` dependency to
+v2 (which has no dependencies on the Pydantic library). If you would like to go
+through the changes in the `kfp` library, you can
+find [the migration guide here](https://www.kubeflow.org/docs/components/pipelines/v2/migration/). (
+We also are considering adding an alternative version of this integration so our
+users can keep using `kfp` V1 in their environment. Stay tuned for any updates.)
+
+### MLflow
+
+`mlflow` is compatible with both Pydantic V1 and v2. However, due to a known
+issue, if you install `zenml` first and then
+do `zenml integration install mlflow -y`, it downgrades `pydantic` to V1. This
+is why we manually added the same duplicated `pydantic` requirement in the
+integration definition as well. Keep in mind that the `mlflow` library is still
+using some features of `pydantic` V1 which are deprecated. So, if the
+integration is installed in your environment, you might run into some
+deprecation warnings.
+
+### Label Studio
+
+While we were working on updating our `pydantic` dependency,
+the `label-studio-sdk` has released its 1.0 version. In this new
+version, `pydantic` v2 is also supported. The implementation and documentation
+of our Label Studio integration have been updated accordingly.
+
+### Skypilot
+
+With the switch to `pydantic` v2, the implementation of our `skypilot`
+integration mostly remained untouched. However, due to an incompatibility
+between the new version `pydantic` and the `azurecli`, the `skypilot[azure]`
+flavor can not be installed at the same time, thus our `skypilot_azure`
+integration is currently deactivated. We are working on fixing this issue and if
+you are using this integration in your workflows, we recommend staying on the
+previous version of ZenML until we can solve this issue.
+
+### Tensorflow
+
+The new version of `pydantic` creates a drift between `tensorflow`
+and  `typing_extensions`  packages and relaxing the dependencies here resolves
+the issue. At the same time, the upgrade to `kfp` v2 (in integrations
+like `kubeflow`, `tekton`, or `gcp`) bumps our `protobuf` dependency from `3.X`
+to `4.X`. To stay compatible with this requirement, the installed version
+of `tensorflow` needs to be `>=2.12.0`. While this change solves the dependency
+issues in most settings, we have bumped into some errors while
+using `tensorflow` 2.12.0 on Python 3.8 on Ubuntu. If you would like to use this
+integration, please consider using a higher Python version.
+
+### Tekton
+
+Similar to the `gcp` and `kubeflow` integrations, the old version of
+our `tekton` integration was not compatible with `pydantic` V1 due to its `kfp`
+dependency. With the switch from `kfp` V1 to v2, we have adapted our
+implementation to use the new version of `kfp` library and updated our
+documentation accordingly.
+
+## Additional Changes
+
+* We have also released a new version of `mlstacks` with Pydantic v2 support. 
+If you are using it in your development environment, you have to upgrade your 
+`mlstacks` package as well.
+* Added `zenml.integrations.huggingface.steps.run_with_accelerate` to enable running any step using [`accelerate`](https://huggingface.co/docs/accelerate/en/index). This function is supported by a utility that wraps any step function into a CLI script (which is required by most distributed training tools).
+* Fixed a memory leak that was observed while using the ZenML dashboard to view pipeline logs or artifact visualizations logged through an S3 Artifact Store linked to an AWS Service Connector.
+* Previously, we had an option called `build_options` that allowed users to pass arguments to the docker build command. However, these options were only applied when building the parent image. On macOS with ARM architecture, one needs to specify `platform=linux/amd64` to the build command to leverage local caching of Docker image layers. We have added a way to specify these build options for the "main" ZenML build as well, not just the parent image build. Additionally, users can now specify a `.dockerignore` file for the parent image build, which was previously not possible.
+
+## What's Changed
+
+* Extend migration testing by @avishniakov in https://github.com/zenml-io/zenml/pull/2768
+* Add retry docs by @htahir1 in https://github.com/zenml-io/zenml/pull/2770
+* Fix nightly Docker build by @strickvl in https://github.com/zenml-io/zenml/pull/2769
+* Start CTA and Cloud -> Pro renaming by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2773
+* Add star CTA to `README` by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2777
+* Use build python version if available by @schustmi in https://github.com/zenml-io/zenml/pull/2775
+* Introduced Legacy env var in docs by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2783
+* Fixing the nlp template for the upcoming pydantic upgrade by @bcdurak in https://github.com/zenml-io/zenml/pull/2778
+* Full renaming away from cloud to pro by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2782
+* Adjust docs url for flavors by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2772
+* Fixed broken unit test on develop and fixed duplicate / by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2785
+* Added timeout by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2786
+* Bump NLP template by @avishniakov in https://github.com/zenml-io/zenml/pull/2787
+* Raise error if Dockerfile does not exist by @schustmi in https://github.com/zenml-io/zenml/pull/2776
+* Pin `numpy<2.0.0` by @avishniakov in https://github.com/zenml-io/zenml/pull/2789
+* Fix partial logs loss in step operators with immutable FS in the backend by @avishniakov in https://github.com/zenml-io/zenml/pull/2788
+* Upgrading to `pydantic` v2 by @bcdurak in https://github.com/zenml-io/zenml/pull/2543
+* New CI/CD docs by @AlexejPenner in https://github.com/zenml-io/zenml/pull/2784
+* Improvements for running pipelines from the dashboard by @schustmi in https://github.com/zenml-io/zenml/pull/2781
+* Accelerate runner helper method by @avishniakov in https://github.com/zenml-io/zenml/pull/2746
+* Add `--ignore-errors` flag for `zenml artifact prune` by @strickvl in https://github.com/zenml-io/zenml/pull/2780
+* Enable running a pipeline through the client by @schustmi in https://github.com/zenml-io/zenml/pull/2736
+* Accelerated template LLMs by @avishniakov in https://github.com/zenml-io/zenml/pull/2797
+* Separate actions from triggers by @schustmi in https://github.com/zenml-io/zenml/pull/2700
+* Fix hook type definition and improve code completion for pipeline decorator by @schustmi in https://github.com/zenml-io/zenml/pull/2793
+* Allow specifying build options for main image build by @schustmi in https://github.com/zenml-io/zenml/pull/2749
+* Small improvements for yaml config files by @schustmi in https://github.com/zenml-io/zenml/pull/2796
+* Docs for the `pydantic` migration guide by @bcdurak in https://github.com/zenml-io/zenml/pull/2801
+* Bump mlflow to v2.14.1 by @christianversloot in https://github.com/zenml-io/zenml/pull/2779
+* Bugfix fixing the installation script to use the right mlstacks branch by @bcdurak in https://github.com/zenml-io/zenml/pull/2803
+* Fix S3 artifact store memory leak and other improvements by @stefannica in https://github.com/zenml-io/zenml/pull/2802
+
+## 🥳 Community Contributions 🥳
+
+We'd like to give a special thanks to @christianversloot who contributed to 
+this release by bumping the `mlflow` version to 2.14.1
+
+**Full Changelog**: https://github.com/zenml-io/zenml/compare/0.58.2...0.60.0
+
 # 0.58.2
 
 The 0.58.2 minor release is packed with a set of improvements to the ZenML logging and ZenML Server.

zenml/VERSION

@@ -1 +1 @@
-0.58.2.dev20240625
+0.60.0.dev20240627

zenml/artifact_stores/base_artifact_store.py

@@ -14,6 +14,7 @@
 """The base interface to extend the ZenML artifact store."""
 
 import inspect
+import os
 import textwrap
 from abc import abstractmethod
 from pathlib import Path
@@ -34,6 +35,7 @@ from typing import (
 
 from pydantic import model_validator
 
+from zenml.constants import ENV_ZENML_SERVER
 from zenml.enums import StackComponentType
 from zenml.exceptions import ArtifactStoreInterfaceError
 from zenml.io import fileio
@@ -433,7 +435,11 @@ class BaseArtifactStore(StackComponent):
             **kwargs: The keyword arguments to pass to the Pydantic object.
         """
         super(BaseArtifactStore, self).__init__(*args, **kwargs)
-        self._register()
+
+        # If running in a ZenML server environment, we don't register
+        # the filesystems. We always use the artifact stores directly.
+        if ENV_ZENML_SERVER not in os.environ:
+            self._register()
 
     def _register(self) -> None:
         """Create and register a filesystem within the filesystem registry."""

zenml/artifacts/utils.py

@@ -436,18 +436,21 @@ def load_artifact_visualization(
     artifact_store = _load_artifact_store(
         artifact_store_id=artifact.artifact_store_id, zen_store=zen_store
     )
-    mode = "rb" if visualization.type == VisualizationType.IMAGE else "r"
-    value = _load_file_from_artifact_store(
-        uri=visualization.uri,
-        artifact_store=artifact_store,
-        mode=mode,
-    )
+    try:
+        mode = "rb" if visualization.type == VisualizationType.IMAGE else "r"
+        value = _load_file_from_artifact_store(
+            uri=visualization.uri,
+            artifact_store=artifact_store,
+            mode=mode,
+        )
 
-    # Encode image visualizations if requested
-    if visualization.type == VisualizationType.IMAGE and encode_image:
-        value = base64.b64encode(bytes(value))
+        # Encode image visualizations if requested
+        if visualization.type == VisualizationType.IMAGE and encode_image:
+            value = base64.b64encode(bytes(value))
 
-    return LoadedVisualization(type=visualization.type, value=value)
+        return LoadedVisualization(type=visualization.type, value=value)
+    finally:
+        artifact_store.cleanup()
 
 
 def load_artifact_from_response(artifact: "ArtifactVersionResponse") -> Any:

zenml/cli/service_connectors.py

@@ -1109,6 +1109,7 @@ def describe_service_connector(
                 name_id_or_prefix=name_id_or_prefix,
                 resource_type=resource_type,
                 resource_id=resource_id,
+                verify=True,
             )
         except (
             KeyError,

zenml/client.py

@@ -5666,6 +5666,7 @@ class Client(metaclass=ClientMetaClass):
             name_id_or_prefix=name_id_or_prefix,
             resource_type=resource_type,
             resource_id=resource_id,
+            verify=False,
         )
 
         connector_client.configure_local_client(
@@ -5679,6 +5680,7 @@ class Client(metaclass=ClientMetaClass):
         name_id_or_prefix: Union[UUID, str],
         resource_type: Optional[str] = None,
         resource_id: Optional[str] = None,
+        verify: bool = False,
     ) -> "ServiceConnector":
         """Get the client side of a service connector instance to use with a local client.
 
@@ -5694,6 +5696,8 @@ class Client(metaclass=ClientMetaClass):
                 equivalent to the one requested, a `ValueError` exception is
                 raised. May be omitted for connectors and resource types that do
                 not support multiple resource instances.
+            verify: Whether to verify that the service connector configuration
+                and credentials can be used to gain access to the resource.
 
         Returns:
             The client side of the indicated service connector instance that can
@@ -5731,9 +5735,10 @@ class Client(metaclass=ClientMetaClass):
                 )
             )
 
-            # Verify the connector client on the local machine, because the
-            # server-side implementation may not be able to do so
-            connector_client.verify()
+            if verify:
+                # Verify the connector client on the local machine, because the
+                # server-side implementation may not be able to do so
+                connector_client.verify()
         else:
             connector_instance = (
                 service_connector_registry.instantiate_connector(

zenml/integrations/gcp/google_credentials_mixin.py

@@ -85,7 +85,7 @@ class GoogleCredentialsMixin(StackComponent):
                     "trying to use the linked connector, but got "
                     f"{type(credentials)}."
                 )
-            return credentials, connector.config.project_id
+            return credentials, connector.config.gcp_project_id
 
         if self.config.service_account_path:
             credentials, project_id = load_credentials_from_file(

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -351,24 +351,72 @@ class GCPOAuth2Token(AuthenticationConfig):
 class GCPBaseConfig(AuthenticationConfig):
     """GCP base configuration."""
 
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        This method must be implemented by subclasses to ensure that the GCP
+        project ID is always available.
+
+        Raises:
+            NotImplementedError: If the method is not implemented.
+        """
+        raise NotImplementedError
+
+
+class GCPBaseProjectIDConfig(GCPBaseConfig):
+    """GCP base configuration with included project ID."""
+
     project_id: str = Field(
         title="GCP Project ID where the target resource is located.",
     )
 
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        Returns:
+            The GCP project ID.
+        """
+        return self.project_id
 
-class GCPUserAccountConfig(GCPBaseConfig, GCPUserAccountCredentials):
+
+class GCPUserAccountConfig(GCPBaseProjectIDConfig, GCPUserAccountCredentials):
     """GCP user account configuration."""
 
 
 class GCPServiceAccountConfig(GCPBaseConfig, GCPServiceAccountCredentials):
     """GCP service account configuration."""
 
+    _project_id: Optional[str] = None
+
+    @property
+    def gcp_project_id(self) -> str:
+        """Get the GCP project ID.
+
+        When a service account JSON is provided, the project ID can be extracted
+        from it instead of being provided explicitly.
 
-class GCPExternalAccountConfig(GCPBaseConfig, GCPExternalAccountCredentials):
+        Returns:
+            The GCP project ID.
+        """
+        if self._project_id is None:
+            self._project_id = json.loads(
+                self.service_account_json.get_secret_value()
+            )["project_id"]
+            # Guaranteed by the field validator
+            assert self._project_id is not None
+
+        return self._project_id
+
+
+class GCPExternalAccountConfig(
+    GCPBaseProjectIDConfig, GCPExternalAccountCredentials
+):
     """GCP external account configuration."""
 
 
-class GCPOAuth2TokenConfig(GCPBaseConfig, GCPOAuth2Token):
+class GCPOAuth2TokenConfig(GCPBaseProjectIDConfig, GCPOAuth2Token):
     """GCP OAuth 2.0 configuration."""
 
     service_account_email: Optional[str] = Field(
@@ -540,7 +588,7 @@ resources in the specified project. When used remotely in a GCP workload, the
 configured project has to be the same as the project of the attached service
 account.
 """,
-            config_class=GCPBaseConfig,
+            config_class=GCPBaseProjectIDConfig,
         ),
         AuthenticationMethodModel(
             name="GCP User Account",
@@ -1006,6 +1054,7 @@ class GCPServiceConnector(ServiceConnector):
                 # service account authentication)
 
                 assert isinstance(cfg, GCPServiceAccountConfig)
+
                 credentials = (
                     gcp_service_account.Credentials.from_service_account_info(
                         json.loads(
@@ -1115,7 +1164,7 @@ class GCPServiceConnector(ServiceConnector):
         #
         # We need to extract the project ID and registry ID from
         # the provided resource ID
-        config_project_id = self.config.project_id
+        config_project_id = self.config.gcp_project_id
         project_id: Optional[str] = None
         # A GCR repository URI uses one of several hostnames (gcr.io, us.gcr.io,
         # eu.gcr.io, asia.gcr.io etc.) and the project ID is the first part of
@@ -1219,9 +1268,9 @@ class GCPServiceConnector(ServiceConnector):
                 authorized.
         """
         if resource_type == GCP_RESOURCE_TYPE:
-            return self.config.project_id
+            return self.config.gcp_project_id
         elif resource_type == DOCKER_REGISTRY_RESOURCE_TYPE:
-            return f"gcr.io/{self.config.project_id}"
+            return f"gcr.io/{self.config.gcp_project_id}"
 
         raise RuntimeError(
             f"Default resource ID not supported for '{resource_type}' resource "
@@ -1278,7 +1327,7 @@ class GCPServiceConnector(ServiceConnector):
 
             # Create an GCS client for the bucket
             client = storage.Client(
-                project=self.config.project_id, credentials=credentials
+                project=self.config.gcp_project_id, credentials=credentials
             )
             return client
 
@@ -1384,7 +1433,7 @@ class GCPServiceConnector(ServiceConnector):
                             "config",
                             "set",
                             "project",
-                            self.config.project_id,
+                            self.config.gcp_project_id,
                         ],
                         check=True,
                         stderr=subprocess.STDOUT,
@@ -1488,7 +1537,7 @@ class GCPServiceConnector(ServiceConnector):
             )
 
         if auth_method == GCPAuthenticationMethods.IMPLICIT:
-            auth_config = GCPBaseConfig(
+            auth_config = GCPBaseProjectIDConfig(
                 project_id=project_id,
             )
         elif auth_method == GCPAuthenticationMethods.OAUTH2_TOKEN:
@@ -1697,7 +1746,7 @@ class GCPServiceConnector(ServiceConnector):
 
         if resource_type == GCS_RESOURCE_TYPE:
             gcs_client = storage.Client(
-                project=self.config.project_id, credentials=credentials
+                project=self.config.gcp_project_id, credentials=credentials
             )
             if not resource_id:
                 # List all GCS buckets
@@ -1736,7 +1785,7 @@ class GCPServiceConnector(ServiceConnector):
             # List all GKE clusters
             try:
                 clusters = gke_client.list_clusters(
-                    parent=f"projects/{self.config.project_id}/locations/-"
+                    parent=f"projects/{self.config.gcp_project_id}/locations/-"
                 )
                 cluster_names = [cluster.name for cluster in clusters.clusters]
             except google.api_core.exceptions.GoogleAPIError as e:
@@ -1810,7 +1859,7 @@ class GCPServiceConnector(ServiceConnector):
             # object
             auth_method: str = GCPAuthenticationMethods.OAUTH2_TOKEN
             config: GCPBaseConfig = GCPOAuth2TokenConfig(
-                project_id=self.config.project_id,
+                project_id=self.config.gcp_project_id,
                 token=credentials.token,
                 service_account_email=credentials.signer_email
                 if hasattr(credentials, "signer_email")
@@ -1884,7 +1933,7 @@ class GCPServiceConnector(ServiceConnector):
             # List all GKE clusters
             try:
                 clusters = gke_client.list_clusters(
-                    parent=f"projects/{self.config.project_id}/locations/-"
+                    parent=f"projects/{self.config.gcp_project_id}/locations/-"
                 )
                 cluster_map = {
                     cluster.name: cluster for cluster in clusters.clusters
@@ -1928,7 +1977,7 @@ class GCPServiceConnector(ServiceConnector):
                 auth_method=KubernetesAuthenticationMethods.TOKEN,
                 resource_type=resource_type,
                 config=KubernetesTokenConfig(
-                    cluster_name=f"gke_{self.config.project_id}_{cluster_name}",
+                    cluster_name=f"gke_{self.config.gcp_project_id}_{cluster_name}",
                     certificate_authority=cluster_ca_cert,
                     server=f"https://{cluster_server}",
                     token=bearer_token,

zenml/integrations/s3/artifact_stores/s3_artifact_store.py

@@ -26,6 +26,7 @@ from typing import (
 )
 
 import s3fs
+from fsspec.asyn import FSTimeoutError, sync, sync_wrapper
 
 from zenml.artifact_stores import BaseArtifactStore
 from zenml.integrations.s3.flavors.s3_artifact_store_flavor import (
@@ -38,10 +39,77 @@ from zenml.stack.authentication_mixin import AuthenticationMixin
 PathType = Union[bytes, str]
 
 
+class ZenMLS3Filesystem(s3fs.S3FileSystem):  # type: ignore[misc]
+    """Modified s3fs.S3FileSystem to disable caching.
+
+    The original s3fs.S3FileSystem caches all class instances based on the
+    constructor input arguments and it never releases them. This is problematic
+    in the context of the ZenML server, because the server is a long-running
+    process that instantiates many S3 filesystems with different credentials,
+    especially when the credentials are generated by service connectors.
+
+    The caching behavior of s3fs causes the server to slowly consume more and
+    more memory over time until it crashes. This class disables the caching
+    behavior of s3fs by setting the `cachable` attribute to `False`.
+
+    In addition to disabling instance caching, this class also provides a
+    correct cleanup implementation by overriding the `close_session` method
+    the S3 aiobotocore client. The original one provided by s3fs was causing
+    memory leaks by creating a new event loop in the destructor instead of
+    using the existing one.
+
+    A `close` method is also provided to allow for synchronous on-demand cleanup
+    of the S3 client.
+    """
+
+    cachable = False
+
+    async def _close(self) -> None:
+        """Close the S3 client."""
+        if self._s3creator is not None:  # type: ignore[has-type]
+            await self._s3creator.__aexit__(None, None, None)  # type: ignore[has-type]
+            self._s3creator = None
+            self._s3 = None
+
+    close = sync_wrapper(_close)
+
+    @staticmethod
+    def close_session(loop: Any, s3: Any) -> None:
+        """Close the S3 client session.
+
+        Args:
+            loop: The event loop to use for closing the session.
+            s3: The S3 client to close.
+        """
+        # IMPORTANT: This method is a copy of the original close_session method
+        # from s3fs.S3FileSystem. The only difference is that it uses the
+        # provided event loop instead of creating a new one.
+        if loop is not None and loop.is_running():
+            try:
+                # NOTE: this is the line in the original method that causes
+                # the memory leak
+                # loop = asyncio.get_event_loop()
+                loop.create_task(s3.__aexit__(None, None, None))
+                return
+            except RuntimeError:
+                pass
+            try:
+                sync(loop, s3.__aexit__, None, None, None, timeout=0.1)
+                return
+            except FSTimeoutError:
+                pass
+        try:
+            # close the actual socket
+            s3._client._endpoint.http_session._connector._close()
+        except AttributeError:
+            # but during shutdown, it may have gone
+            pass
+
+
 class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
     """Artifact Store for S3 based artifacts."""
 
-    _filesystem: Optional[s3fs.S3FileSystem] = None
+    _filesystem: Optional[ZenMLS3Filesystem] = None
 
     @property
     def config(self) -> S3ArtifactStoreConfig:
@@ -98,7 +166,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
             return self.config.key, self.config.secret, self.config.token
 
     @property
-    def filesystem(self) -> s3fs.S3FileSystem:
+    def filesystem(self) -> ZenMLS3Filesystem:
         """The s3 filesystem to access this artifact store.
 
         Returns:
@@ -110,7 +178,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
 
         key, secret, token = self.get_credentials()
 
-        self._filesystem = s3fs.S3FileSystem(
+        self._filesystem = ZenMLS3Filesystem(
             key=key,
             secret=secret,
             token=token,
@@ -120,6 +188,11 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         )
         return self._filesystem
 
+    def cleanup(self) -> None:
+        """Close the filesystem."""
+        if self._filesystem:
+            self._filesystem.close()
+
     def open(self, path: PathType, mode: str = "r") -> Any:
         """Open a file at the given path.