Update

app.py

@@ -2,79 +2,22 @@
 
 from __future__ import annotations
 
-import os
-
 import gradio as gr
-import pandas as pd
-from apscheduler.schedulers.background import BackgroundScheduler
-from huggingface_hub import HfApi
 
+from constants import (HARDWARE_CHOICES, OWNER_CHOICES, SDK_CHOICES,
+                       SLEEP_TIME_CHOICES, STATUS_CHOICES, VISIBILITY_CHOICES)
 from demo_list import DemoList
+from restart_scheduler import RestartScheduler
+from settings import HF_TOKEN, INTERVAL_MINUTES, SPACE_ID
 
 demo_list = DemoList()
 
-api = HfApi()
-WHOAMI = api.whoami()['name']
-if (SPACE_ID := os.getenv('SPACE_ID')) is not None:
-    INTERVAL_MIN = int(os.getenv('INTERVAL_MIN', '10'))
-    scheduler = BackgroundScheduler()
-    scheduler.add_job(func=lambda: api.restart_space(SPACE_ID),
-                      trigger='interval',
-                      seconds=60 * INTERVAL_MIN)
+if SPACE_ID is not None and INTERVAL_MINUTES > 0:
+    scheduler = RestartScheduler(space_id=SPACE_ID,
+                                 interval_minutes=INTERVAL_MINUTES,
+                                 hf_token=HF_TOKEN)
     scheduler.start()
 
-STATUS_CHOICES = [
-    'RUNNING',
-    'PAUSED',
-    'STOPPED',
-    'RUNTIME_ERROR',
-    'BUILD_ERROR',
-    'BUILDING',
-]
-HARDWARE_CHOICES = [
-    'cpu-basic',
-    'cpu-upgrade',
-    't4-small',
-    't4-medium',
-    'zero-a10g',
-    'a10g-small',
-    'a10g-large',
-    'a100-large',
-]
-SDK_CHOICES = [
-    'gradio',
-    'streamlit',
-    'docker',
-]
-SLEEP_TIME_CHOICES = list(demo_list.TO_TIME_STR.values())
-SLEEP_TIME_STR_TO_INT = {v: k for k, v in demo_list.TO_TIME_STR.items()}
-OWNER_CHOICES = [WHOAMI, 'other organizations']
-
-
-def update_df(status: list[str], hardware: list[str], sdk: list[str],
-              sleep_time: list[str], owner: list[str],
-              multiple_replicas: bool) -> pd.DataFrame:
-    df_raw = demo_list.df_raw
-    df = demo_list.df
-
-    if multiple_replicas:
-        df = df[df_raw.replicas > 1]
-
-    df = df[(df_raw.status.isin(status)) & (df_raw.hardware.isin(hardware)) &
-            (df_raw.sdk.isin(sdk))]
-
-    sleep_time_int = [SLEEP_TIME_STR_TO_INT[s] for s in sleep_time]
-    df = df[df_raw.sleep_time.isin(sleep_time_int)]
-
-    if set(owner) == set(OWNER_CHOICES):
-        pass
-    elif WHOAMI in owner:
-        df = df[df_raw.owner == WHOAMI]
-    else:
-        df = df[df_raw.owner != WHOAMI]
-
-    return df
-
 
 def update_status_checkboxes(choices: list[str]) -> list[str]:
     if '(ALL)' in choices:
@@ -123,20 +66,24 @@ with gr.Blocks(css='style.css') as demo:
                                     HARDWARE_CHOICES,
                                     value=HARDWARE_CHOICES,
                                     type='value')
-        sdk = gr.CheckboxGroup(label='SDK',
-                               choices=['(ALL)', '(NONE)'] + SDK_CHOICES,
-                               value=SDK_CHOICES,
-                               type='value')
         sleep_time = gr.CheckboxGroup(label='Sleep time',
                                       choices=['(ALL)', '(NONE)'] +
                                       SLEEP_TIME_CHOICES,
                                       value=SLEEP_TIME_CHOICES,
                                       type='value')
+        multiple_replicas = gr.Checkbox(label='Multiple replicas', value=False)
+        sdk = gr.CheckboxGroup(label='SDK',
+                               choices=['(ALL)', '(NONE)'] + SDK_CHOICES,
+                               value=SDK_CHOICES,
+                               type='value')
+        visibility = gr.CheckboxGroup(label='Visibility',
+                                      choices=VISIBILITY_CHOICES,
+                                      value=VISIBILITY_CHOICES,
+                                      type='value')
         owner = gr.CheckboxGroup(label='Owner',
                                  choices=OWNER_CHOICES,
                                  value=OWNER_CHOICES,
                                  type='value')
-        multiple_replicas = gr.Checkbox(label='Multiple replicas', value=False)
         apply_button = gr.Button('Apply')
     df = gr.Dataframe(value=demo_list.df,
                       datatype=demo_list.column_datatype,
@@ -166,14 +113,15 @@ with gr.Blocks(css='style.css') as demo:
                       queue=False,
                       show_progress=False,
                       api_name=False)
-    apply_button.click(fn=update_df,
+    apply_button.click(fn=demo_list.apply_filter,
                        inputs=[
                            status,
                            hardware,
-                           sdk,
                            sleep_time,
-                           owner,
                            multiple_replicas,
+                           sdk,
+                           visibility,
+                           owner,
                        ],
                        outputs=df,
                        api_name=False)

constants.py

@@ -0,0 +1,47 @@
+from huggingface_hub import HfApi
+
+STATUS_CHOICES = [
+    'RUNNING',
+    'PAUSED',
+    'STOPPED',
+    'RUNTIME_ERROR',
+    'BUILD_ERROR',
+    'BUILDING',
+    'RUNNING_BUILDING',
+    'NO_APP_FILE',
+]
+HARDWARE_CHOICES = [
+    'cpu-basic',
+    'cpu-upgrade',
+    't4-small',
+    't4-medium',
+    'zero-a10g',
+    'a10g-small',
+    'a10g-large',
+    'a100-large',
+]
+SDK_CHOICES = [
+    'gradio',
+    'streamlit',
+    'docker',
+]
+SLEEP_TIME_INT_TO_STR = {
+    -1: 'null',
+    300: '5 minutes',
+    900: '15 minutes',
+    1800: '30 minutes',
+    3600: '1 hour',
+    36000: '10 hours',
+    86400: '24 hours',
+    172800: '48 hours',
+    259200: '72 hours',
+    604800: '1 week',
+}
+SLEEP_TIME_CHOICES = list(SLEEP_TIME_INT_TO_STR.values())
+SLEEP_TIME_STR_TO_INT = {v: k for k, v in SLEEP_TIME_INT_TO_STR.items()}
+
+VISIBILITY_CHOICES = ['public', 'private']
+
+api = HfApi()
+WHOAMI = api.whoami()['name']
+OWNER_CHOICES = [WHOAMI, 'other organizations']

demo_list.py

@@ -7,6 +7,9 @@ import tqdm.auto
 import yaml
 from huggingface_hub import HfApi
 
+from constants import (OWNER_CHOICES, SLEEP_TIME_INT_TO_STR,
+                       SLEEP_TIME_STR_TO_INT, WHOAMI)
+
 repo_dir = pathlib.Path(__file__).parent
 
 
@@ -29,20 +32,6 @@ class DemoList:
         ['replicas', 'markdown'],
     ]
 
-    TO_TIME_STR = {
-        -1: 'null',
-        300: '5 minutes',
-        600: '10 minutes',
-        900: '15 minutes',
-        1800: '30 minutes',
-        3600: '1 hour',
-        36000: '10 hours',
-        86400: '24 hours',
-        172800: '48 hours',
-        259200: '72 hours',
-        604800: '1 week',
-    }
-
     def __init__(self):
         self.api = HfApi()
         self._raw_data = self.load_data()
@@ -69,31 +58,39 @@ class DemoList:
             space_id = self.get_space_id(url)
             space_info = self.api.space_info(repo_id=space_id)
             card = space_info.cardData
-            info = data[url]
+
+            info: dict = data[url] | {
+                'url': url,
+                'title': card['title'] if 'title' in card else space_id,
+                'owner': space_id.split('/')[0],
+                'sdk': card['sdk'],
+                'sdk_version': card.get('sdk_version', ''),
+                'likes': space_info.likes,
+                'private': space_info.private,
+                'last_modified': space_info.lastModified,
+                'status': space_info.runtime['stage'],
+                'suggested_hardware': card.get('suggested_hardware', ''),
+            }
 
             for tag in ['arxiv', 'github', 'tags']:
                 if tag not in info:
                     info[tag] = []
 
-            info['url'] = url
-            info['owner'] = space_id.split('/')[0]
-            info['title'] = card['title']
-            info['sdk'] = card['sdk']
-            info['sdk_version'] = card.get('sdk_version', '')
-            info['likes'] = space_info.likes
-            info['last_modified'] = space_info.lastModified
-            info['status'] = space_info.runtime['stage']
-
-            info['suggested_hardware'] = card.get('suggested_hardware', '')
+            # `current` of paused Spaces is `None`, but `requested` is not
             info['hardware'] = space_info.runtime['hardware']['current']
             if info['hardware'] is None:
                 info['hardware'] = space_info.runtime['hardware']['requested']
 
-            if info['hardware'] == 'cpu-basic':
-                info['sleep_time'] = 172800
-            else:
-                info['sleep_time'] = space_info.runtime['gcTimeout'] or -1
+            # `gcTimeout` is `None` for `cpu-basic` Spaces and Spaces
+            # with "Don't sleep" sleep time.
+            # We use `-1` to represent it.
+            info['sleep_time'] = space_info.runtime['gcTimeout'] or -1
+            if info['sleep_time'] not in SLEEP_TIME_INT_TO_STR:
+                print(space_id)
+                print(f'Unknown sleep time: {info["sleep_time"]}')
+                continue
 
+            # `resources` of paused Spaces is `None`
             resources = space_info.runtime['resources']
             info['replicas'] = -1 if resources is None else resources[
                 'replicas']
@@ -175,10 +172,46 @@ class DemoList:
                 row.sdk_version,
                 'sleep_time':
                 self.add_div_tag_to_sleep_time(
-                    self.TO_TIME_STR[row.sleep_time], row.hardware),
+                    SLEEP_TIME_INT_TO_STR[row.sleep_time], row.hardware),
                 'replicas':
                 self.add_div_tag_to_replicas(row.replicas),
             }
             new_rows.append(new_row)
         df = pd.DataFrame(new_rows).loc[:, self.column_names]
         return df
+
+    def apply_filter(
+        self,
+        status: list[str],
+        hardware: list[str],
+        sleep_time: list[str],
+        multiple_replicas: bool,
+        sdk: list[str],
+        visibility: list[str],
+        owner: list[str],
+    ) -> pd.DataFrame:
+        df_raw = self.df_raw
+        df = self.df
+
+        if multiple_replicas:
+            df = df[df_raw.replicas > 1]
+
+        if visibility == ['public']:
+            df = df[~df_raw.private]
+        elif visibility == ['private']:
+            df = df[df_raw.private]
+
+        df = df[(df_raw.status.isin(status)) & (df_raw.hardware.isin(hardware))
+                & (df_raw.sdk.isin(sdk))]
+
+        sleep_time_int = [SLEEP_TIME_STR_TO_INT[s] for s in sleep_time]
+        df = df[df_raw.sleep_time.isin(sleep_time_int)]
+
+        if set(owner) == set(OWNER_CHOICES):
+            pass
+        elif WHOAMI in owner:
+            df = df[df_raw.owner == WHOAMI]
+        else:
+            df = df[df_raw.owner != WHOAMI]
+
+        return df

restart_scheduler.py

@@ -0,0 +1,25 @@
+from apscheduler.schedulers.background import BackgroundScheduler
+from huggingface_hub import HfApi
+from huggingface_hub.utils import RepositoryNotFoundError
+
+
+class RestartScheduler:
+    def __init__(self, space_id: str, interval_minutes: int,
+                 hf_token: str | None):
+        api = HfApi(token=hf_token)
+        if api.get_token_permission() != 'write':
+            raise ValueError('The HF token must have write permission.')
+        try:
+            api.space_info(repo_id=space_id)
+        except RepositoryNotFoundError:
+            raise ValueError('The Space ID does not exist.')
+        if interval_minutes <= 0:
+            raise ValueError('The interval must be positive.')
+
+        self.scheduler = BackgroundScheduler()
+        self.scheduler.add_job(func=lambda: api.restart_space(space_id),
+                               trigger='interval',
+                               seconds=60 * interval_minutes)
+
+    def start(self):
+        self.scheduler.start()

settings.py

@@ -0,0 +1,5 @@
+import os
+
+HF_TOKEN = os.getenv('HUGGING_FACE_HUB_TOKEN')
+SPACE_ID = os.getenv('SPACE_ID')
+INTERVAL_MINUTES = int(os.getenv('INTERVAL_MINUTES', '30'))