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

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/v2/core/artifact_version.py

@@ -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,
@@ -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

@@ -31,7 +31,7 @@ from zenml.models.v2.base.scoped import (
 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
         ]
 

zenml/models/v2/core/model_version.py

@@ -32,7 +32,7 @@ from zenml.models.v2.base.scoped import (
 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,

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