zenml-nightly 0.55.0.dev20240124__py3-none-any.whl → 0.72.0.dev20250116__py3-none-any.whl

zenml/model/model.py

@@ -23,11 +23,14 @@ from typing import (
 )
 from uuid import UUID
 
-from pydantic import BaseModel, PrivateAttr, root_validator
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator
 
 from zenml.enums import MetadataResourceTypes, ModelStages
 from zenml.exceptions import EntityExistsError
 from zenml.logger import get_logger
+from zenml.pipelines.pipeline_context import get_pipeline_context
+from zenml.utils.pydantic_utils import before_validator_handler
+from zenml.utils.string_utils import format_name_template
 
 if TYPE_CHECKING:
     from zenml.metadata.metadata_types import MetadataType
@@ -36,7 +39,6 @@ if TYPE_CHECKING:
         ModelResponse,
         ModelVersionResponse,
         PipelineRunResponse,
-        RunMetadataResponse,
     )
 
 logger = get_logger(__name__)
@@ -55,7 +57,9 @@ class Model(BaseModel):
     ethics: The ethical implications of the model.
     tags: Tags associated with the model.
     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.
+        to a specific version/stage. If skipped new version will be created. `version`
+        also supports placeholders: standard `{date}` and `{time}` and any custom placeholders
+        that are passed as substitutions in the pipeline or step decorators.
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
     """
@@ -69,15 +73,25 @@ class Model(BaseModel):
     trade_offs: Optional[str] = None
     ethics: Optional[str] = None
     tags: Optional[List[str]] = None
-    version: Optional[Union[ModelStages, int, str]] = None
+    version: Optional[Union[ModelStages, int, str]] = Field(
+        default=None, union_mode="smart"
+    )
     save_models_to_registry: bool = True
 
+    # technical attributes
+    model_version_id: Optional[UUID] = None
     suppress_class_validation_warnings: bool = False
-    was_created_in_this_run: bool = False
-
     _model_id: UUID = PrivateAttr(None)
-    _id: UUID = PrivateAttr(None)
-    _number: int = PrivateAttr(None)
+    _number: Optional[int] = PrivateAttr(None)
+    _created_model_version: bool = PrivateAttr(False)
+
+    # TODO: In Pydantic v2, the `model_` is a protected namespaces for all
+    #  fields defined under base models. If not handled, this raises a warning.
+    #  It is possible to suppress this warning message with the following
+    #  configuration, however the ultimate solution is to rename these fields.
+    #  Even though they do not cause any problems right now, if we are not
+    #  careful we might overwrite some fields protected by pydantic.
+    model_config = ConfigDict(protected_namespaces=())
 
     #########################
     #    Public methods     #
@@ -91,16 +105,21 @@ class Model(BaseModel):
                 doesn't exist and can only be read given current
                 config (you used stage name or number as
                 a version name).
+
+        Raises:
+            RuntimeError: if model version doesn't exist and
+                cannot be fetched from the Model Control Plane.
         """
-        if self._id is None:
+        if self.model_version_id is None:
             try:
-                self._get_or_create_model_version()
-            except RuntimeError:
-                logger.info(
-                    f"Version `{self.version}` of `{self.name}` model doesn't exist "
-                    "and cannot be fetched from the Model Control Plane."
-                )
-        return self._id
+                mv = self._get_or_create_model_version()
+                self.model_version_id = mv.id
+            except RuntimeError as e:
+                raise RuntimeError(
+                    f"Version `{self.version}` of `{self.name}` model doesn't "
+                    "exist and cannot be fetched from the Model Control Plane."
+                ) from e
+        return self.model_version_id
 
     @property
     def model_id(self) -> UUID:
@@ -122,15 +141,20 @@ class Model(BaseModel):
                 doesn't exist and can only be read given current
                 config (you used stage name or number as
                 a version name).
+
+        Raises:
+            KeyError: if model version doesn't exist and
+                cannot be fetched from the Model Control Plane.
         """
         if self._number is None:
             try:
-                self._get_or_create_model_version()
-            except RuntimeError:
-                logger.info(
-                    f"Version `{self.version}` of `{self.name}` model doesn't exist "
-                    "and cannot be fetched from the Model Control Plane."
-                )
+                mv = self._get_or_create_model_version()
+                self._number = mv.number
+            except RuntimeError as e:
+                raise KeyError(
+                    f"Version `{self.version}` of `{self.name}` model doesn't "
+                    "exist and cannot be fetched from the Model Control Plane."
+                ) from e
         return self._number
 
     @property
@@ -149,8 +173,8 @@ class Model(BaseModel):
                 return ModelStages(stage)
         except RuntimeError:
             logger.info(
-                f"Version `{self.version}` of `{self.name}` model doesn't exist "
-                "and cannot be fetched from the Model Control Plane."
+                f"Version `{self.version}` of `{self.name}` model doesn't "
+                "exist and cannot be fetched from the Model Control Plane."
             )
         return None
 
@@ -190,10 +214,12 @@ class Model(BaseModel):
 
         Args:
             name: The name of the artifact to retrieve.
-            version: The version of the artifact to retrieve (None for latest/non-versioned)
+            version: The version of the artifact to retrieve (None for
+                latest/non-versioned)
 
         Returns:
-            Specific version of the artifact or placeholder in the design time of the pipeline.
+            Specific version of the artifact or placeholder in the design time
+                of the pipeline.
         """
         if lazy := self._lazy_artifact_get(name, version):
             return lazy
@@ -212,10 +238,12 @@ class Model(BaseModel):
 
         Args:
             name: The name of the model artifact to retrieve.
-            version: The version of the model artifact to retrieve (None for latest/non-versioned)
+            version: The version of the model artifact to retrieve (None for
+                latest/non-versioned)
 
         Returns:
-            Specific version of the model artifact or placeholder in the design time of the pipeline.
+            Specific version of the model artifact or placeholder in the design
+                time of the pipeline.
         """
         if lazy := self._lazy_artifact_get(name, version):
             return lazy
@@ -234,10 +262,12 @@ class Model(BaseModel):
 
         Args:
             name: The name of the data artifact to retrieve.
-            version: The version of the data artifact to retrieve (None for latest/non-versioned)
+            version: The version of the data artifact to retrieve (None for
+                latest/non-versioned)
 
         Returns:
-            Specific version of the data artifact or placeholder in the design time of the pipeline.
+            Specific version of the data artifact or placeholder in the design
+            time of the pipeline.
         """
         if lazy := self._lazy_artifact_get(name, version):
             return lazy
@@ -256,10 +286,12 @@ class Model(BaseModel):
 
         Args:
             name: The name of the deployment artifact to retrieve.
-            version: The version of the deployment artifact to retrieve (None for latest/non-versioned)
+            version: The version of the deployment artifact to retrieve (None
+                for latest/non-versioned)
 
         Returns:
-            Specific version of the deployment artifact or placeholder in the design time of the pipeline.
+            Specific version of the deployment artifact or placeholder in the
+            design time of the pipeline.
         """
         if lazy := self._lazy_artifact_get(name, version):
             return lazy
@@ -287,7 +319,8 @@ class Model(BaseModel):
 
         Args:
             stage: the target stage for model version.
-            force: whether to force archiving of current model version in target stage or raise.
+            force: whether to force archiving of current model version in
+                target stage or raise.
         """
         self._get_or_create_model_version().set_stage(stage=stage, force=force)
 
@@ -303,16 +336,20 @@ class Model(BaseModel):
             metadata: The metadata to log.
         """
         from zenml.client import Client
+        from zenml.models import RunMetadataResource
 
         response = self._get_or_create_model_version()
         Client().create_run_metadata(
             metadata=metadata,
-            resource_id=response.id,
-            resource_type=MetadataResourceTypes.MODEL_VERSION,
+            resources=[
+                RunMetadataResource(
+                    id=response.id, type=MetadataResourceTypes.MODEL_VERSION
+                )
+            ],
         )
 
     @property
-    def run_metadata(self) -> Dict[str, "RunMetadataResponse"]:
+    def run_metadata(self) -> Dict[str, "MetadataType"]:
         """Get model version run metadata.
 
         Returns:
@@ -322,17 +359,13 @@ class Model(BaseModel):
             RuntimeError: If the model version run metadata cannot be fetched.
         """
         from zenml.metadata.lazy_load import RunMetadataLazyGetter
-        from zenml.new.pipelines.pipeline_context import (
-            get_pipeline_context,
-        )
 
         try:
             get_pipeline_context()
             # avoid exposing too much of internal details by keeping the return type
             return RunMetadataLazyGetter(  # type: ignore[return-value]
-                self,
-                None,
-                None,
+                self.name,
+                self._lazy_version,
             )
         except RuntimeError:
             pass
@@ -356,10 +389,12 @@ class Model(BaseModel):
 
         Args:
             name: The name of the artifact to delete.
-            version: The version of the artifact to delete (None for latest/non-versioned)
+            version: The version of the artifact to delete (None for
+                latest/non-versioned)
             only_link: Whether to only delete the link to the artifact.
             delete_metadata: Whether to delete the metadata of the artifact.
-            delete_from_artifact_store: Whether to delete the artifact from the artifact store.
+            delete_from_artifact_store: Whether to delete the artifact from the
+                artifact store.
         """
         from zenml.client import Client
         from zenml.models import ArtifactVersionResponse
@@ -387,7 +422,8 @@ class Model(BaseModel):
 
         Args:
             only_link: Whether to only delete the link to the artifact.
-            delete_from_artifact_store: Whether to delete the artifact from the artifact store.
+            delete_from_artifact_store: Whether to delete the artifact from
+                the artifact store.
         """
         from zenml.client import Client
 
@@ -407,21 +443,11 @@ class Model(BaseModel):
 
         client.delete_all_model_version_artifact_links(self.id, only_link)
 
-    #########################
-    #   Internal methods    #
-    #########################
-
-    class Config:
-        """Config class."""
-
-        smart_union = True
-
     def _lazy_artifact_get(
         self,
         name: str,
         version: Optional[str] = None,
     ) -> Optional["ArtifactVersionResponse"]:
-        from zenml import get_pipeline_context
         from zenml.models.v2.core.artifact_version import (
             LazyArtifactVersionResponse,
         )
@@ -429,11 +455,10 @@ class Model(BaseModel):
         try:
             get_pipeline_context()
             return LazyArtifactVersionResponse(
-                _lazy_load_name=name,
-                _lazy_load_version=version,
-                _lazy_load_model=Model(
-                    name=self.name, version=self.version or self.number
-                ),
+                lazy_load_name=name,
+                lazy_load_version=version,
+                lazy_load_model_name=self.name,
+                lazy_load_model_version=self._lazy_version,
             )
         except RuntimeError:
             pass
@@ -459,38 +484,33 @@ class Model(BaseModel):
         other_mv = other._get_or_create_model_version()
         return self_mv.id == other_mv.id
 
-    @root_validator(pre=True)
-    def _root_validator(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+    @model_validator(mode="before")
+    @classmethod
+    @before_validator_handler
+    def _root_validator(cls, data: Dict[str, Any]) -> Dict[str, Any]:
         """Validate all in one.
 
         Args:
-            values: Dict of values.
+            data: Dict of values.
 
         Returns:
             Dict of validated values.
+
+        Raises:
+            ValueError: If the model version id, but call is not internal.
         """
-        suppress_class_validation_warnings = values.get(
-            "suppress_class_validation_warnings", False
+        suppress_class_validation_warnings = data.get(
+            "suppress_class_validation_warnings",
+            False,
         )
-        version = values.get("version", None)
-
-        if (
-            version in [stage.value for stage in ModelStages]
-            and not suppress_class_validation_warnings
+        if not suppress_class_validation_warnings and data.get(
+            "model_version_id", None
         ):
-            logger.info(
-                f"`version` `{version}` matches one of the possible `ModelStages` and will be fetched using stage."
-            )
-        if str(version).isnumeric() and not suppress_class_validation_warnings:
-            logger.info(
-                f"`version` `{version}` is numeric and will be fetched using version number."
+            raise ValueError(
+                "`model_version_id` field is for internal use only"
             )
-        values["suppress_class_validation_warnings"] = True
-        return values
-
-    def _validate_config_in_runtime(self) -> None:
-        """Validate that config doesn't conflict with runtime environment."""
-        self._get_or_create_model_version()
+        data["suppress_class_validation_warnings"] = True
+        return data
 
     def _get_or_create_model(self) -> "ModelResponse":
         """This method should get or create a model from Model Control Plane.
@@ -504,96 +524,95 @@ class Model(BaseModel):
         from zenml.models import ModelRequest
 
         zenml_client = Client()
-        try:
-            model = zenml_client.zen_store.get_model(
-                model_name_or_id=self.name
+        # backup logic, if the Model class is used directly from the code
+        self.name = format_name_template(self.name)
+        if self.model_version_id:
+            mv = zenml_client.get_model_version(
+                model_version_name_or_number_or_id=self.model_version_id,
             )
-
-            difference: Dict[str, Any] = {}
-            for key in (
-                "license",
-                "audience",
-                "use_cases",
-                "limitations",
-                "trade_offs",
-                "ethics",
-                "save_models_to_registry",
-            ):
-                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 "
-                    f"existing model `{self.name}` with the "
-                    f"following changes: {difference}. If you want to "
-                    "update the model configuration, please use the "
-                    "`zenml model update` command."
-                )
-        except KeyError:
-            model_request = ModelRequest(
-                name=self.name,
-                license=self.license,
-                description=self.description,
-                audience=self.audience,
-                use_cases=self.use_cases,
-                limitations=self.limitations,
-                trade_offs=self.trade_offs,
-                ethics=self.ethics,
-                tags=self.tags,
-                user=zenml_client.active_user.id,
-                workspace=zenml_client.active_workspace.id,
-                save_models_to_registry=self.save_models_to_registry,
-            )
-            model_request = ModelRequest.parse_obj(model_request)
+            model = mv.model
+        else:
             try:
-                model = zenml_client.zen_store.create_model(
-                    model=model_request
-                )
-                logger.info(f"New model `{self.name}` was created implicitly.")
-            except EntityExistsError:
-                # this is backup logic, if model was created somehow in between get and create calls
-                pass
-            finally:
                 model = zenml_client.zen_store.get_model(
                     model_name_or_id=self.name
                 )
+            except KeyError:
+                model_request = ModelRequest(
+                    name=self.name,
+                    license=self.license,
+                    description=self.description,
+                    audience=self.audience,
+                    use_cases=self.use_cases,
+                    limitations=self.limitations,
+                    trade_offs=self.trade_offs,
+                    ethics=self.ethics,
+                    user=zenml_client.active_user.id,
+                    workspace=zenml_client.active_workspace.id,
+                    save_models_to_registry=self.save_models_to_registry,
+                )
+                model_request = ModelRequest.model_validate(model_request)
+                try:
+                    model = zenml_client.zen_store.create_model(
+                        model=model_request
+                    )
+                    logger.info(
+                        f"New model `{self.name}` was created implicitly."
+                    )
+                except EntityExistsError:
+                    model = zenml_client.zen_store.get_model(
+                        model_name_or_id=self.name
+                    )
+
         self._model_id = model.id
         return model
 
-    def _get_model_version(self) -> "ModelVersionResponse":
+    def _get_model_version(
+        self, hydrate: bool = True
+    ) -> "ModelVersionResponse":
         """This method gets a model version from Model Control Plane.
 
+        Args:
+            hydrate: Flag deciding whether to hydrate the output model(s)
+                by including metadata fields in the response.
+
         Returns:
             The model version based on configuration.
         """
         from zenml.client import Client
 
         zenml_client = Client()
-        mv = zenml_client.get_model_version(
-            model_name_or_id=self.name,
-            model_version_name_or_number_or_id=self.version,
-        )
-        if not self._id:
-            self._id = mv.id
+        if self.model_version_id:
+            mv = zenml_client.get_model_version(
+                model_version_name_or_number_or_id=self.model_version_id,
+                hydrate=hydrate,
+            )
+        else:
+            if self.version in ModelStages.values():
+                logger.info(
+                    f"Version `{self.version}` for model {self.name} matches "
+                    "one of the possible `ModelStages` and will be fetched "
+                    "using stage."
+                )
+            if str(self.version).isnumeric():
+                logger.info(
+                    f"Version `{self.version}` for model {self.name} is "
+                    "numeric and will be fetched using version number."
+                )
+
+            mv = zenml_client.get_model_version(
+                model_name_or_id=self.name,
+                model_version_name_or_number_or_id=self.version,
+                hydrate=hydrate,
+            )
+            self.model_version_id = mv.id
 
         difference: Dict[str, Any] = {}
-        if self.description and mv.description != self.description:
-            difference["description"] = {
-                "config": self.description,
-                "db": mv.description,
-            }
+        if mv.metadata:
+            if self.description and mv.description != self.description:
+                difference["description"] = {
+                    "config": self.description,
+                    "db": mv.description,
+                }
         if self.tags:
             configured_tags = set(self.tags)
             db_tags = {t.name for t in mv.tags}
@@ -615,15 +634,20 @@ class Model(BaseModel):
     ) -> "ModelVersionResponse":
         """This method should get or create a model and a model version from Model Control Plane.
 
-        A new model is created implicitly if missing, otherwise existing model is fetched. Model
-        name is controlled by the `name` parameter.
+        A new model is created implicitly if missing, otherwise existing model
+        is fetched. Model name is controlled by the `name` parameter.
 
         Model Version returned by this method is resolved based on model version:
-        - If `version` is None, a new model version is created, if not created by other steps in same run.
-        - If `version` is not None a model version will be fetched based on the version:
-            - If `version` is set to an integer or digit string, the model version with the matching number will be fetched.
-            - If `version` is set to a string, the model version with the matching version will be fetched.
-            - If `version` is set to a `ModelStage`, the model version with the matching stage will be fetched.
+        - If `version` is None, a new model version is created, if not created
+            by other steps in same run.
+        - If `version` is not None a model version will be fetched based on the
+            version:
+            - If `version` is set to an integer or digit string, the model
+                version with the matching number will be fetched.
+            - If `version` is set to a string, the model version with the
+                matching version will be fetched.
+            - If `version` is set to a `ModelStage`, the model version with the
+                matching stage will be fetched.
 
         Args:
             hydrate: Whether to return a hydrated version of the model version.
@@ -632,57 +656,21 @@ class Model(BaseModel):
             The model version based on configuration.
 
         Raises:
-            RuntimeError: if the model version needs to be created, but provided name is reserved
+            RuntimeError: if the model version needs to be created, but
+                provided name is reserved.
+            RuntimeError: if the model version cannot be created.
         """
         from zenml.client import Client
         from zenml.models import ModelVersionRequest
 
         model = self._get_or_create_model()
 
-        zenml_client = Client()
-        model_version_request = ModelVersionRequest(
-            user=zenml_client.active_user.id,
-            workspace=zenml_client.active_workspace.id,
-            name=self.version,
-            description=self.description,
-            model=model.id,
-            tags=self.tags,
-        )
-        mv_request = ModelVersionRequest.parse_obj(model_version_request)
+        # backup logic, if the Model class is used directly from the code
+        if isinstance(self.version, str):
+            self.version = format_name_template(self.version)
+
         try:
-            if not self.version:
-                try:
-                    from zenml import get_step_context
-
-                    context = get_step_context()
-                except RuntimeError:
-                    pass
-                else:
-                    # if inside a step context we loop over all
-                    # model version configuration to find, if the
-                    # 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
-                    if (
-                        pipeline_mv
-                        and pipeline_mv.was_created_in_this_run
-                        and pipeline_mv.name == self.name
-                        and pipeline_mv.version is not None
-                    ):
-                        self.version = pipeline_mv.version
-                    else:
-                        for step in context.pipeline_run.steps.values():
-                            step_mv = step.config.model
-                            if (
-                                step_mv
-                                and step_mv.was_created_in_this_run
-                                and step_mv.name == self.name
-                                and step_mv.version is not None
-                            ):
-                                self.version = step_mv.version
-                                break
-            if self.version:
+            if self.version or self.model_version_id:
                 model_version = self._get_model_version()
             else:
                 raise KeyError
@@ -698,7 +686,7 @@ class Model(BaseModel):
                     "model version in given stage exists. It might be missing, if "
                     "the pipeline promoting model version to this stage failed,"
                     " as an example. You can explore model versions using "
-                    f"`zenml model version list {self.name}` CLI command."
+                    f"`zenml model version list -n {self.name}` CLI command."
                 )
             if str(self.version).isnumeric():
                 raise RuntimeError(
@@ -708,32 +696,36 @@ class Model(BaseModel):
                     "model version with given number exists. It might be missing, if "
                     "the pipeline creating model version failed,"
                     " as an example. You can explore model versions using "
-                    f"`zenml model version list {self.name}` CLI command."
+                    f"`zenml model version list -n {self.name}` CLI command."
                 )
-            model_version = zenml_client.zen_store.create_model_version(
-                model_version=mv_request
+
+            client = Client()
+            model_version_request = ModelVersionRequest(
+                user=client.active_user.id,
+                workspace=client.active_workspace.id,
+                name=str(self.version) if self.version else None,
+                description=self.description,
+                model=model.id,
+                tags=self.tags,
+            )
+            model_version = client.zen_store.create_model_version(
+                model_version=model_version_request
             )
-            self.version = model_version.name
-            self.was_created_in_this_run = True
-            logger.info(f"New model version `{self.version}` was created.")
-        self._id = model_version.id
+
+            self._created_model_version = True
+
+            logger.info(
+                "Created new model version `%s` for model `%s`.",
+                model_version.name,
+                self.name,
+            )
+
+        self.version = model_version.name
+        self.model_version_id = model_version.id
         self._model_id = model_version.model.id
         self._number = model_version.number
         return model_version
 
-    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.tags))
-            )
-
     def __hash__(self) -> int:
         """Get hash of the `Model`.
 
@@ -751,3 +743,21 @@ class Model(BaseModel):
                 )
             )
         )
+
+    @property
+    def _lazy_version(self) -> Optional[str]:
+        """Get version name for lazy loader.
+
+        This getter ensures that new model version
+        creation is never triggered here.
+
+        Returns:
+            Version name or None if it was not set
+        """
+        if self._number is not None:
+            return str(self._number)
+        elif self.version is not None:
+            if isinstance(self.version, ModelStages):
+                return self.version.value
+            return str(self.version)
+        return None