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
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
16 changes: 13 additions & 3 deletions examples/chatgpt_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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"
Expand Down
12 changes: 7 additions & 5 deletions examples/transformers_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.",
]
Expand All @@ -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",
Expand All @@ -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.",
Expand Down Expand Up @@ -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"])

Expand Down
2 changes: 1 addition & 1 deletion tests/compat/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
import pytest

import wildedge
from wildedge.device import DeviceInfo
from wildedge.platforms.device_info import DeviceInfo


class _DummyConsumer:
Expand Down
33 changes: 31 additions & 2 deletions tests/conftest.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
2 changes: 1 addition & 1 deletion tests/test_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
17 changes: 9 additions & 8 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import json
import os
from pathlib import Path

import pytest

Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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):
Expand All @@ -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):
Expand All @@ -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]
Expand Down
8 changes: 7 additions & 1 deletion tests/test_consumer.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
Loading
Loading