zenml-nightly 0.66.0.dev20240923__py3-none-any.whl → 0.66.0.dev20240925__py3-none-any.whl

zenml/models/v2/core/pipeline_run.py

@@ -21,6 +21,7 @@ from typing import (
     List,
     Optional,
     Union,
+    cast,
 )
 from uuid import UUID
 
@@ -28,8 +29,7 @@ from pydantic import BaseModel, ConfigDict, Field
 
 from zenml.config.pipeline_configurations import PipelineConfiguration
 from zenml.constants import STR_FIELD_MAX_LENGTH
-from zenml.enums import ExecutionStatus, GenericFilterOps
-from zenml.models.v2.base.filter import StrFilter
+from zenml.enums import ExecutionStatus
 from zenml.models.v2.base.scoped import (
     WorkspaceScopedFilter,
     WorkspaceScopedRequest,
@@ -137,7 +137,8 @@ class PipelineRunUpdate(BaseModel):
         "configured by this pipeline run explicitly.",
         default=None,
     )
-    # TODO: we should maybe have a different update model here, the upper three attributes should only be for internal use
+    # TODO: we should maybe have a different update model here, the upper
+    #  three attributes should only be for internal use
     add_tags: Optional[List[str]] = Field(
         default=None, title="New tags to add to the pipeline run."
     )
@@ -236,6 +237,7 @@ class PipelineRunResponseMetadata(WorkspaceScopedResponseMetadata):
         description="Template used for the pipeline run.",
     )
     is_templatable: bool = Field(
+        default=False,
         description="Whether a template can be created from this run.",
     )
 
@@ -308,6 +310,64 @@ class PipelineRunResponse(
 
         return get_artifacts_versions_of_pipeline_run(self, only_produced=True)
 
+    def refresh_run_status(self) -> "PipelineRunResponse":
+        """Method to refresh the status of a run if it is initializing/running.
+
+        Returns:
+            The updated pipeline.
+
+        Raises:
+            ValueError: If the stack of the run response is None.
+        """
+        if self.status in [
+            ExecutionStatus.INITIALIZING,
+            ExecutionStatus.RUNNING,
+        ]:
+            # Check if the stack still accessible
+            if self.stack is None:
+                raise ValueError(
+                    "The stack that this pipeline run response was executed on"
+                    "has been deleted."
+                )
+
+            # Create the orchestrator instance
+            from zenml.enums import StackComponentType
+            from zenml.orchestrators.base_orchestrator import BaseOrchestrator
+            from zenml.stack.stack_component import StackComponent
+
+            # Check if the stack still accessible
+            orchestrator_list = self.stack.components.get(
+                StackComponentType.ORCHESTRATOR, []
+            )
+            if len(orchestrator_list) == 0:
+                raise ValueError(
+                    "The orchestrator that this pipeline run response was "
+                    "executed with has been deleted."
+                )
+
+            orchestrator = cast(
+                BaseOrchestrator,
+                StackComponent.from_model(
+                    component_model=orchestrator_list[0]
+                ),
+            )
+
+            # Fetch the status
+            status = orchestrator.fetch_status(run=self)
+
+            # If it is different from the current status, update it
+            if status != self.status:
+                from zenml.client import Client
+                from zenml.models import PipelineRunUpdate
+
+                client = Client()
+                return client.zen_store.update_run(
+                    run_id=self.id,
+                    run_update=PipelineRunUpdate(status=status),
+                )
+
+        return self
+
     # Body and metadata properties
     @property
     def status(self) -> ExecutionStatus:
@@ -522,6 +582,11 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "schedule_id",
         "stack_id",
         "template_id",
+        "user",
+        "pipeline",
+        "stack",
+        "code_repository",
+        "model",
         "pipeline_name",
         "templatable",
     ]
@@ -538,10 +603,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         description="Pipeline associated with the Pipeline Run",
         union_mode="left_to_right",
     )
-    pipeline_name: Optional[str] = Field(
-        default=None,
-        description="Name of the pipeline associated with the run",
-    )
     workspace_id: Optional[Union[UUID, str]] = Field(
         default=None,
         description="Workspace of the Pipeline Run",
@@ -582,6 +643,11 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         description="Template used for the pipeline run.",
         union_mode="left_to_right",
     )
+    model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Model version associated with the pipeline run.",
+        union_mode="left_to_right",
+    )
     status: Optional[str] = Field(
         default=None,
         description="Name of the Pipeline Run",
@@ -597,7 +663,37 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         union_mode="left_to_right",
     )
     unlisted: Optional[bool] = None
-    templatable: Optional[bool] = None
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the run.",
+    )
+    # TODO: Remove once frontend is ready for it. This is replaced by the more
+    # generic `pipeline` filter below.
+    pipeline_name: Optional[str] = Field(
+        default=None,
+        description="Name of the pipeline associated with the run",
+    )
+    pipeline: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the pipeline associated with the run.",
+    )
+    stack: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the stack associated with the run.",
+    )
+    code_repository: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the code repository associated with the run.",
+    )
+    model: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the model associated with the run.",
+    )
+    templatable: Optional[bool] = Field(
+        default=None, description="Whether the run is templatable."
+    )
+
+    model_config = ConfigDict(protected_namespaces=())
 
     def get_custom_filters(
         self,
@@ -613,12 +709,16 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
 
         from zenml.zen_stores.schemas import (
             CodeReferenceSchema,
+            CodeRepositorySchema,
+            ModelSchema,
+            ModelVersionSchema,
             PipelineBuildSchema,
             PipelineDeploymentSchema,
             PipelineRunSchema,
             PipelineSchema,
             ScheduleSchema,
             StackSchema,
+            UserSchema,
         )
 
         if self.unlisted is not None:
@@ -628,19 +728,6 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
                 unlisted_filter = PipelineRunSchema.pipeline_id.is_not(None)  # type: ignore[union-attr]
             custom_filters.append(unlisted_filter)
 
-        if self.pipeline_name is not None:
-            value, filter_operator = self._resolve_operator(self.pipeline_name)
-            filter_ = StrFilter(
-                operation=GenericFilterOps(filter_operator),
-                column="name",
-                value=value,
-            )
-            pipeline_name_filter = and_(
-                PipelineRunSchema.pipeline_id == PipelineSchema.id,
-                filter_.generate_query_conditions(PipelineSchema),
-            )
-            custom_filters.append(pipeline_name_filter)
-
         if self.code_repository_id:
             code_repo_filter = and_(
                 PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
@@ -682,6 +769,70 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(run_template_filter)
 
+        if self.user:
+            user_filter = and_(
+                PipelineRunSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        if self.pipeline:
+            pipeline_filter = and_(
+                PipelineRunSchema.pipeline_id == PipelineSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.pipeline, table=PipelineSchema
+                ),
+            )
+            custom_filters.append(pipeline_filter)
+
+        if self.stack:
+            stack_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.stack,
+                    table=StackSchema,
+                ),
+            )
+            custom_filters.append(stack_filter)
+
+        if self.code_repository:
+            code_repo_filter = and_(
+                PipelineRunSchema.deployment_id == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.code_reference_id
+                == CodeReferenceSchema.id,
+                CodeReferenceSchema.code_repository_id
+                == CodeRepositorySchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.code_repository,
+                    table=CodeRepositorySchema,
+                ),
+            )
+            custom_filters.append(code_repo_filter)
+
+        if self.model:
+            model_filter = and_(
+                PipelineRunSchema.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)
+
+        if self.pipeline_name:
+            pipeline_name_filter = and_(
+                PipelineRunSchema.pipeline_id == PipelineSchema.id,
+                self.generate_custom_query_conditions_for_column(
+                    value=self.pipeline_name,
+                    table=PipelineSchema,
+                    column="name",
+                ),
+            )
+            custom_filters.append(pipeline_name_filter)
+
         if self.templatable is not None:
             if self.templatable is True:
                 templatable_filter = and_(

zenml/models/v2/core/run_template.py

@@ -299,7 +299,11 @@ class RunTemplateFilter(WorkspaceScopedTaggableFilter):
         *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "code_repository_id",
         "stack_id",
-        "build_id" "pipeline_id",
+        "build_id",
+        "pipeline_id",
+        "user",
+        "pipeline",
+        "stack",
     ]
 
     name: Optional[str] = Field(
@@ -336,6 +340,18 @@ class RunTemplateFilter(WorkspaceScopedTaggableFilter):
         description="Code repository associated with the template.",
         union_mode="left_to_right",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the template.",
+    )
+    pipeline: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the pipeline associated with the template.",
+    )
+    stack: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the stack associated with the template.",
+    )
 
     def get_custom_filters(
         self,
@@ -352,7 +368,10 @@ class RunTemplateFilter(WorkspaceScopedTaggableFilter):
         from zenml.zen_stores.schemas import (
             CodeReferenceSchema,
             PipelineDeploymentSchema,
+            PipelineSchema,
             RunTemplateSchema,
+            StackSchema,
+            UserSchema,
         )
 
         if self.code_repository_id:
@@ -390,4 +409,37 @@ class RunTemplateFilter(WorkspaceScopedTaggableFilter):
             )
             custom_filters.append(pipeline_filter)
 
+        if self.user:
+            user_filter = and_(
+                RunTemplateSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        if self.pipeline:
+            pipeline_filter = and_(
+                RunTemplateSchema.source_deployment_id
+                == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.pipeline_id == PipelineSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.pipeline,
+                    table=PipelineSchema,
+                ),
+            )
+            custom_filters.append(pipeline_filter)
+
+        if self.stack:
+            stack_filter = and_(
+                RunTemplateSchema.source_deployment_id
+                == PipelineDeploymentSchema.id,
+                PipelineDeploymentSchema.stack_id == StackSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.stack,
+                    table=StackSchema,
+                ),
+            )
+            custom_filters.append(stack_filter)
+
         return custom_filters

zenml/models/v2/core/stack.py

@@ -318,12 +318,11 @@ class StackFilter(WorkspaceScopedFilter):
     scoping.
     """
 
-    # `component_id` refers to a relationship through a link-table
-    #  rather than a field in the db, hence it needs to be handled
-    #  explicitly
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
         *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
-        "component_id",  # This is a relationship, not a field
+        "component_id",
+        "user",
+        "component",
     ]
 
     name: Optional[str] = Field(
@@ -348,6 +347,13 @@ class StackFilter(WorkspaceScopedFilter):
         description="Component in the stack",
         union_mode="left_to_right",
     )
+    user: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Name/ID of the user that created the stack.",
+    )
+    component: Optional[Union[UUID, str]] = Field(
+        default=None, description="Name/ID of a component in the stack."
+    )
 
     def get_custom_filters(self) -> List["ColumnElement[bool]"]:
         """Get custom filters.
@@ -357,9 +363,11 @@ class StackFilter(WorkspaceScopedFilter):
         """
         custom_filters = super().get_custom_filters()
 
-        from zenml.zen_stores.schemas.stack_schemas import (
+        from zenml.zen_stores.schemas import (
+            StackComponentSchema,
             StackCompositionSchema,
             StackSchema,
+            UserSchema,
         )
 
         if self.component_id:
@@ -369,4 +377,24 @@ class StackFilter(WorkspaceScopedFilter):
             )
             custom_filters.append(component_id_filter)
 
+        if self.user:
+            user_filter = and_(
+                StackSchema.user_id == UserSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.user, table=UserSchema
+                ),
+            )
+            custom_filters.append(user_filter)
+
+        if self.component:
+            component_filter = and_(
+                StackCompositionSchema.stack_id == StackSchema.id,
+                StackCompositionSchema.component_id == StackComponentSchema.id,
+                self.generate_name_or_id_query_conditions(
+                    value=self.component,
+                    table=StackComponentSchema,
+                ),
+            )
+            custom_filters.append(component_filter)
+
         return custom_filters

zenml/models/v2/core/step_run.py

@@ -536,3 +536,10 @@ class StepRunFilter(WorkspaceScopedFilter):
         description="Workspace of this step run",
         union_mode="left_to_right",
     )
+    model_version_id: Optional[Union[UUID, str]] = Field(
+        default=None,
+        description="Model version associated with the pipeline run.",
+        union_mode="left_to_right",
+    )
+
+    model_config = ConfigDict(protected_namespaces=())

zenml/new/pipelines/pipeline.py

@@ -84,6 +84,7 @@ from zenml.utils import (
     code_utils,
     dashboard_utils,
     dict_utils,
+    env_utils,
     pydantic_utils,
     settings_utils,
     source_utils,
@@ -1030,12 +1031,14 @@ To avoid this consider setting pipeline parameters only in one place (config or
 
         # Update with the values in code so they take precedence
         run_config = pydantic_utils.update_model(run_config, update=update)
+        run_config = env_utils.substitute_env_variable_placeholders(run_config)
 
         deployment = Compiler().compile(
             pipeline=self,
             stack=Client().active_stack,
             run_configuration=run_config,
         )
+        deployment = env_utils.substitute_env_variable_placeholders(deployment)
 
         return deployment, run_config.schedule, run_config.build
 
@@ -1252,6 +1255,7 @@ To avoid this consider setting pipeline parameters only in one place (config or
         if config_path:
             with open(config_path, "r") as f:
                 _from_config_file = yaml.load(f, Loader=yaml.SafeLoader)
+
             _from_config_file = dict_utils.remove_none_values(
                 {k: v for k, v in _from_config_file.items() if k in matcher}
             )

zenml/new/pipelines/run_utils.py

@@ -138,7 +138,10 @@ def deploy_pipeline(
     previous_value = constants.SHOULD_PREVENT_PIPELINE_EXECUTION
     constants.SHOULD_PREVENT_PIPELINE_EXECUTION = True
     try:
-        stack.deploy_pipeline(deployment=deployment)
+        stack.deploy_pipeline(
+            deployment=deployment,
+            placeholder_run=placeholder_run,
+        )
     except Exception as e:
         if (
             placeholder_run

zenml/orchestrators/base_orchestrator.py

@@ -14,12 +14,14 @@
 """Base orchestrator class."""
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, Optional, Type, cast
+from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Type, cast
 
 from pydantic import model_validator
 
-from zenml.enums import StackComponentType
+from zenml.enums import ExecutionStatus, StackComponentType
 from zenml.logger import get_logger
+from zenml.metadata.metadata_types import MetadataType
+from zenml.orchestrators.publish_utils import publish_pipeline_run_metadata
 from zenml.orchestrators.step_launcher import StepLauncher
 from zenml.orchestrators.utils import get_config_environment_vars
 from zenml.stack import Flavor, Stack, StackComponent, StackComponentConfig
@@ -27,7 +29,7 @@ from zenml.utils.pydantic_utils import before_validator_handler
 
 if TYPE_CHECKING:
     from zenml.config.step_configurations import Step
-    from zenml.models import PipelineDeploymentResponse
+    from zenml.models import PipelineDeploymentResponse, PipelineRunResponse
 
 logger = get_logger(__name__)
 
@@ -124,7 +126,7 @@ class BaseOrchestrator(StackComponent, ABC):
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
         environment: Dict[str, str],
-    ) -> Any:
+    ) -> Optional[Iterator[Dict[str, MetadataType]]]:
         """The method needs to be implemented by the respective orchestrator.
 
         Depending on the type of orchestrator you'll have to perform slightly
@@ -169,29 +171,41 @@ class BaseOrchestrator(StackComponent, ABC):
         self,
         deployment: "PipelineDeploymentResponse",
         stack: "Stack",
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Runs a pipeline on a stack.
 
         Args:
             deployment: The pipeline deployment.
             stack: The stack on which to run the pipeline.
-
-        Returns:
-            Orchestrator-specific return value.
+            placeholder_run: An optional placeholder run for the deployment.
+                This will be deleted in case the pipeline deployment failed.
         """
         self._prepare_run(deployment=deployment)
 
         environment = get_config_environment_vars(deployment=deployment)
 
         try:
-            result = self.prepare_or_run_pipeline(
-                deployment=deployment, stack=stack, environment=environment
-            )
+            if metadata_iterator := self.prepare_or_run_pipeline(
+                deployment=deployment,
+                stack=stack,
+                environment=environment,
+            ):
+                for metadata_dict in metadata_iterator:
+                    try:
+                        if placeholder_run:
+                            publish_pipeline_run_metadata(
+                                pipeline_run_id=placeholder_run.id,
+                                pipeline_run_metadata={self.id: metadata_dict},
+                            )
+                    except Exception as e:
+                        logger.debug(
+                            "Something went went wrong trying to publish the"
+                            f"run metadata: {e}"
+                        )
         finally:
             self._cleanup_run()
 
-        return result
-
     def run_step(self, step: "Step") -> None:
         """Runs the given step.
 
@@ -239,6 +253,21 @@ class BaseOrchestrator(StackComponent, ABC):
         """Cleans up the active run."""
         self._active_deployment = None
 
+    def fetch_status(self, run: "PipelineRunResponse") -> ExecutionStatus:
+        """Refreshes the status of a specific pipeline run.
+
+        Args:
+            run: A pipeline run response to fetch its status.
+
+        Raises:
+            NotImplementedError: If any orchestrator inheriting from the base
+                class does not implement this logic.
+        """
+        raise NotImplementedError(
+            "The fetch status functionality is not implemented for the "
+            f"'{self.__class__.__name__}' orchestrator."
+        )
+
 
 class BaseOrchestratorFlavor(Flavor):
     """Base orchestrator flavor class."""

zenml/stack/stack.py

@@ -62,7 +62,11 @@ if TYPE_CHECKING:
     from zenml.image_builders import BaseImageBuilder
     from zenml.model_deployers import BaseModelDeployer
     from zenml.model_registries import BaseModelRegistry
-    from zenml.models import PipelineDeploymentBase, PipelineDeploymentResponse
+    from zenml.models import (
+        PipelineDeploymentBase,
+        PipelineDeploymentResponse,
+        PipelineRunResponse,
+    )
     from zenml.orchestrators import BaseOrchestrator
     from zenml.stack import StackComponent
     from zenml.step_operators import BaseStepOperator
@@ -826,16 +830,21 @@ class Stack:
     def deploy_pipeline(
         self,
         deployment: "PipelineDeploymentResponse",
+        placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Any:
         """Deploys a pipeline on this stack.
 
         Args:
             deployment: The pipeline deployment.
+            placeholder_run: An optional placeholder run for the deployment.
+                This will be deleted in case the pipeline deployment failed.
 
         Returns:
             The return value of the call to `orchestrator.run_pipeline(...)`.
         """
-        return self.orchestrator.run(deployment=deployment, stack=self)
+        return self.orchestrator.run(
+            deployment=deployment, stack=self, placeholder_run=placeholder_run
+        )
 
     def _get_active_components_for_step(
         self, step_config: "StepConfiguration"

zenml/utils/env_utils.py

@@ -14,7 +14,16 @@
 """Utility functions for handling environment variables."""
 
 import os
-from typing import Dict, List, Optional, cast
+import re
+from typing import Any, Dict, List, Match, Optional, TypeVar, cast
+
+from zenml.logger import get_logger
+from zenml.utils import string_utils
+
+logger = get_logger(__name__)
+
+V = TypeVar("V", bound=Any)
+ENV_VARIABLE_PLACEHOLDER_PATTERN = re.compile(pattern=r"\$\{([a-zA-Z0-9_]+)\}")
 
 ENV_VAR_CHUNK_SUFFIX = "_CHUNK_"
 
@@ -99,3 +108,47 @@ def reconstruct_environment_variables(
         # Remove the chunk environment variables
         for key in chunk_keys:
             env.pop(key)
+
+
+def substitute_env_variable_placeholders(
+    value: V, raise_when_missing: bool = True
+) -> V:
+    """Substitute environment variable placeholders in an object.
+
+    Args:
+        value: The object in which to substitute the placeholders.
+        raise_when_missing: If True, an exception will be raised when an
+            environment variable is missing. Otherwise, a warning will be logged
+            instead.
+
+    Returns:
+        The object with placeholders substituted.
+    """
+
+    def _replace_with_env_variable_value(match: Match[str]) -> str:
+        key = match.group(1)
+        if key in os.environ:
+            return os.environ[key]
+        else:
+            if raise_when_missing:
+                raise KeyError(
+                    "Unable to substitute environment variable placeholder "
+                    f"'{key}' because the environment variable is not set."
+                )
+            else:
+                logger.warning(
+                    "Unable to substitute environment variable placeholder %s "
+                    "because the environment variable is not set, using an "
+                    "empty string instead.",
+                    key,
+                )
+                return ""
+
+    def _substitution_func(v: str) -> str:
+        return ENV_VARIABLE_PLACEHOLDER_PATTERN.sub(
+            _replace_with_env_variable_value, v
+        )
+
+    return string_utils.substitute_string(
+        value=value, substitution_func=_substitution_func
+    )