zenml-nightly 0.54.1.dev20240118__py3-none-any.whl → 0.54.1.dev20240120__py3-none-any.whl

zenml/models/v2/base/scoped.py

@@ -23,10 +23,12 @@ from typing import (
     Optional,
     Type,
     TypeVar,
+    Union,
 )
 from uuid import UUID
 
 from pydantic import Field
+from sqlmodel import col
 
 from zenml.models.v2.base.base import (
     BaseRequest,
@@ -37,6 +39,8 @@ from zenml.models.v2.base.base import (
 from zenml.models.v2.base.filter import AnyQuery, BaseFilter
 
 if TYPE_CHECKING:
+    from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
+
     from zenml.models.v2.core.user import UserResponse
     from zenml.models.v2.core.workspace import WorkspaceResponse
     from zenml.zen_stores.schemas import BaseSchema
@@ -274,3 +278,58 @@ class WorkspaceScopedFilter(BaseFilter):
             query = query.where(scope_filter)
 
         return query
+
+
+class WorkspaceScopedTaggableFilter(WorkspaceScopedFilter):
+    """Model to enable advanced scoping with workspace and tagging."""
+
+    tag: Optional[str] = Field(
+        description="Tag to apply to the filter query.", default=None
+    )
+
+    FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
+        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        "tag",
+    ]
+
+    def apply_filter(
+        self,
+        query: AnyQuery,
+        table: Type["AnySchema"],
+    ) -> AnyQuery:
+        """Applies the filter to a query.
+
+        Args:
+            query: The query to which to apply the filter.
+            table: The query table.
+
+        Returns:
+            The query with filter applied.
+        """
+        from zenml.zen_stores.schemas import TagResourceSchema
+
+        query = super().apply_filter(query=query, table=table)
+        if self.tag:
+            query = (
+                query.join(getattr(table, "tags"))
+                .join(TagResourceSchema.tag)
+                .distinct()
+            )
+
+        return query
+
+    def get_custom_filters(
+        self,
+    ) -> List[Union["BinaryExpression[Any]", "BooleanClauseList[Any]"]]:
+        """Get custom tag filters.
+
+        Returns:
+            A list of custom filters.
+        """
+        from zenml.zen_stores.schemas import TagSchema
+
+        custom_filters = super().get_custom_filters()
+        if self.tag:
+            custom_filters.append(col(TagSchema.name) == self.tag)  # type: ignore[arg-type]
+
+        return custom_filters

zenml/models/v2/core/artifact.py

@@ -24,7 +24,7 @@ from zenml.models.v2.base.base import (
     BaseResponseBody,
     BaseResponseMetadata,
 )
-from zenml.models.v2.base.filter import BaseFilter
+from zenml.models.v2.base.scoped import WorkspaceScopedTaggableFilter
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
@@ -133,7 +133,7 @@ class ArtifactResponse(
 # ------------------ Filter Model ------------------
 
 
-class ArtifactFilter(BaseFilter):
+class ArtifactFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of artifacts."""
 
     name: Optional[str] = None

zenml/models/v2/core/artifact_version.py

@@ -32,11 +32,11 @@ from zenml.enums import ArtifactType, GenericFilterOps
 from zenml.logger import get_logger
 from zenml.models.v2.base.filter import StrFilter
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedTaggableFilter,
 )
 from zenml.models.v2.core.artifact import ArtifactResponse
 from zenml.models.v2.core.tag import TagResponse
@@ -44,7 +44,7 @@ from zenml.models.v2.core.tag import TagResponse
 if TYPE_CHECKING:
     from sqlalchemy.sql.elements import BinaryExpression, BooleanClauseList
 
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
     from zenml.models.v2.core.artifact_visualization import (
         ArtifactVisualizationRequest,
         ArtifactVisualizationResponse,
@@ -347,14 +347,14 @@ class ArtifactVersionResponse(
 # ------------------ Filter Model ------------------
 
 
-class ArtifactVersionFilter(WorkspaceScopedFilter):
+class ArtifactVersionFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of artifact versions."""
 
     # `name` and `only_unused` refer to properties related to other entities
     #  rather than a field in the db, hence they needs to be handled
     #  explicitly
     FILTER_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedFilter.FILTER_EXCLUDE_FIELDS,
+        *WorkspaceScopedTaggableFilter.FILTER_EXCLUDE_FIELDS,
         "name",
         "only_unused",
         "has_custom_name",
@@ -477,7 +477,7 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
     id: Optional[UUID] = None  # type: ignore[assignment]
     _lazy_load_name: Optional[str] = None
     _lazy_load_version: Optional[str] = None
-    _lazy_load_model_version: "ModelVersion"
+    _lazy_load_model: "Model"
 
     def get_body(self) -> None:  # type: ignore[override]
         """Protects from misuse of the lazy loader.
@@ -507,7 +507,7 @@ class LazyArtifactVersionResponse(ArtifactVersionResponse):
         from zenml.metadata.lazy_load import RunMetadataLazyGetter
 
         return RunMetadataLazyGetter(  # type: ignore[return-value]
-            self._lazy_load_model_version,
+            self._lazy_load_model,
             self._lazy_load_name,
             self._lazy_load_version,
         )

zenml/models/v2/core/model.py

@@ -22,16 +22,16 @@ from pydantic import BaseModel, Field
 
 from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedTaggableFilter,
 )
 from zenml.utils.pagination_utils import depaginate
 
 if TYPE_CHECKING:
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
     from zenml.models.v2.core.tag import TagResponse
 
 
@@ -310,7 +310,7 @@ class ModelResponse(
 
     # Helper functions
     @property
-    def versions(self) -> List["ModelVersion"]:
+    def versions(self) -> List["Model"]:
         """List all versions of the model.
 
         Returns:
@@ -323,7 +323,7 @@ class ModelResponse(
             partial(client.list_model_versions, model_name_or_id=self.id)
         )
         return [
-            mv.to_model_version(suppress_class_validation_warnings=True)
+            mv.to_model_class(suppress_class_validation_warnings=True)
             for mv in model_versions
         ]
 
@@ -331,7 +331,7 @@ class ModelResponse(
 # ------------------ Filter Model ------------------
 
 
-class ModelFilter(WorkspaceScopedFilter):
+class ModelFilter(WorkspaceScopedTaggableFilter):
     """Model to enable advanced filtering of all Workspaces."""
 
     name: Optional[str] = Field(
@@ -346,7 +346,7 @@ class ModelFilter(WorkspaceScopedFilter):
     )
 
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = [
-        *WorkspaceScopedFilter.CLI_EXCLUDE_FIELDS,
+        *WorkspaceScopedTaggableFilter.CLI_EXCLUDE_FIELDS,
         "workspace_id",
         "user_id",
     ]

zenml/models/v2/core/model_version.py

@@ -23,16 +23,16 @@ from zenml.constants import STR_FIELD_MAX_LENGTH, TEXT_FIELD_MAX_LENGTH
 from zenml.enums import ModelStages
 from zenml.models.v2.base.filter import AnyQuery
 from zenml.models.v2.base.scoped import (
-    WorkspaceScopedFilter,
     WorkspaceScopedRequest,
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedTaggableFilter,
 )
 from zenml.models.v2.core.tag import TagResponse
 
 if TYPE_CHECKING:
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
     from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
     from zenml.models.v2.core.model import ModelResponse
     from zenml.models.v2.core.pipeline_run import PipelineRunResponse
@@ -311,12 +311,12 @@ class ModelVersionResponse(
         return Client().zen_store.get_model_version(self.id)
 
     # Helper functions
-    def to_model_version(
+    def to_model_class(
         self,
         was_created_in_this_run: bool = False,
         suppress_class_validation_warnings: bool = False,
-    ) -> "ModelVersion":
-        """Convert response model to ModelVersion object.
+    ) -> "Model":
+        """Convert response model to Model object.
 
         Args:
             was_created_in_this_run: Whether model version was created during
@@ -325,11 +325,11 @@ class ModelVersionResponse(
                 repeated warnings.
 
         Returns:
-            ModelVersion object
+            Model object
         """
-        from zenml.model.model_version import ModelVersion
+        from zenml.model.model import Model
 
-        mv = ModelVersion(
+        mv = Model(
             name=self.model.name,
             license=self.model.license,
             description=self.description,
@@ -578,7 +578,7 @@ class ModelVersionResponse(
 # ------------------ Filter Model ------------------
 
 
-class ModelVersionFilter(WorkspaceScopedFilter):
+class ModelVersionFilter(WorkspaceScopedTaggableFilter):
     """Filter model for model versions."""
 
     name: Optional[str] = Field(

zenml/models/v2/core/run_metadata.py

@@ -30,7 +30,7 @@ from zenml.models.v2.base.scoped import (
 )
 
 if TYPE_CHECKING:
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
 
 # ------------------ Request Model ------------------
 
@@ -203,7 +203,7 @@ class LazyRunMetadataResponse(RunMetadataResponse):
     _lazy_load_artifact_name: Optional[str] = None
     _lazy_load_artifact_version: Optional[str] = None
     _lazy_load_metadata_name: Optional[str] = None
-    _lazy_load_model_version: "ModelVersion"
+    _lazy_load_model: "Model"
 
     def get_body(self) -> None:  # type: ignore[override]
         """Protects from misuse of the lazy loader.

zenml/new/pipelines/model_utils.py

@@ -17,14 +17,14 @@ from typing import List, Optional
 
 from pydantic import BaseModel, PrivateAttr
 
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 
 
-class NewModelVersionRequest(BaseModel):
-    """Request to create a new model version."""
+class NewModelRequest(BaseModel):
+    """Request to create a new version of a model."""
 
     class Requester(BaseModel):
-        """Requester of a new model version."""
+        """Requester of a new version of a model."""
 
         source: str
         name: str
@@ -38,35 +38,35 @@ class NewModelVersionRequest(BaseModel):
             return f"{self.source}::{self.name}"
 
     requesters: List[Requester] = []
-    _model_version: Optional[ModelVersion] = PrivateAttr(default=None)
+    _model: Optional[Model] = PrivateAttr(default=None)
 
     @property
-    def model_version(self) -> ModelVersion:
-        """Model version getter.
+    def model(self) -> Model:
+        """Model getter.
 
         Returns:
-            The model version.
+            The model.
 
         Raises:
-            RuntimeError: If the model version is not set.
+            RuntimeError: If the model is not set.
         """
-        if self._model_version is None:
-            raise RuntimeError("Model version is not set.")
-        return self._model_version
+        if self._model is None:
+            raise RuntimeError("Model is not set.")
+        return self._model
 
     def update_request(
         self,
-        model_version: ModelVersion,
-        requester: "NewModelVersionRequest.Requester",
+        model: Model,
+        requester: "NewModelRequest.Requester",
     ) -> None:
-        """Update from `ModelVersion` in place.
+        """Update from `Model` in place.
 
         Args:
-            model_version: `ModelVersion` to use.
-            requester: Requester of a new model version.
+            model: `Model` to use.
+            requester: Requester of a new version of a model.
         """
         self.requesters.append(requester)
-        if self._model_version is None:
-            self._model_version = model_version
+        if self._model is None:
+            self._model = model
 
-        self._model_version._merge(model_version)
+        self._model._merge(model)

zenml/new/pipelines/pipeline.py

@@ -70,7 +70,7 @@ from zenml.models import (
     ScheduleRequest,
 )
 from zenml.new.pipelines import build_utils
-from zenml.new.pipelines.model_utils import NewModelVersionRequest
+from zenml.new.pipelines.model_utils import NewModelRequest
 from zenml.orchestrators.utils import get_run_name
 from zenml.stack import Stack
 from zenml.steps import BaseStep
@@ -93,7 +93,7 @@ if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
     from zenml.config.source import Source
     from zenml.model.lazy_load import ModelVersionDataLazyLoader
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
 
     StepConfigurationUpdateOrDict = Union[
         Dict[str, Any], StepConfigurationUpdate
@@ -126,7 +126,7 @@ class Pipeline:
         extra: Optional[Dict[str, Any]] = None,
         on_failure: Optional["HookSpecification"] = None,
         on_success: Optional["HookSpecification"] = None,
-        model_version: Optional["ModelVersion"] = None,
+        model: Optional["Model"] = None,
     ) -> None:
         """Initializes a pipeline.
 
@@ -147,7 +147,7 @@ class Pipeline:
             on_success: Callback function in event of success of the step. Can
                 be a function with no arguments, or a source path to such a
                 function (e.g. `module.my_function`).
-            model_version: configuration of the model version in the Model Control Plane.
+            model: configuration of the model in the Model Control Plane.
         """
         self._invocations: Dict[str, StepInvocation] = {}
         self._run_args: Dict[str, Any] = {}
@@ -166,7 +166,7 @@ class Pipeline:
                 extra=extra,
                 on_failure=on_failure,
                 on_success=on_success,
-                model_version=model_version,
+                model=model,
             )
         self.entrypoint = entrypoint
         self._parameters: Dict[str, Any] = {}
@@ -305,7 +305,7 @@ class Pipeline:
         extra: Optional[Dict[str, Any]] = None,
         on_failure: Optional["HookSpecification"] = None,
         on_success: Optional["HookSpecification"] = None,
-        model_version: Optional["ModelVersion"] = None,
+        model: Optional["Model"] = None,
         parameters: Optional[Dict[str, Any]] = None,
         merge: bool = True,
     ) -> T:
@@ -341,7 +341,7 @@ class Pipeline:
                 configurations. If `False` the given configurations will
                 overwrite all existing ones. See the general description of this
                 method for an example.
-            model_version: configuration of the model version in the Model Control Plane.
+            model: configuration of the model version in the Model Control Plane.
             parameters: input parameters for the pipeline.
 
         Returns:
@@ -367,7 +367,7 @@ class Pipeline:
                 "extra": extra,
                 "failure_hook_source": failure_hook_source,
                 "success_hook_source": success_hook_source,
-                "model_version": model_version,
+                "model": model,
                 "parameters": parameters,
             }
         )
@@ -872,34 +872,32 @@ To avoid this consider setting pipeline parameters only in one place (config or
     def _update_new_requesters(
         self,
         requester_name: str,
-        model_version: "ModelVersion",
+        model: "Model",
         new_versions_requested: Dict[
-            Tuple[str, Optional[str]], NewModelVersionRequest
+            Tuple[str, Optional[str]], NewModelRequest
         ],
-        other_model_versions: Set["ModelVersion"],
+        other_models: Set["Model"],
     ) -> None:
         key = (
-            model_version.name,
-            str(model_version.version) if model_version.version else None,
+            model.name,
+            str(model.version) if model.version else None,
         )
-        if model_version.version is None:
+        if model.version is None:
             version_existed = False
         else:
             try:
-                model_version._get_model_version()
+                model._get_model_version()
                 version_existed = key not in new_versions_requested
             except KeyError:
                 version_existed = False
         if not version_existed:
-            model_version.was_created_in_this_run = True
+            model.was_created_in_this_run = True
             new_versions_requested[key].update_request(
-                model_version,
-                NewModelVersionRequest.Requester(
-                    source="step", name=requester_name
-                ),
+                model,
+                NewModelRequest.Requester(source="step", name=requester_name),
             )
         else:
-            other_model_versions.add(model_version)
+            other_models.add(model)
 
     def prepare_model_versions(
         self, deployment: "PipelineDeploymentBase"
@@ -910,35 +908,32 @@ To avoid this consider setting pipeline parameters only in one place (config or
             deployment: The pipeline deployment configuration.
         """
         new_versions_requested: Dict[
-            Tuple[str, Optional[str]], NewModelVersionRequest
-        ] = defaultdict(NewModelVersionRequest)
-        other_model_versions: Set["ModelVersion"] = set()
+            Tuple[str, Optional[str]], NewModelRequest
+        ] = defaultdict(NewModelRequest)
+        other_models: Set["Model"] = set()
         all_steps_have_own_config = True
         for step in deployment.step_configurations.values():
-            step_model_version = step.config.model_version
+            step_model = step.config.model
             all_steps_have_own_config = (
-                all_steps_have_own_config
-                and step.config.model_version is not None
+                all_steps_have_own_config and step.config.model is not None
             )
-            if step_model_version:
+            if step_model:
                 self._update_new_requesters(
-                    model_version=step_model_version,
+                    model=step_model,
                     requester_name=step.config.name,
                     new_versions_requested=new_versions_requested,
-                    other_model_versions=other_model_versions,
+                    other_models=other_models,
                 )
         if not all_steps_have_own_config:
-            pipeline_model_version = (
-                deployment.pipeline_configuration.model_version
-            )
-            if pipeline_model_version:
+            pipeline_model = deployment.pipeline_configuration.model
+            if pipeline_model:
                 self._update_new_requesters(
-                    model_version=pipeline_model_version,
+                    model=pipeline_model,
                     requester_name=self.name,
                     new_versions_requested=new_versions_requested,
-                    other_model_versions=other_model_versions,
+                    other_models=other_models,
                 )
-        elif deployment.pipeline_configuration.model_version is not None:
+        elif deployment.pipeline_configuration.model is not None:
             logger.warning(
                 f"ModelConfig of pipeline `{self.name}` is overridden in all "
                 f"steps. "
@@ -946,13 +941,13 @@ To avoid this consider setting pipeline parameters only in one place (config or
 
         self._validate_new_version_requests(new_versions_requested)
 
-        for other_model_version in other_model_versions:
-            other_model_version._validate_config_in_runtime()
+        for other_model in other_models:
+            other_model._validate_config_in_runtime()
 
     def _validate_new_version_requests(
         self,
         new_versions_requested: Dict[
-            Tuple[str, Optional[str]], NewModelVersionRequest
+            Tuple[str, Optional[str]], NewModelRequest
         ],
     ) -> None:
         """Validate the model version that are used in the pipeline run.
@@ -967,13 +962,13 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 logger.warning(
                     f"New version of model version `{model_name}::{model_version or 'NEW'}` "
                     f"requested in multiple decorators:\n{data.requesters}\n We recommend "
-                    "that `ModelVersion` requesting new version is configured only in one "
+                    "that `Model` requesting new version is configured only in one "
                     "place of the pipeline."
                 )
-            data.model_version._validate_config_in_runtime()
+            data.model._validate_config_in_runtime()
             self.__new_unnamed_model_versions_in_current_run__[
-                data.model_version.name
-            ] = data.model_version.number
+                data.model.name
+            ] = data.model.number
 
     def get_runs(self, **kwargs: Any) -> List["PipelineRunResponse"]:
         """(Deprecated) Get runs of this pipeline.
@@ -1400,18 +1395,16 @@ To avoid this consider setting pipeline parameters only in one place (config or
                 {k: v for k, v in _from_config_file.items() if k in matcher}
             )
 
-            if "model_version" in _from_config_file:
-                if "model_version" in self._from_config_file:
-                    _from_config_file[
-                        "model_version"
-                    ] = self._from_config_file["model_version"]
+            if "model" in _from_config_file:
+                if "model" in self._from_config_file:
+                    _from_config_file["model"] = self._from_config_file[
+                        "model"
+                    ]
                 else:
-                    from zenml.model.model_version import ModelVersion
+                    from zenml.model.model import Model
 
-                    _from_config_file[
-                        "model_version"
-                    ] = ModelVersion.parse_obj(
-                        _from_config_file["model_version"]
+                    _from_config_file["model"] = Model.parse_obj(
+                        _from_config_file["model"]
                     )
         self._from_config_file = _from_config_file
 

zenml/new/pipelines/pipeline_context.py

@@ -107,4 +107,4 @@ class PipelineContext:
         self.enable_step_logs = pipeline_configuration.enable_step_logs
         self.settings = pipeline_configuration.settings
         self.extra = pipeline_configuration.extra
-        self.model_version = pipeline_configuration.model_version
+        self.model = pipeline_configuration.model

zenml/new/pipelines/pipeline_decorator.py

@@ -26,7 +26,7 @@ from typing import (
 
 if TYPE_CHECKING:
     from zenml.config.base_settings import SettingsOrDict
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
     from zenml.new.pipelines.pipeline import Pipeline
 
     HookSpecification = Union[str, FunctionType]
@@ -62,7 +62,7 @@ def pipeline(
     extra: Optional[Dict[str, Any]] = None,
     on_failure: Optional["HookSpecification"] = None,
     on_success: Optional["HookSpecification"] = None,
-    model_version: Optional["ModelVersion"] = None,
+    model: Optional["Model"] = None,
 ) -> Union["Pipeline", Callable[["F"], "Pipeline"]]:
     """Decorator to create a pipeline.
 
@@ -81,7 +81,7 @@ def pipeline(
         on_success: Callback function in event of success of the step. Can be a
             function with no arguments, or a source path to such a function
             (e.g. `module.my_function`).
-        model_version: configuration of the model version in the Model Control Plane.
+        model: configuration of the model in the Model Control Plane.
 
     Returns:
         A pipeline instance.
@@ -99,7 +99,7 @@ def pipeline(
             extra=extra,
             on_failure=on_failure,
             on_success=on_success,
-            model_version=model_version,
+            model=model,
             entrypoint=func,
         )