Feature/revamp model endpoints

src/zenml/cli/model.py

@@ -495,9 +495,12 @@ def delete_model_version(
             return
 
     try:
-        Client().delete_model_version(
+        model_version = Client().get_model_version(
             model_name_or_id=model_name_or_id,
-            model_version_name_or_id=model_version_name_or_number_or_id,
+            model_version_name_or_number_or_id=model_version_name_or_number_or_id,
+        )
+        Client().delete_model_version(
+            model_version_id=model_version.id,
         )
     except (KeyError, ValueError) as e:
         cli_utils.error(str(e))
@@ -551,9 +554,13 @@ def _print_artifacts_links_generic(
         f"{type_} linked to the model version `{model_version.name}[{model_version.number}]`:"
     )
 
+    model_version = Client().get_model_version(
+        model_name_or_id=model_name_or_id,
+        model_version_name_or_number_or_id=model_version_name_or_number_or_id,
+    )
+
     links = Client().list_model_version_artifact_links(
-        model_name_or_id=model_version.model.id,
-        model_version_name_or_number_or_id=model_version.id,
+        model_version_id=model_version.id,
         model_version_artifact_link_filter_model=ModelVersionArtifactFilterModel(
             only_artifacts=only_artifact_objects,
             only_deployments=only_deployments,
@@ -692,10 +699,13 @@ def list_model_version_pipeline_runs(
     cli_utils.title(
         f"Pipeline runs linked to the model version `{model_version.name}[{model_version.number}]`:"
     )
+    model_version = Client().get_model_version(
+        model_name_or_id=model_name_or_id,
+        model_version_name_or_number_or_id=model_version_name_or_number_or_id,
+    )
 
     links = Client().list_model_version_pipeline_run_links(
-        model_name_or_id=model_version.model.id,
-        model_version_name_or_number_or_id=model_version.id,
+        model_version_id=model_version.id,
         model_version_pipeline_run_link_filter_model=ModelVersionPipelineRunFilterModel(
             **kwargs,
         ),

src/zenml/client.py

@@ -57,6 +57,7 @@
     OAuthDeviceStatus,
     PermissionType,
     SecretScope,
+    SorterOps,
     StackComponentType,
     StoreType,
 )
@@ -5486,25 +5487,22 @@ def create_model_version(
 
     def delete_model_version(
         self,
-        model_name_or_id: Union[str, UUID],
-        model_version_name_or_id: Union[str, UUID],
+        model_version_id: UUID,
     ) -> None:
         """Deletes a model version from Model Control Plane.
 
         Args:
-            model_name_or_id: name or id of the model containing the model version.
-            model_version_name_or_id: name or id of the model version to be deleted.
+            model_version_id: Id of the model version to be deleted.
         """
         self.zen_store.delete_model_version(
-            model_name_or_id=model_name_or_id,
-            model_version_name_or_id=model_version_name_or_id,
+            model_version_id=model_version_id,
         )
 
     def get_model_version(
         self,
         model_name_or_id: Union[str, UUID],
         model_version_name_or_number_or_id: Optional[
-            Union[str, int, UUID, ModelStages]
+            Union[str, int, ModelStages, UUID]
         ] = None,
     ) -> ModelVersionResponseModel:
         """Get an existing model version from Model Control Plane.
@@ -5516,17 +5514,78 @@ def get_model_version(
 
         Returns:
             The model version of interest.
+
+        Raises:
+            RuntimeError: In case method inputs don't adhere to restrictions.
+            KeyError: In case no model version with the identifiers exists.
         """
-        return self.zen_store.get_model_version(
-            model_name_or_id=model_name_or_id,
-            model_version_name_or_number_or_id=model_version_name_or_number_or_id
-            or ModelStages.LATEST,
-        )
+        if model_version_name_or_number_or_id is None:
+            model_version_name_or_number_or_id = ModelStages.LATEST
+
+        if isinstance(model_version_name_or_number_or_id, UUID):
+            return self.zen_store.get_model_version(
+                model_version_id=model_version_name_or_number_or_id
+            )
+        elif isinstance(model_version_name_or_number_or_id, int):
+            model_versions = self.list_model_versions(
+                model_name_or_id=model_name_or_id,
+                model_version_filter_model=ModelVersionFilterModel(
+                    number=model_version_name_or_number_or_id
+                ),
+            ).items
+        elif isinstance(model_version_name_or_number_or_id, str):
+            if model_version_name_or_number_or_id == ModelStages.LATEST:
+                model_versions_page = self.list_model_versions(
+                    model_name_or_id=model_name_or_id,
+                    model_version_filter_model=ModelVersionFilterModel(
+                        sort_by=f"{SorterOps.DESCENDING}:number",
+                    ),
+                )
+
+                if model_versions_page.size > 0:
+                    model_versions = [model_versions_page.items[0]]
+                else:
+                    model_versions = []
+            elif model_version_name_or_number_or_id in ModelStages.values():
+                model_versions = self.list_model_versions(
+                    model_name_or_id=model_name_or_id,
+                    model_version_filter_model=ModelVersionFilterModel(
+                        stage=model_version_name_or_number_or_id
+                    ),
+                ).items
+            else:
+                model_versions = self.list_model_versions(
+                    model_name_or_id=model_name_or_id,
+                    model_version_filter_model=ModelVersionFilterModel(
+                        name=model_version_name_or_number_or_id
+                    ),
+                ).items
+        else:
+            raise RuntimeError(
+                f"The model version identifier "
+                f"`{model_version_name_or_number_or_id}` is not"
+                f"of the correct type."
+            )
+
+        if len(model_versions) == 1:
+            return model_versions[0]
+        elif len(model_versions) == 0:
+            raise KeyError(
+                f"No model version found for model "
+                f"`{model_name_or_id}` with version identifier "
+                f"`{model_version_name_or_number_or_id}`."
+            )
+        else:
+            raise RuntimeError(
+                f"The model version identifier "
+                f"`{model_version_name_or_number_or_id}` is not"
+                f"unique for model `{model_name_or_id}`."
+            )
 
     def list_model_versions(
         self,
-        model_name_or_id: Union[str, UUID],
         model_version_filter_model: ModelVersionFilterModel,
+        model_name_or_id: Optional[Union[str, UUID]] = None,
     ) -> Page[ModelVersionResponseModel]:
         """Get model versions by filter from Model Control Plane.
 
@@ -5570,28 +5629,21 @@ def update_model_version(
 
     def list_model_version_artifact_links(
         self,
-        model_name_or_id: Union[str, UUID],
         model_version_artifact_link_filter_model: ModelVersionArtifactFilterModel,
-        model_version_name_or_number_or_id: Union[str, int, UUID, ModelStages],
+        model_version_id: UUID,
     ) -> Page[ModelVersionArtifactResponseModel]:
         """Get model version to artifact links by filter in Model Control Plane.
 
         Args:
-            model_name_or_id: name or id of the model containing the model version.
-            model_version_name_or_number_or_id: name, id, stage or number of the model version to be retrieved.
+            model_version_id: id of the model version to be retrieved.
             model_version_artifact_link_filter_model: All filter parameters including pagination
                 params.
 
         Returns:
             A page of all model version to artifact links.
         """
-        mv = self.zen_store.get_model_version(
-            model_name_or_id=model_name_or_id,
-            model_version_name_or_number_or_id=model_version_name_or_number_or_id,
-        )
         return self.zen_store.list_model_version_artifact_links(
-            model_name_or_id=mv.model.id,
-            model_version_name_or_id=mv.id,
+            model_version_id=model_version_id,
             model_version_artifact_link_filter_model=model_version_artifact_link_filter_model,
         )
 
@@ -5603,28 +5655,21 @@ def list_model_version_artifact_links(
 
     def list_model_version_pipeline_run_links(
         self,
-        model_name_or_id: Union[str, UUID],
         model_version_pipeline_run_link_filter_model: ModelVersionPipelineRunFilterModel,
-        model_version_name_or_number_or_id: Union[str, int, UUID, ModelStages],
+        model_version_id: UUID,
     ) -> Page[ModelVersionPipelineRunResponseModel]:
         """Get all model version to pipeline run links by filter.
 
         Args:
-            model_name_or_id: name or id of the model containing the model version.
-            model_version_name_or_number_or_id: name, id, stage or number of the model version to be retrieved.
+            model_version_id: id of the model version to be retrieved.
             model_version_pipeline_run_link_filter_model: All filter parameters including pagination
                 params.
 
         Returns:
             A page of all model version to pipeline run links.
         """
-        mv = self.zen_store.get_model_version(
-            model_name_or_id=model_name_or_id,
-            model_version_name_or_number_or_id=model_version_name_or_number_or_id,
-        )
         return self.zen_store.list_model_version_pipeline_run_links(
-            model_name_or_id=mv.model.id,
-            model_version_name_or_id=mv.id,
+            model_version_id=model_version_id,
             model_version_pipeline_run_link_filter_model=model_version_pipeline_run_link_filter_model,
         )
 

src/zenml/model/artifact_config.py

@@ -143,8 +143,7 @@ def _link_to_model_version(
 
         # Create the model version artifact link using the ZenML client
         existing_links = client.list_model_version_artifact_links(
-            model_name_or_id=model_version.model.id,
-            model_version_name_or_number_or_id=model_version.id,
+            model_version_id=model_version.id,
             model_version_artifact_link_filter_model=ModelVersionArtifactFilterModel(
                 user_id=client.active_user.id,
                 workspace_id=client.active_workspace.id,
@@ -162,9 +161,9 @@ def _link_to_model_version(
                 logger.warning(
                     f"Existing artifact link(s) `{artifact_name}` found and will be deleted."
                 )
+
                 client.zen_store.delete_model_version_artifact_link(
-                    model_name_or_id=model_version.model.id,
-                    model_version_name_or_id=model_version.id,
+                    model_version_id=model_version.id,
                     model_version_artifact_link_name_or_id=artifact_name,
                 )
             else:

src/zenml/models/model_models.py

@@ -452,7 +452,7 @@ def _update_default_running_version_name(self) -> None:
 class ModelVersionFilterModel(ModelScopedFilterModel):
     """Filter Model for Model Version."""
 
-    name: Optional[Union[str, UUID]] = Field(
+    name: Optional[str] = Field(
         default=None,
         description="The name of the Model Version",
     )

src/zenml/new/pipelines/pipeline.py

@@ -945,14 +945,13 @@ def delete_running_versions_without_recovery(
                 new_version_request.model_config.delete_new_version_on_failure
                 and new_version_request.model_config.was_created_in_this_run
             ):
-                model = Client().get_model_version(
+                model_version_model = Client().get_model_version(
                     model_name_or_id=model_name,
                     model_version_name_or_number_or_id=model_version
                     or constants.RUNNING_MODEL_VERSION,
                 )
                 Client().delete_model_version(
-                    model_name_or_id=model_name,
-                    model_version_name_or_id=model.id,
+                    model_version_id=model_version_model.id
                 )
 
     def get_runs(self, **kwargs: Any) -> List[PipelineRunResponseModel]: