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

zenml/__init__.py

@@ -42,15 +42,19 @@ from zenml.artifacts.utils import (
     save_artifact,
     load_artifact,
 )
-from zenml.model.utils import log_model_metadata, link_artifact_to_model
+from zenml.model.utils import (
+    log_model_metadata,
+    link_artifact_to_model,
+    log_model_version_metadata,
+)
 from zenml.artifacts.artifact_config import ArtifactConfig
 from zenml.artifacts.external_artifact import ExternalArtifact
-from zenml.model.model_version import ModelVersion
-from zenml.model.utils import log_model_version_metadata
+from zenml.model.model import Model
 from zenml.new.pipelines.pipeline_context import get_pipeline_context
 from zenml.new.pipelines.pipeline_decorator import pipeline
 from zenml.new.steps.step_decorator import step
 from zenml.new.steps.step_context import get_step_context
+from zenml.steps.utils import log_step_metadata
 
 __all__ = [
     "ArtifactConfig",
@@ -60,8 +64,10 @@ __all__ = [
     "load_artifact",
     "log_artifact_metadata",
     "log_model_metadata",
+    "log_model_version_metadata",
+    "log_step_metadata",
+    "Model",
     "link_artifact_to_model",
-    "ModelVersion",
     "pipeline",
     "save_artifact",
     "show",

zenml/artifacts/artifact_config.py

@@ -22,7 +22,7 @@ from zenml.logger import get_logger
 from zenml.new.steps.step_context import get_step_context
 
 if TYPE_CHECKING:
-    from zenml.model.model_version import ModelVersion
+    from zenml.model.model import Model
 
 
 logger = get_logger(__name__)
@@ -52,7 +52,7 @@ class ArtifactConfig(BaseModel):
         version: The version of the artifact.
         tags: The tags of the artifact.
         model_name: The name of the model to link artifact to.
-        model_version: The identifier of the model version to link the artifact
+        model_version: The identifier of a version of the model to link the artifact
             to. It can be an exact version ("my_version"), exact version number
             (42), stage (ModelStages.PRODUCTION or "production"), or
             (ModelStages.LATEST or None) for the latest version (default).
@@ -88,26 +88,26 @@ class ArtifactConfig(BaseModel):
         smart_union = True
 
     @property
-    def _model_version(self) -> Optional["ModelVersion"]:
-        """The model version linked to this artifact.
+    def _model(self) -> Optional["Model"]:
+        """The model linked to this artifact.
 
         Returns:
-            The model version or None if the model version cannot be determined.
+            The model or None if the model version cannot be determined.
         """
         try:
-            model_version = get_step_context().model_version
+            model_ = get_step_context().model
         except (StepContextError, RuntimeError):
-            model_version = None
+            model_ = None
         # Check if another model name was specified
         if (self.model_name is not None) and (
-            model_version is None or model_version.name != self.model_name
+            model_ is None or model_.name != self.model_name
         ):
-            # Create a new ModelConfig instance with the provided model name and version
-            from zenml.model.model_version import ModelVersion
+            # Create a new Model instance with the provided model name and version
+            from zenml.model.model import Model
 
-            on_the_fly_config = ModelVersion(
+            on_the_fly_config = Model(
                 name=self.model_name, version=self.model_version
             )
             return on_the_fly_config
 
-        return model_version
+        return model_

zenml/artifacts/external_artifact.py

@@ -56,8 +56,8 @@ class ExternalArtifact(ExternalArtifactConfiguration):
             `version`, `pipeline_run_name`, or `pipeline_name` are set, the
             latest version of the artifact will be used.
         version: Version of the artifact to search. Only used when `name` is
-            provided. Cannot be used together with `model_version`.
-        model_version: The model version to search in. Only used when `name`
+            provided. Cannot be used together with `model`.
+        model: The model to search in. Only used when `name`
             is provided. Cannot be used together with `version`.
         materializer: The materializer to use for saving the artifact value
             to the artifact store. Only used when `value` is provided.
@@ -149,5 +149,5 @@ class ExternalArtifact(ExternalArtifactConfiguration):
             id=self.id,
             name=self.name,
             version=self.version,
-            model_version=self.model_version,
+            model=self.model,
         )

zenml/artifacts/external_artifact_config.py

@@ -18,7 +18,7 @@ from uuid import UUID
 from pydantic import BaseModel, root_validator
 
 from zenml.logger import get_logger
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.models.v2.core.artifact_version import ArtifactVersionResponse
 
 logger = get_logger(__name__)
@@ -33,13 +33,13 @@ class ExternalArtifactConfiguration(BaseModel):
     id: Optional[UUID] = None
     name: Optional[str] = None
     version: Optional[str] = None
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
 
     @root_validator
     def _validate_all_eac(cls, values: Dict[str, Any]) -> Dict[str, Any]:
-        if values.get("version", None) and values.get("model_version", None):
+        if values.get("version", None) and values.get("model", None):
             raise ValueError(
-                "Cannot provide both `version` and `model_version` when "
+                "Cannot provide both `version` and `model` when "
                 "creating an external artifact."
             )
         return values
@@ -67,13 +67,13 @@ class ExternalArtifactConfiguration(BaseModel):
                 response = client.get_artifact_version(
                     self.name, version=self.version
                 )
-            elif self.model_version:
-                response_ = self.model_version.get_artifact(self.name)
+            elif self.model:
+                response_ = self.model.get_artifact(self.name)
                 if not isinstance(response_, ArtifactVersionResponse):
                     raise RuntimeError(
                         f"Failed to pull artifact `{self.name}` from the Model "
-                        f"Version (name=`{self.model_version.name}`, version="
-                        f"`{self.model_version.version}`). Please validate the "
+                        f"(name=`{self.model.name}`, version="
+                        f"`{self.model.version}`). Please validate the "
                         "input and try again."
                     )
                 response = response_

zenml/artifacts/utils.py

@@ -233,14 +233,14 @@ def save_artifact(
                     saved_artifact_versions={name: response.id}
                 ),
             )
-            error_message = "model version"
-            model_version = step_context.model_version
-            if model_version:
+            error_message = "model"
+            model = step_context.model
+            if model:
                 from zenml.model.utils import link_artifact_to_model
 
                 link_artifact_to_model(
                     artifact_version_id=response.id,
-                    model_version=model_version,
+                    model=model,
                     is_model_artifact=is_model_artifact,
                     is_deployment_artifact=is_deployment_artifact,
                 )

zenml/cli/__init__.py

@@ -796,7 +796,7 @@ Administering your Models
 ----------------------------
 
 ZenML provides several CLI commands to help you administer your models and
-model versions as part of the Model Control Plane.
+their versions as part of the Model Control Plane.
 
 To register a new model, you can use the following CLI command:
 ```bash

zenml/cli/base.py

@@ -73,15 +73,15 @@ class ZenMLProjectTemplateLocation(BaseModel):
 ZENML_PROJECT_TEMPLATES = dict(
     e2e_batch=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-e2e-batch",
-        github_tag="2024.01.17",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.01.18",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     starter=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-starter",
-        github_tag="2023.12.18",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.01.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
     nlp=ZenMLProjectTemplateLocation(
         github_url="zenml-io/template-nlp",
-        github_tag="0.45.0",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
+        github_tag="2024.01.12",  # Make sure it is aligned with .github/workflows/update-templates-to-examples.yml
     ),
 )
 

zenml/cli/utils.py

@@ -58,8 +58,8 @@ from zenml.constants import (
 from zenml.enums import GenericFilterOps, StackComponentType
 from zenml.logger import get_logger
 from zenml.model_registries.base_model_registry import (
-    ModelVersion,
     RegisteredModel,
+    RegistryModelVersion,
 )
 from zenml.models import (
     BaseFilter,
@@ -1182,7 +1182,7 @@ def pretty_print_registered_model_table(
 
 
 def pretty_print_model_version_table(
-    model_versions: List["ModelVersion"],
+    model_versions: List["RegistryModelVersion"],
 ) -> None:
     """Given a list of model_versions, print all associated key-value pairs.
 
@@ -1206,7 +1206,7 @@ def pretty_print_model_version_table(
 
 
 def pretty_print_model_version_details(
-    model_version: "ModelVersion",
+    model_version: "RegistryModelVersion",
 ) -> None:
     """Given a model_version, print all associated key-value pairs.
 

zenml/config/compiler.py

@@ -204,7 +204,7 @@ class Compiler:
                 enable_step_logs=config.enable_step_logs,
                 settings=config.settings,
                 extra=config.extra,
-                model_version=config.model_version,
+                model=config.model,
                 parameters=config.parameters,
             )
 

zenml/config/pipeline_configurations.py

@@ -19,7 +19,7 @@ from pydantic import validator
 from zenml.config.constants import DOCKER_SETTINGS_KEY
 from zenml.config.source import Source, convert_source_validator
 from zenml.config.strict_base_model import StrictBaseModel
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 
 if TYPE_CHECKING:
     from zenml.config import DockerSettings
@@ -40,7 +40,7 @@ class PipelineConfigurationUpdate(StrictBaseModel):
     extra: Dict[str, Any] = {}
     failure_hook_source: Optional[Source] = None
     success_hook_source: Optional[Source] = None
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None
 
     _convert_source = convert_source_validator(

zenml/config/pipeline_run_configuration.py

@@ -19,7 +19,7 @@ from zenml.config.base_settings import BaseSettings
 from zenml.config.schedule import Schedule
 from zenml.config.step_configurations import StepConfigurationUpdate
 from zenml.config.strict_base_model import StrictBaseModel
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.models import PipelineBuildBase
 from zenml.utils import pydantic_utils
 
@@ -39,5 +39,5 @@ class PipelineRunConfiguration(
     steps: Dict[str, StepConfigurationUpdate] = {}
     settings: Dict[str, BaseSettings] = {}
     extra: Dict[str, Any] = {}
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
     parameters: Optional[Dict[str, Any]] = None

zenml/config/step_configurations.py

@@ -35,7 +35,7 @@ from zenml.config.source import Source, convert_source_validator
 from zenml.config.strict_base_model import StrictBaseModel
 from zenml.logger import get_logger
 from zenml.model.lazy_load import ModelVersionDataLazyLoader
-from zenml.model.model_version import ModelVersion
+from zenml.model.model import Model
 from zenml.utils import deprecation_utils
 
 if TYPE_CHECKING:
@@ -135,7 +135,7 @@ class StepConfigurationUpdate(StrictBaseModel):
     extra: Dict[str, Any] = {}
     failure_hook_source: Optional[Source] = None
     success_hook_source: Optional[Source] = None
-    model_version: Optional[ModelVersion] = None
+    model: Optional[Model] = None
 
     outputs: Mapping[str, PartialArtifactConfiguration] = {}
 

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",
@@ -525,7 +525,14 @@ class ModelVersion(BaseModel):
                             "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 "
@@ -588,7 +595,7 @@ class ModelVersion(BaseModel):
                 "db": mv.description,
             }
         if self.tags:
-            configured_tags = set(self.tags or [])
+            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)
@@ -656,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
@@ -666,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
@@ -714,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