Skip to content

Commit c8b7c85

Browse files
author
Ronald Tse
committed
fix(api): resolve model ids through the models.yaml index
_zip_path hardcoded "<id>-fp32.zip" — every int8/int4 variant (ara-diac-small-1.0-int8, tha-g2p-small-1.0-int8/-int4) 404'd as "unknown model", and heb-diac-1.1 (index filename heb.zip, volume copy heb-diac-1.1-fp32.zip) missed too. Resolution now goes through the index contract with precision-convention fallbacks, matched by equality against the volume listing (no user input in path construction). Health reports zips + model dirs.
1 parent 9267c20 commit c8b7c85

3 files changed

Lines changed: 107 additions & 5 deletions

File tree

‎src/api/inference.py‎

Lines changed: 16 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -38,11 +38,16 @@ def _zip_path(model_id: str) -> Path:
3838
# the server's own volume listing; the user value is used solely in
3939
# an equality comparison, never in path construction (CWE-22)
4040
import glob
41+
import os
42+
import sys
43+
44+
sys.path.insert(0, "/root/interscript-ml")
45+
from src.api.model_resolution import load_index, resolve_zip_filename
4146

42-
wanted = f"{model_id}-fp32.zip"
43-
for z in glob.glob("/v/imf/*/*-fp32.zip"):
44-
if z.rsplit("/", 1)[1] == wanted:
45-
return Path(z)
47+
volume_files = [os.path.basename(z) for z in glob.glob("/v/imf/*/*.zip")]
48+
filename = resolve_zip_filename(model_id, load_index("models.yaml"), volume_files)
49+
for z in glob.glob(f"/v/imf/*/{filename}"):
50+
return Path(z)
4651
raise KeyError(model_id)
4752

4853

@@ -124,9 +129,15 @@ async def _run_infer(body):
124129
@api.get("/health")
125130
def health() -> dict:
126131
import glob
132+
import os
127133

128134
models_volume.reload()
129-
return {"ok": True, "models": len(glob.glob("/v/imf/*/*-fp32.zip"))}
135+
zips = glob.glob("/v/imf/*/*.zip")
136+
return {
137+
"ok": True,
138+
"zips": len(zips),
139+
"model_dirs": len({os.path.dirname(z).rsplit("/", 1)[1] for z in zips}),
140+
}
130141

131142
return api
132143

‎src/api/model_resolution.py‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
"""Index-driven model-id resolution: the models.yaml contract, shared
2+
by every consumer. A model id resolves to the exact index filename
3+
when the volume carries it, falling back to the precision convention
4+
(``{id}-{precision}.zip``) and then ``{id}-fp32.zip`` — volume copies
5+
have historically landed under convention names even when the release
6+
artifact uses another name (heb-diac-1.1 ships as heb.zip but the
7+
volume copy is heb-diac-1.1-fp32.zip).
8+
9+
Callers pass the volume listing; candidate names are built server-side
10+
and matched by equality only — user input never touches path
11+
construction (CWE-22).
12+
"""
13+
14+
from __future__ import annotations
15+
16+
import yaml
17+
from pathlib import Path
18+
19+
20+
def load_index(path: str | Path = "models.yaml") -> dict:
21+
with open(path, encoding="utf-8") as fh:
22+
return yaml.safe_load(fh)["models"]
23+
24+
25+
def resolve_zip_filename(model_id: str, index: dict, volume_files: list[str]) -> str:
26+
"""Return the volume filename serving `model_id`.
27+
28+
Raises KeyError when the id is unknown or no volume file matches.
29+
"""
30+
entry = index.get(model_id)
31+
if entry is None:
32+
raise KeyError(model_id)
33+
34+
candidates: list[str] = []
35+
filename = entry.get("filename")
36+
if filename:
37+
candidates.append(filename)
38+
precision = entry.get("precision")
39+
if precision:
40+
candidates.append(f"{model_id}-{precision}.zip")
41+
candidates.append(f"{model_id}-fp32.zip")
42+
43+
available = set(volume_files)
44+
for candidate in candidates:
45+
if candidate in available:
46+
return candidate
47+
raise KeyError(model_id)

‎tests/test_model_resolution.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
"""Index-driven model-id → volume-filename resolution (the harness
2+
contract: model ids resolve through models.yaml, never hardcoded
3+
naming conventions)."""
4+
5+
from pathlib import Path
6+
7+
import yaml
8+
9+
from src.api.model_resolution import resolve_zip_filename
10+
11+
INDEX = yaml.safe_load(Path("models.yaml").read_text())["models"]
12+
13+
14+
def test_exact_index_filename_wins():
15+
volume = ["ara-diac-small-1.0-fp32.zip", "ara-diac-small-1.0-int8.zip"]
16+
assert resolve_zip_filename("ara-diac-small-1.0-int8", INDEX, volume) == (
17+
"ara-diac-small-1.0-int8.zip"
18+
)
19+
20+
21+
def test_precision_fallback_when_volume_name_differs_from_index():
22+
# heb-diac-1.1 ships as heb.zip on GH Releases but the volume copy
23+
# landed as heb-diac-1.1-fp32.zip
24+
volume = ["heb-diac-1.1-fp32.zip", "heb-diac-1.1-fp16.zip"]
25+
assert resolve_zip_filename("heb-diac-1.1", INDEX, volume) == "heb-diac-1.1-fp32.zip"
26+
27+
28+
def test_int4_variant_resolves():
29+
volume = ["tha-g2p-small-1.0-int8.zip", "tha-g2p-small-1.0-int4.zip"]
30+
assert resolve_zip_filename("tha-g2p-small-1.0-int4", INDEX, volume) == (
31+
"tha-g2p-small-1.0-int4.zip"
32+
)
33+
34+
35+
def test_fp32_convention_fallback():
36+
volume = ["ara-diac-1.0-fp32.zip"]
37+
assert resolve_zip_filename("ara-diac-1.0", INDEX, volume) == "ara-diac-1.0-fp32.zip"
38+
39+
40+
def test_unknown_id_raises_keyerror():
41+
import pytest
42+
43+
with pytest.raises(KeyError):
44+
resolve_zip_filename("nope-1.0", INDEX, ["a-fp32.zip"])

0 commit comments

Comments
 (0)