zenml-nightly 0.62.0.dev20240729__py3-none-any.whl → 0.64.0.dev20240809__py3-none-any.whl

zenml/integrations/databricks/orchestrators/databricks_orchestrator.py

@@ -52,6 +52,7 @@ from zenml.orchestrators.utils import get_orchestrator_run_name
 from zenml.orchestrators.wheeled_orchestrator import WheeledOrchestrator
 from zenml.stack import StackValidator
 from zenml.utils import io_utils
+from zenml.utils.package_utils import clean_requirements
 from zenml.utils.pipeline_docker_image_builder import (
     PipelineDockerImageBuilder,
 )
@@ -229,6 +230,9 @@ class DatabricksOrchestrator(WheeledOrchestrator):
             ValueError: If the schedule is not set or if the cron expression
                 is not set.
         """
+        settings = cast(
+            DatabricksOrchestratorSettings, self.get_settings(deployment)
+        )
         if deployment.schedule:
             if (
                 deployment.schedule.catchup
@@ -246,7 +250,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
                 )
             if (
                 deployment.schedule.cron_expression
-                and self.settings_class().schedule_timezone is None
+                and settings.schedule_timezone is None
             ):
                 raise ValueError(
                     "Property `schedule_timezone` must be set when passing "
@@ -321,7 +325,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
                     f"{deployment_id}_{step_name}",
                     ZENML_STEP_DEFAULT_ENTRYPOINT_COMMAND,
                     arguments,
-                    requirements,
+                    clean_requirements(requirements),
                     depends_on=upstream_steps,
                     zenml_project_wheel=zenml_project_wheel,
                     job_cluster_key=job_cluster_key,
@@ -366,7 +370,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
 
         # Construct the env variables for the pipeline
         env_vars = environment.copy()
-        spark_env_vars = self.settings_class().spark_env_vars
+        spark_env_vars = settings.spark_env_vars
         if spark_env_vars:
             for key, value in spark_env_vars.items():
                 env_vars[key] = value
@@ -385,6 +389,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
         job_cluster_key = self.sanitize_name(f"{deployment_id}")
         self._upload_and_run_pipeline(
             pipeline_name=orchestrator_run_name,
+            settings=settings,
             tasks=_construct_databricks_pipeline(
                 databricks_wheel_path, job_cluster_key
             ),
@@ -396,6 +401,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
     def _upload_and_run_pipeline(
         self,
         pipeline_name: str,
+        settings: DatabricksOrchestratorSettings,
         tasks: List[DatabricksTask],
         env_vars: Dict[str, str],
         job_cluster_key: str,
@@ -409,6 +415,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
             env_vars: The environment variables.
             job_cluster_key: The ID of the Databricks job cluster.
             schedule: The schedule to run the pipeline
+            settings: The settings for the Databricks orchestrator.
 
         Raises:
             ValueError: If the `Job Compute` policy is not found.
@@ -416,12 +423,12 @@ class DatabricksOrchestrator(WheeledOrchestrator):
 
         """
         databricks_client = self._get_databricks_client()
-        spark_conf = self.settings_class().spark_conf or {}
+        spark_conf = settings.spark_conf or {}
         spark_conf[
             "spark.databricks.driver.dbfsLibraryInstallationAllowed"
         ] = "true"
 
-        policy_id = self.settings_class().policy_id or None
+        policy_id = settings.policy_id or None
         for policy in databricks_client.cluster_policies.list():
             if policy.name == "Job Compute":
                 policy_id = policy.policy_id
@@ -432,17 +439,16 @@ class DatabricksOrchestrator(WheeledOrchestrator):
         job_cluster = JobCluster(
             job_cluster_key=job_cluster_key,
             new_cluster=ClusterSpec(
-                spark_version=self.settings_class().spark_version
+                spark_version=settings.spark_version
                 or DATABRICKS_SPARK_DEFAULT_VERSION,
-                num_workers=self.settings_class().num_workers,
-                node_type_id=self.settings_class().node_type_id
-                or "Standard_D4s_v5",
+                num_workers=settings.num_workers,
+                node_type_id=settings.node_type_id or "Standard_D4s_v5",
                 policy_id=policy_id,
                 autoscale=AutoScale(
-                    min_workers=self.settings_class().autoscale[0],
-                    max_workers=self.settings_class().autoscale[1],
+                    min_workers=settings.autoscale[0],
+                    max_workers=settings.autoscale[1],
                 ),
-                single_user_name=self.settings_class().single_user_name,
+                single_user_name=settings.single_user_name,
                 spark_env_vars=env_vars,
                 spark_conf=spark_conf,
                 workload_type=WorkloadType(
@@ -451,7 +457,7 @@ class DatabricksOrchestrator(WheeledOrchestrator):
             ),
         )
         if schedule and schedule.cron_expression:
-            schedule_timezone = self.settings_class().schedule_timezone
+            schedule_timezone = settings.schedule_timezone
             if schedule_timezone:
                 databricks_schedule = CronSchedule(
                     quartz_cron_expression=schedule.cron_expression,

zenml/integrations/gcp/orchestrators/vertex_orchestrator.py

@@ -574,6 +574,9 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
             run_name: Orchestrator run name.
             settings: Pipeline level settings for this orchestrator.
             schedule: The schedule the pipeline will run on.
+
+        Raises:
+            RuntimeError: If the Vertex Orchestrator fails to provision or any other Runtime errors
         """
         # We have to replace the hyphens in the run name with underscores
         # and lower case the string, because the Vertex AI Pipelines service
@@ -656,13 +659,15 @@ class VertexOrchestrator(ContainerizedOrchestrator, GoogleCredentialsMixin):
                     run.wait()
 
         except google_exceptions.ClientError as e:
-            logger.warning(
-                "Failed to create the Vertex AI Pipelines job: %s", e
+            logger.error("Failed to create the Vertex AI Pipelines job: %s", e)
+            raise RuntimeError(
+                f"Failed to create the Vertex AI Pipelines job: {e}"
             )
         except RuntimeError as e:
             logger.error(
                 "The Vertex AI Pipelines job execution has failed: %s", e
             )
+            raise
 
     def get_orchestrator_run_id(self) -> str:
         """Returns the active orchestrator run id.

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

@@ -199,12 +199,7 @@ from zenml.models.v2.core.pipeline import (
     PipelineResponse,
     PipelineResponseBody,
     PipelineResponseMetadata,
-)
-from zenml.models.v2.core.pipeline_namespace import (
-    PipelineNamespaceResponseBody,
-    PipelineNamespaceResponseMetadata,
-    PipelineNamespaceResponse,
-    PipelineNamespaceFilter,
+    PipelineResponseResources
 )
 from zenml.models.v2.core.pipeline_build import (
     PipelineBuildBase,
@@ -230,6 +225,16 @@ from zenml.models.v2.core.pipeline_run import (
     PipelineRunResponse,
     PipelineRunResponseBody,
     PipelineRunResponseMetadata,
+    PipelineRunResponseResources
+)
+from zenml.models.v2.core.run_template import (
+    RunTemplateRequest,
+    RunTemplateUpdate,
+    RunTemplateResponse,
+    RunTemplateResponseBody,
+    RunTemplateResponseMetadata,
+    RunTemplateResponseResources,
+    RunTemplateFilter,
 )
 from zenml.models.v2.base.base_plugin_flavor import BasePluginFlavorResponse
 from zenml.models.v2.core.run_metadata import (
@@ -357,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,
@@ -424,6 +422,7 @@ ModelVersionPipelineRunResponseBody.model_rebuild()
 OAuthDeviceResponseBody.model_rebuild()
 PipelineResponseBody.model_rebuild()
 PipelineResponseMetadata.model_rebuild()
+PipelineResponseResources.model_rebuild()
 PipelineBuildBase.model_rebuild()
 PipelineBuildResponseBody.model_rebuild()
 PipelineBuildResponseMetadata.model_rebuild()
@@ -433,6 +432,11 @@ PipelineDeploymentResponseMetadata.model_rebuild()
 PipelineDeploymentResponseResources.model_rebuild()
 PipelineRunResponseBody.model_rebuild()
 PipelineRunResponseMetadata.model_rebuild()
+PipelineRunResponseResources.model_rebuild()
+RunTemplateResponseBody.model_rebuild()
+RunTemplateResponseMetadata.model_rebuild()
+RunTemplateResponseResources.model_rebuild()
+RunTemplateResponseBody.model_rebuild()
 RunMetadataResponseBody.model_rebuild()
 RunMetadataResponseMetadata.model_rebuild()
 ScheduleResponseBody.model_rebuild()
@@ -590,10 +594,7 @@ __all__ = [
     "PipelineResponse",
     "PipelineResponseBody",
     "PipelineResponseMetadata",
-    "PipelineNamespaceFilter",
-    "PipelineNamespaceResponse",
-    "PipelineNamespaceResponseBody",
-    "PipelineNamespaceResponseMetadata",
+    "PipelineResponseResources",
     "PipelineBuildBase",
     "PipelineBuildRequest",
     "PipelineBuildFilter",
@@ -612,6 +613,14 @@ __all__ = [
     "PipelineRunResponse",
     "PipelineRunResponseBody",
     "PipelineRunResponseMetadata",
+    "PipelineRunResponseResources",
+    "RunTemplateRequest",
+    "RunTemplateUpdate",
+    "RunTemplateResponse",
+    "RunTemplateResponseBody",
+    "RunTemplateResponseMetadata",
+    "RunTemplateResponseResources",
+    "RunTemplateFilter",
     "RunMetadataRequest",
     "RunMetadataFilter",
     "RunMetadataResponse",
@@ -719,11 +728,6 @@ __all__ = [
     "ExternalUserModel",
     "BuildItem",
     "LoadedVisualization",
-    "HubPluginRequestModel",
-    "HubPluginResponseModel",
-    "HubUserResponseModel",
-    "HubPluginBaseModel",
-    "PluginStatus",
     "ServerModel",
     "ServerDatabaseType",
     "ServerDeploymentType",

zenml/models/v2/base/filter.py

@@ -36,6 +36,7 @@ from pydantic import (
     field_validator,
     model_validator,
 )
+from sqlalchemy import asc, desc
 from sqlmodel import SQLModel
 
 from zenml.constants import (
@@ -267,6 +268,7 @@ class BaseFilter(BaseModel):
         "size",
         "logical_operator",
     ]
+    CUSTOM_SORTING_OPTIONS: ClassVar[List[str]] = []
 
     # List of fields that are not even mentioned as options in the CLI.
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = []
@@ -352,6 +354,8 @@ class BaseFilter(BaseModel):
             )
         elif column in cls.model_fields:
             return value
+        elif column in cls.CUSTOM_SORTING_OPTIONS:
+            return value
         else:
             raise ValueError(
                 "You can only sort by valid fields of this resource"
@@ -861,3 +865,31 @@ class BaseFilter(BaseModel):
             query = query.where(filters)
 
         return query
+
+    def apply_sorting(
+        self,
+        query: AnyQuery,
+        table: Type["AnySchema"],
+    ) -> AnyQuery:
+        """Apply sorting to the query.
+
+        Args:
+            query: The query to which to apply the sorting.
+            table: The query table.
+
+        Returns:
+            The query with sorting applied.
+        """
+        column, operand = self.sorting_params
+
+        if operand == SorterOps.DESCENDING:
+            sort_clause = desc(getattr(table, column))  # type: ignore[var-annotated]
+        else:
+            sort_clause = asc(getattr(table, column))
+
+        # We always add the `id` column as a tiebreaker to ensure a stable,
+        # repeatable order of items, otherwise subsequent pages might contain
+        # the same items.
+        query = query.order_by(sort_clause, asc(table.id))  # type: ignore[arg-type]
+
+        return query