run garbage collector to delete empty models periodically

app.py

@@ -3,6 +3,8 @@ import pathlib
 import random
 import string
 import tempfile
+import time
+from concurrent.futures import ThreadPoolExecutor
 from typing import Iterable, List
 
 import gradio as gr
@@ -12,6 +14,8 @@ import yaml
 from gradio_logsview.logsview import Log, LogsView, LogsViewRunner
 from mergekit.config import MergeConfiguration
 
+from clean_community_org import garbage_collect_empty_models
+
 has_gpu = torch.cuda.is_available()
 
 # Running directly from Python doesn't work well with Gradio+run_process because of:
@@ -164,7 +168,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str) -> Iterable[List[Log]
             return
 
         # Set tmp HF_HOME to avoid filling up disk Space
-        tmp_env = os.environ.copy() # taken from https://stackoverflow.com/a/4453495
+        tmp_env = os.environ.copy()  # taken from https://stackoverflow.com/a/4453495
         tmp_env["HF_HOME"] = f"{tmpdirname}/.cache"
         yield from runner.run_command(cli.split(), cwd=merged_path, env=tmp_env)
 
@@ -215,4 +219,19 @@ with gr.Blocks() as demo:
 
     button.click(fn=merge, inputs=[config, token, repo_name], outputs=[logs])
 
+
+# Run garbage collection every hour to keep the community org clean.
+# Empty models might exists if the merge fails abruptly (e.g. if user leaves the Space).
+def _garbage_collect_every_hour():
+    while True:
+        try:
+            garbage_collect_empty_models(token=COMMUNITY_HF_TOKEN)
+        except Exception as e:
+            print("Error running garbage collection", e)
+        time.sleep(3600)
+
+
+pool = ThreadPoolExecutor()
+pool.submit(_garbage_collect_every_hour)
+
 demo.queue(default_concurrency_limit=1).launch()

clean_community_org.py

@@ -0,0 +1,36 @@
+# Script to delete empty models from the community org.
+# Can be run manually or scheduled to run periodically in the Space.
+# Usage: python clean_community_org.py
+#
+# 1. List models from https://huggingface.co/mergekit-community
+# 2. Filter out models with no files.
+# 3. Filter out models that are newer than 1 hour.
+# 4. Delete the remaining models.
+from datetime import datetime, timezone
+
+from huggingface_hub import HfApi
+
+
+def garbage_collect_empty_models(token: str | None = None):
+    api = HfApi(token=token)
+    now = datetime.now(timezone.utc)
+    print("Running garbage collection on mergekit-community.")
+    for model in api.list_models(author="mergekit-community", full=True):
+        if model.siblings and len(model.siblings) > 1:
+            # If model has files, then it's not empty
+            print("Skipping", model.modelId, "(not empty)")
+            continue
+        if (now - model.last_modified).total_seconds() < 3600:
+            # If model was updated in the last hour, then keep it
+            # to avoid deleting models that are being uploaded
+            print("Skipping", model.modelId, "(recently updated)")
+            continue
+        try:
+            print(f"Deleting {model.modelId}")
+            api.delete_repo(model.modelId, missing_ok=True)
+        except Exception as e:
+            print(f"Error deleting {model.modelId}: {e}")
+
+
+if __name__ == "__main__":
+    garbage_collect_empty_models()