zenml-nightly 0.68.0.dev20241027__py3-none-any.whl → 0.68.1.dev20241101__py3-none-any.whl

zenml/models/v2/core/pipeline_run.py

@@ -586,6 +586,7 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "stack",
         "code_repository",
         "model",
+        "stack_component",
         "pipeline_name",
         "templatable",
     ]
@@ -688,6 +689,10 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         default=None,
         description="Name/ID of the model associated with the run.",
     )
+    stack_component: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the stack component associated with the run.",
+    )
     templatable: Optional[bool] = Field(
         default=None, description="Whether the run is templatable."
     )
@@ -716,6 +721,8 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             PipelineRunSchema,
             PipelineSchema,
             ScheduleSchema,
+            StackComponentSchema,
+            StackCompositionSchema,
             StackSchema,
             UserSchema,
         )
@@ -772,7 +779,9 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 PipelineRunSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)
@@ -821,6 +830,19 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(model_filter)
 
+        if self.stack_component:
+            component_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+                StackSchema.id == StackCompositionSchema.stack_id,
+                StackCompositionSchema.component_id == StackComponentSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.stack_component,
+                    table=StackComponentSchema,
+                ),
+            )
+            custom_filters.append(component_filter)
+
         if self.pipeline_name:
             pipeline_name_filter = and_(
                 PipelineRunSchema.pipeline_id == PipelineSchema.id,

zenml/models/v2/core/run_template.py

@@ -413,7 +413,9 @@ class RunTemplateFilter(WorkspaceScopedTaggableFilter):
             user_filter = and_(
                 RunTemplateSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/stack.py

@@ -228,7 +228,7 @@ class StackResponse(
             component_dict = dict(
                 name=component.name,
                 type=str(component.type),
-                flavor=component.flavor,
+                flavor=component.flavor_name,
             )
             configuration = json.loads(
                 component.get_metadata().model_dump_json(
@@ -255,7 +255,9 @@ class StackResponse(
             Dict of analytics metadata.
         """
         metadata = super().get_analytics_metadata()
-        metadata.update({ct: c[0].flavor for ct, c in self.components.items()})
+        metadata.update(
+            {ct: c[0].flavor_name for ct, c in self.components.items()}
+        )
 
         if self.labels is not None:
             metadata.update(
@@ -381,7 +383,9 @@ class StackFilter(WorkspaceScopedFilter):
             user_filter = and_(
                 StackSchema.user_id == UserSchema.id,
                 self.generate_name_or_id_query_conditions(
-                    value=self.user, table=UserSchema
+                    value=self.user,
+                    table=UserSchema,
+                    additional_columns=["full_name"],
                 ),
             )
             custom_filters.append(user_filter)

zenml/models/v2/core/step_run.py

@@ -14,7 +14,7 @@
 """Models representing steps runs."""
 
 from datetime import datetime
-from typing import TYPE_CHECKING, Dict, List, Optional, Union
+from typing import TYPE_CHECKING, ClassVar, Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -33,6 +33,8 @@ from zenml.models.v2.base.scoped import (
 from zenml.models.v2.core.model_version import ModelVersionResponse
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import ColumnElement
+
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.logs import (
         LogsRequest,
@@ -488,6 +490,11 @@ class StepRunResponse(
 class StepRunFilter(WorkspaceScopedFilter):
     """Model to enable advanced filtering of step runs."""
 
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "model",
+    ]
+
     name: Optional[str] = Field(
         default=None,
         description="Name of the step run",
@@ -541,8 +548,42 @@ class StepRunFilter(WorkspaceScopedFilter):
     )
     model_version_id: Optional[Union[UUID, str]] = Field(
         default=None,
-        description="Model version associated with the pipeline run.",
+        description="Model version associated with the step run.",
         union_mode="left_to_right",
     )
+    model: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the model associated with the step run.",
+    )
 
     model_config = ConfigDict(protected_namespaces=())
+
+    def get_custom_filters(
+        self,
+    ) -> List["ColumnElement[bool]"]:
+        """Get custom filters.
+
+        Returns:
+            A list of custom filters.
+        """
+        custom_filters = super().get_custom_filters()
+
+        from sqlmodel import and_
+
+        from zenml.zen_stores.schemas import (
+            ModelSchema,
+            ModelVersionSchema,
+            StepRunSchema,
+        )
+
+        if self.model:
+            model_filter = and_(
+                StepRunSchema.model_version_id == ModelVersionSchema.id,
+                ModelVersionSchema.model_id == ModelSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.model, table=ModelSchema
+                ),
+            )
+            custom_filters.append(model_filter)
+
+        return custom_filters

zenml/models/v2/misc/auth_models.py

@@ -13,10 +13,10 @@
 #  permissions and limitations under the License.
 """Models representing OAuth2 requests and responses."""
 
-from typing import Optional
+from typing import Any, Dict, Optional
 from uuid import UUID
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict
 
 from zenml.enums import OAuthGrantTypes
 
@@ -29,6 +29,7 @@ class OAuthDeviceAuthorizationRequest(BaseModel):
     """OAuth2 device authorization grant request."""
 
     client_id: UUID
+    device_id: Optional[UUID] = None
 
 
 class OAuthDeviceVerificationRequest(BaseModel):
@@ -119,6 +120,14 @@ class OAuthTokenResponse(BaseModel):
     expires_in: Optional[int] = None
     refresh_token: Optional[str] = None
     scope: Optional[str] = None
+    cookie_name: Optional[str] = None
+    device_id: Optional[UUID] = None
+    device_metadata: Optional[Dict[str, Any]] = None
+
+    model_config = ConfigDict(
+        # Allow extra attributes to allow compatibility with different versions
+        extra="allow",
+    )
 
 
 class OAuthRedirectResponse(BaseModel):

zenml/models/v2/misc/server_models.py

@@ -52,6 +52,8 @@ class ServerModel(BaseModel):
 
     id: UUID = Field(default_factory=uuid4, title="The unique server id.")
 
+    name: Optional[str] = Field(None, title="The name of the ZenML server.")
+
     version: str = Field(
         title="The ZenML version that the server is running.",
     )

zenml/orchestrators/base_orchestrator.py

@@ -166,9 +166,8 @@ class BaseOrchestrator(StackComponent, ABC):
             environment: Environment variables to set in the orchestration
                 environment. These don't need to be set if running locally.
 
-        Returns:
-            The optional return value from this method will be returned by the
-            `pipeline_instance.run()` call when someone is running a pipeline.
+        Yields:
+            Metadata for the pipeline run.
         """
 
     def run(
@@ -176,7 +175,7 @@ class BaseOrchestrator(StackComponent, ABC):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> Any:
+    ) -> None:
         """Runs a pipeline on a stack.
 
         Args:
@@ -211,6 +210,11 @@ class BaseOrchestrator(StackComponent, ABC):
                 # the orchestrator does not try to run them
                 deployment.step_configurations.pop(invocation_id)
 
+            for step in deployment.step_configurations.values():
+                for invocation_id in cached_invocations:
+                    if invocation_id in step.spec.upstream_steps:
+                        step.spec.upstream_steps.remove(invocation_id)
+
             if len(deployment.step_configurations) == 0:
                 # All steps were cached, we update the pipeline run status and
                 # don't actually use the orchestrator to run the pipeline

zenml/orchestrators/step_launcher.py

@@ -331,6 +331,7 @@ class StepLauncher:
             status=ExecutionStatus.RUNNING,
             orchestrator_environment=get_run_environment_dict(),
             start_time=datetime.utcnow(),
+            tags=self._deployment.pipeline_configuration.tags,
         )
         return client.zen_store.get_or_create_run(pipeline_run)
 

zenml/orchestrators/step_run_utils.py

@@ -35,7 +35,7 @@ from zenml.models import (
 )
 from zenml.orchestrators import cache_utils, input_utils, utils
 from zenml.stack import Stack
-from zenml.utils import string_utils
+from zenml.utils import pagination_utils, string_utils
 
 if TYPE_CHECKING:
     from zenml.model.model import Model
@@ -424,7 +424,15 @@ def get_model_version_created_by_pipeline_run(
         ):
             return pipeline_run.model_version
 
-    for step_run in pipeline_run.steps.values():
+    # We fetch a list of hydrated step runs here in order to avoid hydration
+    # calls for each step separately.
+    candidate_step_runs = pagination_utils.depaginate(
+        Client().list_run_steps,
+        pipeline_run_id=pipeline_run.id,
+        model=model_name,
+        hydrate=True,
+    )
+    for step_run in candidate_step_runs:
         if step_run.config.model and step_run.model_version:
             if (
                 step_run.config.model.name == model_name

zenml/orchestrators/step_runner.py

@@ -57,6 +57,7 @@ from zenml.utils import materializer_utils, source_utils
 from zenml.utils.typing_utils import get_origin, is_union
 
 if TYPE_CHECKING:
+    from zenml.artifact_stores import BaseArtifactStore
     from zenml.config.source import Source
     from zenml.config.step_configurations import Step
     from zenml.models import (
@@ -190,61 +191,62 @@ class StepRunner:
                         )
                 raise
             finally:
-                step_run_metadata = self._stack.get_step_run_metadata(
-                    info=step_run_info,
-                )
-                publish_step_run_metadata(
-                    step_run_id=step_run_info.step_run_id,
-                    step_run_metadata=step_run_metadata,
-                )
-                self._stack.cleanup_step_run(
-                    info=step_run_info, step_failed=step_failed
-                )
-                if not step_failed:
-                    if (
-                        success_hook_source
-                        := self.configuration.success_hook_source
-                    ):
-                        logger.info("Detected success hook. Running...")
-                        self.load_and_run_hook(
-                            success_hook_source,
-                            step_exception=None,
-                        )
-
-                    # Store and publish the output artifacts of the step function.
-                    output_data = self._validate_outputs(
-                        return_values, output_annotations
+                try:
+                    step_run_metadata = self._stack.get_step_run_metadata(
+                        info=step_run_info,
                     )
-                    artifact_metadata_enabled = is_setting_enabled(
-                        is_enabled_on_step=step_run_info.config.enable_artifact_metadata,
-                        is_enabled_on_pipeline=step_run_info.pipeline.enable_artifact_metadata,
+                    publish_step_run_metadata(
+                        step_run_id=step_run_info.step_run_id,
+                        step_run_metadata=step_run_metadata,
                     )
-                    artifact_visualization_enabled = is_setting_enabled(
-                        is_enabled_on_step=step_run_info.config.enable_artifact_visualization,
-                        is_enabled_on_pipeline=step_run_info.pipeline.enable_artifact_visualization,
+                    self._stack.cleanup_step_run(
+                        info=step_run_info, step_failed=step_failed
                     )
-                    output_artifacts = self._store_output_artifacts(
-                        output_data=output_data,
-                        output_artifact_uris=output_artifact_uris,
-                        output_materializers=output_materializers,
-                        output_annotations=output_annotations,
-                        artifact_metadata_enabled=artifact_metadata_enabled,
-                        artifact_visualization_enabled=artifact_visualization_enabled,
-                    )
-
-                    if (
-                        model_version := step_run.model_version
-                        or pipeline_run.model_version
-                    ):
-                        from zenml.orchestrators import step_run_utils
-
-                        step_run_utils.link_output_artifacts_to_model_version(
-                            artifacts=output_artifacts,
-                            output_configurations=step_run.config.outputs,
-                            model_version=model_version,
+                    if not step_failed:
+                        if (
+                            success_hook_source
+                            := self.configuration.success_hook_source
+                        ):
+                            logger.info("Detected success hook. Running...")
+                            self.load_and_run_hook(
+                                success_hook_source,
+                                step_exception=None,
+                            )
+
+                        # Store and publish the output artifacts of the step function.
+                        output_data = self._validate_outputs(
+                            return_values, output_annotations
+                        )
+                        artifact_metadata_enabled = is_setting_enabled(
+                            is_enabled_on_step=step_run_info.config.enable_artifact_metadata,
+                            is_enabled_on_pipeline=step_run_info.pipeline.enable_artifact_metadata,
+                        )
+                        artifact_visualization_enabled = is_setting_enabled(
+                            is_enabled_on_step=step_run_info.config.enable_artifact_visualization,
+                            is_enabled_on_pipeline=step_run_info.pipeline.enable_artifact_visualization,
                         )
+                        output_artifacts = self._store_output_artifacts(
+                            output_data=output_data,
+                            output_artifact_uris=output_artifact_uris,
+                            output_materializers=output_materializers,
+                            output_annotations=output_annotations,
+                            artifact_metadata_enabled=artifact_metadata_enabled,
+                            artifact_visualization_enabled=artifact_visualization_enabled,
+                        )
+
+                        if (
+                            model_version := step_run.model_version
+                            or pipeline_run.model_version
+                        ):
+                            from zenml.orchestrators import step_run_utils
 
-                StepContext._clear()  # Remove the step context singleton
+                            step_run_utils.link_output_artifacts_to_model_version(
+                                artifacts=output_artifacts,
+                                output_configurations=step_run.config.outputs,
+                                model_version=model_version,
+                            )
+                finally:
+                    StepContext._clear()  # Remove the step context singleton
 
             # Update the status and output artifacts of the step run.
             output_artifact_ids = {
@@ -406,15 +408,25 @@ class StepRunner:
             )
         )
 
-        with register_artifact_store_filesystem(
-            artifact.artifact_store_id
-        ) as target_artifact_store:
+        def _load_artifact(artifact_store: "BaseArtifactStore") -> Any:
             materializer: BaseMaterializer = materializer_class(
-                uri=artifact.uri, artifact_store=target_artifact_store
+                uri=artifact.uri, artifact_store=artifact_store
             )
-            materializer.validate_type_compatibility(data_type)
+            materializer.validate_load_type_compatibility(data_type)
             return materializer.load(data_type=data_type)
 
+        if artifact.artifact_store_id == self._stack.artifact_store.id:
+            # Register the artifact store of the active stack here to avoid
+            # unnecessary component/flavor calls when using
+            # `register_artifact_store_filesystem(...)`
+            self._stack.artifact_store._register()
+            return _load_artifact(artifact_store=self._stack.artifact_store)
+        else:
+            with register_artifact_store_filesystem(
+                artifact.artifact_store_id
+            ) as target_artifact_store:
+                return _load_artifact(artifact_store=target_artifact_store)
+
     def _validate_outputs(
         self,
         return_values: Any,

zenml/orchestrators/utils.py

@@ -25,11 +25,12 @@ from zenml.config.global_config import (
 from zenml.constants import (
     ENV_ZENML_ACTIVE_STACK_ID,
     ENV_ZENML_ACTIVE_WORKSPACE_ID,
+    ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING,
     ENV_ZENML_SERVER,
     ENV_ZENML_STORE_PREFIX,
     PIPELINE_API_TOKEN_EXPIRES_MINUTES,
 )
-from zenml.enums import StackComponentType, StoreType
+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
@@ -93,32 +94,54 @@ def get_config_environment_vars(
     Returns:
         Environment variable dict.
     """
+    from zenml.login.credentials_store import get_credentials_store
     from zenml.zen_stores.rest_zen_store import RestZenStore
 
     global_config = GlobalConfiguration()
     environment_vars = global_config.get_config_environment_vars()
 
-    if deployment and global_config.store_configuration.type == StoreType.REST:
-        # When connected to a 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.
-        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
-        api_token = global_config.zen_store.get_api_token(
-            pipeline_id=pipeline_id,
-            schedule_id=schedule_id,
-            expires_minutes=expires_minutes,
-        )
-        environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = api_token
+    if (
+        global_config.store_configuration.type == StoreType.REST
+        and global_config.zen_store.get_store_info().auth_scheme
+        != AuthScheme.NO_AUTH
+    ):
+        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.
+            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
+            new_api_token = global_config.zen_store.get_api_token(
+                pipeline_id=pipeline_id,
+                schedule_id=schedule_id,
+                expires_minutes=expires_minutes,
+            )
+            environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
+                new_api_token
+            )
+        elif api_token:
+            environment_vars[ENV_ZENML_STORE_PREFIX + "API_TOKEN"] = (
+                api_token.access_token
+            )
+
+    # Disable credentials caching to avoid storing sensitive information
+    # in the pipeline run environment
+    environment_vars[ENV_ZENML_DISABLE_CREDENTIALS_DISK_CACHING] = "true"
 
     # Make sure to use the correct active stack/workspace which might come
     # from a .zen repository and not the global config

zenml/pipelines/pipeline_decorator.py

@@ -18,6 +18,7 @@ from typing import (
     Any,
     Callable,
     Dict,
+    List,
     Optional,
     TypeVar,
     Union,
@@ -49,6 +50,7 @@ def pipeline(
     enable_artifact_metadata: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
     settings: Optional[Dict[str, "SettingsOrDict"]] = None,
+    tags: Optional[List[str]] = None,
     extra: Optional[Dict[str, Any]] = None,
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
@@ -64,6 +66,7 @@ def pipeline(
     enable_artifact_metadata: Optional[bool] = None,
     enable_step_logs: Optional[bool] = None,
     settings: Optional[Dict[str, "SettingsOrDict"]] = None,
+    tags: Optional[List[str]] = None,
     extra: Optional[Dict[str, Any]] = None,
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
@@ -79,6 +82,7 @@ def pipeline(
         enable_artifact_metadata: Whether to enable artifact metadata or not.
         enable_step_logs: If step logs should be enabled for this pipeline.
         settings: Settings for this pipeline.
+        tags: Tags to apply to runs of the pipeline.
         extra: Extra configurations for this pipeline.
         on_failure: Callback function in event of failure of the step. Can be a
             function with a single argument of type `BaseException`, or a source
@@ -101,6 +105,7 @@ def pipeline(
             enable_artifact_metadata=enable_artifact_metadata,
             enable_step_logs=enable_step_logs,
             settings=settings,
+            tags=tags,
             extra=extra,
             on_failure=on_failure,
             on_success=on_success,