Model Hub support

Author: piotrekno1

examples/gguf_example.py

@@ -15,8 +15,7 @@
 client = wildedge.WildEdge(
     app_version="1.0.0",  # set WILDEDGE_DSN env var
 )
-client.instrument("huggingface")
-client.instrument("gguf")
+client.instrument("gguf", hubs=["huggingface"])
 
 model_path = hf_hub_download(
     "bartowski/Llama-3.2-1B-Instruct-GGUF",

examples/onnx_example.py

@@ -16,8 +16,7 @@
 client = wildedge.WildEdge(
     app_version="1.0.0",  # set WILDEDGE_DSN env var
 )
-client.instrument("huggingface")
-client.instrument("onnx")
+client.instrument("onnx", hubs=["huggingface"])
 
 model_path = hf_hub_download("Xenova/resnet-50", "onnx/model.onnx")
 session = ort.InferenceSession(model_path)

examples/timm_example.py

@@ -22,7 +22,7 @@
 client = wildedge.WildEdge(
     app_version="1.0.0",  # set WILDEDGE_DSN env var
 )
-client.instrument("timm")
+client.instrument("timm", hubs=["huggingface", "torchhub"])
 
 model = timm.create_model("resnet18", pretrained=True)
 model.eval()

tests/test_client_flows.py

@@ -18,11 +18,10 @@ def test_register_model_fallback_requires_id_when_no_extractor(
             client.register_model(object())
 
 
-def test_on_model_auto_loaded_uses_hf_records_when_downloads_missing(
+def test_on_model_auto_loaded_uses_hub_records_when_downloads_missing(
     client_with_stubbed_runtime, dummy_handle
 ):
     client = client_with_stubbed_runtime
-    client._hf_instrumented = True
 
     records = [
         {
@@ -37,7 +36,7 @@ def test_on_model_auto_loaded_uses_hf_records_when_downloads_missing(
     ]
 
     with (
-        patch("wildedge.client.drain_downloads", return_value=records),
+        patch.object(client, "_drain_hub_trackers", return_value=records),
         patch.object(client, "register_model", return_value=dummy_handle),
     ):
         client._on_model_auto_loaded(DummyModel(), load_ms=5)
@@ -101,7 +100,7 @@ def test_load_skips_duplicate_track_load_for_auto_loaded_model(
 ):
     client = client_with_stubbed_runtime
     dummy_handle.model_id = "dup-model"
-    client._auto_loaded.add("dup-model")
+    client.auto_loaded.add("dup-model")
 
     with patch.object(client, "register_model", return_value=dummy_handle):
         client.load(DummyModel)

tests/test_hubs.py

@@ -0,0 +1,187 @@
+"""Tests for wildedge.hubs hub trackers."""
+
+from __future__ import annotations
+
+import os
+import types
+from unittest.mock import patch
+
+from wildedge.hubs.huggingface import HuggingFaceHubTracker
+from wildedge.hubs.torchhub import TorchHubTracker
+
+# ---------------------------------------------------------------------------
+# BaseHubTracker.scan_cache
+# ---------------------------------------------------------------------------
+
+
+def test_scan_cache_returns_real_files_skips_symlinks(tmp_path):
+    real_file = tmp_path / "blob"
+    real_file.write_bytes(b"x" * 100)
+    link = tmp_path / "link"
+    link.symlink_to(real_file)
+
+    tracker = HuggingFaceHubTracker()
+    with patch.object(tracker, "cache_dir", return_value=str(tmp_path)):
+        result = tracker.scan_cache()
+
+    assert str(real_file) in result
+    assert str(link) not in result
+    assert result[str(real_file)] == 100
+
+
+def test_scan_cache_returns_empty_when_no_cache_dir():
+    tracker = HuggingFaceHubTracker()
+    with patch.object(tracker, "cache_dir", return_value=None):
+        assert tracker.scan_cache() == {}
+
+
+def test_scan_cache_returns_empty_when_dir_missing():
+    tracker = HuggingFaceHubTracker()
+    with patch.object(tracker, "cache_dir", return_value="/nonexistent/path/xyz"):
+        assert tracker.scan_cache() == {}
+
+
+# ---------------------------------------------------------------------------
+# HuggingFaceHubTracker.diff_to_records
+# ---------------------------------------------------------------------------
+
+
+def test_hf_diff_to_records_groups_by_repo():
+    tracker = HuggingFaceHubTracker()
+    sep = os.sep
+    before = {}
+    after = {
+        f"{sep}cache{sep}hub{sep}models--facebook--opt-125m{sep}blobs{sep}sha1": 200_000_000,
+        f"{sep}cache{sep}hub{sep}models--facebook--opt-125m{sep}snapshots{sep}abc{sep}config.json": 1_000,
+        f"{sep}cache{sep}hub{sep}models--bert-base-uncased{sep}blobs{sep}sha2": 400_000_000,
+    }
+    records = tracker.diff_to_records(before, after, elapsed_ms=5000)
+
+    assert len(records) == 2
+    repo_ids = {r["repo_id"] for r in records}
+    assert repo_ids == {"facebook/opt-125m", "bert-base-uncased"}
+    for r in records:
+        assert r["source_type"] == "huggingface"
+        assert r["source_url"] == f"hf://{r['repo_id']}"
+        assert r["cache_hit"] is False
+        assert r["duration_ms"] == 5000
+
+    opt = next(r for r in records if r["repo_id"] == "facebook/opt-125m")
+    assert opt["size"] == 200_000_000 + 1_000
+
+
+def test_hf_diff_to_records_returns_empty_when_no_new_files():
+    tracker = HuggingFaceHubTracker()
+    snapshot = {"/cache/blobs/sha1": 100}
+    assert tracker.diff_to_records(snapshot, snapshot, elapsed_ms=1000) == []
+
+
+def test_hf_diff_to_records_ignores_files_outside_models_dirs():
+    tracker = HuggingFaceHubTracker()
+    before = {}
+    after = {"/cache/hub/some_other_file.txt": 500}
+    # Files not under a models-- directory are silently dropped (no repo_id)
+    records = tracker.diff_to_records(before, after, elapsed_ms=1000)
+    assert records == []
+
+
+# ---------------------------------------------------------------------------
+# TorchHubTracker.diff_to_records
+# ---------------------------------------------------------------------------
+
+
+def test_torch_hub_diff_to_records_checkpoints():
+    tracker = TorchHubTracker()
+    hub_dir = "/home/user/.cache/torch/hub"
+    before = {}
+    after = {f"{hub_dir}/checkpoints/resnet50-0676ba61.pth": 97_781_926}
+
+    with patch.object(tracker, "cache_dir", return_value=hub_dir):
+        records = tracker.diff_to_records(before, after, elapsed_ms=3000)
+
+    assert len(records) == 1
+    r = records[0]
+    assert r["source_type"] == "torchhub"
+    assert r["source_url"] == "torchhub://checkpoints/resnet50-0676ba61.pth"
+    assert r["repo_id"] == "resnet50.pth"  # hash suffix stripped
+    assert r["cache_hit"] is False
+    assert r["size"] == 97_781_926
+
+
+def test_torch_hub_diff_to_records_repo_clone_dir():
+    tracker = TorchHubTracker()
+    hub_dir = "/home/user/.cache/torch/hub"
+    before = {}
+    after = {
+        f"{hub_dir}/pytorch_vision_v0.10.0/hubconf.py": 2_000,
+        f"{hub_dir}/pytorch_vision_v0.10.0/torchvision/models/resnet.py": 30_000,
+    }
+
+    with patch.object(tracker, "cache_dir", return_value=hub_dir):
+        records = tracker.diff_to_records(before, after, elapsed_ms=2000)
+
+    assert len(records) == 2
+    for r in records:
+        assert r["source_type"] == "torchhub"
+        assert r["source_url"] == "torchhub://pytorch/vision"
+        assert r["repo_id"] == "pytorch/vision"
+
+
+def test_torch_hub_diff_to_records_empty_when_no_new_files():
+    tracker = TorchHubTracker()
+    snapshot = {"/cache/torch/hub/checkpoints/model.pth": 1000}
+    with patch.object(tracker, "cache_dir", return_value="/cache/torch/hub"):
+        assert tracker.diff_to_records(snapshot, snapshot, elapsed_ms=500) == []
+
+
+# ---------------------------------------------------------------------------
+# TorchHubTracker.install_patch idempotency
+# ---------------------------------------------------------------------------
+
+
+def test_torch_hub_install_patch_is_idempotent(monkeypatch):
+    import wildedge.hubs.torchhub as torchhub_mod
+
+    original_load_calls = []
+
+    class FakeHub:
+        @staticmethod
+        def load(repo_or_dir, model, *args, **kwargs):
+            original_load_calls.append((repo_or_dir, model))
+            return object()
+
+        @staticmethod
+        def get_dir():
+            return "/tmp/hub"
+
+    fake_torch = types.SimpleNamespace(hub=FakeHub)
+    monkeypatch.setattr(torchhub_mod, "_torch", fake_torch)
+    monkeypatch.setattr(torchhub_mod, "_torch_hub_load_patched", False)
+
+    tracker = TorchHubTracker()
+    tracker.install_patch(lambda: None)
+    first_patched = fake_torch.hub.load
+    tracker.install_patch(lambda: None)
+    second_patched = fake_torch.hub.load
+
+    assert first_patched is second_patched
+    assert getattr(first_patched, "__wildedge_patch_name__", None) == "torchhub_load"
+
+
+# ---------------------------------------------------------------------------
+# HuggingFaceHubTracker.drain (thread-local buffer)
+# ---------------------------------------------------------------------------
+
+
+def test_hf_tracker_drain_returns_and_clears_buffer(monkeypatch):
+    import wildedge.hubs.huggingface as hf_mod
+
+    tracker = HuggingFaceHubTracker()
+    # Directly inject a record into the thread-local buffer
+    hf_mod._buffer().append({"repo_id": "test/model", "source_type": "huggingface"})
+
+    result = tracker.drain()
+    assert len(result) == 1
+    assert result[0]["repo_id"] == "test/model"
+    # Buffer should be cleared
+    assert tracker.drain() == []

tests/test_integration_patching.py

@@ -1,15 +1,20 @@
 from __future__ import annotations
 
 import types
+from unittest.mock import patch
 
+import pytest
+
+from wildedge.client import WildEdge
+from wildedge.device import DeviceInfo
+from wildedge.hubs.huggingface import HuggingFaceHubTracker
 from wildedge.integrations.gguf import GgufExtractor
-from wildedge.integrations.hf import install_patch as install_hf_patch
 from wildedge.integrations.onnx import OnnxExtractor
 from wildedge.integrations.pytorch import PytorchExtractor
 
 
 def test_hf_install_patch_is_idempotent(monkeypatch):
-    import wildedge.integrations.hf as hf_mod
+    import wildedge.hubs.huggingface as hf_mod
 
     def orig_hf_hub_download(repo_id, filename, **kwargs):
         return f"/tmp/{repo_id}/{filename}"
@@ -28,15 +33,16 @@ def orig_hf_hub_download(repo_id, filename, **kwargs):
         hf_mod.sys.modules, "test_hf_consumer_mod", fake_consumer_module
     )
 
-    install_hf_patch()
+    tracker = HuggingFaceHubTracker()
+    tracker.install_patch(None)
     first = fake_consumer_module.hf_hub_download
-    install_hf_patch()
+    tracker.install_patch(None)
     second = fake_consumer_module.hf_hub_download
     assert first is second
 
 
 def test_hf_install_patch_retries_unpatched_part(monkeypatch):
-    import wildedge.integrations.hf as hf_mod
+    import wildedge.hubs.huggingface as hf_mod
 
     monkeypatch.setattr(hf_mod, "_hf", object())
     monkeypatch.setattr(hf_mod, "_fd", object())
@@ -50,9 +56,10 @@ def fake_install_hf():
         return calls["hf"] > 1
 
     monkeypatch.setattr(hf_mod, "_install_hf_hub_download_patch", fake_install_hf)
-    install_hf_patch()
+    tracker = HuggingFaceHubTracker()
+    tracker.install_patch(None)
     assert hf_mod._hf_hub_download_patched is False
-    install_hf_patch()
+    tracker.install_patch(None)
     assert hf_mod._hf_hub_download_patched is True
     assert calls["hf"] == 2
 
@@ -117,3 +124,51 @@ def client_ref():
     GgufExtractor.install_auto_load_patch(client_ref)
     second = fake_llama_cpp.Llama.__init__
     assert first is second
+
+
+# ---------------------------------------------------------------------------
+# instrument() hubs= parameter
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def stub_client():
+    with (
+        patch("wildedge.client.detect_device", return_value=DeviceInfo("id", "linux")),
+        patch("wildedge.client.Transmitter"),
+        patch("wildedge.client.Consumer"),
+    ):
+        yield WildEdge(dsn="https://secret@ingest.wildedge.dev/key")
+
+
+def test_instrument_hubs_activates_requested_trackers(stub_client):
+    activated = []
+    with (
+        patch.object(stub_client, "_activate_hub", side_effect=activated.append),
+        patch.dict(stub_client.PATCH_INSTALLERS, {"gguf": lambda ref: None}),
+    ):
+        stub_client.instrument("gguf", hubs=["huggingface"])
+
+    assert activated == ["huggingface"]
+
+
+def test_instrument_hubs_unknown_hub_raises(stub_client):
+    with pytest.raises(ValueError, match="Unknown hub"):
+        stub_client.instrument("gguf", hubs=["nonexistent"])
+
+
+def test_instrument_hub_name_directly_raises(stub_client):
+    with pytest.raises(ValueError, match="is a hub"):
+        stub_client.instrument("huggingface")
+
+
+def test_instrument_none_without_hubs_raises(stub_client):
+    with pytest.raises(ValueError, match="requires hubs="):
+        stub_client.instrument(None)
+
+
+def test_instrument_none_activates_hub(stub_client):
+    activated = []
+    with patch.object(stub_client, "_activate_hub", side_effect=activated.append):
+        stub_client.instrument(None, hubs=["huggingface"])
+    assert activated == ["huggingface"]

wildedge/cli.py

@@ -17,6 +17,7 @@
 from wildedge import constants
 from wildedge.client import parse_dsn
 from wildedge.device import get_device_id_path
+from wildedge.hubs.registry import HUBS_BY_NAME, supported_hubs
 from wildedge.integrations.registry import INTEGRATIONS_BY_NAME, supported_integrations
 from wildedge.paths import default_dead_letter_dir, default_pending_queue_dir
 from wildedge.runtime.bootstrap import (
@@ -260,7 +261,7 @@ def run_command(parsed: argparse.Namespace) -> int:
 
 def integration_list(value: str | None) -> list[str]:
     if not value or value == "all":
-        return sorted(supported_integrations())
+        return sorted(supported_integrations() | supported_hubs())
     return [item.strip() for item in value.split(",") if item.strip()]
 
 
@@ -449,7 +450,7 @@ def doctor_report(parsed: argparse.Namespace) -> dict:
         )
 
     for integration in integration_list(parsed.integrations):
-        spec = INTEGRATIONS_BY_NAME.get(integration)
+        spec = INTEGRATIONS_BY_NAME.get(integration) or HUBS_BY_NAME.get(integration)
         if spec is None:
             ok = False
             integrations.append(