zenml-nightly 0.72.0.dev20250120__py3-none-any.whl → 0.73.0.dev20250123__py3-none-any.whl

zenml/VERSION

@@ -1 +1 @@
-0.72.0.dev20250120
+0.73.0.dev20250123

zenml/cli/service_connectors.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Service connector CLI commands."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Union, cast
 from uuid import UUID
 
@@ -291,7 +291,9 @@ def prompt_expires_at(
     while True:
         default_str = ""
         if default is not None:
-            seconds = int((default - datetime.utcnow()).total_seconds())
+            seconds = int(
+                (default - datetime.now(timezone.utc)).total_seconds()
+            )
             default_str = (
                 f" [{str(default)} i.e. in "
                 f"{seconds_to_human_readable(seconds)}]"
@@ -307,14 +309,16 @@ def prompt_expires_at(
 
         assert expires_at is not None
         assert isinstance(expires_at, datetime)
-        if expires_at < datetime.utcnow():
+        if expires_at < datetime.now(timezone.utc):
             cli_utils.warning(
                 "The expiration time must be in the future. Please enter a "
                 "later date and time."
             )
             continue
 
-        seconds = int((expires_at - datetime.utcnow()).total_seconds())
+        seconds = int(
+            (expires_at - datetime.now(timezone.utc)).total_seconds()
+        )
 
         confirm = click.confirm(
             f"Credentials will be valid until {str(expires_at)} UTC (i.e. "

zenml/cli/stack.py

@@ -17,7 +17,7 @@ import getpass
 import re
 import time
 import webbrowser
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -1575,7 +1575,7 @@ def deploy(
         ):
             raise click.Abort()
 
-        date_start = datetime.utcnow()
+        date_start = datetime.now(timezone.utc)
 
         webbrowser.open(deployment_config.deployment_url)
         console.print(

zenml/config/pipeline_configurations.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Pipeline configuration classes."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from pydantic import SerializeAsAny, field_validator
@@ -61,7 +61,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
             The full substitutions dict including date and time.
         """
         if start_time is None:
-            start_time = datetime.utcnow()
+            start_time = datetime.now(timezone.utc)
         ret = self.substitutions.copy()
         ret.setdefault("date", start_time.strftime("%Y_%m_%d"))
         ret.setdefault("time", start_time.strftime("%H_%M_%S_%f"))

zenml/config/server_config.py

@@ -432,6 +432,26 @@ class ServerConfiguration(BaseModel):
 
         return data
 
+    @field_validator("reportable_resources", mode="before")
+    @classmethod
+    def _convert_reportable_resources(cls, value: Any) -> Any:
+        """Convert reportable resources value if necessary.
+
+        This was previously set via an environment variable as a JSON-formatted
+        string. In case this still exists somewhere, this method converts the
+        JSON string to a list.
+
+        Args:
+            value: The reportable resources values.
+
+        Returns:
+            The potentially converted value.
+        """
+        if isinstance(value, str):
+            value = json.loads(value)
+
+        return value
+
     @property
     def deployment_id(self) -> UUID:
         """Get the ZenML server deployment ID.

zenml/constants.py

@@ -174,6 +174,11 @@ ENV_ZENML_PIPELINE_RUN_API_TOKEN_EXPIRATION = (
     "ZENML_PIPELINE_API_TOKEN_EXPIRATION"
 )
 
+# Materializer environment variables
+ENV_ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS = (
+    "ZENML_MATERIALIZER_ALLOW_NON_ASCII_JSON_DUMPS"
+)
+
 # ZenML Server environment variables
 ENV_ZENML_SERVER_PREFIX = "ZENML_SERVER_"
 ENV_ZENML_SERVER_PRO_PREFIX = "ZENML_SERVER_PRO_"

zenml/event_hub/base_event_hub.py

@@ -14,7 +14,7 @@
 """Base class for event hub implementations."""
 
 from abc import ABC, abstractmethod
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Tuple
 
 from zenml import EventSourceResponse
@@ -134,7 +134,7 @@ class BaseEventHub(ABC):
         )
         expires: Optional[datetime] = None
         if trigger.action.auth_window:
-            expires = datetime.utcnow() + timedelta(
+            expires = datetime.now(timezone.utc) + timedelta(
                 minutes=trigger.action.auth_window
             )
         encoded_token = token.encode(expires=expires)

zenml/integrations/airflow/orchestrators/airflow_orchestrator.py

@@ -408,7 +408,8 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
         if schedule:
             if schedule.cron_expression:
                 start_time = schedule.start_time or (
-                    datetime.datetime.utcnow() - datetime.timedelta(7)
+                    datetime.datetime.now(datetime.timezone.utc)
+                    - datetime.timedelta(7)
                 )
                 return {
                     "schedule": schedule.cron_expression,
@@ -428,6 +429,7 @@ class AirflowOrchestrator(ContainerizedOrchestrator):
             "schedule": "@once",
             # set a start time in the past and disable catchup so airflow
             # runs the dag immediately
-            "start_date": datetime.datetime.utcnow() - datetime.timedelta(7),
+            "start_date": datetime.datetime.now(datetime.timezone.utc)
+            - datetime.timedelta(7),
             "catchup": False,
         }

zenml/integrations/airflow/orchestrators/dag_generator.py

@@ -119,6 +119,22 @@ def get_operator_init_kwargs(
     except ImportError:
         pass
 
+    try:
+        # Support for apache-airflow-providers-cncf-kubernetes>=10.0.0 where
+        # the import changed
+        from airflow.providers.cncf.kubernetes.operators.pod import (
+            KubernetesPodOperator,
+        )
+
+        if issubclass(operator_class, KubernetesPodOperator):
+            init_kwargs.update(
+                get_kubernetes_pod_operator_init_kwargs(
+                    dag_config=dag_config, task_config=task_config
+                )
+            )
+    except ImportError:
+        pass
+
     init_kwargs.update(task_config.operator_args)
     return init_kwargs
 

zenml/integrations/gcp/__init__.py

@@ -30,6 +30,7 @@ from zenml.stack import Flavor
 
 GCP_ARTIFACT_STORE_FLAVOR = "gcp"
 GCP_IMAGE_BUILDER_FLAVOR = "gcp"
+GCP_VERTEX_EXPERIMENT_TRACKER_FLAVOR = "vertex"
 GCP_VERTEX_ORCHESTRATOR_FLAVOR = "vertex"
 GCP_VERTEX_STEP_OPERATOR_FLAVOR = "vertex"
 
@@ -70,6 +71,7 @@ class GcpIntegration(Integration):
         from zenml.integrations.gcp.flavors import (
             GCPArtifactStoreFlavor,
             GCPImageBuilderFlavor,
+            VertexExperimentTrackerFlavor,
             VertexOrchestratorFlavor,
             VertexStepOperatorFlavor,
         )
@@ -77,6 +79,7 @@ class GcpIntegration(Integration):
         return [
             GCPArtifactStoreFlavor,
             GCPImageBuilderFlavor,
+            VertexExperimentTrackerFlavor,
             VertexOrchestratorFlavor,
             VertexStepOperatorFlavor,
         ]

zenml/integrations/gcp/experiment_trackers/__init__.py

@@ -0,0 +1,18 @@
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Initialization for the VertexAI experiment tracker."""
+
+from zenml.integrations.gcp.experiment_trackers.vertex_experiment_tracker import (  # noqa
+    VertexExperimentTracker,
+)
+
+__all__ = ["VertexExperimentTracker"]

zenml/integrations/gcp/experiment_trackers/vertex_experiment_tracker.py

@@ -0,0 +1,214 @@
+#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Implementation of the VertexAI experiment tracker for ZenML."""
+
+import re
+from typing import TYPE_CHECKING, Dict, Optional, Type, cast
+
+from google.api_core import exceptions
+from google.cloud import aiplatform
+from google.cloud.aiplatform.compat.types import execution
+
+from zenml.constants import METADATA_EXPERIMENT_TRACKER_URL
+from zenml.experiment_trackers.base_experiment_tracker import (
+    BaseExperimentTracker,
+)
+from zenml.integrations.gcp.flavors.vertex_experiment_tracker_flavor import (
+    VertexExperimentTrackerConfig,
+    VertexExperimentTrackerSettings,
+)
+from zenml.integrations.gcp.google_credentials_mixin import (
+    GoogleCredentialsMixin,
+)
+from zenml.logger import get_logger
+from zenml.metadata.metadata_types import Uri
+
+if TYPE_CHECKING:
+    from zenml.config.step_run_info import StepRunInfo
+    from zenml.metadata.metadata_types import MetadataType
+
+logger = get_logger(__name__)
+
+
+class VertexExperimentTracker(BaseExperimentTracker, GoogleCredentialsMixin):
+    """Track experiments using VertexAI."""
+
+    @property
+    def config(self) -> VertexExperimentTrackerConfig:
+        """Returns the `VertexExperimentTrackerConfig` config.
+
+        Returns:
+            The configuration.
+        """
+        return cast(VertexExperimentTrackerConfig, self._config)
+
+    @property
+    def settings_class(self) -> Type[VertexExperimentTrackerSettings]:
+        """Returns the `BaseSettings` settings class.
+
+        Returns:
+            The settings class.
+        """
+        return VertexExperimentTrackerSettings
+
+    def prepare_step_run(self, info: "StepRunInfo") -> None:
+        """Configures a VertexAI run.
+
+        Args:
+            info: Info about the step that will be executed.
+        """
+        self._initialize_vertex(info=info)
+        self.experiment_name = self._get_experiment_name(info=info)
+        self.run_name = self._get_run_name(info=info)
+
+    def get_step_run_metadata(
+        self, info: "StepRunInfo"
+    ) -> Dict[str, "MetadataType"]:
+        """Get component- and step-specific metadata after a step ran.
+
+        Args:
+            info: Info about the step that was executed.
+
+        Returns:
+            A dictionary of metadata.
+        """
+        experiment_name = self._get_experiment_name(info=info)
+        run_name = self._get_run_name(info=info)
+        tensorboard_resource_name = self._get_tensorboard_resource_name(
+            experiment=experiment_name
+        )
+        dashboard_url = self._get_dashboard_url(experiment=experiment_name)
+        return {
+            METADATA_EXPERIMENT_TRACKER_URL: Uri(dashboard_url),
+            "tensorboard_resource_name": tensorboard_resource_name or "",
+            "vertex_run_name": run_name,
+        }
+
+    def _format_name(self, name: str) -> str:
+        return re.sub(r"[^a-z0-9-]", "-", name.strip().lower())[:128].rstrip(
+            "-"
+        )
+
+    def _get_experiment_name(self, info: "StepRunInfo") -> str:
+        """Gets the experiment name.
+
+        Args:
+            info: Info about the step.
+
+        Returns:
+            The experiment name.
+        """
+        settings = cast(
+            VertexExperimentTrackerSettings, self.get_settings(info)
+        )
+        name = settings.experiment or info.pipeline.name
+        return self._format_name(name)
+
+    def _get_run_name(self, info: "StepRunInfo") -> str:
+        """Gets the run name.
+
+        Args:
+            info: Info about the step that will be executed.
+
+        Returns:
+            The run name.
+        """
+        return self._format_name(info.run_name)
+
+    def _get_dashboard_url(self, experiment: str) -> str:
+        """Gets the run URL.
+
+        Args:
+            experiment: The name of the experiment.
+
+        Returns:
+            The run URL.
+        """
+        resource = aiplatform.Experiment(experiment_name=experiment)
+        return cast(str, resource.dashboard_url)
+
+    def _get_tensorboard_resource_name(self, experiment: str) -> Optional[str]:
+        resource = aiplatform.Experiment(
+            experiment_name=experiment
+        ).get_backing_tensorboard_resource()
+        resource_name = (
+            str(resource.resource_name) if resource is not None else None
+        )
+        return resource_name
+
+    def _initialize_vertex(self, info: "StepRunInfo") -> None:
+        """Initializes a VertexAI run.
+
+        Args:
+            info: Info about the step that will be executed.
+        """
+        settings = cast(
+            VertexExperimentTrackerSettings, self.get_settings(info)
+        )
+        experiment = self._get_experiment_name(info=info)
+        run_name = self._get_run_name(info=info)
+        credentials, project = self._get_authentication()
+        logger.info(
+            f"Initializing VertexAI with experiment name {experiment} "
+            f"and run name {run_name}."
+        )
+
+        aiplatform.init(
+            project=project,
+            location=self.config.location,
+            experiment=experiment,
+            experiment_tensorboard=settings.experiment_tensorboard,
+            staging_bucket=self.config.staging_bucket,
+            credentials=credentials,
+            encryption_spec_key_name=self.config.encryption_spec_key_name,
+            network=self.config.network,
+            api_endpoint=self.config.api_endpoint,
+            api_key=self.config.api_key,
+            api_transport=self.config.api_transport,
+            request_metadata=self.config.request_metadata,
+        )
+
+        try:
+            aiplatform.start_run(
+                run=run_name,
+                tensorboard=settings.experiment_tensorboard,
+                resume=True,
+            )
+        except exceptions.NotFound:
+            aiplatform.start_run(
+                run=run_name,
+                tensorboard=settings.experiment_tensorboard,
+                resume=False,
+            )
+
+        logger.info(
+            f"VertexAI experiment dashboard: {self._get_dashboard_url(experiment=experiment)}"
+        )
+        logger.info(
+            f"Tensorboard resource name: {self._get_tensorboard_resource_name(experiment=experiment)}"
+        )
+
+    def cleanup_step_run(self, info: "StepRunInfo", step_failed: bool) -> None:
+        """Stops the VertexAI run.
+
+        Args:
+            info: Info about the step that was executed.
+            step_failed: Whether the step failed or not.
+        """
+        state = (
+            execution.Execution.State.FAILED
+            if step_failed
+            else execution.Execution.State.COMPLETE
+        )
+        aiplatform.end_run(state=state)

zenml/integrations/gcp/flavors/__init__.py

@@ -21,6 +21,10 @@ from zenml.integrations.gcp.flavors.gcp_image_builder_flavor import (
     GCPImageBuilderConfig,
     GCPImageBuilderFlavor,
 )
+from zenml.integrations.gcp.flavors.vertex_experiment_tracker_flavor import (
+    VertexExperimentTrackerConfig,
+    VertexExperimentTrackerFlavor,
+)
 from zenml.integrations.gcp.flavors.vertex_orchestrator_flavor import (
     VertexOrchestratorConfig,
     VertexOrchestratorFlavor,
@@ -35,6 +39,8 @@ __all__ = [
     "GCPArtifactStoreConfig",
     "GCPImageBuilderFlavor",
     "GCPImageBuilderConfig",
+    "VertexExperimentTrackerFlavor",
+    "VertexExperimentTrackerConfig",
     "VertexOrchestratorFlavor",
     "VertexOrchestratorConfig",
     "VertexStepOperatorFlavor",

zenml/integrations/gcp/flavors/vertex_experiment_tracker_flavor.py

@@ -0,0 +1,199 @@
+#  Copyright (c) ZenML GmbH 2022. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Vertex experiment tracker flavor."""
+
+import re
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type, Union
+
+from pydantic import field_validator
+
+from zenml.config.base_settings import BaseSettings
+from zenml.experiment_trackers.base_experiment_tracker import (
+    BaseExperimentTrackerConfig,
+    BaseExperimentTrackerFlavor,
+)
+from zenml.integrations.gcp import (
+    GCP_RESOURCE_TYPE,
+    GCP_VERTEX_EXPERIMENT_TRACKER_FLAVOR,
+)
+from zenml.integrations.gcp.google_credentials_mixin import (
+    GoogleCredentialsConfigMixin,
+)
+from zenml.models import ServiceConnectorRequirements
+from zenml.utils.secret_utils import SecretField
+
+if TYPE_CHECKING:
+    from zenml.integrations.gcp.experiment_trackers import (
+        VertexExperimentTracker,
+    )
+
+
+class VertexExperimentTrackerSettings(BaseSettings):
+    """Settings for the VertexAI experiment tracker.
+
+    Attributes:
+        experiment: The VertexAI experiment name.
+        experiment_tensorboard: The VertexAI experiment tensorboard.
+    """
+
+    experiment: Optional[str] = None
+    experiment_tensorboard: Optional[Union[str, bool]] = None
+
+    @field_validator("experiment", mode="before")
+    def _validate_experiment(cls, value: str) -> str:
+        """Validates the experiment name matches the the allowed format.
+
+        Args:
+            value: The experiment.
+
+        Raises:
+            ValueError: If the experiment name does not match the expected
+                format.
+
+        Returns:
+            The experiment.
+        """
+        if value and not re.match(r"^[a-z0-9][a-z0-9-]{0,127}$", value):
+            raise ValueError(
+                "Experiment name must match regex [a-z0-9][a-z0-9-]{0,127}"
+            )
+        return value
+
+
+class VertexExperimentTrackerConfig(
+    BaseExperimentTrackerConfig,
+    GoogleCredentialsConfigMixin,
+    VertexExperimentTrackerSettings,
+):
+    """Config for the VertexAI experiment tracker.
+
+    Attributes:
+        location: Optional. The default location to use when making API calls. If not
+            set defaults to us-central1.
+        staging_bucket: Optional. The default staging bucket to use to stage artifacts
+            when making API calls. In the form gs://...
+        network:
+            Optional. The full name of the Compute Engine network to which jobs
+            and resources should be peered. E.g. "projects/12345/global/networks/myVPC".
+            Private services access must already be configured for the network.
+            If specified, all eligible jobs and resources created will be peered
+            with this VPC.
+        encryption_spec_key_name:
+            Optional. The Cloud KMS resource identifier of the customer
+            managed encryption key used to protect a resource. Has the
+            form:
+            ``projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key``.
+            The key needs to be in the same region as where the compute
+            resource is created.
+        api_endpoint (str):
+                Optional. The desired API endpoint,
+                e.g., us-central1-aiplatform.googleapis.com
+        api_key (str):
+            Optional. The API key to use for service calls.
+            NOTE: Not all services support API keys.
+        api_transport (str):
+            Optional. The transport method which is either 'grpc' or 'rest'.
+            NOTE: "rest" transport functionality is currently in a
+            beta state (preview).
+        request_metadata:
+            Optional. Additional gRPC metadata to send with every client request.
+    """
+
+    location: Optional[str] = None
+    staging_bucket: Optional[str] = None
+    network: Optional[str] = None
+    encryption_spec_key_name: Optional[str] = SecretField(default=None)
+    api_endpoint: Optional[str] = SecretField(default=None)
+    api_key: Optional[str] = SecretField(default=None)
+    api_transport: Optional[str] = None
+    request_metadata: Optional[Dict[str, Any]] = None
+
+
+class VertexExperimentTrackerFlavor(BaseExperimentTrackerFlavor):
+    """Flavor for the VertexAI experiment tracker."""
+
+    @property
+    def name(self) -> str:
+        """Name of the flavor.
+
+        Returns:
+            The name of the flavor.
+        """
+        return GCP_VERTEX_EXPERIMENT_TRACKER_FLAVOR
+
+    @property
+    def docs_url(self) -> Optional[str]:
+        """A URL to point at docs explaining this flavor.
+
+        Returns:
+            A flavor docs url.
+        """
+        return self.generate_default_docs_url()
+
+    @property
+    def sdk_docs_url(self) -> Optional[str]:
+        """A URL to point at SDK docs explaining this flavor.
+
+        Returns:
+            A flavor SDK docs url.
+        """
+        return self.generate_default_sdk_docs_url()
+
+    @property
+    def logo_url(self) -> str:
+        """A URL to represent the flavor in the dashboard.
+
+        Returns:
+            The flavor logo.
+        """
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/experiment_tracker/vertexai.png"
+
+    @property
+    def config_class(self) -> Type[VertexExperimentTrackerConfig]:
+        """Returns `VertexExperimentTrackerConfig` config class.
+
+        Returns:
+            The config class.
+        """
+        return VertexExperimentTrackerConfig
+
+    @property
+    def implementation_class(self) -> Type["VertexExperimentTracker"]:
+        """Implementation class for this flavor.
+
+        Returns:
+            The implementation class.
+        """
+        from zenml.integrations.gcp.experiment_trackers import (
+            VertexExperimentTracker,
+        )
+
+        return VertexExperimentTracker
+
+    @property
+    def service_connector_requirements(
+        self,
+    ) -> Optional[ServiceConnectorRequirements]:
+        """Service connector resource requirements for service connectors.
+
+        Specifies resource requirements that are used to filter the available
+        service connector types that are compatible with this flavor.
+
+        Returns:
+            Requirements for compatible service connectors, if a service
+            connector is required for this flavor.
+        """
+        return ServiceConnectorRequirements(
+            resource_type=GCP_RESOURCE_TYPE,
+        )

zenml/integrations/kubernetes/orchestrators/kube_utils.py

@@ -248,7 +248,7 @@ def wait_pod(
     Returns:
         The pod object which meets the exit condition.
     """
-    start_time = datetime.datetime.utcnow()
+    start_time = datetime.datetime.now(datetime.timezone.utc)
 
     # Link to exponential back-off algorithm used here:
     # https://cloud.google.com/storage/docs/exponential-backoff
@@ -288,7 +288,7 @@ def wait_pod(
             return resp
 
         # Check if wait timed out.
-        elapse_time = datetime.datetime.utcnow() - start_time
+        elapse_time = datetime.datetime.now(datetime.timezone.utc) - start_time
         if elapse_time.seconds >= timeout_sec and timeout_sec != 0:
             raise RuntimeError(
                 f"Waiting for pod `{namespace}:{pod_name}` timed out after "

zenml/integrations/mlflow/experiment_trackers/mlflow_experiment_tracker.py

@@ -236,7 +236,6 @@ class MLFlowExperimentTracker(BaseExperimentTracker):
         """Disables MLflow autologging for all supported frameworks."""
         frameworks = [
             "tensorflow",
-            "gluon",
             "xgboost",
             "lightgbm",
             "statsmodels",