Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -110,3 +110,6 @@ omlx/_build_info.py

# Tailwind CSS standalone CLI binary (download on demand via build_css.py)
omlx/admin/tailwindcss-*

# Git worktrees
.worktrees/
62 changes: 29 additions & 33 deletions omlx/admin/routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -2952,6 +2952,24 @@ async def list_hf_models(is_admin: bool = Depends(require_admin)):

model_dirs = global_settings.model.get_model_dirs(global_settings.base_path)

from ..model_discovery import _resolve_hf_cache_entry

def _add_model(model_path: Path, model_name: str) -> None:
if model_name in seen_names:
return
seen_names.add(model_name)
total_size = sum(
f.stat().st_size for f in model_path.rglob("*") if f.is_file()
)
models.append(
{
"name": model_name,
"path": str(model_path),
"size": total_size,
"size_formatted": format_size(total_size),
}
)

models = []
seen_names: set[str] = set()
for model_dir in model_dirs:
Expand All @@ -2963,44 +2981,22 @@ async def list_hf_models(is_admin: bool = Depends(require_admin)):

if (subdir / "config.json").exists():
# Level 1: direct model folder
if subdir.name in seen_names:
continue
seen_names.add(subdir.name)
total_size = sum(
f.stat().st_size for f in subdir.rglob("*") if f.is_file()
)
models.append(
{
"name": subdir.name,
"path": str(subdir),
"size": total_size,
"size_formatted": format_size(total_size),
}
)
_add_model(subdir, subdir.name)
else:
# HF Hub cache entry: models--Org--Name/snapshots/<hash>/
hf_resolved = _resolve_hf_cache_entry(subdir)
if hf_resolved is not None:
snapshot_path, model_name = hf_resolved
if (snapshot_path / "config.json").exists():
_add_model(snapshot_path, model_name)
continue

# Level 2: organization folder — scan children
for child in sorted(subdir.iterdir()):
if not child.is_dir() or child.name.startswith("."):
continue
if not (child / "config.json").exists():
continue
if child.name in seen_names:
continue
seen_names.add(child.name)

total_size = sum(
f.stat().st_size
for f in child.rglob("*")
if f.is_file()
)
models.append(
{
"name": child.name,
"path": str(child),
"size": total_size,
"size_formatted": format_size(total_size),
}
)
if (child / "config.json").exists():
_add_model(child, child.name)

return {"models": models}

Expand Down
32 changes: 32 additions & 0 deletions omlx/model_discovery.py
Original file line number Diff line number Diff line change
Expand Up @@ -488,6 +488,30 @@ def _is_model_dir(path: Path) -> bool:
return (path / "config.json").exists() and not _is_adapter_dir(path)


def _resolve_hf_cache_entry(path: Path) -> tuple[Path, str] | None:
"""Resolve an HF Hub cache entry (models--Org--Name/) to its active snapshot.

Returns (snapshot_path, model_name) or None if not a valid HF cache entry.
"""
name = path.name
if not name.startswith("models--") or name.count("--") < 2:
return None

# "models--Org--Name" → "Name"
model_name = name.split("--", 2)[2]

try:
commit_hash = (path / "refs" / "main").read_text().strip()
except OSError:
return None

snapshot = path / "snapshots" / commit_hash
if not snapshot.is_dir():
return None

return snapshot, model_name


def _register_model(
models: dict[str, DiscoveredModel],
model_dir: Path,
Expand Down Expand Up @@ -594,6 +618,14 @@ def discover_models(model_dir: Path) -> dict[str, DiscoveredModel]:
# Level 1: direct model folder
_register_model(models, subdir, subdir.name)
else:
# HF Hub cache entry: models--Org--Name/snapshots/<hash>/
hf_resolved = _resolve_hf_cache_entry(subdir)
if hf_resolved is not None:
snapshot_path, model_name = hf_resolved
if _is_model_dir(snapshot_path):
_register_model(models, snapshot_path, model_name)
continue

# Level 2: organization folder — scan children
has_children = False
for child in sorted(subdir.iterdir()):
Expand Down
137 changes: 137 additions & 0 deletions tests/test_model_discovery.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
DiscoveredModel,
_is_adapter_dir,
_is_unsupported_model,
_resolve_hf_cache_entry,
detect_model_type,
discover_models,
discover_models_from_dirs,
Expand Down Expand Up @@ -839,3 +840,139 @@ def test_audio_models_included_in_discovery(self, tmp_path):
assert models["whisper-large-v3"].model_type == "audio_stt"
assert "Qwen3-TTS" in models
assert models["Qwen3-TTS"].model_type == "audio_tts"


class TestHfCacheDiscovery:
"""Tests for HF Hub cache entry resolution and discovery."""

FAKE_COMMIT = "abc123def456"

def _make_model(self, path: Path, model_type: str = "llama"):
"""Helper to create a valid model directory."""
path.mkdir(parents=True, exist_ok=True)
(path / "config.json").write_text(json.dumps({"model_type": model_type}))
(path / "model.safetensors").write_bytes(b"0" * 1000)

def _make_hf_cache_entry(self, parent: Path, org: str, name: str):
"""Helper to create a bare HF Hub cache directory layout (no model files)."""
entry = parent / f"models--{org}--{name}"
refs = entry / "refs"
refs.mkdir(parents=True)
(refs / "main").write_text(self.FAKE_COMMIT)
snapshot = entry / "snapshots" / self.FAKE_COMMIT
snapshot.mkdir(parents=True)
return entry, snapshot

def _make_hf_cache_model(self, parent: Path, org: str, name: str, model_type: str = "llama"):
"""Helper to create an HF cache entry with a valid model in the snapshot."""
_, snapshot = self._make_hf_cache_entry(parent, org, name)
(snapshot / "config.json").write_text(json.dumps({"model_type": model_type}))
(snapshot / "model.safetensors").write_bytes(b"0" * 1000)

def test_resolve_valid_entry(self, tmp_path):
"""Valid HF cache entry resolves to snapshot path and model name."""
entry, snapshot = self._make_hf_cache_entry(tmp_path, "mlx-community", "Qwen3-8B-4bit")

result = _resolve_hf_cache_entry(entry)
assert result is not None
assert result[0] == snapshot
assert result[1] == "Qwen3-8B-4bit"

def test_resolve_regular_dir_returns_none(self, tmp_path):
"""Regular directory without models-- prefix returns None."""
regular = tmp_path / "mlx-community"
regular.mkdir()
assert _resolve_hf_cache_entry(regular) is None

def test_resolve_single_separator_returns_none(self, tmp_path):
"""models--Name (no org separator) returns None."""
entry = tmp_path / "models--NoOrg"
entry.mkdir()
assert _resolve_hf_cache_entry(entry) is None

def test_resolve_missing_refs_main_returns_none(self, tmp_path):
"""Missing refs/main returns None."""
entry = tmp_path / "models--mlx-community--Qwen3-8B"
entry.mkdir(parents=True)
assert _resolve_hf_cache_entry(entry) is None

def test_resolve_missing_snapshot_returns_none(self, tmp_path):
"""Valid refs/main but missing snapshot directory returns None."""
entry = tmp_path / "models--mlx-community--Qwen3-8B"
refs = entry / "refs"
refs.mkdir(parents=True)
(refs / "main").write_text("deadbeef")
assert _resolve_hf_cache_entry(entry) is None

def test_resolve_strips_whitespace_from_refs(self, tmp_path):
"""Trailing newline in refs/main is stripped (matches real HF cache)."""
entry, snapshot = self._make_hf_cache_entry(tmp_path, "mlx-community", "Qwen3-8B")
# Overwrite with trailing newline (like real HF cache)
(entry / "refs" / "main").write_text(self.FAKE_COMMIT + "\n")

result = _resolve_hf_cache_entry(entry)
assert result is not None
assert result[0] == snapshot

def test_discover_hf_cache_model(self, tmp_path):
"""HF cache entries are discovered as models."""
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")

models = discover_models(tmp_path)
assert len(models) == 1
assert "Qwen3-8B-4bit" in models
assert models["Qwen3-8B-4bit"].model_type == "llm"

def test_discover_multiple_hf_cache_models(self, tmp_path):
"""Multiple HF cache entries are all discovered."""
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
self._make_hf_cache_model(tmp_path, "mlx-community", "Mistral-7B-v0.3")

models = discover_models(tmp_path)
assert len(models) == 2
assert "Qwen3-8B-4bit" in models
assert "Mistral-7B-v0.3" in models

def test_hf_cache_model_path_points_to_snapshot(self, tmp_path):
"""model_path points to the snapshot dir, not the cache entry."""
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")

models = discover_models(tmp_path)
assert models["Qwen3-8B-4bit"].model_path == str(
tmp_path / "models--mlx-community--Qwen3-8B-4bit" / "snapshots" / self.FAKE_COMMIT
)

def test_hf_cache_without_config_json_skipped(self, tmp_path):
"""HF cache entries without config.json in snapshot are skipped."""
self._make_hf_cache_entry(tmp_path, "mlx-community", "NoConfig")

models = discover_models(tmp_path)
assert len(models) == 0

def test_mixed_flat_and_hf_cache(self, tmp_path):
"""Mix of flat models and HF cache entries."""
self._make_model(tmp_path / "mistral-7b")
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")

models = discover_models(tmp_path)
assert len(models) == 2
assert "mistral-7b" in models
assert "Qwen3-8B-4bit" in models

def test_mixed_org_and_hf_cache(self, tmp_path):
"""Mix of org folders and HF cache entries."""
self._make_model(tmp_path / "Qwen" / "Qwen3-8B", model_type="qwen2")
self._make_hf_cache_model(tmp_path, "mlx-community", "Mistral-7B")

models = discover_models(tmp_path)
assert len(models) == 2
assert "Qwen3-8B" in models
assert "Mistral-7B" in models

def test_hf_cache_does_not_fall_through_to_org_scan(self, tmp_path):
"""HF cache entries don't get scanned as org folders."""
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")

models = discover_models(tmp_path)
assert len(models) == 1
assert "Qwen3-8B-4bit" in models