[Model WatchTower] Delete running versions on failed pipelines with delete_new_version_on_failure option

src/zenml/model/model_config.py

@@ -38,7 +38,7 @@ class ModelConfig(ModelConfigModel):
     create_new_model_version: Whether to create a new model version during execution
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
-    recovery: Whether to keep failed runs with new versions for later recovery from it.
+    delete_new_version_on_failure: Whether to delete failed runs with new versions for later recovery from it.
     """
 
     _model: Optional["ModelResponseModel"] = PrivateAttr(default=None)
@@ -61,16 +61,16 @@ def _validate_create_new_model_version(
             values["version"] = RUNNING_MODEL_VERSION
         return create_new_model_version
 
-    @validator("recovery")
+    @validator("delete_new_version_on_failure")
     def _validate_recovery(
-        cls, recovery: bool, values: Dict[str, Any]
+        cls, delete_new_version_on_failure: bool, values: Dict[str, Any]
     ) -> bool:
-        if recovery:
+        if not delete_new_version_on_failure:
             if not values.get("create_new_model_version", False):
                 logger.warning(
-                    "Using `recovery` flag without `create_new_model_version=True` has no effect."
+                    "Using `delete_new_version_on_failure=False` without `create_new_model_version=True` has no effect."
                 )
-        return recovery
+        return delete_new_version_on_failure
 
     @validator("version")
     def _validate_version(
@@ -157,15 +157,12 @@ def _create_model_version(
         )
         mv_request = ModelVersionRequestModel.parse_obj(model_version_request)
         try:
-            self._model_version = zenml_client.zen_store.get_model_version(
+            mv = zenml_client.zen_store.get_model_version(
                 model_name_or_id=self.name,
                 model_version_name_or_id=self.version,
             )
+            self._model_version = mv
         except KeyError:
-            if self.recovery:
-                logger.warning(
-                    f"Recovery mode: No `{self.version}` model version found."
-                )
             self._model_version = zenml_client.zen_store.create_model_version(
                 model_version=mv_request
             )
@@ -201,10 +198,18 @@ def _get_model_version(
     def get_or_create_model_version(self) -> "ModelVersionResponseModel":
         """This method should get or create a model and a model version from Model WatchTower.
 
-        New model is created implicitly, if missing, otherwise fetched.
-
-        New version will be created if `create_new_model_version`, otherwise
-        will try to fetch based on `model_version`.
+        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 configuration:
+        - If there is an existing model version leftover from the previous failed run with
+        `delete_new_version_on_failure` is set to False and `create_new_model_version` is True,
+        leftover model version will be reused.
+        - Otherwise if `create_new_model_version` is True, a new model version is created.
+        - If `create_new_model_version` is False a model version will be fetched based on the version:
+            - If `version` is not set, the latest model version 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.
 
         Returns:
             The model version based on configuration.

src/zenml/models/model_base_model.py

@@ -68,7 +68,7 @@ class ModelConfigModel(ModelBaseModel):
     create_new_model_version: Whether to create a new model version during execution
     save_models_to_registry: Whether to save all ModelArtifacts to Model Registry,
         if available in active stack.
-    recovery: Whether to keep new model versions from failed runs for later recovery.
+    delete_new_version_on_failure: Whether to delete failed runs with new versions for later recovery from it.
     """
 
     version: Optional[Union[ModelStages, str]] = Field(
@@ -78,4 +78,4 @@ class ModelConfigModel(ModelBaseModel):
     )
     create_new_model_version: bool = False
     save_models_to_registry: bool = True
-    recovery: bool = False
+    delete_new_version_on_failure: bool = True

src/zenml/models/model_models.py

@@ -68,15 +68,15 @@ class ModelVersionResponseModel(
     model: "ModelResponseModel" = Field(
         title="The model containing version",
     )
-    model_object_ids: Dict[str, UUID] = Field(
+    model_object_ids: Dict[str, Dict[str, UUID]] = Field(
         title="Model Objects linked to the model version",
         default={},
     )
-    artifact_object_ids: Dict[str, UUID] = Field(
+    artifact_object_ids: Dict[str, Dict[str, UUID]] = Field(
         title="Artifacts linked to the model version",
         default={},
     )
-    deployment_ids: Dict[str, UUID] = Field(
+    deployment_ids: Dict[str, Dict[str, UUID]] = Field(
         title="Deployments linked to the model version",
         default={},
     )
@@ -86,45 +86,54 @@ class ModelVersionResponseModel(
     )
 
     @property
-    def model_objects(self) -> Dict[str, ArtifactResponseModel]:
-        """Get all model objects linked to this version.
+    def model_objects(self) -> Dict[str, Dict[str, ArtifactResponseModel]]:
+        """Get all model objects linked to this model version.
 
         Returns:
-            Dictionary of Model Objects as ArtifactResponseModel
+            Dictionary of Model Objects with versions as Dict[str, ArtifactResponseModel]
         """
         from zenml.client import Client
 
         return {
-            name: Client().get_artifact(a)
-            for name, a in self.model_object_ids.items()
+            name: {
+                version: Client().get_artifact(a)
+                for version, a in self.model_object_ids[name].items()
+            }
+            for name in self.model_object_ids
         }
 
     @property
-    def artifact_objects(self) -> Dict[str, ArtifactResponseModel]:
-        """Get all artifacts linked to this version.
+    def artifacts(self) -> Dict[str, Dict[str, ArtifactResponseModel]]:
+        """Get all artifacts linked to this model version.
 
         Returns:
-            Dictionary of Artifact Objects as ArtifactResponseModel
+            Dictionary of Artifacts with versions as Dict[str, ArtifactResponseModel]
         """
         from zenml.client import Client
 
         return {
-            name: Client().get_artifact(a)
-            for name, a in self.artifact_object_ids.items()
+            name: {
+                version: Client().get_artifact(a)
+                for version, a in self.artifact_object_ids[name].items()
+            }
+            for name in self.artifact_object_ids
         }
 
     @property
-    def deployments(self) -> Dict[str, ArtifactResponseModel]:
-        """Get all deployments linked to this version.
+    def deployments(self) -> Dict[str, Dict[str, ArtifactResponseModel]]:
+        """Get all deployments linked to this model version.
 
         Returns:
-            Dictionary of Deployments as ArtifactResponseModel
+            Dictionary of Deployments with versions as Dict[str, ArtifactResponseModel]
         """
         from zenml.client import Client
 
         return {
-            name: Client().get_artifact(a)
-            for name, a in self.deployment_ids.items()
+            name: {
+                version: Client().get_artifact(a)
+                for version, a in self.deployment_ids[name].items()
+            }
+            for name in self.deployment_ids
         }
 
     @property
@@ -141,44 +150,59 @@ def pipeline_runs(self) -> Dict[str, PipelineRunResponseModel]:
             for name, pr in self.pipeline_run_ids.items()
         }
 
-    def get_model_object(self, name: str) -> ArtifactResponseModel:
-        """Get model object linked to this version.
+    def get_model_object(
+        self, name: str, version: Optional[str] = None
+    ) -> ArtifactResponseModel:
+        """Get model object linked to this model version.
 
         Args:
             name: The name of the model object to retrieve.
+            version: The version of the model object to retrieve (None for latest/non-versioned)
 
         Returns:
-            Model Object as ArtifactResponseModel
+            Specific version of Model Object
         """
         from zenml.client import Client
 
-        return Client().get_artifact(self.model_object_ids[name])
+        if version is None:
+            version = max(self.model_object_ids[name].keys())
+        return Client().get_artifact(self.model_object_ids[name][version])
 
-    def get_artifact_object(self, name: str) -> ArtifactResponseModel:
-        """Get artifact linked to this version.
+    def get_artifact_object(
+        self, name: str, version: Optional[str] = None
+    ) -> ArtifactResponseModel:
+        """Get artifact linked to this model version.
 
         Args:
             name: The name of the artifact to retrieve.
+            version: The version of the model object to retrieve (None for latest/non-versioned)
 
         Returns:
-            Artifact Object as ArtifactResponseModel
+            Specific version of Artifact
         """
         from zenml.client import Client
 
-        return Client().get_artifact(self.artifact_object_ids[name])
+        if version is None:
+            version = max(self.artifact_object_ids[name].keys())
+        return Client().get_artifact(self.artifact_object_ids[name][version])
 
-    def get_deployment(self, name: str) -> ArtifactResponseModel:
-        """Get deployment linked to this version.
+    def get_deployment(
+        self, name: str, version: Optional[str] = None
+    ) -> ArtifactResponseModel:
+        """Get deployment linked to this model version.
 
         Args:
             name: The name of the deployment to retrieve.
+            version: The version of the model object to retrieve (None for latest/non-versioned)
 
         Returns:
-            Deployment as ArtifactResponseModel
+            Specific version of Deployment
         """
         from zenml.client import Client
 
-        return Client().get_artifact(self.deployment_ids[name])
+        if version is None:
+            version = max(self.deployment_ids[name].keys())
+        return Client().get_artifact(self.deployment_ids[name][version])
 
     def get_pipeline_run(self, name: str) -> PipelineRunResponseModel:
         """Get pipeline run linked to this version.