Schedule Space restart to update list of models

app.py

@@ -1,4 +1,5 @@
 import gradio as gr
+from apscheduler.schedulers.background import BackgroundScheduler
 
 import src.constants as constants
 from src.details import (
@@ -12,6 +13,7 @@ from src.details import (
     update_task_description_component,
 )
 from src.env_impact import plot_env_impact
+from src.hub import restart_space
 from src.model_tree import load_model_tree
 from src.results import (
     clear_results,
@@ -298,4 +300,9 @@ with gr.Blocks(fill_height=True, fill_width=True) as demo:
         ],
     )
 
+# Start scheduler
+scheduler = BackgroundScheduler()
+scheduler.add_job(restart_space, "interval", hours=1)  # Restart every 1h
+scheduler.start()
+
 demo.launch()

requirements.txt

@@ -1,2 +1,2 @@
+apscheduler
 plotly
-

src/hub.py

@@ -1,9 +1,10 @@
 import asyncio
 import io
 import json
+import os
 
 import httpx
-from huggingface_hub import HfFileSystem, ModelCard, hf_hub_url
+from huggingface_hub import HfApi, HfFileSystem, ModelCard, hf_hub_url
 from huggingface_hub.utils import build_hf_headers
 
 import src.constants as constants
@@ -39,6 +40,7 @@ class Client:
                     return
 
 
+api = HfApi()
 client = Client()
 fs = HfFileSystem()
 
@@ -87,3 +89,9 @@ async def list_models(filtering=None):
     if r is None:
         return
     return r.json()
+
+
+def restart_space():
+    space_id = os.getenv("SPACE_ID")
+    if space_id:
+        api.restart_space(repo_id=space_id, token=os.getenv("HF_TOKEN"))