From 7b62d8b557bb2d93fb11249c27c7d20241b9ad2a Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 09:35:10 +0100 Subject: [PATCH 01/10] Support for inference hardware context --- tests/conftest.py | 17 +++ tests/test_hardware_context.py | 82 +++++++++++++ tests/test_platform_adapter.py | 61 ++++++++++ tests/test_platform_linux.py | 215 +++++++++++++++++++++++++++++++++ tests/test_platform_macos.py | 67 ++++++++++ tests/test_platform_windows.py | 37 ++++++ wildedge/__init__.py | 3 + wildedge/events/inference.py | 5 + wildedge/model.py | 3 + wildedge/platforms/base.py | 7 +- wildedge/platforms/hardware.py | 71 +++++++++++ wildedge/platforms/linux.py | 129 +++++++++++++++++++- wildedge/platforms/macos.py | 159 ++++++++++++++++++++++-- wildedge/platforms/unknown.py | 4 + wildedge/platforms/windows.py | 84 ++++++++++--- 15 files changed, 907 insertions(+), 37 deletions(-) create mode 100644 tests/test_hardware_context.py create mode 100644 tests/test_platform_adapter.py create mode 100644 tests/test_platform_linux.py create mode 100644 tests/test_platform_macos.py create mode 100644 tests/test_platform_windows.py create mode 100644 wildedge/platforms/hardware.py diff --git a/tests/conftest.py b/tests/conftest.py index fce8f2e..e887aec 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,6 @@ """Shared fixtures for the WildEdge SDK test suite.""" +import sys from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -9,6 +10,22 @@ from wildedge.device import DeviceInfo from wildedge.model import ModelInfo +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 def device_info(): diff --git a/tests/test_hardware_context.py b/tests/test_hardware_context.py new file mode 100644 index 0000000..0a10f53 --- /dev/null +++ b/tests/test_hardware_context.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +from wildedge.events.inference import InferenceEvent +from wildedge.platforms.hardware import HardwareContext, ThermalContext, capture +from wildedge.platforms.linux import LinuxPlatform + + +def test_thermal_context_to_dict_omits_none(): + t = ThermalContext(state="fair", state_raw="passive") + assert t.to_dict() == {"state": "fair", "state_raw": "passive"} + + +def test_thermal_context_to_dict_includes_temp(): + t = ThermalContext(state="nominal", state_raw="active", cpu_temp_celsius=42.0) + assert t.to_dict()["cpu_temp_celsius"] == 42.0 + + +def test_hardware_context_to_dict_omits_none_fields(): + hw = HardwareContext(memory_available_bytes=1_000_000) + assert hw.to_dict() == {"memory_available_bytes": 1_000_000} + + +def test_hardware_context_to_dict_includes_thermal(): + hw = HardwareContext( + thermal=ThermalContext( + state="nominal", state_raw="active", cpu_temp_celsius=42.0 + ), + cpu_freq_mhz=2400, + cpu_freq_max_mhz=3200, + ) + d = hw.to_dict() + assert d["thermal"] == { + "state": "nominal", + "state_raw": "active", + "cpu_temp_celsius": 42.0, + } + assert d["cpu_freq_mhz"] == 2400 + + +def test_hardware_context_empty_thermal_not_serialised(): + assert "thermal" not in HardwareContext(thermal=ThermalContext()).to_dict() + + +def test_hardware_context_accelerator_actual_serialised(): + assert ( + HardwareContext(accelerator_actual="cpu").to_dict()["accelerator_actual"] + == "cpu" + ) + + +def test_inference_event_hardware_included_in_dict(): + hw = HardwareContext(memory_available_bytes=512_000_000, accelerator_actual="mps") + inference = InferenceEvent(model_id="m1", duration_ms=10, hardware=hw).to_dict()[ + "inference" + ] + assert inference["hardware"]["memory_available_bytes"] == 512_000_000 + assert inference["hardware"]["accelerator_actual"] == "mps" + + +def test_inference_event_no_hardware_key_when_none(): + assert ( + "hardware" + not in InferenceEvent(model_id="m1", duration_ms=10).to_dict()["inference"] + ) + + +def test_capture_sets_accelerator_actual(monkeypatch): + monkeypatch.setattr("wildedge.platforms.CURRENT_PLATFORM", LinuxPlatform()) + monkeypatch.setattr( + LinuxPlatform, "hardware_context", lambda self: HardwareContext() + ) + assert capture(accelerator_actual="gpu").accelerator_actual == "gpu" + + +def test_capture_preserves_existing_accelerator_when_not_passed(monkeypatch): + monkeypatch.setattr("wildedge.platforms.CURRENT_PLATFORM", LinuxPlatform()) + monkeypatch.setattr( + LinuxPlatform, + "hardware_context", + lambda self: HardwareContext(accelerator_actual="mps"), + ) + assert capture().accelerator_actual == "mps" diff --git a/tests/test_platform_adapter.py b/tests/test_platform_adapter.py new file mode 100644 index 0000000..6c754cf --- /dev/null +++ b/tests/test_platform_adapter.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import pytest + +from wildedge.platforms.hardware import HardwareContext +from wildedge.platforms.linux import LinuxPlatform +from wildedge.platforms.macos import MacOSPlatform +from wildedge.platforms.windows import WindowsPlatform + +PLATFORMS = [ + pytest.param(LinuxPlatform, id="linux", marks=pytest.mark.requires_linux), + pytest.param(MacOSPlatform, id="macos", marks=pytest.mark.requires_macos), + pytest.param(WindowsPlatform, id="windows", marks=pytest.mark.requires_windows), +] + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_device_model(platform_cls): + model = platform_cls().device_model() + if model is not None: + assert isinstance(model, str) and model + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_os_version(platform_cls): + ver = platform_cls().os_version() + if ver is not None: + assert isinstance(ver, str) and ver + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_ram_bytes(platform_cls): + total = platform_cls().ram_bytes() + assert isinstance(total, int) and total > 0 + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_disk_bytes(platform_cls): + total = platform_cls().disk_bytes() + assert isinstance(total, int) and total > 0 + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_gpu_accelerators(platform_cls): + accs, gpu_name = platform_cls().gpu_accelerators() + assert isinstance(accs, list) + assert all(isinstance(a, str) for a in accs) + if gpu_name is not None: + assert isinstance(gpu_name, str) and gpu_name + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_gpu_accelerator_for_offload(platform_cls): + offload = platform_cls().gpu_accelerator_for_offload() + assert isinstance(offload, str) and offload + + +@pytest.mark.parametrize("platform_cls", PLATFORMS) +def test_real_hardware_context(platform_cls): + ctx = platform_cls().hardware_context() + assert isinstance(ctx, HardwareContext) diff --git a/tests/test_platform_linux.py b/tests/test_platform_linux.py new file mode 100644 index 0000000..973f7d6 --- /dev/null +++ b/tests/test_platform_linux.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from wildedge.platforms.hardware import ThermalContext +from wildedge.platforms.linux import ( + CPU_THERMAL_ZONE_TYPES, + TRIP_POINT_STATES, + LinuxPlatform, +) + + +def test_meminfo_parses_total_and_available(monkeypatch): + content = "MemTotal: 16384000 kB\nMemFree: 1000000 kB\nMemAvailable: 8192000 kB\n" + monkeypatch.setattr( + "builtins.open", + lambda p, **kw: ( + __import__("io").StringIO(content) + if p == "/proc/meminfo" + else open(p, **kw) + ), + ) + total, available = LinuxPlatform().meminfo() + assert total == 16384000 * 1024 + assert available == 8192000 * 1024 + + +def test_ram_bytes_delegates_to_meminfo(monkeypatch): + monkeypatch.setattr( + LinuxPlatform, "meminfo", lambda self: (8_000_000_000, 4_000_000_000) + ) + assert LinuxPlatform().ram_bytes() == 8_000_000_000 + + +def test_cpu_freq_converts_khz_to_mhz(tmp_path, monkeypatch): + cpufreq = tmp_path / "cpufreq" + cpufreq.mkdir() + (cpufreq / "scaling_cur_freq").write_text("2400000\n") + (cpufreq / "cpuinfo_max_freq").write_text("3200000\n") + + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: cpufreq if str(p).endswith("cpu0/cpufreq") else Path(p), + ) + cur, max_f = LinuxPlatform().cpu_freq() + assert cur == 2400 + assert max_f == 3200 + + +def test_battery_reads_capacity_and_status(tmp_path, monkeypatch): + ps = tmp_path / "power_supply" + bat = ps / "BAT0" + bat.mkdir(parents=True) + (bat / "capacity").write_text("78\n") + (bat / "status").write_text("Discharging\n") + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: ps if p == "/sys/class/power_supply" else Path(p), + ) + level, charging = LinuxPlatform().battery() + assert level == pytest.approx(0.78) + assert charging is False + + +def test_battery_charging_status(tmp_path, monkeypatch): + ps = tmp_path / "power_supply" + bat = ps / "BAT0" + bat.mkdir(parents=True) + (bat / "capacity").write_text("95\n") + (bat / "status").write_text("Charging\n") + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: ps if p == "/sys/class/power_supply" else Path(p), + ) + _, charging = LinuxPlatform().battery() + assert charging is True + + +def test_battery_no_supply_returns_none(tmp_path, monkeypatch): + ps = tmp_path / "empty_ps" + ps.mkdir() + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: ps if p == "/sys/class/power_supply" else Path(p), + ) + level, charging = LinuxPlatform().battery() + assert level is None and charging is None + + +def test_thermal_nominal(tmp_path, monkeypatch): + zone = tmp_path / "thermal_zone0" + zone.mkdir() + (zone / "type").write_text("x86_pkg_temp\n") + (zone / "temp").write_text("45000\n") + (zone / "trip_point_0_type").write_text("passive\n") + (zone / "trip_point_0_temp").write_text("80000\n") + (zone / "trip_point_1_type").write_text("critical\n") + (zone / "trip_point_1_temp").write_text("100000\n") + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: tmp_path if p == "/sys/class/thermal" else Path(p), + ) + ctx = LinuxPlatform().thermal() + assert ctx is not None + assert ctx.state == "nominal" + assert ctx.state_raw == "active" + assert ctx.cpu_temp_celsius == pytest.approx(45.0) + + +def test_thermal_serious(tmp_path, monkeypatch): + zone = tmp_path / "thermal_zone0" + zone.mkdir() + (zone / "type").write_text("cpu\n") + (zone / "temp").write_text("90000\n") + (zone / "trip_point_0_type").write_text("passive\n") + (zone / "trip_point_0_temp").write_text("80000\n") + (zone / "trip_point_1_type").write_text("hot\n") + (zone / "trip_point_1_temp").write_text("85000\n") + (zone / "trip_point_2_type").write_text("critical\n") + (zone / "trip_point_2_temp").write_text("100000\n") + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: tmp_path if p == "/sys/class/thermal" else Path(p), + ) + ctx = LinuxPlatform().thermal() + assert ctx is not None + assert ctx.state == "serious" + assert ctx.state_raw == "hot" + + +def test_thermal_skips_non_cpu_zone(tmp_path, monkeypatch): + zone = tmp_path / "thermal_zone0" + zone.mkdir() + (zone / "type").write_text("iwlwifi_1\n") + (zone / "temp").write_text("50000\n") + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: tmp_path if p == "/sys/class/thermal" else Path(p), + ) + assert LinuxPlatform().thermal() is None + + +def test_thermal_missing_sysfs_returns_none(tmp_path, monkeypatch): + monkeypatch.setattr( + "wildedge.platforms.linux.Path", + lambda p: tmp_path / "no_thermal" if p == "/sys/class/thermal" else Path(p), + ) + assert LinuxPlatform().thermal() is None + + +def test_hardware_context_assembles_fields(monkeypatch): + monkeypatch.setattr(LinuxPlatform, "meminfo", lambda self: (None, 2_000_000_000)) + monkeypatch.setattr(LinuxPlatform, "battery", lambda self: (0.5, True)) + monkeypatch.setattr(LinuxPlatform, "cpu_freq", lambda self: (1800, 3200)) + monkeypatch.setattr( + LinuxPlatform, + "thermal", + lambda self: ThermalContext(state="fair", state_raw="passive"), + ) + ctx = LinuxPlatform().hardware_context() + assert ctx.memory_available_bytes == 2_000_000_000 + assert ctx.battery_level == pytest.approx(0.5) + assert ctx.battery_charging is True + assert ctx.cpu_freq_mhz == 1800 + assert ctx.cpu_freq_max_mhz == 3200 + assert ctx.thermal.state == "fair" + + +def test_cpu_thermal_zone_types_is_tuple_of_strings(): + assert isinstance(CPU_THERMAL_ZONE_TYPES, tuple) + assert all(isinstance(k, str) for k in CPU_THERMAL_ZONE_TYPES) + + +def test_trip_point_states_covers_all_severity_levels(): + trip_types = {row[0] for row in TRIP_POINT_STATES} + assert {"critical", "hot", "passive"} <= trip_types + + +@pytest.mark.requires_linux +def test_real_meminfo(): + total, available = LinuxPlatform().meminfo() + assert isinstance(total, int) and total > 0 + assert isinstance(available, int) and 0 < available <= total + + +@pytest.mark.requires_linux +def test_real_cpu_freq(): + cur, max_f = LinuxPlatform().cpu_freq() + if cur is not None: + assert 0 < cur <= 10_000 + if max_f is not None: + assert 0 < max_f <= 10_000 + if cur is not None and max_f is not None: + assert cur <= max_f + + +@pytest.mark.requires_linux +def test_real_battery(): + level, charging = LinuxPlatform().battery() + if level is not None: + assert 0.0 <= level <= 1.0 + if charging is not None: + assert isinstance(charging, bool) + + +@pytest.mark.requires_linux +def test_real_thermal(): + ctx = LinuxPlatform().thermal() + if ctx is not None: + assert ctx.state in ("nominal", "fair", "serious", "critical") + assert isinstance(ctx.state_raw, str) + if ctx.cpu_temp_celsius is not None: + assert -10 <= ctx.cpu_temp_celsius <= 150 diff --git a/tests/test_platform_macos.py b/tests/test_platform_macos.py new file mode 100644 index 0000000..9506321 --- /dev/null +++ b/tests/test_platform_macos.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import pytest + +from wildedge.platforms.macos import MacOSPlatform + + +def test_meminfo_computes_available(monkeypatch): + from wildedge.platforms import macos as macos_mod + + monkeypatch.setattr( + macos_mod, + "_sysctl_uint64", + lambda name: 16_000_000_000 if name == b"hw.memsize" else None, + ) + monkeypatch.setattr( + macos_mod, + "_sysctl_uint32", + lambda name: {b"hw.pagesize": 4096, b"vm.page_free_count": 100_000}.get(name), + ) + total, available = MacOSPlatform().meminfo() + assert total == 16_000_000_000 + assert available == 4096 * 100_000 + + +def test_meminfo_available_none_when_page_size_missing(monkeypatch): + from wildedge.platforms import macos as macos_mod + + monkeypatch.setattr(macos_mod, "_sysctl_uint64", lambda name: 16_000_000_000) + monkeypatch.setattr(macos_mod, "_sysctl_uint32", lambda name: None) + _, available = MacOSPlatform().meminfo() + assert available is None + + +def test_ram_bytes_delegates_to_meminfo(monkeypatch): + monkeypatch.setattr( + MacOSPlatform, "meminfo", lambda self: (8_000_000_000, 1_000_000_000) + ) + assert MacOSPlatform().ram_bytes() == 8_000_000_000 + + +def test_hardware_context_fields(monkeypatch): + monkeypatch.setattr(MacOSPlatform, "meminfo", lambda self: (None, 3_000_000_000)) + monkeypatch.setattr(MacOSPlatform, "battery", lambda self: (0.8, False)) + ctx = MacOSPlatform().hardware_context() + assert ctx.memory_available_bytes == 3_000_000_000 + assert ctx.battery_level == pytest.approx(0.8) + assert ctx.battery_charging is False + assert ctx.cpu_freq_mhz is None + assert ctx.thermal is None + + +@pytest.mark.requires_macos +def test_real_meminfo(): + total, available = MacOSPlatform().meminfo() + assert isinstance(total, int) and total > 0 + if available is not None: + assert 0 < available <= total + + +@pytest.mark.requires_macos +def test_real_battery(): + level, charging = MacOSPlatform().battery() + if level is not None: + assert 0.0 <= level <= 1.0 + if charging is not None: + assert isinstance(charging, bool) diff --git a/tests/test_platform_windows.py b/tests/test_platform_windows.py new file mode 100644 index 0000000..862d7eb --- /dev/null +++ b/tests/test_platform_windows.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +import pytest + +from wildedge.platforms.windows import WindowsPlatform + + +def test_ram_bytes_delegates_to_meminfo(monkeypatch): + monkeypatch.setattr( + WindowsPlatform, "meminfo", lambda self: (16_000_000_000, 4_000_000_000) + ) + assert WindowsPlatform().ram_bytes() == 16_000_000_000 + + +def test_hardware_context_fields(monkeypatch): + monkeypatch.setattr(WindowsPlatform, "meminfo", lambda self: (None, 6_000_000_000)) + monkeypatch.setattr(WindowsPlatform, "battery", lambda self: (0.6, True)) + ctx = WindowsPlatform().hardware_context() + assert ctx.memory_available_bytes == 6_000_000_000 + assert ctx.battery_level == pytest.approx(0.6) + assert ctx.battery_charging is True + + +@pytest.mark.requires_windows +def test_real_meminfo(): + total, available = WindowsPlatform().meminfo() + assert isinstance(total, int) and total > 0 + assert isinstance(available, int) and 0 < available <= total + + +@pytest.mark.requires_windows +def test_real_battery(): + level, charging = WindowsPlatform().battery() + if level is not None: + assert 0.0 <= level <= 1.0 + if charging is not None: + assert isinstance(charging, bool) diff --git a/wildedge/__init__.py b/wildedge/__init__.py index 087492d..c153200 100644 --- a/wildedge/__init__.py +++ b/wildedge/__init__.py @@ -17,11 +17,14 @@ ImageInputMeta, TextInputMeta, ) +from wildedge.platforms.hardware import HardwareContext, ThermalContext from wildedge.queue import QueuePolicy from wildedge.timing import Timer __all__ = [ "WildEdge", + "HardwareContext", + "ThermalContext", "track", "QueuePolicy", "DeviceInfo", diff --git a/wildedge/events/inference.py b/wildedge/events/inference.py index 8a4aa7c..b730020 100644 --- a/wildedge/events/inference.py +++ b/wildedge/events/inference.py @@ -5,6 +5,8 @@ from datetime import datetime, timezone from typing import Any +from wildedge.platforms.hardware import HardwareContext + @dataclass class HistogramSummary: @@ -278,6 +280,7 @@ class InferenceEvent: | None ) = None generation_config: GenerationConfig | None = None + hardware: HardwareContext | None = None event_id: str = field(default_factory=lambda: str(uuid.uuid4())) timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) inference_id: str = field(default_factory=lambda: str(uuid.uuid4())) @@ -302,6 +305,8 @@ def to_dict(self) -> dict: inference_data["output_meta"] = self.output_meta.to_dict() if self.generation_config is not None: inference_data["generation_config"] = self.generation_config.to_dict() + if self.hardware is not None: + inference_data["hardware"] = self.hardware.to_dict() return { "event_id": self.event_id, diff --git a/wildedge/model.py b/wildedge/model.py index 5b66549..ec9de62 100644 --- a/wildedge/model.py +++ b/wildedge/model.py @@ -25,6 +25,7 @@ TextInputMeta, ) from wildedge.logging import logger +from wildedge.platforms.hardware import HardwareContext @dataclass @@ -141,6 +142,7 @@ def track_inference( | EmbeddingOutputMeta | None = None, generation_config: GenerationConfig | None = None, + hardware: HardwareContext | None = None, ) -> str: event = InferenceEvent( model_id=self.model_id, @@ -153,6 +155,7 @@ def track_inference( input_meta=input_meta, output_meta=output_meta, generation_config=generation_config, + hardware=hardware, ) self.last_inference_id = event.inference_id self.publish(event.to_dict()) diff --git a/wildedge/platforms/base.py b/wildedge/platforms/base.py index 692ca9d..edaf7cc 100644 --- a/wildedge/platforms/base.py +++ b/wildedge/platforms/base.py @@ -2,10 +2,13 @@ import ctypes from pathlib import Path -from typing import Protocol +from typing import TYPE_CHECKING, Protocol from wildedge.logging import logger +if TYPE_CHECKING: + from wildedge.platforms.hardware import HardwareContext + class PlatformAdapter(Protocol): wire_type: str @@ -26,6 +29,8 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: ... def gpu_accelerator_for_offload(self) -> str: ... + def hardware_context(self) -> HardwareContext: ... + def debug_detection_failure(context: str, exc: BaseException) -> None: logger.debug("wildedge: device detection failed for %s: %s", context, exc) diff --git a/wildedge/platforms/hardware.py b/wildedge/platforms/hardware.py new file mode 100644 index 0000000..09ec36f --- /dev/null +++ b/wildedge/platforms/hardware.py @@ -0,0 +1,71 @@ +"""Hardware context dataclasses and inference-time capture.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + + +@dataclass +class ThermalContext: + state: str | None = None # nominal|fair|serious|critical + state_raw: str | None = None # platform-native value + cpu_temp_celsius: float | None = None + + def to_dict(self) -> dict: + return { + k: v + for k, v in { + "state": self.state, + "state_raw": self.state_raw, + "cpu_temp_celsius": self.cpu_temp_celsius, + }.items() + if v is not None + } + + +@dataclass +class HardwareContext: + """Device hardware state at the time of inference.""" + + thermal: ThermalContext | None = None + battery_level: float | None = None + battery_charging: bool | None = None + memory_available_bytes: int | None = None + cpu_freq_mhz: int | None = None + cpu_freq_max_mhz: int | None = None + accelerator_actual: str | None = None + + def to_dict(self) -> dict: + d: dict[str, Any] = {} + if self.thermal is not None: + thermal_dict = self.thermal.to_dict() + if thermal_dict: + d["thermal"] = thermal_dict + for k, v in [ + ("battery_level", self.battery_level), + ("battery_charging", self.battery_charging), + ("memory_available_bytes", self.memory_available_bytes), + ("cpu_freq_mhz", self.cpu_freq_mhz), + ("cpu_freq_max_mhz", self.cpu_freq_max_mhz), + ("accelerator_actual", self.accelerator_actual), + ]: + if v is not None: + d[k] = v + return d + + +def capture(accelerator_actual: str | None = None) -> HardwareContext: + """Read current hardware state. Never raises. + + Pass ``accelerator_actual`` if the runtime reports which accelerator was + used; it may differ from the accelerator requested at load time. + """ + # Lazy import avoids a circular dependency: __init__.py imports platform + # modules which import this file. + from wildedge.platforms import CURRENT_PLATFORM # noqa: PLC0415 + + ctx = CURRENT_PLATFORM.hardware_context() + if accelerator_actual is not None: + ctx.accelerator_actual = accelerator_actual + return ctx diff --git a/wildedge/platforms/linux.py b/wildedge/platforms/linux.py index 95c4daf..eb3d4f7 100644 --- a/wildedge/platforms/linux.py +++ b/wildedge/platforms/linux.py @@ -11,6 +11,19 @@ hip_device_count, nvml_gpu_name, ) +from wildedge.platforms.hardware import HardwareContext, ThermalContext + +# Thermal zone type substrings that identify a CPU zone. +CPU_THERMAL_ZONE_TYPES = ("cpu", "x86_pkg", "acpi", "soc", "pkg") + +# Maps sysfs trip-point type -> (normalized state, state_raw). +# Evaluated in priority order (highest severity first). +TRIP_POINT_STATES: tuple[tuple[str, str, str], ...] = ( + ("critical", "critical", "critical"), + ("hot", "serious", "hot"), + ("passive", "fair", "passive"), +) +TRIP_POINT_DEFAULT = ("nominal", "active") class LinuxPlatform: @@ -52,15 +65,25 @@ def os_version(self) -> str | None: except Exception: return None - def ram_bytes(self) -> int | None: + def meminfo(self) -> tuple[int | None, int | None]: + """Return (total_bytes, available_bytes) from a single /proc/meminfo read.""" + total: int | None = None + available: int | None = None try: with open("/proc/meminfo") as f: for line in f: if line.startswith("MemTotal:"): - return int(line.split()[1]) * 1024 + total = int(line.split()[1]) * 1024 + elif line.startswith("MemAvailable:"): + available = int(line.split()[1]) * 1024 + if total is not None and available is not None: + break except Exception as exc: - debug_detection_failure("linux ram_bytes", exc) - return None + debug_detection_failure("linux meminfo", exc) + return total, available + + def ram_bytes(self) -> int | None: + return self.meminfo()[0] def disk_bytes(self) -> int | None: try: @@ -82,3 +105,101 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: def gpu_accelerator_for_offload(self) -> str: accs, _ = self.gpu_accelerators() return accs[0] if accs else "cpu" + + def hardware_context(self) -> HardwareContext: + bat_level, bat_charging = self.battery() + cpu_cur, cpu_max = self.cpu_freq() + _, mem_available = self.meminfo() + return HardwareContext( + thermal=self.thermal(), + battery_level=bat_level, + battery_charging=bat_charging, + memory_available_bytes=mem_available, + cpu_freq_mhz=cpu_cur, + cpu_freq_max_mhz=cpu_max, + ) + + def cpu_freq(self) -> tuple[int | None, int | None]: + """Read current and max CPU frequency (MHz) from cpufreq for cpu0.""" + base = Path("/sys/devices/system/cpu/cpu0/cpufreq") + try: + cur_path = base / "scaling_cur_freq" + max_path = base / "cpuinfo_max_freq" + cur = ( + int(cur_path.read_text().strip()) // 1000 if cur_path.exists() else None + ) + max_f = ( + int(max_path.read_text().strip()) // 1000 if max_path.exists() else None + ) + return cur, max_f + except Exception as exc: + debug_detection_failure("linux cpu_freq", exc) + return None, None + + def battery(self) -> tuple[float | None, bool | None]: + """Read level (0.0-1.0) and charging state from the first BAT* supply.""" + try: + for bat in sorted(Path("/sys/class/power_supply").glob("BAT*")): + cap_path = bat / "capacity" + status_path = bat / "status" + level = ( + int(cap_path.read_text().strip()) / 100.0 + if cap_path.exists() + else None + ) + status_raw = ( + status_path.read_text().strip().lower() + if status_path.exists() + else None + ) + charging = status_raw in ("charging", "full") if status_raw else None + return level, charging + except Exception as exc: + debug_detection_failure("linux battery", exc) + return None, None + + def thermal(self) -> ThermalContext | None: + """Read the first CPU thermal zone and map trip points to a normalized state.""" + try: + thermal_dir = Path("/sys/class/thermal") + if not thermal_dir.exists(): + return None + + for zone in sorted(thermal_dir.glob("thermal_zone*")): + try: + zone_type = (zone / "type").read_text().strip().lower() + except Exception: + continue + if not any(k in zone_type for k in CPU_THERMAL_ZONE_TYPES): + continue + + try: + temp_c = int((zone / "temp").read_text().strip()) / 1000.0 + except Exception: + continue + + trip_temps: dict[str, float] = {} + for trip_path in sorted(zone.glob("trip_point_*_type")): + try: + ttype = trip_path.read_text().strip().lower() + idx = trip_path.name[len("trip_point_") : -len("_type")] + ttemp = ( + int((zone / f"trip_point_{idx}_temp").read_text().strip()) + / 1000.0 + ) + trip_temps[ttype] = ttemp + except Exception: + continue + + state, state_raw = TRIP_POINT_DEFAULT + for trip_type, norm_state, norm_raw in TRIP_POINT_STATES: + if trip_type in trip_temps and temp_c >= trip_temps[trip_type]: + state, state_raw = norm_state, norm_raw + break + + return ThermalContext( + state=state, state_raw=state_raw, cpu_temp_celsius=temp_c + ) + except Exception as exc: + debug_detection_failure("linux thermal", exc) + return None diff --git a/wildedge/platforms/macos.py b/wildedge/platforms/macos.py index 6949cc1..1b6b1ee 100644 --- a/wildedge/platforms/macos.py +++ b/wildedge/platforms/macos.py @@ -6,6 +6,34 @@ from pathlib import Path from wildedge.platforms.base import debug_detection_failure +from wildedge.platforms.hardware import HardwareContext + +_libc = ctypes.CDLL(None) + +_CF_STRING_ENCODING_UTF8 = 0x08000100 +_CF_NUMBER_SINT32_TYPE = 3 + + +def _sysctl_uint32(name: bytes) -> int | None: + try: + buf = ctypes.c_uint32(0) + size = ctypes.c_size_t(ctypes.sizeof(buf)) + ret = _libc.sysctlbyname(name, ctypes.byref(buf), ctypes.byref(size), None, 0) + return int(buf.value) if ret == 0 else None + except Exception as exc: + debug_detection_failure(f"macos sysctl_uint32({name!r})", exc) + return None + + +def _sysctl_uint64(name: bytes) -> int | None: + try: + buf = ctypes.c_uint64(0) + size = ctypes.c_size_t(ctypes.sizeof(buf)) + ret = _libc.sysctlbyname(name, ctypes.byref(buf), ctypes.byref(size), None, 0) + return int(buf.value) if ret == 0 else None + except Exception as exc: + debug_detection_failure(f"macos sysctl_uint64({name!r})", exc) + return None class MacOSPlatform: @@ -24,9 +52,7 @@ def device_model(self) -> str | None: try: buf = ctypes.create_string_buffer(128) size = ctypes.c_size_t(ctypes.sizeof(buf)) - ret = ctypes.CDLL(None).sysctlbyname( - b"hw.model", buf, ctypes.byref(size), None, 0 - ) + ret = _libc.sysctlbyname(b"hw.model", buf, ctypes.byref(size), None, 0) if ret == 0: val = buf.value.decode().strip() return val or None @@ -34,17 +60,20 @@ def device_model(self) -> str | None: debug_detection_failure("macos device_model", exc) return None + def meminfo(self) -> tuple[int | None, int | None]: + """Return (total_bytes, available_bytes) from sysctl.""" + total = _sysctl_uint64(b"hw.memsize") + page_size = _sysctl_uint32(b"hw.pagesize") + free_count = _sysctl_uint32(b"vm.page_free_count") + available = ( + free_count * page_size + if free_count is not None and page_size is not None + else None + ) + return total, available + def ram_bytes(self) -> int | None: - try: - buf = ctypes.c_uint64(0) - size = ctypes.c_size_t(ctypes.sizeof(buf)) - ret = ctypes.CDLL(None).sysctlbyname( - b"hw.memsize", ctypes.byref(buf), ctypes.byref(size), None, 0 - ) - return buf.value if ret == 0 else None - except Exception as exc: - debug_detection_failure("macos ram_bytes", exc) - return None + return self.meminfo()[0] def disk_bytes(self) -> int | None: try: @@ -68,3 +97,107 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: def gpu_accelerator_for_offload(self) -> str: return "mps" if platform.machine() == "arm64" else "cpu" + + def battery(self) -> tuple[float | None, bool | None]: + """Read battery level and charging state via IOKit AppleSmartBattery.""" + try: + iokit = ctypes.cdll.LoadLibrary( + "/System/Library/Frameworks/IOKit.framework/IOKit" + ) + cf = ctypes.cdll.LoadLibrary( + "/System/Library/Frameworks/CoreFoundation.framework/CoreFoundation" + ) + + iokit.IOServiceMatching.restype = ctypes.c_void_p + iokit.IOIteratorNext.restype = ctypes.c_uint32 + iokit.IORegistryEntryCreateCFProperties.restype = ctypes.c_int32 + iokit.IOObjectRelease.argtypes = [ctypes.c_uint32] + + cf.CFDictionaryGetValue.restype = ctypes.c_void_p + cf.CFNumberGetValue.argtypes = [ + ctypes.c_void_p, + ctypes.c_int32, + ctypes.c_void_p, + ] + cf.CFBooleanGetValue.restype = ctypes.c_bool + cf.CFStringCreateWithCString.restype = ctypes.c_void_p + cf.CFRelease.argtypes = [ctypes.c_void_p] + + matching = iokit.IOServiceMatching(b"AppleSmartBattery") + if not matching: + return None, None + + iterator = ctypes.c_uint32(0) + if ( + iokit.IOServiceGetMatchingServices( + ctypes.c_uint32(0), + ctypes.c_void_p(matching), + ctypes.byref(iterator), + ) + != 0 + ): + return None, None + + service = iokit.IOIteratorNext(iterator) + iokit.IOObjectRelease(iterator) + if not service: + return None, None + + props = ctypes.c_void_p(0) + ret = iokit.IORegistryEntryCreateCFProperties( + service, ctypes.byref(props), None, 0 + ) + iokit.IOObjectRelease(service) + if ret != 0: + return None, None + + def cfstr(s: bytes) -> ctypes.c_void_p: + return ctypes.c_void_p( + cf.CFStringCreateWithCString(None, s, _CF_STRING_ENCODING_UTF8) + ) + + cap_key = cfstr(b"CurrentCapacity") + max_key = cfstr(b"MaxCapacity") + chg_key = cfstr(b"IsCharging") + + level: float | None = None + charging: bool | None = None + + cap_ref = cf.CFDictionaryGetValue(props, cap_key) + max_ref = cf.CFDictionaryGetValue(props, max_key) + if cap_ref and max_ref: + cur = ctypes.c_int32(0) + mx = ctypes.c_int32(0) + cf.CFNumberGetValue( + ctypes.c_void_p(cap_ref), _CF_NUMBER_SINT32_TYPE, ctypes.byref(cur) + ) + cf.CFNumberGetValue( + ctypes.c_void_p(max_ref), _CF_NUMBER_SINT32_TYPE, ctypes.byref(mx) + ) + if mx.value > 0: + level = cur.value / mx.value + + chg_ref = cf.CFDictionaryGetValue(props, chg_key) + if chg_ref: + charging = bool(cf.CFBooleanGetValue(ctypes.c_void_p(chg_ref))) + + cf.CFRelease(props) + for k in (cap_key, max_key, chg_key): + if k: + cf.CFRelease(k) + + return level, charging + except Exception as exc: + debug_detection_failure("macos battery", exc) + return None, None + + def hardware_context(self) -> HardwareContext: + _, mem_available = self.meminfo() + bat_level, bat_charging = self.battery() + return HardwareContext( + memory_available_bytes=mem_available, + battery_level=bat_level, + battery_charging=bat_charging, + # cpu_freq not available on macOS (no public API on Apple Silicon) + # thermal not available via public Python API + ) diff --git a/wildedge/platforms/unknown.py b/wildedge/platforms/unknown.py index dc72a36..43039fe 100644 --- a/wildedge/platforms/unknown.py +++ b/wildedge/platforms/unknown.py @@ -6,6 +6,7 @@ from pathlib import Path from wildedge.platforms.base import debug_detection_failure +from wildedge.platforms.hardware import HardwareContext class UnknownPlatform: @@ -45,3 +46,6 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: def gpu_accelerator_for_offload(self) -> str: return "cpu" + + def hardware_context(self) -> HardwareContext: + return HardwareContext() diff --git a/wildedge/platforms/windows.py b/wildedge/platforms/windows.py index aabce70..4d2504f 100644 --- a/wildedge/platforms/windows.py +++ b/wildedge/platforms/windows.py @@ -11,6 +11,7 @@ debug_detection_failure, nvml_gpu_name, ) +from wildedge.platforms.hardware import HardwareContext try: import winreg as _winreg # type: ignore[import] # Windows only @@ -18,6 +19,31 @@ _winreg = None # type: ignore[assignment] +class _MEMORYSTATUSEX(ctypes.Structure): + _fields_ = [ + ("dwLength", ctypes.c_ulong), + ("dwMemoryLoad", ctypes.c_ulong), + ("ullTotalPhys", ctypes.c_ulonglong), + ("ullAvailPhys", ctypes.c_ulonglong), + ("ullTotalPageFile", ctypes.c_ulonglong), + ("ullAvailPageFile", ctypes.c_ulonglong), + ("ullTotalVirtual", ctypes.c_ulonglong), + ("ullAvailVirtual", ctypes.c_ulonglong), + ("ullAvailExtendedVirtual", ctypes.c_ulonglong), + ] + + +class _SystemPowerStatus(ctypes.Structure): + _fields_ = [ + ("ACLineStatus", ctypes.c_byte), + ("BatteryFlag", ctypes.c_byte), + ("BatteryLifePercent", ctypes.c_byte), + ("SystemStatusFlag", ctypes.c_byte), + ("BatteryLifeTime", ctypes.c_ulong), + ("BatteryFullLifeTime", ctypes.c_ulong), + ] + + class WindowsPlatform: wire_type = "windows" @@ -52,29 +78,19 @@ def os_version(self) -> str | None: debug_detection_failure("windows os_version", exc) return None - def ram_bytes(self) -> int | None: + def meminfo(self) -> tuple[int | None, int | None]: + """Return (total_bytes, available_bytes) from a single GlobalMemoryStatusEx call.""" try: - - class MEMORYSTATUSEX(ctypes.Structure): - _fields_ = [ - ("dwLength", ctypes.c_ulong), - ("dwMemoryLoad", ctypes.c_ulong), - ("ullTotalPhys", ctypes.c_ulonglong), - ("ullAvailPhys", ctypes.c_ulonglong), - ("ullTotalPageFile", ctypes.c_ulonglong), - ("ullAvailPageFile", ctypes.c_ulonglong), - ("ullTotalVirtual", ctypes.c_ulonglong), - ("ullAvailVirtual", ctypes.c_ulonglong), - ("ullAvailExtendedVirtual", ctypes.c_ulonglong), - ] - - stat = MEMORYSTATUSEX() + stat = _MEMORYSTATUSEX() stat.dwLength = ctypes.sizeof(stat) ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat)) # type: ignore[attr-defined] - return stat.ullTotalPhys + return stat.ullTotalPhys, stat.ullAvailPhys except Exception as exc: - debug_detection_failure("windows ram_bytes", exc) - return None + debug_detection_failure("windows meminfo", exc) + return None, None + + def ram_bytes(self) -> int | None: + return self.meminfo()[0] def disk_bytes(self) -> int | None: try: @@ -92,3 +108,33 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: def gpu_accelerator_for_offload(self) -> str: accs, _ = self.gpu_accelerators() return accs[0] if accs else "cpu" + + def battery(self) -> tuple[float | None, bool | None]: + """Read battery level and charging state via GetSystemPowerStatus.""" + try: + status = _SystemPowerStatus() + if not ctypes.windll.kernel32.GetSystemPowerStatus(ctypes.byref(status)): # type: ignore[attr-defined] + return None, None + # BatteryLifePercent == 255 means unknown + level = ( + status.BatteryLifePercent / 100.0 + if status.BatteryLifePercent != 255 + else None + ) + # ACLineStatus: 0=offline, 1=online, 255=unknown + charging = ( + bool(status.ACLineStatus == 1) if status.ACLineStatus != 255 else None + ) + return level, charging + except Exception as exc: + debug_detection_failure("windows battery", exc) + return None, None + + def hardware_context(self) -> HardwareContext: + _, mem_available = self.meminfo() + bat_level, bat_charging = self.battery() + return HardwareContext( + memory_available_bytes=mem_available, + battery_level=bat_level, + battery_charging=bat_charging, + ) From de13dd5383dcac9553e6f42b3fb6b56f898e0e5b Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 09:44:44 +0100 Subject: [PATCH 02/10] fix: tox isolation (wsl fix) --- tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/tox.ini b/tox.ini index c7cfa53..93b2f6b 100644 --- a/tox.ini +++ b/tox.ini @@ -1,6 +1,7 @@ [tox] envlist = py310, py311, py312, py313 skip_missing_interpreters = true +isolated_build = true [testenv] package = wheel From 20baa3308ff5c5c51e82335fefe803d8e6ee9d76 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 10:07:11 +0100 Subject: [PATCH 03/10] fix: windows fixes --- tests/test_cli.py | 9 +++++---- tests/test_consumer.py | 6 ++++++ tests/test_hubs.py | 7 ++++++- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/tests/test_cli.py b/tests/test_cli.py index 4db4ea9..b294547 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -2,6 +2,7 @@ import json import os +from pathlib import Path import pytest @@ -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): diff --git a/tests/test_consumer.py b/tests/test_consumer.py index 5553a47..ca0a771 100644 --- a/tests/test_consumer.py +++ b/tests/test_consumer.py @@ -3,6 +3,8 @@ import os from unittest.mock import MagicMock, patch +import pytest + from wildedge import constants from wildedge.consumer import Consumer from wildedge.device import DeviceInfo @@ -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 diff --git a/tests/test_hubs.py b/tests/test_hubs.py index c2b6a2d..bdb1108 100644 --- a/tests/test_hubs.py +++ b/tests/test_hubs.py @@ -6,6 +6,8 @@ import types from unittest.mock import patch +import pytest + from wildedge.hubs.huggingface import HuggingFaceHubTracker from wildedge.hubs.torchhub import TorchHubTracker @@ -18,7 +20,10 @@ 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) + try: + link.symlink_to(real_file) + except OSError: + pytest.skip("symlinks require elevated privileges on Windows") tracker = HuggingFaceHubTracker() with patch.object(tracker, "cache_dir", return_value=str(tmp_path)): From bba7b376f31b4f19b16182ddd346428acd5be29a Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 10:11:41 +0100 Subject: [PATCH 04/10] fix: windows fixes --- wildedge/platforms/macos.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/wildedge/platforms/macos.py b/wildedge/platforms/macos.py index 1b6b1ee..7f1deb6 100644 --- a/wildedge/platforms/macos.py +++ b/wildedge/platforms/macos.py @@ -8,7 +8,15 @@ from wildedge.platforms.base import debug_detection_failure from wildedge.platforms.hardware import HardwareContext -_libc = ctypes.CDLL(None) +_libc: ctypes.CDLL | None = None + + +def _get_libc() -> ctypes.CDLL: + global _libc + if _libc is None: + _libc = ctypes.CDLL(None) + return _libc + _CF_STRING_ENCODING_UTF8 = 0x08000100 _CF_NUMBER_SINT32_TYPE = 3 @@ -18,7 +26,9 @@ def _sysctl_uint32(name: bytes) -> int | None: try: buf = ctypes.c_uint32(0) size = ctypes.c_size_t(ctypes.sizeof(buf)) - ret = _libc.sysctlbyname(name, ctypes.byref(buf), ctypes.byref(size), None, 0) + ret = _get_libc().sysctlbyname( + name, ctypes.byref(buf), ctypes.byref(size), None, 0 + ) return int(buf.value) if ret == 0 else None except Exception as exc: debug_detection_failure(f"macos sysctl_uint32({name!r})", exc) @@ -29,7 +39,9 @@ def _sysctl_uint64(name: bytes) -> int | None: try: buf = ctypes.c_uint64(0) size = ctypes.c_size_t(ctypes.sizeof(buf)) - ret = _libc.sysctlbyname(name, ctypes.byref(buf), ctypes.byref(size), None, 0) + ret = _get_libc().sysctlbyname( + name, ctypes.byref(buf), ctypes.byref(size), None, 0 + ) return int(buf.value) if ret == 0 else None except Exception as exc: debug_detection_failure(f"macos sysctl_uint64({name!r})", exc) @@ -52,7 +64,9 @@ def device_model(self) -> str | None: try: buf = ctypes.create_string_buffer(128) size = ctypes.c_size_t(ctypes.sizeof(buf)) - ret = _libc.sysctlbyname(b"hw.model", buf, ctypes.byref(size), None, 0) + ret = _get_libc().sysctlbyname( + b"hw.model", buf, ctypes.byref(size), None, 0 + ) if ret == 0: val = buf.value.decode().strip() return val or None From 5b3decbe3f73d83039e573d90a20b4b4eb3d35dc Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 10:17:49 +0100 Subject: [PATCH 05/10] fix: windows: unknown value handling --- wildedge/platforms/windows.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/wildedge/platforms/windows.py b/wildedge/platforms/windows.py index 4d2504f..558230f 100644 --- a/wildedge/platforms/windows.py +++ b/wildedge/platforms/windows.py @@ -35,10 +35,10 @@ class _MEMORYSTATUSEX(ctypes.Structure): class _SystemPowerStatus(ctypes.Structure): _fields_ = [ - ("ACLineStatus", ctypes.c_byte), - ("BatteryFlag", ctypes.c_byte), - ("BatteryLifePercent", ctypes.c_byte), - ("SystemStatusFlag", ctypes.c_byte), + ("ACLineStatus", ctypes.c_ubyte), + ("BatteryFlag", ctypes.c_ubyte), + ("BatteryLifePercent", ctypes.c_ubyte), + ("SystemStatusFlag", ctypes.c_ubyte), ("BatteryLifeTime", ctypes.c_ulong), ("BatteryFullLifeTime", ctypes.c_ulong), ] From bc12252c5a6e0f2ebb70022622679f947c99ea7f Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 10:33:51 +0100 Subject: [PATCH 06/10] Cleanup: avoid duplication --- tests/test_hardware_context.py | 3 +- wildedge/platforms/__init__.py | 14 ++++++++-- wildedge/platforms/base.py | 50 ++++++++++++++++++++++++++++------ wildedge/platforms/hardware.py | 18 +----------- wildedge/platforms/linux.py | 21 ++------------ wildedge/platforms/macos.py | 19 ++----------- wildedge/platforms/unknown.py | 11 ++------ wildedge/platforms/windows.py | 16 ++--------- 8 files changed, 65 insertions(+), 87 deletions(-) diff --git a/tests/test_hardware_context.py b/tests/test_hardware_context.py index 0a10f53..c2eee41 100644 --- a/tests/test_hardware_context.py +++ b/tests/test_hardware_context.py @@ -1,7 +1,8 @@ from __future__ import annotations from wildedge.events.inference import InferenceEvent -from wildedge.platforms.hardware import HardwareContext, ThermalContext, capture +from wildedge.platforms import capture +from wildedge.platforms.hardware import HardwareContext, ThermalContext from wildedge.platforms.linux import LinuxPlatform diff --git a/wildedge/platforms/__init__.py b/wildedge/platforms/__init__.py index 8a574c5..01b6d0e 100644 --- a/wildedge/platforms/__init__.py +++ b/wildedge/platforms/__init__.py @@ -2,21 +2,29 @@ import sys -from wildedge.platforms.base import PlatformAdapter +from wildedge.platforms.base import Platform +from wildedge.platforms.hardware import HardwareContext from wildedge.platforms.linux import LinuxPlatform from wildedge.platforms.macos import MacOSPlatform from wildedge.platforms.unknown import UnknownPlatform from wildedge.platforms.windows import WindowsPlatform -PLATFORMS: dict[str, PlatformAdapter] = { +PLATFORMS: dict[str, Platform] = { "linux": LinuxPlatform(), "darwin": MacOSPlatform(), "win32": WindowsPlatform(), } -def get_current_platform() -> PlatformAdapter: +def get_current_platform() -> Platform: return PLATFORMS.get(sys.platform, UnknownPlatform()) CURRENT_PLATFORM = get_current_platform() + + +def capture(accelerator_actual: str | None = None) -> HardwareContext: + ctx = CURRENT_PLATFORM.hardware_context() + if accelerator_actual is not None: + ctx.accelerator_actual = accelerator_actual + return ctx diff --git a/wildedge/platforms/base.py b/wildedge/platforms/base.py index edaf7cc..c530feb 100644 --- a/wildedge/platforms/base.py +++ b/wildedge/platforms/base.py @@ -1,35 +1,69 @@ from __future__ import annotations import ctypes +from abc import ABC, abstractmethod from pathlib import Path -from typing import TYPE_CHECKING, Protocol from wildedge.logging import logger +from wildedge.platforms.hardware import HardwareContext, ThermalContext -if TYPE_CHECKING: - from wildedge.platforms.hardware import HardwareContext - -class PlatformAdapter(Protocol): +class Platform(ABC): wire_type: str + @abstractmethod def config_base(self) -> Path: ... + + @abstractmethod def state_base(self) -> Path: ... + + @abstractmethod def cache_base(self) -> Path: ... + @abstractmethod def device_model(self) -> str | None: ... + @abstractmethod def os_version(self) -> str | None: ... - def ram_bytes(self) -> int | None: ... - + @abstractmethod def disk_bytes(self) -> int | None: ... + @abstractmethod def gpu_accelerators(self) -> tuple[list[str], str | None]: ... + @abstractmethod def gpu_accelerator_for_offload(self) -> str: ... - def hardware_context(self) -> HardwareContext: ... + # Optional hardware introspection — subclasses override what they support. + + def meminfo(self) -> tuple[int | None, int | None]: + return None, None + + def ram_bytes(self) -> int | None: + return self.meminfo()[0] + + def battery(self) -> tuple[float | None, bool | None]: + return None, None + + def cpu_freq(self) -> tuple[int | None, int | None]: + return None, None + + def thermal(self) -> ThermalContext | None: + return None + + def hardware_context(self) -> HardwareContext: + _, mem_available = self.meminfo() + bat_level, bat_charging = self.battery() + cpu_cur, cpu_max = self.cpu_freq() + return HardwareContext( + thermal=self.thermal(), + battery_level=bat_level, + battery_charging=bat_charging, + memory_available_bytes=mem_available, + cpu_freq_mhz=cpu_cur, + cpu_freq_max_mhz=cpu_max, + ) def debug_detection_failure(context: str, exc: BaseException) -> None: diff --git a/wildedge/platforms/hardware.py b/wildedge/platforms/hardware.py index 09ec36f..078407b 100644 --- a/wildedge/platforms/hardware.py +++ b/wildedge/platforms/hardware.py @@ -1,4 +1,4 @@ -"""Hardware context dataclasses and inference-time capture.""" +"""Hardware context dataclasses.""" from __future__ import annotations @@ -53,19 +53,3 @@ def to_dict(self) -> dict: if v is not None: d[k] = v return d - - -def capture(accelerator_actual: str | None = None) -> HardwareContext: - """Read current hardware state. Never raises. - - Pass ``accelerator_actual`` if the runtime reports which accelerator was - used; it may differ from the accelerator requested at load time. - """ - # Lazy import avoids a circular dependency: __init__.py imports platform - # modules which import this file. - from wildedge.platforms import CURRENT_PLATFORM # noqa: PLC0415 - - ctx = CURRENT_PLATFORM.hardware_context() - if accelerator_actual is not None: - ctx.accelerator_actual = accelerator_actual - return ctx diff --git a/wildedge/platforms/linux.py b/wildedge/platforms/linux.py index eb3d4f7..c189e4a 100644 --- a/wildedge/platforms/linux.py +++ b/wildedge/platforms/linux.py @@ -6,12 +6,13 @@ from pathlib import Path from wildedge.platforms.base import ( + Platform, cuda_device_count, debug_detection_failure, hip_device_count, nvml_gpu_name, ) -from wildedge.platforms.hardware import HardwareContext, ThermalContext +from wildedge.platforms.hardware import ThermalContext # Thermal zone type substrings that identify a CPU zone. CPU_THERMAL_ZONE_TYPES = ("cpu", "x86_pkg", "acpi", "soc", "pkg") @@ -26,7 +27,7 @@ TRIP_POINT_DEFAULT = ("nominal", "active") -class LinuxPlatform: +class LinuxPlatform(Platform): wire_type = "linux" def config_base(self) -> Path: @@ -82,9 +83,6 @@ def meminfo(self) -> tuple[int | None, int | None]: debug_detection_failure("linux meminfo", exc) return total, available - def ram_bytes(self) -> int | None: - return self.meminfo()[0] - def disk_bytes(self) -> int | None: try: return shutil.disk_usage("/").total @@ -106,19 +104,6 @@ def gpu_accelerator_for_offload(self) -> str: accs, _ = self.gpu_accelerators() return accs[0] if accs else "cpu" - def hardware_context(self) -> HardwareContext: - bat_level, bat_charging = self.battery() - cpu_cur, cpu_max = self.cpu_freq() - _, mem_available = self.meminfo() - return HardwareContext( - thermal=self.thermal(), - battery_level=bat_level, - battery_charging=bat_charging, - memory_available_bytes=mem_available, - cpu_freq_mhz=cpu_cur, - cpu_freq_max_mhz=cpu_max, - ) - def cpu_freq(self) -> tuple[int | None, int | None]: """Read current and max CPU frequency (MHz) from cpufreq for cpu0.""" base = Path("/sys/devices/system/cpu/cpu0/cpufreq") diff --git a/wildedge/platforms/macos.py b/wildedge/platforms/macos.py index 7f1deb6..bcbb71b 100644 --- a/wildedge/platforms/macos.py +++ b/wildedge/platforms/macos.py @@ -5,8 +5,7 @@ import shutil from pathlib import Path -from wildedge.platforms.base import debug_detection_failure -from wildedge.platforms.hardware import HardwareContext +from wildedge.platforms.base import Platform, debug_detection_failure _libc: ctypes.CDLL | None = None @@ -48,7 +47,7 @@ def _sysctl_uint64(name: bytes) -> int | None: return None -class MacOSPlatform: +class MacOSPlatform(Platform): wire_type = "macos" def config_base(self) -> Path: @@ -86,9 +85,6 @@ def meminfo(self) -> tuple[int | None, int | None]: ) return total, available - def ram_bytes(self) -> int | None: - return self.meminfo()[0] - def disk_bytes(self) -> int | None: try: return shutil.disk_usage("/").total @@ -204,14 +200,3 @@ def cfstr(s: bytes) -> ctypes.c_void_p: except Exception as exc: debug_detection_failure("macos battery", exc) return None, None - - def hardware_context(self) -> HardwareContext: - _, mem_available = self.meminfo() - bat_level, bat_charging = self.battery() - return HardwareContext( - memory_available_bytes=mem_available, - battery_level=bat_level, - battery_charging=bat_charging, - # cpu_freq not available on macOS (no public API on Apple Silicon) - # thermal not available via public Python API - ) diff --git a/wildedge/platforms/unknown.py b/wildedge/platforms/unknown.py index 43039fe..973dd4e 100644 --- a/wildedge/platforms/unknown.py +++ b/wildedge/platforms/unknown.py @@ -5,11 +5,10 @@ import sys from pathlib import Path -from wildedge.platforms.base import debug_detection_failure -from wildedge.platforms.hardware import HardwareContext +from wildedge.platforms.base import Platform, debug_detection_failure -class UnknownPlatform: +class UnknownPlatform(Platform): wire_type = sys.platform def config_base(self) -> Path: @@ -31,9 +30,6 @@ def os_version(self) -> str | None: debug_detection_failure("unknown os_version", exc) return None - def ram_bytes(self) -> int | None: - return None - def disk_bytes(self) -> int | None: try: return shutil.disk_usage("/").total @@ -46,6 +42,3 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: def gpu_accelerator_for_offload(self) -> str: return "cpu" - - def hardware_context(self) -> HardwareContext: - return HardwareContext() diff --git a/wildedge/platforms/windows.py b/wildedge/platforms/windows.py index 558230f..7748816 100644 --- a/wildedge/platforms/windows.py +++ b/wildedge/platforms/windows.py @@ -7,11 +7,11 @@ from pathlib import Path from wildedge.platforms.base import ( + Platform, cuda_device_count, debug_detection_failure, nvml_gpu_name, ) -from wildedge.platforms.hardware import HardwareContext try: import winreg as _winreg # type: ignore[import] # Windows only @@ -44,7 +44,7 @@ class _SystemPowerStatus(ctypes.Structure): ] -class WindowsPlatform: +class WindowsPlatform(Platform): wire_type = "windows" def config_base(self) -> Path: @@ -89,9 +89,6 @@ def meminfo(self) -> tuple[int | None, int | None]: debug_detection_failure("windows meminfo", exc) return None, None - def ram_bytes(self) -> int | None: - return self.meminfo()[0] - def disk_bytes(self) -> int | None: try: drive = os.environ.get("SystemDrive", "C:\\") @@ -129,12 +126,3 @@ def battery(self) -> tuple[float | None, bool | None]: except Exception as exc: debug_detection_failure("windows battery", exc) return None, None - - def hardware_context(self) -> HardwareContext: - _, mem_available = self.meminfo() - bat_level, bat_charging = self.battery() - return HardwareContext( - memory_available_bytes=mem_available, - battery_level=bat_level, - battery_charging=bat_charging, - ) From 94c8d8992bdb3b7532400146287dcb0ba29a0188 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 11:01:12 +0100 Subject: [PATCH 07/10] Cleanup: device info, hardware info code cleanup --- examples/chatgpt_example.py | 9 +- examples/gguf_gemma_manual_example.py | 3 +- tests/compat/conftest.py | 2 +- tests/conftest.py | 2 +- tests/test_batch.py | 2 +- tests/test_consumer.py | 2 +- tests/test_device.py | 107 +++++++++------- tests/test_hardware_context.py | 6 +- tests/test_integration_patching.py | 2 +- tests/test_offline_replay.py | 2 +- tests/test_transmitter.py | 2 +- wildedge/__init__.py | 4 +- wildedge/batch.py | 2 +- wildedge/cli.py | 2 +- wildedge/client.py | 3 +- wildedge/consumer.py | 2 +- wildedge/device.py | 175 -------------------------- wildedge/integrations/gguf.py | 2 +- wildedge/integrations/keras.py | 2 +- wildedge/integrations/tensorflow.py | 2 +- wildedge/platforms/__init__.py | 14 ++- wildedge/platforms/base.py | 82 ++++++++++++ wildedge/platforms/device_info.py | 43 +++++++ 23 files changed, 232 insertions(+), 240 deletions(-) delete mode 100644 wildedge/device.py create mode 100644 wildedge/platforms/device_info.py diff --git a/examples/chatgpt_example.py b/examples/chatgpt_example.py index edb6cb3..42b4db9 100644 --- a/examples/chatgpt_example.py +++ b/examples/chatgpt_example.py @@ -18,7 +18,13 @@ 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" @@ -72,6 +78,7 @@ inference_id = handle.track_inference( duration_ms=t.elapsed_ms, + hardware=capture_hardware(), input_modality="text", output_modality="text", success=True, diff --git a/examples/gguf_gemma_manual_example.py b/examples/gguf_gemma_manual_example.py index b1f0c0a..3bddef4 100644 --- a/examples/gguf_gemma_manual_example.py +++ b/examples/gguf_gemma_manual_example.py @@ -19,7 +19,7 @@ from llama_cpp import Llama import wildedge -from wildedge.events.inference import GenerationOutputMeta, TextInputMeta +from wildedge import GenerationOutputMeta, TextInputMeta, capture_hardware from wildedge.timing import Timer REPO = "bartowski/gemma-2-2b-it-GGUF" @@ -92,6 +92,7 @@ handle.track_inference( duration_ms=t.elapsed_ms, + hardware=capture_hardware(), input_modality="text", output_modality="text", success=True, diff --git a/tests/compat/conftest.py b/tests/compat/conftest.py index d806cf2..a42f63a 100644 --- a/tests/compat/conftest.py +++ b/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: diff --git a/tests/conftest.py b/tests/conftest.py index e887aec..46f7c77 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,8 +7,8 @@ import pytest from wildedge.client import WildEdge -from wildedge.device import DeviceInfo from wildedge.model import ModelInfo +from wildedge.platforms.device_info import DeviceInfo PLATFORM_MARKS = { "requires_linux": "linux", diff --git a/tests/test_batch.py b/tests/test_batch.py index efb5ff7..45865d5 100644 --- a/tests/test_batch.py +++ b/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: diff --git a/tests/test_consumer.py b/tests/test_consumer.py index ca0a771..a257a07 100644 --- a/tests/test_consumer.py +++ b/tests/test_consumer.py @@ -7,7 +7,7 @@ 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 diff --git a/tests/test_device.py b/tests/test_device.py index 7c5da5f..70fa84d 100644 --- a/tests/test_device.py +++ b/tests/test_device.py @@ -3,14 +3,11 @@ from pathlib import Path from unittest.mock import patch -from wildedge.device import ( - DeviceInfo, - LinuxPlatform, - MacOSPlatform, - detect_device, - get_device_id_path, - hmac_device_id, -) +from wildedge.platforms import detect_device, get_device_id_path +from wildedge.platforms.base import Platform, hmac_device_id +from wildedge.platforms.device_info import DeviceInfo +from wildedge.platforms.linux import LinuxPlatform +from wildedge.platforms.macos import MacOSPlatform class TestHmacDeviceId: @@ -31,48 +28,47 @@ def test_different_api_keys_produce_different_ids(self): class TestDetectDevice: - def test_returns_device_info(self): - with patch( - "wildedge.device.load_or_create_device_uuid", return_value="test-uuid" - ): - info = detect_device(api_key="test-key", app_version="1.0.0") + def test_returns_device_info(self, monkeypatch): + monkeypatch.setattr( + Platform, "load_or_create_device_uuid", lambda self: "test-uuid" + ) + info = detect_device(api_key="test-key", app_version="1.0.0") assert isinstance(info, DeviceInfo) assert info.app_version == "1.0.0" assert info.sdk_version == "wildedge-python-0.1.0" - def test_device_type_is_normalised(self): - with patch( - "wildedge.device.load_or_create_device_uuid", return_value="test-uuid" - ): - with patch("wildedge.device.CURRENT_PLATFORM", MacOSPlatform()): - info = detect_device(api_key="k", app_version="1.0") + def test_device_type_is_normalised(self, monkeypatch): + monkeypatch.setattr( + Platform, "load_or_create_device_uuid", lambda self: "test-uuid" + ) + with patch("wildedge.platforms.CURRENT_PLATFORM", MacOSPlatform()): + info = detect_device(api_key="k", app_version="1.0") assert info.device_type == "macos" - def test_linux_device_type(self): - with patch( - "wildedge.device.load_or_create_device_uuid", return_value="test-uuid" - ): - with patch("wildedge.device.CURRENT_PLATFORM", LinuxPlatform()): - info = detect_device(api_key="k", app_version="1.0") + def test_linux_device_type(self, monkeypatch): + monkeypatch.setattr( + Platform, "load_or_create_device_uuid", lambda self: "test-uuid" + ) + with patch("wildedge.platforms.CURRENT_PLATFORM", LinuxPlatform()): + info = detect_device(api_key="k", app_version="1.0") assert info.device_type == "linux" - def test_device_id_is_hmac_of_uuid(self): - with patch( - "wildedge.device.load_or_create_device_uuid", return_value="fixed-uuid" - ): - info = detect_device(api_key="fixed-key", app_version="1.0") - expected = hmac_device_id("fixed-key", "fixed-uuid") - assert info.device_id == expected - - def test_overrides_applied(self): - with patch( - "wildedge.device.load_or_create_device_uuid", return_value="test-uuid" - ): - info = detect_device( - api_key="k", - app_version="1.0", - overrides={"device_model": "Jetson AGX"}, - ) + def test_device_id_is_hmac_of_uuid(self, monkeypatch): + monkeypatch.setattr( + Platform, "load_or_create_device_uuid", lambda self: "fixed-uuid" + ) + info = detect_device(api_key="fixed-key", app_version="1.0") + assert info.device_id == hmac_device_id("fixed-key", "fixed-uuid") + + def test_overrides_applied(self, monkeypatch): + monkeypatch.setattr( + Platform, "load_or_create_device_uuid", lambda self: "test-uuid" + ) + info = detect_device( + api_key="k", + app_version="1.0", + overrides={"device_model": "Jetson AGX"}, + ) assert info.device_model == "Jetson AGX" def test_to_dict_returns_protocol_fields(self): @@ -91,10 +87,33 @@ def test_to_dict_returns_protocol_fields(self): assert "sdk_version" in d def test_get_device_id_path_uses_config_constants(self): - class _FakePlatform: + class FakePlatform(Platform): + wire_type = "test" + def config_base(self): return Path("/tmp/test-config") - with patch("wildedge.device.CURRENT_PLATFORM", _FakePlatform()): + def state_base(self): + return Path("/tmp") + + def cache_base(self): + return Path("/tmp") + + def device_model(self): + return None + + def os_version(self): + return None + + def disk_bytes(self): + return None + + def gpu_accelerators(self): + return [], None + + def gpu_accelerator_for_offload(self): + return "cpu" + + with patch("wildedge.platforms.CURRENT_PLATFORM", FakePlatform()): path = get_device_id_path() assert path == Path("/tmp/test-config/wildedge/device_id") diff --git a/tests/test_hardware_context.py b/tests/test_hardware_context.py index c2eee41..d35f0a8 100644 --- a/tests/test_hardware_context.py +++ b/tests/test_hardware_context.py @@ -1,7 +1,7 @@ from __future__ import annotations from wildedge.events.inference import InferenceEvent -from wildedge.platforms import capture +from wildedge.platforms import capture_hardware from wildedge.platforms.hardware import HardwareContext, ThermalContext from wildedge.platforms.linux import LinuxPlatform @@ -70,7 +70,7 @@ def test_capture_sets_accelerator_actual(monkeypatch): monkeypatch.setattr( LinuxPlatform, "hardware_context", lambda self: HardwareContext() ) - assert capture(accelerator_actual="gpu").accelerator_actual == "gpu" + assert capture_hardware(accelerator_actual="gpu").accelerator_actual == "gpu" def test_capture_preserves_existing_accelerator_when_not_passed(monkeypatch): @@ -80,4 +80,4 @@ def test_capture_preserves_existing_accelerator_when_not_passed(monkeypatch): "hardware_context", lambda self: HardwareContext(accelerator_actual="mps"), ) - assert capture().accelerator_actual == "mps" + assert capture_hardware().accelerator_actual == "mps" diff --git a/tests/test_integration_patching.py b/tests/test_integration_patching.py index 02a5561..cf75de5 100644 --- a/tests/test_integration_patching.py +++ b/tests/test_integration_patching.py @@ -6,11 +6,11 @@ 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.onnx import OnnxExtractor from wildedge.integrations.pytorch import PytorchExtractor +from wildedge.platforms.device_info import DeviceInfo def test_hf_install_patch_is_idempotent(monkeypatch): diff --git a/tests/test_offline_replay.py b/tests/test_offline_replay.py index c3970f7..869bf62 100644 --- a/tests/test_offline_replay.py +++ b/tests/test_offline_replay.py @@ -3,7 +3,7 @@ from unittest.mock import patch from wildedge.client import WildEdge -from wildedge.device import DeviceInfo +from wildedge.platforms.device_info import DeviceInfo class _DummyConsumer: diff --git a/tests/test_transmitter.py b/tests/test_transmitter.py index b4729d9..68fb9fb 100644 --- a/tests/test_transmitter.py +++ b/tests/test_transmitter.py @@ -202,7 +202,7 @@ def test_unexpected_status_keeps_events_in_consumer_queue(): # End-to-end regression: TransmitError from an unexpected status must cause the # consumer to retain events for retry, not lose them. from wildedge.consumer import Consumer - from wildedge.device import DeviceInfo + from wildedge.platforms.device_info import DeviceInfo from wildedge.queue import EventQueue queue = EventQueue(max_size=100) diff --git a/wildedge/__init__.py b/wildedge/__init__.py index c153200..73f74ef 100644 --- a/wildedge/__init__.py +++ b/wildedge/__init__.py @@ -2,7 +2,6 @@ from wildedge.client import WildEdge from wildedge.decorators import track -from wildedge.device import DeviceInfo from wildedge.events import ( AdapterDownload, AdapterLoad, @@ -17,12 +16,15 @@ ImageInputMeta, TextInputMeta, ) +from wildedge.platforms import capture_hardware +from wildedge.platforms.device_info import DeviceInfo from wildedge.platforms.hardware import HardwareContext, ThermalContext from wildedge.queue import QueuePolicy from wildedge.timing import Timer __all__ = [ "WildEdge", + "capture_hardware", "HardwareContext", "ThermalContext", "track", diff --git a/wildedge/batch.py b/wildedge/batch.py index 9afa3df..239a4a5 100644 --- a/wildedge/batch.py +++ b/wildedge/batch.py @@ -4,7 +4,7 @@ from datetime import datetime, timezone from wildedge import constants -from wildedge.device import DeviceInfo +from wildedge.platforms.device_info import DeviceInfo def _sanitize_event(event: dict) -> dict: diff --git a/wildedge/cli.py b/wildedge/cli.py index 06a7470..f3992d2 100644 --- a/wildedge/cli.py +++ b/wildedge/cli.py @@ -16,10 +16,10 @@ 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.platforms import get_device_id_path from wildedge.settings import read_client_env, resolve_app_identity diff --git a/wildedge/client.py b/wildedge/client.py index 92de53e..a2a485d 100644 --- a/wildedge/client.py +++ b/wildedge/client.py @@ -12,7 +12,6 @@ from wildedge import constants from wildedge.consumer import Consumer from wildedge.dead_letters import DeadLetterStore -from wildedge.device import DeviceInfo, detect_device from wildedge.hubs.base import BaseHubTracker from wildedge.hubs.huggingface import HuggingFaceHubTracker from wildedge.hubs.registry import supported_hubs @@ -34,6 +33,8 @@ default_model_registry_path, default_pending_queue_dir, ) +from wildedge.platforms import detect_device +from wildedge.platforms.device_info import DeviceInfo from wildedge.queue import EventQueue, QueuePolicy from wildedge.settings import read_client_env, resolve_app_identity from wildedge.timing import Timer, elapsed_ms diff --git a/wildedge/consumer.py b/wildedge/consumer.py index 068a793..ef8c15d 100644 --- a/wildedge/consumer.py +++ b/wildedge/consumer.py @@ -16,7 +16,7 @@ from wildedge.transmitter import TransmitError, Transmitter if TYPE_CHECKING: - from wildedge.device import DeviceInfo + from wildedge.platforms.device_info import DeviceInfo class Consumer: diff --git a/wildedge/device.py b/wildedge/device.py deleted file mode 100644 index a5e2519..0000000 --- a/wildedge/device.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Device info auto-detection (stdlib-only).""" - -from __future__ import annotations - -import hashlib -import hmac -import locale -import os -import platform -import uuid -from dataclasses import dataclass, field -from datetime import datetime -from pathlib import Path - -from wildedge import constants -from wildedge.platforms import CURRENT_PLATFORM -from wildedge.platforms import PLATFORMS as _PLATFORMS -from wildedge.platforms.base import debug_detection_failure -from wildedge.platforms.linux import LinuxPlatform -from wildedge.platforms.macos import MacOSPlatform -from wildedge.platforms.unknown import UnknownPlatform -from wildedge.platforms.windows import WindowsPlatform - -PLATFORMS = _PLATFORMS - -__all__ = [ - "CURRENT_PLATFORM", - "PLATFORMS", - "LinuxPlatform", - "MacOSPlatform", - "WindowsPlatform", - "UnknownPlatform", - "DeviceInfo", - "get_device_id_path", - "load_or_create_device_uuid", - "hmac_device_id", - "detect_gpu_info", - "detect_locale", - "detect_timezone", - "detect_device", -] - - -def get_device_id_path() -> Path: - """ - Returns the path to the device ID file for the current platform. - """ - return ( - CURRENT_PLATFORM.config_base() - / constants.DEVICE_ID_DIR - / constants.DEVICE_ID_FILE - ) - - -def load_or_create_device_uuid() -> str: - """ - Loads a persistent UUID for this device from disk, or creates one if not found. - """ - path = get_device_id_path() - try: - if path.exists(): - stored = path.read_text().strip() - if stored: - return stored - except OSError as exc: - debug_detection_failure("device_uuid read", exc) - - new_id = str(uuid.uuid4()) - try: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(new_id) - except OSError as exc: - debug_detection_failure("device_uuid write", exc) - return new_id - - -def hmac_device_id(api_key: str, raw_id: str) -> str: - return hmac.new( - key=api_key.encode(), - msg=raw_id.encode(), - digestmod=hashlib.sha256, - ).hexdigest() - - -def detect_gpu_info() -> tuple[list[str], str | None]: - """Returns (accelerators, gpu_name). Always includes 'cpu'. Stdlib + ctypes only.""" - gpu_accs, gpu_name = CURRENT_PLATFORM.gpu_accelerators() - return ["cpu", *gpu_accs], gpu_name - - -def detect_locale() -> str | None: - try: - loc = locale.getlocale() - return loc[0] if loc else None - except Exception as exc: - debug_detection_failure("locale", exc) - return None - - -def detect_timezone() -> str | None: - try: - return datetime.now().astimezone().tzname() - except Exception as exc: - debug_detection_failure("timezone", exc) - return None - - -@dataclass -class DeviceInfo: - device_id: str - device_type: str - sdk_version: str = constants.SDK_VERSION - device_model: str | None = None - os_version: str | None = None - locale: str | None = None - timezone: str | None = None - cpu_arch: str | None = None - cpu_cores: int | None = None - ram_total_bytes: int | None = None - disk_total_bytes: int | None = None - accelerators: list[str] = field(default_factory=list) - gpu_name: str | None = None - app_version: str | None = None - - def to_dict(self) -> dict: - d = { - "device_id": self.device_id, - "device_type": self.device_type, - "device_model": self.device_model, - "os_version": self.os_version, - "sdk_version": self.sdk_version, - "locale": self.locale, - "timezone": self.timezone, - "cpu_arch": self.cpu_arch, - "cpu_cores": self.cpu_cores, - "ram_total_bytes": self.ram_total_bytes, - "disk_total_bytes": self.disk_total_bytes, - "accelerators": self.accelerators, - "gpu_name": self.gpu_name, - } - if self.app_version is not None: - d["app_version"] = self.app_version - return d - - -def detect_device( - api_key: str, app_version: str | None, overrides: dict | None = None -) -> DeviceInfo: - """Auto-detect device info and HMAC the stored UUID with api_key.""" - raw_uuid = load_or_create_device_uuid() - device_id = hmac_device_id(api_key, raw_uuid) - accelerators, gpu_name = detect_gpu_info() - - info = DeviceInfo( - app_version=app_version, - device_id=device_id, - device_type=CURRENT_PLATFORM.wire_type, - device_model=CURRENT_PLATFORM.device_model(), - os_version=CURRENT_PLATFORM.os_version(), - locale=detect_locale(), - timezone=detect_timezone(), - cpu_arch=platform.machine() or None, - cpu_cores=os.cpu_count(), - ram_total_bytes=CURRENT_PLATFORM.ram_bytes(), - disk_total_bytes=CURRENT_PLATFORM.disk_bytes(), - accelerators=accelerators, - gpu_name=gpu_name, - ) - - if overrides: - for k, v in overrides.items(): - if hasattr(info, k): - setattr(info, k, v) - - return info diff --git a/wildedge/integrations/gguf.py b/wildedge/integrations/gguf.py index ff5ebe2..9937cdd 100644 --- a/wildedge/integrations/gguf.py +++ b/wildedge/integrations/gguf.py @@ -10,12 +10,12 @@ from typing import TYPE_CHECKING from wildedge import constants -from wildedge.device import CURRENT_PLATFORM from wildedge.events.inference import GenerationOutputMeta, TextInputMeta from wildedge.integrations.base import BaseExtractor, patch_instance_call_once from wildedge.integrations.common import debug_failure from wildedge.logging import logger from wildedge.model import ModelInfo +from wildedge.platforms import CURRENT_PLATFORM from wildedge.timing import elapsed_ms try: diff --git a/wildedge/integrations/keras.py b/wildedge/integrations/keras.py index 2a2dbf4..51b55d5 100644 --- a/wildedge/integrations/keras.py +++ b/wildedge/integrations/keras.py @@ -12,7 +12,6 @@ _np = None # type: ignore[assignment] from wildedge import constants -from wildedge.device import CURRENT_PLATFORM from wildedge.events.inference import ClassificationOutputMeta, TopKPrediction from wildedge.integrations.base import BaseExtractor, patch_instance_call_once from wildedge.integrations.common import ( @@ -21,6 +20,7 @@ num_classes_from_output_shape, ) from wildedge.model import ModelInfo +from wildedge.platforms import CURRENT_PLATFORM from wildedge.timing import elapsed_ms if TYPE_CHECKING: diff --git a/wildedge/integrations/tensorflow.py b/wildedge/integrations/tensorflow.py index e4f4681..03d2416 100644 --- a/wildedge/integrations/tensorflow.py +++ b/wildedge/integrations/tensorflow.py @@ -13,7 +13,6 @@ _np = None # type: ignore[assignment] from wildedge import constants -from wildedge.device import CURRENT_PLATFORM from wildedge.events.inference import ClassificationOutputMeta, TopKPrediction from wildedge.integrations.base import BaseExtractor, patch_instance_call_once from wildedge.integrations.common import ( @@ -22,6 +21,7 @@ num_classes_from_output_shape, ) from wildedge.model import ModelInfo +from wildedge.platforms import CURRENT_PLATFORM from wildedge.timing import elapsed_ms if TYPE_CHECKING: diff --git a/wildedge/platforms/__init__.py b/wildedge/platforms/__init__.py index 01b6d0e..e43ab41 100644 --- a/wildedge/platforms/__init__.py +++ b/wildedge/platforms/__init__.py @@ -1,8 +1,10 @@ from __future__ import annotations import sys +from pathlib import Path from wildedge.platforms.base import Platform +from wildedge.platforms.device_info import DeviceInfo from wildedge.platforms.hardware import HardwareContext from wildedge.platforms.linux import LinuxPlatform from wildedge.platforms.macos import MacOSPlatform @@ -23,7 +25,17 @@ def get_current_platform() -> Platform: CURRENT_PLATFORM = get_current_platform() -def capture(accelerator_actual: str | None = None) -> HardwareContext: +def get_device_id_path() -> Path: + return CURRENT_PLATFORM.get_device_id_path() + + +def detect_device( + api_key: str, app_version: str | None, overrides: dict | None = None +) -> DeviceInfo: + return CURRENT_PLATFORM.detect_device(api_key, app_version, overrides) + + +def capture_hardware(accelerator_actual: str | None = None) -> HardwareContext: ctx = CURRENT_PLATFORM.hardware_context() if accelerator_actual is not None: ctx.accelerator_actual = accelerator_actual diff --git a/wildedge/platforms/base.py b/wildedge/platforms/base.py index c530feb..62df959 100644 --- a/wildedge/platforms/base.py +++ b/wildedge/platforms/base.py @@ -1,10 +1,19 @@ from __future__ import annotations import ctypes +import hashlib +import hmac +import locale +import os +import platform +import uuid from abc import ABC, abstractmethod +from datetime import datetime from pathlib import Path +from wildedge import constants from wildedge.logging import logger +from wildedge.platforms.device_info import DeviceInfo from wildedge.platforms.hardware import HardwareContext, ThermalContext @@ -65,6 +74,79 @@ def hardware_context(self) -> HardwareContext: cpu_freq_max_mhz=cpu_max, ) + def get_device_id_path(self) -> Path: + return self.config_base() / constants.DEVICE_ID_DIR / constants.DEVICE_ID_FILE + + def load_or_create_device_uuid(self) -> str: + path = self.get_device_id_path() + try: + if path.exists(): + stored = path.read_text().strip() + if stored: + return stored + except OSError as exc: + debug_detection_failure("device_uuid read", exc) + + new_id = str(uuid.uuid4()) + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(new_id) + except OSError as exc: + debug_detection_failure("device_uuid write", exc) + return new_id + + def detect_device( + self, api_key: str, app_version: str | None, overrides: dict | None = None + ) -> DeviceInfo: + raw_uuid = self.load_or_create_device_uuid() + device_id = hmac_device_id(api_key, raw_uuid) + gpu_accs, gpu_name = self.gpu_accelerators() + info = DeviceInfo( + app_version=app_version, + device_id=device_id, + device_type=self.wire_type, + device_model=self.device_model(), + os_version=self.os_version(), + locale=detect_locale(), + timezone=detect_timezone(), + cpu_arch=platform.machine() or None, + cpu_cores=os.cpu_count(), + ram_total_bytes=self.ram_bytes(), + disk_total_bytes=self.disk_bytes(), + accelerators=["cpu", *gpu_accs], + gpu_name=gpu_name, + ) + if overrides: + for k, v in overrides.items(): + if hasattr(info, k): + setattr(info, k, v) + return info + + +def hmac_device_id(api_key: str, raw_id: str) -> str: + return hmac.new( + key=api_key.encode(), + msg=raw_id.encode(), + digestmod=hashlib.sha256, + ).hexdigest() + + +def detect_locale() -> str | None: + try: + loc = locale.getlocale() + return loc[0] if loc else None + except Exception as exc: + debug_detection_failure("locale", exc) + return None + + +def detect_timezone() -> str | None: + try: + return datetime.now().astimezone().tzname() + except Exception as exc: + debug_detection_failure("timezone", exc) + return None + def debug_detection_failure(context: str, exc: BaseException) -> None: logger.debug("wildedge: device detection failed for %s: %s", context, exc) diff --git a/wildedge/platforms/device_info.py b/wildedge/platforms/device_info.py new file mode 100644 index 0000000..14b0b0f --- /dev/null +++ b/wildedge/platforms/device_info.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +from wildedge import constants + + +@dataclass +class DeviceInfo: + device_id: str + device_type: str + sdk_version: str = constants.SDK_VERSION + device_model: str | None = None + os_version: str | None = None + locale: str | None = None + timezone: str | None = None + cpu_arch: str | None = None + cpu_cores: int | None = None + ram_total_bytes: int | None = None + disk_total_bytes: int | None = None + accelerators: list[str] = field(default_factory=list) + gpu_name: str | None = None + app_version: str | None = None + + def to_dict(self) -> dict: + d = { + "device_id": self.device_id, + "device_type": self.device_type, + "device_model": self.device_model, + "os_version": self.os_version, + "sdk_version": self.sdk_version, + "locale": self.locale, + "timezone": self.timezone, + "cpu_arch": self.cpu_arch, + "cpu_cores": self.cpu_cores, + "ram_total_bytes": self.ram_total_bytes, + "disk_total_bytes": self.disk_total_bytes, + "accelerators": self.accelerators, + "gpu_name": self.gpu_name, + } + if self.app_version is not None: + d["app_version"] = self.app_version + return d From e3fab5f4671b3774d54e644e7d03fd73be2b20d2 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 11:25:29 +0100 Subject: [PATCH 08/10] background hardware sampler --- README.md | 1 + examples/chatgpt_example.py | 7 +- ...anual_example.py => gguf_gemma_example.py} | 11 +- tests/conftest.py | 14 +- tests/test_sampler.py | 136 ++++++++++++++++++ wildedge/client.py | 7 +- wildedge/model.py | 3 + wildedge/platforms/__init__.py | 26 +++- wildedge/platforms/sampler.py | 41 ++++++ 9 files changed, 235 insertions(+), 11 deletions(-) rename examples/{gguf_gemma_manual_example.py => gguf_gemma_example.py} (88%) create mode 100644 tests/test_sampler.py create mode 100644 wildedge/platforms/sampler.py diff --git a/README.md b/README.md index 524df87..8824b24 100644 --- a/README.md +++ b/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` | - | Seconds between background hardware snapshots; `None` to disable | ## Privacy diff --git a/examples/chatgpt_example.py b/examples/chatgpt_example.py index 42b4db9..ed4a81e 100644 --- a/examples/chatgpt_example.py +++ b/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. @@ -32,6 +34,7 @@ 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 diff --git a/examples/gguf_gemma_manual_example.py b/examples/gguf_gemma_example.py similarity index 88% rename from examples/gguf_gemma_manual_example.py rename to examples/gguf_gemma_example.py index 3bddef4..f3962cf 100644 --- a/examples/gguf_gemma_manual_example.py +++ b/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 import GenerationOutputMeta, TextInputMeta, capture_hardware +from wildedge import GenerationOutputMeta, TextInputMeta from wildedge.timing import Timer REPO = "bartowski/gemma-2-2b-it-GGUF" @@ -92,7 +94,6 @@ handle.track_inference( duration_ms=t.elapsed_ms, - hardware=capture_hardware(), input_modality="text", output_modality="text", success=True, diff --git a/tests/conftest.py b/tests/conftest.py index 46f7c77..e8c8199 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,6 +10,15 @@ 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", @@ -88,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 diff --git a/tests/test_sampler.py b/tests/test_sampler.py new file mode 100644 index 0000000..df71696 --- /dev/null +++ b/tests/test_sampler.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +from wildedge.platforms import ( + capture_hardware, + is_sampling, + start_sampler, + stop_sampler, +) +from wildedge.platforms.hardware import HardwareContext +from wildedge.platforms.sampler import HardwareSampler + + +def make_platform(snapshot: HardwareContext | None = None) -> MagicMock: + platform = MagicMock() + platform.hardware_context.return_value = snapshot or HardwareContext( + memory_available_bytes=1_000_000 + ) + return platform + + +class TestHardwareSampler: + def test_snapshot_is_warm_after_start(self): + platform = make_platform() + sampler = HardwareSampler(platform=platform, interval_s=60) + sampler.start() + assert sampler.snapshot().memory_available_bytes == 1_000_000 + sampler.stop() + + def test_start_calls_hardware_context_once(self): + platform = make_platform() + sampler = HardwareSampler(platform=platform, interval_s=60) + sampler.start() + platform.hardware_context.assert_called_once() + sampler.stop() + + def test_stop_signals_thread(self): + platform = make_platform() + sampler = HardwareSampler(platform=platform, interval_s=60) + sampler.start() + sampler.stop() + assert sampler.done.is_set() + + def test_snapshot_before_start_returns_empty_context(self): + platform = make_platform() + sampler = HardwareSampler(platform=platform, interval_s=60) + assert sampler.snapshot() == HardwareContext() + + +class TestSamplerModule: + def test_is_sampling_false_by_default(self): + assert not is_sampling() + + def test_is_sampling_true_after_start(self): + platform = make_platform() + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + assert is_sampling() + + def test_is_sampling_false_after_stop(self): + platform = make_platform() + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + stop_sampler() + assert not is_sampling() + + def test_stop_sampler_is_idempotent(self): + stop_sampler() + stop_sampler() # should not raise + + def test_capture_hardware_uses_snapshot_when_sampling(self): + platform = make_platform(HardwareContext(memory_available_bytes=42)) + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + ctx = capture_hardware() + assert ctx.memory_available_bytes == 42 + + def test_capture_hardware_live_read_when_not_sampling(self, monkeypatch): + platform = make_platform(HardwareContext(memory_available_bytes=99)) + monkeypatch.setattr("wildedge.platforms.CURRENT_PLATFORM", platform) + ctx = capture_hardware() + assert ctx.memory_available_bytes == 99 + platform.hardware_context.assert_called_once() + + def test_capture_hardware_does_not_mutate_snapshot(self): + platform = make_platform(HardwareContext(memory_available_bytes=1)) + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + from wildedge.platforms import _sampler + + original = _sampler.snapshot() + capture_hardware(accelerator_actual="cuda") + assert original.accelerator_actual is None + + def test_capture_hardware_accelerator_override(self): + platform = make_platform() + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + ctx = capture_hardware(accelerator_actual="mps") + assert ctx.accelerator_actual == "mps" + + +class TestTrackInferenceHardware: + def test_auto_attaches_when_sampling(self, monkeypatch): + from wildedge.model import ModelHandle + + platform = make_platform(HardwareContext(memory_available_bytes=512)) + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + + events = [] + handle = ModelHandle("m1", MagicMock(), events.append) + handle.track_inference(duration_ms=10) + assert events[0]["inference"]["hardware"]["memory_available_bytes"] == 512 + + def test_no_hardware_when_not_sampling(self): + from wildedge.model import ModelHandle + + events = [] + handle = ModelHandle("m1", MagicMock(), events.append) + handle.track_inference(duration_ms=10) + assert "hardware" not in events[0]["inference"] + + def test_explicit_hardware_takes_precedence_over_sampler(self): + from wildedge.model import ModelHandle + + platform = make_platform(HardwareContext(memory_available_bytes=1)) + with patch("wildedge.platforms.CURRENT_PLATFORM", platform): + start_sampler(interval_s=60) + + explicit = HardwareContext(memory_available_bytes=999) + events = [] + handle = ModelHandle("m1", MagicMock(), events.append) + handle.track_inference(duration_ms=10, hardware=explicit) + assert events[0]["inference"]["hardware"]["memory_available_bytes"] == 999 diff --git a/wildedge/client.py b/wildedge/client.py index a2a485d..eeda7fc 100644 --- a/wildedge/client.py +++ b/wildedge/client.py @@ -33,7 +33,7 @@ default_model_registry_path, default_pending_queue_dir, ) -from wildedge.platforms import detect_device +from wildedge.platforms import detect_device, start_sampler, stop_sampler from wildedge.platforms.device_info import DeviceInfo from wildedge.queue import EventQueue, QueuePolicy from wildedge.settings import read_client_env, resolve_app_identity @@ -146,6 +146,7 @@ def __init__( dead_letter_dir: str | None = None, max_dead_letter_batches: int = constants.DEFAULT_MAX_DEAD_LETTER_BATCHES, on_delivery_failure: Callable[[str, int, int], None] | None = None, + sampling_interval_s: float | None = 30.0, ): env = read_client_env(dsn=dsn, debug=debug, app_identity=app_identity) dsn = env.dsn @@ -232,6 +233,9 @@ def __init__( on_delivery_failure=on_delivery_failure, ) + if sampling_interval_s: + start_sampler(interval_s=sampling_interval_s) + self.auto_loaded: set[str] = set() # Active hub trackers keyed by hub name. Populated by _activate_hub() # when instrument() is called with a hub name. @@ -622,6 +626,7 @@ def flush(self, timeout: float = 5.0) -> None: def close(self, timeout: float | None = None) -> None: """Best-effort shutdown; pass timeout to attempt bounded flush first.""" self.closed = True + stop_sampler() if timeout is None: self.consumer.close() else: diff --git a/wildedge/model.py b/wildedge/model.py index ec9de62..e5b355c 100644 --- a/wildedge/model.py +++ b/wildedge/model.py @@ -25,6 +25,7 @@ TextInputMeta, ) from wildedge.logging import logger +from wildedge.platforms import capture_hardware, is_sampling from wildedge.platforms.hardware import HardwareContext @@ -144,6 +145,8 @@ def track_inference( generation_config: GenerationConfig | None = None, hardware: HardwareContext | None = None, ) -> str: + if hardware is None and is_sampling(): + hardware = capture_hardware() event = InferenceEvent( model_id=self.model_id, duration_ms=duration_ms, diff --git a/wildedge/platforms/__init__.py b/wildedge/platforms/__init__.py index e43ab41..5d6ff3d 100644 --- a/wildedge/platforms/__init__.py +++ b/wildedge/platforms/__init__.py @@ -1,5 +1,6 @@ from __future__ import annotations +import dataclasses import sys from pathlib import Path @@ -8,9 +9,30 @@ from wildedge.platforms.hardware import HardwareContext from wildedge.platforms.linux import LinuxPlatform from wildedge.platforms.macos import MacOSPlatform +from wildedge.platforms.sampler import HardwareSampler from wildedge.platforms.unknown import UnknownPlatform from wildedge.platforms.windows import WindowsPlatform +_sampler: HardwareSampler | None = None + + +def start_sampler(interval_s: float) -> None: + global _sampler + _sampler = HardwareSampler(platform=CURRENT_PLATFORM, interval_s=interval_s) + _sampler.start() + + +def stop_sampler() -> None: + global _sampler + if _sampler is not None: + _sampler.stop() + _sampler = None + + +def is_sampling() -> bool: + return _sampler is not None + + PLATFORMS: dict[str, Platform] = { "linux": LinuxPlatform(), "darwin": MacOSPlatform(), @@ -36,7 +58,7 @@ def detect_device( def capture_hardware(accelerator_actual: str | None = None) -> HardwareContext: - ctx = CURRENT_PLATFORM.hardware_context() + ctx = _sampler.snapshot() if _sampler else CURRENT_PLATFORM.hardware_context() if accelerator_actual is not None: - ctx.accelerator_actual = accelerator_actual + return dataclasses.replace(ctx, accelerator_actual=accelerator_actual) return ctx diff --git a/wildedge/platforms/sampler.py b/wildedge/platforms/sampler.py new file mode 100644 index 0000000..167662c --- /dev/null +++ b/wildedge/platforms/sampler.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +import threading +from typing import TYPE_CHECKING + +from wildedge.platforms.hardware import HardwareContext + +if TYPE_CHECKING: + from wildedge.platforms.base import Platform + + +class HardwareSampler: + """Samples hardware context on a background thread at a fixed interval. + + The platform is injected so this module has no dependency on platforms/__init__.py. + The first snapshot is taken synchronously in start() so capture_hardware() never + returns an empty context immediately after initialisation. + """ + + def __init__(self, platform: Platform, interval_s: float = 30.0): + self.platform = platform + self.interval_s = interval_s + self.current: HardwareContext = HardwareContext() + self.done = threading.Event() + self.thread = threading.Thread( + target=self._run, daemon=True, name="wildedge-hw-sampler" + ) + + def start(self) -> None: + self.current = self.platform.hardware_context() + self.thread.start() + + def stop(self) -> None: + self.done.set() + + def snapshot(self) -> HardwareContext: + return self.current + + def _run(self) -> None: + while not self.done.wait(self.interval_s): + self.current = self.platform.hardware_context() From c30745627c6fa7c5d33be8e4f37431ac28b0b96b Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 12:23:23 +0100 Subject: [PATCH 09/10] Testing Testing Testing --- README.md | 2 +- examples/transformers_example.py | 23 ++-- tests/test_cli.py | 8 +- tests/test_integration_patching.py | 80 +++++++++++++ wildedge/cli.py | 15 +++ wildedge/client.py | 33 +++++- wildedge/constants.py | 2 + wildedge/integrations/transformers.py | 161 ++++++++++++++------------ wildedge/runtime/bootstrap.py | 8 +- wildedge/settings.py | 9 ++ 10 files changed, 250 insertions(+), 91 deletions(-) diff --git a/README.md b/README.md index 8824b24..d1889d7 100644 --- a/README.md +++ b/README.md @@ -99,7 +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` | - | Seconds between background hardware snapshots; `None` to disable | +| `sampling_interval_s` | `30.0` | `WILDEDGE_SAMPLING_INTERVAL_S` | Seconds between background hardware snapshots; `0` or `None` to disable | ## Privacy diff --git a/examples/transformers_example.py b/examples/transformers_example.py index f8bdc68..c1d0739 100644 --- a/examples/transformers_example.py +++ b/examples/transformers_example.py @@ -21,19 +21,23 @@ from __future__ import annotations import argparse +import platform -from transformers import pipeline +import transformers import wildedge +_DEVICE = "mps" if platform.machine() == "arm64" else "cpu" + def run_classify() -> None: - pipe = pipeline( + pipe = transformers.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,9 @@ def run_classify() -> None: def run_generate() -> None: - pipe = pipeline("text-generation", model="gpt2", max_new_tokens=40) + pipe = transformers.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 +66,9 @@ def run_generate() -> None: def run_embed() -> None: - pipe = pipeline("feature-extraction", model="bert-base-uncased") + pipe = transformers.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 +96,9 @@ 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. + # instrument() patches transformers.pipeline and AutoModel.from_pretrained. + # Import transformers as a module (not `from transformers import pipeline`) + # so attribute lookups happen after patching and the correct device is captured. client = wildedge.WildEdge(app_version="1.0.0") # set WILDEDGE_DSN env var client.instrument("transformers", hubs=["huggingface"]) diff --git a/tests/test_cli.py b/tests/test_cli.py index b294547..4fbf3c8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -153,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] @@ -183,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 @@ -219,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] @@ -369,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] diff --git a/tests/test_integration_patching.py b/tests/test_integration_patching.py index cf75de5..7390d4c 100644 --- a/tests/test_integration_patching.py +++ b/tests/test_integration_patching.py @@ -10,6 +10,7 @@ from wildedge.integrations.gguf import GgufExtractor from wildedge.integrations.onnx import OnnxExtractor from wildedge.integrations.pytorch import PytorchExtractor +from wildedge.integrations.transformers import TransformersExtractor from wildedge.platforms.device_info import DeviceInfo @@ -105,6 +106,85 @@ def client_ref(): assert first is second +def test_transformers_install_patch_is_idempotent(monkeypatch): + import wildedge.integrations.transformers as tf_mod + + class FakePipeline: + def __init__(self, *args, **kwargs): + pass + + def __call__(self, inputs, *args, **kwargs): + return [] + + class FakePreTrainedModel: + @classmethod + def from_pretrained(cls, *args, **kwargs): + return cls() + + fake_transformers = types.SimpleNamespace( + Pipeline=FakePipeline, + PreTrainedModel=FakePreTrainedModel, + ) + monkeypatch.setattr(tf_mod, "_transformers", fake_transformers) + monkeypatch.setattr(tf_mod, "_transformers_patched", False) + + def client_ref(): + return None + + TransformersExtractor.install_auto_load_patch(client_ref) + first_init = fake_transformers.Pipeline.__init__ + first_call = fake_transformers.Pipeline.__call__ + TransformersExtractor.install_auto_load_patch(client_ref) + assert fake_transformers.Pipeline.__init__ is first_init + assert fake_transformers.Pipeline.__call__ is first_call + + +def test_transformers_from_import_style_intercepted(monkeypatch): + """from transformers import pipeline; pipeline() must be intercepted.""" + import wildedge.integrations.transformers as tf_mod + + loaded = [] + + class FakePipeline: + def __init__(self, *args, **kwargs): + loaded.append(self) + + def __call__(self, inputs, *args, **kwargs): + return [] + + class FakePreTrainedModel: + @classmethod + def from_pretrained(cls, *args, **kwargs): + return cls() + + fake_transformers = types.SimpleNamespace( + Pipeline=FakePipeline, + PreTrainedModel=FakePreTrainedModel, + ) + + client = types.SimpleNamespace( + closed=False, + _snapshot_hub_caches=lambda: {}, + _diff_hub_caches=lambda before, ms: None, + _on_model_auto_loaded=lambda obj, **kw: loaded.append(("tracked", obj)), + ) + + monkeypatch.setattr(tf_mod, "_transformers", fake_transformers) + monkeypatch.setattr(tf_mod, "_transformers_patched", False) + + TransformersExtractor.install_auto_load_patch(lambda: client) + + # Simulate `from transformers import pipeline; pipeline = FakePipeline` + # The user holds a direct reference to the original class, but __init__ + # is now patched at the class level, so construction is still intercepted. + original_cls = FakePipeline + original_cls("text-classification", model="distilbert") + + assert any( + event == "tracked" for event, _ in [t for t in loaded if isinstance(t, tuple)] + ) + + def test_gguf_install_auto_load_patch_is_idempotent(monkeypatch): import wildedge.integrations.gguf as gguf_mod diff --git a/wildedge/cli.py b/wildedge/cli.py index f3992d2..3d44b75 100644 --- a/wildedge/cli.py +++ b/wildedge/cli.py @@ -46,6 +46,12 @@ def build_parser() -> argparse.ArgumentParser: default="none", help="Comma-separated hub trackers to enable (huggingface, torchhub). Default: none.", ) + run.add_argument( + "--sampling-interval", + type=float, + default=constants.DEFAULT_SAMPLING_INTERVAL_S, + help="Hardware sampling interval in seconds. 0 to disable. Default: 30.", + ) run.add_argument( "--flush-timeout", type=float, @@ -112,6 +118,12 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Attempt TCP reachability check to DSN host:port.", ) + doctor.add_argument( + "--sampling-interval", + type=float, + default=constants.DEFAULT_SAMPLING_INTERVAL_S, + help="Validate intended hardware sampling interval (seconds). 0 to disable.", + ) doctor.add_argument( "--batch-size", type=int, @@ -228,6 +240,7 @@ def run_command(parsed: argparse.Namespace) -> int: env[constants.ENV_DEBUG] = "1" env[constants.ENV_INTEGRATIONS] = parsed.integrations env[constants.ENV_HUBS] = parsed.hubs + env[constants.ENV_SAMPLING_INTERVAL] = str(parsed.sampling_interval) env[constants.ENV_FLUSH_TIMEOUT] = str(parsed.flush_timeout) env[constants.ENV_PROPAGATE] = "1" if parsed.propagate else "0" env[constants.ENV_STRICT_INTEGRATIONS] = "1" if parsed.strict_integrations else "0" @@ -293,6 +306,8 @@ def validate_runtime_config(parsed: argparse.Namespace) -> tuple[bool, str]: return False, "max_event_age_sec must be > 0" if parsed.max_dead_letter_batches < 0: return False, "max_dead_letter_batches must be >= 0" + if parsed.sampling_interval < 0: + return False, "sampling_interval must be >= 0 (0 = disabled)" return True, "OK" diff --git a/wildedge/client.py b/wildedge/client.py index eeda7fc..dc9c9d2 100644 --- a/wildedge/client.py +++ b/wildedge/client.py @@ -146,7 +146,7 @@ def __init__( dead_letter_dir: str | None = None, max_dead_letter_batches: int = constants.DEFAULT_MAX_DEAD_LETTER_BATCHES, on_delivery_failure: Callable[[str, int, int], None] | None = None, - sampling_interval_s: float | None = 30.0, + sampling_interval_s: float | None = constants.DEFAULT_SAMPLING_INTERVAL_S, ): env = read_client_env(dsn=dsn, debug=debug, app_identity=app_identity) dsn = env.dsn @@ -237,6 +237,10 @@ def __init__( start_sampler(interval_s=sampling_interval_s) self.auto_loaded: set[str] = set() + # Stores the active weakref.finalize for each auto-loaded model so that + # when a wrapping object (e.g. a Pipeline) supersedes the inner model, the + # original finalizer can be detached and re-registered on the outer object. + self._auto_load_finalizers: dict[str, weakref.ref] = {} # Active hub trackers keyed by hub name. Populated by _activate_hub() # when instrument() is called with a hub name. self.hub_trackers: dict[str, BaseHubTracker] = {} @@ -535,8 +539,32 @@ def _on_model_auto_loaded( downloads = thread_records handle = self.register_model(obj, model_id=model_id) + already_tracked = handle.model_id in self.auto_loaded self.auto_loaded.add(handle.model_id) + if already_tracked: + # A wrapping object (e.g. Pipeline) is superseding the inner model + # that was already registered (e.g. via from_pretrained). Wire up + # hooks on the new object so inference is tracked, transfer the + # unload finalizer so only one unload event fires, and skip + # duplicate download/load events. + extractor = self._find_extractor(obj) + if extractor is not None: + extractor.install_hooks(obj, handle) + prev_fin = self._auto_load_finalizers.pop(handle.model_id, None) + if prev_fin is not None: + prev_fin.detach() + loaded_at = time.perf_counter() + + def _on_unload_superseded() -> None: + handle.track_unload( + duration_ms=0, reason="gc", uptime_ms=elapsed_ms(loaded_at) + ) + + fin = weakref.finalize(obj, _on_unload_superseded) + self._auto_load_finalizers[handle.model_id] = fin + return + memory = self._memory_bytes_for(obj) # Emit one download event per repo_id, aggregating per-file records. @@ -576,7 +604,8 @@ def _on_unload() -> None: duration_ms=0, reason="gc", uptime_ms=elapsed_ms(loaded_at) ) - weakref.finalize(obj, _on_unload) + fin = weakref.finalize(obj, _on_unload) + self._auto_load_finalizers[handle.model_id] = fin def load(self, model_class: type, *args: Any, **kwargs: Any) -> object: """ diff --git a/wildedge/constants.py b/wildedge/constants.py index 07cd45b..91a8e28 100644 --- a/wildedge/constants.py +++ b/wildedge/constants.py @@ -13,8 +13,10 @@ ENV_STRICT_INTEGRATIONS = "WILDEDGE_STRICT_INTEGRATIONS" ENV_PROPAGATE = "WILDEDGE_PROPAGATE" ENV_PRINT_STARTUP_REPORT = "WILDEDGE_PRINT_STARTUP_REPORT" +ENV_SAMPLING_INTERVAL = "WILDEDGE_SAMPLING_INTERVAL_S" # Defaults +DEFAULT_SAMPLING_INTERVAL_S = 30.0 DEFAULT_MAX_QUEUE_SIZE = 200 DEFAULT_BATCH_SIZE = 10 DEFAULT_FLUSH_INTERVAL_SEC = 60.0 diff --git a/wildedge/integrations/transformers.py b/wildedge/integrations/transformers.py index 0528ef8..e5a3b2c 100644 --- a/wildedge/integrations/transformers.py +++ b/wildedge/integrations/transformers.py @@ -16,7 +16,7 @@ TextInputMeta, TopKPrediction, ) -from wildedge.integrations.base import BaseExtractor, patch_instance_call_once +from wildedge.integrations.base import BaseExtractor from wildedge.integrations.common import debug_failure, dtype_to_quantization from wildedge.model import ModelInfo from wildedge.timing import elapsed_ms @@ -32,13 +32,17 @@ # --- Patch state --- _transformers_patched = False _TRANSFORMERS_PATCH_LOCK = threading.Lock() -TRANSFORMERS_AUTO_LOAD_PATCH_NAME = "transformers_auto_load" -# --- Pipeline instance patching --- +# Patch names used to guard against double-patching. +TRANSFORMERS_AUTO_LOAD_PATCH_NAME = "transformers_auto_load" +PIPELINE_INIT_PATCH_NAME = "transformers_pipeline_init" PIPELINE_CALL_PATCH_NAME = "transformers_pipeline_call" + +# Per-instance attribute that stores the ModelHandle on a pipeline object. PIPELINE_HANDLE_ATTR = "__wildedge_pipeline_handle__" -# Thread-local flag: suppress from_pretrained tracking when called inside pipeline() +# Thread-local flag: suppress from_pretrained tracking when called inside +# Pipeline.__init__ so the pipeline-level tracking takes precedence. _tl = threading.local() @@ -120,8 +124,8 @@ def detect_accelerator(obj: object) -> str: model = getattr(obj, "model", obj) first = next(model.parameters()) # type: ignore[union-attr] return str(getattr(first.device, "type", "cpu") or "cpu") - except Exception: - pass + except Exception as exc: + debug_transformers_failure("accelerator detection", exc) return "cpu" @@ -142,7 +146,7 @@ def infer_task_from_arch(arch: str | None) -> str | None: # --------------------------------------------------------------------------- -# Pipeline call patching +# Pipeline input/output helpers # --------------------------------------------------------------------------- @@ -240,46 +244,6 @@ def pipeline_modalities(task: str | None) -> tuple[str | None, str | None]: return "text", None -def build_pipeline_patched_call(original_call): # type: ignore[no-untyped-def] - def patched_call(self_inner, inputs, *args, **kwargs): # type: ignore[no-untyped-def] - handle = getattr(self_inner, PIPELINE_HANDLE_ATTR, None) - if handle is None: - return original_call(self_inner, inputs, *args, **kwargs) - - task = getattr(self_inner, "task", None) - batch_size: int | None = ( - len(inputs) - if isinstance(inputs, list) - else (1 if isinstance(inputs, str) else None) - ) - input_meta = pipeline_input_meta(inputs) - input_modality, output_modality = pipeline_modalities(task) - - t0 = time.perf_counter() - try: - outputs = original_call(self_inner, inputs, *args, **kwargs) - duration_ms = elapsed_ms(t0) - output_meta = pipeline_output_meta(task, outputs) - handle.track_inference( - duration_ms=duration_ms, - batch_size=batch_size, - input_modality=input_modality, - output_modality=output_modality, - input_meta=input_meta, - output_meta=output_meta, - success=True, - ) - return outputs - except Exception as exc: - handle.track_error( - error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], - ) - raise - - return patched_call - - # --------------------------------------------------------------------------- # Extractor # --------------------------------------------------------------------------- @@ -331,12 +295,9 @@ def install_hooks(self, obj: object, handle: ModelHandle) -> None: handle.detected_accelerator = detect_accelerator(obj) if is_pipeline(obj): + # Attach handle to the instance; the class-level __call__ patch + # picks it up and no-ops for instances without this attribute. setattr(obj, PIPELINE_HANDLE_ATTR, handle) - patch_instance_call_once( - obj, - patch_name=PIPELINE_CALL_PATCH_NAME, - make_patched_call=build_pipeline_patched_call, - ) else: # PreTrainedModel: use PyTorch forward hooks _local = threading.local() @@ -385,14 +346,13 @@ def post_hook(module, args, output): # type: ignore[no-untyped-def] @classmethod def install_auto_load_patch(cls, client_ref: object) -> None: - """Patch transformers.pipeline and PreTrainedModel.from_pretrained. - - Called once at WildEdge client initialisation. Any subsequent - ``pipeline(...)`` or ``AutoModel.from_pretrained(...)`` call is timed - and registered automatically. HuggingFace Hub downloads are intercepted - for the duration of the call and emitted as a model_download event. - A thread-local guard prevents double-tracking when pipeline() calls - from_pretrained() internally. + """Patch Pipeline.__init__, Pipeline.__call__, and from_pretrained. + + Patches at the class level so both ``import transformers; transformers.pipeline()`` + and ``from transformers import pipeline; pipeline()`` are intercepted correctly. + The Pipeline.__init__ patch times load and calls _on_model_auto_loaded with the + fully-constructed pipeline (model already on its target device). The __call__ + patch dispatches inference tracking via the per-instance handle attribute. """ global _transformers_patched if _transformers_patched or _transformers is None: @@ -401,39 +361,88 @@ def install_auto_load_patch(cls, client_ref: object) -> None: with _TRANSFORMERS_PATCH_LOCK: if _transformers_patched: return - cls._patch_pipeline(client_ref) + cls._patch_pipeline_init(client_ref) + cls._patch_pipeline_call() cls._patch_from_pretrained(client_ref) _transformers_patched = True @classmethod - def _patch_pipeline(cls, client_ref: object) -> None: - original_pipeline = _transformers.pipeline + def _patch_pipeline_init(cls, client_ref: object) -> None: + original_init = _transformers.Pipeline.__init__ if ( - getattr(original_pipeline, "__wildedge_patch_name__", None) - == TRANSFORMERS_AUTO_LOAD_PATCH_NAME + getattr(original_init, "__wildedge_patch_name__", None) + == PIPELINE_INIT_PATCH_NAME ): return - def patched_pipeline(*args, **kwargs): # type: ignore[no-untyped-def] + def patched_init(self, *args, **kwargs): # type: ignore[no-untyped-def] c = client_ref() # type: ignore[call-arg] hub_before = ( c._snapshot_hub_caches() if c is not None and not c.closed else {} ) t0 = time.perf_counter() - _tl.inside_pipeline = True + _tl.inside_pipeline_init = True try: - pipe = original_pipeline(*args, **kwargs) + original_init(self, *args, **kwargs) finally: - _tl.inside_pipeline = False + _tl.inside_pipeline_init = False load_ms = elapsed_ms(t0) if c is not None and not c.closed: downloads = c._diff_hub_caches(hub_before, load_ms) or None - c._on_model_auto_loaded(pipe, load_ms=load_ms, downloads=downloads) - return pipe + c._on_model_auto_loaded(self, load_ms=load_ms, downloads=downloads) + + patched_init.__wildedge_patch_name__ = PIPELINE_INIT_PATCH_NAME # type: ignore[attr-defined] + patched_init.__wildedge_original_call__ = original_init # type: ignore[attr-defined] + _transformers.Pipeline.__init__ = patched_init + + @classmethod + def _patch_pipeline_call(cls) -> None: + original_call = _transformers.Pipeline.__call__ + if ( + getattr(original_call, "__wildedge_patch_name__", None) + == PIPELINE_CALL_PATCH_NAME + ): + return + + def patched_call(self, inputs, *args, **kwargs): # type: ignore[no-untyped-def] + handle = getattr(self, PIPELINE_HANDLE_ATTR, None) + if handle is None: + return original_call(self, inputs, *args, **kwargs) + + task = getattr(self, "task", None) + batch_size: int | None = ( + len(inputs) + if isinstance(inputs, list) + else (1 if isinstance(inputs, str) else None) + ) + input_meta = pipeline_input_meta(inputs) + input_modality, output_modality = pipeline_modalities(task) + + t0 = time.perf_counter() + try: + outputs = original_call(self, inputs, *args, **kwargs) + duration_ms = elapsed_ms(t0) + output_meta = pipeline_output_meta(task, outputs) + handle.track_inference( + duration_ms=duration_ms, + batch_size=batch_size, + input_modality=input_modality, + output_modality=output_modality, + input_meta=input_meta, + output_meta=output_meta, + success=True, + ) + return outputs + except Exception as exc: + handle.track_error( + error_code="UNKNOWN", + error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + ) + raise - patched_pipeline.__wildedge_patch_name__ = TRANSFORMERS_AUTO_LOAD_PATCH_NAME # type: ignore[attr-defined] - patched_pipeline.__wildedge_original_call__ = original_pipeline # type: ignore[attr-defined] - _transformers.pipeline = patched_pipeline + patched_call.__wildedge_patch_name__ = PIPELINE_CALL_PATCH_NAME # type: ignore[attr-defined] + patched_call.__wildedge_original_call__ = original_call # type: ignore[attr-defined] + _transformers.Pipeline.__call__ = patched_call @classmethod def _patch_from_pretrained(cls, client_ref: object) -> None: @@ -447,8 +456,8 @@ def _patch_from_pretrained(cls, client_ref: object) -> None: original_func = original_bound.__func__ def patched_from_pretrained(model_cls, *args, **kwargs): # type: ignore[no-untyped-def] - # Don't double-track models loaded inside pipeline() - if getattr(_tl, "inside_pipeline", False): + # Don't double-track models loaded inside Pipeline.__init__ + if getattr(_tl, "inside_pipeline_init", False): return original_func(model_cls, *args, **kwargs) c = client_ref() # type: ignore[call-arg] hub_before = ( diff --git a/wildedge/runtime/bootstrap.py b/wildedge/runtime/bootstrap.py index 151a5e4..e9e1154 100644 --- a/wildedge/runtime/bootstrap.py +++ b/wildedge/runtime/bootstrap.py @@ -38,6 +38,7 @@ def clear_runtime_env() -> None: constants.ENV_STRICT_INTEGRATIONS, constants.ENV_PROPAGATE, constants.ENV_PRINT_STARTUP_REPORT, + constants.ENV_SAMPLING_INTERVAL, ): os.environ.pop(key, None) @@ -106,7 +107,12 @@ def install_runtime(*, install_signal_handlers: bool = True) -> RuntimeContext: if not env.dsn: raise RuntimeConfigError(f"{ENV_DSN} must be set to use `wildedge run`.") - client = WildEdge(dsn=env.dsn, app_version=env.app_version, debug=env.debug) + client = WildEdge( + dsn=env.dsn, + app_version=env.app_version, + debug=env.debug, + sampling_interval_s=env.sampling_interval_s, + ) statuses: list[dict[str, str]] = [] for integration in env.integrations: diff --git a/wildedge/settings.py b/wildedge/settings.py index fc32606..a01b88c 100644 --- a/wildedge/settings.py +++ b/wildedge/settings.py @@ -27,6 +27,7 @@ class RuntimeEnv: integrations: list[str] hubs: list[str] propagate: bool + sampling_interval_s: float | None @dataclass(frozen=True) @@ -98,6 +99,13 @@ def read_runtime_env( str(constants.DEFAULT_SHUTDOWN_FLUSH_TIMEOUT_SEC), ) ) + raw_sampling = env.get(constants.ENV_SAMPLING_INTERVAL) + sampling_interval_s: float | None + if raw_sampling is None: + sampling_interval_s = constants.DEFAULT_SAMPLING_INTERVAL_S + else: + parsed_sampling = float(raw_sampling) + sampling_interval_s = parsed_sampling if parsed_sampling > 0 else None return RuntimeEnv( dsn=env.get(constants.ENV_DSN), app_version=env.get(constants.ENV_APP_VERSION), @@ -110,6 +118,7 @@ def read_runtime_env( ), hubs=parse_hub_list(env.get(constants.ENV_HUBS), all_hubs), propagate=parse_bool(env.get(constants.ENV_PROPAGATE, "1")), + sampling_interval_s=sampling_interval_s, ) From 37237532b253e02f5f3cd0149c8967ffca2cd1e3 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Mon, 16 Mar 2026 12:27:43 +0100 Subject: [PATCH 10/10] Removed stale comment for the transformer sample --- examples/transformers_example.py | 15 ++++----------- 1 file changed, 4 insertions(+), 11 deletions(-) diff --git a/examples/transformers_example.py b/examples/transformers_example.py index c1d0739..bac9df4 100644 --- a/examples/transformers_example.py +++ b/examples/transformers_example.py @@ -23,7 +23,7 @@ import argparse import platform -import transformers +from transformers import pipeline import wildedge @@ -31,7 +31,7 @@ def run_classify() -> None: - pipe = transformers.pipeline( + pipe = pipeline( "text-classification", model="distilbert-base-uncased-finetuned-sst-2-english", device=_DEVICE, @@ -51,9 +51,7 @@ def run_classify() -> None: def run_generate() -> None: - pipe = transformers.pipeline( - "text-generation", model="gpt2", max_new_tokens=40, device=_DEVICE - ) + 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", @@ -66,9 +64,7 @@ def run_generate() -> None: def run_embed() -> None: - pipe = transformers.pipeline( - "feature-extraction", model="bert-base-uncased", device=_DEVICE - ) + pipe = pipeline("feature-extraction", model="bert-base-uncased", device=_DEVICE) sentences = [ "Machine learning is transforming every industry.", "On-device inference keeps your data private.", @@ -96,9 +92,6 @@ def main() -> None: ) args = parser.parse_args() - # instrument() patches transformers.pipeline and AutoModel.from_pretrained. - # Import transformers as a module (not `from transformers import pipeline`) - # so attribute lookups happen after patching and the correct device is captured. client = wildedge.WildEdge(app_version="1.0.0") # set WILDEDGE_DSN env var client.instrument("transformers", hubs=["huggingface"])