app.py

#!/usr/bin/env python

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 demo_list import DemoList

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)
    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
    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 multiple_replicas:
        df = df[df_raw.replicas > 1]
    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:
        return STATUS_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


def update_hardware_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return HARDWARE_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


def update_sdk_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return SDK_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


def update_sleep_time_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return SLEEP_TIME_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


with gr.Blocks(css='style.css') as demo:
    with gr.Accordion(label='Filter', open=False):
        status = gr.CheckboxGroup(label='Status',
                                  choices=['(ALL)', '(NONE)'] + STATUS_CHOICES,
                                  value=STATUS_CHOICES,
                                  type='value')
        hardware = gr.CheckboxGroup(label='Hardware',
                                    choices=['(ALL)', '(NONE)'] +
                                    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')
        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,
                      type='pandas')

    status.change(fn=update_status_checkboxes,
                  inputs=status,
                  outputs=status,
                  queue=False,
                  show_progress=False,
                  api_name=False)
    hardware.change(fn=update_hardware_checkboxes,
                    inputs=hardware,
                    outputs=hardware,
                    queue=False,
                    show_progress=False,
                    api_name=False)
    sdk.change(fn=update_sdk_checkboxes,
               inputs=sdk,
               outputs=sdk,
               queue=False,
               show_progress=False,
               api_name=False)
    sleep_time.change(fn=update_sleep_time_checkboxes,
                      inputs=sleep_time,
                      outputs=sleep_time,
                      queue=False,
                      show_progress=False,
                      api_name=False)
    apply_button.click(fn=update_df,
                       inputs=[
                           status,
                           hardware,
                           sdk,
                           sleep_time,
                           owner,
                           multiple_replicas,
                       ],
                       outputs=df,
                       api_name=False)
demo.queue(api_open=False).launch()