zenml-nightly 0.58.2.dev20240623__py3-none-any.whl → 0.61.0.dev20240712__py3-none-any.whl

zenml/integrations/kubeflow/__init__.py

@@ -31,6 +31,9 @@ class KubeflowIntegration(Integration):
 
     NAME = KUBEFLOW
     REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]  # Only 1.x version that supports pyyaml 6
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by GCP as well
+    ]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/kubeflow/flavors/kubeflow_orchestrator_flavor.py

@@ -149,7 +149,7 @@ class KubeflowOrchestratorConfig(
 
     kubeflow_hostname: Optional[str] = None
     kubeflow_namespace: str = "kubeflow"
-    kubernetes_context: Optional[str]  # TODO: Potential setting
+    kubernetes_context: Optional[str] = None  # TODO: Potential setting
 
     @model_validator(mode="before")
     @classmethod

zenml/integrations/kubeflow/orchestrators/kubeflow_orchestrator.py

@@ -433,7 +433,6 @@ class KubeflowOrchestrator(ContainerizedOrchestrator):
             The dynamic container component.
         """
 
-        @dsl.container_component  # type: ignore[misc]
         def dynamic_container_component() -> dsl.ContainerSpec:
             """Dynamic container component.
 

zenml/integrations/kubernetes/__init__.py

@@ -31,7 +31,9 @@ class KubernetesIntegration(Integration):
 
     NAME = KUBERNETES
     REQUIREMENTS = ["kubernetes>=21.7,<26"]
-
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "kfp", # it is used by many others
+    ]
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:
         """Declare the stack component flavors for the Kubernetes integration.

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -225,8 +225,11 @@ def wait_pod(
             response = core_api.read_namespaced_pod_log(
                 name=pod_name,
                 namespace=namespace,
+                _preload_content=False,
             )
-            logs = response.splitlines()
+            raw_data = response.data
+            decoded_log = raw_data.decode("utf-8", errors="replace")
+            logs = decoded_log.splitlines()
             if len(logs) > logged_lines:
                 for line in logs[logged_lines:]:
                     logger.info(line)

zenml/integrations/label_studio/annotators/label_studio_annotator.py

@@ -527,6 +527,7 @@ class LabelStudioAnnotator(BaseAnnotator, AuthenticationMixin):
                 aws_access_key_id,
                 aws_secret_access_key,
                 aws_session_token,
+                _,
             ) = artifact_store.get_credentials()
 
             if aws_access_key_id and aws_secret_access_key:

zenml/integrations/langchain/__init__.py

@@ -30,6 +30,7 @@ class LangchainIntegration(Integration):
         "pyyaml>=6.0.1",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyyaml","tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/mlflow/__init__.py

@@ -33,7 +33,7 @@ class MlflowIntegration(Integration):
     NAME = MLFLOW
 
     REQUIREMENTS = [
-        "mlflow>=2.1.1,<=2.12.2",
+        "mlflow>=2.1.1,<=2.14.1",
         "mlserver>=1.3.3",
         "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed
@@ -45,9 +45,11 @@ class MlflowIntegration(Integration):
         # This downgrades pydantic to v1 even though mlflow does not have
         # any issues with v2. This is why we have to pin it here so a downgrade
         # will not happen.
-        "pydantic>=2.7.0,<2.8.0"
+        "pydantic>=2.7.0,<2.8.0",
     ]
 
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["python-rapidjson", "pydantic"]
+
     @classmethod
     def activate(cls) -> None:
         """Activate the MLflow integration."""

zenml/integrations/neural_prophet/__init__.py

@@ -26,6 +26,7 @@ class NeuralProphetIntegration(Integration):
         "holidays>=0.4.1,<0.25.0",
         "tenacity!=8.4.0",  # https://github.com/jd/tenacity/issues/471
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["tenacity"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/polars/__init__.py

@@ -25,6 +25,7 @@ class PolarsIntegration(Integration):
         "polars>=0.19.5",
         "pyarrow>=12.0.0"
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["pyarrow"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/prodigy/__init__.py

@@ -29,6 +29,7 @@ class ProdigyIntegration(Integration):
         "prodigy",
         "urllib3<2",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["urllib3"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/pycaret/__init__.py

@@ -28,6 +28,12 @@ class PyCaretIntegration(Integration):
         "catboost",
         "lightgbm",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = [
+        "scikit-learn",
+        "xgboost",
+        "catboost",
+        "lightgbm",
+    ]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/registry.py

@@ -124,6 +124,43 @@ class IntegrationRegistry(object):
                 )
             ]
 
+    def select_uninstall_requirements(
+        self,
+        integration_name: Optional[str] = None,
+        target_os: Optional[str] = None,
+    ) -> List[str]:
+        """Select the uninstall requirements for a given integration or all integrations.
+
+        Args:
+            integration_name: Name of the integration to check.
+            target_os: Target OS for the requirements.
+
+        Returns:
+            List of requirements for the integration uninstall.
+
+        Raises:
+            KeyError: If the integration is not found.
+        """
+        if integration_name:
+            if integration_name in self.list_integration_names:
+                return self._integrations[
+                    integration_name
+                ].get_uninstall_requirements(target_os=target_os)
+            else:
+                raise KeyError(
+                    f"Version {integration_name} does not exist. "
+                    f"Currently the following integrations are implemented. "
+                    f"{self.list_integration_names}"
+                )
+        else:
+            return [
+                requirement
+                for name in self.list_integration_names
+                for requirement in self._integrations[
+                    name
+                ].get_uninstall_requirements(target_os=target_os)
+            ]
+
     def is_installed(self, integration_name: Optional[str] = None) -> bool:
         """Checks if all requirements for an integration are installed.
 

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:
@@ -54,7 +122,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
 
     def get_credentials(
         self,
-    ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
         """Gets authentication credentials.
 
         If an authentication secret is configured, the secret values are
@@ -62,8 +130,8 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         attributes.
 
         Returns:
-            Tuple (key, secret, token) of credentials used to authenticate with
-            the S3 filesystem.
+            Tuple (key, secret, token, region) of credentials used to
+            authenticate with the S3 filesystem.
 
         Raises:
             RuntimeError: If the AWS connector behaves unexpectedly.
@@ -83,6 +151,7 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 credentials.access_key,
                 credentials.secret_key,
                 credentials.token,
+                client.meta.region_name,
             )
 
         secret = self.get_typed_authentication_secret(
@@ -93,12 +162,13 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
                 secret.aws_access_key_id,
                 secret.aws_secret_access_key,
                 secret.aws_session_token,
+                None,
             )
         else:
-            return self.config.key, self.config.secret, self.config.token
+            return self.config.key, self.config.secret, self.config.token, None
 
     @property
-    def filesystem(self) -> s3fs.S3FileSystem:
+    def filesystem(self) -> ZenMLS3Filesystem:
         """The s3 filesystem to access this artifact store.
 
         Returns:
@@ -108,18 +178,32 @@ class S3ArtifactStore(BaseArtifactStore, AuthenticationMixin):
         if self._filesystem and not self.connector_has_expired():
             return self._filesystem
 
-        key, secret, token = self.get_credentials()
+        key, secret, token, region = self.get_credentials()
+
+        # Use the region from the connector if available, otherwise some
+        # remote workloads (e.g. Sagemaker) might not work correctly because
+        # they look for the bucket in the wrong region
+        client_kwargs = {}
+        if region:
+            client_kwargs["region_name"] = region
+        if self.config.client_kwargs:
+            client_kwargs.update(self.config.client_kwargs)
 
-        self._filesystem = s3fs.S3FileSystem(
+        self._filesystem = ZenMLS3Filesystem(
             key=key,
             secret=secret,
             token=token,
-            client_kwargs=self.config.client_kwargs,
+            client_kwargs=client_kwargs,
             config_kwargs=self.config.config_kwargs,
             s3_additional_kwargs=self.config.s3_additional_kwargs,
         )
         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.
 

zenml/integrations/seldon/__init__.py

@@ -33,6 +33,7 @@ class SeldonIntegration(Integration):
     REQUIREMENTS = [
         "kubernetes==18.20.0",
     ]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kubernetes"]
 
     @classmethod
     def activate(cls) -> None:

zenml/integrations/seldon/model_deployers/seldon_model_deployer.py

@@ -290,6 +290,7 @@ class SeldonModelDeployer(BaseModelDeployer):
             assert isinstance(artifact_store, S3ArtifactStore)
 
             (
+                region_name,
                 aws_access_key_id,
                 aws_secret_access_key,
                 aws_session_token,

zenml/integrations/skypilot/flavors/skypilot_orchestrator_base_vm_config.py

@@ -43,7 +43,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
             `{'tpu_vm': True, 'runtime_version': 'tpu-vm-base'}` for TPUs.
         use_spot: whether to use spot instances. If None, defaults to
             False.
-        spot_recovery: the spot recovery strategy to use for the managed
+        job_recovery: the spot recovery strategy to use for the managed
             spot to recover the cluster from preemption. Refer to
             `recovery_strategy module <https://github.com/skypilot-org/skypilot/blob/master/sky/spot/recovery_strategy.py>`__ # pylint: disable=line-too-long
             for more details.
@@ -103,7 +103,7 @@ class SkypilotBaseOrchestratorSettings(BaseSettings):
     )
     accelerator_args: Optional[Dict[str, str]] = None
     use_spot: Optional[bool] = None
-    spot_recovery: Optional[str] = None
+    job_recovery: Optional[str] = None
     region: Optional[str] = None
     zone: Optional[str] = None
     image_id: Union[Dict[str, str], str, None] = Field(

zenml/integrations/skypilot/orchestrators/skypilot_base_vm_orchestrator.py

@@ -303,7 +303,7 @@ class SkypilotBaseOrchestrator(ContainerizedOrchestrator):
                     accelerators=settings.accelerators,
                     accelerator_args=settings.accelerator_args,
                     use_spot=settings.use_spot,
-                    spot_recovery=settings.spot_recovery,
+                    job_recovery=settings.job_recovery,
                     region=settings.region,
                     zone=settings.zone,
                     image_id=settings.image_id,

zenml/integrations/skypilot/orchestrators/skypilot_orchestrator_entrypoint.py

@@ -136,7 +136,7 @@ def main() -> None:
             settings.disk_size,  # Assuming disk_size is part of the settings
             settings.disk_tier,  # Assuming disk_tier is part of the settings
             settings.use_spot,
-            settings.spot_recovery,
+            settings.job_recovery,
             settings.region,
             settings.zone,
             accelerators_hashable,
@@ -213,7 +213,7 @@ def main() -> None:
                 accelerators=settings.accelerators,
                 accelerator_args=settings.accelerator_args,
                 use_spot=settings.use_spot,
-                spot_recovery=settings.spot_recovery,
+                job_recovery=settings.job_recovery,
                 region=settings.region,
                 zone=settings.zone,
                 image_id=settings.image_id,

zenml/integrations/skypilot_aws/__init__.py

@@ -31,7 +31,8 @@ class SkypilotAWSIntegration(Integration):
     """Definition of Skypilot AWS Integration for ZenML."""
 
     NAME = SKYPILOT_AWS
-    REQUIREMENTS = ["skypilot[aws]<=0.5.0"]
+    # all 0.6.x versions of skypilot[aws] are compatible
+    REQUIREMENTS = ["skypilot[aws]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/integrations/skypilot_azure/__init__.py

@@ -33,7 +33,7 @@ class SkypilotAzureIntegration(Integration):
     """Definition of Skypilot (Azure) Integration for ZenML."""
 
     NAME = SKYPILOT_AZURE
-    REQUIREMENTS = ["skypilot[azure]<=0.5.0"]
+    REQUIREMENTS = ["skypilot[azure]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/integrations/skypilot_gcp/__init__.py

@@ -31,7 +31,7 @@ class SkypilotGCPIntegration(Integration):
     """Definition of Skypilot (GCP) Integration for ZenML."""
 
     NAME = SKYPILOT_GCP
-    REQUIREMENTS = ["skypilot[gcp]<=0.5.0"]
+    REQUIREMENTS = ["skypilot[gcp]~=0.6.0"]
     APT_PACKAGES = ["openssh-client", "rsync"]
 
     @classmethod

zenml/integrations/skypilot_lambda/__init__.py

@@ -31,7 +31,7 @@ class SkypilotLambdaIntegration(Integration):
     """Definition of Skypilot Lambda Integration for ZenML."""
 
     NAME = SKYPILOT_LAMBDA
-    REQUIREMENTS = ["skypilot[lambda]<=0.5.0"]
+    REQUIREMENTS = ["skypilot[lambda]<=0.6.0"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/skypilot_lambda/flavors/skypilot_orchestrator_lambda_vm_flavor.py

@@ -40,7 +40,7 @@ class SkypilotLambdaOrchestratorSettings(SkypilotBaseOrchestratorSettings):
 
     _UNSUPPORTED_FEATURES = {
         "use_spot": "Spot instances not supported for Lambda orchestrator.",
-        "spot_recovery": "Spot recovery not supported for Lambda orchestrator.",
+        "job_recovery": "Job recovery not supported for Lambda orchestrator.",
         "image_id": "Custom image IDs not supported for Lambda orchestrator.",
         # Add other unsupported features as needed
     }

zenml/integrations/slack/__init__.py

@@ -31,6 +31,7 @@ class SlackIntegration(Integration):
 
     NAME = SLACK
     REQUIREMENTS = ["slack-sdk>=3.16.1", "aiohttp>=3.8.1"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["aiohttp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tekton/__init__.py

@@ -32,6 +32,7 @@ class TektonIntegration(Integration):
 
     NAME = TEKTON
     REQUIREMENTS = ["kfp>=2.6.0", "kfp-kubernetes>=1.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["kfp"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]:

zenml/integrations/tensorboard/__init__.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Initialization for TensorBoard integration."""
 
-import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORBOARD
 from zenml.integrations.integration import Integration

zenml/integrations/tensorflow/__init__.py

@@ -18,6 +18,9 @@ import sys
 from typing import List, Optional
 from zenml.integrations.constants import TENSORFLOW
 from zenml.integrations.integration import Integration
+from zenml.logger import get_logger
+
+logger = get_logger(__name__)
 
 
 class TensorflowIntegration(Integration):
@@ -25,20 +28,27 @@ class TensorflowIntegration(Integration):
 
     NAME = TENSORFLOW
     REQUIREMENTS = []
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["typing-extensions"]
 
     @classmethod
     def activate(cls) -> None:
         """Activates the integration."""
         # need to import this explicitly to load the Tensorflow file IO support
         # for S3 and other file systems
-        if (
-            not platform.system() == "Darwin"
-            or not platform.machine() == "arm64"
-        ):
+        if not platform.system() == "Darwin" or not platform.machine() == "arm64":
             import tensorflow_io  # type: ignore
 
         from zenml.integrations.tensorflow import materializers  # noqa
 
+        if sys.version_info.minor == 8:
+            logger.warning(
+                "Python 3.8 with TensorFlow is not fully "
+                "compatible with Pydantic 2 requirements. "
+                "Consider upgrading to a higher Python "
+                "version if you would like to use the "
+                "Tensorflow integration."
+            )
+
     @classmethod
     def get_requirements(cls, target_os: Optional[str] = None) -> List[str]:
         """Defines platform specific requirements for the integration.
@@ -52,13 +62,15 @@ class TensorflowIntegration(Integration):
         target_os = target_os or platform.system()
         if target_os == "Darwin" and platform.machine() == "arm64":
             requirements = [
-                f"tensorflow-macos>=2.12,<=2.15",
+                "tensorflow-macos>=2.12,<=2.15",
             ]
         else:
             requirements = [
-                f"tensorflow>=2.12,<=2.15",
+                "tensorflow>=2.12,<=2.15",
                 "tensorflow_io>=0.24.0",
             ]
+        if sys.version_info.minor == 8:
+            requirements.append("typing-extensions>=4.6.1")
         return requirements
 
 

zenml/integrations/wandb/__init__.py

@@ -31,6 +31,7 @@ class WandbIntegration(Integration):
 
     NAME = WANDB
     REQUIREMENTS = ["wandb>=0.12.12", "Pillow>=9.1.0"]
+    REQUIREMENTS_IGNORED_ON_UNINSTALL = ["Pillow"]
 
     @classmethod
     def flavors(cls) -> List[Type[Flavor]]: