diff --git a/docs/manual-tracking.md b/docs/manual-tracking.md index 8968c16..decc7eb 100644 --- a/docs/manual-tracking.md +++ b/docs/manual-tracking.md @@ -9,6 +9,42 @@ Use manual tracking when your framework is not covered by a WildEdge integration - You want to attach input/output metadata (token counts, image dimensions, confidence scores, etc.) - You want to record user feedback tied to a specific inference +## The default client + +Every process has at most one default client, and the module-level API +delegates to it. You rarely need to hold a client instance yourself: + +```python +import wildedge + +with wildedge.trace(run_id="run-123"): + with wildedge.span(kind="agent_step", name="plan", step_index=1): + ... +``` + +The default client comes from wherever initialization happened first: + +- `wildedge run` installs one before your code loads. A later + `wildedge.init()` without `dsn` reuses it, so the same code works with or + without the CLI wrapper. +- `wildedge.init(...)` creates one (or reuses the existing default) and + registers it. +- Any module-level call (`wildedge.span`, `wildedge.register_model`, ...) + creates one lazily from the environment: `WILDEDGE_DSN` for the endpoint + (without it the client is a no-op), plus `WILDEDGE_INTEGRATIONS` and + `WILDEDGE_HUBS` for auto-instrumentation. Unset means nothing is patched. + +`wildedge.get_client()` returns the default client when you need direct +access; `wildedge.set_default_client()` overrides it (useful in tests). +`wildedge.trace` is the exception to the delegation rule: a trace only sets +correlation contextvars and emits nothing, so it never touches or creates a +client. + +Auto-instrumentation patches at creation time, so ordering matters: the CLI +patches before any user code runs, an early `init()` covers everything created +after it, and a lazily created client only covers what is created after the +first module-level call. When in doubt, call `init()` at application startup. + ## torch and keras For `torch` and `keras`, models are user-defined subclasses with no constructor to patch. Use `client.load()` to get load, unload, and inference tracking automatically: @@ -296,6 +332,30 @@ with client.trace(run_id="run-123", agent_id="agent-1"): handle.track_inference(duration_ms=t.elapsed_ms, input_modality="text", output_modality="generation") ``` +Both are also available module-level on the default client, so this is +equivalent without holding a client: + +```python +with wildedge.trace(run_id="run-123", agent_id="agent-1"): + with wildedge.span(kind="agent_step", name="plan", step_index=1): + ... +``` + +### Recording results on an open span + +The span object is mutable until the block exits; set outcomes you only know +mid-block directly on it. `set_attributes()` merges into the span's +attributes, `fail()` marks it failed with an optional summary, and an +exception escaping the block sets `status="error"` automatically: + +```python +with wildedge.span(kind="eval", name="lint") as span: + errors = lint(result) + span.set_attributes(lint_errors=len(errors)) + if errors: + span.fail(f"{len(errors)} lint errors") +``` + If you need to set correlation fields without emitting a span event, use the lower-level `span_context()` directly: diff --git a/tests/conftest.py b/tests/conftest.py index e8c8199..d57f383 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,6 +19,15 @@ def reset_hardware_sampler(): stop_sampler() +@pytest.fixture(autouse=True) +def reset_default_client(): + from wildedge.defaults import set_default_client + + set_default_client(None) + yield + set_default_client(None) + + PLATFORM_MARKS = { "requires_linux": "linux", "requires_macos": "darwin", diff --git a/tests/test_cli.py b/tests/test_cli.py index 4fbf3c8..1d2ba1f 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -472,3 +472,32 @@ def test_parse_run_args_only_double_dash_raises(): """['--'] with nothing after raises ValueError.""" with pytest.raises(ValueError, match="missing command"): cli.parse_run_args(["--"]) + + +def test_install_runtime_registers_default_client(monkeypatch): + from wildedge.defaults import peek_default_client + + class FakeWildEdge: + SUPPORTED_INTEGRATIONS = {"onnx"} + + 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] + pass + + def flush(self, timeout): # type: ignore[no-untyped-def] + pass + + def close(self): # type: ignore[no-untyped-def] + pass + + monkeypatch.setattr(bootstrap, "WildEdge", FakeWildEdge) + monkeypatch.setenv(constants.ENV_DSN, "https://secret@ingest.wildedge.dev/key") + monkeypatch.setattr(bootstrap.importlib.util, "find_spec", lambda _: object()) + + context = bootstrap.install_runtime() + try: + assert peek_default_client() is context.client + finally: + context.shutdown() diff --git a/tests/test_defaults.py b/tests/test_defaults.py new file mode 100644 index 0000000..a8cbe9e --- /dev/null +++ b/tests/test_defaults.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +from unittest.mock import MagicMock + +import wildedge +import wildedge.defaults as defaults +from wildedge import constants + + +class DummyClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.instrument_calls: list[tuple[str | None, list[str] | None]] = [] + + def instrument(self, integration, *, hubs=None): + self.instrument_calls.append((integration, hubs)) + + +def test_get_client_lazily_creates_noop_without_dsn(monkeypatch): + monkeypatch.delenv(constants.ENV_DSN, raising=False) + monkeypatch.delenv(constants.ENV_INTEGRATIONS, raising=False) + client = wildedge.get_client() + assert client.noop is True + assert wildedge.get_client() is client + + +def test_set_default_client_wins(): + client = DummyClient() + defaults.set_default_client(client) + assert wildedge.get_client() is client + assert defaults.peek_default_client() is client + + +def test_peek_does_not_create(): + assert defaults.peek_default_client() is None + + +def test_lazy_client_reads_integrations_env(monkeypatch): + monkeypatch.setattr(defaults, "WildEdge", DummyClient) + monkeypatch.setenv(constants.ENV_INTEGRATIONS, "openai") + monkeypatch.delenv(constants.ENV_HUBS, raising=False) + assert wildedge.get_client().instrument_calls == [("openai", None)] + + +def test_lazy_client_without_env_instruments_nothing(monkeypatch): + monkeypatch.setattr(defaults, "WildEdge", DummyClient) + monkeypatch.delenv(constants.ENV_INTEGRATIONS, raising=False) + monkeypatch.delenv(constants.ENV_HUBS, raising=False) + assert wildedge.get_client().instrument_calls == [] + + +def test_lazy_client_hubs_only(monkeypatch): + monkeypatch.setattr(defaults, "WildEdge", DummyClient) + monkeypatch.delenv(constants.ENV_INTEGRATIONS, raising=False) + monkeypatch.setenv(constants.ENV_HUBS, "huggingface") + assert wildedge.get_client().instrument_calls == [(None, ["huggingface"])] + + +def test_lazy_client_env_errors_do_not_raise(monkeypatch): + class Exploding(DummyClient): + def instrument(self, integration, *, hubs=None): + raise ValueError("unknown integration") + + monkeypatch.setattr(defaults, "WildEdge", Exploding) + monkeypatch.setenv(constants.ENV_INTEGRATIONS, "bogus") + client = wildedge.get_client() + assert isinstance(client, Exploding) + + +def test_span_delegates_to_default_client(): + client = MagicMock() + defaults.set_default_client(client) + wildedge.span(kind="tool", name="lint", step_index=2) + kwargs = client.span.call_args.kwargs + assert kwargs["kind"] == "tool" + assert kwargs["name"] == "lint" + assert kwargs["step_index"] == 2 + + +def test_track_span_delegates_to_default_client(): + client = MagicMock() + defaults.set_default_client(client) + wildedge.track_span(kind="tool", name="call", duration_ms=15, status="ok") + kwargs = client.track_span.call_args.kwargs + assert kwargs["duration_ms"] == 15 + + +def test_register_model_delegates_to_default_client(): + client = MagicMock() + defaults.set_default_client(client) + model = object() + wildedge.register_model(model, model_id="org/model", source="openrouter") + args, kwargs = client.register_model.call_args + assert args == (model,) + assert kwargs["model_id"] == "org/model" + assert kwargs["source"] == "openrouter" + + +def test_flush_delegates_to_default_client(): + client = MagicMock() + defaults.set_default_client(client) + wildedge.flush(timeout=2.0) + client.flush.assert_called_once_with(timeout=2.0) + + +def test_trace_is_pure_context_and_creates_no_client(monkeypatch): + monkeypatch.delenv(constants.ENV_DSN, raising=False) + with wildedge.trace(run_id="r1", agent_id="a1"): + ctx = wildedge.get_trace_context() + assert ctx is not None + assert ctx.run_id == "r1" + assert ctx.agent_id == "a1" + assert wildedge.get_trace_context() is None + assert defaults.peek_default_client() is None + + +def test_module_api_works_end_to_end_without_dsn(monkeypatch): + monkeypatch.delenv(constants.ENV_DSN, raising=False) + monkeypatch.delenv(constants.ENV_INTEGRATIONS, raising=False) + with wildedge.trace(run_id="r"): + with wildedge.span(kind="eval", name="lint") as span: + span.set_attributes(lint_errors=3) + span.fail("3 lint errors") + assert wildedge.get_client().noop is True diff --git a/tests/test_init.py b/tests/test_init.py index 3450446..c8d6a7a 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -67,3 +67,84 @@ def instrument(self, integration, *, hubs=None): assert isinstance(client, DummyClient) assert logs == ["wildedge: init called without integrations or hubs"] + + +def test_init_reuses_existing_default_client(monkeypatch): + from wildedge.defaults import peek_default_client, set_default_client + + calls: list[tuple[str | None, list[str] | None]] = [] + + class ExistingClient: + def instrument(self, integration, *, hubs=None): + calls.append((integration, hubs)) + + existing = ExistingClient() + set_default_client(existing) + + client = wildedge.init(integrations=["onnx"]) + + assert client is existing + assert calls == [("onnx", None)] + assert peek_default_client() is existing + + +def test_init_with_dsn_replaces_default_client(monkeypatch): + from wildedge.defaults import peek_default_client, set_default_client + + class ExistingClient: + def instrument(self, integration, *, hubs=None): + raise AssertionError("existing client must not be touched") + + class DummyClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def instrument(self, integration, *, hubs=None): + pass + + existing = ExistingClient() + set_default_client(existing) + monkeypatch.setattr(convenience, "WildEdge", DummyClient) + + client = wildedge.init(dsn="https://secret@ingest.wildedge.dev/key") + + assert isinstance(client, DummyClient) + assert client is not existing + assert peek_default_client() is client + + +def test_init_reuse_warns_about_ignored_kwargs(monkeypatch): + from wildedge.defaults import set_default_client + + class ExistingClient: + def instrument(self, integration, *, hubs=None): + pass + + logs: list[str] = [] + set_default_client(ExistingClient()) + monkeypatch.setattr( + convenience.logger, "warning", lambda msg, *args: logs.append(msg % args) + ) + + wildedge.init(integrations=["onnx"], app_version="1.0.0") + + assert logs == [ + "wildedge: init() reusing the existing default client; ignoring app_version" + ] + + +def test_init_registers_new_client_as_default(monkeypatch): + from wildedge.defaults import peek_default_client + + class DummyClient: + def __init__(self, **kwargs): + self.kwargs = kwargs + + def instrument(self, integration, *, hubs=None): + pass + + monkeypatch.setattr(convenience, "WildEdge", DummyClient) + + client = wildedge.init(dsn="https://secret@ingest.wildedge.dev/key") + + assert peek_default_client() is client diff --git a/tests/test_tracing.py b/tests/test_tracing.py index 944f1e0..b5de0fd 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -91,3 +91,50 @@ def test_nested_spans_correct_parent_chain(): assert inner_ev["parent_span_id"] == outer.span_id assert outer_ev["span_id"] == outer.span_id assert outer_ev["parent_span_id"] is None + + +def test_span_set_attributes_merges(): + events: list[dict] = [] + client = _FakeClient(events) + + with SpanContextManager( + client, kind="eval", name="lint", attributes={"a": 1} + ) as span: + span.set_attributes(b=2) + span.set_attributes(a=3) + + assert events[0]["attributes"] == {"a": 3, "b": 2} + + +def test_span_set_attributes_from_none(): + events: list[dict] = [] + client = _FakeClient(events) + + with SpanContextManager(client, kind="eval", name="lint") as span: + span.set_attributes(lint_errors=4) + + assert events[0]["attributes"] == {"lint_errors": 4} + + +def test_span_fail_sets_status_and_summary(): + events: list[dict] = [] + client = _FakeClient(events) + + with SpanContextManager(client, kind="eval", name="lint") as span: + span.fail("4 lint errors") + + assert events[0]["status"] == "error" + assert events[0]["output_summary"] == "4 lint errors" + + +def test_span_fail_without_detail_keeps_summary(): + events: list[dict] = [] + client = _FakeClient(events) + + with SpanContextManager( + client, kind="eval", name="lint", output_summary="kept" + ) as span: + span.fail() + + assert events[0]["status"] == "error" + assert events[0]["output_summary"] == "kept" diff --git a/wildedge/__init__.py b/wildedge/__init__.py index a2d68b9..afc3feb 100644 --- a/wildedge/__init__.py +++ b/wildedge/__init__.py @@ -4,6 +4,15 @@ from wildedge.client import SpanContextManager, WildEdge from wildedge.convenience import init from wildedge.decorators import track +from wildedge.defaults import ( + flush, + get_client, + register_model, + set_default_client, + span, + trace, + track_span, +) from wildedge.events import ( AdapterDownload, AdapterLoad, @@ -36,6 +45,13 @@ __all__ = [ "WildEdge", "init", + "get_client", + "set_default_client", + "trace", + "span", + "track_span", + "register_model", + "flush", "Attachment", "capture_hardware", "HardwareContext", diff --git a/wildedge/client.py b/wildedge/client.py index 48fbfe2..9a46357 100644 --- a/wildedge/client.py +++ b/wildedge/client.py @@ -216,6 +216,18 @@ async def __aenter__(self): async def __aexit__(self, exc_type, exc_val, exc_tb): return self.__exit__(exc_type, exc_val, exc_tb) + def set_attributes(self, **attributes: Any) -> None: + """Merge ``attributes`` into the span's attribute dict.""" + if self.attributes is None: + self.attributes = {} + self.attributes.update(attributes) + + def fail(self, detail: str | None = None) -> None: + """Mark the span as failed; ``detail`` becomes the output summary.""" + self.status = "error" + if detail is not None: + self.output_summary = detail + class WildEdge: """ diff --git a/wildedge/convenience.py b/wildedge/convenience.py index a9e61aa..bfa2f1a 100644 --- a/wildedge/convenience.py +++ b/wildedge/convenience.py @@ -4,6 +4,7 @@ from typing import Any from wildedge.client import WildEdge +from wildedge.defaults import peek_default_client, set_default_client from wildedge.logging import logger @@ -22,11 +23,26 @@ def init( **kwargs: Any, ) -> WildEdge: """ - Convenience initializer: construct a WildEdge client and instrument integrations. + Initialize the process default client and instrument integrations. - Additional keyword arguments are forwarded to WildEdge(...). + When a default client already exists (typically installed by ``wildedge + run`` before user code loaded) and no ``dsn`` is passed, that client is + reused: the requested integrations are applied to it (instrumenting is + idempotent per process) and it is returned. Passing ``dsn`` always + constructs a fresh client and makes it the new default. + + Additional keyword arguments are forwarded to WildEdge(...) when a new + client is constructed. """ - client = WildEdge(**kwargs) + client = None if "dsn" in kwargs else peek_default_client() + if client is not None and kwargs: + logger.warning( + "wildedge: init() reusing the existing default client; ignoring %s", + ", ".join(sorted(kwargs)), + ) + if client is None: + client = WildEdge(**kwargs) + normalized_integrations = _normalize_list(integrations) normalized_hubs = _normalize_list(hubs) @@ -38,4 +54,5 @@ def init( elif getattr(client, "debug", False): logger.debug("wildedge: init called without integrations or hubs") + set_default_client(client) return client diff --git a/wildedge/defaults.py b/wildedge/defaults.py new file mode 100644 index 0000000..1a29dda --- /dev/null +++ b/wildedge/defaults.py @@ -0,0 +1,198 @@ +"""Process-wide default client and the module-level tracking API. + +One client per process is the common case: `wildedge run` installs one before +user code loads, `wildedge.init()` creates or reuses one, and the module-level +functions below (`wildedge.trace`, `wildedge.span`, ...) delegate to whichever +client is current. Application code never has to thread a client instance +through call sites. +""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +from wildedge import constants +from wildedge.client import SpanContextManager, WildEdge +from wildedge.events.span import SpanKind, SpanStatus +from wildedge.hubs.registry import supported_hubs +from wildedge.integrations.registry import supported_integrations +from wildedge.logging import logger +from wildedge.model import ModelHandle +from wildedge.settings import parse_hub_list, parse_integration_list +from wildedge.tracing import trace_context + +_lock = threading.Lock() +_default_client: WildEdge | None = None + + +def set_default_client(client: WildEdge | None) -> None: + """Register ``client`` as the process default; ``None`` clears it.""" + global _default_client + with _lock: + _default_client = client + + +def peek_default_client() -> WildEdge | None: + """Return the process default client without creating one.""" + return _default_client + + +def get_client() -> WildEdge: + """ + Return the process default client, creating one from the environment on + first use. + + The lazily created client reads ``WILDEDGE_DSN`` exactly like a directly + constructed ``WildEdge`` (without a DSN it is a no-op) and enables only + the integrations and hubs named in ``WILDEDGE_INTEGRATIONS`` and + ``WILDEDGE_HUBS``. When those variables are unset nothing is patched: + auto-instrumentation stays opt-in. + """ + global _default_client + with _lock: + if _default_client is None: + _default_client = _client_from_env() + return _default_client + + +def _client_from_env() -> WildEdge: + client = WildEdge() + raw_integrations = os.environ.get(constants.ENV_INTEGRATIONS) + raw_hubs = os.environ.get(constants.ENV_HUBS) + integrations = ( + parse_integration_list(raw_integrations, sorted(supported_integrations())) + if raw_integrations + else [] + ) + hubs = parse_hub_list(raw_hubs, sorted(supported_hubs())) if raw_hubs else [] + + targets: list[tuple[str | None, list[str] | None]] = [] + if integrations: + targets = [(integration, hubs or None) for integration in integrations] + elif hubs: + targets = [(None, hubs)] + # A bad environment value must not break the caller that triggered lazy + # creation, so failures downgrade to a warning here; explicit init() with + # the same names would raise. + for integration, target_hubs in targets: + try: + client.instrument(integration, hubs=target_hubs) + except Exception as exc: + logger.warning( + "wildedge: instrument(%r) from environment failed: %s", + integration, + exc, + ) + return client + + +# A trace sets correlation contextvars and emits nothing, so it needs no +# client; the module-level name is the context manager itself. +trace = trace_context + + +def span( + *, + kind: SpanKind, + name: str, + status: SpanStatus = "ok", + model_id: str | None = None, + input_summary: str | None = None, + output_summary: str | None = None, + attributes: dict[str, Any] | None = None, + trace_id: str | None = None, + span_id: str | None = None, + parent_span_id: str | None = None, + run_id: str | None = None, + agent_id: str | None = None, + step_index: int | None = None, + conversation_id: str | None = None, + context: dict[str, Any] | None = None, +) -> SpanContextManager: + """``WildEdge.span`` on the default client.""" + return get_client().span( + kind=kind, + name=name, + status=status, + model_id=model_id, + input_summary=input_summary, + output_summary=output_summary, + attributes=attributes, + trace_id=trace_id, + span_id=span_id, + parent_span_id=parent_span_id, + run_id=run_id, + agent_id=agent_id, + step_index=step_index, + conversation_id=conversation_id, + context=context, + ) + + +def track_span( + *, + kind: SpanKind, + name: str, + duration_ms: int, + status: SpanStatus = "ok", + model_id: str | None = None, + input_summary: str | None = None, + output_summary: str | None = None, + attributes: dict[str, Any] | None = None, + trace_id: str | None = None, + span_id: str | None = None, + parent_span_id: str | None = None, + run_id: str | None = None, + agent_id: str | None = None, + step_index: int | None = None, + conversation_id: str | None = None, + context: dict[str, Any] | None = None, +) -> str: + """``WildEdge.track_span`` on the default client.""" + return get_client().track_span( + kind=kind, + name=name, + duration_ms=duration_ms, + status=status, + model_id=model_id, + input_summary=input_summary, + output_summary=output_summary, + attributes=attributes, + trace_id=trace_id, + span_id=span_id, + parent_span_id=parent_span_id, + run_id=run_id, + agent_id=agent_id, + step_index=step_index, + conversation_id=conversation_id, + context=context, + ) + + +def register_model( + model_obj: object, + *, + model_id: str | None = None, + source: str | None = None, + family: str | None = None, + version: str | None = None, + quantization: str | None = None, + auto_instrument: bool = True, +) -> ModelHandle: + """``WildEdge.register_model`` on the default client.""" + return get_client().register_model( + model_obj, + model_id=model_id, + source=source, + family=family, + version=version, + quantization=quantization, + auto_instrument=auto_instrument, + ) + + +def flush(timeout: float = 5.0) -> None: + """Flush the default client's queued events.""" + get_client().flush(timeout=timeout) diff --git a/wildedge/runtime/bootstrap.py b/wildedge/runtime/bootstrap.py index fbc47d7..f94fc52 100644 --- a/wildedge/runtime/bootstrap.py +++ b/wildedge/runtime/bootstrap.py @@ -14,6 +14,7 @@ from wildedge import constants from wildedge.client import WildEdge from wildedge.constants import ENV_DSN, WILDEDGE_AUTOLOAD +from wildedge.defaults import set_default_client from wildedge.hubs.registry import HUBS_BY_NAME, supported_hubs from wildedge.integrations.registry import INTEGRATIONS_BY_NAME, supported_integrations from wildedge.settings import read_runtime_env @@ -115,6 +116,7 @@ def install_runtime(*, install_signal_handlers: bool = True) -> RuntimeContext: debug=env.debug, sampling_interval_s=env.sampling_interval_s, ) + set_default_client(client) statuses: list[dict[str, str]] = [] for integration in env.integrations: