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

zenml/integrations/mlflow/model_registries/mlflow_model_registry.py

@@ -19,7 +19,9 @@ from typing import Any, Dict, List, Optional, Tuple, cast
 
 import mlflow
 from mlflow import MlflowClient
-from mlflow.entities.model_registry import ModelVersion as MLflowModelVersion
+from mlflow.entities.model_registry import (
+    ModelVersion as MLflowModelVersion,
+)
 from mlflow.exceptions import MlflowException
 from mlflow.pyfunc import load_model
 
@@ -36,9 +38,9 @@ from zenml.logger import get_logger
 from zenml.model_registries.base_model_registry import (
     BaseModelRegistry,
     ModelRegistryModelMetadata,
-    ModelVersion,
     ModelVersionStage,
     RegisteredModel,
+    RegistryModelVersion,
 )
 from zenml.stack.stack import Stack
 from zenml.stack.stack_validator import StackValidator
@@ -348,7 +350,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
         description: Optional[str] = None,
         metadata: Optional[ModelRegistryModelMetadata] = None,
         **kwargs: Any,
-    ) -> ModelVersion:
+    ) -> RegistryModelVersion:
         """Register a model version to the MLflow model registry.
 
         Args:
@@ -442,7 +444,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
         metadata: Optional[ModelRegistryModelMetadata] = None,
         remove_metadata: Optional[List[str]] = None,
         stage: Optional[ModelVersionStage] = None,
-    ) -> ModelVersion:
+    ) -> RegistryModelVersion:
         """Update a model version in the MLflow model registry.
 
         Args:
@@ -521,7 +523,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
         self,
         name: str,
         version: str,
-    ) -> ModelVersion:
+    ) -> RegistryModelVersion:
         """Get a model version from the MLflow model registry.
 
         Args:
@@ -561,7 +563,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
         created_before: Optional[datetime] = None,
         order_by_date: Optional[str] = None,
         **kwargs: Any,
-    ) -> List[ModelVersion]:
+    ) -> List[RegistryModelVersion]:
         """List model versions from the MLflow model registry.
 
         Args:
@@ -704,7 +706,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
 
     def get_model_uri_artifact_store(
         self,
-        model_version: ModelVersion,
+        model_version: RegistryModelVersion,
     ) -> str:
         """Get the model URI artifact store.
 
@@ -723,7 +725,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
     def _cast_mlflow_version_to_model_version(
         self,
         mlflow_model_version: MLflowModelVersion,
-    ) -> ModelVersion:
+    ) -> RegistryModelVersion:
         """Cast an MLflow model version to a model version.
 
         Args:
@@ -748,7 +750,7 @@ class MLFlowModelRegistry(BaseModelRegistry):
             )
         except ImportError:
             model_library = None
-        return ModelVersion(
+        return RegistryModelVersion(
             registered_model=RegisteredModel(name=mlflow_model_version.name),
             model_format=MLFLOW_MODEL_FORMAT,
             model_library=model_library,

zenml/metadata/lazy_load.py

@@ -16,7 +16,7 @@
 from typing import TYPE_CHECKING, Optional
 
 if TYPE_CHECKING:
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
     from zenml.models import RunMetadataResponse
 
 
@@ -30,18 +30,18 @@ class RunMetadataLazyGetter:
 
     def __init__(
         self,
-        _lazy_load_model_version: "ModelVersion",
+        _lazy_load_model: "Model",
         _lazy_load_artifact_name: Optional[str],
         _lazy_load_artifact_version: Optional[str],
     ):
         """Initialize a RunMetadataLazyGetter.
 
         Args:
-            _lazy_load_model_version: The model version.
+            _lazy_load_model: The model version.
             _lazy_load_artifact_name: The artifact name.
             _lazy_load_artifact_version: The artifact version.
         """
-        self._lazy_load_model_version = _lazy_load_model_version
+        self._lazy_load_model = _lazy_load_model
         self._lazy_load_artifact_name = _lazy_load_artifact_name
         self._lazy_load_artifact_version = _lazy_load_artifact_version
 
@@ -57,7 +57,7 @@ class RunMetadataLazyGetter:
         from zenml.models.v2.core.run_metadata import LazyRunMetadataResponse
 
         return LazyRunMetadataResponse(
-            _lazy_load_model_version=self._lazy_load_model_version,
+            _lazy_load_model=self._lazy_load_model,
             _lazy_load_artifact_name=self._lazy_load_artifact_name,
             _lazy_load_artifact_version=self._lazy_load_artifact_version,
             _lazy_load_metadata_name=key,

zenml/model/lazy_load.py

@@ -17,7 +17,7 @@ from typing import Optional
 
 from pydantic import BaseModel
 
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 
 
 class ModelVersionDataLazyLoader(BaseModel):
@@ -28,7 +28,7 @@ class ModelVersionDataLazyLoader(BaseModel):
     model version during runtime time of the step.
     """
 
-    model_version: ModelVersion
+    model: Model
     artifact_name: Optional[str] = None
     artifact_version: Optional[str] = None
     metadata_name: Optional[str] = None

zenml/model/{model_version.py → model.py}

@@ -11,7 +11,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
 #  or implied. See the License for the specific language governing
 #  permissions and limitations under the License.
-"""ModelVersion user facing interface to pass into pipeline or step."""
+"""Model user facing interface to pass into pipeline or step."""
 
 from typing import (
     TYPE_CHECKING,
@@ -42,8 +42,8 @@ if TYPE_CHECKING:
 logger = get_logger(__name__)
 
 
-class ModelVersion(BaseModel):
-    """ModelVersion class to pass into pipeline or step to set it into a model context.
+class Model(BaseModel):
+    """Model class to pass into pipeline or step to set it into a model context.
 
     name: The name of the model.
     license: The license under which the model is created.
@@ -54,8 +54,8 @@ class ModelVersion(BaseModel):
     trade_offs: The tradeoffs of the model.
     ethics: The ethical implications of the model.
     tags: Tags associated with the model.
-    version: The model version name, number or stage is optional and points model context
-        to a specific version/stage. If skipped new model version will be created.
+    version: The version name, version number or stage is optional and points model context
+        to a specific version/stage. If skipped new version will be created.
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
     """
@@ -97,7 +97,7 @@ class ModelVersion(BaseModel):
                 self._get_or_create_model_version()
             except RuntimeError:
                 logger.info(
-                    f"Model version `{self.version}` doesn't exist "
+                    f"Version `{self.version}` of `{self.name}` model doesn't exist "
                     "and cannot be fetched from the Model Control Plane."
                 )
         return self._id
@@ -128,7 +128,7 @@ class ModelVersion(BaseModel):
                 self._get_or_create_model_version()
             except RuntimeError:
                 logger.info(
-                    f"Model version `{self.version}` doesn't exist "
+                    f"Version `{self.version}` of `{self.name}` model doesn't exist "
                     "and cannot be fetched from the Model Control Plane."
                 )
         return self._number
@@ -149,7 +149,7 @@ class ModelVersion(BaseModel):
                 return ModelStages(stage)
         except RuntimeError:
             logger.info(
-                f"Model version `{self.version}` doesn't exist "
+                f"Version `{self.version}` of `{self.name}` model doesn't exist "
                 "and cannot be fetched from the Model Control Plane."
             )
         return None
@@ -283,7 +283,7 @@ class ModelVersion(BaseModel):
     def set_stage(
         self, stage: Union[str, ModelStages], force: bool = False
     ) -> None:
-        """Sets this Model Version to a desired stage.
+        """Sets this Model to a desired stage.
 
         Args:
             stage: the target stage for model version.
@@ -431,7 +431,7 @@ class ModelVersion(BaseModel):
             return LazyArtifactVersionResponse(
                 _lazy_load_name=name,
                 _lazy_load_version=version,
-                _lazy_load_model_version=ModelVersion(
+                _lazy_load_model=Model(
                     name=self.name, version=self.version or self.number
                 ),
             )
@@ -441,7 +441,7 @@ class ModelVersion(BaseModel):
         return None
 
     def __eq__(self, other: object) -> bool:
-        """Check two ModelVersions for equality.
+        """Check two Models for equality.
 
         Args:
             other: object to compare with
@@ -449,7 +449,7 @@ class ModelVersion(BaseModel):
         Returns:
             True, if equal, False otherwise.
         """
-        if not isinstance(other, ModelVersion):
+        if not isinstance(other, Model):
             return NotImplemented
         if self.name != other.name:
             return False
@@ -509,7 +509,7 @@ class ModelVersion(BaseModel):
                 model_name_or_id=self.name
             )
 
-            difference = {}
+            difference: Dict[str, Any] = {}
             for key in (
                 "license",
                 "audience",
@@ -519,12 +519,20 @@ class ModelVersion(BaseModel):
                 "ethics",
                 "save_models_to_registry",
             ):
-                if getattr(self, key) != getattr(model, key):
-                    difference[key] = {
-                        "config": getattr(self, key),
-                        "db": getattr(model, key),
-                    }
-
+                if self_attr := getattr(self, key, None):
+                    if self_attr != getattr(model, key):
+                        difference[key] = {
+                            "config": getattr(self, key),
+                            "db": getattr(model, key),
+                        }
+            if self.tags:
+                configured_tags = set(self.tags)
+                db_tags = {t.name for t in model.tags}
+                if db_tags != configured_tags:
+                    difference["tags added"] = list(configured_tags - db_tags)
+                    difference["tags removed"] = list(
+                        db_tags - configured_tags
+                    )
             if difference:
                 logger.warning(
                     "Provided model configuration does not match "
@@ -581,16 +589,17 @@ class ModelVersion(BaseModel):
             self._id = mv.id
 
         difference: Dict[str, Any] = {}
-        if mv.description != self.description:
+        if self.description and mv.description != self.description:
             difference["description"] = {
                 "config": self.description,
                 "db": mv.description,
             }
-        configured_tags = set(self.tags or [])
-        db_tags = {t.name for t in mv.tags}
-        if db_tags != configured_tags:
-            difference["tags added"] = list(configured_tags - db_tags)
-            difference["tags removed"] = list(db_tags - configured_tags)
+        if self.tags:
+            configured_tags = set(self.tags)
+            db_tags = {t.name for t in mv.tags}
+            if db_tags != configured_tags:
+                difference["tags added"] = list(configured_tags - db_tags)
+                difference["tags removed"] = list(db_tags - configured_tags)
         if difference:
             logger.warning(
                 "Provided model version configuration does not match existing model "
@@ -654,7 +663,7 @@ class ModelVersion(BaseModel):
                     # model version for current model was already
                     # created in the current run, not to create
                     # new model versions
-                    pipeline_mv = context.pipeline_run.config.model_version
+                    pipeline_mv = context.pipeline_run.config.model
                     if (
                         pipeline_mv
                         and pipeline_mv.was_created_in_this_run
@@ -664,7 +673,7 @@ class ModelVersion(BaseModel):
                         self.version = pipeline_mv.version
                     else:
                         for step in context.pipeline_run.steps.values():
-                            step_mv = step.config.model_version
+                            step_mv = step.config.model
                             if (
                                 step_mv
                                 and step_mv.was_created_in_this_run
@@ -712,21 +721,21 @@ class ModelVersion(BaseModel):
         self._number = model_version.number
         return model_version
 
-    def _merge(self, model_version: "ModelVersion") -> None:
-        self.license = self.license or model_version.license
-        self.description = self.description or model_version.description
-        self.audience = self.audience or model_version.audience
-        self.use_cases = self.use_cases or model_version.use_cases
-        self.limitations = self.limitations or model_version.limitations
-        self.trade_offs = self.trade_offs or model_version.trade_offs
-        self.ethics = self.ethics or model_version.ethics
-        if model_version.tags is not None:
+    def _merge(self, model: "Model") -> None:
+        self.license = self.license or model.license
+        self.description = self.description or model.description
+        self.audience = self.audience or model.audience
+        self.use_cases = self.use_cases or model.use_cases
+        self.limitations = self.limitations or model.limitations
+        self.trade_offs = self.trade_offs or model.trade_offs
+        self.ethics = self.ethics or model.ethics
+        if model.tags is not None:
             self.tags = list(
-                {t for t in self.tags or []}.union(set(model_version.tags))
+                {t for t in self.tags or []}.union(set(model.tags))
             )
 
     def __hash__(self) -> int:
-        """Get hash of the `ModelVersion`.
+        """Get hash of the `Model`.
 
         Returns:
             Hash function results

zenml/model/utils.py

@@ -22,7 +22,7 @@ from zenml.enums import ModelStages
 from zenml.exceptions import StepContextError
 from zenml.logger import get_logger
 from zenml.metadata.metadata_types import MetadataType
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.models import ModelVersionArtifactRequest
 from zenml.new.steps.step_context import get_step_context
 
@@ -48,9 +48,9 @@ def link_step_artifacts_to_model(
             "step."
         )
     try:
-        model_version = step_context.model_version
+        model = step_context.model
     except StepContextError:
-        model_version = None
+        model = None
         logger.debug("No model context found, unable to auto-link artifacts.")
 
     for artifact_name, artifact_version_id in artifact_version_ids.items():
@@ -59,47 +59,47 @@ def link_step_artifacts_to_model(
         ).artifact_config
 
         # Implicit linking
-        if artifact_config is None and model_version is not None:
+        if artifact_config is None and model is not None:
             artifact_config = ArtifactConfig(name=artifact_name)
             logger.info(
                 f"Implicitly linking artifact `{artifact_name}` to model "
-                f"`{model_version.name}` version `{model_version.version}`."
+                f"`{model.name}` version `{model.version}`."
             )
 
         if artifact_config:
-            link_artifact_config_to_model_version(
+            link_artifact_config_to_model(
                 artifact_config=artifact_config,
                 artifact_version_id=artifact_version_id,
-                model_version=model_version,
+                model=model,
             )
 
 
-def link_artifact_config_to_model_version(
+def link_artifact_config_to_model(
     artifact_config: ArtifactConfig,
     artifact_version_id: UUID,
-    model_version: Optional["ModelVersion"] = None,
+    model: Optional["Model"] = None,
 ) -> None:
     """Link an artifact config to its model version.
 
     Args:
         artifact_config: The artifact config to link.
         artifact_version_id: The ID of the artifact to link.
-        model_version: The model version from the step or pipeline context.
+        model: The model version from the step or pipeline context.
     """
     client = Client()
 
     # If the artifact config specifies a model itself then always use that
     if artifact_config.model_name is not None:
-        from zenml.model.model_version import ModelVersion
+        from zenml.model.model import Model
 
-        model_version = ModelVersion(
+        model = Model(
             name=artifact_config.model_name,
             version=artifact_config.model_version,
         )
 
-    if model_version:
-        model_version._get_or_create_model_version()
-        model_version_response = model_version._get_model_version()
+    if model:
+        model._get_or_create_model_version()
+        model_version_response = model._get_model_version()
         request = ModelVersionArtifactRequest(
             user=client.active_user.id,
             workspace=client.active_workspace.id,
@@ -125,10 +125,10 @@ def log_model_version_metadata(
         metadata: The metadata to log.
         model_name: The name of the model to log metadata for. Can
             be omitted when being called inside a step with configured
-            `model_version` in decorator.
+            `model` in decorator.
         model_version: The version of the model to log metadata for. Can
             be omitted when being called inside a step with configured
-            `model_version` in decorator.
+            `model` in decorator.
     """
     logger.warning(
         "`log_model_version_metadata` is deprecated. Please use "
@@ -152,38 +152,38 @@ def log_model_metadata(
         metadata: The metadata to log.
         model_name: The name of the model to log metadata for. Can
             be omitted when being called inside a step with configured
-            `model_version` in decorator.
+            `model` in decorator.
         model_version: The version of the model to log metadata for. Can
             be omitted when being called inside a step with configured
-            `model_version` in decorator.
+            `model` in decorator.
 
     Raises:
         ValueError: If no model name/version is provided and the function is not
-            called inside a step with configured `model_version` in decorator.
+            called inside a step with configured `model` in decorator.
     """
     mv = None
     try:
         step_context = get_step_context()
-        mv = step_context.model_version
+        mv = step_context.model
     except RuntimeError:
         step_context = None
 
     if not step_context and not (model_name and model_version):
         raise ValueError(
             "Model name and version must be provided unless the function is "
-            "called inside a step with configured `model_version` in decorator."
+            "called inside a step with configured `model` in decorator."
         )
     if mv is None:
-        from zenml import ModelVersion
+        from zenml import Model
 
-        mv = ModelVersion(name=model_name, version=model_version)
+        mv = Model(name=model_name, version=model_version)
 
     mv.log_metadata(metadata)
 
 
 def link_artifact_to_model(
     artifact_version_id: UUID,
-    model_version: Optional["ModelVersion"] = None,
+    model: Optional["Model"] = None,
     is_model_artifact: bool = False,
     is_deployment_artifact: bool = False,
 ) -> None:
@@ -191,34 +191,34 @@ def link_artifact_to_model(
 
     Args:
         artifact_version_id: The ID of the artifact version.
-        model_version: The model version to link to.
+        model: The model to link to.
         is_model_artifact: Whether the artifact is a model artifact.
         is_deployment_artifact: Whether the artifact is a deployment artifact.
 
     Raises:
         RuntimeError: If called outside of a step.
     """
-    if not model_version:
+    if not model:
         is_issue = False
         try:
             step_context = get_step_context()
-            model_version = step_context.model_version
+            model = step_context.model
         except StepContextError:
             is_issue = True
 
-        if model_version is None or is_issue:
+        if model is None or is_issue:
             raise RuntimeError(
-                "`link_artifact_to_model` called without `model_version` parameter "
+                "`link_artifact_to_model` called without `model` parameter "
                 "and configured model context cannot be identified. Consider "
-                "passing the `model_version` explicitly or configuring it in "
+                "passing the `model` explicitly or configuring it in "
                 "@step or @pipeline decorator."
             )
 
-    link_artifact_config_to_model_version(
+    link_artifact_config_to_model(
         artifact_config=ArtifactConfig(
             is_model_artifact=is_model_artifact,
             is_deployment_artifact=is_deployment_artifact,
         ),
         artifact_version_id=artifact_version_id,
-        model_version=model_version,
+        model=model,
     )

zenml/model_registries/base_model_registry.py

@@ -123,10 +123,10 @@ class ModelRegistryModelMetadata(BaseModel):
         extra = "allow"
 
 
-class ModelVersion(BaseModel):
+class RegistryModelVersion(BaseModel):
     """Base class for all ZenML model versions.
 
-    The `ModelVersion` class represents a version or snapshot of a registered
+    The `RegistryModelVersion` class represents a version or snapshot of a registered
     model, including information such as the associated `ModelBundle`, version
     number, creation time, pipeline run information, and metadata. It serves as
     a blueprint for creating concrete model version implementations in a registry,
@@ -288,7 +288,7 @@ class BaseModelRegistry(StackComponent, ABC):
         description: Optional[str] = None,
         metadata: Optional[ModelRegistryModelMetadata] = None,
         **kwargs: Any,
-    ) -> ModelVersion:
+    ) -> RegistryModelVersion:
         """Registers a model version in the model registry.
 
         Args:
@@ -333,7 +333,7 @@ class BaseModelRegistry(StackComponent, ABC):
         metadata: Optional[ModelRegistryModelMetadata] = None,
         remove_metadata: Optional[List[str]] = None,
         stage: Optional[ModelVersionStage] = None,
-    ) -> ModelVersion:
+    ) -> RegistryModelVersion:
         """Updates a model version in the model registry.
 
         Args:
@@ -364,7 +364,7 @@ class BaseModelRegistry(StackComponent, ABC):
         created_before: Optional[datetime] = None,
         order_by_date: Optional[str] = None,
         **kwargs: Any,
-    ) -> Optional[List[ModelVersion]]:
+    ) -> Optional[List[RegistryModelVersion]]:
         """Lists all model versions for a registered model.
 
         Args:
@@ -387,7 +387,7 @@ class BaseModelRegistry(StackComponent, ABC):
         self,
         name: str,
         stage: Optional[ModelVersionStage] = None,
-    ) -> Optional[ModelVersion]:
+    ) -> Optional[RegistryModelVersion]:
         """Gets the latest model version for a registered model.
 
         This method is used to get the latest model version for a registered
@@ -410,7 +410,9 @@ class BaseModelRegistry(StackComponent, ABC):
         return None
 
     @abstractmethod
-    def get_model_version(self, name: str, version: str) -> ModelVersion:
+    def get_model_version(
+        self, name: str, version: str
+    ) -> RegistryModelVersion:
         """Gets a model version for a registered model.
 
         Args:
@@ -450,7 +452,7 @@ class BaseModelRegistry(StackComponent, ABC):
     @abstractmethod
     def get_model_uri_artifact_store(
         self,
-        model_version: ModelVersion,
+        model_version: RegistryModelVersion,
     ) -> str:
         """Gets the URI artifact store for a model version.
 

zenml/models/__init__.py

@@ -34,6 +34,7 @@ from zenml.models.v2.base.scoped import (
     WorkspaceScopedResponse,
     WorkspaceScopedResponseBody,
     WorkspaceScopedResponseMetadata,
+    WorkspaceScopedTaggableFilter
 )
 from zenml.models.v2.base.filter import (
     BaseFilter,
@@ -478,6 +479,7 @@ __all__ = [
     "WorkspaceScopedResponse",
     "WorkspaceScopedResponseBody",
     "WorkspaceScopedResponseMetadata",
+    "WorkspaceScopedTaggableFilter",
     "BaseFilter",
     "StrFilter",
     "BoolFilter",

zenml/models/v2/base/filter.py

@@ -264,6 +264,9 @@ class BaseFilter(BaseModel):
     # List of fields that are not even mentioned as options in the CLI.
     CLI_EXCLUDE_FIELDS: ClassVar[List[str]] = []
 
+    # List of fields that are wrapped with `fastapi.Query(default)` in API.
+    API_MULTI_INPUT_PARAMS: ClassVar[List[str]] = []
+
     sort_by: str = Field(
         default="created", description="Which column to sort by."
     )