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

zenml/models/v2/core/pipeline_run.py

@@ -28,8 +28,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,
@@ -522,6 +521,11 @@ class PipelineRunFilter(WorkspaceScopedTaggableFilter):
         "schedule_id",
         "stack_id",
         "template_id",
+        "user",
+        "pipeline",
+        "stack",
+        "code_repository",
+        "model",
         "pipeline_name",
         "templatable",
     ]
@@ -538,10 +542,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 +582,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 +602,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 +648,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 +667,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 +708,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/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
+    )

zenml/utils/string_utils.py

@@ -15,13 +15,17 @@
 
 import base64
 import datetime
+import functools
 import random
 import string
+from typing import Any, Callable, Dict, TypeVar, cast
 
 from pydantic import BaseModel
 
 from zenml.constants import BANNED_NAME_CHARACTERS
 
+V = TypeVar("V", bound=Any)
+
 
 def get_human_readable_time(seconds: float) -> str:
     """Convert seconds into a human-readable string.
@@ -167,3 +171,49 @@ def format_name_template(
         datetime.datetime.now(datetime.timezone.utc).strftime("%H_%M_%S_%f"),
     )
     return name_template.format(**kwargs)
+
+
+def substitute_string(value: V, substitution_func: Callable[[str], str]) -> V:
+    """Recursively substitute strings in objects.
+
+    Args:
+        value: An object in which the strings should be recursively substituted.
+            This can be a pydantic model, dict, set, list, tuple or any
+            primitive type.
+        substitution_func: The function that does the actual string
+            substitution.
+
+    Returns:
+        The object with the substitution function applied to all string values.
+    """
+    substitute_ = functools.partial(
+        substitute_string, substitution_func=substitution_func
+    )
+
+    if isinstance(value, BaseModel):
+        model_values = {}
+
+        for k, v in value.__iter__():
+            new_value = substitute_(v)
+
+            if k not in value.model_fields_set and new_value == getattr(
+                value, k
+            ):
+                # This is a default value on the model and was not set
+                # explicitly. In this case, we don't include it in the model
+                # values to keep the `exclude_unset` behavior the same
+                continue
+
+            model_values[k] = new_value
+
+        return cast(V, type(value).model_validate(model_values))
+    elif isinstance(value, Dict):
+        return cast(
+            V, {substitute_(k): substitute_(v) for k, v in value.items()}
+        )
+    elif isinstance(value, (list, set, tuple)):
+        return cast(V, type(value)(substitute_(v) for v in value))
+    elif isinstance(value, str):
+        return cast(V, substitution_func(value))
+
+    return value

zenml/zen_stores/sql_zen_store.py

@@ -973,6 +973,7 @@ class SqlZenStore(BaseZenStore):
             ValueError: if the filtered page number is out of bounds.
             RuntimeError: if the schema does not have a `to_model` method.
         """
+        query = query.distinct()
         query = filter_model.apply_filter(query=query, table=table)
         query = query.distinct()
 

{zenml_nightly-0.66.0.dev20240923.dist-info → zenml_nightly-0.66.0.dev20240924.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zenml-nightly
-Version: 0.66.0.dev20240923
+Version: 0.66.0.dev20240924
 Summary: ZenML: Write production-ready ML code.
 Home-page: https://zenml.io
 License: Apache-2.0