Model Hub support

examples/cli/demo.sh

@@ -14,8 +14,8 @@ cd "${PROJECT_DIR}"
 
 uv sync
 
-uv run wildedge doctor --integrations timm
+uv run wildedge doctor --integrations timm --hubs huggingface,torchhub
 
 uv run wildedge run \
-  --print-startup-report --integrations timm -- \
+  --print-startup-report --integrations timm --hubs huggingface,torchhub -- \
   "${EXAMPLE_SCRIPT}"

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_cli.py

@@ -228,10 +228,10 @@ def test_doctor_passes_for_available_module(monkeypatch, capsys):
     monkeypatch.setenv("WILDEDGE_DSN", "https://secret@ingest.wildedge.dev/key")
     monkeypatch.setattr(cli.importlib.util, "find_spec", lambda _: object())
     monkeypatch.setattr(cli, "check_writable_dir", lambda _: (True, "ok"))
-    rc = cli.main(["doctor", "--integrations", "huggingface"])
+    rc = cli.main(["doctor", "--hubs", "huggingface"])
     out = capsys.readouterr().out
     assert rc == 0
-    assert "integration[huggingface]: OK" in out
+    assert "hub[huggingface]: OK" in out
     assert "doctor: PASS" in out
 
 
@@ -245,6 +245,7 @@ def test_doctor_json_output_schema(monkeypatch, capsys):
     assert rc == 0
     assert sorted(payload.keys()) == [
         "checks",
+        "hubs",
         "integrations",
         "platform",
         "python",

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"]

tests/test_settings.py

@@ -45,7 +45,7 @@ def test_read_runtime_env_uses_run_over_base_dsn():
         RUN_FLUSH_TIMEOUT_ENV: "7.5",
         RUN_INTEGRATIONS_ENV: "onnx,timm",
     }
-    s = read_runtime_env(all_integrations=["onnx", "timm"], environ=env)
+    s = read_runtime_env(all_integrations=["onnx", "timm"], all_hubs=[], environ=env)
     assert s.dsn == "https://run@ingest.wildedge.dev/run"
     assert s.flush_timeout == 7.5
     assert s.integrations == ["onnx", "timm"]