app.py

from typing import Callable
import gradio as gr
import asyncio
import time
import threading
from src.retrieve_data import (
    get_gpus_for_leaderboard,
    get_leaderboard_names,
    get_leaderboard_submissions,
)

from src.envs import CACHE_TIMEOUT, BACKGROUND_REFRESH_INTERVAL

# key: func_name:args:kwargs, value: (timestamp, data)
cache = {}


active_selections = {
    "leaderboard": None,
    "gpu": None,
}

loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)

background_refresh_running = True


def cached_fetch(func: Callable, *args, force_refresh=False, **kwargs):
    """Fetch data with caching to avoid redundant API calls"""
    cache_key = f"{func.__name__}:{str(args)}:{str(kwargs)}"

    current_time = time.time()

    if not force_refresh and cache_key in cache:
        timestamp, data = cache[cache_key]
        if current_time - timestamp < CACHE_TIMEOUT:
            return data

    result = loop.run_until_complete(func(*args, **kwargs))
    cache[cache_key] = (current_time, result)
    return result


def invalidate_cache(prefix=None):
    """Invalidate all cache entries or those matching a prefix"""
    global cache
    if prefix is None:
        cache = {}
    else:
        cache = {k: v for k, v in cache.items() if not k.startswith(prefix)}


def background_refresh():
    """Background thread to refresh active data periodically"""
    while background_refresh_running:
        try:
            time.sleep(BACKGROUND_REFRESH_INTERVAL)

            lb_name = active_selections["leaderboard"]
            gpu_name = active_selections["gpu"]

            if lb_name and gpu_name:
                cached_fetch(
                    get_leaderboard_submissions, lb_name, gpu_name, force_refresh=True
                )

                cached_fetch(get_gpus_for_leaderboard, lb_name, force_refresh=True)

            cached_fetch(get_leaderboard_names, force_refresh=True)

        except Exception as e:
            print(f"Background refresh error: {e}")


background_thread = threading.Thread(target=background_refresh, daemon=True)
background_thread.start()


def create_table_for_lb(lb_data):
    headers = [
        "Rank",
        "Discord User ID",
        "Submission Name",
        "Runtime (ms)",
        "Submission Date",
    ]

    rows = []
    for i, result in enumerate(lb_data.results, 1):
        rank_display = i
        if i == 1:
            rank_display = "🥇 1"
        elif i == 2:
            rank_display = "🥈 2"
        elif i == 3:
            rank_display = "🥉 3"

        rows.append(
            [
                rank_display,
                result.user_id,
                result.submission_name,
                f"{float(result.submission_score):.4f}",
                result.submission_time,
            ]
        )

    df = gr.Dataframe(
        headers=headers,
        datatype=[
            "str",
            "int",
            "str",
            "str",
            "datetime",
        ],
        value=rows,
        interactive=False,
    )

    return df


def on_lb_change(lb_name):
    gpu_choices = cached_fetch(get_gpus_for_leaderboard, lb_name)

    active_selections["leaderboard"] = lb_name
    if gpu_choices:
        active_selections["gpu"] = gpu_choices[0]

    return (
        gr.update(choices=gpu_choices, value=gpu_choices[0] if gpu_choices else None),
        update_table(lb_name, gpu_choices[0] if gpu_choices else None),
    )


def update_table(lb_name, gpu_name):
    if not gpu_name:
        return None

    active_selections["gpu"] = gpu_name

    data = cached_fetch(get_leaderboard_submissions, lb_name, gpu_name)
    return create_table_for_lb(data)


def build_ui():
    with gr.Blocks(
        title="ML Leaderboards",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container table tr:nth-child(1) {
            background-color: rgba(255, 215, 0, 0.2) !important;  /* Gold */
        }
        .gradio-container table tr:nth-child(2) {
            background-color: rgba(192, 192, 192, 0.2) !important;  /* Silver */
        }
        .gradio-container table tr:nth-child(3) {
            background-color: rgba(205, 127, 50, 0.2) !important;  /* Bronze */
        }
    """,
    ) as app:
        gr.Markdown("# 🍿 KernelBot Leaderboard 🍿")

        lb_names = cached_fetch(get_leaderboard_names)
        selected_lb = lb_names[0]
        gpu_names = cached_fetch(get_gpus_for_leaderboard, selected_lb)
        selected_gpu = gpu_names[0]

        data = cached_fetch(get_leaderboard_submissions, selected_lb, selected_gpu)

        with gr.Row():
            with gr.Column(scale=1):
                lb_dropdown = gr.Dropdown(
                    choices=lb_names,
                    label="Select Leaderboard",
                    interactive=True,
                    value=selected_lb,
                )
                gpu_dropdown = gr.Dropdown(
                    choices=gpu_names,
                    label="Select GPU",
                    interactive=True,
                    value=selected_gpu,
                )

        with gr.Row():
            results_table = create_table_for_lb(data)

        lb_dropdown.change(
            fn=on_lb_change,
            inputs=[lb_dropdown],
            outputs=[gpu_dropdown, results_table],
        )

        gpu_dropdown.change(
            fn=update_table, inputs=[lb_dropdown, gpu_dropdown], outputs=results_table
        )

    return app


if __name__ == "__main__":
    try:
        app = build_ui()
        app.launch()
    finally:
        background_refresh_running = False
        background_thread.join(timeout=1.0)
        loop.close()