zenml-nightly 0.68.1.dev20241112__py3-none-any.whl → 0.70.0.dev20241114__py3-none-any.whl

zenml/models/v2/core/pipeline_run.py

@@ -587,6 +587,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "stack_component",
         "pipeline_name",
         "templatable",
+        "run_metadata",
     ]
     name: Optional[str] = Field(
         default=None,
@@ -665,6 +666,10 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         default=None,
         description="Name/ID of the user that created the run.",
     )
+    run_metadata: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="The run_metadata to filter the pipeline runs by.",
+    )
     # TODO: Remove once frontend is ready for it. This is replaced by the more
     # generic `pipeline` filter below.
     pipeline_name: Optional[str] = Field(
@@ -694,7 +699,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
     templatable: Optional[bool] = Field(
         default=None, description="Whether the run is templatable."
     )
-
     model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
@@ -718,6 +722,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             PipelineDeploymentSchema,
             PipelineRunSchema,
             PipelineSchema,
+            RunMetadataSchema,
             ScheduleSchema,
             StackComponentSchema,
             StackCompositionSchema,
@@ -887,5 +892,21 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
                 )
 
             custom_filters.append(templatable_filter)
+        if self.run_metadata is not None:
+            from zenml.enums import MetadataResourceTypes
+
+            for key, value in self.run_metadata.items():
+                additional_filter = and_(
+                    RunMetadataSchema.resource_id == PipelineRunSchema.id,
+                    RunMetadataSchema.resource_type
+                    == MetadataResourceTypes.PIPELINE_RUN,
+                    RunMetadataSchema.key == key,
+                    self.generate_custom_query_conditions_for_column(
+                        value=value,
+                        table=RunMetadataSchema,
+                        column="value",
+                    ),
+                )
+                custom_filters.append(additional_filter)
 
         return custom_filters

zenml/models/v2/core/step_run.py

@@ -509,6 +509,7 @@ class StepRunFilter(WorkspaceScopedFilter):
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
         "model",
+        "run_metadata",
     ]
 
     name: Optional[str] = Field(
@@ -571,6 +572,10 @@ class StepRunFilter(WorkspaceScopedFilter):
         default=None,
         description="Name/ID of the model associated with the step run.",
     )
+    run_metadata: Optional[Dict[str, str]] = Field(
+        default=None,
+        description="The run_metadata to filter the step runs by.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
 
@@ -589,6 +594,7 @@ class StepRunFilter(WorkspaceScopedFilter):
         from zenml.zen_stores.schemas import (
             ModelSchema,
             ModelVersionSchema,
+            RunMetadataSchema,
             StepRunSchema,
         )
 
@@ -601,5 +607,21 @@ class StepRunFilter(WorkspaceScopedFilter):
                 ),
             )
             custom_filters.append(model_filter)
+        if self.run_metadata is not None:
+            from zenml.enums import MetadataResourceTypes
+
+            for key, value in self.run_metadata.items():
+                additional_filter = and_(
+                    RunMetadataSchema.resource_id == StepRunSchema.id,
+                    RunMetadataSchema.resource_type
+                    == MetadataResourceTypes.STEP_RUN,
+                    RunMetadataSchema.key == key,
+                    self.generate_custom_query_conditions_for_column(
+                        value=value,
+                        table=RunMetadataSchema,
+                        column="value",
+                    ),
+                )
+                custom_filters.append(additional_filter)
 
         return custom_filters

zenml/orchestrators/base_orchestrator.py

@@ -15,6 +15,7 @@
 
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
+from uuid import UUID
 
 from pydantic import model_validator
 
@@ -186,7 +187,17 @@ class BaseOrchestrator(StackComponent, ABC):
         """
         self._prepare_run(deployment=deployment)
 
-        environment = get_config_environment_vars(deployment=deployment)
+        pipeline_run_id: Optional[UUID] = None
+        schedule_id: Optional[UUID] = None
+        if deployment.schedule:
+            schedule_id = deployment.schedule.id
+        if placeholder_run:
+            pipeline_run_id = placeholder_run.id
+
+        environment = get_config_environment_vars(
+            schedule_id=schedule_id,
+            pipeline_run_id=pipeline_run_id,
+        )
 
         prevent_client_side_caching = handle_bool_env_var(
             ENV_ZENML_PREVENT_CLIENT_SIDE_CACHING, default=False

zenml/orchestrators/step_launcher.py

@@ -422,7 +422,8 @@ class StepLauncher:
             )
         )
         environment = orchestrator_utils.get_config_environment_vars(
-            deployment=self._deployment
+            pipeline_run_id=step_run_info.run_id,
+            step_run_id=step_run_info.step_run_id,
         )
         if last_retry:
             environment[ENV_ZENML_IGNORE_FAILURE_HOOK] = str(False)

zenml/orchestrators/utils.py

@@ -28,16 +28,16 @@ from zenml.constants import (
     ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
     ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
-    PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
 from zenml.enums import AuthScheme, StackComponentType, StoreType
 from zenml.logger import get_logger
 from zenml.stack import StackComponent
 from zenml.utils.string_utils import format_name_template
 
+logger = get_logger(__name__)
+
 if TYPE_CHECKING:
     from zenml.artifact_stores.base_artifact_store import BaseArtifactStore
-    from zenml.models import PipelineDeploymentResponse
 
 
 def get_orchestrator_run_name(pipeline_name: str) -> str:
@@ -80,16 +80,23 @@ def is_setting_enabled(
 
 
 def get_config_environment_vars(
-    deployment: Optional["PipelineDeploymentResponse"] = None,
+    schedule_id: Optional[UUID] = None,
+    pipeline_run_id: Optional[UUID] = None,
+    step_run_id: Optional[UUID] = None,
 ) -> Dict[str, str]:
     """Gets environment variables to set for mirroring the active config.
 
-    If a pipeline deployment is given, the environment variables will be set to
-    include a newly generated API token valid for the duration of the pipeline
-    run instead of the API token from the global config.
+    If a schedule ID, pipeline run ID or step run ID is given, and the current
+    client is not authenticated to a server with an API key, the environment
+    variables will be updated to include a newly generated workload API token
+    that will be valid for the duration of the schedule, pipeline run, or step
+    run instead of the current API token used to authenticate the client.
 
     Args:
-        deployment: Optional deployment to use for the environment variables.
+        schedule_id: Optional schedule ID to use to generate a new API token.
+        pipeline_run_id: Optional pipeline run ID to use to generate a new API
+            token.
+        step_run_id: Optional step run ID to use to generate a new API token.
 
     Returns:
         Environment variable dict.
@@ -107,34 +114,46 @@ def get_config_environment_vars(
     ):
         credentials_store = get_credentials_store()
         url = global_config.store_configuration.url
-        api_key = credentials_store.get_api_key(url)
         api_token = credentials_store.get_token(url, allow_expired=False)
-        if api_key:
-            environment_vars[ENV_ZENML_STORE_PREFIX + "API_KEY"] = api_key
-        elif deployment:
-            # When connected to an authenticated ZenML server, if a pipeline
-            # deployment is supplied, we need to fetch an API token that will be
-            # valid for the duration of the pipeline run.
+        if schedule_id or pipeline_run_id or step_run_id:
+            # When connected to an authenticated ZenML server, if a schedule ID,
+            # pipeline run ID or step run ID is supplied, we need to fetch a new
+            # workload API token scoped to the schedule, pipeline run or step
+            # run.
             assert isinstance(global_config.zen_store, RestZenStore)
-            pipeline_id: Optional[UUID] = None
-            if deployment.pipeline:
-                pipeline_id = deployment.pipeline.id
-            schedule_id: Optional[UUID] = None
-            expires_minutes: Optional[int] = PIPELINE_API_TOKEN_EXPIRES_MINUTES
-            if deployment.schedule:
-                schedule_id = deployment.schedule.id
-                # If a schedule is given, this is a long running pipeline that
-                # should not have an API token that expires.
-                expires_minutes = None
+
+            # If only a schedule is given, the pipeline run credentials will
+            # be valid for the entire duration of the schedule.
+            api_key = credentials_store.get_api_key(url)
+            if not api_key and not pipeline_run_id and not step_run_id:
+                logger.warning(
+                    "An API token without an expiration time will be generated "
+                    "and used to run this pipeline on a schedule. This is very "
+                    "insecure because the API token will be valid for the "
+                    "entire lifetime of the schedule and can be used to access "
+                    "your user account if accidentally leaked. When deploying "
+                    "a pipeline on a schedule, it is strongly advised to use a "
+                    "service account API key to authenticate to the ZenML "
+                    "server instead of your regular user account. For more "
+                    "information, see "
+                    "https://docs.zenml.io/how-to/connecting-to-zenml/connect-with-a-service-account"
+                )
+
+            # The schedule, pipeline run or step run credentials are scoped to
+            # the schedule, pipeline run or step run and will only be valid for
+            # the duration of the schedule/pipeline run/step run.
             new_api_token = global_config.zen_store.get_api_token(
-                pipeline_id=pipeline_id,
                 schedule_id=schedule_id,
-                expires_minutes=expires_minutes,
+                pipeline_run_id=pipeline_run_id,
+                step_run_id=step_run_id,
             )
+
             environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
                 new_api_token
             )
         elif api_token:
+            # For all other cases, the pipeline run environment is configured
+            # with the current access token.
             environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
                 api_token.access_token
             )

zenml/steps/utils.py

@@ -442,6 +442,11 @@ def log_step_metadata(
             from within a step or if no pipeline name or ID is provided and
             the function is not called from within a step.
     """
+    logger.warning(
+        "The `log_step_metadata` function is deprecated and will soon be "
+        "removed. Please use `log_metadata` instead."
+    )
+
     step_context = None
     if not step_name:
         with contextlib.suppress(RuntimeError):

zenml/utils/metadata_utils.py

@@ -0,0 +1,335 @@
+#  Copyright (c) ZenML GmbH 2024. 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.
+"""Utility functions to handle metadata for ZenML entities."""
+
+import contextlib
+from typing import Dict, Optional, Union, overload
+from uuid import UUID
+
+from zenml.client import Client
+from zenml.enums import MetadataResourceTypes
+from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.steps.step_context import get_step_context
+
+logger = get_logger(__name__)
+
+
+@overload
+def log_metadata(metadata: Dict[str, MetadataType]) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    artifact_version_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    artifact_name: str,
+    artifact_version: Optional[str] = None,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    model_version_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    model_name: str,
+    model_version: str,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    step_id: UUID,
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    run_id_name_or_prefix: Union[UUID, str],
+) -> None: ...
+
+
+@overload
+def log_metadata(
+    *,
+    metadata: Dict[str, MetadataType],
+    step_name: str,
+    run_id_name_or_prefix: Union[UUID, str],
+) -> None: ...
+
+
+def log_metadata(
+    metadata: Dict[str, MetadataType],
+    # Parameters to manually log metadata for steps and runs
+    step_id: Optional[UUID] = None,
+    step_name: Optional[str] = None,
+    run_id_name_or_prefix: Optional[Union[UUID, str]] = None,
+    # Parameters to manually log metadata for artifacts
+    artifact_version_id: Optional[UUID] = None,
+    artifact_name: Optional[str] = None,
+    artifact_version: Optional[str] = None,
+    # Parameters to manually log metadata for models
+    model_version_id: Optional[UUID] = None,
+    model_name: Optional[str] = None,
+    model_version: Optional[str] = None,
+) -> None:
+    """Logs metadata for various resource types in a generalized way.
+
+    Args:
+        metadata: The metadata to log.
+        step_id: The ID of the step.
+        step_name: The name of the step.
+        run_id_name_or_prefix: The id, name or prefix of the run
+        artifact_version_id: The ID of the artifact version
+        artifact_name: The name of the artifact.
+        artifact_version: The version of the artifact.
+        model_version_id: The ID of the model version.
+        model_name: The name of the model.
+        model_version: The version of the model
+
+    Raises:
+        ValueError: If no identifiers are provided and the function is not
+            called from within a step.
+    """
+    client = Client()
+
+    # If a step name is provided, we need a run_id_name_or_prefix and will log
+    # metadata for the steps pipeline and model accordingly.
+    if step_name is not None and run_id_name_or_prefix is not None:
+        run_model = client.get_pipeline_run(
+            name_id_or_prefix=run_id_name_or_prefix
+        )
+        step_model = run_model.steps[step_name]
+
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=run_model.id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_model.id,
+            resource_type=MetadataResourceTypes.STEP_RUN,
+        )
+        if step_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If a step is identified by id, fetch it directly through the client,
+    # follow a similar procedure and log metadata for its pipeline and model
+    # as well.
+    elif step_id is not None:
+        step_model = client.get_run_step(step_run_id=step_id)
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_model.pipeline_run_id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_model.id,
+            resource_type=MetadataResourceTypes.STEP_RUN,
+        )
+        if step_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If a pipeline run id is identified, we need to log metadata to it and its
+    # model as well.
+    elif run_id_name_or_prefix is not None:
+        run_model = client.get_pipeline_run(
+            name_id_or_prefix=run_id_name_or_prefix
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=run_model.id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        if run_model.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=run_model.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    # If the user provides a model name and version, we use to model abstraction
+    # to fetch the model version and attach the corresponding metadata to it.
+    elif model_name is not None and model_version is not None:
+        from zenml import Model
+
+        mv = Model(name=model_name, version=model_version)
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=mv.id,
+            resource_type=MetadataResourceTypes.MODEL_VERSION,
+        )
+
+    # If the user provides a model version id, we use the client to fetch it and
+    # attach the metadata to it.
+    elif model_version_id is not None:
+        model_version_id = client.get_model_version(
+            model_version_name_or_number_or_id=model_version_id
+        ).id
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=model_version_id,
+            resource_type=MetadataResourceTypes.MODEL_VERSION,
+        )
+
+    # If the user provides an artifact name, there are three possibilities. If
+    # an artifact version is also provided with the name, we use both to fetch
+    # the artifact version and use it to log the metadata. If no version is
+    # provided, if the function is called within a step we search the artifacts
+    # of the step if not we fetch the latest version and attach the metadata
+    # to the latest version.
+    elif artifact_name is not None:
+        if artifact_version:
+            artifact_version_model = client.get_artifact_version(
+                name_id_or_prefix=artifact_name, version=artifact_version
+            )
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=artifact_version_model.id,
+                resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+            )
+        else:
+            step_context = None
+            with contextlib.suppress(RuntimeError):
+                step_context = get_step_context()
+
+            if step_context:
+                step_context.add_output_metadata(
+                    metadata=metadata, output_name=artifact_name
+                )
+            else:
+                artifact_version_model = client.get_artifact_version(
+                    name_id_or_prefix=artifact_name
+                )
+                client.create_run_metadata(
+                    metadata=metadata,
+                    resource_id=artifact_version_model.id,
+                    resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+                )
+
+    # If the user directly provides an artifact_version_id, we use the client to
+    # fetch is and attach the metadata accordingly.
+    elif artifact_version_id is not None:
+        artifact_version_model = client.get_artifact_version(
+            name_id_or_prefix=artifact_version_id,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=artifact_version_model.id,
+            resource_type=MetadataResourceTypes.ARTIFACT_VERSION,
+        )
+
+    # If every additional value is None, that means we are calling it bare bones
+    # and this call needs to happen during a step execution. We will use the
+    # step context to fetch the step, run and possibly the model version and
+    # attach the metadata accordingly.
+    elif all(
+        v is None
+        for v in [
+            step_id,
+            step_name,
+            run_id_name_or_prefix,
+            artifact_version_id,
+            artifact_name,
+            artifact_version,
+            model_version_id,
+            model_name,
+            model_version,
+        ]
+    ):
+        try:
+            step_context = get_step_context()
+        except RuntimeError:
+            raise ValueError(
+                "You are calling 'log_metadata()' outside of a step execution. "
+                "If you would like to add metadata to a ZenML entity outside "
+                "of the step execution, please provide the required "
+                "identifiers."
+            )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_context.pipeline_run.id,
+            resource_type=MetadataResourceTypes.PIPELINE_RUN,
+        )
+        client.create_run_metadata(
+            metadata=metadata,
+            resource_id=step_context.step_run.id,
+            resource_type=MetadataResourceTypes.STEP_RUN,
+        )
+        if step_context.model_version:
+            client.create_run_metadata(
+                metadata=metadata,
+                resource_id=step_context.model_version.id,
+                resource_type=MetadataResourceTypes.MODEL_VERSION,
+            )
+
+    else:
+        raise ValueError(
+            """
+            Unsupported way to call the `log_metadata`. Possible combinations "
+            include:
+            
+            # Inside a step
+            # Logs the metadata to the step, its run and possibly its model
+            log_metadata(metadata={})
+            
+            # Manually logging for a step
+            # Logs the metadata to the step, its run and possibly its model
+            log_metadata(metadata={}, step_name=..., run_id_name_or_prefix=...)
+            log_metadata(metadata={}, step_id=...)
+            
+            # Manually logging for a run 
+            # Logs the metadata to the run, possibly its model
+            log_metadata(metadata={}, run_id_name_or_prefix=...)
+            
+            # Manually logging for a model
+            log_metadata(metadata={}, model_name=..., model_version=...)
+            log_metadata(metadata={}, model_version_id=...)
+            
+            # Manually logging for an artifact
+            log_metadata(metadata={}, artifact_name=...)  # inside a step 
+            log_metadata(metadata={}, artifact_name=..., artifact_version=...)
+            log_metadata(metadata={}, artifact_version_id=...)
+            """
+        )