zenml-nightly 0.58.2.dev20240626__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

@@ -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

@@ -122,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
@@ -130,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.
@@ -151,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(
@@ -161,9 +162,10 @@ 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) -> ZenMLS3Filesystem:
@@ -176,13 +178,22 @@ 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 = 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,
         )

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]]:

zenml/models/__init__.py

@@ -327,6 +327,7 @@ from zenml.models.v2.misc.service_connector_type import (
     ResourceTypeModel,
 )
 from zenml.models.v2.misc.server_models import ServerDatabaseType, ServerModel
+from zenml.models.v2.misc.full_stack import FullStackRequest
 from zenml.models.v2.core.trigger import (
     TriggerRequest,
     TriggerFilter,
@@ -386,6 +387,11 @@ from zenml.models.v2.core.server_settings import (
     ServerSettingsResponseMetadata,
     ServerSettingsUpdate,
 )
+from zenml.models.v2.misc.stack_deployment import (
+    DeployedStack,
+    StackDeploymentConfig,
+    StackDeploymentInfo,
+)
 
 # ----------------------------- Forward References -----------------------------
 
@@ -405,6 +411,7 @@ EventSourceResponseMetadata.model_rebuild()
 EventSourceResponseResources.model_rebuild()
 FlavorResponseBody.model_rebuild()
 FlavorResponseMetadata.model_rebuild()
+FullStackRequest.model_rebuild()
 LazyArtifactVersionResponse.model_rebuild()
 LazyRunMetadataResponse.model_rebuild()
 ModelResponseBody.model_rebuild()
@@ -701,11 +708,13 @@ __all__ = [
     "WorkspaceResponseMetadata",
     # V2 Misc
     "AuthenticationMethodModel",
+    "DeployedStack",
     "ServiceConnectorResourcesModel",
     "ServiceConnectorTypeModel",
     "ServiceConnectorTypedResourcesModel",
     "ServiceConnectorRequirements",
     "ResourceTypeModel",
+    "FullStackRequest",
     "UserAuthModel",
     "ExternalUserModel",
     "BuildItem",
@@ -718,6 +727,8 @@ __all__ = [
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",
+    "StackDeploymentConfig",
+    "StackDeploymentInfo",
     "OAuthDeviceAuthorizationRequest",
     "OAuthDeviceAuthorizationResponse",
     "OAuthDeviceTokenRequest",

zenml/models/v2/core/component.py

@@ -238,6 +238,24 @@ class ComponentResponse(
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the component labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        if self.labels is not None:
+            metadata.update(
+                {
+                    label[6:]: value
+                    for label, value in self.labels.items()
+                    if label.startswith("zenml:")
+                }
+            )
+        return metadata
+
     def get_hydrated_version(self) -> "ComponentResponse":
         """Get the hydrated version of this component.
 

zenml/models/v2/core/model.py

@@ -13,7 +13,6 @@
 #  permissions and limitations under the License.
 """Models representing models."""
 
-from functools import partial
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Union
 from uuid import UUID
 
@@ -303,7 +302,7 @@ class ModelResponse(
 
         client = Client()
         model_versions = depaginate(
-            partial(client.list_model_versions, model_name_or_id=self.id)
+            client.list_model_versions, model_name_or_id=self.id
         )
         return [
             mv.to_model_class(suppress_class_validation_warnings=True)

zenml/models/v2/core/service_connector.py

@@ -493,6 +493,23 @@ class ServiceConnectorResponse(
         max_length=STR_FIELD_MAX_LENGTH,
     )
 
+    def get_analytics_metadata(self) -> Dict[str, Any]:
+        """Add the service connector labels to analytics metadata.
+
+        Returns:
+            Dict of analytics metadata.
+        """
+        metadata = super().get_analytics_metadata()
+
+        metadata.update(
+            {
+                label[6:]: value
+                for label, value in self.labels.items()
+                if label.startswith("zenml:")
+            }
+        )
+        return metadata
+
     def get_hydrated_version(self) -> "ServiceConnectorResponse":
         """Get the hydrated version of this service connector.
 

zenml/models/v2/core/stack.py

@@ -36,6 +36,7 @@ from zenml.models.v2.core.component import ComponentResponse
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import ColumnElement
 
+
 # ------------------ Request Model ------------------
 
 
@@ -59,6 +60,10 @@ class StackRequest(WorkspaceScopedRequest):
         title="A mapping of stack component types to the actual"
         "instances of components of this type.",
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
     @property
     def is_valid(self) -> bool:
@@ -109,6 +114,10 @@ class StackUpdate(BaseUpdate):
         "instances of components of this type.",
         default=None,
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
 
 # ------------------ Response Model ------------------
@@ -134,6 +143,10 @@ class StackResponseMetadata(WorkspaceScopedResponseMetadata):
         default=None,
         title="The path to the stack spec used for mlstacks deployments.",
     )
+    labels: Optional[Dict[str, Any]] = Field(
+        default=None,
+        title="The stack labels.",
+    )
 
 
 class StackResponseResources(WorkspaceScopedResponseResources):
@@ -214,6 +227,15 @@ class StackResponse(
         """
         metadata = super().get_analytics_metadata()
         metadata.update({ct: c[0].flavor for ct, c in self.components.items()})
+
+        if self.labels is not None:
+            metadata.update(
+                {
+                    label[6:]: value
+                    for label, value in self.labels.items()
+                    if label.startswith("zenml:")
+                }
+            )
         return metadata
 
     @property
@@ -243,6 +265,15 @@ class StackResponse(
         """
         return self.get_metadata().components
 
+    @property
+    def labels(self) -> Optional[Dict[str, Any]]:
+        """The `labels` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().labels
+
 
 # ------------------ Filter Model ------------------