src/retrieve_data.py

from collections import defaultdict

from httpx import AsyncClient

from src.envs import API_URL, TIMEOUT
from src.result import LbData, Result


async def get_leaderboard_names() -> list[str]:
    async with AsyncClient(timeout=TIMEOUT) as client:
        response = await client.get(f"{API_URL}/leaderboards")
        response.raise_for_status()
        return [lb["name"] for lb in response.json()]


async def get_gpus_for_leaderboard(lb_name: str) -> list[str]:
    async with AsyncClient(timeout=TIMEOUT) as client:
        response = await client.get(f"{API_URL}/gpus/{lb_name}")
        response.raise_for_status()
        return response.json()


async def get_leaderboard_submissions(lb_name: str, gpu: str) -> LbData:
    async with AsyncClient(timeout=TIMEOUT) as client:
        response = await client.get(f"{API_URL}/submissions/{lb_name}/{gpu}")
        response.raise_for_status()
        return LbData(
            gpu=gpu,
            name=lb_name,
            results=[Result.from_dict(result) for result in response.json()],
        )


async def populate_lb_data():
    leaderboards: dict[str, dict[str, LbData]] = defaultdict(dict)
    lb_names = await get_leaderboard_names()
    for lb_name in lb_names:
        gpus = await get_gpus_for_leaderboard(lb_name)
        for gpu in gpus:
            lb_data = await get_leaderboard_submissions(lb_name, gpu)
            leaderboards[lb_name][gpu] = lb_data

    return leaderboards