diff --git a/src/api/inference.py b/src/api/inference.py index f30e9af..7a6a98b 100644 --- a/src/api/inference.py +++ b/src/api/inference.py @@ -38,11 +38,16 @@ def _zip_path(model_id: str) -> Path: # the server's own volume listing; the user value is used solely in # an equality comparison, never in path construction (CWE-22) import glob + import os + import sys + + sys.path.insert(0, "/root/interscript-ml") + from src.api.model_resolution import load_index, resolve_zip_filename - wanted = f"{model_id}-fp32.zip" - for z in glob.glob("/v/imf/*/*-fp32.zip"): - if z.rsplit("/", 1)[1] == wanted: - return Path(z) + volume_files = [os.path.basename(z) for z in glob.glob("/v/imf/*/*.zip")] + filename = resolve_zip_filename(model_id, load_index("models.yaml"), volume_files) + for z in glob.glob(f"/v/imf/*/{filename}"): + return Path(z) raise KeyError(model_id) @@ -124,9 +129,15 @@ async def _run_infer(body): @api.get("/health") def health() -> dict: import glob + import os models_volume.reload() - return {"ok": True, "models": len(glob.glob("/v/imf/*/*-fp32.zip"))} + zips = glob.glob("/v/imf/*/*.zip") + return { + "ok": True, + "zips": len(zips), + "model_dirs": len({os.path.dirname(z).rsplit("/", 1)[1] for z in zips}), + } return api diff --git a/src/api/model_resolution.py b/src/api/model_resolution.py new file mode 100644 index 0000000..629faeb --- /dev/null +++ b/src/api/model_resolution.py @@ -0,0 +1,47 @@ +"""Index-driven model-id resolution: the models.yaml contract, shared +by every consumer. A model id resolves to the exact index filename +when the volume carries it, falling back to the precision convention +(``{id}-{precision}.zip``) and then ``{id}-fp32.zip`` — volume copies +have historically landed under convention names even when the release +artifact uses another name (heb-diac-1.1 ships as heb.zip but the +volume copy is heb-diac-1.1-fp32.zip). + +Callers pass the volume listing; candidate names are built server-side +and matched by equality only — user input never touches path +construction (CWE-22). +""" + +from __future__ import annotations + +import yaml +from pathlib import Path + + +def load_index(path: str | Path = "models.yaml") -> dict: + with open(path, encoding="utf-8") as fh: + return yaml.safe_load(fh)["models"] + + +def resolve_zip_filename(model_id: str, index: dict, volume_files: list[str]) -> str: + """Return the volume filename serving `model_id`. + + Raises KeyError when the id is unknown or no volume file matches. + """ + entry = index.get(model_id) + if entry is None: + raise KeyError(model_id) + + candidates: list[str] = [] + filename = entry.get("filename") + if filename: + candidates.append(filename) + precision = entry.get("precision") + if precision: + candidates.append(f"{model_id}-{precision}.zip") + candidates.append(f"{model_id}-fp32.zip") + + available = set(volume_files) + for candidate in candidates: + if candidate in available: + return candidate + raise KeyError(model_id) diff --git a/tests/test_model_resolution.py b/tests/test_model_resolution.py new file mode 100644 index 0000000..c72bc8a --- /dev/null +++ b/tests/test_model_resolution.py @@ -0,0 +1,44 @@ +"""Index-driven model-id → volume-filename resolution (the harness +contract: model ids resolve through models.yaml, never hardcoded +naming conventions).""" + +from pathlib import Path + +import yaml + +from src.api.model_resolution import resolve_zip_filename + +INDEX = yaml.safe_load(Path("models.yaml").read_text())["models"] + + +def test_exact_index_filename_wins(): + volume = ["ara-diac-small-1.0-fp32.zip", "ara-diac-small-1.0-int8.zip"] + assert resolve_zip_filename("ara-diac-small-1.0-int8", INDEX, volume) == ( + "ara-diac-small-1.0-int8.zip" + ) + + +def test_precision_fallback_when_volume_name_differs_from_index(): + # heb-diac-1.1 ships as heb.zip on GH Releases but the volume copy + # landed as heb-diac-1.1-fp32.zip + volume = ["heb-diac-1.1-fp32.zip", "heb-diac-1.1-fp16.zip"] + assert resolve_zip_filename("heb-diac-1.1", INDEX, volume) == "heb-diac-1.1-fp32.zip" + + +def test_int4_variant_resolves(): + volume = ["tha-g2p-small-1.0-int8.zip", "tha-g2p-small-1.0-int4.zip"] + assert resolve_zip_filename("tha-g2p-small-1.0-int4", INDEX, volume) == ( + "tha-g2p-small-1.0-int4.zip" + ) + + +def test_fp32_convention_fallback(): + volume = ["ara-diac-1.0-fp32.zip"] + assert resolve_zip_filename("ara-diac-1.0", INDEX, volume) == "ara-diac-1.0-fp32.zip" + + +def test_unknown_id_raises_keyerror(): + import pytest + + with pytest.raises(KeyError): + resolve_zip_filename("nope-1.0", INDEX, ["a-fp32.zip"])