zenml-nightly 0.63.0.dev20240801__py3-none-any.whl → 0.64.0.dev20240809__py3-none-any.whl

zenml/integrations/gcp/service_connectors/gcp_service_connector.py

@@ -496,6 +496,108 @@ class GCPAuthenticationMethods(StrEnum):
     IMPERSONATION = "impersonation"
 
 
+try:
+    from google.auth.aws import _DefaultAwsSecurityCredentialsSupplier
+
+    class ZenMLAwsSecurityCredentialsSupplier(
+        _DefaultAwsSecurityCredentialsSupplier  # type: ignore[misc]
+    ):
+        """An improved version of the GCP external account credential supplier for AWS.
+
+        The original GCP external account credential supplier only provides
+        rudimentary support for extracting AWS credentials from environment
+        variables or the AWS metadata service. This version improves on that by
+        using the boto3 library itself (if available), which uses the entire range
+        of implicit authentication features packed into it.
+
+        Without this improvement, `sts.AssumeRoleWithWebIdentity` authentication is
+        not supported for EKS pods and the EC2 attached role credentials are
+        used instead (see: https://medium.com/@derek10cloud/gcp-workload-identity-federation-doesnt-yet-support-eks-irsa-in-aws-a3c71877671a).
+        """
+
+        def get_aws_security_credentials(
+            self, context: Any, request: Any
+        ) -> gcp_aws.AwsSecurityCredentials:
+            """Get the security credentials from the local environment.
+
+            This method is a copy of the original method from the
+            `google.auth.aws._DefaultAwsSecurityCredentialsSupplier` class. It has
+            been modified to use the boto3 library to extract the AWS credentials
+            from the local environment.
+
+            Args:
+                context: The context to use to get the security credentials.
+                request: The request to use to get the security credentials.
+
+            Returns:
+                The AWS temporary security credentials.
+            """
+            try:
+                import boto3
+
+                session = boto3.Session()
+                credentials = session.get_credentials()
+                if credentials is not None:
+                    creds = credentials.get_frozen_credentials()
+                    return gcp_aws.AwsSecurityCredentials(
+                        creds.access_key,
+                        creds.secret_key,
+                        creds.token,
+                    )
+            except ImportError:
+                pass
+
+            logger.debug(
+                "Failed to extract AWS credentials from the local environment "
+                "using the boto3 library. Falling back to the original "
+                "implementation."
+            )
+
+            return super().get_aws_security_credentials(context, request)
+
+        def get_aws_region(self, context: Any, request: Any) -> str:
+            """Get the AWS region from the local environment.
+
+            This method is a copy of the original method from the
+            `google.auth.aws._DefaultAwsSecurityCredentialsSupplier` class. It has
+            been modified to use the boto3 library to extract the AWS
+            region from the local environment.
+
+            Args:
+                context: The context to use to get the security credentials.
+                request: The request to use to get the security credentials.
+
+            Returns:
+                The AWS region.
+            """
+            try:
+                import boto3
+
+                session = boto3.Session()
+                if session.region_name:
+                    return session.region_name  # type: ignore[no-any-return]
+            except ImportError:
+                pass
+
+            logger.debug(
+                "Failed to extract AWS region from the local environment "
+                "using the boto3 library. Falling back to the original "
+                "implementation."
+            )
+
+            return super().get_aws_region(  # type: ignore[no-any-return]
+                context, request
+            )
+
+except ImportError:
+    # The `google.auth.aws._DefaultAwsSecurityCredentialsSupplier`
+    # class has been introduced in the `google-auth` library version 2.29.0.
+    # Before that, the AWS logic was part of the `google.auth.awsCredentials`
+    # class itself.
+    ZenMLAwsSecurityCredentialsSupplier = None  # type: ignore[assignment,misc]
+    pass
+
+
 class ZenMLGCPAWSExternalAccountCredentials(gcp_aws.Credentials):  # type: ignore[misc]
     """An improved version of the GCP external account credential for AWS.
 
@@ -508,6 +610,13 @@ class ZenMLGCPAWSExternalAccountCredentials(gcp_aws.Credentials):  # type: ignor
     Without this improvement, `sts.AssumeRoleWithWebIdentity` authentication is
     not supported for EKS pods and the EC2 attached role credentials are
     used instead (see: https://medium.com/@derek10cloud/gcp-workload-identity-federation-doesnt-yet-support-eks-irsa-in-aws-a3c71877671a).
+
+    IMPORTANT: subclassing this class only works with the `google-auth` library
+    version lower than 2.29.0. Starting from version 2.29.0, the AWS logic
+    has been moved to a separate `google.auth.aws._DefaultAwsSecurityCredentialsSupplier`
+    class that can be subclassed instead and supplied as the
+    `aws_security_credentials_supplier` parameter to the
+    `google.auth.aws.Credentials` class.
     """
 
     def _get_security_credentials(
@@ -539,12 +648,14 @@ class ZenMLGCPAWSExternalAccountCredentials(gcp_aws.Credentials):  # type: ignor
                     "secret_access_key": creds.secret_key,
                     "security_token": creds.token,
                 }
-        except Exception:
-            logger.debug(
-                "Failed to extract AWS credentials from the local environment "
-                "using the boto3 library. Falling back to the original "
-                "implementation."
-            )
+        except ImportError:
+            pass
+
+        logger.debug(
+            "Failed to extract AWS credentials from the local environment "
+            "using the boto3 library. Falling back to the original "
+            "implementation."
+        )
 
         return super()._get_security_credentials(  # type: ignore[no-any-return]
             request, imdsv2_session_token
@@ -1126,6 +1237,12 @@ class GCPServiceConnector(ServiceConnector):
                     account_info.get("subject_token_type")
                     == _AWS_SUBJECT_TOKEN_TYPE
                 ):
+                    if ZenMLAwsSecurityCredentialsSupplier is not None:
+                        account_info["aws_security_credentials_supplier"] = (
+                            ZenMLAwsSecurityCredentialsSupplier(
+                                account_info.pop("credential_source"),
+                            )
+                        )
                     credentials = (
                         ZenMLGCPAWSExternalAccountCredentials.from_info(
                             account_info,

zenml/integrations/kaniko/image_builders/kaniko_image_builder.py

@@ -295,7 +295,7 @@ class KanikoImageBuilder(BaseImageBuilder):
         logger.debug("Writing build context to process stdin.")
         assert process.stdin
         with process.stdin as _, tempfile.TemporaryFile(mode="w+b") as f:
-            build_context.write_archive(f, gzip=True)
+            build_context.write_archive(f, use_gzip=True)
             while True:
                 data = f.read(1024)
                 if not data:

zenml/integrations/mlflow/__init__.py

@@ -33,7 +33,7 @@ class MlflowIntegration(Integration):
     NAME = MLFLOW
 
     REQUIREMENTS = [
-        "mlflow>=2.1.1,<=2.14.2",
+        "mlflow>=2.1.1,<3",
         "mlserver>=1.3.3",
         "mlserver-mlflow>=1.3.3",
         # TODO: remove this requirement once rapidjson is fixed

zenml/integrations/mlflow/experiment_trackers/mlflow_experiment_tracker.py

@@ -57,6 +57,7 @@ DATABRICKS_HOST = "DATABRICKS_HOST"
 DATABRICKS_USERNAME = "DATABRICKS_USERNAME"
 DATABRICKS_PASSWORD = "DATABRICKS_PASSWORD"
 DATABRICKS_TOKEN = "DATABRICKS_TOKEN"
+DATABRICKS_UNITY_CATALOG = "databricks-uc"
 
 
 class MLFlowExperimentTracker(BaseExperimentTracker):
@@ -285,7 +286,6 @@ class MLFlowExperimentTracker(BaseExperimentTracker):
         """Configures the MLflow tracking URI and any additional credentials."""
         tracking_uri = self.get_tracking_uri()
         mlflow.set_tracking_uri(tracking_uri)
-        mlflow.set_registry_uri(tracking_uri)
 
         if is_databricks_tracking_uri(tracking_uri):
             if self.config.databricks_host:
@@ -296,6 +296,8 @@ class MLFlowExperimentTracker(BaseExperimentTracker):
                 os.environ[DATABRICKS_PASSWORD] = self.config.tracking_password
             if self.config.tracking_token:
                 os.environ[DATABRICKS_TOKEN] = self.config.tracking_token
+            if self.config.enable_unity_catalog:
+                mlflow.set_registry_uri(DATABRICKS_UNITY_CATALOG)
         else:
             os.environ[MLFLOW_TRACKING_URI] = tracking_uri
             if self.config.tracking_username:

zenml/integrations/mlflow/flavors/mlflow_experiment_tracker_flavor.py

@@ -98,6 +98,8 @@ class MLFlowExperimentTrackerConfig(
         databricks_host: The host of the Databricks workspace with the MLflow
             managed server to connect to. This is only required if
             `tracking_uri` value is set to `"databricks"`.
+        enable_unity_catalog: If `True`, will enable the Databricks Unity Catalog for
+            logging and registering models.
     """
 
     tracking_uri: Optional[str] = None
@@ -106,6 +108,7 @@ class MLFlowExperimentTrackerConfig(
     tracking_token: Optional[str] = SecretField(default=None)
     tracking_insecure_tls: bool = False
     databricks_host: Optional[str] = None
+    enable_unity_catalog: bool = False
 
     @model_validator(mode="after")
     def _ensure_authentication_if_necessary(

zenml/logger.py

@@ -46,6 +46,7 @@ class CustomFormatter(logging.Formatter):
     cyan: str = "\x1b[1;36m"
     bold_red: str = "\x1b[31;1m"
     purple: str = "\x1b[1;35m"
+    blue: str = "\x1b[34m"
     reset: str = "\x1b[0m"
 
     format_template: str = (
@@ -94,6 +95,18 @@ class CustomFormatter(logging.Formatter):
                     + quoted
                     + self.COLORS.get(LoggingLevels(record.levelno)),
                 )
+
+            # Format URLs
+            url_pattern = r"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+"
+            urls = re.findall(url_pattern, formatted_message)
+            for url in urls:
+                formatted_message = formatted_message.replace(
+                    url,
+                    self.reset
+                    + self.blue
+                    + url
+                    + self.COLORS.get(LoggingLevels(record.levelno)),
+                )
             return formatted_message
 
 

zenml/models/__init__.py

@@ -362,13 +362,6 @@ from zenml.models.v2.core.event_source import (
 from zenml.models.v2.misc.user_auth import UserAuthModel
 from zenml.models.v2.misc.build_item import BuildItem
 from zenml.models.v2.misc.loaded_visualization import LoadedVisualization
-from zenml.models.v2.misc.hub_plugin_models import (
-    HubPluginRequestModel,
-    HubPluginResponseModel,
-    HubUserResponseModel,
-    HubPluginBaseModel,
-    PluginStatus,
-)
 from zenml.models.v2.misc.external_user import ExternalUserModel
 from zenml.models.v2.misc.auth_models import (
     OAuthDeviceAuthorizationRequest,
@@ -735,11 +728,6 @@ __all__ = [
     "ExternalUserModel",
     "BuildItem",
     "LoadedVisualization",
-    "HubPluginRequestModel",
-    "HubPluginResponseModel",
-    "HubUserResponseModel",
-    "HubPluginBaseModel",
-    "PluginStatus",
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",

zenml/models/v2/core/pipeline_deployment.py

@@ -18,7 +18,6 @@ from uuid import UUID
 
 from pydantic import Field
 
-from zenml.config.docker_settings import SourceFileMode
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.config.pipeline_spec import PipelineSpec
 from zenml.config.step_configurations import Step
@@ -82,26 +81,14 @@ class PipelineDeploymentBase(BaseZenModel):
     )
 
     @property
-    def requires_included_files(self) -> bool:
-        """Whether the deployment requires included files.
+    def should_prevent_build_reuse(self) -> bool:
+        """Whether the deployment prevents a build reuse.
 
         Returns:
-            Whether the deployment requires included files.
+            Whether the deployment prevents a build reuse.
         """
         return any(
-            step.config.docker_settings.source_files == SourceFileMode.INCLUDE
-            for step in self.step_configurations.values()
-        )
-
-    @property
-    def requires_code_download(self) -> bool:
-        """Whether the deployment requires downloading some code files.
-
-        Returns:
-            Whether the deployment requires downloading some code files.
-        """
-        return any(
-            step.config.docker_settings.source_files == SourceFileMode.DOWNLOAD
+            step.config.docker_settings.prevent_build_reuse
             for step in self.step_configurations.values()
         )
 
@@ -125,6 +112,10 @@ class PipelineDeploymentRequest(
         default=None,
         title="The code reference associated with the deployment.",
     )
+    code_path: Optional[str] = Field(
+        default=None,
+        title="Optional path where the code is stored in the artifact store.",
+    )
     template: Optional[UUID] = Field(
         default=None,
         description="Template used for the deployment.",
@@ -169,6 +160,10 @@ class PipelineDeploymentResponseMetadata(WorkspaceScopedResponseMetadata):
     pipeline_spec: Optional[PipelineSpec] = Field(
         default=None, title="The pipeline spec of the deployment."
     )
+    code_path: Optional[str] = Field(
+        default=None,
+        title="Optional path where the code is stored in the artifact store.",
+    )
 
     pipeline: Optional[PipelineResponse] = Field(
         default=None, title="The pipeline associated with the deployment."
@@ -293,6 +288,15 @@ class PipelineDeploymentResponse(
         """
         return self.get_metadata().pipeline_spec
 
+    @property
+    def code_path(self) -> Optional[str]:
+        """The `code_path` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().code_path
+
     @property
     def pipeline(self) -> Optional[PipelineResponse]:
         """The `pipeline` property.
@@ -347,18 +351,6 @@ class PipelineDeploymentResponse(
         """
         return self.get_metadata().template_id
 
-    @property
-    def requires_code_download(self) -> bool:
-        """Whether the deployment requires downloading some code files.
-
-        Returns:
-            Whether the deployment requires downloading some code files.
-        """
-        return any(
-            step.config.docker_settings.source_files == SourceFileMode.DOWNLOAD
-            for step in self.step_configurations.values()
-        )
-
 
 # ------------------ Filter Model ------------------
 

zenml/models/v2/core/pipeline_run.py

@@ -206,6 +206,10 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         max_length=STR_FIELD_MAX_LENGTH,
         default=None,
     )
+    code_path: Optional[str] = Field(
+        default=None,
+        title="Optional path where the code is stored in the artifact store.",
+    )
     template_id: Optional[UUID] = Field(
         default=None,
         description="Template used for the pipeline run.",
@@ -425,6 +429,15 @@ class PipelineRunResponse(
         """
         return self.get_metadata().orchestrator_run_id
 
+    @property
+    def code_path(self) -> Optional[str]:
+        """The `code_path` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().code_path
+
     @property
     def template_id(self) -> Optional[UUID]:
         """The `template_id` property.

zenml/models/v2/core/server_settings.py

@@ -82,6 +82,9 @@ class ServerSettingsResponseBody(BaseResponseBody):
     display_updates: Optional[bool] = Field(
         title="Whether to display notifications about ZenML updates in the dashboard.",
     )
+    last_user_activity: datetime = Field(
+        title="The timestamp when the last user activity was detected.",
+    )
     updated: datetime = Field(
         title="The timestamp when this resource was last updated."
     )
@@ -179,6 +182,15 @@ class ServerSettingsResponse(
         """
         return self.get_body().active
 
+    @property
+    def last_user_activity(self) -> datetime:
+        """The `last_user_activity` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_body().last_user_activity
+
     @property
     def updated(self) -> datetime:
         """The `updated` property.

zenml/models/v2/core/user.py

@@ -65,12 +65,6 @@ class UserBase(BaseModel):
         description="`null` if not answered, `true` if agreed, "
         "`false` if skipped.",
     )
-    hub_token: Optional[str] = Field(
-        default=None,
-        title="JWT Token for the connected Hub account. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     password: Optional[str] = Field(
         default=None,
         title="A password for the user.",
@@ -296,12 +290,6 @@ class UserResponseMetadata(BaseResponseMetadata):
         "for user accounts.",
         max_length=STR_FIELD_MAX_LENGTH,
     )
-    hub_token: Optional[str] = Field(
-        default=None,
-        title="JWT Token for the connected Hub account. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
     external_user_id: Optional[UUID] = Field(
         default=None,
         title="The external user ID associated with the account. Only relevant "
@@ -416,15 +404,6 @@ class UserResponse(
         """
         return self.get_metadata().email
 
-    @property
-    def hub_token(self) -> Optional[str]:
-        """The `hub_token` property.
-
-        Returns:
-            the value of the property.
-        """
-        return self.get_metadata().hub_token
-
     @property
     def external_user_id(self) -> Optional[UUID]:
         """The `external_user_id` property.

zenml/models/v2/misc/server_models.py

@@ -13,7 +13,8 @@
 #  permissions and limitations under the License.
 """Model definitions for ZenML servers."""
 
-from typing import Dict
+from datetime import datetime
+from typing import Dict, Optional
 from uuid import UUID, uuid4
 
 from pydantic import BaseModel, Field
@@ -103,6 +104,11 @@ class ServerModel(BaseModel):
         title="Flag to indicate whether the server is using the legacy dashboard.",
     )
 
+    last_user_activity: Optional[datetime] = Field(
+        None,
+        title="Timestamp of latest user activity traced on the server.",
+    )
+
     def is_local(self) -> bool:
         """Return whether the server is running locally.
 

zenml/models/v2/misc/user_auth.py

@@ -73,13 +73,6 @@ class UserAuthModel(BaseZenModel):
         "`false` if skipped.",
     )
 
-    hub_token: Optional[str] = Field(
-        default=None,
-        title="JWT Token for the connected Hub account. Only relevant for user "
-        "accounts.",
-        max_length=STR_FIELD_MAX_LENGTH,
-    )
-
     @classmethod
     def _get_crypt_context(cls) -> "CryptContext":
         """Returns the password encryption context.