SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
7.2 KB · 123 lines python
Raw Blame History
1"""/hardware listing, /hardware/fit (ESTIMATED memory fit) and /hardware/{slug} alias."""2from __future__ import annotations34from typing import Any56from fastapi import APIRouter, Query, Request78from aiatlas.api.common import (9    ENTITY_COLS,10    ENTITY_FROM,11    PAGINATION,12    ApiError,13    Pagination,14    attr_num,15    cached,16    entity_summary,17    page,18)19from aiatlas.api.routers.entities import detail_for_type20from aiatlas.db import connection, fetch_all, fetch_val21from aiatlas.services import hardware_fit as hf2223router = APIRouter(prefix="/api/v1/hardware", tags=["hardware"])2425MEMORY = ("(case when jsonb_typeof(e.attributes->'memory_gb') = 'array' then (select max(x::text::double precision) from jsonb_array_elements(e.attributes->'memory_gb') x where jsonb_typeof(x) = 'number') "26          "when jsonb_typeof(e.attributes->'memory_gb') = 'number' then (e.attributes->>'memory_gb')::double precision else " + attr_num("memory_gb") + " end)")27SORTS = {"memory": f"{MEMORY} desc nulls last", "name": "e.canonical_name asc", "release": "e.attributes->>'release_date' desc nulls last", "updated": "e.updated_at desc",28         "bandwidth": attr_num("memory_bandwidth_gbs") + " desc nulls last"}293031@router.get("")32@cached(300)33async def list_hardware(request: Request, p: Pagination = PAGINATION, kind: str | None = None, manufacturer: str | None = None, min_memory: float | None = Query(None, ge=0),34                        q: str | None = Query(None, max_length=200), sort: str = "memory") -> dict[str, Any]:35    if sort not in SORTS:36        raise ApiError(400, f"sort must be one of {', '.join(SORTS)}")37    where = ["e.entity_type = 'hardware'", "e.merged_into is null"]38    params: dict[str, Any] = {}39    if kind:40        where.append("e.attributes->>'kind' ilike :kind")41        params["kind"] = kind42    if manufacturer:43        where.append("(e.attributes->>'manufacturer' ilike :man or eo.slug = :man or eo.canonical_name ilike :man)")44        params["man"] = manufacturer45    if min_memory is not None:46        where.append(f"{MEMORY} >= :min_memory")47        params["min_memory"] = float(min_memory)48    if q:49        where.append("e.canonical_name ilike :qlike")50        params["qlike"] = f"%{q}%"51    where_sql = " and ".join(where)52    async with connection() as conn:53        rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where_sql} order by {SORTS[sort]}, e.id limit :lim offset :off", lim=p.limit, off=p.offset, **params)54        total = await fetch_val(conn, f"select count(*) from {ENTITY_FROM} where {where_sql}", **params)55        kinds = await fetch_all(conn, "select e.attributes->>'kind' as value, count(*) as count from entities e where e.entity_type = 'hardware' and e.merged_into is null and e.attributes ? 'kind' group by 1 order by 2 desc")56        mans = await fetch_all(conn, "select e.attributes->>'manufacturer' as value, count(*) as count from entities e where e.entity_type = 'hardware' and e.merged_into is null and e.attributes ? 'manufacturer' group by 1 order by 2 desc")57    out = page([entity_summary(r) for r in rows], int(total or 0), p)58    out["facets"] = {"kinds": [{"value": r["value"], "count": int(r["count"])} for r in kinds], "manufacturers": [{"value": r["value"], "count": int(r["count"])} for r in mans]}59    return out606162@router.get("/fit")63@cached(300)64async def hardware_fit(request: Request, memory_gb: float = Query(..., gt=0, le=100000), quant: str = Query("4bit"), context: int = Query(8192, ge=0, le=10_000_000),65                       limit: int = Query(100, ge=1, le=500), openness: str | None = None) -> dict[str, Any]:66    if quant not in hf.BYTES_PER_PARAM:67        raise ApiError(400, f"quant must be one of {', '.join(hf.BYTES_PER_PARAM)}")68    where = "e.entity_type = 'model' and e.merged_into is null and e.attributes ? 'parameter_count'"69    params: dict[str, Any] = {}70    if openness:71        vals = [v.strip() for v in openness.split(",") if v.strip()]72        if "open" in vals:73            vals += ["open-weights", "open-source"]74        where += " and e.attributes->>'openness' = any(cast(:openness as text[]))"75        params["openness"] = vals76    async with connection() as conn:77        rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where} order by {attr_num('parameter_count')} desc nulls last limit 2000", **params)78    items = []79    for r in rows:80        pc = hf.parameter_count(r["attributes"])81        if pc is None:82            continue83        items.append({"model": entity_summary(r), "parameter_count": pc, **hf.fit(pc, memory_gb, quant, context)})84    items.sort(key=lambda x: (not x["fits"], -x["parameter_count"] if x["fits"] else x["parameter_count"]))85    return {"inputs": {"memory_gb": memory_gb, "quant": quant, "context": context}, "estimated": True, "assumptions": hf.ASSUMPTIONS,86            "counts": {"fits": sum(1 for i in items if i["fits"]), "evaluated": len(items)}, "items": items[:limit]}878889@router.get("/{slug}/fit")90@cached(300)91async def hardware_entity_fit(request: Request, slug: str, quant: str = Query("4bit"), context: int = Query(8192, ge=0, le=10_000_000), memory_gb: float | None = Query(None, gt=0),92                              gpu_count: int = Query(1, ge=1, le=8), limit: int = Query(100, ge=1, le=500), openness: str | None = None) -> dict[str, Any]:93    """ESTIMATED fit of canonical models on ONE hardware entity (largest memory configuration unless `memory_gb` picks one)."""94    from aiatlas.api.common import openness_values, resolve_entity9596    q = hf.normalize_quant(quant)97    async with connection() as conn:98        hw = await resolve_entity(conn, slug, ("hardware",))99        options = hf.hardware_memory_options(hw.get("attributes"))100        if not options:101            return {"hardware": entity_summary(hw), "memory_options_gb": [], "items": [], "estimated": True, "note": "no memory_gb recorded for this hardware — nothing is estimated"}102        mem = memory_gb if memory_gb is not None else max(options)103        where = "e.entity_type = 'model' and e.merged_into is null and e.attributes ? 'parameter_count'"104        params: dict[str, Any] = {}105        if openness:106            where += " and e.attributes->>'openness' = any(cast(:openness as text[]))"107            params["openness"] = openness_values(openness)108        rows = await fetch_all(conn, f"select {ENTITY_COLS} from {ENTITY_FROM} where {where} order by {attr_num('parameter_count')} desc nulls last limit 3000", **params)109    items = []110    for r in rows:111        f = hf.fit_detailed(r["attributes"], mem, quant=q, context=context, gpu_count=gpu_count)112        if f:113            items.append({"model": entity_summary(r), **f})114    items.sort(key=lambda x: (not x["fits"], -(x.get("parameter_count") or 0) if x["fits"] else (x.get("parameter_count") or 0)))115    return {"hardware": entity_summary(hw), "memory_options_gb": options, "inputs": {"memory_gb": mem, "quant": q, "context": context, "gpu_count": gpu_count}, "estimated": True,116            "assumptions": hf.ASSUMPTIONS, "runtimes": (hw.get("attributes") or {}).get("runtimes"), "counts": {"fits": sum(1 for i in items if i["fits"]), "evaluated": len(items)}, "items": items[:limit]}117118119@router.get("/{slug}")120@cached(300)121async def get_hardware(request: Request, slug: str) -> dict[str, Any]:122    return await detail_for_type(slug, ("hardware",))123