Track unused models and automatically delete them : issue-08

ersilia/cli/commands/__init__.py

@@ -4,7 +4,6 @@
 from ... import logger
 from ..echo import Silencer
 
-
 @click.group(cls=BentoMLCommandGroup)
 @click.version_option(version=__version__)
 @click.option(

ersilia/cli/commands/delete.py

@@ -1,4 +1,5 @@
 import click
+import csv
 
 from . import ersilia_cli
 from .. import echo
@@ -28,7 +29,7 @@ def delete(model):
             echo(
                 ":collision: Model {0} deleted successfully!".format(model_id),
                 fg="green",
-            )
+            )  
         else:
             echo(
                 ":person_tipping_hand: Model {0} is not available locally. No delete is necessary".format(

ersilia/cli/commands/fetch.py

@@ -1,4 +1,5 @@
 import click
+import time
 
 import os
 from . import ersilia_cli

ersilia/cli/commands/serve.py

@@ -1,6 +1,5 @@
 import click
 import os
-
 from ...utils import tmp_pid_file
 from . import ersilia_cli
 from .. import echo

ersilia/core/model.py

@@ -5,6 +5,8 @@
 import collections
 import importlib
 import __main__ as main
+import time
+import csv
 
 from .. import logger
 from .base import ErsiliaBase
@@ -321,6 +323,19 @@ def api_task(self, api_name, input, output, batch_size):
         else:
             # Result is a dict, a numpy array, a dataframe...
             return result
+
+    def update_model_usage_time(model_id):
+
+        ts_str = str(time.time())
+        with open("fetched_models.txt") as infile:
+            models = dict(csv.reader(infile))
+        infile.close()
+        if model_id in models.keys():
+            models[model_id] = ts_str
+
+        with open('fetched_models.txt', 'w') as f:
+            for key, values in models.items():
+                f.write(f"{key},{values}\n")
 
     def serve(self):
         self.close()
@@ -330,6 +345,7 @@ def serve(self):
         self.url = self.autoservice.service.url
         self.pid = self.autoservice.service.pid
         self.scl = self.autoservice._service_class
+        self.update_model_usage_time(self.model_id)
 
     def close(self):
         self.autoservice.close()

ersilia/hub/delete/delete.py

@@ -1,6 +1,9 @@
+
 import os
 import shutil
-
+import csv
+from pathlib import Path
+import os.path 
 from ... import ErsiliaBase
 from ...utils.terminal import run_command
 from ...utils.environment import Environment
@@ -218,6 +221,17 @@ def needs_delete(self, model_id):
             if v:
                 return True
         return False
+    def delete_model_entry(self,model_id):
+        # Remove deleted model's entry from fetched_models.txt
+        if (os.path.exists("fetched_models.txt")):
+            with open("fetched_models.txt") as infile:
+                models = dict(csv.reader(infile))
+            infile.close()
+            del models[model_id]
+            with open('fetched_models.txt', 'w') as f:
+                for key, values in models.items():
+                    f.write(f"{key},{values}\n")
+            print("Model Entry deleted")
 
     def delete(self, model_id):
         self.logger.info("Starting delete of model {0}".format(model_id))
@@ -230,4 +244,5 @@ def delete(self, model_id):
         ModelLakeDeleter(self.config_json).delete(model_id)
         ModelPipDeleter(self.config_json).delete(model_id)
         ModelDockerDeleter(self.config_json).delete(model_id)
+        self.delete_model_entry(model_id)
         self.logger.success("Model {0} deleted successfully".format(model_id))

ersilia/hub/fetch/fetch.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import time
 
 from ... import ErsiliaBase
 from ... import logger
@@ -80,6 +81,14 @@ def _success(self):
         with open(status_file, "w") as f:
             json.dump(done, f, indent=4)
 
+    def _fetchtime(self):
+        ts = time.time()
+        with open("fetched_models.txt","a") as file:
+            file.write(self.model_id)
+            file.write(',')
+            file.write(str(ts))
+            file.write('\n')
+
     def fetch(self, model_id):
         self.model_id = model_id
         self._prepare()
@@ -90,4 +99,5 @@ def fetch(self, model_id):
         self._check()
         self._sniff()
         self._success()
+        self._fetchtime()
         logger.info("Fetching {0} done successfully".format(model_id))

ersilia/utils/cron.py

@@ -0,0 +1,42 @@
+from ersilia.hub.delete.delete import ModelFullDeleter
+import os
+from pathlib import Path
+from ersilia.cli import echo
+import json
+import csv
+import time
+
+
+# in days :
+model_usage_lim   = 30
+model_cleanup_lim = 7
+
+def seconds_to_days(s):
+    return s / (24 * 3600) 
+
+def model_cleanup():
+    ts_dict = {}
+    my_file = Path("last_cleaned.json")
+    if not my_file.is_file():
+        ts_dict['timestamp'] = str(time.time())
+        with open('last_cleaned.json', 'w') as outfile:
+            json.dump(ts_dict, outfile)
+    else:
+        current_ts = time.time()
+        with open('last_cleaned.json') as json_file:
+            ts_dict = json.load(json_file)
+            ts = float(ts_dict['timestamp'])
+            if (seconds_to_days(current_ts - ts))>model_cleanup_lim: 
+                with open("fetched_models.txt") as infile:
+                    fetched_models = dict(csv.reader(infile))
+                for model_id in fetched_models: 
+                    if (seconds_to_days(current_ts-float(fetched_models[model_id])))>model_usage_lim:
+                        del_model(model_id)
+
+def del_model(model_id):
+        md = ModelFullDeleter()
+        if md.needs_delete(model_id):
+            print("Deleting model {0}".format(model_id))
+            md.delete(model_id)
+            print("Model ", model_id, " deleted successfully!")          
+