Skip to content
Open
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
47 changes: 30 additions & 17 deletions mellea/backends/adapters/adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import contextlib
import pathlib
import re
import tempfile
import warnings
from typing import Literal, TypeAlias, TypeVar, cast

Expand Down Expand Up @@ -751,8 +752,18 @@ def from_hub(
) -> list["EmbeddedIntrinsicAdapter"]:
"""Load embedded adapters from a Granite Switch model on Hugging Face Hub.

Downloads `adapter_index.json` and the `io_configs/` directory, then
delegates to :meth:`from_model_directory`.
Downloads `adapter_index.json` and the `io_configs/` directory into a
self-contained local directory, then delegates to
:meth:`from_model_directory`.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

NIT: :meth:from_model_directoryis an RST cross-reference directive. AGENTS.md §5 bans RST markup in docstrings — should be from_model_directory ``. Fix all three occurrences in this docstring (lines 757, 782, 784) or leave for a sweep.


`huggingface_hub.snapshot_download`'s default cache-backed snapshot
directory populates `io_configs/` with symlinks that resolve into a
sibling `blobs/` directory *outside* the snapshot root. That breaks the
contract `from_model_directory` expects (a self-contained model
directory) and trips its path-escape check. To satisfy that contract,
the files are downloaded directly into a temporary directory (via
`local_dir`) instead, so `io_configs/` contains real files rather than
symlinks escaping the directory.

Args:
repo_id (str): Hugging Face Hub repository ID
Expand Down Expand Up @@ -780,22 +791,24 @@ def from_hub(
'Hugging Face Hub. Please install it with: pip install "mellea[switch]"'
) from e

local_root = huggingface_hub.snapshot_download(
repo_id=repo_id,
allow_patterns=["adapter_index.json", "io_configs/**"],
cache_dir=cache_dir,
revision=revision,
)
try:
return EmbeddedIntrinsicAdapter.from_model_directory(
local_root, intrinsic_name=intrinsic_name
with tempfile.TemporaryDirectory() as local_dir:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cache always cold. local_dir prevents cache_dir from ever being written to; nothing in mellea/ pre-populates io_configs/** in cache mode, so try_to_load_from_cache never hits.

Each call is ~51 HTTP round-trips (repo_info + 25 HEAD + 25 GET) with zero reuse across calls. _resolve_adapter (adapter.py:513) calls this once per distinct intrinsic name — 12 times for a full Switch model.

Measured on ibm-granite/granite-switch-4.1-3b-preview: three consecutive calls took 0.88s / 0.44s / 0.50s.

Hard offline failure: HF_HUB_OFFLINE=1 now raises LocalEntryNotFoundError because a fresh empty temp dir has nothing to read from, and the cached-snapshot fallback at _snapshot_download.py:293 is gated on local_dir is None.

Suggested fix: Use a persistent directory under cache_dir/HF_HUB_CACHE + repo + revision instead. Verified: no symlinks (escape check passes), warm online 0.9s → 0.1s, offline works. Path must include revision to avoid cross-revision collisions.

downloaded_dir = huggingface_hub.snapshot_download(
repo_id=repo_id,
allow_patterns=["adapter_index.json", "io_configs/**"],
cache_dir=cache_dir,
local_dir=local_dir,
revision=revision,
)
except ValueError as e:
if intrinsic_name is not None:
raise ValueError(
f"No adapter found for adapter function '{intrinsic_name}' in {repo_id}"
) from e
raise ValueError(f"No adapters found in {repo_id}") from e
try:
return EmbeddedIntrinsicAdapter.from_model_directory(
downloaded_dir, intrinsic_name=intrinsic_name
)
except ValueError as e:
if intrinsic_name is not None:
raise ValueError(
f"No adapter found for adapter function '{intrinsic_name}' in {repo_id}"
) from e
raise ValueError(f"No adapters found in {repo_id}") from e

@staticmethod
def from_source(
Expand Down
14 changes: 13 additions & 1 deletion test/backends/test_adapters/test_embedded_adapter.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
import json
import os
import pathlib
from unittest.mock import MagicMock, patch
from unittest.mock import ANY, MagicMock, patch

import pytest
import yaml
Expand Down Expand Up @@ -308,6 +308,7 @@ def test_downloads_and_delegates(self, model_dir):
repo_id="ibm-granite/granite-switch-micro",
allow_patterns=["adapter_index.json", "io_configs/**"],
cache_dir="/tmp/test-cache",
local_dir=ANY,
revision="test-rev",
)
assert len(adapters) == 2
Expand All @@ -324,11 +325,21 @@ def test_filter_single_intrinsic(self, model_dir):
repo_id="ibm-granite/granite-switch-micro",
allow_patterns=["adapter_index.json", "io_configs/**"],
cache_dir=None,
local_dir=ANY,
revision="main",
)
assert len(adapters) == 1
assert adapters[0].intrinsic_name == "citations"

def test_from_hub_requests_local_dir(self, model_dir):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Test gap: mock returns model_dir (pre-built fixture with real files) — never proves that snapshot_download(..., local_dir=...) actually produces a self-contained directory that passes the escape check.

Suggested: a side_effect-based unit test that materialises a cache-style symlink layout (relative symlinks into a sibling blobs/ dir) when local_dir is None, and real files when it is set. Fails on main, passes on this branch.

"""from_hub requests local_dir to avoid symlinks escaping model directory."""
with patch(
"huggingface_hub.snapshot_download", return_value=str(model_dir)
) as mock_dl:
EmbeddedIntrinsicAdapter.from_hub("ibm-granite/granite-switch-micro")
_, kwargs = mock_dl.call_args
assert "local_dir" in kwargs and kwargs["local_dir"] is not None

def test_missing_huggingface_hub_raises(self):
with patch.dict("sys.modules", {"huggingface_hub": None}):
with pytest.raises(ImportError, match="huggingface_hub is required"):
Expand Down Expand Up @@ -377,6 +388,7 @@ def test_hub_passes_revision_and_cache(self, model_dir):
repo_id="ibm-granite/granite-switch-micro",
allow_patterns=["adapter_index.json", "io_configs/**"],
cache_dir="/tmp/cache",
local_dir=ANY,
revision="v2",
)

Expand Down
Loading