zenml-nightly 0.72.0.dev20250121__py3-none-any.whl → 0.73.0.dev20250124__py3-none-any.whl

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/huggingface/__init__.py

@@ -47,16 +47,11 @@ class HuggingfaceIntegration(Integration):
             A list of requirements.
         """
         requirements = [
-            "datasets",
+            "datasets>=2.16.0",
             "huggingface_hub>0.19.0",
             "accelerate",
             "bitsandbytes>=0.41.3",
             "peft",
-            # temporary fix for CI issue similar to:
-            # - https://github.com/huggingface/datasets/issues/6737
-            # - https://github.com/huggingface/datasets/issues/6697
-            # TODO try relaxing it back going forward
-            "fsspec<=2023.12.0",
             "transformers",
         ]
 

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",

zenml/integrations/whylogs/data_validators/whylogs_data_validator.py

@@ -97,7 +97,9 @@ class WhylogsDataValidator(BaseDataValidator, AuthenticationMixin):
         """
         results = why.log(pandas=dataset)
         profile = results.profile()
-        dataset_timestamp = dataset_timestamp or datetime.datetime.utcnow()
+        dataset_timestamp = dataset_timestamp or datetime.datetime.now(
+            datetime.timezone.utc
+        )
         profile.set_dataset_timestamp(dataset_timestamp=dataset_timestamp)
         return profile.view()
 

zenml/models/v2/core/api_key.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Models representing API keys."""
 
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, ClassVar, List, Optional, Type, Union
 from uuid import UUID
 
@@ -319,7 +319,7 @@ class APIKeyInternalResponse(APIKeyResponse):
             and self.retain_period_minutes > 0
         ):
             # check if the previous key is still valid
-            if datetime.utcnow() - self.last_rotated < timedelta(
+            if datetime.now(timezone.utc) - self.last_rotated < timedelta(
                 minutes=self.retain_period_minutes
             ):
                 key_hash = self.previous_key

zenml/models/v2/core/schedule.py

@@ -14,13 +14,14 @@
 """Models representing schedules."""
 
 import datetime
-from typing import Optional, Union
+from typing import Dict, Optional, Union
 from uuid import UUID
 
 from pydantic import Field, model_validator
 
 from zenml.constants import STR_FIELD_MAX_LENGTH
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
 from zenml.models.v2.base.base import BaseUpdate
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
@@ -136,6 +137,11 @@ class ScheduleResponseMetadata(WorkspaceScopedResponseMetadata):
     orchestrator_id: Optional[UUID]
     pipeline_id: Optional[UUID]
 
+    run_metadata: Dict[str, MetadataType] = Field(
+        title="Metadata associated with this schedule.",
+        default={},
+    )
+
 
 class ScheduleResponseResources(WorkspaceScopedResponseResources):
     """Class for all resource models associated with the schedule entity."""
@@ -272,6 +278,15 @@ class ScheduleResponse(
         """
         return self.get_metadata().pipeline_id
 
+    @property
+    def run_metadata(self) -> Dict[str, MetadataType]:
+        """The `run_metadata` property.
+
+        Returns:
+            the value of the property.
+        """
+        return self.get_metadata().run_metadata
+
 
 # ------------------ Filter Model ------------------
 

zenml/orchestrators/publish_utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Utilities to publish pipeline and step runs."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Dict, List
 
 from zenml.client import Client
@@ -48,7 +48,7 @@ def publish_successful_step_run(
         step_run_id=step_run_id,
         step_run_update=StepRunUpdate(
             status=ExecutionStatus.COMPLETED,
-            end_time=datetime.utcnow(),
+            end_time=datetime.now(timezone.utc),
             outputs=output_artifact_ids,
         ),
     )
@@ -67,7 +67,7 @@ def publish_failed_step_run(step_run_id: "UUID") -> "StepRunResponse":
         step_run_id=step_run_id,
         step_run_update=StepRunUpdate(
             status=ExecutionStatus.FAILED,
-            end_time=datetime.utcnow(),
+            end_time=datetime.now(timezone.utc),
         ),
     )
 
@@ -87,7 +87,7 @@ def publish_failed_pipeline_run(
         run_id=pipeline_run_id,
         run_update=PipelineRunUpdate(
             status=ExecutionStatus.FAILED,
-            end_time=datetime.utcnow(),
+            end_time=datetime.now(timezone.utc),
         ),
     )
 

zenml/orchestrators/step_launcher.py

@@ -16,7 +16,7 @@
 import os
 import time
 from contextlib import nullcontext
-from datetime import datetime
+from datetime import datetime, timezone
 from functools import partial
 from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple
 
@@ -201,7 +201,7 @@ class StepLauncher:
                         f"Failed preparing step `{self._step_name}`."
                     )
                     step_run_request.status = ExecutionStatus.FAILED
-                    step_run_request.end_time = datetime.utcnow()
+                    step_run_request.end_time = datetime.now(timezone.utc)
                     raise
                 finally:
                     step_run = Client().zen_store.create_run_step(
@@ -305,7 +305,7 @@ class StepLauncher:
             The created or existing pipeline run,
             and a boolean indicating whether the run was created or reused.
         """
-        start_time = datetime.utcnow()
+        start_time = datetime.now(timezone.utc)
         run_name = string_utils.format_name_template(
             name_template=self._deployment.run_name_template,
             substitutions=self._deployment.pipeline_configuration._get_full_substitutions(

zenml/orchestrators/step_run_utils.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Utilities for creating step runs."""
 
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Dict, List, Optional, Set, Tuple
 
 from zenml.client import Client
@@ -75,7 +75,7 @@ class StepRunRequestFactory:
             pipeline_run_id=self.pipeline_run.id,
             deployment=self.deployment.id,
             status=ExecutionStatus.RUNNING,
-            start_time=datetime.utcnow(),
+            start_time=datetime.now(timezone.utc),
             user=Client().active_user.id,
             workspace=Client().active_workspace.id,
         )

zenml/pipelines/run_utils.py

@@ -1,7 +1,7 @@
 """Utility functions for running pipelines."""
 
 import time
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Union
 from uuid import UUID
 
@@ -65,7 +65,7 @@ def create_placeholder_run(
 
     if deployment.schedule:
         return None
-    start_time = datetime.utcnow()
+    start_time = datetime.now(timezone.utc)
     run_request = PipelineRunRequest(
         name=string_utils.format_name_template(
             name_template=deployment.run_name_template,

zenml/service_connectors/service_connector.py

@@ -799,8 +799,8 @@ class ServiceConnector(BaseModel, metaclass=ServiceConnectorMeta):
             name=name,
             body=ServiceConnectorResponseBody(
                 user=user,
-                created=datetime.utcnow(),
-                updated=datetime.utcnow(),
+                created=datetime.now(timezone.utc),
+                updated=datetime.now(timezone.utc),
                 description=description,
                 connector_type=self.get_type(),
                 auth_method=self.auth_method,

zenml/stack/stack.py

@@ -16,7 +16,7 @@
 import itertools
 import json
 import os
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import (
     TYPE_CHECKING,
     AbstractSet,
@@ -751,8 +751,8 @@ class Stack:
                 config=LocalImageBuilderConfig(),
                 user=Client().active_user.id,
                 workspace=Client().active_workspace.id,
-                created=datetime.utcnow(),
-                updated=datetime.utcnow(),
+                created=datetime.now(timezone.utc),
+                updated=datetime.now(timezone.utc),
             )
 
             self._image_builder = image_builder