From 570fcfd061c038940739cc6476baa3e674b74ccd Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Wed, 22 Dec 2021 13:09:23 +0200
Subject: [PATCH] Add Model.remove() to allow removing a model from the model
 repository

---
 clearml/model.py | 66 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/clearml/model.py b/clearml/model.py
index 09ebf7a5..2b1b99a4 100644
--- a/clearml/model.py
+++ b/clearml/model.py
@@ -21,6 +21,7 @@ from .storage.util import get_common_path
 from .utilities.enum import Options
 from .backend_interface import Task as _Task
 from .backend_interface.model import create_dummy_model, Model as _Model
+from .backend_interface.session import SendError
 from .config import running_remotely, get_cache_dir
 
 
@@ -512,6 +513,71 @@ class Model(BaseModel):
         # type: () -> str
         return self._base_model_id if self._base_model_id else super(Model, self).id
 
+    @classmethod
+    def remove(cls, model, delete_weights_file=True, force=False, raise_on_errors=False):
+        # type: (Union[str, Model], bool, bool, bool) -> bool
+        """
+        Remove a model from the model repository.
+        Optional, delete the model weights file from the remote storage.
+
+        :param model: Model ID or Model object to remove
+        :param delete_weights_file: If True (default) delete the weights file from the remote storage
+        :param force: If True, remove model even if other Tasks are using this model. default False.
+        :param raise_on_errors: If True, throw ValueError if something went wrong, default False.
+        :return: True if Model was removed successfully
+            partial removal returns False, i.e. Model was deleted but weights file deletion failed
+        """
+        if isinstance(model, str):
+            model = Model(model_id=model)
+
+        # noinspection PyBroadException
+        try:
+            weights_url = model.url
+        except Exception:
+            if raise_on_errors:
+                raise ValueError("Could not find model id={}".format(model.id))
+            return False
+
+        try:
+            # noinspection PyProtectedMember
+            res = _Model._get_default_session().send(
+                models.DeleteRequest(model.id, force=force),
+            )
+            response = res.wait()
+            if not response.ok():
+                if raise_on_errors:
+                    raise ValueError("Could not remove model id={}: {}".format(model.id, response.meta))
+                return False
+        except SendError as ex:
+            if raise_on_errors:
+                raise ValueError("Could not remove model id={}: {}".format(model.id, ex))
+            return False
+        except ValueError:
+            if raise_on_errors:
+                raise
+            return False
+        except Exception as ex:
+            if raise_on_errors:
+                raise ValueError("Could not remove model id={}: {}".format(model.id, ex))
+            return False
+
+        if not delete_weights_file:
+            return True
+
+        helper = StorageHelper.get(url=weights_url)
+        try:
+            if not helper.delete(weights_url):
+                if raise_on_errors:
+                    raise ValueError("Could not remove model id={} weights file: {}".format(model.id, weights_url))
+                return False
+        except Exception as ex:
+            if raise_on_errors:
+                raise ValueError("Could not remove model id={} weights file \'{}\': {}".format(
+                    model.id, weights_url, ex))
+            return False
+
+        return True
+
 
 class InputModel(Model):
     """