Support for inference hardware context (#16)

Author: piotrekno1

README.md

@@ -99,6 +99,7 @@ def run(input):
 | `enable_offline_persistence` | `true` | - | Persist unsent events to disk and replay on restart |
 | `max_event_age_sec` | `900` | - | Max age before dead-lettering |
 | `enable_dead_letter_persistence` | `false` | - | Persist dropped batches to disk |
+| `sampling_interval_s` | `30.0` | `WILDEDGE_SAMPLING_INTERVAL_S` | Seconds between background hardware snapshots; `0` or `None` to disable |
 
 ## Privacy
 

examples/chatgpt_example.py

@@ -5,11 +5,13 @@
 # [tool.uv.sources]
 # wildedge-sdk = { path = "..", editable = true }
 # ///
-"""ChatGPT (OpenAI API): fully manual integration.
+"""ChatGPT (OpenAI API): fully manual integration with explicit hardware capture.
 
 Shows how to instrument a remote LLM with no local model file.
 Tracks input/output token counts, generation config, latency, errors,
-and user feedback without any auto-instrumentation hooks.
+and user feedback without any auto-instrumentation hooks. The background
+hardware sampler is disabled; hardware context is captured explicitly via
+capture_hardware() and passed to track_inference().
 
 Run with: uv run chatgpt_example.py
 Requires: WILDEDGE_DSN and OPENAI_API_KEY environment variables.
@@ -18,14 +20,21 @@
 from openai import OpenAI
 
 import wildedge
-from wildedge import FeedbackType, GenerationConfig, GenerationOutputMeta, TextInputMeta
+from wildedge import (
+    FeedbackType,
+    GenerationConfig,
+    GenerationOutputMeta,
+    TextInputMeta,
+    capture_hardware,
+)
 from wildedge.timing import Timer
 
 MODEL = "gpt-4o"
 MODEL_VERSION = "2024-08-06"
 
 client = wildedge.WildEdge(
     app_version="1.0.0",  # set WILDEDGE_DSN env var
+    sampling_interval_s=None,  # disabled: hardware captured explicitly per call
 )
 
 # Remote models have no local object to inspect, so register with a
@@ -72,6 +81,7 @@
 
         inference_id = handle.track_inference(
             duration_ms=t.elapsed_ms,
+            hardware=capture_hardware(),
             input_modality="text",
             output_modality="text",
             success=True,

examples/gguf_gemma_manual_example.py examples/gguf_gemma_example.pyexamples/gguf_gemma_manual_example.py renamed to examples/gguf_gemma_example.py

@@ -5,12 +5,14 @@
 # [tool.uv.sources]
 # wildedge-sdk = { path = "..", editable = true }
 # ///
-"""Gemma 2 GGUF: fully manual integration, no auto-instrumentation.
+"""Gemma 2 GGUF: fully manual integration with background hardware sampling.
 
 Shows explicit download / load / inference / error tracking without
-client.instrument() or any automatic hooks.
+client.instrument() or any automatic hooks. Hardware context is captured
+automatically on every track_inference() call via the background sampler
+started by WildEdge (sampling_interval_s=30 by default).
 
-Run with: uv run gguf_gemma_manual_example.py
+Run with: uv run gguf_gemma_example.py
 """
 
 import os
@@ -19,7 +21,7 @@
 from llama_cpp import Llama
 
 import wildedge
-from wildedge.events.inference import GenerationOutputMeta, TextInputMeta
+from wildedge import GenerationOutputMeta, TextInputMeta
 from wildedge.timing import Timer
 
 REPO = "bartowski/gemma-2-2b-it-GGUF"

examples/transformers_example.py

@@ -21,19 +21,23 @@
 from __future__ import annotations
 
 import argparse
+import platform
 
 from transformers import pipeline
 
 import wildedge
 
+_DEVICE = "mps" if platform.machine() == "arm64" else "cpu"
+
 
 def run_classify() -> None:
     pipe = pipeline(
         "text-classification",
         model="distilbert-base-uncased-finetuned-sst-2-english",
+        device=_DEVICE,
     )
     inputs = [
-        "I absolutely loved this film — the performances were outstanding!",
+        "I absolutely loved this film, the performances were outstanding!",
         "The service was awful and the food arrived cold.",
         "An average experience, nothing special either way.",
     ]
@@ -47,7 +51,7 @@ def run_classify() -> None:
 
 
 def run_generate() -> None:
-    pipe = pipeline("text-generation", model="gpt2", max_new_tokens=40)
+    pipe = pipeline("text-generation", model="gpt2", max_new_tokens=40, device=_DEVICE)
     prompts = [
         "The future of on-device AI is",
         "Once upon a time, a small robot learned",
@@ -60,7 +64,7 @@ def run_generate() -> None:
 
 
 def run_embed() -> None:
-    pipe = pipeline("feature-extraction", model="bert-base-uncased")
+    pipe = pipeline("feature-extraction", model="bert-base-uncased", device=_DEVICE)
     sentences = [
         "Machine learning is transforming every industry.",
         "On-device inference keeps your data private.",
@@ -88,8 +92,6 @@ def main() -> None:
     )
     args = parser.parse_args()
 
-    # instrument() patches transformers.pipeline and AutoModel.from_pretrained
-    # before any model is loaded; everything below is tracked automatically.
     client = wildedge.WildEdge(app_version="1.0.0")  # set WILDEDGE_DSN env var
     client.instrument("transformers", hubs=["huggingface"])
 

tests/compat/conftest.py

@@ -5,7 +5,7 @@
 import pytest
 
 import wildedge
-from wildedge.device import DeviceInfo
+from wildedge.platforms.device_info import DeviceInfo
 
 
 class _DummyConsumer:

tests/conftest.py

@@ -1,13 +1,39 @@
 """Shared fixtures for the WildEdge SDK test suite."""
 
+import sys
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
 import pytest
 
 from wildedge.client import WildEdge
-from wildedge.device import DeviceInfo
 from wildedge.model import ModelInfo
+from wildedge.platforms.device_info import DeviceInfo
+
+
+@pytest.fixture(autouse=True)
+def reset_hardware_sampler():
+    yield
+    from wildedge.platforms import stop_sampler
+
+    stop_sampler()
+
+
+PLATFORM_MARKS = {
+    "requires_linux": "linux",
+    "requires_macos": "darwin",
+    "requires_windows": "win32",
+}
+
+
+def pytest_collection_modifyitems(items):
+    for item in items:
+        for mark_name, required_platform in PLATFORM_MARKS.items():
+            if item.get_closest_marker(mark_name) and sys.platform != required_platform:
+                item.add_marker(
+                    pytest.mark.skip(reason=f"requires {required_platform}")
+                )
+                break
 
 
 @pytest.fixture
@@ -71,5 +97,8 @@ def client_with_stubbed_runtime():
         patch("wildedge.client.Transmitter"),
         patch("wildedge.client.Consumer"),
     ):
-        client = WildEdge(dsn="https://secret@ingest.wildedge.dev/key")
+        client = WildEdge(
+            dsn="https://secret@ingest.wildedge.dev/key",
+            sampling_interval_s=None,
+        )
     return client

tests/test_batch.py

@@ -4,7 +4,7 @@
 
 from wildedge import constants
 from wildedge.batch import build_batch
-from wildedge.device import DeviceInfo
+from wildedge.platforms.device_info import DeviceInfo
 
 
 def make_device() -> DeviceInfo:

tests/test_cli.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+from pathlib import Path
 
 import pytest
 
@@ -152,7 +153,7 @@ def test_install_runtime_default_flush_timeout_is_shutdown_budget(monkeypatch):
     class FakeWildEdge:
         SUPPORTED_INTEGRATIONS = {"onnx"}
 
-        def __init__(self, *, dsn, app_version, debug):  # type: ignore[no-untyped-def]
+        def __init__(self, *, dsn, app_version, debug, sampling_interval_s=None):  # type: ignore[no-untyped-def]
             pass
 
         def instrument(self, name):  # type: ignore[no-untyped-def]
@@ -182,7 +183,7 @@ def test_install_runtime_instruments_requested_integrations(monkeypatch):
     class FakeWildEdge:
         SUPPORTED_INTEGRATIONS = {"onnx", "torch"}
 
-        def __init__(self, *, dsn, app_version, debug):  # type: ignore[no-untyped-def]
+        def __init__(self, *, dsn, app_version, debug, sampling_interval_s=None):  # type: ignore[no-untyped-def]
             assert dsn == "https://secret@ingest.wildedge.dev/key"
             assert app_version == "2.0.0"
             assert debug is True
@@ -218,7 +219,7 @@ def test_install_runtime_strict_integrations_raises(monkeypatch):
     class FakeWildEdge:
         SUPPORTED_INTEGRATIONS = {"onnx"}
 
-        def __init__(self, *, dsn, app_version, debug):  # type: ignore[no-untyped-def]
+        def __init__(self, *, dsn, app_version, debug, sampling_interval_s=None):  # type: ignore[no-untyped-def]
             pass
 
         def instrument(self, name):  # type: ignore[no-untyped-def]
@@ -330,8 +331,8 @@ def test_doctor_uses_project_key_for_default_namespace(monkeypatch, capsys):
     rc = cli.main(["doctor", "--integrations", "onnx"])
     out = capsys.readouterr().out
     assert rc == 0
-    assert "/test-prod/pending_queue" in out
-    assert "/test-prod/dead_letters" in out
+    assert str(Path("test-prod") / "pending_queue") in out
+    assert str(Path("test-prod") / "dead_letters") in out
 
 
 def test_doctor_uses_app_identity_override_for_namespace(monkeypatch, capsys):
@@ -342,8 +343,8 @@ def test_doctor_uses_app_identity_override_for_namespace(monkeypatch, capsys):
     rc = cli.main(["doctor", "--integrations", "onnx"])
     out = capsys.readouterr().out
     assert rc == 0
-    assert "/my-app/pending_queue" in out
-    assert "/my-app/dead_letters" in out
+    assert str(Path("my-app") / "pending_queue") in out
+    assert str(Path("my-app") / "dead_letters") in out
 
 
 def test_runner_clears_runtime_env_when_no_propagate(monkeypatch):
@@ -368,7 +369,7 @@ def shutdown(self):  # type: ignore[no-untyped-def]
 
 def test_install_runtime_tracks_missing_dependency_status(monkeypatch):
     class FakeWildEdge:
-        def __init__(self, *, dsn, app_version, debug):  # type: ignore[no-untyped-def]
+        def __init__(self, *, dsn, app_version, debug, sampling_interval_s=None):  # type: ignore[no-untyped-def]
             pass
 
         def instrument(self, name):  # type: ignore[no-untyped-def]

tests/test_consumer.py

@@ -3,9 +3,11 @@
 import os
 from unittest.mock import MagicMock, patch
 
+import pytest
+
 from wildedge import constants
 from wildedge.consumer import Consumer
-from wildedge.device import DeviceInfo
+from wildedge.platforms.device_info import DeviceInfo
 from wildedge.queue import EventQueue
 from wildedge.transmitter import IngestResponse, TransmitError, Transmitter
 
@@ -435,6 +437,10 @@ def test_resume_registers_atexit(self, monkeypatch):
 
 
 class TestForkRegistration:
+    @pytest.mark.skipif(
+        not hasattr(os, "register_at_fork"),
+        reason="os.register_at_fork not available on Windows",
+    )
     def test_register_at_fork_wires_pause_and_resume(self, monkeypatch):
         """WildEdge.__init__ registers _pause and _resume via os.register_at_fork."""
         from wildedge.client import WildEdge