diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 4f942d3..b9f9123 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,6 +7,7 @@ on: - 'wildedge/**' - 'tests/**' - 'scripts/**' + - 'examples/**' - 'pyproject.toml' - 'uv.lock' - '.github/workflows/ci.yml' @@ -16,11 +17,31 @@ on: - 'wildedge/**' - 'tests/**' - 'scripts/**' + - 'examples/**' - 'pyproject.toml' - 'uv.lock' - '.github/workflows/ci.yml' jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + - name: Install uv + run: | + pip install uv + - name: Install dependencies + run: | + uv sync --group dev + - name: Ruff lint + run: uv run ruff check --output-format=github . + - name: Ruff format + run: uv run ruff format --check . + test: runs-on: ${{ matrix.os }} strategy: diff --git a/.github/workflows/compat.yml b/.github/workflows/compat.yml index 3e267e3..46ef03d 100644 --- a/.github/workflows/compat.yml +++ b/.github/workflows/compat.yml @@ -1,5 +1,12 @@ name: Compatibility +# Runs tests/compat against the pinned library versions in +# scripts/compat_matrix.py, which is also the source of the job matrix and of +# docs/compatibility.md. +# +# Every run uses the full matrix. It runs on pull requests into main or a +# release branch, on pushes to them, nightly against devel, and on demand. + on: push: branches: [ main, 'release/**' ] @@ -10,7 +17,7 @@ on: - 'pyproject.toml' - '.github/workflows/compat.yml' pull_request: - branches: [ main ] + branches: [ main, 'release/**' ] types: [opened, synchronize, reopened] paths: - 'wildedge/**' @@ -23,46 +30,23 @@ on: workflow_dispatch: jobs: - compat: - if: github.event_name != 'pull_request' || startsWith(github.head_ref, 'release/') + matrix: runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - python-version: ["3.10", "3.11", "3.12", "3.13"] - integration: ["onnx", "torch", "timm", "tensorflow", "gguf", "openai", "transformers", "ultralytics"] - version-set: ["min", "current"] - exclude: - - python-version: "3.13" - integration: "tensorflow" + outputs: + rows: ${{ steps.rows.outputs.rows }} steps: - uses: actions/checkout@v4 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - - name: Install project - run: | - python -m pip install --upgrade pip - pip install -e . pytest pytest-asyncio pytest-mock - - name: Install compatibility dependencies - run: | - python scripts/compat_matrix.py deps \ - --integration "${{ matrix.integration }}" \ - --version-set "${{ matrix.version-set }}" \ - --python-version "${{ matrix.python-version }}" > compat-requirements.txt - pip install -r compat-requirements.txt - - name: Run compatibility test - run: python -m pytest "tests/compat/test_${{ matrix.integration }}_compat.py" -q + - id: rows + run: echo "rows=$(python scripts/compat_matrix.py github-matrix)" >> "$GITHUB_OUTPUT" - compat-mlx: - if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' - runs-on: macos-14 + compat: + needs: matrix + name: compat (${{ matrix.python-version }}, ${{ matrix.integration }}, ${{ matrix.version-set }}) + runs-on: ${{ matrix.runs-on }} + continue-on-error: ${{ matrix.canary }} strategy: fail-fast: false - matrix: - python-version: ["3.12", "3.13"] - version-set: ["current"] + matrix: ${{ fromJSON(needs.matrix.outputs.rows) }} steps: - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} @@ -76,40 +60,23 @@ jobs: - name: Install compatibility dependencies run: | python scripts/compat_matrix.py deps \ - --integration "mlx" \ + --integration "${{ matrix.integration }}" \ --version-set "${{ matrix.version-set }}" \ --python-version "${{ matrix.python-version }}" > compat-requirements.txt pip install -r compat-requirements.txt - - name: Run compatibility test - run: python -m pytest tests/compat/test_mlx_compat.py -q - - compat-canary-314: - if: github.event_name != 'pull_request' || startsWith(github.head_ref, 'release/') - runs-on: ubuntu-latest - continue-on-error: true - strategy: - fail-fast: false - matrix: - integration: ["torch", "timm", "tensorflow"] - version-set: ["current"] - exclude: - - integration: "tensorflow" - steps: - - uses: actions/checkout@v4 - - name: Set up Python 3.14 - uses: actions/setup-python@v5 - with: - python-version: "3.14" - - name: Install project + - name: Check installed versions match the pins run: | - python -m pip install --upgrade pip - pip install -e . pytest pytest-asyncio pytest-mock - - name: Install compatibility dependencies - run: | - python scripts/compat_matrix.py deps \ + python scripts/compat_matrix.py verify \ --integration "${{ matrix.integration }}" \ --version-set "${{ matrix.version-set }}" \ - --python-version "3.14" > compat-requirements.txt - pip install -r compat-requirements.txt - - name: Run compatibility test - run: python -m pytest "tests/compat/test_${{ matrix.integration }}_compat.py" -q + --python-version "${{ matrix.python-version }}" + - name: Cache downloaded test models + if: matrix.integration == 'llama_cpp' || matrix.integration == 'faster_whisper' + uses: actions/cache@v4 + with: + path: | + tests/compat/.cache + ~/.cache/huggingface/hub + key: compat-models-${{ matrix.integration }}-${{ hashFiles(format('tests/compat/test_{0}_compat.py', matrix.integration)) }} + - name: Run compatibility tests + run: python -m pytest "tests/compat/test_${{ matrix.integration }}_compat.py" -q -rxX diff --git a/.gitignore b/.gitignore index 4dcd2c5..9c6e1cc 100644 --- a/.gitignore +++ b/.gitignore @@ -218,3 +218,6 @@ internal/ # macOS .DS_Store + +# Models downloaded by the compat tests +tests/compat/.cache/ diff --git a/CHANGELOG.md b/CHANGELOG.md index bce4574..4d7b234 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,144 @@ under Unreleased and move into a version section at release time. ## Unreleased +## 0.3.0 - 2026-10-24 + +### Added + +- `mlx_whisper` and `faster_whisper` integrations track local Whisper model + loads and transcriptions, including faster-whisper's + `BatchedInferencePipeline`. Transcriptions record duration, language, word + count and audio metadata. Audio and transcript attachments must be supplied + explicitly through the manual tracking API. See docs/whisper.md. +- The OpenAI integration now tracks `responses.create()` for synchronous and + asynchronous clients, including streaming. Records input/output token usage, + cached and reasoning tokens, response status or incomplete reason, resolved + model metadata, and time to first output delta for streams. Older OpenAI SDKs + without Responses remain supported. +- OpenAI sync and async `embeddings.create()` instrumentation records latency, + input token usage, batch size, vector dimensions, model metadata and errors. + Embedding vectors and raw inputs are never attached by this adapter. + `EmbeddingOutputMeta` now accepts optional `tokens_in` for analytics. +- Anthropic native `messages.stream()` helpers are now instrumented for sync + and async clients, including event iteration, `text_stream`, and final-message + helpers, with the shared streaming lifecycle contract. +- Process-wide integrations and hub patches now support sequential SDK client + replacement after the previous owner closes. Reinitialization replaces a + closed default. Existing models and provider clients keep their original + destination and must be recreated. Attempts to use a closed telemetry + destination warn once instead of silently dropping inference and error + events. + +### Changed + +- The llama.cpp Python integration is now named `llama_cpp` in Python, CLI, + and environment configuration. `gguf` remains a compatibility alias; `all` + selects the integration once under its new name. Model format metadata + remains `gguf`. +- `wildedge.init(dsn=...)` reuses the existing default client when the DSN + points to the same project, for example `dsn=os.environ["WILDEDGE_DSN"]` + under `wildedge run`. It used to create a second client. +- Enabling an integration or hub on a second client while another live client + owns process-wide instrumentation now logs a warning. Automatic events keep + going to the owner's project, as before, but this used to happen without a + word. With `wildedge run --strict-integrations` it raises `RuntimeError`. +- Quantization labels are defined once in `wildedge.quantizations` and used + by every integration. Behaviour changes: + - MLX-LM reports N-bit weights as `int4`/`int8`, like Transformers, instead + of `q4`/`q8`. + - GGUF models labelled from `general.file_type` report the type name (for + example `q4_k_m`) instead of its number, and the `i8` filename suffix + reports `int8`. + - faster-whisper reports weight precision (`int8`, `f16`) instead of + CTranslate2 compute types such as `int8_float32`. +- The ultralytics integration now requires ultralytics 8.0.200 or later + (was 8.0.0). Earlier 8.0.x releases keep their helpers under + `ultralytics.yolo.utils`, and the integration turned itself off on them + without a word. `instrument("ultralytics")` now logs a warning naming the + installed version instead. + +### Fixed + +- Every error event now caps `error_message` at 500 characters, keeping the + start and the end of a longer message. Integrations used to cut it to the + first 200, which lost the cause in messages that end with it, such as + mlx-whisper audio errors that start with the ffmpeg banner. Manual + `track_error()` calls had no limit at all. `provider_error_code` is capped + at 64 characters for every event. +- The llama.cpp integration now tracks `Llama.create_chat_completion()` and + `Llama.create_completion()`, streaming or not. It used to patch only + `Llama.__call__`, so chat apps on llama-cpp-python reported model loads but + no inference events. It now patches `create_completion`, which every entry + point runs, so each call is tracked once. +- llama.cpp inference events now include the stop reason for non-streamed + calls, the generation config (`temperature`, `top_p`, `top_k`, `max_tokens`, + `repeat_penalty`, penalties, `seed`, stop sequence count), and, for chat + calls, `prompt_type="chat"` and `turn_index` measured on the last user + message. Prompts, chat messages and replies are never attached. As before, + attachments upload only when passed to `track_inference(attachments=[...])`. +- `import wildedge` no longer imports every installed ML framework and API + client. Each library is imported when its integration is activated or a + model from it is registered, so SDK import is faster and a broken unrelated + package no longer stops the SDK from starting. +- Changing a model's metadata (for example `register_model(..., version="v2")`) + now takes effect after a restart. With offline persistence on (the default), + the model registry saved by an earlier process kept the first metadata ever + registered for each model on that machine, so later changes were ignored. + This includes anthropic models already saved as `"unknown"`. +- MLX-LM models now report the device they run on (`mps` for the Apple GPU) + instead of always `cpu`. Accelerator names are defined once in + `wildedge.accelerators`, and PyTorch, Transformers, Ultralytics, + TensorFlow and Keras share their device detection. +- timm models now use the requested architecture/variant name for model identity + instead of the Python class name, so models such as resnet18 and resnet50 no + longer share telemetry. Manual registration falls back to timm architecture + metadata, and explicit model IDs retain priority. +- Ultralytics `model.predict()` now records inference telemetry. Instrumentation + patches `predict`, which `model(...)` delegates to, so both entry points are + tracked exactly once. +- Streaming wrappers now finalize once across exhaustion, explicit closure, + cancellation, and iterator errors. Early closure/cancellation records an + unsuccessful inference with partial metadata when available. TTFT ignores + lifecycle/role-only chunks; telemetry callbacks cannot break iteration and + retain the trace context captured at stream creation. Remote API streams + ending without their completion marker record `STREAM_INCOMPLETE`. +- A span attribute whose value is a list, dict or other non-scalar no longer + gets the whole event batch rejected. Ingest only accepts scalar attribute + values and answered 422 for the batch, dropping unrelated loads, inferences + and spans. Such values are now sent as JSON strings, with one warning per + key. +- `close(timeout=...)` and `flush(timeout=...)` now upload buffered + attachments in the time left after sending events. `close()` used to stop + the attachment uploader without a final upload, so a short-lived process + lost the attachments of its last few seconds (the uploader runs every 5 s) + while its events still referenced them. `close()` without a timeout still + leaves them on disk for the next start. +- ONNX models exported with the default `main_graph` graph name (newer + `torch.onnx.export`, Piper TTS voices) are now named after their file, so + different voices no longer show up as one model called `main_graph`. A + session created without an `id` falls back to the file name too, and + `model_version` 0, negative or INT64_MAX (the exporters' "unset") is sent as + `unknown` instead of `9223372036854775807`. +- transformers pipelines called with a string or a list record one inference + event per call instead of two. `pipeline()` loads the model before + `Pipeline.__init__` runs, so the model was registered with its own forward + hooks, and each pipeline call was tracked by the pipeline and again by the + model inside it. That model's hooks now skip forward passes inside a + tracked call of its pipeline; other models the pipeline runs, and direct + calls to the model, are still tracked. Generator and dataset inputs, which + run lazily, are not covered yet (#78). +- `flush()` no longer sends a batch the background sender is already sending. + Both used to send the same events and then each removed a batch from the + queue, so events were duplicated and, depending on timing, later events + could be removed without being sent. This showed up when a process started + with events saved by an earlier run and flushed soon after. +- Enabling the `huggingface` hub no longer imports unused transformers model + modules. To find references to `hf_hub_download` and `snapshot_download`, + the hub tracker looked up the name on every loaded module, and transformers' + lazy modules import code on lookup. With transformers loaded this imported + about 140 extra modules, took about 0.15 s and logged about 430 "failed to + patch" debug lines. The tracker now reads each module's namespace instead. + ## 0.2.2 - 2026-09-30 ### Added @@ -38,8 +176,6 @@ under Unreleased and move into a version section at release time. - Models tracked by the anthropic integration register with `model_format` `"api"`, like the openai integration. The Anthropic extractor was never in the client's extractor list, so they were registered as `"unknown"`. - Machines that already tracked these models keep the saved `"unknown"` until - the SDK's model registry file is deleted. ## 0.2.1 - 2026-09-01 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index baa954d..f01f4fa 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -30,6 +30,9 @@ uv run ruff check . uv run ruff format . ``` + CI runs both (`ruff check` and `ruff format --check`) and fails on any finding. + The rule set is in `pyproject.toml`; prefer fixing over `# noqa`, and give a + reason when a `noqa` is needed. - Follow PEP 8 and PEP 484 (type hints). @@ -46,7 +49,7 @@ Releases follow a `release/` branching flow: 1. Cut a `release/` branch from `devel`. -2. Bump the version in `pyproject.toml`. +2. Bump the version in `pyproject.toml` and move the Unreleased entries in `CHANGELOG.md` into a `## - ` section. The release notes are that section, so the release fails without it. 3. Open a pull request targeting `main`. CI will automatically run the full compatibility matrix and post a changelog preview comment. 4. Once merged, tag the commit as `v`. The release workflow publishes to PyPI and creates a GitHub release. diff --git a/README.md b/README.md index 3d5c3f5..03c63e0 100644 --- a/README.md +++ b/README.md @@ -9,9 +9,33 @@ [![Tested on Linux](https://img.shields.io/badge/tested%20on-linux-blue)](https://github.com/wild-edge/wildedge-python/actions/workflows/ci.yml) [![Tested on macOS](https://img.shields.io/badge/tested%20on-macOS-blue)](https://github.com/wild-edge/wildedge-python/actions/workflows/ci.yml) [![Tested on Windows](https://img.shields.io/badge/tested%20on-windows-blue)](https://github.com/wild-edge/wildedge-python/actions/workflows/ci.yml) +[![PyPI](https://img.shields.io/pypi/v/wildedge-sdk)](https://pypi.org/project/wildedge-sdk/) [![Coverage](https://codecov.io/gh/wild-edge/wildedge-python/branch/main/graph/badge.svg)](https://codecov.io/gh/wild-edge/wildedge-python) -ML inference monitoring for Python: on-device and remote models. Tracks latency, tokens, errors, and model metadata with no code changes, plus traces for agent pipelines. +Monitor local and remote ML inference from Python. Track latency, token usage, +errors, model metadata, and local hardware context across llama.cpp, MLX, +Transformers, OpenAI, and more. + +Instrument supported libraries with one command, or add tracking in your code. +Prompts and outputs aren't uploaded unless you explicitly attach them. + +**[Get started](#install)** · [Explore the demo](https://wildedge.dev/demo) · [Docs](https://docs.wildedge.dev) · [Website](https://wildedge.dev) + +| Model performance | Workflow trace | SQL editor | +|:---:|:---:|:---:| +| Model health, latency, time to first token, and throughput | Named workflow step and inference timing in a trace waterfall | SQL query with stacked hourly inference volumes across local and cloud models | + +*Illustrative sample data across WildEdge SDKs. Click a thumbnail for the full-size image, +or [explore the read-only demo](https://wildedge.dev/demo).* + +
+More screenshots: events and hardware + +| Events | Hardware | +|:---:|:---:| +| Inference and error events with model and device metadata | Thermal state, accelerator usage, and CPU throttling | + +
## Install @@ -19,98 +43,175 @@ ML inference monitoring for Python: on-device and remote models. Tracks latency, uv add wildedge-sdk ``` -Without a DSN everything is a silent no-op, so dev and CI need no configuration. +To send events, sign up at [app.wildedge.dev](https://app.wildedge.dev), create a project and copy its DSN. Without a DSN everything is a silent no-op, so dev and CI need no configuration. ## CLI -Drop `wildedge run` in front of your existing command. WildEdge instruments the runtime before your code starts. No SDK calls required in user code. +Monitor an existing application by adding `wildedge run` to its start command: ```bash WILDEDGE_DSN="https://@ingest.wildedge.dev/" \ -wildedge run --integrations timm -- python app.py +uv run wildedge run --integrations llama_cpp -- python app.py ``` -Validate your environment before deploying. `--send-test-event` proves the -whole pipeline end to end by sending one span event and reporting the ingest -response (exit codes: 0 pass, 1 config failure, 2 connectivity failure): +This assumes your application already uses `llama-cpp-python`. The `llama_cpp` +integration instruments that library; replace it with your +[integration](#supported-integrations). +No SDK calls are required in your application code. + +See the [deployment guide](docs/deployment.md) for startup options, +environment validation, and production configuration. + +## SDK + +To try tracking in your code, run one local text-generation model on CPU. +Install the model's dependencies: ```bash -wildedge doctor --integrations all --send-test-event +uv add transformers torch ``` -Useful flags: +Save this as `app.py`. The first run downloads the model from Hugging Face; +no model-provider API key is needed. -| Flag | Description | -|---|---| -| `--integrations` | Comma-separated list of integrations to activate (or `all`) | -| `--hubs` | Hub trackers to activate: `huggingface`, `torchhub` | -| `--print-startup-report` | Print per-integration status at startup | -| `--strict-integrations` | Exit (code 121) if a requested integration can't be instrumented | -| `--strict` | Exit (120 config, 122 internal) instead of running untracked when bootstrap fails | -| `--attachments` | Enable opt-in raw input/output attachment upload | -| `--no-propagate` | Don't pass WildEdge env vars to child processes | +```python +import wildedge +from transformers import pipeline -## SDK +client = wildedge.init(integrations=["transformers"]) +try: + model = pipeline( + "text-generation", + model="HuggingFaceTB/SmolLM2-360M-Instruct", + device=-1, # CPU + ) + result = model( + "Explain how rain forms in one sentence.", + max_new_tokens=32, + do_sample=False, + return_full_text=False, + ) + print(result[0]["generated_text"]) +finally: + client.close(timeout=10) +``` + +Set the DSN from your WildEdge project and run it: + +```bash +export WILDEDGE_DSN="https://@ingest.wildedge.dev/" +uv run python app.py +``` + +Open your project at [app.wildedge.dev](https://app.wildedge.dev) to inspect the +model load and inference events, including duration and model metadata. The +generated text is printed locally; it isn't uploaded. + +Already have an application? Initialize WildEdge before creating your model or +API client. See [examples for each integration](#supported-integrations), +[traces and spans](docs/manual-tracking.md#track-spans-for-agentic-workflows), +and the [deployment guide](docs/deployment.md) for lifecycle details. -Use the SDK when you can't wrap the process, or when you want traces and spans -around your own code. `wildedge run` covers everything else. +
+Trace a workflow across local and remote models + +Combine integrations and group inference calls into traces and spans. This +example also requires `uv add openai` and an `OPENAI_API_KEY`, alongside the +Transformers dependencies and WildEdge DSN above. ```python import wildedge from openai import OpenAI from transformers import pipeline -wildedge.init(integrations=["transformers", "openai"]) # optional under `wildedge run` - -local = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-360M-Instruct") -remote = OpenAI() - -with wildedge.trace(run_id="run-1"): - with wildedge.span(kind="agent_step", name="draft"): - draft = local(prompt)[0]["generated_text"] - with wildedge.span(kind="agent_step", name="refine"): - remote.chat.completions.create( - model="gpt-4o-mini", - messages=[{"role": "user", "content": draft}], - ) +client = wildedge.init(integrations=["transformers", "openai"]) +try: + local = pipeline( + "text-generation", + model="HuggingFaceTB/SmolLM2-360M-Instruct", + device=-1, + ) + remote = OpenAI() + prompt = "Explain how rain forms in one sentence." + + with wildedge.trace(run_id="run-1"): + with wildedge.span(kind="agent_step", name="draft"): + draft = local( + prompt, max_new_tokens=32, return_full_text=False + )[0]["generated_text"] + with wildedge.span(kind="agent_step", name="refine"): + response = remote.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": draft}], + ) + print(response.choices[0].message.content) +finally: + client.close(timeout=10) ``` -See [Deployment](https://github.com/wild-edge/wildedge-python/blob/main/docs/deployment.md) -for client lifecycle and the full contract. +Both inference calls belong to one trace, with separate `draft` and `refine` +spans. See [manual tracking](docs/manual-tracking.md#track-spans-for-agentic-workflows) +for additional span types and context. + +
## Supported integrations -**On-device** +**Local inference:** llama.cpp, Transformers, MLX-LM, ONNX Runtime, PyTorch, +TensorFlow, Keras, Ultralytics, timm, faster-whisper, and MLX Whisper. -| Integration | Example | -|---|---| -| `transformers` | [transformers_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/transformers_example.py) | -| `mlx` | [mlx_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/mlx_example.py) | -| `timm` | [timm_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/timm_example.py) | -| `gguf` | [gguf_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/gguf_example.py) | -| `onnx` | [onnx_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/onnx_example.py) | -| `ultralytics` | - | -| `tensorflow` | [tensorflow_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/tensorflow_example.py) | -| `torch` | [pytorch_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/pytorch_example.py) | -| `keras` | [keras_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/keras_example.py) | - -**Remote models** - -| Integration | Example | -|---|---| -| `anthropic` | [anthropic_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/anthropic_example.py) | -| `openai` | [openai_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/openai_example.py) | +**API clients:** OpenAI, Anthropic, and OpenAI-compatible providers. + +Most integrations track supported entry points automatically after activation. +PyTorch and standalone Keras models require explicit registration. + +
+Integration keys, setup, and examples + +Enable the integration before creating models or API clients, using +`--integrations` or `wildedge.init(integrations=[...])`. Automatic tracking +covers the entry points listed below. PyTorch and standalone Keras models need +explicit registration through `client.load(...)` or `client.register_model(...)`. + +**Local models** + +| Library | Integration key | How tracking starts | Example | +|---|---|---|---| +| Transformers | `transformers` | Automatically through `pipeline` / `from_pretrained` | [Example](examples/transformers_example.py) | +| llama.cpp (`llama-cpp-python`) | `llama_cpp` | Automatically when constructing `Llama` | [Example](examples/gguf_example.py) | +| MLX-LM | `mlx` | Automatically through `mlx_lm.load` | [Example](examples/mlx_example.py) | +| faster-whisper / MLX Whisper | `faster_whisper` / `mlx_whisper` | Automatically through supported model loaders | [Example](examples/whisper_example.py) | +| ONNX Runtime | `onnx` | Automatically when constructing `InferenceSession` | [Example](examples/onnx_example.py) | +| timm | `timm` | Automatically through `create_model` | [Example](examples/timm_example.py) | +| Ultralytics | `ultralytics` | Automatically when constructing `YOLO` | — | +| TensorFlow | `tensorflow` | Automatically through `keras.models.load_model` / `saved_model.load`; register models created in code | [Example](examples/tensorflow_example.py) | +| PyTorch | `torch` | Register the model explicitly | [Example](examples/pytorch_example.py) | +| Keras | `keras` | Register the model explicitly | [Example](examples/keras_example.py) | + +**API clients** + +| Library | Integration key | How tracking starts | Example | +|---|---|---|---| +| Anthropic | `anthropic` | Automatically for clients created after initialization | [Example](examples/anthropic_example.py) | +| OpenAI | `openai` | Automatically for clients created after initialization | [Example](examples/openai_example.py) | + +
+ +
+API coverage and compatible providers Fireworks, Together, Baseten, xAI, Mistral, Groq and other OpenAI-compatible hosts work through the `openai` integration — pass their `base_url`. See [Providers](https://github.com/wild-edge/wildedge-python/blob/main/docs/providers.md) for every recognized endpoint. -Using httpx, requests or urllib directly instead of the `openai` / `anthropic` -client libraries? Wrap the call in -[`wildedge.llm_api()`](https://github.com/wild-edge/wildedge-python/blob/main/docs/llm_api.md) -to get the same events: -[llm_api_example.py](https://github.com/wild-edge/wildedge-python/blob/main/examples/llm_api_example.py). +See [API tracking](docs/providers.md) for endpoint coverage, streaming, embeddings, +and tracking calls made without the official client libraries. + +
+ +
+Model download tracking **Hub tracking** @@ -121,6 +222,8 @@ Pass `hubs=` to track model download provenance. Hubs are framework-agnostic and | `huggingface` | Downloads via `huggingface_hub` | | `torchhub` | Downloads via `torch.hub` | +
+ For unsupported frameworks, see [Manual tracking](https://github.com/wild-edge/wildedge-python/blob/main/docs/manual-tracking.md). ## Configuration @@ -136,22 +239,13 @@ For unsupported frameworks, see [Manual tracking](https://github.com/wild-edge/w For advanced options (batching, queue tuning, dead-letter storage, attachments), see [Configuration](https://github.com/wild-edge/wildedge-python/blob/main/docs/configuration.md). -## Projects using this SDK - -| Name | Link | -|---|---| -| outfitstudio.app | https://outfitstudio.app/ | -| agntr | [github.com/pmaciolek/agntr](https://github.com/pmaciolek/agntr) | -| demo-app | [github.com/wild-edge/demo-app](https://github.com/wild-edge/demo-app) | -| *(your project here)* | - | - -Using WildEdge in your project? Open a PR to add it to the list. - ## Security & Privacy -By default the SDK transmits only anonymized telemetry, never raw model inputs -or outputs. The one exception is opt-in [attachments](https://github.com/wild-edge/wildedge-python/blob/main/docs/configuration.md#attachments) -(`attachments_enabled`), which upload raw bytes you explicitly pass in. +Automatic integrations send inference metrics, model metadata, and device/hardware +information, including a device identifier, OS information, locale, and timezone. +They do not attach raw inputs or outputs. Raw payload uploads require explicitly +supplied [attachments](docs/configuration.md#attachments) and +`attachments_enabled=True`. Report security and privacy issues to: support@wildedge.dev @@ -161,8 +255,6 @@ Report security and privacy issues to: support@wildedge.dev - [Manual tracking](https://github.com/wild-edge/wildedge-python/blob/main/docs/manual-tracking.md) - [LLM API tracking](https://github.com/wild-edge/wildedge-python/blob/main/docs/llm_api.md) - [Compatibility Matrix](https://github.com/wild-edge/wildedge-python/blob/main/docs/compatibility.md) +- [Streaming lifecycle](docs/streaming.md) - [Changelog](https://github.com/wild-edge/wildedge-python/blob/main/CHANGELOG.md) - [License](https://github.com/wild-edge/wildedge-python/blob/main/LICENSE) - -Each GitHub release ships `llms.txt` and `llms-full.txt`: the full -documentation for that exact version in one file, built for AI assistants. diff --git a/docs/compatibility.md b/docs/compatibility.md index 9f417fa..e5aec96 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -1,23 +1,71 @@ ## Compatibility Matrix -The table below is generated from `scripts/compat_matrix.py`. +Each integration is tested against the pinned library versions below: the +oldest supported release (`min`) and a recent one (`current`). The table is +generated from `scripts/compat_matrix.py`, which also defines the CI jobs. + +### What the tests check + +`tests/compat/test__compat.py` loads a real model with the real +library and checks the events the SDK records: a `model_load` event, then an +`inference` event for the same model with the expected fields (token counts +for LLMs, input modality for vision models). Apart from one download for llama_cpp, nothing +touches the network: + +- openai and anthropic get canned responses through the SDK's own HTTP + client (`MockTransport`). +- onnx loads a tiny model committed in `tests/compat/fixtures`. llama_cpp + downloads a 1.2 MB model on first use, checked against a pinned sha256. +- torch, timm, tensorflow, transformers, ultralytics and mlx build models from + a config, with random weights. + +A test never skips. A library that fails to install or import fails the job, +and a step checks that the installed versions match the pins. A known SDK +bug is marked `xfail(strict=True)` with its issue number, so the job fails +once the bug is fixed and the marker has to come off. + +### When they run + +On pull requests into `main` or a release branch, on pushes to them, nightly +against `devel`, and on manual runs. Every run covers all the combinations +below. Python 3.14 rows are canaries and don't fail the run. mlx runs on +macOS. + +To run the full matrix locally: `python scripts/run_compat_local.py`. +`python scripts/compat_matrix.py table` prints the table below. + +### Matrix | Integration | Version set | Dependencies | Supported Python | |---|---|---|---| | `anthropic` | `min` | `anthropic==0.25.0` | 3.10, 3.11, 3.12, 3.13 | -| `anthropic` | `current` | `anthropic==0.92.0` | 3.10, 3.11, 3.12, 3.13 | +| `anthropic` | `current` | `anthropic==1.15.0` | 3.10, 3.11, 3.12, 3.13 | | `onnx` | `min` | `onnxruntime==1.18.1`, `numpy==1.26.4` | 3.10, 3.11, 3.12, 3.13, 3.14 | | `onnx` | `min` (`py3.13` override) | `onnxruntime==1.20.1`, `numpy==2.1.3` | 3.13 | -| `onnx` | `current` | `onnxruntime==1.20.1`, `numpy==2.1.3` | 3.10, 3.11, 3.12, 3.13, 3.14 | +| `onnx` | `current` | `onnxruntime==1.31.0`, `numpy==2.2.6` | 3.10, 3.11, 3.12, 3.13, 3.14 | +| `onnx` | `current` (`py3.10` override) | `onnxruntime==1.23.2`, `numpy==2.2.6` | 3.10 | | `torch` | `min` | `torch==2.4.1`, `numpy==1.26.4` | 3.10, 3.11, 3.12, 3.13, 3.14 | | `torch` | `min` (`py3.13` override) | `torch==2.6.0`, `numpy==2.1.3` | 3.13 | | `torch` | `min` (`py3.14` override) | `torch==2.5.0`, `numpy==2.1.3` | 3.14 | -| `torch` | `current` | `torch==2.6.0`, `numpy==2.1.3` | 3.10, 3.11, 3.12, 3.13, 3.14 | -| `torch` | `current` (`py3.14` override) | `torch==2.10.0`, `numpy==2.1.3` | 3.14 | +| `torch` | `current` | `torch==2.14.1`, `numpy==2.2.6` | 3.10, 3.11, 3.12, 3.13, 3.14 | +| `torch` | `current` (`py3.14` override) | `torch==2.14.1`, `numpy==2.5.4` | 3.14 | | `timm` | `min` | `torch==2.4.1`, `torchvision==0.19.1`, `timm==1.0.11`, `numpy==1.26.4` | 3.10, 3.11, 3.12, 3.13, 3.14 | | `timm` | `min` (`py3.13` override) | `torch==2.6.0`, `torchvision==0.21.0`, `timm==1.0.11`, `numpy==2.1.3` | 3.13 | | `timm` | `min` (`py3.14` override) | `torch==2.5.0`, `torchvision==0.19.1`, `timm==1.0.11`, `numpy==2.1.3` | 3.14 | -| `timm` | `current` | `torch==2.6.0`, `torchvision==0.21.0`, `timm==1.0.15`, `numpy==2.1.3` | 3.10, 3.11, 3.12, 3.13, 3.14 | -| `timm` | `current` (`py3.14` override) | `torch==2.10.0`, `torchvision==0.25.0`, `timm==1.0.15`, `numpy==2.1.3` | 3.14 | +| `timm` | `current` | `torch==2.14.1`, `torchvision==0.29.1`, `timm==1.0.30`, `numpy==2.2.6` | 3.10, 3.11, 3.12, 3.13, 3.14 | +| `timm` | `current` (`py3.14` override) | `torch==2.14.1`, `torchvision==0.29.1`, `timm==1.0.30`, `numpy==2.5.4` | 3.14 | | `tensorflow` | `min` | `tensorflow==2.16.1`, `keras==3.3.3`, `numpy==1.26.4` | 3.10, 3.11, 3.12 | -| `tensorflow` | `current` | `tensorflow==2.18.0`, `keras==3.8.0`, `numpy==2.0.2` | 3.10, 3.11, 3.12 | +| `tensorflow` | `current` | `tensorflow==2.21.0`, `keras==3.15.1`, `numpy==2.2.6` | 3.10, 3.11, 3.12 | +| `tensorflow` | `current` (`py3.10` override) | `tensorflow==2.21.0`, `keras==3.12.4`, `numpy==2.2.6` | 3.10 | +| `llama_cpp` | `min` | `llama-cpp-python==0.2.90`, `numpy==1.26.4` | 3.10, 3.11, 3.12, 3.13 | +| `llama_cpp` | `current` | `llama-cpp-python==0.3.36`, `numpy==2.2.6` | 3.10, 3.11, 3.12, 3.13 | +| `openai` | `min` | `openai==1.30.0`, `httpx==0.27.2` | 3.10, 3.11, 3.12, 3.13 | +| `openai` | `current` | `openai==3.35.1` | 3.10, 3.11, 3.12, 3.13 | +| `transformers` | `min` | `transformers==4.40.0`, `torch==2.4.1`, `numpy==1.26.4` | 3.10, 3.11, 3.12, 3.13 | +| `transformers` | `min` (`py3.13` override) | `transformers==4.45.0`, `torch==2.5.0`, `numpy==2.1.3` | 3.13 | +| `transformers` | `current` | `transformers==5.20.1`, `torch==2.14.1`, `numpy==2.2.6` | 3.10, 3.11, 3.12, 3.13 | +| `ultralytics` | `min` | `ultralytics==8.0.200`, `numpy==1.26.4` | 3.10, 3.11, 3.12, 3.13 | +| `ultralytics` | `current` | `ultralytics==8.4.182`, `numpy==2.2.6` | 3.10, 3.11, 3.12, 3.13 | +| `mlx` | `current` | `mlx==0.32.4`, `mlx-lm==0.32.2` | 3.12, 3.13 | +| `faster_whisper` | `current` | `faster-whisper==1.2.1`, `ctranslate2==4.8.2`, `numpy==2.2.6` | 3.12, 3.13 | +| `mlx_whisper` | `current` | `mlx-whisper==0.4.3`, `mlx==0.32.3`, `numpy==2.2.6` | 3.12, 3.13 | diff --git a/docs/configuration.md b/docs/configuration.md index 2487ebe..02f6f46 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -41,4 +41,17 @@ event flushes never block on an upload. | `attachment_filter` | `None` | - | `Callable[[list[Attachment]], list[Attachment]]` to redact/drop before buffering | | `attachment_dir` | per `app_identity` | - | Override the local buffer directory | +`attachments_enabled=True` permits uploads of attachments explicitly supplied to +`track_inference(attachments=[...])`. Automatic integrations do not attach raw +inputs or outputs. + See [Track attachments](manual-tracking.md#track-attachments) for usage. + +### When ML libraries are imported + +`import wildedge` does not import any ML framework or API client. The SDK imports a library only when you use it: + +- activating an integration (`wildedge.init(integrations=[...])`, `client.instrument(...)` or `wildedge run`) imports that library; +- registering a model imports the library the model came from, and nothing else. For example, registering a PyTorch model does not import timm (unless it is a 1000-class classifier, for which timm supplies the ImageNet labels), and registering a Keras model does not import TensorFlow. + +Which integrations are activated by default is unchanged. `wildedge.init()` activates only the ones you list, while `wildedge run` activates every installed one unless you pass `--integrations`. diff --git a/docs/deployment.md b/docs/deployment.md index 713c016..d7f1a1f 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -16,8 +16,9 @@ There are three ways a process gets its client, in order of coverage: All three produce the same thing: one process-wide default client that `wildedge.span()`, `wildedge.register_model()` and the rest delegate to. -`init()` without `dsn` reuses whatever client already exists, so code written -for in-process init runs unchanged under `wildedge run`. +`init()` without `dsn`, or with the DSN that client already uses, reuses +whatever client already exists, so code written for in-process init runs +unchanged under `wildedge run`. The lazy mode configures itself from the environment: `WILDEDGE_DSN`, `WILDEDGE_INTEGRATIONS`, `WILDEDGE_HUBS`. Unset integration variables mean @@ -35,6 +36,22 @@ Under `wildedge run`, a missing DSN warns on stderr and your program runs untracked. Telemetry failing must not take production down; that is the default policy. +## CLI options + +| Flag | Description | +|---|---| +| `--integrations` | Comma-separated runtime integrations, such as `llama_cpp`, `onnx`, or `openai`; default: `all` | +| `--hubs` | Download tracking via `huggingface` or `torchhub` | +| `--print-startup-report` | Print activation status for each integration | +| `--strict-integrations` | Exit with 121 if a requested integration cannot be instrumented | +| `--strict` | Exit on bootstrap failure; see the exit codes below | +| `--attachments` | Enable upload of explicitly supplied attachments | +| `--no-propagate` | Scrub run-scoped environment variables after initialization | + +Use `llama_cpp` for the `llama-cpp-python` runtime. The older name `gguf` is +accepted as a compatibility alias in Python, CLI, and environment settings. +The recorded model format remains `gguf`. + ## Strict mode `wildedge run --strict` inverts the failure policy for deployments where @@ -82,16 +99,53 @@ them unwrapped or unset `WILDEDGE_DSN` there. ## One client per process -- Framework patches are installed at most once per process; the first client - to instrument owns the patch. -- `init()` reuses the existing default client unless you pass `dsn` - explicitly, so the CLI-installed client and application code share one - client instead of racing. +- Framework patches are installed at most once per process. One live client + owns process-wide auto-instrumentation. When another client enables an + integration or hub while that owner is live, it logs a warning and + automatic events keep going to the owner's project. The new client still + sends its manual tracking. With `wildedge run --strict-integrations`, the + conflict raises `RuntimeError` instead. Manual tracking does not claim + ownership. +- `init()` reuses the live default client unless you pass a `dsn` for a + different project, so the CLI-installed client and application code share + one client instead of racing. - Trace and span correlation lives in contextvars at module level, not on a client. Auto-instrumented events emitted inside `wildedge.trace(...)` / `wildedge.span(...)` blocks correlate into the same trace no matter which client emits them. +### Replacing a client + +Sequential replacement is supported. Finish any in-flight work and close the +old SDK client before initializing its replacement. Enable integrations and hubs +explicitly on the replacement: previously enabled integrations do not carry over. +Factory patches resolve the new owner without being installed again. Buffered +HuggingFace download provenance is scoped to its owner session and does not +carry over to another project. + +```python +first = wildedge.init(dsn=dsn, integrations=["onnx"]) +session = ort.InferenceSession("model.onnx") +# Run and finish work with session. +first.close(timeout=10) + +second = wildedge.init(dsn=other_dsn, integrations=["onnx"]) +session = ort.InferenceSession("model.onnx") # create after replacement init +``` + +Existing model objects, provider clients, model handles and in-flight streams +remain bound to the original client; they are not automatically transferred to +another project. Recreate auto-instrumented models/provider clients after +initializing the replacement. Attempts to emit inference/error telemetry through +a closed client are dropped with a one-time warning; application inference can +still run. Simply changing the default client does not change instrumentation +ownership. Repeated `close()` calls are harmless. + +`init()` replaces a closed default even when no explicit `dsn` is supplied, +reading the new client's configuration from the environment. If initialization +fails (including an ownership conflict), a newly created client is closed and +the existing default is preserved. + ## Verifying a deployment `wildedge doctor` answers "will events actually flow from this machine" diff --git a/docs/images/README.md b/docs/images/README.md new file mode 100644 index 0000000..c07d95d --- /dev/null +++ b/docs/images/README.md @@ -0,0 +1,33 @@ +# README screenshots + +`demo-model-performance.jpg` is an unaltered 1600 × 1100 desktop capture from the public +WildEdge demo on 2026-10-21. It shows sample data, not measurements from the +Python quickstart. + +To refresh: open https://wildedge.dev/demo, switch to News Reader, then open +Models → Article Summarizer. Use a 1600 × 1100 viewport so the sidebar fits without a scrollbar. Capture +the full viewport, including the model heading, health notice, performance +cards, and charts. Keep metric values as displayed; do not fabricate data. + +The README links to the demo entry point because direct navigation to nested +model routes currently loads the marketing page. + +Additional views in the News Reader demo: + +- `demo-events.jpg`: Events. +- `demo-hardware.jpg`: Hardware. + +These sample views span WildEdge SDKs, including mobile hardware. They do not +represent output from the Python quickstart. The README uses 180-pixel-wide +thumbnails linked to the full-size images, with a separate live-demo link. + +- `demo-trace.jpg`: News Reader → Runs → demo-run-1. Shows a named step and + a local summarizer inference; this sample does not include a cloud call. +- `demo-sql.jpg`: a local screenshot-only composition of the dashboard's SQL + editor and chart result components. It shows query results computed with DuckDB + from generated fictional events, not the deployed offline demo or production + telemetry. See [sql-preview/README.md](sql-preview/README.md) to reproduce it. + +The opening row contains performance, trace, and SQL previews. Events and +hardware remain in a collapsed More screenshots section. All five thumbnails +are 180 pixels wide and open their full-size image. diff --git a/docs/images/demo-events.jpg b/docs/images/demo-events.jpg new file mode 100644 index 0000000..5e4df1e Binary files /dev/null and b/docs/images/demo-events.jpg differ diff --git a/docs/images/demo-hardware.jpg b/docs/images/demo-hardware.jpg new file mode 100644 index 0000000..5e69747 Binary files /dev/null and b/docs/images/demo-hardware.jpg differ diff --git a/docs/images/demo-model-performance.jpg b/docs/images/demo-model-performance.jpg new file mode 100644 index 0000000..a8a22c9 Binary files /dev/null and b/docs/images/demo-model-performance.jpg differ diff --git a/docs/images/demo-sql.jpg b/docs/images/demo-sql.jpg new file mode 100644 index 0000000..2aff736 Binary files /dev/null and b/docs/images/demo-sql.jpg differ diff --git a/docs/images/demo-trace.jpg b/docs/images/demo-trace.jpg new file mode 100644 index 0000000..a95e27f Binary files /dev/null and b/docs/images/demo-trace.jpg differ diff --git a/docs/images/sql-preview/DemoAnalytics.vue b/docs/images/sql-preview/DemoAnalytics.vue new file mode 100644 index 0000000..d69bd34 --- /dev/null +++ b/docs/images/sql-preview/DemoAnalytics.vue @@ -0,0 +1,40 @@ + + + diff --git a/docs/images/sql-preview/README.md b/docs/images/sql-preview/README.md new file mode 100644 index 0000000..ea06efc --- /dev/null +++ b/docs/images/sql-preview/README.md @@ -0,0 +1,27 @@ +# Reproduce the SQL screenshot + +This is an illustrative local preview, not a deployed demo feature or production +telemetry. The SQL runs against deterministic fictional events in DuckDB, and +the existing dashboard SqlExplorer and QueryResult components render the result. +No chart pixels or query results are fabricated independently of the query. + +Prerequisites: a wild-edge checkout with dashboard dependencies installed, Node, +and a Python environment containing DuckDB (the platform's `.venv` includes it). + +```bash +/path/to/wild-edge/.venv/bin/python prepare.py /path/to/wild-edge /private/tmp/sql-preview +cd /private/tmp/sql-preview +npm run typecheck +npm run build:demo +MOCKED_BE=true npm run preview -- --outDir dist-demo --host 127.0.0.1 --port 4180 +``` + +Open `http://127.0.0.1:4180/demo/wildedge-demolabs/news-reader/analytics?mode=sql` +at 1600 × 1100 and capture the full viewport. The preview uses the existing +compact chart layout to keep the query, axes, stacked bars, and legend visible. +The source checkout and production site are unchanged. + +`prepare.py` creates 12,876 events spanning 48 hours and three fictional model +workloads. It executes `query.sql` and checks that its 144 grouped rows account +for every generated event. `DemoAnalytics.vue` is a screenshot-only composition +of existing dashboard components; the editable SQL field does not execute queries. diff --git a/docs/images/sql-preview/prepare.py b/docs/images/sql-preview/prepare.py new file mode 100644 index 0000000..0654891 --- /dev/null +++ b/docs/images/sql-preview/prepare.py @@ -0,0 +1,76 @@ +"""Create an isolated, fictional SQL screenshot preview using the actual dashboard. + +Usage: python prepare.py /path/to/wild-edge /private/tmp/wildedge-sql-preview +Requires duckdb in the Python environment and installed dashboard node_modules. +""" + +import json +import math +import shutil +import sys +from datetime import datetime, timedelta +from pathlib import Path + +import duckdb + +source = Path(__file__).resolve().parent +repository = Path(sys.argv[1]).resolve() +destination = Path(sys.argv[2]).resolve() +if destination.exists(): + raise SystemExit( + "Choose a new preview directory; existing files are never replaced." + ) +shutil.copytree( + repository / "dashboard", + destination, + ignore=shutil.ignore_patterns("node_modules", "dist", "dist-demo", ".vite"), +) +(destination / "node_modules").symlink_to( + repository / "dashboard" / "node_modules", target_is_directory=True +) +connection = duckdb.connect() +connection.execute( + "CREATE TABLE events(event_timestamp TIMESTAMP, event_type VARCHAR, model_name VARCHAR)" +) +events = [] +models = ["Gemma 3 4B · local", "Llama 3.2 3B · local", "GPT-4o mini · cloud"] +for hour in range(48): + timestamp = datetime(2026, 10, 19) + timedelta(hours=hour) + daytime = (1 + math.sin((hour % 24 - 7) * math.pi / 12)) / 2 + volumes = [ + round(45 + 175 * daytime + 20 * math.sin(hour * 0.7)), + round(30 + 115 * daytime + 12 * math.cos(hour * 0.9)), + round(15 + 55 * daytime + (35 if 31 <= hour <= 37 else 0)), + ] + for model, count in zip(models, volumes, strict=True): + events.extend( + (timestamp + timedelta(seconds=second * 7), "inference", model) + for second in range(count) + ) +connection.executemany("INSERT INTO events VALUES (?, ?, ?)", events) +query = (source / "query.sql").read_text() +cursor = connection.execute(query) +rows = [ + [value.isoformat() if isinstance(value, datetime) else value for value in row] + for row in cursor.fetchall() +] +assert len(rows) == 48 * 3 +assert sum(row[2] for row in rows) == len(events) +fixture = { + "columns": [description[0] for description in cursor.description], + "rows": rows, + "truncated": False, + "elapsed_ms": 0, +} + +(destination / "src/services/demo/screenshot-result.json").write_text( + json.dumps(fixture) +) +shutil.copyfile( + source / "query.sql", destination / "src/services/demo/screenshot-query.sql" +) +shutil.copyfile( + source / "DemoAnalytics.vue", destination / "src/views/DemoAnalytics.vue" +) +print(f"Computed {len(rows)} grouped rows from {len(events)} fictional events.") +print(f"Preview prepared in {destination}") diff --git a/docs/images/sql-preview/query.sql b/docs/images/sql-preview/query.sql new file mode 100644 index 0000000..7a914da --- /dev/null +++ b/docs/images/sql-preview/query.sql @@ -0,0 +1,10 @@ +SELECT + date_trunc('hour', event_timestamp) AS hour, + model_name, + count(*) AS inferences +FROM events +WHERE event_type = 'inference' + AND event_timestamp >= TIMESTAMP '2026-10-19 00:00:00' + AND event_timestamp < TIMESTAMP '2026-10-21 00:00:00' +GROUP BY 1, 2 +ORDER BY 1, 2; diff --git a/docs/manual-tracking.md b/docs/manual-tracking.md index 9295c6c..0bcc769 100644 --- a/docs/manual-tracking.md +++ b/docs/manual-tracking.md @@ -26,8 +26,8 @@ with wildedge.trace(run_id="run-123"): 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()` without `dsn`, or with the same 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`, ...) @@ -261,14 +261,19 @@ handle.feedback(FeedbackType.THUMBS_DOWN) Opt-in upload of the raw bytes behind an inference, such as the source image, audio clip, or generated text, so you can inspect or curate them later. This is off by default and requires the attachment feature enabled on your project. -Enable it at init, then pass `attachments=` to `track_inference`: +Enable it at init, then pass `attachments=` to `track_inference`. Automatic +integrations, including GGUF and Whisper, do not supply attachments. When tracking +the inference manually, use `auto_instrument=False` at model registration to avoid +recording the same call twice: ```python import wildedge from wildedge import Attachment client = wildedge.init(attachments_enabled=True) -handle = client.register_model(model, model_id="doc-classifier-v1") +handle = client.register_model( + model, model_id="doc-classifier-v1", auto_instrument=False +) handle.track_inference( duration_ms=120, @@ -292,6 +297,35 @@ Bytes are buffered to disk and survive restarts. Capture is gated by [`examples/attachments_example.py`](../examples/attachments_example.py) for the full reference. +### Explicit GGUF prompt and reply attachments + +Use manual tracking when you want to attach GGUF payloads: + +```python +import wildedge +from llama_cpp import Llama +from wildedge import Attachment +from wildedge.timing import Timer + +client = wildedge.init(integrations=[], attachments_enabled=True) +llm = Llama(model_path="/models/model.gguf") +handle = client.register_model(llm, auto_instrument=False) +prompt = "Explain how rain forms." +with Timer() as timer: + result = llm.create_completion(prompt=prompt, max_tokens=128) +reply = result["choices"][0]["text"] +handle.track_inference( + duration_ms=timer.elapsed_ms, + input_modality="text", + output_modality="generation", + attachments=[ + Attachment(content_type="text/plain", role="input", data=prompt.encode()), + Attachment(content_type="text/plain", role="output", data=reply.encode()), + ], +) +client.close(timeout=10) +``` + ## Track spans for agentic workflows Use span events to track non-inference steps like planning, tool calls, retrieval, or memory updates. @@ -311,6 +345,11 @@ client.track_span( ) ``` +Span attribute values should be strings, numbers, booleans or `None`. Lists, +dicts and other values are sent as JSON strings (`["a", "b"]` arrives as +`'["a", "b"]'`), and the SDK logs a warning the first time it sees each such +key. + You can also attach optional correlation fields (`trace_id`, `span_id`, `parent_span_id`, `run_id`, `agent_id`, `step_index`, `conversation_id`) to any event by passing them into `track_inference`, `track_error`, `track_feedback`, diff --git a/docs/providers.md b/docs/providers.md index 8f78cea..1c29502 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -8,6 +8,12 @@ mapped to a `model_source` name on every event. OpenAI(base_url="https://api.fireworks.ai/inference/v1") ``` +## Direct HTTP calls + +For calls made with httpx, requests, or urllib instead of the OpenAI or +Anthropic client libraries, use [`wildedge.llm_api()`](llm_api.md). +See the [direct HTTP example](../examples/llm_api_example.py) for usage. + ## Recognized providers | Provider | Base URL | `model_source` | @@ -50,3 +56,93 @@ wildedge.register_model(client, source="ollama", model_id="gemma-7b") providers omit them; they record as `None`, not `0`. - Streamed responses carry no usage unless you pass `stream_options={"include_usage": True}`. Timing and TTFT are unaffected. + + +## Responses API + +The `openai` integration instruments `client.responses.create()` and +`await client.responses.create()`, with or without `stream=True`. Enable it +before constructing the OpenAI client, as with Chat Completions. + +```python +import wildedge +from openai import OpenAI + +wildedge.init(integrations=["openai"]) +client = OpenAI() +response = client.responses.create(model="gpt-4o-mini", input="Say hello") +print(response.output_text) +wildedge.flush() +``` + +Usage comes from Responses' `input_tokens` / `output_tokens`, including cached +input and reasoning tokens when provided. The stop reason is the response +status, or the incomplete reason such as `max_output_tokens`. Failed responses +are recorded as unsuccessful. No prompt or response text is uploaded by this +adapter. String input and simple user messages provide text input metadata; +input iterators are passed through without being consumed by instrumentation. + +See [Streaming lifecycle](streaming.md) for completion, closure, cancellation, +and coverage limits. + +For streams, consume the iterator to exhaustion to record the terminal response +and its usage. Time to first output ignores lifecycle events and measures the +first nonempty text, refusal, function argument, or audio delta. Streams without +a terminal response record an unsuccessful inference. Explicit early closure and +async cancellation are tracked; background response polling and individual tool +spans remain outside coverage. API exceptions retain their original behavior and use +the same error tracking as Chat Completions. + +OpenAI-compatible hosts must implement the Responses endpoint to use this path. +Older OpenAI SDK versions without a `responses` resource continue to support +Chat Completions. + + +### End-to-end check + +From the SDK root, with `OPENAI_API_KEY` and `WILDEDGE_DSN` set: + +```bash +uv run examples/openai_responses_example.py +``` + +This uses the local SDK checkout and makes four small, billable requests: +sync, sync streaming, async, and async streaming. Set `OPENAI_MODEL` to override +the default `gpt-4o-mini`. The script checks completion and usage, prints a +unique `responses-e2e-...` run ID, and closes the SDK with a 15-second deadline. +In WildEdge, locate that run and verify four inference events beneath the four +named spans, with input/output token usage and first-output timing on the two +streams. Provider checks alone do not prove telemetry delivery; confirm these +events in the dashboard. + + +## Embeddings + +The `openai` integration also tracks sync and async `client.embeddings.create()` +calls. Enable instrumentation before constructing the client. OpenAI-compatible +providers work through the same integration if they implement the embeddings +endpoint. + +Records latency, resolved model, input token usage, batch size, and dimensions. +Text, batches of strings, token ID arrays, and batches of token ID arrays are +supported. Arbitrary iterables pass through without instrumentation consuming +them; their batch size is omitted. Dimensions come from the returned vector +length, or float32 byte length for explicit base64 output. Usage comes from the +provider, not a local tokenizer. Both float and base64 vectors remain in the +application response and are never included in telemetry. Raw inputs are also +not uploaded by this adapter. + +Input usage is recorded in `input_meta.token_count` and `output_meta.tokens_in` +so existing analytics can query it. Embeddings do not emit generated-token +counts, throughput, or TTFT. API exceptions use the existing error tracking. + +With `OPENAI_API_KEY` and `WILDEDGE_DSN` configured, run from the SDK root: + +```bash +uv run examples/openai_embeddings_example.py +``` + +This makes two small billable calls, one sync and one async. Find the printed +run ID in WildEdge and verify two inference events, each linked to its named +span, with batch size 2, dimensions, and input usage. Set +`OPENAI_EMBEDDING_MODEL` to override `text-embedding-3-small`. diff --git a/docs/streaming.md b/docs/streaming.md new file mode 100644 index 0000000..6d2931e --- /dev/null +++ b/docs/streaming.md @@ -0,0 +1,83 @@ +# Streaming lifecycle and coverage + +The shared wrappers track OpenAI Chat Completions and Responses +`create(stream=True)`, Anthropic `messages.create(stream=True)` (sync and async), +Anthropic `messages.stream()` and its `text_stream` / `get_final_message()` +helpers (sync and async), and llama-cpp-python/GGUF completion iterators (sync). + +| Outcome | Telemetry | +|---|---| +| Iterator exhausted | One inference event; remote APIs require their completion marker | +| Explicit `close()` / `aclose()` or early context exit | One unsuccessful inference with `STREAM_CLOSED`, unless already finalized | +| Async cancellation while awaiting a chunk or exiting its context | One unsuccessful inference with `STREAM_CANCELLED`; cancellation is re-raised | +| Provider/iterator exception | One error event using existing error classification; original exception is re-raised | +| Remote API EOF without completion marker | One unsuccessful inference with `STREAM_INCOMPLETE` | + +Repeated exhaustion and closing after finalization do not emit another outcome. +Partial usage is retained when supplied by the provider, but missing totals are +not estimated. Responses interrupted before a terminal response has no final +usage; first-output timing is retained when available. Telemetry callback failures +are logged at debug level and do not break application iteration. + +TTFT measures from the original call until the first nonempty output delta, +ignoring role-only chunks, usage chunks, and message/response lifecycle events. +Output can be text, refusal, tool arguments, or supported audio/reasoning deltas; +it need not be user-visible text. No output means TTFT is omitted. Callbacks use +the trace context captured when the stream wrapper was created, even when the +application consumes it later under another context. + +Use a context manager or explicitly close the stream when stopping early: + +```python +with client.chat.completions.create(model=model, messages=messages, stream=True) as stream: + for chunk in stream: + if enough_output(chunk): + break +``` + +Async streams support `async with`, `await stream.close()`, and +`await stream.aclose()`. A bare `break` without closing, or simply abandoning an +unconsumed stream, does not provide an observable end of the operation. There is +no garbage-collection finalizer, except for faster-whisper segments (see +[whisper.md](whisper.md)). Keep stream consumption within the lifetime of the +SDK client and do not consume one iterator concurrently from multiple tasks. + +## Audit findings outside these wrappers + +- Provider convenience properties/methods that consume an underlying stream + directly can bypass iteration on the wrapper. They are not covered by the + lifecycle contract above; use the wrapped event iterator. +- MLX instruments `mlx_lm.generate()` and model calls, not + `mlx_lm.stream_generate()`. Streaming generation requires a dedicated adapter + and suppression of nested token-loop model calls. +- Transformers instruments pipeline/model calls; user-supplied streamers and + background generation threads are not wrapped stream iterators. +- `wildedge.llm_api()` is a manual context manager: its lifetime is controlled by + the caller, and the caller must supply usage and first-output marks. + +The audit does not claim coverage for streaming APIs that are not instrumented. + +## Anthropic native helper + +```python +with client.messages.stream(model=model, max_tokens=64, messages=messages) as stream: + for text in stream.text_stream: + print(text, end="") + message = stream.get_final_message() +``` + +The async equivalent uses `async with`, `async for`, and +`await stream.get_final_message()`. Event iteration, text iteration, and final +message helpers share a single tracked iterator and produce one outcome. Early +context exit, cancellation, request/iteration errors, and truncated EOF use the +same lifecycle contract above. HTTP timing starts at context entry, when the +native helper sends its request, rather than when its lazy manager is created. + +The adapter preserves the native stream object and parser. It wraps the SDK's +internal `_iterator` and delegates native close behavior; compatibility is tested +against both the minimum supported Anthropic SDK and the current pinned SDK. + +Run `uv run examples/anthropic_stream_example.py` from the SDK root with +`ANTHROPIC_API_KEY` and `WILDEDGE_DSN` set. It makes two small billable calls, +prints a run ID, and checks provider output and usage. Confirm two inference +events and two spans under that run in WildEdge to verify telemetry delivery. diff --git a/docs/whisper.md b/docs/whisper.md new file mode 100644 index 0000000..e7dff27 --- /dev/null +++ b/docs/whisper.md @@ -0,0 +1,122 @@ +# Local Whisper speech recognition + +WildEdge tracks local Whisper transcription with two integrations: + +- `mlx_whisper` for [mlx-whisper](https://github.com/ml-explore/mlx-examples/tree/main/whisper) on Apple Silicon +- `faster_whisper` for [faster-whisper](https://github.com/SYSTRAN/faster-whisper) (CTranslate2) + +Initialize the SDK before loading a model or importing `transcribe` by name: + +```python +import wildedge + +client = wildedge.init(integrations=["mlx_whisper"]) +import mlx_whisper + +result = mlx_whisper.transcribe("speech.wav", path_or_hf_repo="/models/whisper") +``` + +faster-whisper decodes lazily, so the transcription is recorded once you have +read all of its segments: + +```python +client = wildedge.init(integrations=["faster_whisper"]) +from faster_whisper import WhisperModel + +model = WhisperModel("/models/whisper", device="cpu", compute_type="int8") +segments, info = model.transcribe("speech.wav") +text = "".join(segment.text for segment in segments) +``` + +`BatchedInferencePipeline(model).transcribe(...)` is tracked the same way. + +## What's recorded + +**Model loads** + +- Load time, format (`mlx` or `ctranslate2`), family (`whisper`), accelerator + and compute type or quantization. +- The model ID is the backend plus the model name, for example + `faster_whisper:Systran/faster-whisper-tiny`. A Hugging Face repo keeps its + repo name, also when loaded from a cache path. A local folder is named after + the folder. +- Every load is recorded, including a second `WhisperModel` for the same model. +- mlx-whisper caches one model at a time and reloads when you switch models. + Each reload is recorded. A model that was already cached before the SDK + started is recorded once, with zero load time and `model_cache_hit=true`. + +**Transcriptions** + +- Wall-clock duration and `language`. +- `language_confidence` (faster-whisper only; mlx-whisper doesn't report it). +- `transcript_word_count`, except for languages written without spaces between + words, such as Chinese, Japanese and Thai. +- `time_to_first_segment_ms` (faster-whisper only). +- Audio metadata: + - WAV files: duration, sample rate, channels and bit depth. + - 1-D arrays: treated as 16 kHz mono, as both libraries expect. + - Other formats: duration only, when faster-whisper reports it. The SDK + doesn't decode the audio a second time. + +Token counts aren't reported; neither library exposes them. + +## faster-whisper segment lifecycle + +Timing runs from `transcribe()` until the last segment is read, including time +your code spends between segments. Read the segments inside the trace or span +that should own the transcription. + +| How the segments end | Recorded as | +| --- | --- | +| All segments read | One successful transcription | +| `segments.close()` | One unsuccessful transcription, `STREAM_CLOSED` | +| Loop exits early without `close()` | One unsuccessful transcription, `STREAM_ABANDONED`, sent when the iterator is garbage collected and timed up to the last segment read | +| Never iterated (for example, only `info.language` was used) | Nothing | +| Decoder error | One error event; the exception reaches your code unchanged | + +A `WhisperModel` stays with the SDK client that was active when it was created. +After you replace the client, create the model again; until then, it isn't +tracked and the SDK warns once. + +## Explicit audio and transcript attachments + +Automatic instrumentation records metadata only, even with +`attachments_enabled=True`. To upload audio or transcript text, supply the +attachments explicitly through [manual tracking](manual-tracking.md#track-attachments). +For example, with a faster-whisper model: + +```python +import time +import wildedge +from wildedge import Attachment + +# Use manual tracking for this call so it produces one inference event. +client = wildedge.init(integrations=[], attachments_enabled=True) +from faster_whisper import WhisperModel + +model = WhisperModel("/models/whisper", device="cpu", compute_type="int8") +handle = client.register_model(model, auto_instrument=False) +started = time.perf_counter() +segments, info = model.transcribe("speech.wav") +text = "".join(segment.text for segment in segments) +handle.track_inference( + duration_ms=round((time.perf_counter() - started) * 1000), + input_modality="audio", + output_modality="text", + attachments=[ + Attachment(content_type="audio/wav", role="input", path="speech.wav"), + Attachment(content_type="text/plain", role="output", data=text.encode()), + ], +) +client.close(timeout=10) +``` + +Choose either attachment or both. Files must exist when `track_inference` runs; +normal attachment size limits and the optional `attachment_filter` apply. + +## Examples and tests + +See [the runnable example](../examples/whisper_example.py). Compatibility CI +decodes real audio with a random-weight MLX Whisper model and the pinned public +faster-whisper tiny model. To run the faster-whisper tests offline, set +`WILDEDGE_FASTER_WHISPER_MODEL=/path/to/model`. diff --git a/examples/agentic_example.py b/examples/agentic_example.py index 6874184..e4d17a0 100644 --- a/examples/agentic_example.py +++ b/examples/agentic_example.py @@ -80,7 +80,7 @@ def calculator(expression: str) -> str: # ~60ms to simulate a remote computation call. time.sleep(0.06) try: - result = eval(expression, {"__builtins__": {}}) # noqa: S307 + result = eval(expression, {"__builtins__": {}}) return json.dumps({"expression": expression, "result": result}) except Exception as e: return json.dumps({"error": str(e)}) diff --git a/examples/anthropic_stream_example.py b/examples/anthropic_stream_example.py new file mode 100644 index 0000000..dc41a12 --- /dev/null +++ b/examples/anthropic_stream_example.py @@ -0,0 +1,80 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["wildedge-sdk", "anthropic>=0.25.0"] +# +# [tool.uv.sources] +# wildedge-sdk = { path = "..", editable = true } +# /// +"""Two small billable native streaming calls, sync text and async final message. + +Requires ANTHROPIC_API_KEY and WILDEDGE_DSN; optionally ANTHROPIC_MODEL. +From the SDK root: uv run examples/anthropic_stream_example.py +""" + +from __future__ import annotations + +import asyncio +import os +import uuid + +from anthropic import Anthropic, AsyncAnthropic + +import wildedge + +MODEL = os.environ.get("ANTHROPIC_MODEL", "claude-haiku-4-5") +MESSAGES = [{"role": "user", "content": "Reply with exactly: Hello from WildEdge"}] + + +def run_sync() -> None: + with ( + Anthropic() as client, + wildedge.span(kind="agent_step", name="anthropic-native-text"), + ): + with client.messages.stream( + model=MODEL, max_tokens=64, messages=MESSAGES + ) as stream: + text = "".join(stream.text_stream) + result = stream.get_final_message() + if not text or not result.usage: + raise RuntimeError("Missing output or usage") + print( + f"Sync text: input={result.usage.input_tokens}, output={result.usage.output_tokens}" + ) + + +async def run_async() -> None: + async with AsyncAnthropic() as client: + with wildedge.span(kind="agent_step", name="anthropic-native-final"): + async with client.messages.stream( + model=MODEL, max_tokens=64, messages=MESSAGES + ) as stream: + result = await stream.get_final_message() + if not result.content or not result.usage: + raise RuntimeError("Missing output or usage") + print( + f"Async final: input={result.usage.input_tokens}, output={result.usage.output_tokens}" + ) + + +def main() -> None: + missing = [ + name + for name in ("ANTHROPIC_API_KEY", "WILDEDGE_DSN") + if not os.environ.get(name) + ] + if missing: + raise SystemExit("Set required environment variables: " + ", ".join(missing)) + run_id = "anthropic-native-e2e-" + uuid.uuid4().hex[:12] + print(f"Run ID: {run_id}") + sdk = wildedge.init(integrations=["anthropic"], app_version="anthropic-native-e2e") + try: + with wildedge.trace(run_id=run_id): + run_sync() + asyncio.run(run_async()) + finally: + sdk.close(timeout=15) + print("Provider checks passed. Verify two inferences and two spans in WildEdge.") + + +if __name__ == "__main__": + main() diff --git a/examples/django_gemma/demo.sh b/examples/django_gemma/demo.sh index 78d9ba1..9914f6b 100755 --- a/examples/django_gemma/demo.sh +++ b/examples/django_gemma/demo.sh @@ -11,7 +11,7 @@ cd "${SCRIPT_DIR}" uv sync -uv run wildedge doctor --integrations gguf --hubs huggingface +uv run wildedge doctor --integrations llama_cpp --hubs huggingface # wildedge run replaces this process (os.execle) so sitecustomize.py # auto-installs the runtime before Django loads, patching Llama.__init__ @@ -26,13 +26,13 @@ uv run wildedge doctor --integrations gguf --hubs huggingface if [[ "$(uname)" == "Darwin" ]]; then uv run wildedge run \ --print-startup-report \ - --integrations gguf \ + --integrations llama_cpp \ --hubs huggingface \ -- waitress-serve --port=8100 gemmaapp.wsgi:application else uv run wildedge run \ --print-startup-report \ - --integrations gguf \ + --integrations llama_cpp \ --hubs huggingface \ -- gunicorn gemmaapp.wsgi:application --config gunicorn.conf.py fi diff --git a/examples/django_gemma/gemmaapp/views.py b/examples/django_gemma/gemmaapp/views.py index 2bc3cf2..ebad98e 100644 --- a/examples/django_gemma/gemmaapp/views.py +++ b/examples/django_gemma/gemmaapp/views.py @@ -1,6 +1,6 @@ """Gemma inference view. -The Llama constructor is patched automatically by `wildedge run --integrations gguf` +The Llama constructor is patched automatically by `wildedge run --integrations llama_cpp` via sitecustomize.py; load/unload/inference events are tracked without any wildedge imports here. diff --git a/examples/gguf_example.py b/examples/gguf_example.py index 861603d..4f252b6 100644 --- a/examples/gguf_example.py +++ b/examples/gguf_example.py @@ -5,7 +5,14 @@ # [tool.uv.sources] # wildedge-sdk = { path = "..", editable = true } # /// -"""GGUF / llama.cpp integration example. Run with: uv run gguf_example.py""" +"""GGUF / llama.cpp integration example. Run with: uv run gguf_example.py + +The llama_cpp integration tracks every llama-cpp-python entry point, once per +call: create_chat_completion (streamed or not), create_completion and +llm(...). Each inference event carries token counts, tokens per second, +the stop reason and the sampling settings; chat calls also record +prompt_type="chat" and the turn index. +""" from huggingface_hub import hf_hub_download from llama_cpp import Llama @@ -14,7 +21,7 @@ wildedge.init( app_version="1.0.0", # uses WILDEDGE_DSN if set; otherwise no-op - integrations="gguf", + integrations="llama_cpp", hubs=["huggingface"], ) @@ -24,16 +31,27 @@ ) llm = Llama(model_path, n_ctx=2048, n_gpu_layers=-1, verbose=False) -prompts = [ - "Explain what a transformer architecture is in one sentence.", - "What is the capital of France?", - "Write a haiku about on-device AI.", -] - -for prompt in prompts: - stream = llm(prompt, max_tokens=128, temperature=0.7, stream=True) - print(f"Q: {prompt}\nA: ", end="", flush=True) - for chunk in stream: - token = chunk["choices"][0].get("text", "") +# A short multi-turn chat, streamed. Each turn is one inference event. +messages = [{"role": "system", "content": "You are a concise assistant."}] +for question in [ + "What is a transformer architecture, in one sentence?", + "And why does it suit on-device AI?", + "Now say that as a haiku.", +]: + messages.append({"role": "user", "content": question}) + print(f"Q: {question}\nA: ", end="", flush=True) + reply = "" + for chunk in llm.create_chat_completion( + messages=messages, max_tokens=128, temperature=0.7, top_p=0.9, stream=True + ): + token = chunk["choices"][0]["delta"].get("content") or "" + reply += token print(token, end="", flush=True) print("\n") + messages.append({"role": "assistant", "content": reply}) + +# Plain text completion, not streamed. +result = llm.create_completion( + "The capital of France is", max_tokens=8, temperature=0.0, stop=["."] +) +print(f"Completion: {result['choices'][0]['text'].strip()}") diff --git a/examples/llm_api_example.py b/examples/llm_api_example.py index 1afe2ff..19d92f0 100644 --- a/examples/llm_api_example.py +++ b/examples/llm_api_example.py @@ -51,12 +51,12 @@ def chat(prompt: str) -> dict: with wildedge.trace(agent_id="llm-api-example", run_id=str(uuid.uuid4())): for step, prompt in enumerate(PROMPTS): - with wildedge.span(kind="agent_step", name="ask", step_index=step): - with wildedge.llm_api( - model=MODEL, provider="openrouter", prompt=prompt - ) as call: - data = chat(prompt) - call.response(data) + with ( + wildedge.span(kind="agent_step", name="ask", step_index=step), + wildedge.llm_api(model=MODEL, provider="openrouter", prompt=prompt) as call, + ): + data = chat(prompt) + call.response(data) print(f"Q: {prompt}\nA: {data['choices'][0]['message']['content']}\n") wildedge.flush() diff --git a/examples/openai_embeddings_example.py b/examples/openai_embeddings_example.py new file mode 100644 index 0000000..06a436c --- /dev/null +++ b/examples/openai_embeddings_example.py @@ -0,0 +1,76 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["wildedge-sdk", "openai>=1.30.0"] +# +# [tool.uv.sources] +# wildedge-sdk = { path = "..", editable = true } +# /// +"""Run two small billable embedding calls; requires OPENAI_API_KEY and WILDEDGE_DSN. + +From the SDK root: uv run examples/openai_embeddings_example.py +Optionally set OPENAI_EMBEDDING_MODEL. Find the printed run ID in WildEdge. +""" + +from __future__ import annotations + +import asyncio +import os +import uuid + +from openai import AsyncOpenAI, OpenAI + +import wildedge + +MODEL = os.environ.get("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small") +INPUT = ["Hello from WildEdge", "Embedding instrumentation check"] + + +def check_result(result: object) -> None: + if len(result.data) != len(INPUT) or not result.data[0].embedding: + raise RuntimeError("Expected two nonempty embeddings") + if result.usage is None: + raise RuntimeError("Provider returned no usage") + print( + f"Completed: batch={len(result.data)}, dimensions={len(result.data[0].embedding)}, input_tokens={result.usage.prompt_tokens}" + ) + + +def run_sync() -> None: + with OpenAI() as client, wildedge.span(kind="agent_step", name="embeddings-sync"): + check_result( + client.embeddings.create(model=MODEL, input=INPUT, encoding_format="float") + ) + + +async def run_async() -> None: + async with AsyncOpenAI() as client: + with wildedge.span(kind="agent_step", name="embeddings-async"): + check_result( + await client.embeddings.create( + model=MODEL, input=INPUT, encoding_format="float" + ) + ) + + +def main() -> None: + missing = [ + key for key in ("OPENAI_API_KEY", "WILDEDGE_DSN") if not os.environ.get(key) + ] + if missing: + raise SystemExit("Set required environment variables: " + ", ".join(missing)) + run_id = "embeddings-e2e-" + uuid.uuid4().hex[:12] + print(f"Run ID: {run_id}") + sdk = wildedge.init(integrations=["openai"], app_version="embeddings-e2e") + try: + with wildedge.trace(run_id=run_id): + run_sync() + asyncio.run(run_async()) + finally: + sdk.close(timeout=15) + print( + "Provider checks passed. Verify two inference events and two spans in WildEdge." + ) + + +if __name__ == "__main__": + main() diff --git a/examples/openai_responses_example.py b/examples/openai_responses_example.py new file mode 100644 index 0000000..1ae6baa --- /dev/null +++ b/examples/openai_responses_example.py @@ -0,0 +1,101 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["wildedge-sdk", "openai>=1.66.0"] +# +# [tool.uv.sources] +# wildedge-sdk = { path = "..", editable = true } +# /// +"""Exercise all four Responses paths against OpenAI and WildEdge. + +From the SDK root: + uv run examples/openai_responses_example.py + +Requires OPENAI_API_KEY and WILDEDGE_DSN. Optionally set OPENAI_MODEL. +Makes four small, billable API calls. Find the printed run ID in WildEdge. +""" + +from __future__ import annotations + +import asyncio +import os +import uuid + +from openai import AsyncOpenAI, OpenAI + +import wildedge + +PROMPT = "Reply with exactly: Hello from Responses" +MODEL = os.environ.get("OPENAI_MODEL", "gpt-4o-mini") + + +def check_response(result: object, label: str) -> None: + if result.status != "completed" or not result.output_text: + raise RuntimeError(f"{label}: response did not complete with text") + if result.usage is None: + raise RuntimeError(f"{label}: provider returned no usage") + print( + f"{label}: completed; input={result.usage.input_tokens}, output={result.usage.output_tokens}" + ) + + +def run_sync() -> None: + with OpenAI() as client: + with wildedge.span(kind="agent_step", name="responses-sync"): + result = client.responses.create( + model=MODEL, input=PROMPT, max_output_tokens=64 + ) + check_response(result, "sync") + with wildedge.span(kind="agent_step", name="responses-sync-stream"): + terminal = None + with client.responses.create( + model=MODEL, input=PROMPT, max_output_tokens=64, stream=True + ) as stream: + for event in stream: + if event.type == "response.completed": + terminal = event.response + if terminal is None: + raise RuntimeError("sync stream: missing completed response") + check_response(terminal, "sync stream") + + +async def run_async() -> None: + async with AsyncOpenAI() as client: + with wildedge.span(kind="agent_step", name="responses-async"): + result = await client.responses.create( + model=MODEL, input=PROMPT, max_output_tokens=64 + ) + check_response(result, "async") + with wildedge.span(kind="agent_step", name="responses-async-stream"): + terminal = None + async with await client.responses.create( + model=MODEL, input=PROMPT, max_output_tokens=64, stream=True + ) as stream: + async for event in stream: + if event.type == "response.completed": + terminal = event.response + if terminal is None: + raise RuntimeError("async stream: missing completed response") + check_response(terminal, "async stream") + + +def main() -> None: + missing = [ + key for key in ("OPENAI_API_KEY", "WILDEDGE_DSN") if not os.environ.get(key) + ] + if missing: + raise SystemExit("Set required environment variables: " + ", ".join(missing)) + run_id = f"responses-e2e-{uuid.uuid4().hex[:12]}" + print(f"Run ID: {run_id}") + sdk = wildedge.init(integrations=["openai"], app_version="responses-e2e") + try: + with wildedge.trace(run_id=run_id): + run_sync() + asyncio.run(run_async()) + finally: + sdk.close(timeout=15) + print("Provider checks passed. SDK shutdown finished; verify delivery in WildEdge.") + print("Expect four inference events under this run, one beneath each named span.") + + +if __name__ == "__main__": + main() diff --git a/examples/whisper_example.py b/examples/whisper_example.py new file mode 100644 index 0000000..f00a59d --- /dev/null +++ b/examples/whisper_example.py @@ -0,0 +1,52 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["wildedge-sdk"] +# +# [tool.uv.sources] +# wildedge-sdk = { path = "..", editable = true } +# /// +"""Transcribe a WAV using a local Whisper model with automatic instrumentation. + +Set WILDEDGE_DSN, then add the backend you want to run: + uv run --with mlx-whisper examples/whisper_example.py mlx_whisper MODEL AUDIO + uv run --with faster-whisper examples/whisper_example.py faster_whisper MODEL AUDIO + +MODEL is a local path or Hugging Face repo, AUDIO is a WAV file. + +See docs/whisper.md for lifecycle and explicit attachment examples. +""" + +import argparse + +import wildedge + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("backend", choices=["mlx_whisper", "faster_whisper"]) + parser.add_argument("model") + parser.add_argument("audio") + args = parser.parse_args() + client = wildedge.init(integrations=[args.backend]) + try: + with client.trace(), client.span(kind="agent_step", name="transcribe"): + if args.backend == "mlx_whisper": + import mlx_whisper # noqa: PLC0415 -- optional runtime + + result = mlx_whisper.transcribe(args.audio, path_or_hf_repo=args.model) + print(result["text"]) + else: + from faster_whisper import ( # noqa: PLC0415 -- optional runtime + WhisperModel, + ) + + model = WhisperModel(args.model, device="cpu", compute_type="int8") + segments, _info = model.transcribe(args.audio) + print("".join(segment.text for segment in segments)) + client.flush() + finally: + client.close() + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index a378272..7127eec 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "wildedge-sdk" -version = "0.2.2" +version = "0.3.0" description = "ML inference monitoring for Python: on-device and remote models" readme = "README.md" requires-python = ">=3.10" @@ -52,11 +52,22 @@ packages = ["wildedge"] line-length = 88 [tool.ruff.lint] -select = ["E", "F", "W", "I", "N", "UP", "TID252"] +select = [ + "E", "F", "W", "I", "N", "UP", "TID252", + "B", # likely bugs (bugbear) + "C4", # comprehensions + "PIE", # misc. simplifications + "PLC0415", # imports belong at the top of the module + "PT", # pytest style + "RET", # return statements + "RUF", # ruff-specific checks + "SIM", # simplifications +] ignore = ["E501"] # line-length is set above [tool.ruff.lint.per-file-ignores] "wildedge/decorators.py" = ["N801"] # track is a decorator class, not following CapWords +"tests/**" = ["RUF012"] # test doubles mimic library objects with class-level lists and dicts [tool.ruff.lint.flake8-tidy-imports] ban-relative-imports = "all" @@ -65,6 +76,9 @@ ban-relative-imports = "all" [tool.pytest.ini_options] testpaths = ["tests"] +# tests/compat needs the pinned ML libraries; the Compatibility workflow runs +# it file by file (scripts/compat_matrix.py). +norecursedirs = [".*", "*.egg", "build", "dist", "node_modules", "venv", "compat"] log_level = "ERROR" asyncio_mode = "auto" markers = [ diff --git a/scripts/build_changelog_comment.py b/scripts/build_changelog_comment.py index f05bdc6..e88bec5 100644 --- a/scripts/build_changelog_comment.py +++ b/scripts/build_changelog_comment.py @@ -1,6 +1,9 @@ #!/usr/bin/env python3 """Generate changelog text for release PRs and for tag-build release notes. +The notes are the version's section of CHANGELOG.md, so they list what changed +for users rather than every commit and merge since the previous tag. + The two callers differ only in wording: release-pr.yml renders a preview of notes that do not exist yet, while release.yml (which sets TAG_NAME) renders the notes themselves. @@ -13,6 +16,8 @@ import subprocess from pathlib import Path +CHANGELOG = Path("CHANGELOG.md") + def get_version() -> str: # Regex instead of tomllib so the script also runs on Python 3.10. @@ -24,9 +29,6 @@ def get_version() -> str: return match.group(1) -RELEASE_COMMIT_RE = re.compile(r"^Release \d+\.\d+\.\d+$") - - def get_previous_tag() -> str | None: result = subprocess.run( ["git", "describe", "--tags", "--abbrev=0", "--match", "v*", "HEAD^"], @@ -37,26 +39,28 @@ def get_previous_tag() -> str | None: return tag if result.returncode == 0 and tag else None -def get_commits_since(prev_tag: str | None) -> list[str]: - revision = f"{prev_tag}..HEAD" if prev_tag else "HEAD" - result = subprocess.run( - ["git", "log", revision, "--pretty=format:%s"], - capture_output=True, - text=True, - check=True, - ) - lines = [line for line in result.stdout.splitlines() if line.strip()] - return [line for line in lines if not RELEASE_COMMIT_RE.match(line)] +def changelog_section(text: str, version: str) -> str: + """Body of the ``## `` section, up to the next ``## `` heading.""" + heading = re.compile(rf"^## \[?{re.escape(version)}\]?(?:\s|$)", re.MULTILINE) + match = heading.search(text) + if match is None: + raise SystemExit(f"CHANGELOG.md has no section for {version}") + line_end = text.find("\n", match.start()) + start = len(text) if line_end == -1 else line_end + 1 + end = re.compile(r"^## ", re.MULTILINE).search(text, start) + body = text[start : end.start() if end else len(text)].strip() + if not body: + raise SystemExit(f"CHANGELOG.md section for {version} is empty") + return body def build_comment( tag_name: str, - commits: list[str], + notes: str, prev_tag: str | None, *, preview: bool = True, ) -> str: - items = "\n".join(f"- {c}" for c in commits) if commits else "- No changes." range_str = f"{prev_tag}...{tag_name}" if prev_tag else tag_name heading = "Changelog preview for" if preview else "Changelog for" blurb = ( @@ -67,22 +71,22 @@ def build_comment( return ( f"## {heading} `{tag_name}`\n\n" f"{blurb}" - f"**What's Changed**\n\n{items}\n\n" + f"{notes}\n\n" f"**Full Changelog**: {range_str}\n" ) def main() -> None: output = os.environ.get("OUTPUT", "/tmp/changelog-preview.md") - # release.yml sets TAG_NAME on tag builds; release-pr.yml never does, so its + # release.yml sets TAG_NAME on tag builds and release-pr.yml never does. Its # absence is what distinguishes a PR preview from the real release notes. tag_name = os.environ.get("TAG_NAME") preview = not tag_name tag_name = tag_name or f"v{get_version()}" + notes = changelog_section(CHANGELOG.read_text(), tag_name.removeprefix("v")) prev_tag = get_previous_tag() - commits = get_commits_since(prev_tag) - Path(output).write_text(build_comment(tag_name, commits, prev_tag, preview=preview)) + Path(output).write_text(build_comment(tag_name, notes, prev_tag, preview=preview)) if __name__ == "__main__": diff --git a/scripts/build_llms_txt.py b/scripts/build_llms_txt.py index 66f0cad..d4eaa3e 100644 --- a/scripts/build_llms_txt.py +++ b/scripts/build_llms_txt.py @@ -24,6 +24,8 @@ ("docs/manual-tracking.md", "manual-tracking", "Manual tracking"), ("docs/llm_api.md", "llm_api", "LLM API tracking"), ("docs/providers.md", "providers", "Providers"), + ("docs/streaming.md", "streaming", "Streaming lifecycle"), + ("docs/whisper.md", "whisper", "Whisper"), ("docs/compatibility.md", "compatibility", "Compatibility matrix"), ] diff --git a/scripts/compat_matrix.py b/scripts/compat_matrix.py index 7cb5d22..97f424e 100644 --- a/scripts/compat_matrix.py +++ b/scripts/compat_matrix.py @@ -6,19 +6,20 @@ import argparse import json import sys +from importlib import metadata MATRIX = { "anthropic": { "min": ["anthropic==0.25.0"], - "current": ["anthropic==0.92.0"], + "current": ["anthropic==1.15.0"], }, "onnx": { "min": ["onnxruntime==1.18.1", "numpy==1.26.4"], - "current": ["onnxruntime==1.20.1", "numpy==2.1.3"], + "current": ["onnxruntime==1.31.0", "numpy==2.2.6"], }, "torch": { "min": ["torch==2.4.1", "numpy==1.26.4"], - "current": ["torch==2.6.0", "numpy==2.1.3"], + "current": ["torch==2.14.1", "numpy==2.2.6"], }, "timm": { "min": [ @@ -28,35 +29,39 @@ "numpy==1.26.4", ], "current": [ - "torch==2.6.0", - "torchvision==0.21.0", - "timm==1.0.15", - "numpy==2.1.3", + "torch==2.14.1", + "torchvision==0.29.1", + "timm==1.0.30", + "numpy==2.2.6", ], }, "tensorflow": { "min": ["tensorflow==2.16.1", "keras==3.3.3", "numpy==1.26.4"], - "current": ["tensorflow==2.18.0", "keras==3.8.0", "numpy==2.0.2"], + "current": ["tensorflow==2.21.0", "keras==3.15.1", "numpy==2.2.6"], }, - "gguf": { + "llama_cpp": { "min": ["llama-cpp-python==0.2.90", "numpy==1.26.4"], - "current": ["llama-cpp-python==0.3.4", "numpy==2.1.3"], + "current": ["llama-cpp-python==0.3.36", "numpy==2.2.6"], }, "openai": { - "min": ["openai==1.30.0"], - "current": ["openai==1.61.0"], + "min": ["openai==1.30.0", "httpx==0.27.2"], + "current": ["openai==3.35.1"], }, "transformers": { "min": ["transformers==4.40.0", "torch==2.4.1", "numpy==1.26.4"], - "current": ["transformers==4.47.0", "torch==2.6.0", "numpy==2.1.3"], + "current": ["transformers==5.20.1", "torch==2.14.1", "numpy==2.2.6"], }, "ultralytics": { - "min": ["ultralytics==8.0.0", "numpy==1.26.4"], - "current": ["ultralytics==8.3.4", "numpy==2.1.3"], + "min": ["ultralytics==8.0.200", "numpy==1.26.4"], + "current": ["ultralytics==8.4.182", "numpy==2.2.6"], }, "mlx": { - "current": ["mlx==0.22.0", "mlx-lm==0.21.0"], + "current": ["mlx==0.32.4", "mlx-lm==0.32.2"], }, + "faster_whisper": { + "current": ["faster-whisper==1.2.1", "ctranslate2==4.8.2", "numpy==2.2.6"] + }, + "mlx_whisper": {"current": ["mlx-whisper==0.4.3", "mlx==0.32.3", "numpy==2.2.6"]}, } SUPPORTED_PYTHON = { @@ -65,19 +70,38 @@ "torch": ["3.10", "3.11", "3.12", "3.13", "3.14"], "timm": ["3.10", "3.11", "3.12", "3.13", "3.14"], "tensorflow": ["3.10", "3.11", "3.12"], - "gguf": ["3.10", "3.11", "3.12", "3.13"], + "llama_cpp": ["3.10", "3.11", "3.12", "3.13"], "openai": ["3.10", "3.11", "3.12", "3.13"], "transformers": ["3.10", "3.11", "3.12", "3.13"], "ultralytics": ["3.10", "3.11", "3.12", "3.13"], "mlx": ["3.12", "3.13"], + "faster_whisper": ["3.12", "3.13"], + "mlx_whisper": ["3.12", "3.13"], +} + +# Integrations without a row of their own, and the row that exercises them. +COVERED_BY = { + "keras": "tensorflow", } +# Integrations that need a specific runner; everything else runs on Linux. +RUNNER = { + "mlx": "macos-14", + "mlx_whisper": "macos-14", +} + +# Python versions run as non-blocking canaries, and only with current pins. +CANARY_PYTHON = {"3.14"} + # Interpreter-specific overrides where upstream wheels are unavailable for older pins. PYTHON_DEP_OVERRIDES = { "onnx": { "min": { "3.13": ["onnxruntime==1.20.1", "numpy==2.1.3"], - } + }, + "current": { + "3.10": ["onnxruntime==1.23.2", "numpy==2.2.6"], + }, }, "torch": { "min": { @@ -85,7 +109,12 @@ "3.14": ["torch==2.5.0", "numpy==2.1.3"], }, "current": { - "3.14": ["torch==2.10.0", "numpy==2.1.3"], + "3.14": ["torch==2.14.1", "numpy==2.5.4"], + }, + }, + "tensorflow": { + "current": { + "3.10": ["tensorflow==2.21.0", "keras==3.12.4", "numpy==2.2.6"], }, }, "transformers": { @@ -110,10 +139,10 @@ }, "current": { "3.14": [ - "torch==2.10.0", - "torchvision==0.25.0", - "timm==1.0.15", - "numpy==2.1.3", + "torch==2.14.1", + "torchvision==0.29.1", + "timm==1.0.30", + "numpy==2.5.4", ], }, }, @@ -177,25 +206,88 @@ def print_table() -> int: return 0 +def _strip_local(version: str) -> str: + return version.split("+", 1)[0] + + +def verify(integration: str, version_set: str, python_version: str) -> int: + """Check that the installed distributions match the pins exactly. + + pip can satisfy a pin and then let another package upgrade it, or fall + back to a different release; this makes the job fail instead. + """ + mismatches = [] + for requirement in resolve_deps(integration, version_set, python_version): + name, _, wanted = requirement.partition("==") + try: + installed = metadata.version(name) + except metadata.PackageNotFoundError: + mismatches.append(f"{name}: not installed, pinned {wanted}") + continue + if _strip_local(installed) != wanted: + mismatches.append(f"{name}: installed {installed}, pinned {wanted}") + for line in mismatches: + print(line, file=sys.stderr) + return 1 if mismatches else 0 + + +def github_matrix() -> list[dict]: + """Rows for the compat workflow: every supported combination. + + Canary Pythons run with current pins only and don't fail the run. + """ + rows = [] + for integration, sets in MATRIX.items(): + runner = RUNNER.get(integration, "ubuntu-latest") + for version_set in sets: + pythons = [ + py + for py in SUPPORTED_PYTHON[integration] + if py not in CANARY_PYTHON or version_set == "current" + ] + for py in pythons: + rows.append( + { + "integration": integration, + "version-set": version_set, + "python-version": py, + "runs-on": runner, + "canary": py in CANARY_PYTHON, + } + ) + return rows + + +def print_github_matrix() -> int: + print(json.dumps({"include": github_matrix()})) + return 0 + + def main() -> int: parser = argparse.ArgumentParser() sub = parser.add_subparsers(dest="cmd", required=True) - deps = sub.add_parser("deps") - deps.add_argument("--integration", required=True) - deps.add_argument("--version-set", required=True, choices=["min", "current"]) - deps.add_argument("--python-version", required=True) + for name in ("deps", "verify"): + cmd = sub.add_parser(name) + cmd.add_argument("--integration", required=True) + cmd.add_argument("--version-set", required=True, choices=["min", "current"]) + cmd.add_argument("--python-version", required=True) sub.add_parser("rows") sub.add_parser("table") + sub.add_parser("github-matrix") args = parser.parse_args() if args.cmd == "deps": return print_deps(args.integration, args.version_set, args.python_version) + if args.cmd == "verify": + return verify(args.integration, args.version_set, args.python_version) if args.cmd == "rows": return print_rows() if args.cmd == "table": return print_table() + if args.cmd == "github-matrix": + return print_github_matrix() return 2 diff --git a/scripts/run_compat_local.py b/scripts/run_compat_local.py index 44fe82f..4c2d5b3 100755 --- a/scripts/run_compat_local.py +++ b/scripts/run_compat_local.py @@ -24,14 +24,28 @@ class Row: def load_rows() -> list[Row]: + """The rows the nightly workflow runs; macOS-only rows only on macOS.""" result = subprocess.run( - [sys.executable, str(REPO_ROOT / "scripts" / "compat_matrix.py"), "rows"], + [ + sys.executable, + str(REPO_ROOT / "scripts" / "compat_matrix.py"), + "github-matrix", + ], check=True, capture_output=True, text=True, cwd=REPO_ROOT, ) - return [Row(**r) for r in json.loads(result.stdout)] + on_macos = sys.platform == "darwin" + return [ + Row( + python_version=r["python-version"], + integration=r["integration"], + version_set=r["version-set"], + ) + for r in json.loads(result.stdout)["include"] + if on_macos or not r["runs-on"].startswith("macos") + ] WORKFLOW_ROWS: list[Row] = load_rows() diff --git a/tests/compat/conftest.py b/tests/compat/conftest.py index a42f63a..fbc19e7 100644 --- a/tests/compat/conftest.py +++ b/tests/compat/conftest.py @@ -1,5 +1,18 @@ +"""Shared fixtures for the compatibility suite. + +These tests run against real, pinned library versions (see +scripts/compat_matrix.py). Import the library directly: a missing or broken +install must fail the job, never skip it. +""" + from __future__ import annotations +import hashlib +import importlib +import os +import urllib.request +from pathlib import Path +from types import ModuleType from unittest.mock import patch import pytest @@ -7,6 +20,68 @@ import wildedge from wildedge.platforms.device_info import DeviceInfo +FIXTURES = Path(__file__).parent / "fixtures" +# Downloaded test models; CI caches this directory between runs. +DOWNLOAD_CACHE = Path( + os.environ.get("WILDEDGE_COMPAT_CACHE", Path(__file__).parent / ".cache") +) + + +def downloaded_fixture(name: str, url: str, sha256: str) -> Path: + """Download ``url`` once into the cache and check its sha256. + + A failed download or a different checksum fails the test: update the URL + (and the checksum if the file changed) where the fixture is declared. + """ + path = DOWNLOAD_CACHE / name + if path.exists() and _sha256(path) == sha256: + return path + DOWNLOAD_CACHE.mkdir(parents=True, exist_ok=True) + partial = path.with_suffix(path.suffix + ".part") + try: + urllib.request.urlretrieve(url, partial) + except OSError as exc: + pytest.fail(f"could not download test model {name} from {url}: {exc}") + actual = _sha256(partial) + if actual != sha256: + partial.unlink() + pytest.fail(f"{url} has sha256 {actual}, expected {sha256}") + partial.replace(path) + return path + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as f: + for block in iter(lambda: f.read(1 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +@pytest.hookimpl(hookwrapper=True) +def pytest_runtest_makereport(item, call): + """Fail any skipped test: a skip here means a dependency didn't install. + + Expected failures (xfail) are still allowed; they name a known bug. + """ + outcome = yield + report = outcome.get_result() + if report.skipped and not hasattr(report, "wasxfail"): + report.outcome = "failed" + report.longrepr = f"compat tests must not skip: {report.longrepr}" + + +def sdk_http_module(package: str) -> ModuleType: + """The HTTP library an API SDK is built on: httpx, or httpx2 in newer releases. + + Tests pass a client with a MockTransport from that same library, so no + request leaves the machine. + """ + base_client = importlib.import_module(f"{package}._base_client") + module = getattr(base_client, "httpx2", None) or base_client.httpx + assert hasattr(module, "MockTransport"), module + return module + class _DummyConsumer: def __init__(self, *args, **kwargs): @@ -15,7 +90,7 @@ def __init__(self, *args, **kwargs): def flush(self, timeout: float = 5.0) -> None: pass - def close(self) -> None: + def close(self, timeout: float | None = None) -> None: pass def _pause(self) -> None: @@ -30,8 +105,38 @@ def __init__(self, *args, **kwargs): pass -@pytest.fixture -def compat_client(): +class CompatClient: + """A WildEdge client that records published events instead of sending them.""" + + def __init__(self, client: wildedge.WildEdge) -> None: + self.client = client + self.events: list[dict] = [] + # Handles keep a reference to client.publish, so swap it before any + # model is registered. + original_publish = client.publish + + def publish(event): + original_publish(event) + if not client.closed: + self.events.append(event) + + client.publish = publish # type: ignore[method-assign] + + def __getattr__(self, name: str): + return getattr(self.client, name) + + def of_type(self, event_type: str) -> list[dict]: + return [e for e in self.events if e["event_type"] == event_type] + + def one(self, event_type: str) -> dict: + found = self.of_type(event_type) + types = [e["event_type"] for e in self.events] + assert len(found) == 1, f"expected one {event_type} event, got {types}" + return found[0] + + +def make_compat_client(): + """Construct a recording client without background network services.""" with ( patch( "wildedge.client.detect_device", @@ -43,6 +148,17 @@ def compat_client(): client = wildedge.WildEdge( dsn="https://compat-secret@ingest.wildedge.dev/compat-project", app_version="compat", + enable_offline_persistence=False, + sampling_interval_s=0, ) - yield client - client.close() + return CompatClient(client) + + +@pytest.fixture +def compat_client(): + """Fresh clients exercise replacement through already-installed process patches.""" + wrapped = make_compat_client() + try: + yield wrapped + finally: + wrapped.close() diff --git a/tests/compat/fixtures/README.md b/tests/compat/fixtures/README.md new file mode 100644 index 0000000..91be61c --- /dev/null +++ b/tests/compat/fixtures/README.md @@ -0,0 +1,9 @@ +# Compatibility test fixtures + +- `identity.onnx`: a one-node ONNX graph, generated by `make_onnx_fixture.py`. + +Larger models are downloaded on first use and checked against a sha256 (see +`downloaded_fixture` in `../conftest.py`). They are cached in +`tests/compat/.cache`, or in `$WILDEDGE_COMPAT_CACHE` when set, which CI +keeps between runs. The llama_cpp test declares its model's URL and checksum at +the top of `test_llama_cpp_compat.py`. diff --git a/tests/compat/fixtures/identity.onnx b/tests/compat/fixtures/identity.onnx new file mode 100644 index 0000000..01be959 Binary files /dev/null and b/tests/compat/fixtures/identity.onnx differ diff --git a/tests/compat/fixtures/make_onnx_fixture.py b/tests/compat/fixtures/make_onnx_fixture.py new file mode 100644 index 0000000..006525b --- /dev/null +++ b/tests/compat/fixtures/make_onnx_fixture.py @@ -0,0 +1,26 @@ +"""Regenerate identity.onnx, the model the onnx compat test loads. + + uv run --with onnx python tests/compat/fixtures/make_onnx_fixture.py + +IR version 8 and opset 13 keep it loadable by the oldest onnxruntime in the +compatibility matrix. +""" + +from __future__ import annotations + +from pathlib import Path + +import onnx +from onnx import TensorProto, helper + +graph = helper.make_graph( + [helper.make_node("Identity", ["x"], ["y"])], + "compat_identity", + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4])], +) +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) +model.ir_version = 8 +model.model_version = 3 +onnx.checker.check_model(model) +onnx.save(model, Path(__file__).with_name("identity.onnx")) diff --git a/tests/compat/test_anthropic_compat.py b/tests/compat/test_anthropic_compat.py index 0b7fadc..fe32033 100644 --- a/tests/compat/test_anthropic_compat.py +++ b/tests/compat/test_anthropic_compat.py @@ -1,9 +1,384 @@ from __future__ import annotations +import json + +import anthropic import pytest +from conftest import make_compat_client, sdk_http_module + +httpx = sdk_http_module("anthropic") + +MESSAGE = { + "id": "msg_compat", + "type": "message", + "role": "assistant", + "model": "claude-haiku-4-5", + "content": [{"type": "text", "text": "Hello there"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 3}, +} + + +def _client(handler) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key="sk-ant-compat", + max_retries=0, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + +def test_message_is_tracked(compat_client): + compat_client.instrument("anthropic") + client = _client(lambda request: httpx.Response(200, json=MESSAGE)) + + result = client.messages.create( + model="claude-haiku-4-5", + max_tokens=32, + messages=[{"role": "user", "content": "Say hello"}], + ) + assert result.content[0].text == "Hello there" + + inference = compat_client.one("inference") + assert inference["model_id"] == "claude-haiku-4-5" + output = inference["inference"]["output_meta"] + assert output["tokens_in"] == 12 + assert output["tokens_out"] == 3 + assert output["stop_reason"] == "end_turn" + assert compat_client.registry.models["claude-haiku-4-5"].model_format == "api" + + +def test_api_error_is_tracked(compat_client): + compat_client.instrument("anthropic") + client = _client( + lambda request: httpx.Response( + 529, + json={ + "type": "error", + "error": {"type": "overloaded_error", "message": "Overloaded"}, + }, + ) + ) + try: + client.messages.create( + model="claude-haiku-4-5", + max_tokens=8, + messages=[{"role": "user", "content": "x"}], + ) + except anthropic.APIStatusError: + pass + else: + raise AssertionError("expected APIStatusError") + + error = compat_client.one("error") + assert error["error"]["http_status"] == 529 + assert error["error"]["provider_error_code"] == "overloaded_error" + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("early_close", [False, True]) +@pytest.mark.parametrize("truncated", [False, True]) +async def test_stream_lifecycle(compat_client, asynchronous, early_close, truncated): + compat_client.instrument("anthropic") + events = [ + { + "type": "message_start", + "message": { + **MESSAGE, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + }, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello"}, + }, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, + ] + if truncated: + events = events[:-1] + body = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + + def handler(request): + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + if asynchronous: + async with anthropic.AsyncAnthropic( + api_key="sk-ant-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) as client: + async with await client.messages.create( + model="claude-haiku-4-5", max_tokens=32, messages=[], stream=True + ) as stream: + if early_close: + await anext(stream) + else: + assert len([event async for event in stream]) == len(events) + await stream.close() + else: + with _client(handler) as client: + with client.messages.create( + model="claude-haiku-4-5", max_tokens=32, messages=[], stream=True + ) as stream: + if early_close: + next(stream) + else: + assert len(list(stream)) == len(events) + stream.close() + inference = compat_client.one("inference")["inference"] + assert inference["success"] == (not early_close and not truncated) + assert inference["output_meta"]["tokens_in"] == 12 + if early_close: + assert inference["error_code"] == "STREAM_CLOSED" + assert inference["output_meta"].get("time_to_first_token_ms") is None + else: + if truncated: + assert inference["error_code"] == "STREAM_INCOMPLETE" + assert inference["output_meta"]["tokens_out"] == 3 + assert inference["output_meta"]["time_to_first_token_ms"] is not None + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("mode", ["events", "text", "final", "close", "truncated"]) +async def test_native_stream_helper(compat_client, asynchronous, mode): + compat_client.instrument("anthropic") + events = [ + { + "type": "message_start", + "message": { + **MESSAGE, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + }, + }, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "text", "text": ""}, + }, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "text_delta", "text": "Hello there"}, + }, + {"type": "content_block_stop", "index": 0}, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, + ] + if mode == "truncated": + events = events[:-1] + body = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + requests = [] + + def handler(request): + requests.append(request) + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + if asynchronous: + async with ( + anthropic.AsyncAnthropic( + api_key="sk-ant-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) as client, + client.messages.stream( + model="claude-haiku-4-5", + max_tokens=32, + messages=[{"role": "user", "content": "hi"}], + ) as stream, + ): + if mode == "text": + assert ( + "".join([text async for text in stream.text_stream]) + == "Hello there" + ) + elif mode in ("final", "truncated"): + assert (await stream.get_final_message()).content[ + 0 + ].text == "Hello there" + elif mode == "events": + assert [e.type async for e in stream][-1] == "message_stop" + else: + await anext(stream) + else: + with ( + _client(handler) as client, + client.messages.stream( + model="claude-haiku-4-5", + max_tokens=32, + messages=[{"role": "user", "content": "hi"}], + ) as stream, + ): + if mode == "text": + assert "".join(stream.text_stream) == "Hello there" + elif mode in ("final", "truncated"): + assert stream.get_final_message().content[0].text == "Hello there" + elif mode == "events": + assert [e.type for e in stream][-1] == "message_stop" + else: + next(stream) + assert len(requests) == 1 + inference = compat_client.one("inference")["inference"] + assert inference["success"] == (mode not in ("close", "truncated")) + assert inference["output_meta"]["tokens_in"] == 12 + if mode == "close": + assert inference["error_code"] == "STREAM_CLOSED" + else: + if mode == "truncated": + assert inference["error_code"] == "STREAM_INCOMPLETE" + assert inference["output_meta"]["tokens_out"] == 3 + assert inference["output_meta"]["time_to_first_token_ms"] is not None + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_native_stream_request_error(compat_client, asynchronous): + compat_client.instrument("anthropic") + + def handler(request): + return httpx.Response( + 429, + json={ + "type": "error", + "error": {"type": "rate_limit_error", "message": "slow down"}, + }, + ) + + async def invoke_async(): + async with ( + anthropic.AsyncAnthropic( + api_key="sk-ant-compat", + max_retries=0, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) as client, + client.messages.stream( + model="claude-haiku-4-5", max_tokens=32, messages=[] + ) as stream, + ): + await stream.get_final_message() + + def invoke_sync(): + with ( + _client(handler) as client, + client.messages.stream( + model="claude-haiku-4-5", max_tokens=32, messages=[] + ) as stream, + ): + stream.get_final_message() + + if asynchronous: + with pytest.raises(anthropic.RateLimitError): + await invoke_async() + else: + with pytest.raises(anthropic.RateLimitError): + invoke_sync() + error = compat_client.one("error")["error"] + assert error["http_status"] == 429 + assert error["provider_error_code"] == "rate_limit_error" + assert not compat_client.of_type("inference") + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_native_stream_provider_error_midstream(compat_client, asynchronous): + compat_client.instrument("anthropic") + events = [ + { + "type": "message_start", + "message": {**MESSAGE, "content": [], "stop_reason": None}, + }, + { + "type": "error", + "error": {"type": "overloaded_error", "message": "overloaded"}, + }, + ] + body = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + + def handler(request): + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + async def invoke_async(): + async with ( + anthropic.AsyncAnthropic( + api_key="sk-ant-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) as client, + client.messages.stream( + model="claude-haiku-4-5", max_tokens=32, messages=[] + ) as stream, + ): + await stream.get_final_message() + + def invoke_sync(): + with ( + _client(handler) as client, + client.messages.stream( + model="claude-haiku-4-5", max_tokens=32, messages=[] + ) as stream, + ): + stream.get_final_message() + + if asynchronous: + with pytest.raises(anthropic.APIError): + await invoke_async() + else: + with pytest.raises(anthropic.APIError): + invoke_sync() + error = compat_client.one("error")["error"] + assert error["provider_error_code"] == "overloaded_error" + assert not compat_client.of_type("inference") -def test_anthropic_import_and_instrument(compat_client): - anthropic = pytest.importorskip("anthropic") - assert getattr(anthropic, "__version__", None) +def test_replacement_client_keeps_native_helper_ownership(compat_client): compat_client.instrument("anthropic") + old = _client(lambda request: httpx.Response(200, json=MESSAGE)) + old.messages.create(model="claude-haiku-4-5", max_tokens=32, messages=[]) + compat_client.close() + replacement = make_compat_client() + try: + replacement.instrument("anthropic") + old.messages.create(model="claude-haiku-4-5", max_tokens=32, messages=[]) + events = [ + { + "type": "message_start", + "message": {**MESSAGE, "content": [], "stop_reason": None}, + }, + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + {"type": "message_stop"}, + ] + body = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + with ( + _client( + lambda request: httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + ) as new, + new.messages.stream( + model="claude-haiku-4-5", max_tokens=32, messages=[] + ) as stream, + ): + assert stream.get_final_message().usage.output_tokens == 3 + assert len(compat_client.of_type("inference")) == 1 + replacement.one("inference") + finally: + old.close() + replacement.close() diff --git a/tests/compat/test_faster_whisper_compat.py b/tests/compat/test_faster_whisper_compat.py new file mode 100644 index 0000000..8578b78 --- /dev/null +++ b/tests/compat/test_faster_whisper_compat.py @@ -0,0 +1,53 @@ +"""Real CTranslate2 model load and lazy transcription using Whisper tiny.""" + +import os +from pathlib import Path + +import faster_whisper +import numpy as np +import pytest +from huggingface_hub import snapshot_download + + +@pytest.fixture(scope="module") +def model_path(): + # A local model can be supplied for offline development. CI caches HF downloads. + path = os.environ.get("WILDEDGE_FASTER_WHISPER_MODEL") or snapshot_download( + "Systran/faster-whisper-tiny", + revision="d90ca5fe260221311c53c58e660288d3deb8d356", + allow_patterns=["config.json", "model.bin", "tokenizer.json", "vocabulary.*"], + ) + assert Path(path).is_dir() + return path + + +def test_load_and_lazy_transcribe(compat_client, model_path): + compat_client.instrument("faster_whisper") + model = faster_whisper.WhisperModel(model_path, device="cpu", compute_type="int8") + load = compat_client.one("model_load") + if "WILDEDGE_FASTER_WHISPER_MODEL" not in os.environ: + # Named after the hub repo, not the snapshot hash in the cache path. + assert load["model_id"] == "faster_whisper:Systran/faster-whisper-tiny" + segments, info = model.transcribe( + np.zeros(16000, dtype=np.float32), language="en", beam_size=1 + ) + assert not compat_client.of_type("inference") + text = " ".join(segment.text for segment in segments) + event = compat_client.one("inference") + assert event["model_id"] == load["model_id"] + assert event["attributes"]["transcript_word_count"] == len(text.split()) + assert event["attributes"]["language_confidence"] == info.language_probability + assert event["inference"]["input_meta"]["duration_ms"] == 1000 + + +def test_batched_pipeline(compat_client, model_path): + compat_client.instrument("faster_whisper") + model = faster_whisper.WhisperModel(model_path, device="cpu", compute_type="int8") + load = compat_client.one("model_load") + pipeline = faster_whisper.BatchedInferencePipeline(model) + segments, _info = pipeline.transcribe( + np.zeros(16000, dtype=np.float32), language="en", batch_size=2 + ) + list(segments) + event = compat_client.one("inference") + assert event["model_id"] == load["model_id"] diff --git a/tests/compat/test_gguf_compat.py b/tests/compat/test_gguf_compat.py deleted file mode 100644 index 9512db6..0000000 --- a/tests/compat/test_gguf_compat.py +++ /dev/null @@ -1,9 +0,0 @@ -from __future__ import annotations - -import pytest - - -def test_gguf_import_and_instrument(compat_client): - llama_cpp = pytest.importorskip("llama_cpp") - assert hasattr(llama_cpp, "Llama") - compat_client.instrument("gguf") diff --git a/tests/compat/test_llama_cpp_compat.py b/tests/compat/test_llama_cpp_compat.py new file mode 100644 index 0000000..51cdd63 --- /dev/null +++ b/tests/compat/test_llama_cpp_compat.py @@ -0,0 +1,95 @@ +from __future__ import annotations + +import llama_cpp +import pytest +from conftest import downloaded_fixture + +# The 260K-parameter TinyStories Llama from karpathy/llama2.c (MIT), converted +# to GGUF by ggml-org. 1.2 MB. If the link breaks, point it at another copy. +MODEL_URL = ( + "https://huggingface.co/ggml-org/models-moved/resolve/" + "499bc8821c6b12b4e53c5bffcb21ec206f212d81/tinyllamas/stories260K.gguf" +) +MODEL_SHA256 = "270cba1bd5109f42d03350f60406024560464db173c0e387d91f0426d3bd256d" + + +@pytest.fixture +def llm(compat_client): + model = downloaded_fixture("stories260K-f32.gguf", MODEL_URL, MODEL_SHA256) + compat_client.instrument("llama_cpp") + return llama_cpp.Llama(model_path=str(model), n_ctx=128, seed=1, verbose=False) + + +def test_llama_load_is_tracked(compat_client, llm): + load = compat_client.one("model_load") + assert load["model_id"] == "stories260K-f32" + info = compat_client.registry.models["stories260K-f32"] + assert info.model_format == "gguf" + assert info.model_family == "llama" + assert info.quantization == "f32" + + +def test_llama_call_is_tracked(compat_client, llm): + result = llm("Once upon a time", max_tokens=4) + assert result["choices"][0]["text"] is not None + + inference = compat_client.one("inference") + assert inference["model_id"] == "stories260K-f32" + output = inference["inference"]["output_meta"] + assert output["tokens_in"] == result["usage"]["prompt_tokens"] + assert output["tokens_out"] == result["usage"]["completion_tokens"] + + +def test_create_completion_is_tracked_with_its_config(compat_client, llm): + result = llm.create_completion( + "Once upon a time", max_tokens=4, temperature=0.3, stop=["\n"] + ) + + inference = compat_client.one("inference")["inference"] + output = inference["output_meta"] + assert output["tokens_in"] == result["usage"]["prompt_tokens"] + assert output["tokens_out"] == result["usage"]["completion_tokens"] + assert output["stop_reason"] == result["choices"][0]["finish_reason"] + config = inference["generation_config"] + assert config["temperature"] == 0.3 + assert config["max_tokens"] == 4 + assert config["stop_sequences_count"] == 1 + # Arguments not passed are reported with llama-cpp-python's defaults. + assert "top_p" in config + assert "top_k" in config + + +def test_chat_completion_is_tracked_once_as_chat(compat_client, llm): + result = llm.create_chat_completion( + messages=[ + {"role": "user", "content": "Tell me a story"}, + {"role": "assistant", "content": "Once upon a time"}, + {"role": "user", "content": "Go on"}, + ], + max_tokens=4, + ) + + inference = compat_client.one("inference")["inference"] + assert inference["input_meta"]["prompt_type"] == "chat" + assert inference["input_meta"]["turn_index"] == 1 + assert inference["input_meta"]["word_count"] == 2 + output = inference["output_meta"] + assert output["tokens_in"] == result["usage"]["prompt_tokens"] + assert output["stop_reason"] == result["choices"][0]["finish_reason"] + assert inference["generation_config"]["max_tokens"] == 4 + + +def test_streamed_chat_completion_is_tracked_once(compat_client, llm): + stream = llm.create_chat_completion( + messages=[{"role": "user", "content": "Tell me a story"}], + max_tokens=4, + stream=True, + ) + chunks = list(stream) + assert chunks + + inference = compat_client.one("inference")["inference"] + assert inference["input_meta"]["prompt_type"] == "chat" + output = inference["output_meta"] + assert output["time_to_first_token_ms"] >= 0 + assert output["stop_reason"] in {"stop", "length"} diff --git a/tests/compat/test_mlx_compat.py b/tests/compat/test_mlx_compat.py index d66db56..96e6977 100644 --- a/tests/compat/test_mlx_compat.py +++ b/tests/compat/test_mlx_compat.py @@ -1,17 +1,60 @@ from __future__ import annotations +import json + +import mlx.core as mx +import mlx_lm import pytest +from mlx.utils import tree_flatten +from mlx_lm.models import llama +from tokenizers import Tokenizer, models, pre_tokenizers +from transformers import PreTrainedTokenizerFast + +VOCAB = ["", "", "", "once", "upon", "a", "time", "there", "was"] + + +@pytest.fixture(scope="module") +def tiny_model_dir(tmp_path_factory): + """A random-weight two-layer Llama with a word-level tokenizer, built locally.""" + path = tmp_path_factory.mktemp("tiny-llama") + config = { + "model_type": "llama", + "hidden_size": 16, + "num_hidden_layers": 2, + "intermediate_size": 32, + "num_attention_heads": 2, + "num_key_value_heads": 2, + "rms_norm_eps": 1e-5, + "vocab_size": len(VOCAB), + "tie_word_embeddings": True, + } + (path / "config.json").write_text(json.dumps(config)) + model = llama.Model(llama.ModelArgs.from_dict(config)) + mx.save_safetensors( + str(path / "model.safetensors"), dict(tree_flatten(model.parameters())) + ) + tokenizer = Tokenizer( + models.WordLevel({w: i for i, w in enumerate(VOCAB)}, unk_token="") + ) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + unk_token="", + bos_token="", + eos_token="", + ).save_pretrained(path) + return path -def test_mlx_import_and_instrumentation(compat_client): - pytest.importorskip("mlx_lm") - mx = pytest.importorskip("mlx.core") - nn = pytest.importorskip("mlx.nn") +def test_mlx_lm_load_and_generate_are_tracked(compat_client, tiny_model_dir): compat_client.instrument("mlx") + model, tokenizer = mlx_lm.load(str(tiny_model_dir)) + text = mlx_lm.generate(model, tokenizer, prompt="once upon a time", max_tokens=3) + assert isinstance(text, str) - model = nn.Linear(4, 2) - x = mx.ones((3, 4)) - y = model(x) - mx.eval(y) - assert y.shape == (3, 2) + load = compat_client.one("model_load") + inference = compat_client.one("inference") + assert inference["model_id"] == load["model_id"] + output = inference["inference"]["output_meta"] + assert output["tokens_out"] >= 1 diff --git a/tests/compat/test_mlx_whisper_compat.py b/tests/compat/test_mlx_whisper_compat.py new file mode 100644 index 0000000..386a05b --- /dev/null +++ b/tests/compat/test_mlx_whisper_compat.py @@ -0,0 +1,49 @@ +"""Real MLX Whisper loader and decoding, with locally built random weights.""" + +import importlib +import json +from dataclasses import asdict + +import mlx.core as mx +import mlx_whisper +import numpy as np +from mlx.utils import tree_flatten +from mlx_whisper.whisper import ModelDimensions, Whisper + + +def test_load_and_transcribe(compat_client, tmp_path): + dimensions = ModelDimensions( + n_mels=80, + n_audio_ctx=1500, + n_audio_state=8, + n_audio_head=2, + n_audio_layer=1, + n_vocab=51865, + n_text_ctx=448, + n_text_state=8, + n_text_head=2, + n_text_layer=1, + ) + mx.random.seed(0) + model = Whisper(dimensions, mx.float32) + (tmp_path / "config.json").write_text(json.dumps(asdict(dimensions))) + mx.save_safetensors( + str(tmp_path / "weights.safetensors"), dict(tree_flatten(model.parameters())) + ) + importlib.import_module("mlx_whisper.transcribe").ModelHolder.model = None + compat_client.instrument("mlx_whisper") + result = mlx_whisper.transcribe( + np.zeros(16000, dtype=np.float32), + path_or_hf_repo=str(tmp_path), + language="en", + fp16=False, + sample_len=2, + temperature=0, + condition_on_previous_text=False, + ) + load = compat_client.one("model_load") + event = compat_client.one("inference") + assert event["model_id"] == load["model_id"] + assert event["attributes"]["transcript_word_count"] == len(result["text"].split()) + assert event["inference"]["input_meta"]["duration_ms"] == 1000 + assert event["inference"]["input_meta"]["sample_rate"] == 16000 diff --git a/tests/compat/test_onnx_compat.py b/tests/compat/test_onnx_compat.py index ea4fc66..3436ce6 100644 --- a/tests/compat/test_onnx_compat.py +++ b/tests/compat/test_onnx_compat.py @@ -1,9 +1,46 @@ from __future__ import annotations -import pytest +import numpy as np +import onnxruntime as ort +from conftest import FIXTURES, make_compat_client -def test_onnxruntime_import_and_instrument(compat_client): - onnxruntime = pytest.importorskip("onnxruntime") - assert getattr(onnxruntime, "__version__", None) +def test_inference_session_load_and_run_are_tracked(compat_client): compat_client.instrument("onnx") + session = ort.InferenceSession( + str(FIXTURES / "identity.onnx"), providers=["CPUExecutionProvider"] + ) + x = np.arange(4, dtype=np.float32).reshape(1, 4) + (y,) = session.run(None, {"x": x}) + assert np.array_equal(x, y) + + # Auto-load takes the id from the file name, the rest from the session. + load = compat_client.one("model_load") + assert load["model_id"] == "identity" + info = compat_client.registry.models["identity"] + assert info.model_format == "onnx" + assert info.model_name == "compat_identity" + assert info.model_version == "3" + inference = compat_client.one("inference") + assert inference["model_id"] == "identity" + assert inference["inference"]["success"] is True + + +def test_replacement_client_tracks_new_sessions_only(compat_client): + compat_client.instrument("onnx") + path = str(FIXTURES / "identity.onnx") + old = ort.InferenceSession(path, providers=["CPUExecutionProvider"]) + x = np.arange(4, dtype=np.float32).reshape(1, 4) + old.run(None, {"x": x}) + compat_client.close() + replacement = make_compat_client() + try: + replacement.instrument("onnx") + old.run(None, {"x": x}) + new = ort.InferenceSession(path, providers=["CPUExecutionProvider"]) + assert np.array_equal(new.run(None, {"x": x})[0], x) + assert len(compat_client.of_type("inference")) == 1 + replacement.one("model_load") + replacement.one("inference") + finally: + replacement.close() diff --git a/tests/compat/test_openai_compat.py b/tests/compat/test_openai_compat.py index a704734..8641959 100644 --- a/tests/compat/test_openai_compat.py +++ b/tests/compat/test_openai_compat.py @@ -1,9 +1,410 @@ from __future__ import annotations +import json +import runpy +from pathlib import Path + +import openai import pytest +from conftest import make_compat_client, sdk_http_module + +httpx = sdk_http_module("openai") + +COMPLETION = { + "id": "chatcmpl-compat", + "object": "chat.completion", + "created": 1700000000, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello there"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11}, +} + + +def _client(handler) -> openai.OpenAI: + return openai.OpenAI( + api_key="sk-compat", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + + +def test_chat_completion_is_tracked(compat_client): + compat_client.instrument("openai") + client = _client(lambda request: httpx.Response(200, json=COMPLETION)) + + result = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Say hello"}], + temperature=0.2, + ) + assert result.choices[0].message.content == "Hello there" + + inference = compat_client.one("inference") + assert inference["model_id"] == "gpt-4o-mini" + output = inference["inference"]["output_meta"] + assert output["tokens_in"] == 9 + assert output["tokens_out"] == 2 + assert output["stop_reason"] == "stop" + +def test_streamed_chat_completion_is_tracked(compat_client): + compat_client.instrument("openai") + chunks = [ + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hi"}}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + "choices": [], + "usage": {"prompt_tokens": 9, "completion_tokens": 1, "total_tokens": 10}, + }, + ] + base = { + "id": "c", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + } + body = "".join(f"data: {json.dumps({**base, **c})}\n\n" for c in chunks) + body += "data: [DONE]\n\n" + client = _client( + lambda request: httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + ) + + stream = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Say hi"}], + stream=True, + stream_options={"include_usage": True}, + ) + assert ( + "".join(c.choices[0].delta.content or "" for c in stream if c.choices) == "Hi" + ) + + inference = compat_client.one("inference") + output = inference["inference"]["output_meta"] + assert output["tokens_out"] == 1 + assert output["stop_reason"] == "stop" + assert output["time_to_first_token_ms"] >= 0 -def test_openai_import_and_instrument(compat_client): - openai = pytest.importorskip("openai") - assert getattr(openai, "__version__", None) + +def test_api_error_is_tracked(compat_client): compat_client.instrument("openai") + client = _client( + lambda request: httpx.Response( + 429, + json={ + "error": { + "message": "slow down", + "type": "rate_limit", + "code": "rate_limit_exceeded", + } + }, + ) + ) + try: + client.with_options(max_retries=0).chat.completions.create( + model="gpt-4o-mini", messages=[{"role": "user", "content": "x"}] + ) + except openai.RateLimitError: + pass + else: + raise AssertionError("expected RateLimitError") + + error = compat_client.one("error") + assert error["error"]["http_status"] == 429 + assert error["error"]["provider_error_code"] == "rate_limit_exceeded" + + +RESPONSE = { + "id": "resp-compat", + "object": "response", + "created_at": 1700000000, + "model": "gpt-4o-mini", + "status": "completed", + "error": None, + "incomplete_details": None, + "output": [], + "parallel_tool_calls": True, + "tools": [], + "tool_choice": "auto", + "usage": { + "input_tokens": 9, + "output_tokens": 2, + "total_tokens": 11, + "input_tokens_details": {"cached_tokens": 3}, + "output_tokens_details": {"reasoning_tokens": 1}, + }, +} + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +async def test_responses_is_tracked(compat_client, asynchronous, streaming): + compat_client.instrument("openai") + with openai.OpenAI(api_key="sk-compat") as probe: + if not hasattr(probe, "responses"): + return # The minimum supported openai version predates Responses. + + def handler(request): + assert request.url.path == "/v1/responses" + if not streaming: + return httpx.Response(200, json=RESPONSE) + events = [ + { + "type": "response.created", + "sequence_number": 0, + "response": {**RESPONSE, "status": "in_progress", "usage": None}, + }, + { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg-1", + "output_index": 0, + "content_index": 0, + "delta": "Hi", + "logprobs": [], + }, + {"type": "response.completed", "sequence_number": 2, "response": RESPONSE}, + ] + body = "".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n" for e in events) + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + if asynchronous: + client = openai.AsyncOpenAI( + api_key="sk-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + async with client: + result = await client.responses.create( + model="gpt-4o-mini", input="Say hello", stream=streaming + ) + if streaming: + assert [e.type async for e in result][-1] == "response.completed" + else: + assert result.id == RESPONSE["id"] + else: + with _client(handler) as client: + result = client.responses.create( + model="gpt-4o-mini", input="Say hello", stream=streaming + ) + if streaming: + assert [e.type for e in result][-1] == "response.completed" + else: + assert result.id == RESPONSE["id"] + event = compat_client.one("inference") + assert event["model_id"] == "gpt-4o-mini" + output = event["inference"]["output_meta"] + assert output["tokens_in"] == 9 + assert output["tokens_out"] == 2 + assert output["cached_input_tokens"] == 3 + assert output["reasoning_tokens_out"] == 1 + assert output["stop_reason"] == "completed" + + +async def test_responses_example(compat_client): + compat_client.instrument("openai") + with openai.OpenAI(api_key="sk-compat") as probe: + if not hasattr(probe, "responses"): + return + example = runpy.run_path( + str(Path(__file__).parents[2] / "examples" / "openai_responses_example.py") + ) + payload = { + **RESPONSE, + "output": [ + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [ + { + "type": "output_text", + "text": "Hello from Responses", + "annotations": [], + } + ], + } + ], + } + + def handler(request): + if not json.loads(request.content).get("stream"): + return httpx.Response(200, json=payload) + event = { + "type": "response.completed", + "sequence_number": 0, + "response": payload, + } + return httpx.Response( + 200, + text=f"event: response.completed\ndata: {json.dumps(event)}\n\n", + headers={"content-type": "text/event-stream"}, + ) + + example["run_sync"].__globals__["OpenAI"] = lambda: _client(handler) + example["run_async"].__globals__["AsyncOpenAI"] = lambda: openai.AsyncOpenAI( + api_key="sk-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + with compat_client.trace(run_id="responses-example-test"): + example["run_sync"]() + await example["run_async"]() + events = compat_client.of_type("inference") + assert len(events) == 4 + assert all(e["run_id"] == "responses-example-test" for e in events) + assert len({e["parent_span_id"] for e in events}) == 4 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("encoding", ["float", "base64"]) +async def test_embeddings_is_tracked(compat_client, asynchronous, encoding): + compat_client.instrument("openai") + + def handler(request): + assert request.url.path == "/v1/embeddings" + return httpx.Response( + 200, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [ + { + "object": "embedding", + "index": i, + "embedding": [0.25, 0.5, 0.75] + if encoding == "float" + else "AACAPgAAAD8AAEA/", + } + for i in range(2) + ], + "usage": {"prompt_tokens": 7, "total_tokens": 7}, + }, + ) + + if asynchronous: + client = openai.AsyncOpenAI( + api_key="sk-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + async with client: + result = await client.embeddings.create( + model="text-embedding-3-small", + input=["hello", "world"], + encoding_format=encoding, + ) + else: + with _client(handler) as client: + result = client.embeddings.create( + model="text-embedding-3-small", + input=["hello", "world"], + encoding_format=encoding, + ) + assert len(result.data) == 2 + event = compat_client.one("inference") + inference = event["inference"] + assert inference["batch_size"] == 2 + assert inference["input_meta"]["token_count"] == 7 + assert inference["output_meta"] == { + "task": "embedding", + "dimensions": 3, + "tokens_in": 7, + } + assert inference["api_meta"]["resolved_model_id"] == "text-embedding-3-small" + assert "embedding" not in inference["output_meta"] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_early_stream_close_is_tracked_once(compat_client, asynchronous): + compat_client.instrument("openai") + chunk = { + "id": "c", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {"content": "Hi"}}], + } + body = f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n" + + def handler(request): + return httpx.Response( + 200, text=body, headers={"content-type": "text/event-stream"} + ) + + if asynchronous: + async with openai.AsyncOpenAI( + api_key="sk-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) as client: + async with await client.chat.completions.create( + model="gpt-4o-mini", messages=[], stream=True + ) as stream: + assert (await anext(stream)).choices[0].delta.content == "Hi" + await stream.close() + else: + with _client(handler) as client: + with client.chat.completions.create( + model="gpt-4o-mini", messages=[], stream=True + ) as stream: + assert next(stream).choices[0].delta.content == "Hi" + stream.close() + inference = compat_client.one("inference")["inference"] + assert inference["success"] is False + assert inference["error_code"] == "STREAM_CLOSED" + assert inference["output_meta"]["time_to_first_token_ms"] is not None + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_replacement_client_gets_new_provider_clients( + compat_client, asynchronous +): + compat_client.instrument("openai") + + def handler(request): + return httpx.Response(200, json=COMPLETION) + + if asynchronous: + old = openai.AsyncOpenAI( + api_key="sk-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) + await old.chat.completions.create(model="gpt-4o-mini", messages=[]) + else: + old = _client(handler) + old.chat.completions.create(model="gpt-4o-mini", messages=[]) + compat_client.one("inference") + compat_client.close() + replacement = make_compat_client() + try: + replacement.instrument("openai") + if asynchronous: + await old.chat.completions.create(model="gpt-4o-mini", messages=[]) + async with openai.AsyncOpenAI( + api_key="sk-compat", + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + ) as new: + await new.chat.completions.create(model="gpt-4o-mini", messages=[]) + await new.chat.completions.create(model="gpt-4o-mini", messages=[]) + await old.close() + else: + old.chat.completions.create(model="gpt-4o-mini", messages=[]) + with _client(handler) as new: + new.chat.completions.create(model="gpt-4o-mini", messages=[]) + new.chat.completions.create(model="gpt-4o-mini", messages=[]) + old.close() + assert len(compat_client.of_type("inference")) == 1 + assert len(replacement.of_type("inference")) == 2 + assert ( + replacement.registry.handles["gpt-4o-mini"] + is not compat_client.registry.handles["gpt-4o-mini"] + ) + finally: + replacement.close() diff --git a/tests/compat/test_tensorflow_compat.py b/tests/compat/test_tensorflow_compat.py index 37a0299..82c0d1d 100644 --- a/tests/compat/test_tensorflow_compat.py +++ b/tests/compat/test_tensorflow_compat.py @@ -1,12 +1,10 @@ from __future__ import annotations -import pytest +import numpy as np +import tensorflow as tf -def test_tensorflow_import_and_instrumentation(compat_client): - np = pytest.importorskip("numpy") - tf = pytest.importorskip("tensorflow") - +def test_keras_model_load_and_call_are_tracked(compat_client): compat_client.instrument("tensorflow") model = compat_client.load( tf.keras.Sequential, @@ -18,3 +16,9 @@ def test_tensorflow_import_and_instrumentation(compat_client): ) out = model(np.random.randn(2, 8).astype(np.float32), training=False) assert tuple(out.shape) == (2, 2) + + load = compat_client.one("model_load") + assert load["model_id"] + inference = compat_client.one("inference") + assert inference["model_id"] == load["model_id"] + assert inference["inference"]["success"] is True diff --git a/tests/compat/test_timm_compat.py b/tests/compat/test_timm_compat.py index c5f75a6..2ee3553 100644 --- a/tests/compat/test_timm_compat.py +++ b/tests/compat/test_timm_compat.py @@ -1,14 +1,48 @@ from __future__ import annotations -import pytest +import timm +import torch -def test_timm_import_and_instrumentation(compat_client): - torch = pytest.importorskip("torch") - timm = pytest.importorskip("timm") - +def test_timm_create_model_and_forward_are_tracked(compat_client): compat_client.instrument("timm") model = timm.create_model("resnet18", pretrained=False) - x = torch.randn(1, 3, 224, 224) - y = model(x) + y = model(torch.randn(1, 3, 64, 64)) assert y.shape[0] == 1 + + load = compat_client.one("model_load") + assert load["load"]["duration_ms"] >= 0 + inference = compat_client.one("inference") + assert inference["model_id"] == load["model_id"] + assert inference["inference"]["input_modality"] == "image" + + +def test_timm_architectures_get_their_own_ids(compat_client): + compat_client.instrument("timm") + # Neither is loaded by another test, so each should add one model_load. + models = [ + timm.create_model(name, pretrained=False).eval() + for name in ("resnet34", "resnet50") + ] + with torch.no_grad(): + for model in models: + model(torch.randn(1, 3, 64, 64)) + ids = {e["model_id"] for e in compat_client.of_type("model_load")} + assert ids == {"resnet34", "resnet50"} + assert {e["model_id"] for e in compat_client.of_type("inference")} == ids + assert {compat_client.registry.models[name].model_name for name in ids} == ids + + +def test_repeated_architecture_reuses_identity(compat_client): + compat_client.instrument("timm") + models = [ + timm.create_model(model_name="resnet18", pretrained=False).eval() + for _ in range(2) + ] + with torch.no_grad(): + for model in models: + model(torch.randn(1, 3, 64, 64)) + events = compat_client.of_type("inference") + assert len(events) == 2 + assert {event["model_id"] for event in events} == {"resnet18"} + assert compat_client.registry.models["resnet18"].model_name == "resnet18" diff --git a/tests/compat/test_torch_compat.py b/tests/compat/test_torch_compat.py index 2ea558e..f3f2a9b 100644 --- a/tests/compat/test_torch_compat.py +++ b/tests/compat/test_torch_compat.py @@ -1,11 +1,16 @@ from __future__ import annotations -import pytest +import torch -def test_torch_import_and_load_tracking(compat_client): - torch = pytest.importorskip("torch") +def test_torch_load_and_inference_are_tracked(compat_client): model = compat_client.load(torch.nn.Linear, 4, 2) - x = torch.randn(3, 4) - y = model(x) + y = model(torch.randn(3, 4)) assert tuple(y.shape) == (3, 2) + + load = compat_client.one("model_load") + assert load["load"]["duration_ms"] >= 0 + inference = compat_client.one("inference") + assert inference["model_id"] == load["model_id"] + assert inference["inference"]["success"] is True + assert inference["inference"]["batch_size"] == 3 diff --git a/tests/compat/test_transformers_compat.py b/tests/compat/test_transformers_compat.py index 01b22cf..0ddb803 100644 --- a/tests/compat/test_transformers_compat.py +++ b/tests/compat/test_transformers_compat.py @@ -1,25 +1,52 @@ from __future__ import annotations -import pytest +import torch +import transformers -def test_transformers_import_and_instrumentation(compat_client): - torch = pytest.importorskip("torch") - transformers = pytest.importorskip("transformers") - - compat_client.instrument("transformers") - +def _save_tiny_bert(path) -> None: config = transformers.BertConfig( - hidden_size=32, - num_hidden_layers=2, + vocab_size=32, + hidden_size=16, + num_hidden_layers=1, num_attention_heads=2, - intermediate_size=64, + intermediate_size=32, num_labels=2, ) - model = transformers.BertForSequenceClassification(config) - inputs = { - "input_ids": torch.zeros((1, 4), dtype=torch.long), - "attention_mask": torch.ones((1, 4), dtype=torch.long), - } - out = model(**inputs) + transformers.BertForSequenceClassification(config).save_pretrained(path) + + +def test_from_pretrained_load_and_forward_are_tracked(compat_client, tmp_path): + _save_tiny_bert(tmp_path) + compat_client.instrument("transformers") + + model = transformers.AutoModelForSequenceClassification.from_pretrained(tmp_path) + out = model( + input_ids=torch.zeros((1, 4), dtype=torch.long), + attention_mask=torch.ones((1, 4), dtype=torch.long), + ) assert tuple(out.logits.shape) == (1, 2) + + load = compat_client.one("model_load") + inference = compat_client.one("inference") + assert inference["model_id"] == load["model_id"] + assert inference["inference"]["success"] is True + + +def test_pipeline_load_and_call_are_tracked(compat_client, tmp_path): + _save_tiny_bert(tmp_path) + vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world"] + (tmp_path / "vocab.txt").write_text("\n".join(vocab) + "\n") + transformers.BertTokenizer(str(tmp_path / "vocab.txt")).save_pretrained(tmp_path) + compat_client.instrument("transformers") + + classifier = transformers.pipeline( + "text-classification", model=str(tmp_path), tokenizer=str(tmp_path) + ) + result = classifier("hello world") + assert result[0]["label"] in {"LABEL_0", "LABEL_1"} + + load = compat_client.one("model_load") + inference = compat_client.one("inference") + assert inference["model_id"] == load["model_id"] + assert inference["inference"]["input_modality"] == "text" diff --git a/tests/compat/test_ultralytics_compat.py b/tests/compat/test_ultralytics_compat.py index 6f6e5e5..818aada 100644 --- a/tests/compat/test_ultralytics_compat.py +++ b/tests/compat/test_ultralytics_compat.py @@ -1,9 +1,47 @@ from __future__ import annotations import pytest +import ultralytics +from PIL import Image -def test_ultralytics_import_and_instrument(compat_client): - ultralytics = pytest.importorskip("ultralytics") - assert getattr(ultralytics, "__version__", None) +@pytest.fixture +def yolo(compat_client): compat_client.instrument("ultralytics") + # Built from the architecture config, so no weights are downloaded. + return ultralytics.YOLO("yolov8n.yaml") + + +@pytest.fixture +def image(tmp_path) -> str: + # A file path is a source every supported ultralytics release accepts. + path = tmp_path / "blank.png" + Image.new("RGB", (64, 64)).save(path) + return str(path) + + +def test_yolo_load_is_tracked(compat_client, yolo): + compat_client.one("model_load") + assert compat_client.of_type("model_download") == [] + + +def test_yolo_call_is_tracked(compat_client, yolo, image): + results = yolo(image, imgsz=64, verbose=False) + assert len(results) == 1 + + inference = compat_client.one("inference") + assert inference["model_id"] in compat_client.registry.handles + assert inference["inference"]["input_modality"] == "image" + assert inference["inference"]["success"] is True + + +def test_yolo_predict_is_tracked(compat_client, yolo, image): + yolo.predict(image, imgsz=64, verbose=False) + compat_client.one("inference") + + +def test_predict_and_call_each_track_once(compat_client, yolo, image): + yolo.predict(image, imgsz=64, verbose=False) + assert len(compat_client.of_type("inference")) == 1 + yolo(image, imgsz=64, verbose=False) + assert len(compat_client.of_type("inference")) == 2 diff --git a/tests/conftest.py b/tests/conftest.py index 3c8890a..8b011f3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,23 +6,25 @@ import pytest +import wildedge.hubs.huggingface as huggingface +import wildedge.instrumentation as instrumentation +import wildedge.paths as paths from wildedge.client import WildEdge +from wildedge.defaults import set_default_client from wildedge.model import ModelInfo +from wildedge.platforms import stop_sampler from wildedge.platforms.device_info import DeviceInfo @pytest.fixture(autouse=True) def reset_hardware_sampler(): yield - from wildedge.platforms import stop_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) @@ -37,7 +39,6 @@ def isolate_sdk_state(tmp_path, monkeypatch): between runs and machines (and once masking a real regression locally while CI failed). """ - import wildedge.paths as paths monkeypatch.setattr(paths, "default_sdk_state_dir", lambda: tmp_path / "state") monkeypatch.setattr(paths, "default_sdk_cache_dir", lambda: tmp_path / "cache") @@ -103,7 +104,7 @@ def publish(event): @pytest.fixture def dummy_handle(): - handle = SimpleNamespace( + return SimpleNamespace( model_id="model-1", track_download=MagicMock(), track_load=MagicMock(), @@ -111,7 +112,6 @@ def dummy_handle(): track_inference=MagicMock(), track_error=MagicMock(), ) - return handle @pytest.fixture @@ -121,8 +121,15 @@ def client_with_stubbed_runtime(): patch("wildedge.client.Transmitter"), patch("wildedge.client.Consumer"), ): - client = WildEdge( + return WildEdge( dsn="https://secret@ingest.wildedge.dev/key", sampling_interval_s=None, ) - return client + + +@pytest.fixture(autouse=True) +def isolate_instrumentation_owner(monkeypatch): + monkeypatch.setattr(instrumentation, "_owner", None) + monkeypatch.setattr(instrumentation, "_enabled", set()) + monkeypatch.setattr(instrumentation, "_strict_ownership", False) + monkeypatch.setattr(huggingface, "_client_route", None) diff --git a/tests/test_accelerators.py b/tests/test_accelerators.py new file mode 100644 index 0000000..dccf2c8 --- /dev/null +++ b/tests/test_accelerators.py @@ -0,0 +1,80 @@ +"""Accelerator names are shared across integrations and reported consistently.""" + +import sys +from types import ModuleType, SimpleNamespace +from unittest.mock import patch + +import pytest + +from wildedge import accelerators +from wildedge.integrations.common import ( + mlx_accelerator, + tf_weights_accelerator, + torch_module_accelerator, +) +from wildedge.integrations.mlx import MlxExtractor + + +def torch_like(*device_types): + params = [SimpleNamespace(device=SimpleNamespace(type=t)) for t in device_types] + return SimpleNamespace(parameters=lambda: iter(params)) + + +@pytest.mark.parametrize("device_type", ["cpu", "cuda", "mps", "xpu"]) +def test_torch_module_reports_its_parameter_device(device_type): + assert torch_module_accelerator(torch_like(device_type)) == device_type + + +def test_torch_module_without_parameters_has_no_device(): + assert torch_module_accelerator(torch_like()) is None + + +@pytest.mark.parametrize( + ("device", "expected"), + [ + ("/job:localhost/replica:0/task:0/device:GPU:0", accelerators.CUDA), + ("/job:localhost/replica:0/task:0/device:CPU:0", accelerators.CPU), + ], +) +def test_tf_weights_report_the_platform_gpu(device, expected): + model = SimpleNamespace(weights=[SimpleNamespace(device=device)]) + with patch( + "wildedge.platforms.CURRENT_PLATFORM.gpu_accelerator_for_offload", + return_value=accelerators.CUDA, + ): + assert tf_weights_accelerator(model) == expected + + +def test_tf_model_without_weights_is_cpu(): + assert tf_weights_accelerator(SimpleNamespace(weights=[])) == accelerators.CPU + + +@pytest.fixture +def fake_mlx(monkeypatch): + mlx, core = ModuleType("mlx"), ModuleType("mlx.core") + core.cpu, core.gpu = "cpu-device", "gpu-device" + core.default = core.gpu + core.default_device = lambda: core.default + mlx.core = core + monkeypatch.setitem(sys.modules, "mlx", mlx) + monkeypatch.setitem(sys.modules, "mlx.core", core) + return core + + +def test_mlx_reports_the_apple_gpu_by_default(fake_mlx): + assert mlx_accelerator() == accelerators.MPS + + +def test_mlx_follows_a_cpu_default_device(fake_mlx): + fake_mlx.default = fake_mlx.cpu + assert mlx_accelerator() == accelerators.CPU + + +def test_mlx_lm_models_report_their_device(fake_mlx): + class Model: + def __call__(self, *args): + return None + + handle = SimpleNamespace(detected_accelerator=accelerators.CPU) + MlxExtractor().install_hooks(Model(), handle) + assert handle.detected_accelerator == accelerators.MPS diff --git a/tests/test_anthropic_native_stream.py b/tests/test_anthropic_native_stream.py new file mode 100644 index 0000000..a5f8695 --- /dev/null +++ b/tests/test_anthropic_native_stream.py @@ -0,0 +1,71 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from wildedge.integrations.anthropic import AnthropicStreamManager, wrap_native_stream + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_request_failure_preserves_exception_when_telemetry_fails(asynchronous): + original_error = TimeoutError("request timeout") + handle = MagicMock() + handle.track_error.side_effect = ValueError("telemetry failed") + original = MagicMock() + original.__enter__.side_effect = original_error + original.__aenter__.side_effect = original_error + manager = AnthropicStreamManager(original, handle, []) + + async def enter(): + return await manager.__aenter__() if asynchronous else manager.__enter__() + + with pytest.raises(TimeoutError) as caught: + await enter() + assert caught.value is original_error + handle.track_error.assert_called_once() + + +async def test_cancelled_before_entry_preserves_cancellation(): + original = MagicMock() + original.__aenter__.side_effect = asyncio.CancelledError() + handle = MagicMock() + handle.track_inference.side_effect = ValueError("telemetry failed") + manager = AnthropicStreamManager(original, handle, []) + with pytest.raises(asyncio.CancelledError): + await manager.__aenter__() + assert handle.track_inference.call_args.kwargs["error_code"] == "STREAM_CANCELLED" + + +async def test_cancelled_inside_native_context(): + async def chunks(): + yield SimpleNamespace(type="message_start") + + native = SimpleNamespace(_iterator=chunks(), close=AsyncMock()) + original = MagicMock() + original.__aenter__.return_value = native + + async def exit_context(*args): + await native.close() + + original.__aexit__.side_effect = exit_context + handle = MagicMock() + manager = AnthropicStreamManager(original, handle, []) + + async def consume(): + async with manager as stream: + assert stream is native + await anext(stream._iterator) + raise asyncio.CancelledError + + with pytest.raises(asyncio.CancelledError): + await consume() + handle.track_inference.assert_called_once() + assert handle.track_inference.call_args.kwargs["error_code"] == "STREAM_CANCELLED" + + +def test_closed_client_returns_original_manager(): + original_manager = object() + resource = SimpleNamespace(stream=MagicMock(return_value=original_manager)) + wrap_native_stream(resource, lambda: SimpleNamespace(closed=True)) + assert resource.stream(model="model") is original_manager diff --git a/tests/test_attachments.py b/tests/test_attachments.py index 46ff061..0616c22 100644 --- a/tests/test_attachments.py +++ b/tests/test_attachments.py @@ -2,10 +2,14 @@ import json import time +import types import urllib.error import urllib.request from datetime import datetime, timezone +import pytest + +from wildedge import attachments, constants from wildedge.attachments import ( Attachment, AttachmentManager, @@ -15,7 +19,8 @@ PendingAttachment, UploadOutcome, ) -from wildedge.events.inference import InferenceEvent +from wildedge.client import WildEdge +from wildedge.events.inference import InferenceEvent, TextInputMeta from wildedge.model import ModelHandle, ModelInfo # --------------------------------------------------------------------------- # @@ -41,7 +46,8 @@ def test_store_file_strategy_roundtrip(tmp_path): assert len(pending) == 1 assert pending[0].attachment_id == "a1" assert pending[0].read_bytes() == b"hello" - assert pending[0].bin_path is not None and pending[0].bin_path.exists() + assert pending[0].bin_path is not None + assert pending[0].bin_path.exists() def test_store_inline_strategy_roundtrip(tmp_path): @@ -158,7 +164,7 @@ def test_manager_drops_oversized(tmp_path): def test_manager_filter_hook(tmp_path): - store, mgr = _manager( + _store, mgr = _manager( tmp_path, attachment_filter=lambda atts: [a for a in atts if a.role == "input"], ) @@ -199,8 +205,6 @@ def test_inference_event_omits_attachments_when_absent(): def test_text_input_meta_no_has_attachments(): - from wildedge.events.inference import TextInputMeta - assert not hasattr(TextInputMeta(char_count=1), "has_attachments") @@ -220,17 +224,23 @@ def _make_handle(tmp_path): return handle, store, published -def test_track_inference_wires_attachments(tmp_path): +@pytest.mark.parametrize( + "content_type", + ["text/plain", "image/jpeg", "audio/wav", "video/mp4", "application/octet-stream"], +) +def test_track_inference_wires_attachments(tmp_path, content_type): handle, store, published = _make_handle(tmp_path) handle.track_inference( duration_ms=10, input_modality="image", - attachments=[Attachment(content_type="image/jpeg", role="input", data=b"img")], + attachments=[ + Attachment(content_type=content_type, role="input", data=b"payload") + ], ) assert store.length() == 1 event = published[0] assert len(event["attachments"]) == 1 - assert event["attachments"][0]["content_type"] == "image/jpeg" + assert event["attachments"][0]["content_type"] == content_type # The buffered attachment is keyed to this inference event. assert store.peek_many(1)[0].inference_id == event["inference"]["inference_id"] @@ -361,7 +371,7 @@ def __init__(self, outcomes): self.outcomes = list(outcomes) self.seen: list[str] = [] - def transmit(self, pending): + def transmit(self, pending, deadline=None): self.seen.append(pending.attachment_id) return self.outcomes.pop(0) @@ -435,9 +445,6 @@ def test_uploader_evicts_stale(tmp_path): def test_client_attachments_disabled_by_default(monkeypatch): - from wildedge import constants - from wildedge.client import WildEdge - monkeypatch.delenv(constants.ENV_ATTACHMENTS_ENABLED, raising=False) client = WildEdge(dsn="https://test@test.com/key") try: @@ -448,8 +455,6 @@ def test_client_attachments_disabled_by_default(monkeypatch): def test_client_attachments_enabled_wires_capture(monkeypatch, tmp_path): - from wildedge.client import WildEdge - client = WildEdge( dsn="https://test@test.com/key", attachments_enabled=True, @@ -465,9 +470,6 @@ def test_client_attachments_enabled_wires_capture(monkeypatch, tmp_path): def test_client_attachments_enabled_via_env(monkeypatch, tmp_path): - from wildedge import constants - from wildedge.client import WildEdge - monkeypatch.setenv(constants.ENV_ATTACHMENTS_ENABLED, "1") client = WildEdge( dsn="https://test@test.com/key", @@ -478,3 +480,216 @@ def test_client_attachments_enabled_via_env(monkeypatch, tmp_path): assert client.attachment_manager is not None finally: client.close() + + +# --------------------------------------------------------------------------- # +# Upload on flush / close +# --------------------------------------------------------------------------- # + + +class FakeClock: + """Stands in for time.monotonic() in wildedge.attachments; tests move it.""" + + def __init__(self) -> None: + self.now = 1000.0 + + def monotonic(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch) -> FakeClock: + fake = FakeClock() + monkeypatch.setattr( + attachments, + "time", + types.SimpleNamespace( + monotonic=fake.monotonic, time=time.time, time_ns=time.time_ns + ), + ) + return fake + + +def test_transmit_past_deadline_is_transient_without_requests(monkeypatch, tmp_path): + monkeypatch.setattr(urllib.request, "urlopen", _fake_urlopen([])) + tx = AttachmentTransmitter(api_key="secret", host="https://ingest.example") + outcome = tx.transmit(_pending(tmp_path), deadline=time.monotonic() - 1) + assert outcome is UploadOutcome.TRANSIENT + + +def test_transmit_caps_request_timeout_at_deadline(monkeypatch, tmp_path, clock): + timeouts: list[float] = [] + responses = iter([PRESIGN_OK, _Resp(200)]) + + def opener(req, timeout=None): + timeouts.append(timeout) + clock.advance(0.25) + return next(responses) + + monkeypatch.setattr(urllib.request, "urlopen", opener) + tx = AttachmentTransmitter(api_key="secret", host="https://ingest.example") + outcome = tx.transmit(_pending(tmp_path), deadline=clock.now + 1.0) + assert outcome is UploadOutcome.UPLOADED + # Each request gets what is left of the deadline, below the 10 s default. + assert timeouts == [1.0, 0.75] + + +def test_drain_uploads_more_than_one_batch(tmp_path): + n = constants.ATTACHMENT_PRESIGN_BATCH_SIZE * 2 + 1 + store = _filled_store(tmp_path, n) + tx = _StubTransmitter([UploadOutcome.UPLOADED] * n) + up = AttachmentUploader(store, tx) + up.drain(timeout=5.0) + assert store.length() == 0 + assert up.uploaded_count == n + + +def test_drain_stops_at_transient(tmp_path): + store = _filled_store(tmp_path, 3) + tx = _StubTransmitter([UploadOutcome.UPLOADED, UploadOutcome.TRANSIENT]) + up = AttachmentUploader(store, tx) + up.drain(timeout=5.0) + assert [p.attachment_id for p in store.peek_many(10)] == ["a1", "a2"] + + +def test_drain_respects_timeout(tmp_path, clock): + store = _filled_store(tmp_path, 3) + seen: list[str] = [] + + class Slow: + def transmit(self, pending, deadline=None): + seen.append(pending.attachment_id) + clock.advance(0.1) + if deadline is not None and clock.now >= deadline: + return UploadOutcome.TRANSIENT + return UploadOutcome.UPLOADED + + up = AttachmentUploader(store, Slow()) + up.drain(timeout=0.15) + # The first upload ends inside the deadline, the second past it. + assert seen == ["a0", "a1"] + assert [p.attachment_id for p in store.peek_many(10)] == ["a1", "a2"] + + +def test_drain_waits_for_background_batch_within_timeout(tmp_path): + store = _filled_store(tmp_path, 1) + up = AttachmentUploader(store, _StubTransmitter([UploadOutcome.UPLOADED])) + up.process_lock.acquire() # the background thread is mid-batch + try: + t0 = time.monotonic() + up.drain(timeout=0.1) + assert time.monotonic() - t0 < 0.5 + finally: + up.process_lock.release() + assert store.length() == 1 # left for the thread, not sent twice + + +def test_drain_after_feature_disabled_is_noop(tmp_path): + store = _filled_store(tmp_path, 2) + tx = _StubTransmitter([UploadOutcome.FEATURE_DISABLED]) + up = AttachmentUploader(store, tx) + up.drain(timeout=1.0) + _append(store, b"x", "late") + up.drain(timeout=1.0) + assert tx.seen == ["a0"] + + +def _client_with_pending(tmp_path, n=2): + client = WildEdge( + dsn="https://test@test.com/key", + attachments_enabled=True, + attachment_dir=str(tmp_path), + sampling_interval_s=0, + ) + client.attachment_uploader.flush_interval_s = 3600 + tx = _StubTransmitter([UploadOutcome.UPLOADED] * n) + client.attachment_uploader.transmitter = tx + for i in range(n): + _append(client.attachment_uploader.store, b"x", f"c{i}") + return client, tx + + +def test_close_with_timeout_uploads_pending_attachments(tmp_path): + client, tx = _client_with_pending(tmp_path) + client.close(timeout=2.0) + assert tx.seen == ["c0", "c1"] + assert client.attachment_uploader.store.length() == 0 + + +def test_close_without_timeout_leaves_attachments_buffered(tmp_path): + client, tx = _client_with_pending(tmp_path) + client.close() + assert tx.seen == [] + assert client.attachment_uploader.store.length() == 2 + + +def test_flush_uploads_pending_attachments(tmp_path): + client, tx = _client_with_pending(tmp_path) + try: + client.flush(timeout=2.0) + assert tx.seen == ["c0", "c1"] + finally: + client.close() + + +def test_transmit_deadline_passing_after_presign_skips_put( + monkeypatch, tmp_path, clock +): + calls: list[str] = [] + + def opener(req, timeout=None): + calls.append(req.get_method()) + clock.advance(0.06) + return PRESIGN_OK + + monkeypatch.setattr(urllib.request, "urlopen", opener) + tx = AttachmentTransmitter(api_key="secret", host="https://ingest.example") + outcome = tx.transmit(_pending(tmp_path), deadline=clock.now + 0.05) + assert outcome is UploadOutcome.TRANSIENT + assert calls == ["POST"] + + +def test_drain_stops_once_store_is_empty(tmp_path): + n = constants.ATTACHMENT_PRESIGN_BATCH_SIZE + store = _filled_store(tmp_path, n) + tx = _StubTransmitter([UploadOutcome.UPLOADED] * n) + up = AttachmentUploader(store, tx) + up.drain(timeout=5.0) + assert store.length() == 0 + assert len(tx.seen) == n + + +def test_flush_with_attachments_disabled(monkeypatch): + monkeypatch.delenv(constants.ENV_ATTACHMENTS_ENABLED, raising=False) + client = WildEdge(dsn="https://test@test.com/key") + try: + client.flush(timeout=0.1) # no uploader; must not raise + finally: + client.close() + + +def test_flush_with_no_time_left_skips_attachments(tmp_path): + client, tx = _client_with_pending(tmp_path) + try: + client.flush(timeout=0) + assert tx.seen == [] + finally: + client.close() + + +def test_flush_survives_attachment_drain_error(tmp_path, caplog): + client, _ = _client_with_pending(tmp_path) + + def boom(timeout): + raise RuntimeError("disk gone") + + client.attachment_uploader.drain = boom + try: + with caplog.at_level("WARNING", logger="wildedge"): + client.flush(timeout=1.0) + assert any("disk gone" in r.getMessage() for r in caplog.records) + finally: + client.close() diff --git a/tests/test_autoload.py b/tests/test_autoload.py index 578af6a..7e699ff 100644 --- a/tests/test_autoload.py +++ b/tests/test_autoload.py @@ -4,12 +4,15 @@ import importlib import importlib.util +import os +import pathlib import subprocess import sys import textwrap from unittest.mock import patch import wildedge.autoload.sitecustomize as _sc_mod +from wildedge.runtime.bootstrap import RuntimeConfigError, RuntimeStrictIntegrationError _MARKER = _sc_mod._INSTALLED_MARKER @@ -138,8 +141,6 @@ def test_uvicorn_reload_worker_bootstraps(tmp_path): across exec and blocked the worker's bootstrap. sys.modules is not inherited across exec, so the worker re-initialises correctly. """ - import os - import pathlib autoload_dir = str( pathlib.Path( @@ -179,8 +180,6 @@ def test_uvicorn_reload_worker_bootstraps(tmp_path): def test_strict_missing_dsn_exits_120(monkeypatch): - from wildedge.runtime.bootstrap import RuntimeConfigError - monkeypatch.delitem(sys.modules, _MARKER, raising=False) monkeypatch.setenv("WILDEDGE_AUTOLOAD", "1") monkeypatch.setenv("WILDEDGE_STRICT", "1") @@ -200,7 +199,6 @@ def test_strict_missing_dsn_exits_120(monkeypatch): def test_strict_integration_error_exits_121_without_strict_env(monkeypatch): """--strict-integrations is its own opt-in; it exits even without --strict.""" - from wildedge.runtime.bootstrap import RuntimeStrictIntegrationError monkeypatch.delitem(sys.modules, _MARKER, raising=False) monkeypatch.setenv("WILDEDGE_AUTOLOAD", "1") @@ -236,8 +234,6 @@ def test_strict_internal_error_exits_122(monkeypatch): def test_non_strict_config_error_warns_and_continues(monkeypatch, capsys): - from wildedge.runtime.bootstrap import RuntimeConfigError - monkeypatch.delitem(sys.modules, _MARKER, raising=False) monkeypatch.setenv("WILDEDGE_AUTOLOAD", "1") monkeypatch.delenv("WILDEDGE_STRICT", raising=False) @@ -277,8 +273,6 @@ class FakeContext: def test_no_propagate_clears_runtime_env(monkeypatch): - import os - class FakeContext: debug = False print_startup_report = False @@ -300,8 +294,6 @@ class FakeContext: def test_strict_missing_dsn_exit_code_in_real_interpreter(tmp_path): """sys.exit inside sitecustomize must terminate interpreter startup with our code.""" - import os - import pathlib autoload_dir = str( pathlib.Path( diff --git a/tests/test_build_changelog_comment.py b/tests/test_build_changelog_comment.py index f3c860c..fa5af30 100644 --- a/tests/test_build_changelog_comment.py +++ b/tests/test_build_changelog_comment.py @@ -3,9 +3,34 @@ import importlib.util from pathlib import Path +import pytest + SCRIPT = Path(__file__).parent.parent / "scripts" / "build_changelog_comment.py" -COMMITS = ["Add a thing (#1)", "Fix another thing (#2)"] +CHANGELOG = """# Changelog + +Intro text. + +## Unreleased + +- Not released yet. + +## 1.2.3 - 2026-10-21 + +### Added + +- A thing. + +### Fixed + +- Another thing. + +## 1.2.2 - 2026-09-30 + +### Fixed + +- An older thing. +""" def load_script(): @@ -15,58 +40,74 @@ def load_script(): return module +def test_notes_are_the_version_section_of_the_changelog(): + module = load_script() + notes = module.changelog_section(CHANGELOG, "1.2.3") + + assert notes == "### Added\n\n- A thing.\n\n### Fixed\n\n- Another thing." + + +def test_last_section_runs_to_the_end_of_the_file(): + module = load_script() + assert ( + module.changelog_section(CHANGELOG, "1.2.2") == "### Fixed\n\n- An older thing." + ) + + +def test_bracketed_version_headings_are_found(): + module = load_script() + text = "## [2.0.0] - 2026-11-01\n\n- New.\n" + assert module.changelog_section(text, "2.0.0") == "- New." + + +def test_version_prefix_does_not_match_a_longer_version(): + module = load_script() + with pytest.raises(SystemExit, match=r"no section for 1\.2$"): + module.changelog_section(CHANGELOG, "1.2") + + +def test_missing_or_empty_section_fails_the_build(): + module = load_script() + with pytest.raises(SystemExit, match=r"no section for 9\.9\.9"): + module.changelog_section(CHANGELOG, "9.9.9") + with pytest.raises(SystemExit, match="is empty"): + module.changelog_section("## 1.0.0\n\n## 0.9.0\n\n- Old.\n", "1.0.0") + + def test_preview_wording_on_pull_requests(): module = load_script() - text = module.build_comment("v1.2.3", COMMITS, "v1.2.2", preview=True) + text = module.build_comment("v1.2.3", "- A thing.", "v1.2.2", preview=True) assert text.startswith("## Changelog preview for `v1.2.3`") assert "> Preview of the release notes" in text + assert "- A thing." in text assert "**Full Changelog**: v1.2.2...v1.2.3" in text - for commit in COMMITS: - assert f"- {commit}" in text def test_release_notes_drop_the_preview_wording(): """Tag builds render the notes themselves, not a preview of them.""" module = load_script() - text = module.build_comment("v1.2.3", COMMITS, "v1.2.2", preview=False) + text = module.build_comment("v1.2.3", "- A thing.", "v1.2.2", preview=False) assert text.startswith("## Changelog for `v1.2.3`") assert "preview" not in text.lower() assert "**Full Changelog**: v1.2.2...v1.2.3" in text - for commit in COMMITS: - assert f"- {commit}" in text def test_first_release_falls_back_to_the_tag_alone(): module = load_script() - text = module.build_comment("v0.1.0", COMMITS, None, preview=False) + text = module.build_comment("v0.1.0", "- First.", None, preview=False) assert "**Full Changelog**: v0.1.0" in text -def test_empty_commit_range_is_reported(): - module = load_script() - text = module.build_comment("v1.2.3", [], "v1.2.2", preview=False) - - assert "- No changes." in text - - -def test_release_commits_are_filtered_from_the_log(monkeypatch): +def test_get_version_reads_pyproject_without_tomllib(): module = load_script() - log = "Release 1.2.3\nAdd a thing (#1)\n\nBump version to 1.2.3\n" - monkeypatch.setattr( - module.subprocess, - "run", - lambda *a, **k: type("R", (), {"stdout": log, "returncode": 0})(), - ) - - assert module.get_commits_since("v1.2.2") == [ - "Add a thing (#1)", - "Bump version to 1.2.3", - ] + assert module.get_version().count(".") >= 1 -def test_get_version_reads_pyproject_without_tomllib(): +def test_current_version_has_a_changelog_section(): + """Catches a release whose CHANGELOG.md entries were left under Unreleased.""" module = load_script() - assert module.get_version().count(".") >= 1 + changelog = (SCRIPT.parent.parent / "CHANGELOG.md").read_text() + assert module.changelog_section(changelog, module.get_version()) diff --git a/tests/test_cli.py b/tests/test_cli.py index 8bb3cf3..88e12bf 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -6,9 +6,13 @@ import pytest -from wildedge import cli, constants +from wildedge import cli, constants, instrumentation +from wildedge.defaults import peek_default_client from wildedge.integrations.registry import IntegrationSpec +from wildedge.platforms.device_info import DeviceInfo from wildedge.runtime import bootstrap +from wildedge.settings import parse_integration_list +from wildedge.transmitter import IngestResponse, TransmitError def _fake_execle(captured: dict): @@ -233,6 +237,36 @@ def instrument(self, name): # type: ignore[no-untyped-def] bootstrap.install_runtime() +@pytest.mark.parametrize("strict", ["0", "1"]) +def test_install_runtime_records_strict_ownership(monkeypatch, strict): + 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.setenv(constants.ENV_INTEGRATIONS, "onnx") + monkeypatch.setenv(constants.ENV_STRICT_INTEGRATIONS, strict) + monkeypatch.setattr(bootstrap.importlib.util, "find_spec", lambda _: object()) + + context = bootstrap.install_runtime() + try: + assert instrumentation.strict_ownership() is (strict == "1") + finally: + context.shutdown() + + def test_doctor_reports_missing_dsn_and_unknown_integration(monkeypatch, capsys): monkeypatch.delenv("WILDEDGE_DSN", raising=False) rc = cli.main(["doctor", "--integrations", "madeup"]) @@ -432,8 +466,6 @@ def test_parse_run_args_only_double_dash_raises(): def test_install_runtime_registers_default_client(monkeypatch): - from wildedge.defaults import peek_default_client - class FakeWildEdge: SUPPORTED_INTEGRATIONS = {"onnx"} @@ -467,14 +499,10 @@ def _doctor_env(monkeypatch): def _fake_device(api_key, app_version, overrides=None): - from wildedge.platforms.device_info import DeviceInfo - return DeviceInfo(device_id="d", device_type="linux") def test_doctor_send_test_event_accepted(monkeypatch, capsys): - from wildedge.transmitter import IngestResponse - _doctor_env(monkeypatch) monkeypatch.setattr(cli, "detect_device", _fake_device) sent_batches: list[dict] = [] @@ -504,8 +532,6 @@ def send(self, batch): def test_doctor_send_test_event_unauthorized_exits_2(monkeypatch, capsys): - from wildedge.transmitter import IngestResponse - _doctor_env(monkeypatch) monkeypatch.setattr(cli, "detect_device", _fake_device) @@ -530,8 +556,6 @@ def send(self, batch): def test_doctor_send_test_event_network_error_exits_2(monkeypatch, capsys): - from wildedge.transmitter import TransmitError - _doctor_env(monkeypatch) monkeypatch.setattr(cli, "detect_device", _fake_device) @@ -585,3 +609,12 @@ def test_doctor_json_includes_environment_and_statuses(monkeypatch, capsys): assert data["environment"]["autoload_on_pythonpath"] is False assert data["config_status"] == "PASS" assert data["connectivity_status"] == "SKIP" + + +@pytest.mark.parametrize("names", ["gguf", "llama_cpp", "gguf,llama_cpp"]) +def test_llama_cpp_names_resolve_once_in_cli_and_environment(names): + assert cli.integration_list(names) == ["llama_cpp"] + assert parse_integration_list(names, []) == ["llama_cpp"] + assert cli.INTEGRATIONS_BY_NAME["gguf"] is cli.INTEGRATIONS_BY_NAME["llama_cpp"] + assert "llama_cpp" in cli.integration_list("all") + assert "gguf" not in cli.integration_list("all") diff --git a/tests/test_client.py b/tests/test_client.py index 127dc29..4620aa5 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -4,7 +4,9 @@ import pytest +import wildedge.client as client_module from wildedge import constants +from wildedge.client import WildEdge @pytest.fixture(autouse=True) @@ -22,8 +24,6 @@ def mock_dependencies(): def test_batch_size_too_low(): - from wildedge.client import WildEdge - with pytest.raises( ValueError, match=f"batch_size must be between {constants.BATCH_SIZE_MIN} and {constants.BATCH_SIZE_MAX}", @@ -32,9 +32,6 @@ def test_batch_size_too_low(): def test_no_dsn_is_noop(monkeypatch, caplog): - import wildedge.client as client_module - from wildedge.client import WildEdge - monkeypatch.delenv(constants.ENV_DSN, raising=False) monkeypatch.setattr(client_module, "_noop_notice_logged", False) with caplog.at_level("INFO"): @@ -52,8 +49,6 @@ def test_no_dsn_is_noop(monkeypatch, caplog): def test_no_dsn_instrument_does_not_raise(monkeypatch): - from wildedge.client import WildEdge - monkeypatch.delenv(constants.ENV_DSN, raising=False) client = WildEdge() @@ -62,8 +57,6 @@ def test_no_dsn_instrument_does_not_raise(monkeypatch): def test_no_dsn_publish_does_not_enqueue(monkeypatch): - from wildedge.client import WildEdge - monkeypatch.delenv(constants.ENV_DSN, raising=False) client = WildEdge() @@ -72,8 +65,6 @@ def test_no_dsn_publish_does_not_enqueue(monkeypatch): def test_batch_size_too_high(): - from wildedge.client import WildEdge - with pytest.raises( ValueError, match=f"batch_size must be between {constants.BATCH_SIZE_MIN} and {constants.BATCH_SIZE_MAX}", @@ -82,8 +73,6 @@ def test_batch_size_too_high(): def test_flush_interval_too_low(): - from wildedge.client import WildEdge - with pytest.raises( ValueError, match=f"flush_interval_sec must be between {constants.FLUSH_INTERVAL_MIN} and {constants.FLUSH_INTERVAL_MAX}", @@ -92,8 +81,6 @@ def test_flush_interval_too_low(): def test_flush_interval_too_high(): - from wildedge.client import WildEdge - with pytest.raises( ValueError, match=f"flush_interval_sec must be between {constants.FLUSH_INTERVAL_MIN} and {constants.FLUSH_INTERVAL_MAX}", @@ -102,8 +89,6 @@ def test_flush_interval_too_high(): def test_max_queue_size_too_low(): - from wildedge.client import WildEdge - with pytest.raises( ValueError, match=f"max_queue_size must be between {constants.MAX_QUEUE_SIZE_MIN} and {constants.MAX_QUEUE_SIZE_MAX}", @@ -112,8 +97,6 @@ def test_max_queue_size_too_low(): def test_max_queue_size_too_high(): - from wildedge.client import WildEdge - with pytest.raises( ValueError, match=f"max_queue_size must be between {constants.MAX_QUEUE_SIZE_MIN} and {constants.MAX_QUEUE_SIZE_MAX}", @@ -122,8 +105,6 @@ def test_max_queue_size_too_high(): def test_valid_values(): - from wildedge.client import WildEdge - # Should not raise client = WildEdge( dsn="https://test@test.com/key", @@ -135,22 +116,16 @@ def test_valid_values(): def test_max_event_age_must_be_positive(): - from wildedge.client import WildEdge - with pytest.raises(ValueError, match="max_event_age_sec must be greater than 0"): WildEdge(dsn="https://test@test.com/key", max_event_age_sec=0) def test_max_dead_letter_batches_must_be_non_negative(): - from wildedge.client import WildEdge - with pytest.raises(ValueError, match="max_dead_letter_batches must be >= 0"): WildEdge(dsn="https://test@test.com/key", max_dead_letter_batches=-1) def test_app_identity_defaults_to_project_key(): - from wildedge.client import WildEdge - with ( patch( "wildedge.client.default_pending_queue_dir", return_value="pending-dir" @@ -167,8 +142,6 @@ def test_app_identity_defaults_to_project_key(): def test_app_identity_override_used_for_paths(): - from wildedge.client import WildEdge - with ( patch( "wildedge.client.default_pending_queue_dir", return_value="pending-dir" @@ -185,8 +158,6 @@ def test_app_identity_override_used_for_paths(): def test_app_identity_env_override_used_for_paths(monkeypatch): - from wildedge.client import WildEdge - monkeypatch.setenv(constants.ENV_APP_IDENTITY, "env-app") with ( patch( @@ -204,8 +175,6 @@ def test_app_identity_env_override_used_for_paths(monkeypatch): def test_register_model_defaults_name_to_model_id_and_accepts_format(): - from wildedge.client import WildEdge - client = WildEdge(dsn="https://secret@ingest.wildedge.dev/key") client.registry.register.return_value = (object(), True) @@ -220,8 +189,6 @@ def test_register_model_defaults_name_to_model_id_and_accepts_format(): def test_register_model_format_defaults_to_unknown(): - from wildedge.client import WildEdge - client = WildEdge(dsn="https://secret@ingest.wildedge.dev/key") client.registry.register.return_value = (object(), True) @@ -232,12 +199,11 @@ def test_register_model_format_defaults_to_unknown(): @pytest.mark.parametrize( - "class_name, source", + ("class_name", "source"), [("Messages", "anthropic"), ("Completions", "openrouter")], ) def test_register_model_api_integrations_use_their_extractor(class_name, source): """The anthropic and openai integrations register the patched resource object.""" - from wildedge.client import WildEdge client = WildEdge(dsn="https://secret@ingest.wildedge.dev/key") client.registry.register.return_value = (object(), True) diff --git a/tests/test_client_flows.py b/tests/test_client_flows.py index e0ab7ae..ede625b 100644 --- a/tests/test_client_flows.py +++ b/tests/test_client_flows.py @@ -13,9 +13,11 @@ def test_register_model_fallback_requires_id_when_no_extractor( client_with_stubbed_runtime, ): client = client_with_stubbed_runtime - with patch.object(client, "_find_extractor", return_value=None): - with pytest.raises(ValueError, match="Could not auto-derive a stable model_id"): - client.register_model(object()) + with ( + patch.object(client, "_find_extractor", return_value=None), + pytest.raises(ValueError, match="Could not auto-derive a stable model_id"), + ): + client.register_model(object()) def test_on_model_auto_loaded_uses_hub_records_when_downloads_missing( diff --git a/tests/test_compat_fixtures.py b/tests/test_compat_fixtures.py new file mode 100644 index 0000000..3effa8c --- /dev/null +++ b/tests/test_compat_fixtures.py @@ -0,0 +1,79 @@ +"""The download helper the compat tests use for test models.""" + +from __future__ import annotations + +import hashlib +import importlib.util +from pathlib import Path + +import pytest + +_spec = importlib.util.spec_from_file_location( + "compat_conftest", Path(__file__).parent / "compat" / "conftest.py" +) +compat_conftest = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(compat_conftest) + +URL = "https://example.invalid/model.bin" +DATA = b"tiny model" +SHA = hashlib.sha256(DATA).hexdigest() + + +@pytest.fixture +def cache(tmp_path, monkeypatch): + monkeypatch.setattr(compat_conftest, "DOWNLOAD_CACHE", tmp_path) + return tmp_path + + +@pytest.fixture +def downloads(monkeypatch): + calls: list[str] = [] + served = {"data": DATA, "error": None} + + def urlretrieve(url, filename): + calls.append(url) + if served["error"] is not None: + raise served["error"] + Path(filename).write_bytes(served["data"]) + + monkeypatch.setattr(compat_conftest.urllib.request, "urlretrieve", urlretrieve) + return calls, served + + +def test_downloads_and_verifies(cache, downloads): + calls, _ = downloads + path = compat_conftest.downloaded_fixture("model.bin", URL, SHA) + assert path == cache / "model.bin" + assert path.read_bytes() == DATA + assert calls == [URL] + assert not (cache / "model.bin.part").exists() + + +def test_reuses_a_cached_file(cache, downloads): + calls, _ = downloads + (cache / "model.bin").write_bytes(DATA) + compat_conftest.downloaded_fixture("model.bin", URL, SHA) + assert calls == [] + + +def test_replaces_a_corrupt_cached_file(cache, downloads): + calls, _ = downloads + (cache / "model.bin").write_bytes(b"truncated") + path = compat_conftest.downloaded_fixture("model.bin", URL, SHA) + assert path.read_bytes() == DATA + assert calls == [URL] + + +def test_wrong_checksum_fails_and_keeps_nothing(cache, downloads): + _, served = downloads + served["data"] = b"something else" + with pytest.raises(pytest.fail.Exception, match="expected " + SHA): + compat_conftest.downloaded_fixture("model.bin", URL, SHA) + assert list(cache.iterdir()) == [] + + +def test_download_error_fails_with_the_url(cache, downloads): + _, served = downloads + served["error"] = OSError("404 Not Found") + with pytest.raises(pytest.fail.Exception, match=URL): + compat_conftest.downloaded_fixture("model.bin", URL, SHA) diff --git a/tests/test_compat_matrix.py b/tests/test_compat_matrix.py new file mode 100644 index 0000000..67f9b9a --- /dev/null +++ b/tests/test_compat_matrix.py @@ -0,0 +1,134 @@ +"""Keep scripts/compat_matrix.py, tests/compat and the docs in step.""" + +from __future__ import annotations + +import importlib.util +import re +from pathlib import Path + +import pytest + +from wildedge.integrations.registry import INTEGRATION_SPECS + +ROOT = Path(__file__).resolve().parents[1] +COMPAT_TESTS = ROOT / "tests" / "compat" + +_spec = importlib.util.spec_from_file_location( + "compat_matrix", ROOT / "scripts" / "compat_matrix.py" +) +compat_matrix = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(compat_matrix) + + +def test_every_integration_has_a_compat_row(): + for spec in INTEGRATION_SPECS: + row = compat_matrix.COVERED_BY.get(spec.name, spec.name) + assert row in compat_matrix.MATRIX, f"{spec.name} has no compat row" + + +def test_covered_by_points_at_real_rows(): + for integration, row in compat_matrix.COVERED_BY.items(): + assert integration not in compat_matrix.MATRIX + assert row in compat_matrix.MATRIX + + +def test_every_row_has_a_test_file_and_supported_pythons(): + assert set(compat_matrix.SUPPORTED_PYTHON) == set(compat_matrix.MATRIX) + for integration in compat_matrix.MATRIX: + assert (COMPAT_TESTS / f"test_{integration}_compat.py").exists() + + +def test_no_compat_test_file_without_a_row(): + for path in COMPAT_TESTS.glob("test_*_compat.py"): + integration = path.name[len("test_") : -len("_compat.py")] + assert integration in compat_matrix.MATRIX, path.name + + +def test_compat_tests_never_skip(): + # A skip would turn a broken install into a green job. + for path in COMPAT_TESTS.glob("*.py"): + source = path.read_text() + assert "importorskip" not in source, path.name + assert not re.search(r"pytest\.skip\(|mark\.skip", source), path.name + + +def test_pins_are_exact(): + pin = re.compile(r"^[A-Za-z0-9_.-]+==[0-9][0-9A-Za-z.]*$") + for integration, sets in compat_matrix.MATRIX.items(): + for deps in sets.values(): + assert all(pin.match(d) for d in deps), (integration, deps) + for integration, by_set in compat_matrix.PYTHON_DEP_OVERRIDES.items(): + for version_set, by_python in by_set.items(): + assert version_set in compat_matrix.MATRIX[integration] + for py, deps in by_python.items(): + assert py in compat_matrix.SUPPORTED_PYTHON[integration] + assert all(pin.match(d) for d in deps), (integration, py, deps) + + +def test_docs_table_matches_the_matrix(capsys): + compat_matrix.print_table() + table = capsys.readouterr().out.strip() + docs = (ROOT / "docs" / "compatibility.md").read_text() + assert table in docs, ( + "docs/compatibility.md is out of date; regenerate it with " + "`python scripts/compat_matrix.py table`" + ) + + +def test_github_matrix_covers_every_combination(): + rows = compat_matrix.github_matrix() + keys = {(r["integration"], r["version-set"], r["python-version"]) for r in rows} + assert len(keys) == len(rows) + for integration, sets in compat_matrix.MATRIX.items(): + for version_set in sets: + for py in compat_matrix.SUPPORTED_PYTHON[integration]: + canary = py in compat_matrix.CANARY_PYTHON + if canary and version_set != "current": + assert (integration, version_set, py) not in keys + else: + assert (integration, version_set, py) in keys + for row in rows: + assert row["canary"] == (row["python-version"] in compat_matrix.CANARY_PYTHON) + expected_runner = compat_matrix.RUNNER.get(row["integration"], "ubuntu-latest") + assert row["runs-on"] == expected_runner + + +@pytest.fixture +def installed(monkeypatch): + versions: dict[str, str] = {} + + def version(name): + if name not in versions: + raise compat_matrix.metadata.PackageNotFoundError(name) + return versions[name] + + monkeypatch.setattr(compat_matrix.metadata, "version", version) + return versions + + +def test_verify_accepts_matching_versions(installed, capsys): + installed.update({"onnxruntime": "1.31.0", "numpy": "2.2.6"}) + assert compat_matrix.verify("onnx", "current", "3.12") == 0 + assert capsys.readouterr().err == "" + + +def test_verify_ignores_local_version_suffix(installed): + installed.update({"torch": "2.14.1+cpu", "numpy": "2.2.6"}) + assert compat_matrix.verify("torch", "current", "3.12") == 0 + + +def test_verify_reports_a_different_version(installed, capsys): + installed.update({"onnxruntime": "1.31.0", "numpy": "2.3.0"}) + assert compat_matrix.verify("onnx", "current", "3.12") == 1 + assert "numpy: installed 2.3.0, pinned 2.2.6" in capsys.readouterr().err + + +def test_verify_reports_a_missing_package(installed, capsys): + installed.update({"numpy": "2.2.6"}) + assert compat_matrix.verify("onnx", "current", "3.12") == 1 + assert "onnxruntime: not installed" in capsys.readouterr().err + + +def test_verify_uses_python_overrides(installed): + installed.update({"onnxruntime": "1.23.2", "numpy": "2.2.6"}) + assert compat_matrix.verify("onnx", "current", "3.10") == 0 diff --git a/tests/test_consumer.py b/tests/test_consumer.py index 0db0971..12846e0 100644 --- a/tests/test_consumer.py +++ b/tests/test_consumer.py @@ -1,11 +1,13 @@ """Tests for the background consumer.""" import os +import threading from unittest.mock import MagicMock, patch import pytest from wildedge import constants +from wildedge.client import WildEdge from wildedge.consumer import Consumer from wildedge.platforms.device_info import DeviceInfo from wildedge.queue import EventQueue @@ -42,7 +44,7 @@ def _make_consumer(self, queue, transmitter, **kwargs) -> Consumer: queue=queue, transmitter=transmitter, device=make_device(), - get_models=lambda: {}, + get_models=dict, session_id="sess-1", **kwargs, ) @@ -345,6 +347,48 @@ def fake_register(fn, *args, **kwargs): assert registered["fn"] == consumer.flush assert registered["args"] == (constants.DEFAULT_SHUTDOWN_FLUSH_TIMEOUT_SEC,) + def test_flush_does_not_resend_a_batch_already_in_flight(self, monkeypatch): + monkeypatch.setattr(Consumer, "run", lambda self: None) + queue = EventQueue(max_size=100) + for i in range(10): + queue.add({"event_id": str(i), "event_type": "inference", "model_id": "m"}) + sending = threading.Event() + release = threading.Event() + sent: list[str] = [] + + def slow_send(batch): + sent.extend(e["event_id"] for e in batch["events"]) + sending.set() + release.wait(5) + return make_success_response(len(batch["events"])) + + mock_transmitter = MagicMock(spec=Transmitter) + mock_transmitter.send.side_effect = slow_send + consumer = self._make_consumer(queue, mock_transmitter, batch_size=10) + + background = threading.Thread(target=consumer.drain_once) + background.start() + assert sending.wait(5) + flusher = threading.Thread(target=consumer.flush, kwargs={"timeout": 5}) + flusher.start() + release.set() + background.join(5) + flusher.join(5) + + assert sorted(sent, key=int) == [str(i) for i in range(10)] + assert queue.length() == 0 + + def test_drain_once_gives_up_when_lock_wait_times_out(self, monkeypatch): + monkeypatch.setattr(Consumer, "run", lambda self: None) + queue = EventQueue(max_size=100) + queue.add({"event_id": "e1", "event_type": "inference", "model_id": "m"}) + mock_transmitter = MagicMock(spec=Transmitter) + consumer = self._make_consumer(queue, mock_transmitter) + + with consumer.drain_lock: + assert consumer.drain_once(timeout=0.01) is False + mock_transmitter.send.assert_not_called() + class TestConsumerForkSafety: def setup_method(self): @@ -359,7 +403,7 @@ def _make_consumer(self, queue, transmitter, **kwargs) -> Consumer: queue=queue, transmitter=transmitter, device=DeviceInfo(app_version="1.0", device_id="d", device_type="linux"), - get_models=lambda: {}, + get_models=dict, session_id="sess-fork", **kwargs, ) @@ -441,6 +485,32 @@ def test_resume_registers_atexit(self, monkeypatch): assert any(fn == consumer.flush for fn, _ in registered) + def test_pause_holds_drain_lock_until_resume(self, monkeypatch): + monkeypatch.setattr(Consumer, "run", lambda self: None) + consumer = self._make_consumer( + EventQueue(max_size=100), MagicMock(spec=Transmitter) + ) + + consumer._pause() + assert consumer.drain_lock.locked() + consumer._resume() + assert not consumer.drain_lock.locked() + + def test_resume_replaces_lock_held_by_in_flight_send(self, monkeypatch): + monkeypatch.setattr(Consumer, "run", lambda self: None) + consumer = self._make_consumer( + EventQueue(max_size=100), MagicMock(spec=Transmitter) + ) + stuck = consumer.drain_lock + # A send that outlived _pause() still holds the lock. + stuck.acquire() + consumer.fork_holds_drain_lock = False + consumer._resume() + + assert consumer.drain_lock is not stuck + assert not consumer.drain_lock.locked() + stuck.release() + class TestForkRegistration: @pytest.mark.skipif( @@ -449,7 +519,6 @@ class TestForkRegistration: ) 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 registered = {} @@ -475,7 +544,6 @@ def fake_register_at_fork(*, before, after_in_child, after_in_parent): def test_register_at_fork_noop_on_windows(self, monkeypatch): """When os.register_at_fork is absent (Windows), __init__ must not raise.""" - from wildedge.client import WildEdge monkeypatch.delattr(os, "register_at_fork", raising=False) diff --git a/tests/test_decorators.py b/tests/test_decorators.py index 24ba728..94d48b6 100644 --- a/tests/test_decorators.py +++ b/tests/test_decorators.py @@ -56,7 +56,7 @@ def test_decorator_without_error_capture_does_not_track_error(self, publish_spy) def failing_func(): raise ValueError("oops") - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="oops"): failing_func() assert len(publish_spy.events) == 0 @@ -95,9 +95,11 @@ def test_context_manager_tracks_successful_inference(self, publish_spy): def test_context_manager_tracks_error_and_does_not_suppress(self, publish_spy): handle = make_handle(publish_spy) - with pytest.raises(ValueError, match="boom"): - with track(handle, capture_errors=True): - raise ValueError("boom") + with ( + pytest.raises(ValueError, match="boom"), + track(handle, capture_errors=True), + ): + raise ValueError("boom") assert len(publish_spy.events) == 1 event = publish_spy.events[0] diff --git a/tests/test_defaults.py b/tests/test_defaults.py index a8cbe9e..ff9e1b9 100644 --- a/tests/test_defaults.py +++ b/tests/test_defaults.py @@ -117,8 +117,10 @@ def test_trace_is_pure_context_and_creates_no_client(monkeypatch): 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") + with ( + wildedge.trace(run_id="r"), + 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_event_serialization.py b/tests/test_event_serialization.py index 9d59ae9..d6850e3 100644 --- a/tests/test_event_serialization.py +++ b/tests/test_event_serialization.py @@ -1,5 +1,7 @@ from __future__ import annotations +from wildedge import constants +from wildedge.events import common from wildedge.events.error import ErrorCode, ErrorEvent from wildedge.events.feedback import FeedbackEvent, FeedbackType from wildedge.events.inference import InferenceEvent, TextInputMeta @@ -76,6 +78,33 @@ def test_feedback_event_enum_and_string_forms(): assert string_event.to_dict()["feedback"]["feedback_type"] == "reject" +def test_error_event_caps_message_and_provider_code(): + event = ErrorEvent( + model_id="m1", + error_code=ErrorCode.UNKNOWN, + error_message="x" * 5000, + provider_error_code="c" * 500, + ).to_dict()["error"] + assert len(event["error_message"]) == constants.ERROR_MSG_MAX_LEN + assert len(event["provider_error_code"]) == constants.PROVIDER_ERROR_CODE_MAX_LEN + + +def test_long_error_message_keeps_start_and_end(): + message = "ffmpeg banner " * 100 + "Error opening input file audio.wav" + sent = ErrorEvent( + model_id="m1", error_code=ErrorCode.UNKNOWN, error_message=message + ).to_dict()["error"]["error_message"] + assert len(sent) == constants.ERROR_MSG_MAX_LEN + assert sent.startswith("ffmpeg banner") + assert sent.endswith("Error opening input file audio.wav") + assert " … " in sent + + +def test_short_error_message_is_unchanged(): + event = ErrorEvent(model_id="m1", error_code="UNKNOWN", error_message="boom") + assert event.to_dict()["error"]["error_message"] == "boom" + + def test_span_event_to_dict_includes_required_fields(): event = SpanEvent( kind="tool", @@ -132,3 +161,78 @@ def test_error_event_to_dict_http_status(): "duration_ms": 30_000, } assert "http_status" not in ErrorEvent("m1", "OOM").to_dict()["error"] + + +def test_span_event_sends_non_scalar_attributes_as_json_strings(): + event = SpanEvent( + kind="custom", + name="s", + duration_ms=1, + status="ok", + attributes={ + "rejections": ["English", "too long"], + "scores": {"a": 1}, + "count": 3, + "ratio": 0.5, + "ok": True, + "note": None, + "label": "x", + }, + ) + attrs = event.to_dict()["span"]["attributes"] + assert attrs == { + "rejections": '["English", "too long"]', + "scores": '{"a": 1}', + "count": 3, + "ratio": 0.5, + "ok": True, + "note": None, + "label": "x", + } + # The caller's dict is left as it was. + assert event.attributes["rejections"] == ["English", "too long"] + + +def test_span_attribute_flattening_handles_unserializable_values(): + class Thing: + def __str__(self): + return "thing" + + event = SpanEvent( + kind="custom", + name="s", + duration_ms=1, + status="ok", + attributes={"obj": Thing(), "tags": ("a", Thing())}, + ) + attrs = event.to_dict()["span"]["attributes"] + assert attrs == {"obj": '"thing"', "tags": '["a", "thing"]'} + + +def test_span_attribute_warning_is_logged_once_per_key(caplog): + common._warned_attribute_keys.discard("warn_once_tags") + with caplog.at_level("WARNING", logger="wildedge"): + for _ in range(3): + SpanEvent( + kind="custom", + name="s", + duration_ms=1, + status="ok", + attributes={"warn_once_tags": ["a"]}, + ).to_dict() + warnings = [r for r in caplog.records if "warn_once_tags" in r.getMessage()] + assert len(warnings) == 1 + + +def test_span_attribute_falls_back_to_str_when_json_fails(): + circular: list = [] + circular.append(circular) + event = SpanEvent( + kind="custom", + name="s", + duration_ms=1, + status="ok", + attributes={"loop": circular, 7: "int key"}, + ) + attrs = event.to_dict()["span"]["attributes"] + assert attrs == {"loop": "[[...]]", "7": "int key"} diff --git a/tests/test_hubs.py b/tests/test_hubs.py index bdb1108..c1c4c37 100644 --- a/tests/test_hubs.py +++ b/tests/test_hubs.py @@ -3,11 +3,14 @@ from __future__ import annotations import os +import sys import types from unittest.mock import patch import pytest +import wildedge.hubs.huggingface as hf_mod +import wildedge.hubs.torchhub as torchhub_mod from wildedge.hubs.huggingface import HuggingFaceHubTracker from wildedge.hubs.torchhub import TorchHubTracker @@ -145,8 +148,6 @@ def test_torch_hub_diff_to_records_empty_when_no_new_files(): def test_torch_hub_install_patch_is_idempotent(monkeypatch): - import wildedge.hubs.torchhub as torchhub_mod - original_load_calls = [] class FakeHub: @@ -179,8 +180,6 @@ def get_dir(): def test_hf_tracker_drain_returns_and_clears_buffer(monkeypatch): - import wildedge.hubs.huggingface as hf_mod - tracker = HuggingFaceHubTracker() # Directly inject a record into the thread-local buffer hf_mod._buffer().append({"repo_id": "test/model", "source_type": "huggingface"}) @@ -190,3 +189,35 @@ def test_hf_tracker_drain_returns_and_clears_buffer(monkeypatch): assert result[0]["repo_id"] == "test/model" # Buffer should be cleared assert tracker.drain() == [] + + +class _LazyModule(types.ModuleType): + """Imports on attribute access, like transformers' lazy modules.""" + + def __init__(self, name: str) -> None: + super().__init__(name) + self.lookups: list[str] = [] + + def __getattr__(self, name: str): + self.lookups.append(name) + raise ModuleNotFoundError("No module named 'torchvision'") + + +def test_patch_symbol_references_skips_lazy_attribute_lookup(monkeypatch): + def original(): + pass + + def replacement(): + pass + + eager = types.ModuleType("eager_mod") + eager.hf_hub_download = original + lazy = _LazyModule("lazy_mod") + monkeypatch.setitem(sys.modules, "eager_mod", eager) + monkeypatch.setitem(sys.modules, "lazy_mod", lazy) + + count = hf_mod._patch_symbol_references("hf_hub_download", original, replacement) + + assert count >= 1 + assert eager.hf_hub_download is replacement + assert lazy.lookups == [] diff --git a/tests/test_init.py b/tests/test_init.py index c8d6a7a..37eed7b 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -2,6 +2,7 @@ import wildedge import wildedge.convenience as convenience +from wildedge.defaults import peek_default_client, set_default_client def test_init_calls_instrument_for_integrations(monkeypatch): @@ -70,11 +71,11 @@ def instrument(self, integration, *, hubs=None): 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: + closed = False + def instrument(self, integration, *, hubs=None): calls.append((integration, hubs)) @@ -89,9 +90,9 @@ def instrument(self, integration, *, hubs=None): def test_init_with_dsn_replaces_default_client(monkeypatch): - from wildedge.defaults import peek_default_client, set_default_client - class ExistingClient: + closed = False + def instrument(self, integration, *, hubs=None): raise AssertionError("existing client must not be touched") @@ -114,9 +115,9 @@ def instrument(self, integration, *, hubs=None): def test_init_reuse_warns_about_ignored_kwargs(monkeypatch): - from wildedge.defaults import set_default_client - class ExistingClient: + closed = False + def instrument(self, integration, *, hubs=None): pass @@ -134,8 +135,6 @@ def instrument(self, integration, *, hubs=None): 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 diff --git a/tests/test_instrumentation_ownership.py b/tests/test_instrumentation_ownership.py new file mode 100644 index 0000000..45f3ad1 --- /dev/null +++ b/tests/test_instrumentation_ownership.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import gc +import threading +import weakref +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock, patch + +import pytest + +import wildedge +import wildedge.convenience as convenience +import wildedge.hubs.huggingface as huggingface +import wildedge.instrumentation as instrumentation +from wildedge.client import WildEdge +from wildedge.defaults import set_default_client +from wildedge.instrumentation import BoundClientReference, claim_instrumentation +from wildedge.platforms.device_info import DeviceInfo + + +class Owner: + closed = False + + +def test_sequential_replacement_and_opt_in(): + first = Owner() + route = claim_instrumentation(first, "onnx") + assert route() is first + assert claim_instrumentation(first, "onnx")() is first + first.closed = True + assert route() is None + second = Owner() + other = claim_instrumentation(second, "openai") + assert other() is second + assert route() is None + claim_instrumentation(second, "onnx") + assert route() is second + + +def test_live_owner_conflict_does_not_redirect(): + first, second = Owner(), Owner() + route = claim_instrumentation(first, "onnx") + with pytest.raises(RuntimeError, match="another live"): + claim_instrumentation(second, "openai") + assert route() is first + + +def test_dead_owner_can_be_replaced(): + first = Owner() + ref = weakref.ref(first) + route = claim_instrumentation(first, "onnx") + del first + gc.collect() + assert ref() is None + second = Owner() + claim_instrumentation(second, "onnx") + assert route() is second + + +def test_provider_reference_never_follows_replacement(caplog): + caplog.set_level("WARNING", logger="wildedge") + first, second = Owner(), Owner() + claim_instrumentation(first, "openai") + ref = BoundClientReference(first) + first.closed = True + claim_instrumentation(second, "openai") + assert ref() is None + assert ref() is None + assert len([r for r in caplog.records if "recreate" in r.message]) == 1 + + +@pytest.fixture +def clients(): + created = [] + with ( + patch( + "wildedge.client.detect_device", + return_value=DeviceInfo(device_id="replacement-test", device_type="linux"), + ), + patch("wildedge.client.Consumer"), + patch("wildedge.client.Transmitter"), + ): + + def create(): + client = WildEdge( + dsn="https://test@localhost/test", + sampling_interval_s=0, + enable_offline_persistence=False, + ) + created.append(client) + return client + + yield create + for client in created: + client.close() + + +@pytest.mark.parametrize("integration", sorted(WildEdge.PATCH_INSTALLERS)) +def test_shared_route_reuses_first_factory_binding(clients, monkeypatch, integration): + routes = [] + installer = MagicMock( + side_effect=lambda route: routes.append(route) if not routes else None + ) + monkeypatch.setitem(WildEdge.PATCH_INSTALLERS, integration, installer) + first = clients() + first.instrument(integration) + assert routes[0]() is first + first.close() + second = clients() + second.instrument(integration) + assert routes[0]() is second + assert len(routes) == 1 + + +def test_closed_model_handle_does_not_move_projects(clients, caplog): + caplog.set_level("WARNING", logger="wildedge") + first = clients() + handle = first.register_model(object(), model_id="model", auto_instrument=False) + first.close() + second = clients() + handle.track_inference(duration_ms=1) + handle.track_inference(duration_ms=2) + assert first.pending_count() == second.pending_count() == 0 + assert len([r for r in caplog.records if "closed SDK client" in r.message]) == 1 + + +def test_init_replaces_closed_default(clients, monkeypatch): + first = clients() + set_default_client(first) + first.close() + second = clients() + monkeypatch.setattr(convenience, "WildEdge", lambda **kwargs: second) + assert wildedge.init() is second + + +def test_close_is_idempotent(clients): + client = clients() + client.close() + client.close() + client.consumer.close.assert_called_once() + + +def test_init_with_other_dsn_warns_and_leaves_tracking_with_owner( + clients, monkeypatch, caplog +): + caplog.set_level("WARNING", logger="wildedge") + first = clients() + route = claim_instrumentation(first, "onnx") + set_default_client(first) + second = clients() + monkeypatch.setattr(convenience, "WildEdge", lambda **kwargs: second) + with patch("wildedge.client.stop_sampler") as stop: + result = wildedge.init( + dsn="https://other@localhost/other", integrations=["onnx"] + ) + stop.assert_not_called() + assert result is second + assert not second.closed + assert not first.closed + assert route() is first + assert convenience.peek_default_client() is second + assert [r for r in caplog.records if "already owns automatic tracking" in r.message] + + +def test_strict_ownership_conflict_raises_and_cleans_up_new_client( + clients, monkeypatch +): + instrumentation.set_strict_ownership(True) + first = clients() + claim_instrumentation(first, "onnx") + set_default_client(first) + second = clients() + monkeypatch.setattr(convenience, "WildEdge", lambda **kwargs: second) + with patch("wildedge.client.stop_sampler") as stop: + with pytest.raises(RuntimeError, match="another live"): + wildedge.init(dsn="https://other@localhost/other", integrations=["onnx"]) + stop.assert_not_called() + assert second.closed + assert not first.closed + assert convenience.peek_default_client() is first + + +def test_init_with_owner_dsn_reuses_owner(clients, monkeypatch, caplog): + caplog.set_level("WARNING", logger="wildedge") + first = clients() + first.instrument("onnx") + set_default_client(first) + monkeypatch.setattr( + convenience, "WildEdge", MagicMock(side_effect=AssertionError("new client")) + ) + assert ( + wildedge.init(dsn="https://test@localhost/test", integrations=["onnx"]) is first + ) + assert not [r for r in caplog.records if "ignoring" in r.message] + + +def test_init_with_dsn_none_reuses_owner_of_env_dsn(clients, monkeypatch): + first = clients() + set_default_client(first) + monkeypatch.setenv("WILDEDGE_DSN", "https://test@localhost/test") + monkeypatch.setattr( + convenience, "WildEdge", MagicMock(side_effect=AssertionError("new client")) + ) + assert wildedge.init(dsn=None) is first + + +def test_hub_conflict_warns_and_skips(clients, monkeypatch, caplog): + caplog.set_level("WARNING", logger="wildedge") + installs = [] + + class Tracker: + def can_install(self): + return True + + def install_patch(self, route): + installs.append(route) + + monkeypatch.setitem(WildEdge.HUB_TRACKER_CLASSES, "huggingface", Tracker) + first, second = clients(), clients() + first.instrument(None, hubs=["huggingface"]) + second.instrument(None, hubs=["huggingface"]) + assert len(installs) == 1 + assert "huggingface" not in second.hub_trackers + assert [r for r in caplog.records if "already owns automatic tracking" in r.message] + + +@pytest.mark.parametrize("hub", ["huggingface", "torchhub"]) +def test_hub_factory_route_follows_replacement(clients, monkeypatch, hub): + routes = [] + + class Tracker: + def can_install(self): + return True + + def install_patch(self, route): + if not routes: + routes.append(route) + + monkeypatch.setitem(WildEdge.HUB_TRACKER_CLASSES, hub, Tracker) + first = clients() + first.instrument(None, hubs=[hub]) + assert routes[0]() is first + first.close() + second = clients() + second.instrument(None, hubs=[hub]) + assert routes[0]() is second + assert first.hub_trackers[hub] is not second.hub_trackers[hub] + + +def test_changing_default_does_not_change_instrumentation_owner(clients): + first, second = clients(), clients() + route = claim_instrumentation(first, "openai") + set_default_client(second) + assert route() is first + + +def test_closed_client_cannot_instrument(clients): + client = clients() + client.close() + with pytest.raises(RuntimeError, match="closed"): + client.instrument("onnx") + + +def test_fork_child_replaces_inherited_lock(): + first = Owner() + route = claim_instrumentation(first, "onnx") + with instrumentation._lock: + instrumentation._reset_lock_after_fork() + assert route() is first + + +def test_concurrent_claims_choose_one_owner(): + owners = [Owner(), Owner()] + barrier = threading.Barrier(2, timeout=2) + + def acquire(owner): + barrier.wait() + try: + return claim_instrumentation(owner, "onnx") + except RuntimeError: + return None + + with ThreadPoolExecutor(max_workers=2) as executor: + results = list(executor.map(acquire, owners)) + routes = [route for route in results if route is not None] + assert len(routes) == 1 + assert routes[0]() in owners + + +def test_hub_provenance_buffer_does_not_cross_owners(clients, monkeypatch): + first = clients() + route = claim_instrumentation(first, "hub:huggingface") + monkeypatch.setattr(huggingface, "_client_route", route) + monkeypatch.setattr(huggingface, "_local", threading.local()) + record = { + "repo_id": "first-project/model", + "size": 1, + "duration_ms": 1, + "cache_hit": True, + "bandwidth_bps": None, + } + huggingface._record_download(**record) + first.close() + second = clients() + claim_instrumentation(second, "onnx") + huggingface._record_download(**record) + claim_instrumentation(second, "hub:huggingface") + assert huggingface.HuggingFaceHubTracker().drain() == [] + huggingface._record_download(**{**record, "repo_id": "second-project/model"}) + assert [r["repo_id"] for r in huggingface.HuggingFaceHubTracker().drain()] == [ + "second-project/model" + ] diff --git a/tests/test_integration_patching.py b/tests/test_integration_patching.py index 7390d4c..9482e29 100644 --- a/tests/test_integration_patching.py +++ b/tests/test_integration_patching.py @@ -5,6 +5,11 @@ import pytest +import wildedge.hubs.huggingface as hf_mod +import wildedge.integrations.gguf as gguf_mod +import wildedge.integrations.onnx as onnx_mod +import wildedge.integrations.pytorch as torch_mod +import wildedge.integrations.transformers as tf_mod from wildedge.client import WildEdge from wildedge.hubs.huggingface import HuggingFaceHubTracker from wildedge.integrations.gguf import GgufExtractor @@ -15,8 +20,6 @@ def test_hf_install_patch_is_idempotent(monkeypatch): - import wildedge.hubs.huggingface as hf_mod - def orig_hf_hub_download(repo_id, filename, **kwargs): return f"/tmp/{repo_id}/{filename}" @@ -43,8 +46,6 @@ def orig_hf_hub_download(repo_id, filename, **kwargs): def test_hf_install_patch_retries_unpatched_part(monkeypatch): - import wildedge.hubs.huggingface as hf_mod - monkeypatch.setattr(hf_mod, "_hf", object()) monkeypatch.setattr(hf_mod, "_fd", object()) monkeypatch.setattr(hf_mod, "_hf_hub_download_patched", False) @@ -66,8 +67,6 @@ def fake_install_hf(): def test_onnx_install_auto_load_patch_is_idempotent(monkeypatch): - import wildedge.integrations.onnx as onnx_mod - class FakeSession: def __init__(self, *args, **kwargs): pass @@ -87,8 +86,6 @@ def client_ref(): def test_timm_install_patch_is_idempotent(monkeypatch): - import wildedge.integrations.pytorch as torch_mod - def create_model(*args, **kwargs): return object() @@ -107,8 +104,6 @@ def client_ref(): def test_transformers_install_patch_is_idempotent(monkeypatch): - import wildedge.integrations.transformers as tf_mod - class FakePipeline: def __init__(self, *args, **kwargs): pass @@ -141,7 +136,6 @@ def client_ref(): def test_transformers_from_import_style_intercepted(monkeypatch): """from transformers import pipeline; pipeline() must be intercepted.""" - import wildedge.integrations.transformers as tf_mod loaded = [] @@ -164,7 +158,7 @@ def from_pretrained(cls, *args, **kwargs): client = types.SimpleNamespace( closed=False, - _snapshot_hub_caches=lambda: {}, + _snapshot_hub_caches=dict, _diff_hub_caches=lambda before, ms: None, _on_model_auto_loaded=lambda obj, **kw: loaded.append(("tracked", obj)), ) @@ -186,8 +180,6 @@ def from_pretrained(cls, *args, **kwargs): def test_gguf_install_auto_load_patch_is_idempotent(monkeypatch): - import wildedge.integrations.gguf as gguf_mod - class FakeLlama: def __init__(self, *args, **kwargs): pass @@ -225,7 +217,7 @@ def test_instrument_hubs_activates_requested_trackers(stub_client): activated = [] with ( patch.object(stub_client, "_activate_hub", side_effect=activated.append), - patch.dict(stub_client.PATCH_INSTALLERS, {"gguf": lambda ref: None}), + patch.dict(stub_client.PATCH_INSTALLERS, {"llama_cpp": lambda ref: None}), ): stub_client.instrument("gguf", hubs=["huggingface"]) @@ -252,3 +244,24 @@ def test_instrument_none_activates_hub(stub_client): with patch.object(stub_client, "_activate_hub", side_effect=activated.append): stub_client.instrument(None, hubs=["huggingface"]) assert activated == ["huggingface"] + + +@pytest.mark.parametrize("names", [("gguf", "llama_cpp"), ("llama_cpp", "gguf")]) +def test_llama_cpp_aliases_share_patch_and_client_route( + stub_client, monkeypatch, names +): + class FakeLlama: + def __init__(self, *args, **kwargs): + pass + + runtime = types.SimpleNamespace(Llama=FakeLlama) + monkeypatch.setattr(gguf_mod, "_llama_cpp", runtime) + monkeypatch.setattr(gguf_mod, "_llama_patched", False) + stub_client.instrument(names[0]) + first = runtime.Llama.__init__ + stub_client.instrument(names[1]) + assert runtime.Llama.__init__ is first + with patch.object(stub_client, "_on_model_auto_loaded") as loaded: + model = runtime.Llama("model.gguf") + loaded.assert_called_once() + assert loaded.call_args.args[0] is model diff --git a/tests/test_integrations.py b/tests/test_integrations.py index 9f6f7e3..1f678df 100644 --- a/tests/test_integrations.py +++ b/tests/test_integrations.py @@ -2,12 +2,19 @@ from __future__ import annotations +import inspect from unittest.mock import MagicMock, patch import pytest +from wildedge.attachments import AttachmentManager, AttachmentStore +from wildedge.integrations import gguf as gguf_integration from wildedge.integrations.gguf import ( GgufExtractor, + call_arguments, + make_gguf_generation_config, + make_gguf_output_meta, + make_gguf_stream_callbacks, parse_quantization, ) from wildedge.integrations.gguf import ( @@ -17,6 +24,9 @@ from wildedge.integrations.keras import detect_accelerator as keras_detect_accelerator from wildedge.integrations.onnx import ( OnnxExtractor, + model_file_stem, + model_id_from_path, + parse_model_version, ) from wildedge.integrations.onnx import ( detect_accelerator as onnx_detect_accelerator, @@ -75,16 +85,50 @@ def run(self, output_names, input_feed, run_options=None): class Llama: # must match the name check in GgufExtractor.can_handle + """Mirrors llama-cpp-python's call graph: every entry point runs create_completion.""" + model_path = "/models/llama-3-Q4_K_M.gguf" n_gpu_layers = 0 metadata = {"general.architecture": "llama", "general.version": "3.0"} - def __call__(self, *args, **kwargs): - return {"usage": {"prompt_tokens": 10, "completion_tokens": 20}} + def create_completion(self, prompt="", **kwargs): + if kwargs.get("stream"): + return iter( + [ + {"choices": [{"text": "hi", "finish_reason": None}]}, + { + "choices": [{"text": "!", "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + }, + ] + ) + return { + "choices": [{"text": "hi!", "finish_reason": "length"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 20}, + } + + def __call__(self, prompt="", **kwargs): + return self.create_completion(prompt=prompt, **kwargs) + + def create_chat_completion(self, messages, **kwargs): + # The chat handler formats the messages, then calls llama.create_completion. + prompt = "".join(str(m["content"]) for m in messages) + result = self.create_completion(prompt=prompt, **kwargs) + if kwargs.get("stream"): + return result + return { + "choices": [ + { + "message": {"role": "assistant", "content": "hi!"}, + "finish_reason": "stop", + } + ], + "usage": result["usage"], + } class _FailingLlama(Llama): - def __call__(self, *args, **kwargs): + def create_completion(self, prompt="", **kwargs): raise RuntimeError("cuda oom") @@ -171,6 +215,94 @@ def test_extract_info_none_model_id_when_graph_name_empty(self): model_id, _ = self.extractor.extract_info(session, {}) assert model_id is None + def _piper_session(self, path="/voices/pl_PL-gosia-medium.onnx"): + session = _FakeOrtSession() + session.get_modelmeta = lambda: MagicMock( + graph_name="main_graph", version=2**63 - 1 + ) + session._model_path = path + return session + + def test_main_graph_name_falls_back_to_file_stem(self): + model_id, info = self.extractor.extract_info(self._piper_session(), {}) + assert model_id == "pl_PL-gosia-medium" + assert info.model_name == "pl_PL-gosia-medium" + + def test_int64_max_version_is_unset(self): + _, info = self.extractor.extract_info(self._piper_session(), {}) + assert info.model_version == "unknown" + + def test_different_voices_get_different_ids(self): + a, _ = self.extractor.extract_info(self._piper_session("/v/en_US-amy.onnx"), {}) + b, _ = self.extractor.extract_info( + self._piper_session("/v/de_DE-thorsten.onnx"), {} + ) + assert a != b + + def test_hub_path_keeps_repo_id_and_names_model_after_file(self): + path = ( + "/hf/hub/models--rhasspy--piper-voices/snapshots/abc/pl/" + "pl_PL-gosia-medium.onnx" + ) + model_id, info = self.extractor.extract_info(self._piper_session(path), {}) + assert model_id == "rhasspy/piper-voices" + assert info.model_name == "pl_PL-gosia-medium" + + def test_session_from_bytes_without_name_still_needs_an_id(self): + session = self._piper_session(path=None) + model_id, info = self.extractor.extract_info(session, {}) + assert model_id is None + assert info.model_name == "onnx-model" + + def test_id_override_names_model_when_graph_and_file_are_generic(self): + session = self._piper_session(path="/models/model.onnx") + model_id, info = self.extractor.extract_info(session, {"id": "my-voice"}) + assert model_id == "my-voice" + assert info.model_name == "my-voice" + + def test_real_graph_name_wins_over_file_stem(self): + session = _FakeOrtSession() + session._model_path = "/models/exported-v3.onnx" + model_id, info = self.extractor.extract_info(session, {}) + assert model_id == "resnet50" + assert info.model_name == "resnet50" + + @pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, None), + (0, None), + (-1, None), + (2**63 - 1, None), + (1, "1"), + (42, "42"), + ("7", "7"), + ("v2", "v2"), + (" ", None), + ], + ) + def test_parse_model_version(self, raw, expected): + assert parse_model_version(raw) == expected + + @pytest.mark.parametrize( + ("path", "expected"), + [ + ("/m/pl_PL-gosia-medium.onnx", "pl_PL-gosia-medium"), + ("C:\\voices\\en_US-amy.onnx", "en_US-amy"), + ("/m/model.onnx", None), + ("/m/main_graph.onnx", None), + ("", None), + (None, None), + (b"bytes", None), + ], + ) + def test_model_file_stem(self, path, expected): + assert model_file_stem(path) == expected + + def test_model_id_from_path_skips_generic_stems(self): + assert model_id_from_path("/m/main_graph.onnx") is None + assert model_id_from_path("/m/yolo.onnx") == "yolo" + def test_install_hooks_publishes_inference_on_run(self, publish_spy): session = _FakeOrtSession() handle = make_handle(publish_spy) @@ -264,6 +396,247 @@ def test_install_hooks_tracks_error_on_exception(self, publish_spy): llm() assert publish_spy.events[0]["event_type"] == "error" + def test_install_hooks_tracks_chat_completion(self, publish_spy): + llm = Llama() + handle = make_handle(publish_spy) + self.extractor.install_hooks(llm, handle) + llm.create_chat_completion(messages=[{"role": "user", "content": "hello"}]) + assert len(publish_spy.events) == 1 + event = publish_spy.events[0] + assert event["event_type"] == "inference" + assert event["inference"]["output_meta"]["tokens_out"] == 20 + + def test_install_hooks_tracks_create_completion(self, publish_spy): + llm = Llama() + handle = make_handle(publish_spy) + self.extractor.install_hooks(llm, handle) + llm.create_completion(prompt="hello") + assert len(publish_spy.events) == 1 + assert publish_spy.events[0]["inference"]["input_meta"]["char_count"] == 5 + + def test_install_hooks_tracks_streamed_chat_completion(self, publish_spy): + llm = Llama() + handle = make_handle(publish_spy) + self.extractor.install_hooks(llm, handle) + chunks = list(llm.create_completion(prompt="hello", stream=True)) + assert len(chunks) == 2 + assert len(publish_spy.events) == 1 + meta = publish_spy.events[0]["inference"]["output_meta"] + assert meta["tokens_out"] == 20 + assert meta["stop_reason"] == "stop" + + def test_install_hooks_tracks_chat_completion_error(self, publish_spy): + llm = _FailingLlama() + handle = make_handle(publish_spy) + self.extractor.install_hooks(llm, handle) + with pytest.raises(RuntimeError): + llm.create_chat_completion(messages=[{"role": "user", "content": "hi"}]) + assert [e["event_type"] for e in publish_spy.events] == ["error"] + + def test_install_hooks_twice_does_not_double_track(self, publish_spy): + llm = Llama() + handle = make_handle(publish_spy) + self.extractor.install_hooks(llm, handle) + self.extractor.install_hooks(llm, handle) + llm.create_chat_completion(messages=[{"role": "user", "content": "hello"}]) + assert len(publish_spy.events) == 1 + + def test_non_streaming_stop_reason(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + llm.create_completion(prompt="hello") + meta = publish_spy.events[0]["inference"]["output_meta"] + assert meta["stop_reason"] == "length" + + def test_generation_config_from_kwargs(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + llm.create_chat_completion( + messages=[{"role": "user", "content": "hello"}], + temperature=0.7, + top_p=0.9, + max_tokens=64, + repeat_penalty=1.1, + stop=["", "User:"], + ) + config = publish_spy.events[0]["inference"]["generation_config"] + assert config == { + "temperature": 0.7, + "top_p": 0.9, + "max_tokens": 64, + "repetition_penalty": 1.1, + "stop_sequences_count": 2, + } + + def test_chat_input_meta(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + llm.create_chat_completion( + messages=[ + {"role": "system", "content": "be brief"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "what time is it"}, + ] + ) + meta = publish_spy.events[0]["inference"]["input_meta"] + assert meta["prompt_type"] == "chat" + assert meta["turn_index"] == 1 + assert meta["word_count"] == 4 + assert meta["token_count"] == 10 + + def test_streamed_chat_input_meta(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + stream = llm.create_chat_completion( + messages=[{"role": "user", "content": "hello"}], stream=True + ) + list(stream) + assert len(publish_spy.events) == 1 + inference = publish_spy.events[0]["inference"] + assert inference["input_meta"]["prompt_type"] == "chat" + assert inference["output_meta"]["stop_reason"] == "stop" + + def test_completion_input_meta_is_not_chat(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + llm.create_completion(prompt="hello") + assert "prompt_type" not in publish_spy.events[0]["inference"]["input_meta"] + + def test_no_attachments_when_disabled(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + llm.create_chat_completion(messages=[{"role": "user", "content": "hello"}]) + assert "attachments" not in publish_spy.events[0] + + @pytest.mark.parametrize("streaming", [False, True]) + @pytest.mark.parametrize("chat", [False, True]) + def test_enabled_uploads_never_capture_payloads( + self, publish_spy, tmp_path, streaming, chat + ): + store = AttachmentStore(str(tmp_path)) + manager = AttachmentManager(store) + llm = Llama() + handle = make_handle(publish_spy) + handle.capture_attachments = manager.capture + self.extractor.install_hooks(llm, handle) + if chat: + result = llm.create_chat_completion( + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "private prompt"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,PRIVATE"}, + }, + ], + } + ], + stream=streaming, + ) + else: + result = llm.create_completion(prompt="private prompt", stream=streaming) + if streaming: + assert list(result) + assert len(publish_spy.events) == 1 + assert "attachments" not in publish_spy.events[0] + assert store.length() == 0 + + def test_stop_given_as_string_counts_one_sequence(self): + config = make_gguf_generation_config({"stop": ""}) + assert config.stop_sequences_count == 1 + + def test_generation_config_none_without_sampling_arguments(self): + assert make_gguf_generation_config({"stop": []}) is None + + def test_output_meta_none_without_tokens_ttft_or_stop_reason(self): + assert make_gguf_output_meta(None, None, None, None, 10) is None + + def test_call_arguments_fill_defaults_and_flatten_kwargs(self): + def create_completion(self, prompt, temperature=0.8, **kwargs): + pass + + args = call_arguments( + inspect.signature(create_completion), object(), ("hi",), {"seed": 1} + ) + assert args["prompt"] == "hi" + assert args["temperature"] == 0.8 + assert args["seed"] == 1 + assert "kwargs" not in args + + def test_call_arguments_fall_back_when_binding_fails(self): + def create_completion(self, prompt): + pass + + args = call_arguments( + inspect.signature(create_completion), object(), ("hi", "extra"), {} + ) + assert args == {"prompt": "hi"} + + def test_call_arguments_without_signature(self): + assert call_arguments(None, object(), (), {"prompt": "p", "top_p": 0.5}) == { + "prompt": "p", + "top_p": 0.5, + } + + def test_patched_call_survives_unreadable_signature(self, publish_spy): + with patch.object( + gguf_integration.inspect, "signature", side_effect=ValueError("no sig") + ): + patched = gguf_integration.build_patched_call(Llama.create_completion) + llm = Llama() + setattr(llm, gguf_integration.GGUF_HANDLE_ATTR, make_handle(publish_spy)) + patched(llm, "hello", temperature=0.1) + event = publish_spy.events[0]["inference"] + assert event["input_meta"]["char_count"] == 5 + assert event["generation_config"] == {"temperature": 0.1} + + def test_generation_config_error_does_not_block_tracking(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + with patch.object( + gguf_integration, + "make_gguf_generation_config", + side_effect=RuntimeError("boom"), + ): + llm.create_completion(prompt="hello") + assert len(publish_spy.events) == 1 + assert "generation_config" not in publish_spy.events[0]["inference"] + + def test_malformed_usage_still_tracks(self, publish_spy): + class OddLlama(Llama): + def create_completion(self, prompt="", **kwargs): + return {"usage": "n/a"} + + llm = OddLlama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + llm.create_completion(prompt="hello") + assert len(publish_spy.events) == 1 + + def test_chat_without_handle_is_untracked(self, publish_spy): + llm = Llama() + self.extractor.install_hooks(llm, make_handle(publish_spy)) + setattr(llm, gguf_integration.GGUF_HANDLE_ATTR, None) + result = llm.create_chat_completion(messages=[{"role": "user", "content": "x"}]) + assert result["choices"][0]["message"]["content"] == "hi!" + assert publish_spy.events == [] + + def test_stream_callbacks_read_chat_deltas_and_skip_non_dicts(self, publish_spy): + captured = [] + handle = make_handle(publish_spy) + handle.capture_attachments = lambda a, *_: captured.extend(a) or [] + on_chunk, on_done = make_gguf_stream_callbacks(handle, "p") + on_chunk("not a dict") + on_chunk({"choices": [{"delta": {"role": "assistant"}}]}) + on_chunk({"choices": [{"delta": {"content": "he"}}]}) + on_chunk({"choices": [{"delta": {"content": "y"}, "finish_reason": "stop"}]}) + on_done(10, 2) + assert captured == [] + meta = publish_spy.events[0]["inference"]["output_meta"] + assert meta["stop_reason"] == "stop" + # --------------------------------------------------------------------------- # Keras @@ -292,7 +665,7 @@ def test_detect_accelerator_gpu_from_weight_device(self): model = _FakeKerasModel() model.weights = [MagicMock(device="/GPU:0")] with patch( - "wildedge.integrations.keras.CURRENT_PLATFORM.gpu_accelerator_for_offload", + "wildedge.platforms.CURRENT_PLATFORM.gpu_accelerator_for_offload", return_value="cuda", ): assert keras_detect_accelerator(model) == "cuda" diff --git a/tests/test_integrations_anthropic.py b/tests/test_integrations_anthropic.py index 3652700..43f438f 100644 --- a/tests/test_integrations_anthropic.py +++ b/tests/test_integrations_anthropic.py @@ -128,7 +128,7 @@ async def __anext__(self): try: return next(self._iter) except StopIteration: - raise StopAsyncIteration + raise StopAsyncIteration from None # Named "Anthropic" / "AsyncAnthropic" so can_handle sees the right type name. @@ -426,6 +426,7 @@ def test_streaming_records_inference_on_exhaustion(self): ), make_content_block_delta_event("Hello world"), make_message_delta_event(output_tokens=20, stop_reason="end_turn"), + SimpleNamespace(type="message_stop"), ] messages = FakeStreamingMessages(events) client = make_fake_client() @@ -449,6 +450,7 @@ def test_streaming_captures_token_counts_from_events(self): events = [ make_message_start_event(input_tokens=8, cached=2), make_message_delta_event(output_tokens=15, stop_reason="end_turn"), + SimpleNamespace(type="message_stop"), ] messages = FakeStreamingMessages(events) client = make_fake_client() @@ -555,6 +557,7 @@ async def test_streaming_records_inference_on_exhaustion(self): make_message_start_event(input_tokens=10, model="claude-opus-4-6-20251101"), make_content_block_delta_event("Hello world"), make_message_delta_event(output_tokens=20, stop_reason="end_turn"), + SimpleNamespace(type="message_stop"), ] messages = FakeAsyncStreamingMessages(events) client = make_fake_client() diff --git a/tests/test_integrations_common.py b/tests/test_integrations_common.py index 787f6ff..eb75a16 100644 --- a/tests/test_integrations_common.py +++ b/tests/test_integrations_common.py @@ -41,7 +41,7 @@ def test_debug_failure_logs_correctly(caplog): @pytest.mark.parametrize( - "dtype_str, expected", + ("dtype_str", "expected"), [ # PyTorch repr strings ("torch.float16", "f16"), @@ -56,11 +56,9 @@ def test_debug_failure_logs_correctly(caplog): ("tensor(uint8)", "int8"), # Plain names ("float16", "f16"), - ("bfloat16", "bf16"), + ("bfloat16", "bf16"), # must not match float16 ("int8", "int8"), ("float32", "f32"), - # bfloat16 must not match float16 - ("bfloat16", "bf16"), # Unknown ("float64", None), ("complex64", None), @@ -101,7 +99,7 @@ def test_histogram_uniform_distribution(): np = pytest.importorskip("numpy") # 500 pixels evenly spread [0, 0.2, 0.4, 0.6, 0.8, 1.0) flat = np.linspace(0.0, 1.0, 500, dtype=np.float32) - mean, std, buckets = image_brightness_histogram(flat) + mean, _std, buckets = image_brightness_histogram(flat) assert abs(mean - 0.5) < 0.01 # Each bucket should have roughly equal counts for b in buckets: @@ -139,7 +137,7 @@ def test_histogram_mean_and_std_rounded_to_4dp(): @pytest.mark.parametrize( - "shape, expected", + ("shape", "expected"), [ ((4, 3, 224, 224), "image"), # NCHW image batch ((1, 3, 64, 64), "image"), # single image @@ -159,7 +157,7 @@ def test_infer_input_modality_from_shape(shape, expected): @pytest.mark.parametrize( - "names, expected", + ("names", "expected"), [ (["input_ids", "attention_mask"], "text"), (["input_ids"], "text"), @@ -185,7 +183,7 @@ def test_infer_input_modality_from_names(names, expected): @pytest.mark.parametrize( - "layer_types, expected", + ("layer_types", "expected"), [ (["InputLayer", "Conv2D", "Dense"], "image"), (["InputLayer", "DepthwiseConv2D", "GlobalAveragePooling2D", "Dense"], "image"), @@ -207,7 +205,7 @@ def test_infer_input_modality_from_layer_types(layer_types, expected): @pytest.mark.parametrize( - "shape, expected", + ("shape", "expected"), [ ((None, 10), 10), # standard classification head ((None, 1000), 1000), # ImageNet @@ -274,7 +272,7 @@ class APITimeoutError(APIConnectionError): @pytest.mark.parametrize( - "exc, expected", + ("exc", "expected"), [ (APITimeoutError(), ErrorCode.INFERENCE_TIMEOUT), (TimeoutError(), ErrorCode.INFERENCE_TIMEOUT), @@ -306,7 +304,7 @@ def test_provider_error_code_of(): @pytest.mark.parametrize( - "exc, expected", + ("exc", "expected"), [ # anthropic, at request time and mid-stream: type in the wrapped body. ( diff --git a/tests/test_integrations_openai.py b/tests/test_integrations_openai.py index 9902731..10fdf82 100644 --- a/tests/test_integrations_openai.py +++ b/tests/test_integrations_openai.py @@ -92,7 +92,7 @@ async def create(self, *args, **kwargs): def make_stream_chunk(content=None, finish_reason=None, usage=None): - chunk = SimpleNamespace( + return SimpleNamespace( choices=[ SimpleNamespace( delta=SimpleNamespace(content=content), @@ -104,7 +104,6 @@ def make_stream_chunk(content=None, finish_reason=None, usage=None): system_fingerprint=None, service_tier=None, ) - return chunk class FakeStreamingCompletions: @@ -138,7 +137,7 @@ async def __anext__(self): try: return next(self._iter) except StopIteration: - raise StopAsyncIteration + raise StopAsyncIteration from None # Named "OpenAI" / "AsyncOpenAI" so can_handle sees the right type name. @@ -184,7 +183,7 @@ def register_model(obj, *, model_id=None, source=None, **kwargs): @pytest.mark.parametrize( - "url,expected", + ("url", "expected"), [ (None, "openai"), ("", "openai"), diff --git a/tests/test_integrations_openai_embeddings.py b/tests/test_integrations_openai_embeddings.py new file mode 100644 index 0000000..b64fc32 --- /dev/null +++ b/tests/test_integrations_openai_embeddings.py @@ -0,0 +1,102 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from wildedge.integrations.openai import wrap_async_embeddings, wrap_sync_embeddings + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("value", "batch_size", "chars"), + [ + ("hello", 1, 5), + (["hello", "world"], 2, 10), + ([1, 2, 3], 1, None), + ([[1, 2], [3]], 2, None), + ], +) +async def test_metadata_and_passthrough(asynchronous, value, batch_size, chars): + result = SimpleNamespace( + model="resolved-model", + usage=SimpleNamespace(prompt_tokens=7), + data=[SimpleNamespace(embedding=[0.1, 0.2, 0.3])], + ) + resource = SimpleNamespace( + create=(AsyncMock if asynchronous else MagicMock)(return_value=result) + ) + handle = MagicMock() + client = SimpleNamespace( + closed=False, register_model=MagicMock(return_value=handle) + ) + (wrap_async_embeddings if asynchronous else wrap_sync_embeddings)( + resource, "openai", lambda: client + ) + call = resource.create(model="model", input=value) + assert (await call if asynchronous else call) is result + handle.track_inference.assert_called_once() + event = handle.track_inference.call_args.kwargs + assert event["batch_size"] == batch_size + assert event["input_meta"].token_count == 7 + assert event["input_meta"].char_count == chars + assert event["output_meta"].tokens_in == 7 + assert event["output_meta"].dimensions == 3 + assert event["api_meta"].resolved_model_id == "resolved-model" + assert "hello" not in str(event) + client.closed = True + call = resource.create(model="model", input=value) + assert (await call if asynchronous else call) is result + assert handle.track_inference.call_count == 1 + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_error_and_telemetry_failure(asynchronous): + original = (AsyncMock if asynchronous else MagicMock)( + side_effect=TimeoutError("timeout") + ) + resource = SimpleNamespace(create=original) + handle = MagicMock() + client = SimpleNamespace( + closed=False, register_model=MagicMock(return_value=handle) + ) + (wrap_async_embeddings if asynchronous else wrap_sync_embeddings)( + resource, "openai", lambda: client + ) + + async def invoke(): + call = resource.create(model="model", input="hello") + return await call if asynchronous else call + + with pytest.raises(TimeoutError, match="timeout"): + await invoke() + handle.track_error.assert_called_once() + handle.track_inference.assert_not_called() + result = SimpleNamespace(model="model", usage=None, data=[]) + original.side_effect = None + original.return_value = result + handle.track_inference.side_effect = RuntimeError("telemetry broken") + assert await invoke() is result + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_input_iterator_is_not_consumed(asynchronous): + def values(): + raise AssertionError("instrumentation must not iterate inputs") + yield "unreachable" + + result = SimpleNamespace(model="model", usage=None, data=[]) + resource = SimpleNamespace( + create=(AsyncMock if asynchronous else MagicMock)(return_value=result) + ) + handle = MagicMock() + client = SimpleNamespace( + closed=False, register_model=MagicMock(return_value=handle) + ) + (wrap_async_embeddings if asynchronous else wrap_sync_embeddings)( + resource, "openai", lambda: client + ) + call = resource.create(model="model", input=values()) + assert (await call if asynchronous else call) is result + event = handle.track_inference.call_args.kwargs + assert event["batch_size"] is None + assert event["output_meta"].dimensions is None diff --git a/tests/test_integrations_openai_responses.py b/tests/test_integrations_openai_responses.py new file mode 100644 index 0000000..19c48d8 --- /dev/null +++ b/tests/test_integrations_openai_responses.py @@ -0,0 +1,144 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from wildedge.integrations.openai import ( + make_responses_stream_callbacks, + wrap_async_responses, + wrap_sync_responses, +) + + +def response(status="completed"): + return SimpleNamespace( + model="resolved", + status=status, + error=None, + incomplete_details=SimpleNamespace(reason="max_output_tokens") + if status == "incomplete" + else None, + usage=SimpleNamespace( + input_tokens=12, + output_tokens=4, + input_tokens_details=SimpleNamespace(cached_tokens=3), + output_tokens_details=SimpleNamespace(reasoning_tokens=2), + ), + ) + + +class AsyncItems: + def __init__(self, items): + self.items = items + + async def __aiter__(self): + for item in self.items: + yield item + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("streaming", [False, True]) +@pytest.mark.parametrize("status", ["completed", "incomplete", "failed"]) +async def test_responses(asynchronous, streaming, status): + result = response(status) + events = [ + SimpleNamespace(type="response.created"), + SimpleNamespace(type="response.output_text.delta", delta="hello"), + SimpleNamespace(type=f"response.{status}", response=result), + ] + returned = ( + (AsyncItems(events) if asynchronous else iter(events)) if streaming else result + ) + original = (AsyncMock if asynchronous else MagicMock)(return_value=returned) + resource = SimpleNamespace(create=original) + handle = MagicMock() + client = SimpleNamespace( + closed=False, register_model=MagicMock(return_value=handle) + ) + (wrap_async_responses if asynchronous else wrap_sync_responses)( + resource, "openai", lambda: client + ) + call = resource.create(model="model", input="hello world", stream=streaming) + actual = await call if asynchronous else call + if streaming: + assert ( + [event async for event in actual] if asynchronous else list(actual) + ) == events + else: + assert actual is result + handle.track_inference.assert_called_once() + data = handle.track_inference.call_args.kwargs + assert data["success"] == (status != "failed") + assert data["input_meta"].char_count == 11 + output = data["output_meta"] + assert ( + output.tokens_in, + output.tokens_out, + output.cached_input_tokens, + output.reasoning_tokens_out, + ) == (12, 4, 3, 2) + assert output.stop_reason == ( + "max_output_tokens" if status == "incomplete" else status + ) + assert data["api_meta"].resolved_model_id == "resolved" + + +def test_ttft_and_once(monkeypatch): + handle = MagicMock() + on_chunk, on_done = make_responses_stream_callbacks(handle, [], 0) + on_chunk(SimpleNamespace(type="response.created")) + monkeypatch.setattr("wildedge.integrations.openai.elapsed_ms", lambda _: 42) + on_chunk(SimpleNamespace(type="response.output_text.delta", delta="hi")) + on_chunk(SimpleNamespace(type="response.completed", response=response())) + on_done(100, 1) + on_done(101, 1) + handle.track_inference.assert_called_once() + assert ( + handle.track_inference.call_args.kwargs["output_meta"].time_to_first_token_ms + == 42 + ) + + +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_error_and_closed_client(asynchronous): + error = TimeoutError("timeout") + original = (AsyncMock if asynchronous else MagicMock)(side_effect=error) + resource = SimpleNamespace(create=original) + handle = MagicMock() + client = SimpleNamespace( + closed=False, register_model=MagicMock(return_value=handle) + ) + (wrap_async_responses if asynchronous else wrap_sync_responses)( + resource, "openai", lambda: client + ) + + async def invoke(): + call = resource.create(model="model", input="hello") + return await call if asynchronous else call + + with pytest.raises(TimeoutError) as caught: + await invoke() + assert caught.value is error + handle.track_error.assert_called_once() + handle.track_inference.assert_not_called() + client.closed = True + original.side_effect = None + original.return_value = "unchanged" + call = resource.create(model="model", input="hello") + assert (await call if asynchronous else call) == "unchanged" + assert client.register_model.call_count == 1 + + +def test_unfinished_and_error_streams(): + handle = MagicMock() + on_chunk, on_done = make_responses_stream_callbacks(handle, [], 0) + on_chunk(SimpleNamespace(type="response.created")) + on_done(100, 1) + handle.track_inference.assert_called_once() + assert handle.track_inference.call_args.kwargs["error_code"] == "STREAM_INCOMPLETE" + handle = MagicMock() + on_chunk, on_done = make_responses_stream_callbacks(handle, [], 0) + on_chunk(SimpleNamespace(type="error", code="server_error")) + on_done(100, 1) + handle.track_error.assert_called_once() + handle.track_inference.assert_not_called() diff --git a/tests/test_integrations_pytorch.py b/tests/test_integrations_pytorch.py index 357079b..91d3041 100644 --- a/tests/test_integrations_pytorch.py +++ b/tests/test_integrations_pytorch.py @@ -2,8 +2,12 @@ from __future__ import annotations +from types import SimpleNamespace from unittest.mock import MagicMock +import pytest + +import wildedge.integrations.pytorch as pytorch_mod from wildedge.integrations.pytorch import ( PytorchExtractor, ) @@ -163,3 +167,51 @@ def test_install_hooks_sets_detected_accelerator(self, publish_spy): handle = make_handle(publish_spy) self.extractor.install_hooks(model, handle) assert handle.detected_accelerator == "cpu" + + +def test_timm_config_identity_and_explicit_override(): + model = _FakeTorchModel() + model.pretrained_cfg = {"architecture": "resnet50"} + extractor = PytorchExtractor() + model_id, info = extractor.extract_info(model, {}) + assert model_id == info.model_name == "resnet50" + model_id, info = extractor.extract_info(model, {"id": "custom-model"}) + assert model_id == "custom-model" + assert info.model_name == "resnet50" + + +def test_timm_requested_variant_takes_priority_over_config(): + model = _FakeTorchModel() + model.pretrained_cfg = {"architecture": "resnet50"} + model.__wildedge_timm_model_name__ = "resnet50.a1_in1k" + model_id, info = PytorchExtractor().extract_info(model, {}) + assert model_id == info.model_name == "resnet50.a1_in1k" + + +def test_timm_invalid_config_falls_back_to_class(): + model = _FakeTorchModel() + model.pretrained_cfg = {"architecture": None} + model.default_cfg = {"architecture": 123} + model_id, info = PytorchExtractor().extract_info(model, {}) + assert model_id == info.model_name == "_FakeTorchModel" + + +@pytest.mark.parametrize("keyword", [False, True]) +@pytest.mark.parametrize("name", ["resnet50.a1_in1k", "hf-hub:example/resnet50"]) +def test_timm_patch_retains_requested_name(monkeypatch, keyword, name): + model = _FakeTorchModel() + fake_timm = SimpleNamespace(create_model=MagicMock(return_value=model)) + monkeypatch.setattr(pytorch_mod, "_timm", fake_timm) + monkeypatch.setattr(pytorch_mod, "_timm_patched", False) + client = MagicMock(closed=False) + client._diff_hub_caches.return_value = [] + PytorchExtractor.install_timm_patch(lambda: client) + result = ( + fake_timm.create_model(model_name=name) + if keyword + else fake_timm.create_model(name) + ) + assert result is model + model_id, info = PytorchExtractor().extract_info(result, {}) + assert model_id == info.model_name == name + client._on_model_auto_loaded.assert_called_once() diff --git a/tests/test_integrations_transformers.py b/tests/test_integrations_transformers.py new file mode 100644 index 0000000..c11193b --- /dev/null +++ b/tests/test_integrations_transformers.py @@ -0,0 +1,222 @@ +"""Inference tracking for transformers pipelines and models (#70).""" + +from __future__ import annotations + +import threading +import types + +import pytest + +import wildedge.integrations.transformers as tf_mod +from wildedge.integrations.transformers import TransformersExtractor +from wildedge.model import ModelHandle, ModelInfo + + +class FakeModel: + """Enough of a torch module for install_hooks: forward hooks and a config.""" + + config = types.SimpleNamespace( + name_or_path="tiny-bert", + model_type="bert", + architectures=["BertForSequenceClassification"], + ) + + def __init__(self) -> None: + self.pre_hooks: list = [] + self.post_hooks: list = [] + + def register_forward_pre_hook(self, hook): + self.pre_hooks.append(hook) + + def register_forward_hook(self, hook): + self.post_hooks.append(hook) + + def __call__(self, input_ids): + args = (input_ids,) + for hook in self.pre_hooks: + hook(self, args) + output = [[0.1, 0.9]] + for hook in self.post_hooks: + hook(self, args, output) + return output + + +class FakeTensor: + shape = (1, 4) + + +def make_pipeline_class(): + class Pipeline: + task = "text-classification" + + def __init__(self, model, fail=False, extra_models=()): + self.model = model + self.fail = fail + # Other models a custom pipeline runs, e.g. from _forward. + self.extra_models = list(extra_models) + + def __call__(self, inputs, *args, **kwargs): + if self.fail: + raise RuntimeError("tokenizer exploded") + self.model(FakeTensor()) + for extra in self.extra_models: + extra(FakeTensor()) + return [{"label": "LABEL_1", "score": 0.9}] + + # is_pipeline() looks for a class named Pipeline from a transformers module. + Pipeline.__module__ = "transformers.pipelines.base" + return Pipeline + + +@pytest.fixture +def patched_pipeline(monkeypatch): + pipeline_cls = make_pipeline_class() + fake_transformers = types.SimpleNamespace(Pipeline=pipeline_cls) + monkeypatch.setattr(tf_mod, "_transformers", fake_transformers) + TransformersExtractor._patch_pipeline_call() + return pipeline_cls + + +def make_handle(events: list, model_id: str = "tiny-bert") -> ModelHandle: + info = ModelInfo( + model_name=model_id, + model_version="1", + model_source="local", + model_format="transformers", + ) + return ModelHandle(model_id=model_id, info=info, publish=events.append) + + +def tracked_pipeline(pipeline_cls, events, fail=False): + """A pipeline wrapping a model, both with hooks, as after pipeline() loads.""" + extractor = TransformersExtractor() + handle = make_handle(events) + model = FakeModel() + extractor.install_hooks(model, handle) # from_pretrained ran first + pipe = pipeline_cls(model, fail=fail) + extractor.install_hooks(pipe, handle) # then the pipeline superseded it + return pipe, model + + +def inference_events(events: list) -> list: + return [e for e in events if e["event_type"] == "inference"] + + +def test_pipeline_call_is_tracked_once(patched_pipeline): + events: list = [] + pipe, _ = tracked_pipeline(patched_pipeline, events) + + pipe("hello world") + + inferences = inference_events(events) + assert len(inferences) == 1 + inference = inferences[0]["inference"] + # The pipeline-level event, with what only the pipeline knows. + assert inference["batch_size"] == 1 + assert inference["input_meta"]["word_count"] == 2 + assert inference["output_meta"]["task"] == "classification" + + +def test_direct_model_call_is_still_tracked(patched_pipeline): + events: list = [] + _, model = tracked_pipeline(patched_pipeline, events) + + model(FakeTensor()) + + inferences = inference_events(events) + assert len(inferences) == 1 + assert inferences[0]["inference"]["input_meta"]["token_count"] == 4 + + +def test_failed_pipeline_call_resets_suppression(patched_pipeline): + events: list = [] + pipe, model = tracked_pipeline(patched_pipeline, events, fail=True) + + with pytest.raises(RuntimeError, match="tokenizer exploded"): + pipe("hello") + assert [e["event_type"] for e in events] == ["error"] + + model(FakeTensor()) + assert len(inference_events(events)) == 1 + + +def test_untracked_pipeline_does_not_suppress_model_hooks(patched_pipeline): + events: list = [] + model = FakeModel() + TransformersExtractor().install_hooks(model, make_handle(events)) + pipe = patched_pipeline(model) # no handle on the pipeline itself + + pipe("hello") + + assert len(inference_events(events)) == 1 + + +def call_on_another_thread(model: FakeModel): + def run(input_ids): + other = threading.Thread(target=model, args=(input_ids,)) + other.start() + other.join() + + return run + + +def test_model_call_on_another_thread_during_pipeline_call_is_tracked( + patched_pipeline, +): + events: list = [] + extractor = TransformersExtractor() + handle = make_handle(events) + model = FakeModel() + extractor.install_hooks(model, handle) + pipe = patched_pipeline(model, extra_models=[call_on_another_thread(model)]) + extractor.install_hooks(pipe, handle) + + pipe("hello") + + # The pipeline's own event, and the other thread's model call; the model + # call on the pipeline's thread is part of the pipeline event. + metas = [e["inference"]["input_meta"] for e in inference_events(events)] + assert len(metas) == 2 + assert {"char_count": 5, "word_count": 1} in metas # the pipeline call + assert {"token_count": 4} in metas # the other thread's model call + + +def test_other_model_called_during_pipeline_call_is_tracked(patched_pipeline): + """A custom pipeline that also runs a second tracked model keeps both.""" + events: list = [] + extractor = TransformersExtractor() + main_model, reranker = FakeModel(), FakeModel() + main_handle = make_handle(events) + reranker_handle = make_handle(events, model_id="reranker") + extractor.install_hooks(main_model, main_handle) + extractor.install_hooks(reranker, reranker_handle) + pipe = patched_pipeline(main_model, extra_models=[reranker]) + extractor.install_hooks(pipe, main_handle) + + pipe("hello") + + ids = sorted(e["model_id"] for e in inference_events(events)) + assert ids == ["reranker", "tiny-bert"] + + +def test_nested_pipeline_calls_keep_suppressing_until_the_outer_ends( + patched_pipeline, +): + events: list = [] + extractor = TransformersExtractor() + handle = make_handle(events) + model = FakeModel() + extractor.install_hooks(model, handle) + inner = patched_pipeline(model) + extractor.install_hooks(inner, handle) + # The outer pipeline wraps the same model and calls the inner one, then + # runs the model again after the inner call has returned. + outer = patched_pipeline(model, extra_models=[lambda _: inner("x"), model]) + extractor.install_hooks(outer, handle) + + outer("hello") + + # One event per pipeline call; no model events from inside either. + inferences = inference_events(events) + assert len(inferences) == 2 + assert all("word_count" in e["inference"]["input_meta"] for e in inferences) diff --git a/tests/test_integrations_ultralytics.py b/tests/test_integrations_ultralytics.py index 1a05458..885c430 100644 --- a/tests/test_integrations_ultralytics.py +++ b/tests/test_integrations_ultralytics.py @@ -2,10 +2,12 @@ from __future__ import annotations +import types from unittest.mock import MagicMock, patch import pytest +import wildedge.integrations.ultralytics as ultralytics_mod from wildedge.integrations.ultralytics import ( YOLO_FAMILY_RE, UltralyticsExtractor, @@ -82,18 +84,21 @@ class YOLO: model = FakeInnerModel() def __call__(self, source=None, stream=False, **kwargs): + return self.predict(source, stream=stream, **kwargs) + + def predict(self, source=None, stream=False, **kwargs): return [FakeResult([(10, 20, 100, 200, 0.9, 0)])] class FailingYOLO(YOLO): - def __call__(self, source=None, stream=False, **kwargs): + def predict(self, source=None, stream=False, **kwargs): raise RuntimeError("cuda oom") class ClassifyYOLO(YOLO): task = "classify" - def __call__(self, source=None, stream=False, **kwargs): + def predict(self, source=None, stream=False, **kwargs): return [FakeClassifyResult()] @@ -130,7 +135,7 @@ def test_is_yolo_false_for_string(): @pytest.mark.parametrize( - "name,expected", + ("name", "expected"), [ ("yolov8n", "yolov8"), ("yolov8s-seg", "yolov8"), @@ -208,8 +213,7 @@ def test_image_input_meta_returns_none_when_numpy_absent(): def test_image_input_meta_extracts_hwc_array(): - pytest.importorskip("numpy") - import numpy as np + np = pytest.importorskip("numpy") arr = np.zeros((480, 640, 3), dtype=np.uint8) meta = image_input_meta(arr) @@ -446,3 +450,53 @@ def test_extractor_install_hooks_batch_size_from_list_source(publish_spy): frames = [np.zeros((480, 640, 3), dtype=np.uint8)] * 3 model(frames) assert publish_spy.events[0]["inference"]["batch_size"] == 3 + + +def test_unsupported_ultralytics_version_warns_and_does_not_patch(monkeypatch, caplog): + fake_yolo = type("YOLO", (), {"__init__": lambda self, *a, **k: None}) + original_init = fake_yolo.__init__ + monkeypatch.setattr( + ultralytics_mod, "_ultralytics", types.SimpleNamespace(YOLO=fake_yolo) + ) + monkeypatch.setattr(ultralytics_mod, "_unsupported_ultralytics_version", "8.0.0") + monkeypatch.setattr(ultralytics_mod, "_ultralytics_patched", False) + + with caplog.at_level("WARNING", logger="wildedge"): + UltralyticsExtractor.install_auto_load_patch(lambda: None) + + assert fake_yolo.__init__ is original_init + messages = [r.getMessage() for r in caplog.records] + assert any("ultralytics 8.0.0 is older than 8.0.200" in m for m in messages) + + +@pytest.mark.parametrize("entrypoint", ["predict", "__call__"]) +def test_predict_and_call_track_once_with_repeated_registration( + publish_spy, entrypoint +): + model = YOLO() + handle = make_handle(publish_spy) + _extractor.install_hooks(model, handle) + _extractor.install_hooks(model, handle) + results = getattr(model, entrypoint)(source="image.png", verbose=False) + assert len(results) == 1 + assert len(publish_spy.events) == 1 + assert publish_spy.events[0]["event_type"] == "inference" + assert publish_spy.events[0]["inference"]["output_meta"]["num_predictions"] == 1 + + +def test_predict_error_is_preserved_and_tracked_once(publish_spy): + model = FailingYOLO() + _extractor.install_hooks(model, make_handle(publish_spy)) + with pytest.raises(RuntimeError, match="cuda oom"): + model.predict(source="image.png") + assert len(publish_spy.events) == 1 + assert publish_spy.events[0]["event_type"] == "error" + + +def test_predict_does_not_instrument_unregistered_instances(publish_spy): + tracked, untouched = YOLO(), YOLO() + _extractor.install_hooks(tracked, make_handle(publish_spy)) + untouched.predict() + assert not publish_spy.events + tracked.predict() + assert len(publish_spy.events) == 1 diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py new file mode 100644 index 0000000..c06982b --- /dev/null +++ b/tests/test_lazy_loading.py @@ -0,0 +1,249 @@ +"""Import isolation must hold in a fresh interpreter, regardless of installed ML packages.""" + +from __future__ import annotations + +import importlib +import json +import os +import subprocess +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import pytest + +from wildedge.hubs.base import BaseHubTracker +from wildedge.hubs.registry import HUB_SPECS +from wildedge.integrations.base import BaseExtractor +from wildedge.integrations.registry import EXTRACTOR_SPECS, INTEGRATION_SPECS +from wildedge.loading import resolve + +ROOT = Path(__file__).resolve().parents[1] +OPTIONAL = ( + "torch", + "timm", + "tensorflow", + "keras", + "mlx", + "mlx_lm", + "mlx_whisper", + "faster_whisper", + "ctranslate2", + "onnxruntime", + "ultralytics", + "transformers", + "llama_cpp", + "openai", + "anthropic", + "huggingface_hub", + "numpy", +) + + +def isolated(tmp_path, code, *, working=(), sources=None): + """Make every optional package discoverable but broken unless explicitly supplied.""" + for name in OPTIONAL: + path = tmp_path / name + path.mkdir() + (path / "__init__.py").write_text( + (sources or {}).get( + name, + "" + if name in working + else f'raise RuntimeError("unexpected import: {name}")\n', + ) + ) + env = dict(os.environ) + env.pop("WILDEDGE_AUTOLOAD", None) + env.pop("WILDEDGE_DSN", None) + bootstrap = f"import sys; sys.path[:0] = [{str(tmp_path)!r}, {str(ROOT)!r}]\n" + result = subprocess.run( + [sys.executable, "-c", bootstrap + code], + cwd=tmp_path, + env=env, + capture_output=True, + text=True, + timeout=30, + check=False, + ) + assert result.returncode == 0, result.stdout + result.stderr + return result.stdout + + +def test_import_and_unconfigured_init_ignore_broken_installed_frameworks(tmp_path): + isolated( + tmp_path, + f""" +import importlib.util +# These are installed/discoverable, rather than merely absent. +for name in {OPTIONAL!r}: + assert importlib.util.find_spec(name) is not None +import wildedge +import wildedge.cli +from wildedge.hubs import BaseHubTracker +client = wildedge.init(dsn="", integrations=["openai"]) +assert client.noop +client.register_model(object(), model_id="manual-model", auto_instrument=False) +assert not ({set(OPTIONAL)!r} & set(sys.modules)) +assert not any(name.startswith("wildedge.integrations.") and name.rsplit(".", 1)[-1] not in ("base", "registry", "common") for name in sys.modules) +client.close() +""", + ) + + +def test_selected_integration_does_not_import_siblings(tmp_path): + isolated( + tmp_path, + """ +from unittest.mock import patch +import wildedge +from wildedge.platforms.device_info import DeviceInfo +with patch("wildedge.client.Consumer"), patch("wildedge.client.Transmitter"), patch("wildedge.paths.default_sdk_state_dir", return_value=__import__("pathlib").Path("state")): + client = wildedge.WildEdge(dsn="https://secret@example.invalid/project", device=DeviceInfo("test", "linux"), sampling_interval_s=None) + client.instrument("openai") + assert "openai" in sys.modules + assert "wildedge.integrations.openai" in sys.modules + assert not ({"torch", "timm", "mlx", "tensorflow", "huggingface_hub", "anthropic"} & set(sys.modules)) + client.close() +""", + sources={"openai": "class OpenAI: pass\nclass AsyncOpenAI: pass\n"}, + ) + + +@pytest.mark.parametrize( + ("hub", "package"), [("huggingface", "huggingface_hub"), ("torchhub", "torch")] +) +def test_public_hub_reexports_load_only_the_requested_tracker(tmp_path, hub, package): + output = isolated( + tmp_path, + f""" +import json +from wildedge import hubs +tracker = getattr(hubs, {("HuggingFaceHubTracker" if hub == "huggingface" else "TorchHubTracker")!r}) +assert tracker().name == {hub!r} +print(json.dumps(sorted(set(sys.modules) & {set(OPTIONAL)!r}))) +""", + working=(package,), + ) + # HF's file_download is absent in this minimal fixture, but importing its root + # still proves the optional package is touched only for the requested tracker. + assert json.loads(output) == [package] + + +def test_manual_pytorch_subclass_registration_does_not_need_timm(tmp_path): + isolated( + tmp_path, + """ +from unittest.mock import patch +from pathlib import Path +import wildedge +from wildedge.platforms.device_info import DeviceInfo +# The base belongs to torch; the application subclass belongs to __main__. +Module = type("Module", (), {"__module__": "torch.nn"}) +class ApplicationModel(Module): + pass +with patch("wildedge.client.Consumer"), patch("wildedge.client.Transmitter"), patch("wildedge.paths.default_sdk_state_dir", return_value=Path("state")): + client = wildedge.WildEdge(dsn="https://secret@example.invalid/project", device=DeviceInfo("test", "linux"), sampling_interval_s=None) + handle = client.register_model(ApplicationModel(), model_id="application-model", auto_instrument=False) + assert handle.info.model_format == "pytorch" + assert "torch" in sys.modules + assert not ({"timm", "tensorflow", "transformers", "huggingface_hub"} & set(sys.modules)) + client.close() +""", + working=("torch", "numpy"), + ) + + +@pytest.mark.parametrize( + ("base_name", "base_module", "format_name", "working"), + [ + ("Model", "keras.src.models.model", "tensorflow", ("numpy",)), + ("Model", "tf_keras.src.engine.training", "tensorflow", ("numpy",)), + ("Module", "mlx.nn.layers.base", "mlx", ("mlx", "numpy")), + ], +) +def test_manual_registration_does_not_load_companion_runtime( + tmp_path, base_name, base_module, format_name, working +): + isolated( + tmp_path, + f""" +from unittest.mock import patch +from pathlib import Path +import wildedge +from wildedge.platforms.device_info import DeviceInfo +Base = type({base_name!r}, (), {{"__module__": {base_module!r}}}) +class ApplicationModel(Base): + name = "application-model" + weights = [] +with patch("wildedge.client.Consumer"), patch("wildedge.client.Transmitter"), patch("wildedge.paths.default_sdk_state_dir", return_value=Path("state")): + client = wildedge.WildEdge(dsn="https://secret@example.invalid/project", device=DeviceInfo("test", "linux"), sampling_interval_s=None) + handle = client.register_model(ApplicationModel(), model_id="application-model", auto_instrument=False) + assert handle.info.model_format == {format_name!r} + assert not ({{"tensorflow", "mlx_lm", "timm", "transformers"}} & set(sys.modules)) + client.close() +""", + working=working, + ) + + +@pytest.mark.parametrize( + ("adapter", "runtime_name", "runtime_attr", "flag", "extractor", "method"), + [ + ( + "pytorch", + "timm", + "_timm", + "_timm_patched", + "PytorchExtractor", + "install_timm_patch", + ), + ( + "mlx", + "mlx_lm", + "_mlx_lm", + "_mlx_patched", + "MlxExtractor", + "install_auto_load_patch", + ), + ( + "tensorflow", + "tensorflow", + "_tf", + "_tf_auto_patched", + "TensorflowExtractor", + "install_auto_load_patch", + ), + ], +) +def test_activation_imports_runtime_after_extractor_was_used( + monkeypatch, adapter, runtime_name, runtime_attr, flag, extractor, method +): + module = importlib.import_module(f"wildedge.integrations.{adapter}") + runtime = ModuleType(runtime_name) + if adapter == "pytorch": + runtime.create_model = lambda *_a, **_kw: None + elif adapter == "mlx": + runtime.load = lambda *_a, **_kw: None + runtime.generate = lambda *_a, **_kw: None + else: + runtime.keras = SimpleNamespace( + models=SimpleNamespace(load_model=lambda *_a, **_kw: None) + ) + runtime.saved_model = SimpleNamespace(load=lambda *_a, **_kw: None) + monkeypatch.setitem(sys.modules, runtime_name, runtime) + monkeypatch.setattr(module, runtime_attr, None) + monkeypatch.setattr(module, flag, False) + getattr(getattr(module, extractor), method)(lambda: None) + assert getattr(module, runtime_attr) is runtime + assert getattr(module, flag) is True + + +def test_registry_entries_point_at_real_code(): + for spec in INTEGRATION_SPECS: + if spec.installer is not None: + assert callable(resolve(spec.installer)), spec.name + for spec in EXTRACTOR_SPECS: + assert isinstance(spec.load(), BaseExtractor), spec.class_name + for spec in HUB_SPECS: + assert issubclass(resolve(spec.tracker), BaseHubTracker), spec.name diff --git a/tests/test_llm_api.py b/tests/test_llm_api.py index d9a5863..3b214f2 100644 --- a/tests/test_llm_api.py +++ b/tests/test_llm_api.py @@ -136,9 +136,11 @@ def test_response_anthropic_shape(client): def test_exception_records_error_not_inference(client): - with pytest.raises(ValueError): - with wildedge.llm_api(model="m", provider="p"): - raise ValueError("bad payload") + with ( + pytest.raises(ValueError, match="bad payload"), + wildedge.llm_api(model="m", provider="p"), + ): + raise ValueError("bad payload") assert client.handle.inferences == [] assert client.handle.errors == [ @@ -150,9 +152,8 @@ def test_http_error_leaves_code_to_server(client): class HTTPStatusError(Exception): response = SimpleNamespace(status_code=401) - with pytest.raises(HTTPStatusError): - with wildedge.llm_api(model="m", provider="p"): - raise HTTPStatusError("Client error '401 Unauthorized'") + with pytest.raises(HTTPStatusError), wildedge.llm_api(model="m", provider="p"): + raise HTTPStatusError("Client error '401 Unauthorized'") assert client.handle.errors[0]["error_code"] == "UNKNOWN" @@ -161,9 +162,8 @@ def test_provider_error_without_status_leaves_code_to_server(client): class APIError(Exception): code = "server_error" - with pytest.raises(APIError): - with wildedge.llm_api(model="m", provider="p"): - raise APIError("An error occurred during streaming") + with pytest.raises(APIError), wildedge.llm_api(model="m", provider="p"): + raise APIError("An error occurred during streaming") assert client.handle.errors[0]["error_code"] == "UNKNOWN" diff --git a/tests/test_model.py b/tests/test_model.py index ae5683d..e103414 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -1,15 +1,17 @@ """Tests for ModelRegistry and ModelHandle.""" +from unittest.mock import patch + from wildedge.model import ModelHandle, ModelInfo, ModelRegistry def make_info(**kwargs) -> ModelInfo: - defaults = dict( - model_name="test", - model_version="1.0", - model_source="local", - model_format="onnx", - ) + defaults = { + "model_name": "test", + "model_version": "1.0", + "model_source": "local", + "model_format": "onnx", + } return ModelInfo(**{**defaults, **kwargs}) @@ -52,6 +54,42 @@ def test_snapshot_reflects_registered_models(self): assert "model-1" in snap assert snap["model-1"]["model_name"] == "yolo" + def test_new_process_registration_replaces_persisted_info(self, tmp_path): + path = tmp_path / "model_registry.json" + ModelRegistry(path).register( + "model-1", make_info(model_version="1.0"), noop_publish + ) + + restarted = ModelRegistry(path) + handle, _ = restarted.register( + "model-1", make_info(model_version="2.0"), noop_publish + ) + + assert handle.info.model_version == "2.0" + assert restarted.snapshot()["model-1"]["model_version"] == "2.0" + assert ModelRegistry(path).models["model-1"].model_version == "2.0" + + def test_unchanged_info_does_not_rewrite_persisted_file(self, tmp_path): + path = tmp_path / "model_registry.json" + ModelRegistry(path).register("model-1", make_info(), noop_publish) + + restarted = ModelRegistry(path) + with patch.object(restarted, "save_to_disk") as save: + restarted.register("model-1", make_info(), noop_publish) + + save.assert_not_called() + + def test_persisted_model_not_registered_again_is_kept(self, tmp_path): + path = tmp_path / "model_registry.json" + ModelRegistry(path).register( + "model-1", make_info(model_name="yolo"), noop_publish + ) + + restarted = ModelRegistry(path) + restarted.register("model-2", make_info(), noop_publish) + + assert restarted.snapshot()["model-1"]["model_name"] == "yolo" + def test_snapshot_excludes_duplicate_registration(self): registry = ModelRegistry() registry.register("model-1", make_info(model_name="original"), noop_publish) diff --git a/tests/test_platform_adapter.py b/tests/test_platform_adapter.py index 6c754cf..f54598b 100644 --- a/tests/test_platform_adapter.py +++ b/tests/test_platform_adapter.py @@ -18,26 +18,30 @@ def test_real_device_model(platform_cls): model = platform_cls().device_model() if model is not None: - assert isinstance(model, str) and model + assert isinstance(model, str) + assert 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 + assert isinstance(ver, str) + assert 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 + assert isinstance(total, int) + assert 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 + assert isinstance(total, int) + assert total > 0 @pytest.mark.parametrize("platform_cls", PLATFORMS) @@ -46,13 +50,15 @@ def test_real_gpu_accelerators(platform_cls): 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 + assert isinstance(gpu_name, str) + assert 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 + assert isinstance(offload, str) + assert offload @pytest.mark.parametrize("platform_cls", PLATFORMS) diff --git a/tests/test_platform_adapters.py b/tests/test_platform_adapters.py index 05a1de4..a9b3d85 100644 --- a/tests/test_platform_adapters.py +++ b/tests/test_platform_adapters.py @@ -1,5 +1,6 @@ from __future__ import annotations +import wildedge.platforms.macos as macos_mod from wildedge.platforms.base import cuda_device_count from wildedge.platforms.linux import LinuxPlatform from wildedge.platforms.macos import MACOS_THERMAL_STATES, MacOSPlatform @@ -67,7 +68,7 @@ def cuInit(self, _flags): # noqa: N802 return 0 def cuDeviceGetCount(self, ptr): # noqa: N802 - ptr._obj.value = 2 # noqa: SLF001 + ptr._obj.value = 2 return 0 monkeypatch.setattr("ctypes.CDLL", lambda _name: FakeCudaLib()) @@ -142,7 +143,6 @@ def test_macos_thermal_states_covers_all_levels(): def test_macos_cpu_freq_intel(monkeypatch): """On Intel, both hw.cpufrequency and hw.cpufrequency_max are available.""" - import wildedge.platforms.macos as macos_mod calls = [] @@ -161,7 +161,6 @@ def fake_sysctl(name): def test_macos_cpu_freq_apple_silicon(monkeypatch): """On Apple Silicon, hw.cpufrequency is absent; max falls back to hw.perflevel0.cpufrequency_max.""" - import wildedge.platforms.macos as macos_mod def fake_sysctl(name): return {b"hw.perflevel0.cpufrequency_max": 4_056_000_000}.get(name) @@ -173,8 +172,6 @@ def fake_sysctl(name): def test_macos_thermal_nominal(monkeypatch): - import wildedge.platforms.macos as macos_mod - monkeypatch.setattr(macos_mod.ctypes.cdll, "LoadLibrary", lambda _: FakeObjc()) monkeypatch.setattr( macos_mod.ctypes, @@ -191,8 +188,6 @@ def test_macos_thermal_nominal(monkeypatch): def test_macos_thermal_serious(monkeypatch): - import wildedge.platforms.macos as macos_mod - monkeypatch.setattr(macos_mod.ctypes.cdll, "LoadLibrary", lambda _: FakeObjc()) monkeypatch.setattr( macos_mod.ctypes, @@ -210,8 +205,6 @@ def test_macos_thermal_serious(monkeypatch): def test_macos_thermal_returns_none_on_error(monkeypatch): - import wildedge.platforms.macos as macos_mod - monkeypatch.setattr( macos_mod.ctypes.cdll, "LoadLibrary", diff --git a/tests/test_platform_linux.py b/tests/test_platform_linux.py index 5c7d431..0a01a6a 100644 --- a/tests/test_platform_linux.py +++ b/tests/test_platform_linux.py @@ -19,7 +19,7 @@ def test_meminfo_parses_total_and_available(monkeypatch): lambda p, **kw: ( __import__("io").StringIO(content) if p == "/proc/meminfo" - else open(p, **kw) + else open(p, **kw) # noqa: SIM115 - the caller closes it ), ) total, available = LinuxPlatform().meminfo() @@ -86,7 +86,8 @@ def test_battery_no_supply_returns_none(tmp_path, monkeypatch): lambda p: ps if p == "/sys/class/power_supply" else Path(p), ) level, charging = LinuxPlatform().battery() - assert level is None and charging is None + assert level is None + assert charging is None def test_thermal_nominal(tmp_path, monkeypatch): @@ -181,8 +182,10 @@ def test_trip_point_states_covers_all_severity_levels(): @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 + assert isinstance(total, int) + assert total > 0 + assert isinstance(available, int) + assert 0 < available <= total @pytest.mark.requires_linux diff --git a/tests/test_platform_macos.py b/tests/test_platform_macos.py index 096030c..281d79c 100644 --- a/tests/test_platform_macos.py +++ b/tests/test_platform_macos.py @@ -2,13 +2,12 @@ import pytest +from wildedge.platforms import macos as macos_mod from wildedge.platforms.hardware import ThermalContext 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", @@ -25,8 +24,6 @@ def test_meminfo_computes_available(monkeypatch): 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() @@ -62,7 +59,8 @@ def test_hardware_context_fields(monkeypatch): @pytest.mark.requires_macos def test_real_meminfo(): total, available = MacOSPlatform().meminfo() - assert isinstance(total, int) and total > 0 + assert isinstance(total, int) + assert total > 0 if available is not None: assert 0 < available <= total diff --git a/tests/test_platform_windows.py b/tests/test_platform_windows.py index 862d7eb..71d1a14 100644 --- a/tests/test_platform_windows.py +++ b/tests/test_platform_windows.py @@ -24,8 +24,10 @@ def test_hardware_context_fields(monkeypatch): @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 + assert isinstance(total, int) + assert total > 0 + assert isinstance(available, int) + assert 0 < available <= total @pytest.mark.requires_windows diff --git a/tests/test_quantizations.py b/tests/test_quantizations.py new file mode 100644 index 0000000..d16db31 --- /dev/null +++ b/tests/test_quantizations.py @@ -0,0 +1,75 @@ +"""Quantization labels are shared across integrations.""" + +from types import SimpleNamespace + +import pytest + +from wildedge import quantizations +from wildedge.integrations.common import dtype_to_quantization, mlx_layer_quantization +from wildedge.integrations.faster_whisper import COMPUTE_TYPE_QUANTIZATION +from wildedge.integrations.gguf import file_type_quantization, parse_quantization + + +@pytest.mark.parametrize( + ("dtype", "expected"), + [ + ("torch.bfloat16", quantizations.BF16), + ("torch.float16", quantizations.F16), + ("tensor(float16)", quantizations.F16), + ("mlx.core.float32", quantizations.F32), + ("torch.qint8", quantizations.INT8), + ], +) +def test_dtype_names(dtype, expected): + assert dtype_to_quantization(dtype) == expected + + +@pytest.mark.parametrize( + ("compute_type", "expected"), + [ + ("int8_float16", quantizations.INT8), + ("int8_bfloat16", quantizations.INT8), + ("float16", quantizations.F16), + ("int16", "int16"), + ], +) +def test_ctranslate2_compute_types_use_weight_precision(compute_type, expected): + assert COMPUTE_TYPE_QUANTIZATION[compute_type] == expected + + +@pytest.mark.parametrize( + ("filename", "expected"), + [("llama-7b.Q4_K_M", "q4_k_m"), ("model-f16", "f16"), ("model.i8", "int8")], +) +def test_gguf_filename_suffixes(filename, expected): + assert parse_quantization(filename) == expected + + +@pytest.mark.parametrize( + ("file_type", "expected"), + [("15", "q4_k_m"), (0, quantizations.F32), ("32", quantizations.BF16)], +) +def test_gguf_file_types(file_type, expected): + assert file_type_quantization(file_type) == expected + + +@pytest.mark.parametrize("file_type", ["1024", "not-a-number", "5"]) +def test_unknown_gguf_file_types_are_not_reported(file_type): + assert file_type_quantization(file_type) is None + + +class QuantizedLinear: + pass + + +@pytest.mark.parametrize( + ("layers", "expected"), + [ + ([SimpleNamespace(bits=4)], "int4"), + ([QuantizedLinear()], quantizations.QUANTIZED), + ([SimpleNamespace()], None), + ], +) +def test_mlx_layers(layers, expected): + model = SimpleNamespace(named_modules=lambda: [("", layer) for layer in layers]) + assert mlx_layer_quantization(model) == expected diff --git a/tests/test_sampler.py b/tests/test_sampler.py index df71696..1f305eb 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -2,6 +2,8 @@ from unittest.mock import MagicMock, patch +from wildedge import platforms +from wildedge.model import ModelHandle from wildedge.platforms import ( capture_hardware, is_sampling, @@ -87,9 +89,8 @@ 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() + original = platforms._sampler.snapshot() capture_hardware(accelerator_actual="cuda") assert original.accelerator_actual is None @@ -103,8 +104,6 @@ def test_capture_hardware_accelerator_override(self): 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) @@ -115,16 +114,12 @@ def test_auto_attaches_when_sampling(self, monkeypatch): 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) diff --git a/tests/test_streaming_lifecycle.py b/tests/test_streaming_lifecycle.py new file mode 100644 index 0000000..dd03fbd --- /dev/null +++ b/tests/test_streaming_lifecycle.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +import asyncio +import time +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from wildedge.integrations.anthropic import make_anthropic_stream_callbacks +from wildedge.integrations.common import ( + AsyncStreamWrapper, + SyncStreamWrapper, + stream_chunk_has_output, +) +from wildedge.integrations.gguf import make_gguf_stream_callbacks +from wildedge.integrations.openai import ( + make_openai_stream_callbacks, + make_responses_stream_callbacks, +) +from wildedge.tracing import get_trace_context, trace_context + + +def wrapper(original, on_chunk=None, on_done=None): + handle = MagicMock() + done = on_done or MagicMock() + return ( + SyncStreamWrapper(original, handle, time.perf_counter(), on_chunk, done), + handle, + done, + ) + + +def test_single_pass_completion_and_close(): + stream, handle, done = wrapper(iter([1, 2])) + assert next(stream) == 1 + assert list(stream) == [2] + assert list(stream) == [] + stream.close() + done.assert_called_once() + handle.track_error.assert_not_called() + + +def test_explicit_close_and_context_exit(): + original = MagicMock() + original.__iter__.return_value = iter([1, 2]) + stream, _, done = wrapper(original) + with stream: + assert next(stream) == 1 + original.__exit__.assert_called_once() + stream.close() + done.assert_called_once() + assert done.call_args.kwargs == {"outcome": "closed"} + + +def test_midstream_error_preserved_once(): + error = TimeoutError("broken stream") + + def chunks(): + yield 1 + raise error + + stream, handle, done = wrapper(chunks()) + assert next(stream) == 1 + with pytest.raises(TimeoutError) as caught: + next(stream) + assert caught.value is error + stream.close() + handle.track_error.assert_called_once() + done.assert_not_called() + + +def test_callback_failures_do_not_break_iteration(): + stream, handle, _ = wrapper( + iter([1, 2]), + MagicMock(side_effect=ValueError()), + MagicMock(side_effect=ValueError()), + ) + assert list(stream) == [1, 2] + handle.track_error.assert_not_called() + + +def test_correlation_uses_creation_context(): + seen = [] + with trace_context(trace_id="origin"): + stream, _, _ = wrapper( + iter([1]), on_done=lambda *_: seen.append(get_trace_context().trace_id) + ) + with trace_context(trace_id="consumer"): + list(stream) + assert seen == ["origin"] + + +class AsyncChunks: + def __init__(self, items, error=None): + self.items = iter(items) + self.error = error + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + try: + return next(self.items) + except StopIteration: + if self.error is not None: + raise self.error from None + raise StopAsyncIteration from None + + async def aclose(self): + self.closed = True + + +@pytest.mark.parametrize( + "error", [None, TimeoutError("timeout"), asyncio.CancelledError()] +) +async def test_async_lifecycle(error): + original = AsyncChunks([1], error) + handle, done = MagicMock(), MagicMock() + stream = AsyncStreamWrapper(original, handle, time.perf_counter(), None, done) + assert await anext(stream) == 1 + if error is None: + assert [item async for item in stream] == [] + done.assert_called_once() + else: + with pytest.raises(type(error)) as caught: + await anext(stream) + assert caught.value is error + if isinstance(error, asyncio.CancelledError): + assert done.call_args.kwargs == {"outcome": "cancelled"} + else: + handle.track_error.assert_called_once() + done.assert_not_called() + await stream.aclose() + assert original.closed + assert [item async for item in stream] == [] + assert done.call_count + handle.track_error.call_count == 1 + + +async def test_async_context_early_exit(): + original = AsyncChunks([1, 2]) + done = MagicMock() + stream = AsyncStreamWrapper(original, MagicMock(), time.perf_counter(), None, done) + async with stream: + assert await anext(stream) == 1 + assert original.closed + assert done.call_args.kwargs == {"outcome": "closed"} + + +@pytest.mark.parametrize( + "factory", + [ + lambda h: make_openai_stream_callbacks(h, []), + lambda h: make_anthropic_stream_callbacks(h, []), + lambda h: make_gguf_stream_callbacks(h, "hello"), + lambda h: make_responses_stream_callbacks(h, [], time.perf_counter()), + ], +) +@pytest.mark.parametrize("outcome", ["closed", "cancelled"]) +def test_all_integrations_record_interruption(factory, outcome): + handle = MagicMock() + _, on_done = factory(handle) + on_done(100, None, outcome=outcome) + handle.track_inference.assert_called_once() + event = handle.track_inference.call_args.kwargs + assert event["success"] is False + assert event["error_code"] == "STREAM_" + outcome.upper() + + +@pytest.mark.parametrize( + ("chunk", "expected"), + [ + ({"choices": [{"delta": {"role": "assistant"}}]}, False), + ({"choices": [], "usage": {}}, False), + ({"choices": [{"delta": {"content": "hello"}}]}, True), + ({"choices": [{"text": "hello"}]}, True), + (SimpleNamespace(type="message_start"), False), + ( + SimpleNamespace( + type="content_block_delta", delta=SimpleNamespace(text="hello") + ), + True, + ), + (SimpleNamespace(type="response.created"), False), + ], +) +def test_first_output_detection(chunk, expected): + assert stream_chunk_has_output(chunk) is expected + + +async def test_task_cancellation_while_waiting(): + entered = asyncio.Event() + + async def chunks(): + entered.set() + await asyncio.Event().wait() + yield 1 + + handle, done = MagicMock(), MagicMock() + stream = AsyncStreamWrapper(chunks(), handle, time.perf_counter(), None, done) + task = asyncio.create_task(anext(stream)) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await stream.aclose() + done.assert_called_once() + assert done.call_args.kwargs == {"outcome": "cancelled"} + handle.track_error.assert_not_called() + + +def test_ttft_waits_for_output(monkeypatch): + clock = [1] + monkeypatch.setattr("wildedge.integrations.common.elapsed_ms", lambda _: clock[0]) + chunks = iter( + [ + {"choices": [{"delta": {"role": "assistant"}}]}, + {"choices": [{"delta": {"content": "hello"}}]}, + ] + ) + stream, _, done = wrapper(chunks) + next(stream) + clock[0] = 42 + next(stream) + clock[0] = 100 + list(stream) + assert done.call_args.args == (100, 42) + + +def test_underlying_close_error_is_preserved(): + original = MagicMock() + original.close.side_effect = ValueError("close failed") + stream, _, done = wrapper(original) + with pytest.raises(ValueError, match="close failed"): + stream.close() + done.assert_called_once() + assert done.call_args.kwargs == {"outcome": "closed"} + + +async def test_async_telemetry_and_context(): + seen = [] + with trace_context(trace_id="origin"): + stream = AsyncStreamWrapper( + AsyncChunks([1]), + MagicMock(), + time.perf_counter(), + MagicMock(side_effect=ValueError()), + lambda *_: seen.append(get_trace_context().trace_id), + ) + with trace_context(trace_id="consumer"): + assert [item async for item in stream] == [1] + assert seen == ["origin"] + + +@pytest.mark.parametrize( + "factory", [make_openai_stream_callbacks, make_anthropic_stream_callbacks] +) +def test_remote_eof_without_completion_is_incomplete(factory): + handle = MagicMock() + callbacks = factory(handle, []) + stream = SyncStreamWrapper(iter([]), handle, time.perf_counter(), *callbacks) + list(stream) + stream.close() + handle.track_inference.assert_called_once() + event = handle.track_inference.call_args.kwargs + assert event["success"] is False + assert event["error_code"] == "STREAM_INCOMPLETE" diff --git a/tests/test_tensorflow_integration.py b/tests/test_tensorflow_integration.py index f27d54d..b1bedf2 100644 --- a/tests/test_tensorflow_integration.py +++ b/tests/test_tensorflow_integration.py @@ -3,6 +3,7 @@ import types from unittest.mock import MagicMock, patch +import wildedge.integrations.tensorflow as tf_mod from wildedge.integrations.tensorflow import TensorflowExtractor from wildedge.model import ModelHandle, ModelInfo @@ -60,8 +61,6 @@ def test_install_hooks_tracks_predict(self, publish_spy): assert publish_spy.events[0]["event_type"] == "inference" def test_install_auto_load_patch_is_idempotent(self, monkeypatch): - import wildedge.integrations.tensorflow as tf_mod - def load_model(*args, **kwargs): return object() diff --git a/tests/test_tracing.py b/tests/test_tracing.py index b5de0fd..24d7a2c 100644 --- a/tests/test_tracing.py +++ b/tests/test_tracing.py @@ -22,14 +22,16 @@ def publish(event: dict) -> None: publish, ) - with trace_context( - trace_id="trace-123", - run_id="run-1", - agent_id="agent-1", - attributes={"trace_key": "trace_val"}, + with ( + trace_context( + trace_id="trace-123", + run_id="run-1", + agent_id="agent-1", + attributes={"trace_key": "trace_val"}, + ), + span_context(span_id="span-abc", step_index=2, attributes={"span_key": 2}), ): - with span_context(span_id="span-abc", step_index=2, attributes={"span_key": 2}): - handle.track_inference(duration_ms=5) + handle.track_inference(duration_ms=5) assert events[0]["trace_id"] == "trace-123" assert events[0]["parent_span_id"] == "span-abc" @@ -81,9 +83,11 @@ def test_nested_spans_correct_parent_chain(): events: list[dict] = [] client = _FakeClient(events) - with SpanContextManager(client, kind="agent_step", name="outer") as outer: - with SpanContextManager(client, kind="tool", name="inner") as inner: - pass + with ( + SpanContextManager(client, kind="agent_step", name="outer") as outer, + SpanContextManager(client, kind="tool", name="inner") as inner, + ): + pass assert len(events) == 2 inner_ev, outer_ev = events[0], events[1] diff --git a/tests/test_transmitter.py b/tests/test_transmitter.py index 68fb9fb..35a37b1 100644 --- a/tests/test_transmitter.py +++ b/tests/test_transmitter.py @@ -8,6 +8,9 @@ import pytest +from wildedge.consumer import Consumer +from wildedge.platforms.device_info import DeviceInfo +from wildedge.queue import EventQueue from wildedge.transmitter import IngestResponse, TransmitError, Transmitter @@ -75,29 +78,35 @@ def test_send_401_returns_unauthorized(self): def test_send_429_raises_transmit_error(self): t = Transmitter(api_key="test-key", host="https://app.wildedge.dev") - with patch.object( - t._opener, "open", side_effect=_make_http_error(429, "Rate limited") + with ( + patch.object( + t._opener, "open", side_effect=_make_http_error(429, "Rate limited") + ), + pytest.raises(TransmitError), ): - with pytest.raises(TransmitError): - t.send({"batch_id": "b-1", "events": []}) + t.send({"batch_id": "b-1", "events": []}) def test_send_500_raises_transmit_error(self): t = Transmitter(api_key="test-key", host="https://app.wildedge.dev") - with patch.object( - t._opener, "open", side_effect=_make_http_error(500, "Server error") + with ( + patch.object( + t._opener, "open", side_effect=_make_http_error(500, "Server error") + ), + pytest.raises(TransmitError), ): - with pytest.raises(TransmitError): - t.send({"batch_id": "b-1", "events": []}) + t.send({"batch_id": "b-1", "events": []}) def test_network_error_raises_transmit_error(self): t = Transmitter(api_key="test-key", host="https://app.wildedge.dev") - with patch.object( - t._opener, - "open", - side_effect=urllib.error.URLError("Connection refused"), + with ( + patch.object( + t._opener, + "open", + side_effect=urllib.error.URLError("Connection refused"), + ), + pytest.raises(TransmitError), ): - with pytest.raises(TransmitError): - t.send({"batch_id": "b-1", "events": []}) + t.send({"batch_id": "b-1", "events": []}) def test_auth_header_sent(self): captured = {} @@ -191,19 +200,20 @@ def test_send_unexpected_status_raises_transmit_error(status_code): # cause AttributeError in the consumer (response.status), silently killing the # background thread and freezing the queue. t = Transmitter(api_key="test-key", host="https://app.wildedge.dev") - with patch.object( - t._opener, "open", return_value=_make_response(status_code, {"status": "ok"}) + with ( + patch.object( + t._opener, + "open", + return_value=_make_response(status_code, {"status": "ok"}), + ), + pytest.raises(TransmitError, match=f"Unexpected HTTP {status_code}"), ): - with pytest.raises(TransmitError, match=f"Unexpected HTTP {status_code}"): - t.send({"batch_id": "b-1", "events": []}) + t.send({"batch_id": "b-1", "events": []}) 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.platforms.device_info import DeviceInfo - from wildedge.queue import EventQueue queue = EventQueue(max_size=100) queue.add({"event_id": "e1", "event_type": "inference", "model_id": "m"}) @@ -215,7 +225,7 @@ def test_unexpected_status_keeps_events_in_consumer_queue(): queue=queue, transmitter=mock_transmitter, device=DeviceInfo(device_id="d", device_type="linux"), - get_models=lambda: {}, + get_models=dict, session_id="sess-1", ) result = consumer.drain_once() diff --git a/tests/test_whisper.py b/tests/test_whisper.py new file mode 100644 index 0000000..e0510ba --- /dev/null +++ b/tests/test_whisper.py @@ -0,0 +1,493 @@ +"""Whisper adapters preserve results, lazy lifecycles, and capture privacy.""" + +import gc +import io +import sys +import wave +from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from wildedge.attachments import AttachmentManager, AttachmentStore +from wildedge.client import WildEdge +from wildedge.instrumentation import claim_instrumentation +from wildedge.integrations import faster_whisper as faster +from wildedge.integrations.faster_whisper import ( + install_auto_load_patch, + wrap_transcribe, +) +from wildedge.integrations.mlx_whisper import install_auto_load_patch as install_mlx +from wildedge.integrations.whisper import ( + CLIENT_ATTR, + DESCRIPTOR_ATTR, + HANDLE_ATTR, + WhisperDescriptor, + WhisperExtractor, + model_name_and_source, + track_transcription, +) +from wildedge.platforms.device_info import DeviceInfo + + +@pytest.fixture +def handle(): + return SimpleNamespace( + capture_attachments=None, track_inference=MagicMock(), track_error=MagicMock() + ) + + +@pytest.fixture +def wav(tmp_path): + path = tmp_path / "input.wav" + with wave.open(str(path), "wb") as audio: + audio.setnchannels(1) + audio.setsampwidth(2) + audio.setframerate(16000) + audio.writeframes(b"\0\0" * 8000) + return path + + +def tracked_model(handle): + """A faster-whisper model object as the constructor patch leaves it.""" + model = SimpleNamespace() + setattr(model, HANDLE_ATTR, handle) + setattr(model, CLIENT_ATTR, lambda: object()) + return model + + +def stubbed_client(): + """A second client, set up like the client_with_stubbed_runtime fixture.""" + with ( + patch("wildedge.client.detect_device", return_value=DeviceInfo("id", "linux")), + patch("wildedge.client.Transmitter"), + patch("wildedge.client.Consumer"), + ): + return WildEdge( + dsn="https://secret@ingest.wildedge.dev/key", sampling_interval_s=None + ) + + +def events_of(events, event_type): + return [e for e in events if e["event_type"] == event_type] + + +# Shared metadata ------------------------------------------------------------ + + +def test_transcription_metadata(wav, handle): + track_transcription( + handle, + wav, + "hello world", + duration_ms=123, + language="en", + language_probability=0.9, + ) + event = handle.track_inference.call_args.kwargs + assert event["duration_ms"] == 123 + assert event["input_meta"].to_dict() == { + "duration_ms": 500, + "sample_rate": 16000, + "channels": 1, + "bit_depth": 16, + "format": "wav", + } + assert event["context"] == { + "transcript_word_count": 2, + "language": "en", + "language_confidence": 0.9, + } + + +@pytest.mark.parametrize("uploads_enabled", [False, True]) +def test_transcription_never_attaches_payloads(wav, handle, uploads_enabled): + if uploads_enabled: + handle.capture_attachments = MagicMock() + track_transcription(handle, wav, "hello world", duration_ms=1) + assert "attachments" not in handle.track_inference.call_args.kwargs + + +def test_file_like_audio_is_not_read(handle): + source = io.BytesIO(b"RIFF....") + handle.capture_attachments = MagicMock() + track_transcription(handle, source, "hello", duration_ms=1) + event = handle.track_inference.call_args.kwargs + assert event["input_meta"].to_dict() == {} + assert "attachments" not in event + assert source.tell() == 0 + + +def test_word_count_is_left_out_for_unspaced_languages(handle): + track_transcription(handle, None, "你好世界", duration_ms=1, language="zh") + context = handle.track_inference.call_args.kwargs["context"] + assert "transcript_word_count" not in context + + +@pytest.mark.parametrize( + ("path", "expected"), + [ + ( + "/cache/hub/models--Systran--faster-whisper-tiny/snapshots/d90ca5fe", + ("Systran/faster-whisper-tiny", "huggingface"), + ), + ("mlx-community/whisper-tiny", ("mlx-community/whisper-tiny", "huggingface")), + ("small", ("small", "huggingface")), + ], +) +def test_model_name_and_source(path, expected): + assert model_name_and_source(path) == expected + + +def test_local_model_is_named_after_its_folder(tmp_path, monkeypatch): + (tmp_path / "whisper-small").mkdir() + monkeypatch.chdir(tmp_path / "whisper-small") + assert model_name_and_source(".") == ("whisper-small", "local") + + +def test_extractor(): + model = SimpleNamespace() + setattr( + model, + DESCRIPTOR_ATTR, + WhisperDescriptor("whisper-small", "local", "faster_whisper", "cpu", "int8"), + ) + ident, info = WhisperExtractor().extract_info(model, {}) + assert ident == "faster_whisper:whisper-small" + assert info.model_source == "local" + assert info.model_format == "ctranslate2" + assert info.model_family == "whisper" + assert info.quantization == "int8" + + +# faster-whisper segment stream --------------------------------------------- + + +@pytest.mark.parametrize("outcome", ["completed", "closed", "error"]) +def test_lazy_segments(handle, outcome): + segment = SimpleNamespace(text="hello world") + info = SimpleNamespace(language="en", language_probability=0.8, duration=2) + + def segments(): + yield segment + if outcome == "error": + raise ValueError("decoder failure") + + call = wrap_transcribe(lambda *_: (segments(), info)) + stream, returned_info = call(tracked_model(handle), "audio.wav") + assert returned_info is info + handle.track_inference.assert_not_called() + assert next(stream) is segment + if outcome == "closed": + stream.close() + stream.close() + elif outcome == "error": + with pytest.raises(ValueError, match="decoder failure"): + next(stream) + handle.track_error.assert_called_once() + handle.track_inference.assert_not_called() + return + else: + assert list(stream) == [] + handle.track_inference.assert_called_once() + event = handle.track_inference.call_args.kwargs + assert event["success"] == (outcome == "completed") + assert event["context"]["transcript_word_count"] == 2 + assert "time_to_first_segment_ms" in event["context"] + assert event["input_meta"].duration_ms == 2000 + + +def test_loop_that_stops_early_is_reported_as_abandoned(handle): + segments = [SimpleNamespace(text="one"), SimpleNamespace(text="two")] + stream, _ = wrap_transcribe(lambda *_: (iter(segments), SimpleNamespace()))( + tracked_model(handle), "audio.wav" + ) + for _segment in stream: + break + handle.track_inference.assert_not_called() + del stream + gc.collect() + event = handle.track_inference.call_args.kwargs + assert event["error_code"] == "STREAM_ABANDONED" + assert event["context"]["transcript_word_count"] == 1 + + +def test_stream_that_was_never_iterated_sends_nothing(handle): + stream, _ = wrap_transcribe(lambda *_: (iter([]), SimpleNamespace()))( + tracked_model(handle), "audio.wav" + ) + del stream + gc.collect() + handle.track_inference.assert_not_called() + + +def test_stream_counts_words_without_attaching_transcript(handle): + handle.capture_attachments = MagicMock() + segments = [SimpleNamespace(text="hello "), SimpleNamespace(text="world")] + stream, _ = wrap_transcribe(lambda *_: (iter(segments), SimpleNamespace()))( + tracked_model(handle), object() + ) + assert list(stream) == segments + event = handle.track_inference.call_args.kwargs + assert "attachments" not in event + assert event["context"]["transcript_word_count"] == 2 + + +def test_transcribe_errors_preserve_exception(handle): + failure = ValueError("invalid audio") + + def original(*_args): + raise failure + + handle.track_error.side_effect = RuntimeError("telemetry failed") + with pytest.raises(ValueError, match="invalid audio") as raised: + wrap_transcribe(original)(tracked_model(handle), "missing.wav") + assert raised.value is failure + handle.track_inference.assert_not_called() + + +def test_lazy_telemetry_failure_does_not_break_iterator(handle): + handle.track_inference.side_effect = RuntimeError("telemetry failed") + segment = SimpleNamespace(text="hello") + stream, _ = wrap_transcribe(lambda *_: (iter([segment]), SimpleNamespace()))( + tracked_model(handle), "audio.wav" + ) + assert list(stream) == [segment] + assert list(stream) == [] + handle.track_inference.assert_called_once() + + +# faster-whisper model patch ------------------------------------------------- + + +@pytest.fixture +def fake_faster_whisper(monkeypatch): + module = ModuleType("faster_whisper") + + class WhisperModel: + def __init__(self, model_size_or_path, device="cpu", compute_type="int8"): + self.model = SimpleNamespace(device=device, compute_type=compute_type) + + def transcribe(self, audio): + return iter([SimpleNamespace(text="hello")]), SimpleNamespace() + + class BatchedInferencePipeline: + def __init__(self, model): + self.model = model + + def transcribe(self, audio): + return iter([SimpleNamespace(text="hello")]), SimpleNamespace() + + module.WhisperModel = WhisperModel + module.BatchedInferencePipeline = BatchedInferencePipeline + monkeypatch.setitem(sys.modules, "faster_whisper", module) + return module + + +@pytest.fixture +def instrumented_client(client_with_stubbed_runtime): + client = client_with_stubbed_runtime + client.events = [] + client.publish = client.events.append + return client + + +def test_faster_constructor_describes_the_runtime( + fake_faster_whisper, instrumented_client +): + route = claim_instrumentation(instrumented_client, "faster_whisper") + install_auto_load_patch(route) + install_auto_load_patch(route) + model = fake_faster_whisper.WhisperModel( + "small", device="cuda", compute_type="float16" + ) + info = getattr(model, DESCRIPTOR_ATTR) + assert info.model_id == "faster_whisper:small" + assert info.accelerator == "cuda" + assert info.quantization == "f16" + assert len(events_of(instrumented_client.events, "model_load")) == 1 + + +def test_each_faster_model_instance_reports_its_load( + fake_faster_whisper, instrumented_client +): + install_auto_load_patch( + claim_instrumentation(instrumented_client, "faster_whisper") + ) + fake_faster_whisper.WhisperModel("small") + fake_faster_whisper.WhisperModel("small") + assert len(events_of(instrumented_client.events, "model_load")) == 2 + + +def test_batched_pipeline_is_tracked(fake_faster_whisper, instrumented_client): + install_auto_load_patch( + claim_instrumentation(instrumented_client, "faster_whisper") + ) + model = fake_faster_whisper.WhisperModel("small") + pipeline = fake_faster_whisper.BatchedInferencePipeline(model) + segments, _ = pipeline.transcribe("audio.wav") + assert [s.text for s in segments] == ["hello"] + inference = events_of(instrumented_client.events, "inference") + assert [e["model_id"] for e in inference] == ["faster_whisper:small"] + + +def test_faster_model_stays_with_its_original_client( + fake_faster_whisper, instrumented_client +): + install_auto_load_patch( + claim_instrumentation(instrumented_client, "faster_whisper") + ) + model = fake_faster_whisper.WhisperModel("small") + instrumented_client.close() + segments, _ = model.transcribe("audio.wav") + assert not isinstance(segments, faster.SegmentStream) + assert [s.text for s in segments] == ["hello"] + assert not events_of(instrumented_client.events, "inference") + + +# mlx-whisper ---------------------------------------------------------------- + + +class FakeMlxModel: + """Stands in for an MLX Whisper model; weak-referenceable like the real one.""" + + def __init__(self, bits=None): + self.layers = [("", SimpleNamespace(bits=bits))] if bits else [] + + def named_modules(self): + return self.layers + + +@pytest.fixture +def fake_mlx_whisper(monkeypatch): + """mlx_whisper with ModelHolder's caching: one model, reloaded on change.""" + root, module, loaders = ( + ModuleType(n) + for n in ("mlx_whisper", "mlx_whisper.transcribe", "mlx_whisper.load_models") + ) + + def load_model(path_or_hf_repo, dtype="float32"): + return FakeMlxModel(bits=4) + + class ModelHolder: + model = None + model_path = None + + def transcribe(_audio, *, path_or_hf_repo="mlx-community/whisper-tiny", **_kw): + if path_or_hf_repo == "broken": + raise ValueError("cannot load") + if ModelHolder.model is None or ModelHolder.model_path != path_or_hf_repo: + ModelHolder.model = module.load_model(path_or_hf_repo) + ModelHolder.model_path = path_or_hf_repo + return {"text": "hello world", "language": "en"} + + module.ModelHolder = ModelHolder + module.transcribe = transcribe + module.load_model = loaders.load_model = load_model + for mod in (root, module, loaders): + monkeypatch.setitem(sys.modules, mod.__name__, mod) + root.holder = ModelHolder + return root + + +def test_mlx_returns_the_original_result(fake_mlx_whisper, instrumented_client): + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + result = fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model") + assert result == {"text": "hello world", "language": "en"} + assert len(events_of(instrumented_client.events, "inference")) == 1 + + +def test_mlx_model_cached_before_init_is_reported_as_a_hit( + fake_mlx_whisper, instrumented_client +): + fake_mlx_whisper.holder.model = FakeMlxModel() + fake_mlx_whisper.holder.model_path = "org/model" + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model") + (load,) = events_of(instrumented_client.events, "model_load") + assert load["model_id"] == "mlx_whisper:org/model" + assert load["attributes"]["model_cache_hit"] is True + registered = instrumented_client.registry.models["mlx_whisper:org/model"] + assert registered.quantization == "f16" + + +def test_mlx_load_reads_quantization(fake_mlx_whisper, instrumented_client): + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model-q4") + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model-q4") + (load,) = events_of(instrumented_client.events, "model_load") + assert load["load"]["accelerator"] == "mps" + assert len(events_of(instrumented_client.events, "inference")) == 2 + registered = instrumented_client.registry.models["mlx_whisper:org/model-q4"] + assert registered.quantization == "int4" + + +def test_mlx_reload_after_switching_models_is_reported( + fake_mlx_whisper, instrumented_client +): + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + for repo in ("org/a", "org/b", "org/a"): + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo=repo) + loads = events_of(instrumented_client.events, "model_load") + assert [e["model_id"] for e in loads] == [ + "mlx_whisper:org/a", + "mlx_whisper:org/b", + "mlx_whisper:org/a", + ] + + +def test_mlx_error_is_tracked_and_reraised(fake_mlx_whisper, instrumented_client): + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + with pytest.raises(ValueError, match="cannot load"): + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="broken") + errors = events_of(instrumented_client.events, "error") + assert [e["model_id"] for e in errors] == ["mlx_whisper:broken"] + + +def test_mlx_sends_nothing_after_client_closes(fake_mlx_whisper, instrumented_client): + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model") + instrumented_client.close() + assert fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model") + assert len(events_of(instrumented_client.events, "inference")) == 1 + + +def test_mlx_replacement_client_reports_the_cached_model( + fake_mlx_whisper, instrumented_client +): + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model") + instrumented_client.close() + replacement = stubbed_client() + replacement_events = [] + replacement.publish = replacement_events.append + claim_instrumentation(replacement, "mlx_whisper") + fake_mlx_whisper.transcribe("audio.wav", path_or_hf_repo="org/model") + (load,) = events_of(replacement_events, "model_load") + assert load["attributes"]["model_cache_hit"] is True + assert len(events_of(replacement_events, "inference")) == 1 + assert len(events_of(instrumented_client.events, "inference")) == 1 + + +@pytest.mark.parametrize("backend", ["faster_whisper", "batched", "mlx_whisper"]) +def test_instrumentation_buffers_no_payloads_with_uploads_enabled( + backend, fake_faster_whisper, fake_mlx_whisper, instrumented_client, tmp_path, wav +): + store = AttachmentStore(str(tmp_path / "attachments")) + instrumented_client.attachment_manager = AttachmentManager(store) + if backend == "mlx_whisper": + install_mlx(claim_instrumentation(instrumented_client, "mlx_whisper")) + result = fake_mlx_whisper.transcribe(str(wav), path_or_hf_repo="org/model") + assert result["text"] + else: + install_auto_load_patch( + claim_instrumentation(instrumented_client, "faster_whisper") + ) + model = fake_faster_whisper.WhisperModel("small") + if backend == "batched": + model = fake_faster_whisper.BatchedInferencePipeline(model) + segments, _ = model.transcribe(str(wav)) + assert list(segments) + (event,) = events_of(instrumented_client.events, "inference") + assert "attachments" not in event + assert store.length() == 0 diff --git a/uv.lock b/uv.lock index 8ff0ac1..a33da49 100644 --- a/uv.lock +++ b/uv.lock @@ -419,7 +419,7 @@ wheels = [ [[package]] name = "wildedge-sdk" -version = "0.2.2" +version = "0.3.0" source = { editable = "." } [package.dev-dependencies] diff --git a/wildedge/__init__.py b/wildedge/__init__.py index 83aa9ec..63deebc 100644 --- a/wildedge/__init__.py +++ b/wildedge/__init__.py @@ -44,44 +44,44 @@ ) __all__ = [ - "WildEdge", - "init", - "get_client", - "set_default_client", - "trace", - "span", - "track_span", - "register_model", - "flush", - "llm_api", - "LLMCall", + "AdapterDownload", + "AdapterLoad", "Attachment", - "capture_hardware", - "HardwareContext", - "ThermalContext", - "track", - "QueuePolicy", - "DeviceInfo", - "ImageInputMeta", "AudioInputMeta", - "TextInputMeta", - "DetectionOutputMeta", "ClassificationOutputMeta", - "GenerationOutputMeta", + "DetectionOutputMeta", + "DeviceInfo", "EmbeddingOutputMeta", + "ErrorCode", + "FeedbackType", "GenerationConfig", - "AdapterLoad", - "AdapterDownload", + "GenerationOutputMeta", + "HardwareContext", + "ImageInputMeta", + "LLMCall", + "QueuePolicy", + "SpanContext", + "SpanContextManager", "SpanEvent", - "FeedbackType", - "ErrorCode", + "SpanKind", + "SpanStatus", + "TextInputMeta", + "ThermalContext", "Timer", - "span_context", "TraceContext", - "SpanContext", - "get_trace_context", + "WildEdge", + "capture_hardware", + "flush", + "get_client", "get_span_context", - "SpanKind", - "SpanStatus", - "SpanContextManager", + "get_trace_context", + "init", + "llm_api", + "register_model", + "set_default_client", + "span", + "span_context", + "trace", + "track", + "track_span", ] diff --git a/wildedge/accelerators.py b/wildedge/accelerators.py new file mode 100644 index 0000000..617e395 --- /dev/null +++ b/wildedge/accelerators.py @@ -0,0 +1,13 @@ +"""Accelerator names reported in model load events and hardware context. + +Integrations report the most specific name they know. PyTorch device types +(for example ``xpu``) pass through as they are, so this list is not closed. +""" + +CPU = "cpu" +CUDA = "cuda" # NVIDIA GPU +ROCM = "rocm" # AMD GPU +MPS = "mps" # Apple GPU (Metal) +COREML = "coreml" # Apple Core ML, which may use the GPU or Neural Engine +DIRECTML = "directml" # Windows DirectML +OPENVINO = "openvino" # Intel OpenVINO diff --git a/wildedge/attachments.py b/wildedge/attachments.py index 9636120..594cf54 100644 --- a/wildedge/attachments.py +++ b/wildedge/attachments.py @@ -349,8 +349,18 @@ def __init__( self.host = host.rstrip("/") self.timeout = timeout - def transmit(self, pending: PendingAttachment) -> UploadOutcome: - outcome, upload_url = self._presign(pending) + def transmit( + self, pending: PendingAttachment, deadline: float | None = None + ) -> UploadOutcome: + """Presign and upload ``pending``. + + ``deadline`` (a ``time.monotonic()`` value) caps each request's timeout; + past it the upload counts as transient and stays buffered. + """ + timeout = self._request_timeout(deadline) + if timeout is None: + return UploadOutcome.TRANSIENT + outcome, upload_url = self._presign(pending, timeout) if upload_url is None: return outcome try: @@ -358,9 +368,22 @@ def transmit(self, pending: PendingAttachment) -> UploadOutcome: except Exception as exc: logger.warning("wildedge: attachment bytes unreadable: %s", exc) return UploadOutcome.PERMANENT - return self._put(upload_url, data, pending.content_type) + timeout = self._request_timeout(deadline) + if timeout is None: + return UploadOutcome.TRANSIENT + return self._put(upload_url, data, pending.content_type, timeout) - def _presign(self, pending: PendingAttachment) -> tuple[UploadOutcome, str | None]: + def _request_timeout(self, deadline: float | None) -> float | None: + if deadline is None: + return self.timeout + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + return min(self.timeout, remaining) + + def _presign( + self, pending: PendingAttachment, timeout: float + ) -> tuple[UploadOutcome, str | None]: url = f"{self.host}/api/attachments/presign" body = json.dumps( [ @@ -384,7 +407,7 @@ def _presign(self, pending: PendingAttachment) -> tuple[UploadOutcome, str | Non method="POST", ) try: - with urllib.request.urlopen(req, timeout=self.timeout) as resp: + with urllib.request.urlopen(req, timeout=timeout) as resp: status = resp.status raw = resp.read() except urllib.error.HTTPError as exc: @@ -421,7 +444,9 @@ def _presign(self, pending: PendingAttachment) -> tuple[UploadOutcome, str | Non # 429 / 5xx / anything else: retry later. return UploadOutcome.TRANSIENT, None - def _put(self, upload_url: str, data: bytes, content_type: str) -> UploadOutcome: + def _put( + self, upload_url: str, data: bytes, content_type: str, timeout: float + ) -> UploadOutcome: req = urllib.request.Request( upload_url, data=data, @@ -429,7 +454,7 @@ def _put(self, upload_url: str, data: bytes, content_type: str) -> UploadOutcome method="PUT", ) try: - with urllib.request.urlopen(req, timeout=self.timeout) as resp: + with urllib.request.urlopen(req, timeout=timeout) as resp: status = resp.status except urllib.error.HTTPError as exc: status = exc.code @@ -467,6 +492,10 @@ def __init__( self.uploaded_count = 0 self.dropped_count = 0 self.stop_event = threading.Event() + # Serialises the background thread and drain() so no attachment is + # sent twice. + self.process_lock = threading.Lock() + self.disabled = False self.thread = self._make_thread() def _make_thread(self) -> threading.Thread: @@ -501,10 +530,50 @@ def _run(self) -> None: except Exception as exc: # pragma: no cover - never kill the thread logger.warning("wildedge: attachment uploader error: %s", exc) - def process_pending(self) -> None: + def drain(self, timeout: float) -> None: + """Upload buffered attachments until none are left or ``timeout`` runs out. + + Stops early when an upload fails transiently; what is left stays on + disk for the next start. + """ + deadline = time.monotonic() + timeout + while self.process_pending(deadline=deadline): + pass + remaining = self.store.length() + if remaining and not self.disabled: + logger.debug( + "wildedge: %d attachment(s) left buffered for the next start", + remaining, + ) + + def process_pending(self, deadline: float | None = None) -> bool: + """Upload one batch, oldest first. Returns True when more may be pending. + + ``deadline`` (a ``time.monotonic()`` value) bounds the call, including + the wait for a batch the background thread is sending. + """ + if deadline is None: + acquired = self.process_lock.acquire() + else: + acquired = self.process_lock.acquire( + timeout=max(deadline - time.monotonic(), 0) + ) + if not acquired: + return False + try: + return self._process_batch(deadline) + finally: + self.process_lock.release() + + def _process_batch(self, deadline: float | None) -> bool: + if self.disabled: + return False # Evict stale buffered bytes from the front of the queue. cutoff = time.time() - self.max_age_s batch = self.store.peek_many(constants.ATTACHMENT_PRESIGN_BATCH_SIZE) + if not batch: + return False + full_batch = len(batch) == constants.ATTACHMENT_PRESIGN_BATCH_SIZE stale = [p for p in batch if p.registered_at < cutoff] if stale: self.store.remove(stale) @@ -517,8 +586,9 @@ def process_pending(self) -> None: batch = [p for p in batch if p.registered_at >= cutoff] done: list[PendingAttachment] = [] + more = full_batch for pending in batch: - outcome = self.transmitter.transmit(pending) + outcome = self.transmitter.transmit(pending, deadline=deadline) if outcome is UploadOutcome.UPLOADED: done.append(pending) self.uploaded_count += 1 @@ -530,10 +600,13 @@ def process_pending(self) -> None: self.on_feature_disabled() self.store.remove(done) self.store.clear() + self.disabled = True self.stop_event.set() - return + return False else: # TRANSIENT - stop, preserve order, retry next tick + more = False break if done: self.store.remove(done) + return more diff --git a/wildedge/autoload/sitecustomize.py b/wildedge/autoload/sitecustomize.py index e3a9dbd..86fa03f 100644 --- a/wildedge/autoload/sitecustomize.py +++ b/wildedge/autoload/sitecustomize.py @@ -15,6 +15,7 @@ from __future__ import annotations +import importlib.util as _iutil import os import sys @@ -37,14 +38,14 @@ def _die(code: int) -> None: os._exit(code) -if _INSTALLED_MARKER not in sys.modules: +if _INSTALLED_MARKER not in sys.modules: # noqa: SIM102 - re-entry guard kept separate # Activate if the CLI launched this process, or if WILDEDGE_DSN is set # and the user has manually prepended wildedge/autoload/ to PYTHONPATH. if os.environ.get(_AUTOLOAD) or os.environ.get(_DSN): # Set the guard before importing wildedge to prevent re-entry. sys.modules[_INSTALLED_MARKER] = True # type: ignore[assignment] try: - from wildedge.runtime.bootstrap import ( # noqa: PLC0415 + from wildedge.runtime.bootstrap import ( EXIT_BOOTSTRAP_INTERNAL_ERROR, EXIT_CONFIG_ERROR, EXIT_STRICT_INTEGRATION_ERROR, @@ -54,7 +55,7 @@ def _die(code: int) -> None: format_startup_report, install_runtime, ) - from wildedge.settings import read_runner_env # noqa: PLC0415 + from wildedge.settings import read_runner_env _runner_env = read_runner_env() _context = None @@ -99,8 +100,6 @@ def _die(code: int) -> None: # Use importlib to find and exec it directly; avoids sys.modules manipulation # which can trigger CPython's module GC and clear globals mid-execution. def _load_existing_sitecustomize() -> None: - import importlib.util as _iutil - _autoload_dir = os.path.dirname(os.path.abspath(__file__)) _saved = sys.path[:] sys.path = [p for p in sys.path if p != _autoload_dir] diff --git a/wildedge/cli.py b/wildedge/cli.py index 2199335..5b6e11a 100644 --- a/wildedge/cli.py +++ b/wildedge/cli.py @@ -21,7 +21,11 @@ from wildedge.client import parse_dsn from wildedge.events.span import SpanEvent from wildedge.hubs.registry import HUBS_BY_NAME, supported_hubs -from wildedge.integrations.registry import INTEGRATIONS_BY_NAME, supported_integrations +from wildedge.integrations.registry import ( + INTEGRATIONS_BY_NAME, + canonical_integration, + supported_integrations, +) from wildedge.paths import default_dead_letter_dir, default_pending_queue_dir from wildedge.platforms import detect_device, get_device_id_path from wildedge.settings import read_client_env, resolve_app_identity @@ -44,7 +48,7 @@ def build_parser() -> argparse.ArgumentParser: run.add_argument( "--integrations", default="all", - help="Comma-separated framework integrations to enable (gguf, onnx, timm, …). Default: all.", + help="Comma-separated framework integrations to enable (llama_cpp, onnx, timm, …). Default: all.", ) run.add_argument( "--hubs", @@ -291,7 +295,13 @@ def run_command(parsed: argparse.Namespace) -> int: def integration_list(value: str | None) -> list[str]: if not value or value == "all": return sorted(supported_integrations()) - return [item.strip() for item in value.split(",") if item.strip()] + return list( + dict.fromkeys( + canonical_integration(item.strip()) + for item in value.split(",") + if item.strip() + ) + ) def hub_list(value: str | None) -> list[str]: diff --git a/wildedge/client.py b/wildedge/client.py index 736c467..83862ee 100644 --- a/wildedge/client.py +++ b/wildedge/client.py @@ -6,7 +6,7 @@ import weakref from collections.abc import Callable from datetime import datetime, timezone -from typing import Any +from typing import Any, ClassVar from urllib.parse import urlparse from wildedge import constants @@ -22,21 +22,23 @@ from wildedge.events import SpanEvent from wildedge.events.span import SpanKind, SpanStatus from wildedge.hubs.base import BaseHubTracker -from wildedge.hubs.huggingface import HuggingFaceHubTracker -from wildedge.hubs.registry import supported_hubs -from wildedge.hubs.torchhub import TorchHubTracker -from wildedge.integrations.anthropic import AnthropicExtractor +from wildedge.hubs.registry import HUB_SPECS, supported_hubs +from wildedge.instrumentation import ( + ClientRoute, + InstrumentationOwnedError, + claim_instrumentation, + has_other_live_owner, + strict_ownership, +) from wildedge.integrations.base import BaseExtractor -from wildedge.integrations.gguf import GgufExtractor -from wildedge.integrations.keras import KerasExtractor -from wildedge.integrations.mlx import MlxExtractor -from wildedge.integrations.onnx import OnnxExtractor -from wildedge.integrations.openai import OpenAIExtractor -from wildedge.integrations.pytorch import PytorchExtractor -from wildedge.integrations.registry import noop_integrations, supported_integrations -from wildedge.integrations.tensorflow import TensorflowExtractor -from wildedge.integrations.transformers import TransformersExtractor -from wildedge.integrations.ultralytics import UltralyticsExtractor +from wildedge.integrations.registry import ( + INTEGRATION_SPECS, + canonical_integration, + find_extractor, + noop_integrations, + supported_integrations, +) +from wildedge.loading import lazy_callable from wildedge.logging import enable_debug, logger from wildedge.model import ModelHandle, ModelInfo, ModelRegistry from wildedge.paths import ( @@ -112,19 +114,6 @@ def parse_dsn(dsn: str) -> tuple[str, str, str]: # One informative log line per process for the intentional no-DSN mode. _noop_notice_logged = False -DEFAULT_EXTRACTORS: list[BaseExtractor] = [ - OnnxExtractor(), - GgufExtractor(), - OpenAIExtractor(), - AnthropicExtractor(), - UltralyticsExtractor(), - TransformersExtractor(), - MlxExtractor(), - PytorchExtractor(), - TensorflowExtractor(), - KerasExtractor(), -] - class SpanContextManager: def __init__( @@ -248,23 +237,16 @@ class WildEdge: # Framework integrations: patch ML runtime constructors for inference tracking. SUPPORTED_INTEGRATIONS = supported_integrations() NOOP_INTEGRATIONS = noop_integrations() - PATCH_INSTALLERS = { - "anthropic": AnthropicExtractor.install_auto_load_patch, - "gguf": GgufExtractor.install_auto_load_patch, - "mlx": MlxExtractor.install_auto_load_patch, - "onnx": OnnxExtractor.install_auto_load_patch, - "openai": OpenAIExtractor.install_auto_load_patch, - "timm": PytorchExtractor.install_timm_patch, - "tensorflow": TensorflowExtractor.install_auto_load_patch, - "transformers": TransformersExtractor.install_auto_load_patch, - "ultralytics": UltralyticsExtractor.install_auto_load_patch, + PATCH_INSTALLERS: ClassVar[dict[str, Callable[[object], None]]] = { + spec.name: lazy_callable(spec.installer) + for spec in INTEGRATION_SPECS + if spec.installer is not None } # Hub trackers: record download provenance (where models came from). SUPPORTED_HUBS = supported_hubs() - HUB_TRACKER_CLASSES: dict[str, type[BaseHubTracker]] = { - "huggingface": HuggingFaceHubTracker, - "torchhub": TorchHubTracker, + HUB_TRACKER_CLASSES: ClassVar[dict[str, Callable[[], BaseHubTracker]]] = { + spec.name: lazy_callable(spec.tracker) for spec in HUB_SPECS } def __init__( @@ -304,10 +286,14 @@ def __init__( dsn = env.dsn debug = env.debug self.noop = False + # (secret, host, project_key), so init(dsn=...) can tell when a + # requested destination matches this client. None for a no-op client. + self.destination: tuple[str, str, str] | None = None if not dsn: self._init_noop(debug=debug, device=device) return api_key, host, project_key = parse_dsn(dsn) + self.destination = (api_key, host, project_key) app_identity = resolve_app_identity( explicit=env.app_identity, project_key=project_key, @@ -482,6 +468,15 @@ def _init_noop(self, *, debug: bool, device: DeviceInfo | None) -> None: def publish(self, event_dict: dict) -> None: if self.closed or self.noop: + if ( + not self.noop + and event_dict.get("event_type") in ("inference", "error") + and not getattr(self, "_closed_publish_warned", False) + ): + self._closed_publish_warned = True + logger.warning( + "wildedge: event belongs to a closed SDK client; recreate models after initializing its replacement" + ) return if self.debug: @@ -686,10 +681,7 @@ def track_span( return correlation["span_id"] def _find_extractor(self, model_obj: object) -> BaseExtractor | None: - for candidate in DEFAULT_EXTRACTORS: - if candidate.can_handle(model_obj): - return candidate - return None + return find_extractor(model_obj) def _memory_bytes_for(self, model_obj: object) -> int | None: extractor = self._find_extractor(model_obj) @@ -707,11 +699,31 @@ def _activate_hub(self, hub_name: str) -> None: "wildedge: hub '%s' library not installed; skipping", hub_name ) return - tracker.install_patch(weakref.ref(self)) + route = self._claim_instrumentation("hub:" + hub_name) + if route is None: + return + tracker.install_patch(route) self.hub_trackers[hub_name] = tracker if self.debug: logger.debug("wildedge: hub tracker '%s' activated", hub_name) + def _claim_instrumentation(self, name: str) -> ClientRoute | None: + """Claim process-wide instrumentation, or warn and skip if owned.""" + try: + return claim_instrumentation(self, name) + except InstrumentationOwnedError: + if strict_ownership(): + raise + logger.warning( + "wildedge: another WildEdge client (for example the one " + "`wildedge run` installed) already owns automatic tracking, so " + "'%s' events keep going to that client's project. Call " + "wildedge.init() without dsn to reuse that client, or close it " + "before enabling integrations on this one.", + name, + ) + return None + def _snapshot_hub_caches(self) -> dict[str, dict[str, int]]: """Return a before-snapshot of all active hub cache directories.""" return { @@ -758,7 +770,7 @@ def instrument( Pass ``hubs=`` alongside a framework name to activate download provenance tracking for the hub(s) your code downloads from:: - client.instrument("gguf", hubs=["huggingface"]) + client.instrument("llama_cpp", hubs=["huggingface"]) client.instrument("onnx", hubs=["huggingface"]) client.instrument("timm", hubs=["huggingface", "torchhub"]) @@ -772,11 +784,13 @@ def instrument( Supported framework integrations --------------------------------- - ``"gguf"`` + ``"llama_cpp"`` (``"gguf"`` remains a compatibility alias) Patches ``llama_cpp.Llama.__init__``. Requires ``llama-cpp-python``. ``"onnx"`` Replaces ``ort.InferenceSession`` with a timed subclass. Requires ``onnxruntime``. + ``"mlx_whisper"`` / ``"faster_whisper"`` + Tracks local speech recognition, including lazy segment consumption. ``"timm"`` Patches ``timm.create_model``. Requires ``timm``. ``"tensorflow"`` @@ -814,6 +828,8 @@ def instrument( if self.debug: logger.debug("wildedge: instrument skipped (no DSN configured)") return + if self.closed: + raise RuntimeError("Cannot instrument with a closed WildEdge client") if integration is None: if not hubs: raise ValueError( @@ -831,6 +847,7 @@ def instrument( self._activate_hub(hub_name) return + integration = canonical_integration(integration) if integration in self.SUPPORTED_HUBS: raise ValueError( f"{integration!r} is a hub, not a framework integration. " @@ -859,7 +876,9 @@ def instrument( logger.debug(LOG_INSTRUMENT_TORCH_KERAS, integration) else: installer = self.PATCH_INSTALLERS[integration] - installer(weakref.ref(self)) + route = self._claim_instrumentation(integration) + if route is not None: + installer(route) for hub_name in hubs or []: self._activate_hub(hub_name) @@ -871,8 +890,14 @@ def _on_model_auto_loaded( downloads: list[dict] | None = None, model_id: str | None = None, load_kwargs: dict | None = None, + reload: bool = False, ) -> None: - """Callback invoked by auto-patches after a model is constructed.""" + """Callback invoked by auto-patches after a model is constructed. + + A model ID seen before normally means a wrapper around a model that is + already loaded, so no second load event is sent. ``reload=True`` says + the weights were loaded again and records the load anyway. + """ if self.closed: return # Always drain thread-local hub buffers to keep them clean. @@ -884,7 +909,7 @@ 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 + already_tracked = handle.model_id in self.auto_loaded and not reload self.auto_loaded.add(handle.model_id) if already_tracked: @@ -994,20 +1019,44 @@ def _on_unload() -> None: return obj def flush(self, timeout: float = 5.0) -> None: - """Block until the event queue drains or timeout expires.""" + """Block until queued events and attachments are sent or timeout expires. + + Events go first; buffered attachments use the time that is left. + """ + deadline = time.monotonic() + timeout self.consumer.flush(timeout=timeout) + self._drain_attachments(deadline) def close(self, timeout: float | None = None) -> None: - """Best-effort shutdown; pass timeout to attempt bounded flush first.""" + """Best-effort shutdown; pass timeout to attempt bounded flush first. + + With a timeout, buffered attachments are uploaded in the time left after + the event flush. Without one, they stay on disk for the next start. + """ + if self.closed: + return self.closed = True - if not self.noop: + if not self.noop and not has_other_live_owner(self): stop_sampler() if self.attachment_uploader is not None: self.attachment_uploader.stop() if timeout is None: self.consumer.close() else: + deadline = time.monotonic() + timeout self.consumer.close(timeout=timeout) + self._drain_attachments(deadline) + + def _drain_attachments(self, deadline: float) -> None: + if self.attachment_uploader is None: + return + remaining = deadline - time.monotonic() + if remaining <= 0: + return + try: + self.attachment_uploader.drain(remaining) + except Exception as exc: + logger.warning("wildedge: attachment upload on flush failed: %s", exc) def pending_count(self) -> int: """Return the number of events currently buffered.""" diff --git a/wildedge/constants.py b/wildedge/constants.py index 22ef327..acd0fcc 100644 --- a/wildedge/constants.py +++ b/wildedge/constants.py @@ -53,7 +53,7 @@ IDLE_POLL_INTERVAL = 5.0 # Error message truncation -ERROR_MSG_MAX_LEN = 200 +ERROR_MSG_MAX_LEN = 500 PROVIDER_ERROR_CODE_MAX_LEN = 64 # Device ID persistence diff --git a/wildedge/consumer.py b/wildedge/consumer.py index ef8c15d..44ed595 100644 --- a/wildedge/consumer.py +++ b/wildedge/consumer.py @@ -50,6 +50,10 @@ def __init__( self.stop_event = threading.Event() self.stopped = False + # One drain at a time: flush() on the caller's thread and the background + # thread would otherwise send the same batch and both remove it. + self.drain_lock = threading.Lock() + self.fork_holds_drain_lock = False self.backoff = constants.BACKOFF_MIN self.created_at = datetime.now(timezone.utc) @@ -133,7 +137,16 @@ def dead_letter_and_drop( self.queue.remove_first_n(len(events)) self.notify_delivery_failure(reason, len(events)) - def drain_once(self) -> bool: + def drain_once(self, timeout: float = -1) -> bool: + """Send one batch. Returns False if no progress was made.""" + if not self.drain_lock.acquire(timeout=timeout): + return False + try: + return self.drain_once_locked() + finally: + self.drain_lock.release() + + def drain_once_locked(self) -> bool: events = self.queue.peek_many(self.batch_size) if not events: return False @@ -218,7 +231,7 @@ def flush(self, timeout: float = 5.0) -> None: deadline = time.monotonic() + timeout backoff = constants.BACKOFF_MIN while self.queue.length() > 0 and time.monotonic() < deadline: - progressed = self.drain_once() + progressed = self.drain_once(timeout=max(0.0, deadline - time.monotonic())) if self.queue.length() == 0: break remaining = deadline - time.monotonic() @@ -243,6 +256,9 @@ def _pause(self) -> None: """ self.stop_event.set() self.thread.join(timeout=1.0) + # Hold the drain lock across fork() so the child never inherits it + # locked by a thread that does not exist there. + self.fork_holds_drain_lock = self.drain_lock.acquire(timeout=1.0) # Reset so _resume() can start a fresh thread in parent and child. self.stopped = False @@ -253,6 +269,12 @@ def _resume(self) -> None: Creates a new Event and Thread so parent and child each get an independent consumer. """ + if self.fork_holds_drain_lock: + self.drain_lock.release() + self.fork_holds_drain_lock = False + else: + # A send outlived _pause(). Start over rather than wait on it. + self.drain_lock = threading.Lock() self.stop_event = threading.Event() self.backoff = constants.BACKOFF_MIN self._held_snapshot = None diff --git a/wildedge/convenience.py b/wildedge/convenience.py index bfa2f1a..1f15de9 100644 --- a/wildedge/convenience.py +++ b/wildedge/convenience.py @@ -3,9 +3,10 @@ from collections.abc import Iterable from typing import Any -from wildedge.client import WildEdge +from wildedge.client import WildEdge, parse_dsn from wildedge.defaults import peek_default_client, set_default_client from wildedge.logging import logger +from wildedge.settings import read_client_env def _normalize_list(value: str | Iterable[str] | None) -> list[str]: @@ -16,6 +17,20 @@ def _normalize_list(value: str | Iterable[str] | None) -> list[str]: return [item for item in value if item] +def _same_destination(client: WildEdge, dsn: str | None) -> bool: + """True when ``dsn`` (or WILDEDGE_DSN) points where ``client`` sends.""" + destination = getattr(client, "destination", None) + if destination is None: + return False + resolved = read_client_env(dsn=dsn).dsn + if not resolved: + return False + try: + return parse_dsn(resolved) == destination + except ValueError: + return False + + def init( *, integrations: str | Iterable[str] | None = None, @@ -25,34 +40,53 @@ def init( """ Initialize the process default client and instrument integrations. - 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. + When a live default client already exists (typically installed by ``wildedge + run`` before user code loaded) and no ``dsn`` is passed, or ``dsn`` points + to the same project, that client is reused: the requested integrations are + applied to it (instrumenting is idempotent per process) and it is returned. + A ``dsn`` for another project constructs a fresh client and makes it the new + default. A closed default is replaced. While another live client owns + process-wide instrumentation, the new client logs a warning and leaves + automatic tracking with that owner. With ``wildedge run + --strict-integrations`` it raises ``RuntimeError`` instead. Additional keyword arguments are forwarded to WildEdge(...) when a new client is constructed. """ - client = None if "dsn" in kwargs else peek_default_client() - if client is not None and kwargs: + client = peek_default_client() + if client is not None and client.closed: + client = None + if ( + client is not None + and "dsn" in kwargs + and not _same_destination(client, kwargs["dsn"]) + ): + client = None + ignored = sorted(key for key in kwargs if key != "dsn") + if client is not None and ignored: logger.warning( "wildedge: init() reusing the existing default client; ignoring %s", - ", ".join(sorted(kwargs)), + ", ".join(ignored), ) - if client is None: + created = client is None + if created: client = WildEdge(**kwargs) normalized_integrations = _normalize_list(integrations) normalized_hubs = _normalize_list(hubs) - if normalized_integrations: - for integration in normalized_integrations: - client.instrument(integration, hubs=normalized_hubs or None) - elif normalized_hubs: - client.instrument(None, hubs=normalized_hubs) - elif getattr(client, "debug", False): - logger.debug("wildedge: init called without integrations or hubs") + try: + if normalized_integrations: + for integration in normalized_integrations: + client.instrument(integration, hubs=normalized_hubs or None) + elif normalized_hubs: + client.instrument(None, hubs=normalized_hubs) + elif getattr(client, "debug", False): + logger.debug("wildedge: init called without integrations or hubs") + except Exception: + if created: + client.close() + raise set_default_client(client) return client diff --git a/wildedge/decorators.py b/wildedge/decorators.py index f953159..23e90b2 100644 --- a/wildedge/decorators.py +++ b/wildedge/decorators.py @@ -4,7 +4,6 @@ import time from typing import TYPE_CHECKING, Any -from wildedge import constants from wildedge.timing import elapsed_ms if TYPE_CHECKING: @@ -54,16 +53,16 @@ def __init__( self.input_meta = input_meta self.output_meta = output_meta self.generation_config = generation_config - self._correlation = dict( - 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, - ) + self._correlation = { + "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, + } self.start_time: float | None = None def __call__(self, func): @@ -87,7 +86,7 @@ def wrapper(*args, **kwargs): if self.capture_errors: self.handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), **self._correlation, ) raise @@ -106,9 +105,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): if self.capture_errors: self.handle.track_error( error_code="UNKNOWN", - error_message=str(exc_val)[: constants.ERROR_MSG_MAX_LEN] - if exc_val - else None, + error_message=str(exc_val) if exc_val else None, **self._correlation, ) else: diff --git a/wildedge/events/__init__.py b/wildedge/events/__init__.py index 6648c40..3a75e29 100644 --- a/wildedge/events/__init__.py +++ b/wildedge/events/__init__.py @@ -22,9 +22,9 @@ from wildedge.events.span import SpanEvent __all__ = [ - "ApiMeta", "AdapterDownload", "AdapterLoad", + "ApiMeta", "AudioInputMeta", "ClassificationOutputMeta", "DetectionOutputMeta", diff --git a/wildedge/events/common.py b/wildedge/events/common.py index 399b55b..2d8daef 100644 --- a/wildedge/events/common.py +++ b/wildedge/events/common.py @@ -1,7 +1,15 @@ from __future__ import annotations +import json +import threading from typing import Any +from wildedge.logging import logger + +_SCALAR_TYPES = (str, int, float, bool, type(None)) +_warned_attribute_keys: set[str] = set() +_warned_lock = threading.Lock() + def add_optional_fields(event: dict, fields: dict[str, Any]) -> dict: """Add non-None fields to an event payload.""" @@ -9,3 +17,33 @@ def add_optional_fields(event: dict, fields: dict[str, Any]) -> dict: if value is not None: event[key] = value return event + + +def flatten_attributes(attributes: dict[str, Any]) -> dict[str, Any]: + """Return attributes with every value a scalar the ingest API accepts. + + The server rejects the whole batch when one attribute value is a list, + dict or other non-scalar. Such values are sent as JSON strings instead + (``["a", "b"]`` becomes ``'["a", "b"]'``), with one warning per key. + """ + flat: dict[str, Any] = {} + for key, value in attributes.items(): + key = str(key) + if isinstance(value, _SCALAR_TYPES): + flat[key] = value + continue + try: + flat[key] = json.dumps(value, default=str, ensure_ascii=False) + except (TypeError, ValueError): + flat[key] = str(value) + with _warned_lock: + first = key not in _warned_attribute_keys + _warned_attribute_keys.add(key) + if first: + logger.warning( + "wildedge: span attribute %r is a %s, not a str, int, float or " + "bool; sending it as a JSON string", + key, + type(value).__name__, + ) + return flat diff --git a/wildedge/events/error.py b/wildedge/events/error.py index c5272c8..a3d6bd2 100644 --- a/wildedge/events/error.py +++ b/wildedge/events/error.py @@ -6,6 +6,7 @@ from enum import Enum from typing import Any +from wildedge import constants from wildedge.events.common import add_optional_fields @@ -19,6 +20,29 @@ class ErrorCode(str, Enum): UNKNOWN = "UNKNOWN" +ELLIPSIS = " … " + + +def truncate(value: str | None, limit: int) -> str | None: + return None if value is None else str(value)[:limit] + + +def truncate_message(value: str | None, limit: int) -> str | None: + """Keep the start and the end of a long message. + + Errors put their detail at either end. mlx-whisper prints the ffmpeg banner + first and the failure last, and Hugging Face Hub starts with a request id. + """ + if value is None: + return None + value = str(value) + if len(value) <= limit: + return value + keep = limit - len(ELLIPSIS) + head = keep - keep // 2 + return value[:head] + ELLIPSIS + value[len(value) - keep // 2 :] + + @dataclass class ErrorEvent: model_id: str @@ -49,11 +73,17 @@ def to_dict(self) -> dict: error_data = add_optional_fields( {"error_code": code}, { - "error_message": self.error_message, + # Capped here, so every error event gets the same limit, + # from an integration or a manual track_error() call. + "error_message": truncate_message( + self.error_message, constants.ERROR_MSG_MAX_LEN + ), "stack_trace_hash": self.stack_trace_hash, "related_event_id": self.related_event_id, "http_status": self.http_status, - "provider_error_code": self.provider_error_code, + "provider_error_code": truncate( + self.provider_error_code, constants.PROVIDER_ERROR_CODE_MAX_LEN + ), "duration_ms": self.duration_ms, }, ) diff --git a/wildedge/events/inference.py b/wildedge/events/inference.py index d8acc40..2c64d19 100644 --- a/wildedge/events/inference.py +++ b/wildedge/events/inference.py @@ -275,11 +275,14 @@ def to_dict(self) -> dict: class EmbeddingOutputMeta: task: str = "embedding" dimensions: int | None = None + tokens_in: int | None = None def to_dict(self) -> dict: d: dict[str, Any] = {"task": self.task} if self.dimensions is not None: d["dimensions"] = self.dimensions + if self.tokens_in is not None: + d["tokens_in"] = self.tokens_in return d diff --git a/wildedge/events/model_load.py b/wildedge/events/model_load.py index 43f1ad2..b58ed16 100644 --- a/wildedge/events/model_load.py +++ b/wildedge/events/model_load.py @@ -5,6 +5,7 @@ from datetime import datetime, timezone from typing import Any +from wildedge import accelerators from wildedge.events.common import add_optional_fields @@ -42,7 +43,7 @@ class ModelLoadEvent: model_id: str duration_ms: int memory_bytes: int | None = None - accelerator: str = "cpu" + accelerator: str = accelerators.CPU success: bool = True error_code: str | None = None peak_memory_bytes: int | None = None diff --git a/wildedge/events/span.py b/wildedge/events/span.py index 9d5be3c..b0785b5 100644 --- a/wildedge/events/span.py +++ b/wildedge/events/span.py @@ -5,7 +5,7 @@ from datetime import datetime, timezone from typing import Any, Literal -from wildedge.events.common import add_optional_fields +from wildedge.events.common import add_optional_fields, flatten_attributes SpanKind = Literal[ "agent_step", @@ -54,7 +54,7 @@ def to_dict(self) -> dict: if self.output_summary is not None: span_data["output_summary"] = self.output_summary if self.attributes is not None: - span_data["attributes"] = self.attributes + span_data["attributes"] = flatten_attributes(self.attributes) event = { "event_id": self.event_id, diff --git a/wildedge/hubs/__init__.py b/wildedge/hubs/__init__.py index 0de302a..cb12a1e 100644 --- a/wildedge/hubs/__init__.py +++ b/wildedge/hubs/__init__.py @@ -32,16 +32,25 @@ """ from wildedge.hubs.base import BaseHubTracker -from wildedge.hubs.huggingface import HuggingFaceHubTracker from wildedge.hubs.registry import HUB_SPECS, HUBS_BY_NAME, HubSpec, supported_hubs -from wildedge.hubs.torchhub import TorchHubTracker __all__ = [ + "HUBS_BY_NAME", + "HUB_SPECS", "BaseHubTracker", + "HubSpec", "HuggingFaceHubTracker", "TorchHubTracker", - "HubSpec", - "HUB_SPECS", - "HUBS_BY_NAME", "supported_hubs", ] + + +def __getattr__(name: str): + # Tracker classes are re-exported here but imported only when accessed, + # because importing a tracker imports its hub library. + from wildedge.loading import resolve # noqa: PLC0415 + + for spec in HUB_SPECS: + if spec.tracker_class == name: + return resolve(spec.tracker) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/wildedge/hubs/base.py b/wildedge/hubs/base.py index c4a40e8..ca2ad73 100644 --- a/wildedge/hubs/base.py +++ b/wildedge/hubs/base.py @@ -2,6 +2,7 @@ from __future__ import annotations +import contextlib import os from abc import ABC, abstractmethod @@ -44,9 +45,9 @@ def install_patch(self, client_ref: object) -> None: """ Monkey-patch the hub library's download or load functions. - Called at most once per process. ``client_ref`` is a ``weakref.ref`` - to the ``WildEdge`` client; dereference it before use and guard against - ``None`` (client may have been closed or garbage-collected). + Patches are installed at most once per process. ``client_ref`` is a + callable resolving the current enabled owner, or ``None`` when inactive. + Call it before publishing; activation may repeat for a replacement owner. """ @abstractmethod @@ -77,10 +78,8 @@ def scan_cache(self) -> dict[str, int]: path = os.path.join(dirpath, filename) if os.path.islink(path): continue - try: + with contextlib.suppress(OSError): snapshot[path] = os.path.getsize(path) - except OSError: - pass except Exception: pass return snapshot diff --git a/wildedge/hubs/huggingface.py b/wildedge/hubs/huggingface.py index da91503..d531ecc 100644 --- a/wildedge/hubs/huggingface.py +++ b/wildedge/hubs/huggingface.py @@ -15,6 +15,7 @@ import time from wildedge.hubs.base import BaseHubTracker +from wildedge.instrumentation import ClientRoute from wildedge.logging import logger from wildedge.timing import elapsed_ms @@ -28,9 +29,16 @@ _local = threading.local() _hf_hub_download_patched = False _snapshot_download_patched = False +_client_route: ClientRoute | None = None def _buffer() -> list[dict]: + if _client_route is not None: + client = _client_route() + session = client.session_id if client is not None else None + if getattr(_local, "owner_session", None) != session: + _local.records = [] + _local.owner_session = session if not hasattr(_local, "records"): _local.records = [] return _local.records @@ -44,7 +52,10 @@ def _record_download( cache_hit: bool, bandwidth_bps: int | None, ) -> None: - _buffer().append( + records = _buffer() + if _client_route is not None and _client_route() is None: + return + records.append( { "repo_id": repo_id, "size": size, @@ -61,7 +72,16 @@ def _patch_symbol_references(symbol: str, original: object, replacement: object) patched_count = 0 for mod in list(sys.modules.values()): try: - if getattr(mod, symbol, None) is original: + # Read the module dict, not getattr: lazy modules (transformers) + # import submodules on attribute access. A name a lazy module has + # not loaded yet cannot hold a reference to the original. The lazy + # huggingface_hub package resolves names from its submodules on + # each access, so patching those covers `from huggingface_hub import`. + namespace = vars(mod) + except TypeError: + continue + try: + if namespace.get(symbol) is original: setattr(mod, symbol, replacement) patched_count += 1 except Exception as exc: @@ -242,16 +262,16 @@ def drain(self) -> list[dict]: records.clear() return result - def install_patch(self, client_ref: object) -> None: # noqa: ARG002 + def install_patch(self, client_ref: object) -> None: """ Patch ``hf_hub_download`` and ``snapshot_download`` for thread-local tracking. - ``client_ref`` is accepted for interface consistency but not used: the HF - hub tracker buffers records in a thread-local store and relies on the - client draining them via ``_drain_hub_trackers()`` in - ``_on_model_auto_loaded``. + Thread-local records are scoped to the enabled owner's session so a + replacement cannot drain provenance buffered for a previous project. """ - global _hf_hub_download_patched, _snapshot_download_patched + global _hf_hub_download_patched, _snapshot_download_patched, _client_route + if isinstance(client_ref, ClientRoute): + _client_route = client_ref if _hf_hub_download_patched and _snapshot_download_patched: return diff --git a/wildedge/hubs/registry.py b/wildedge/hubs/registry.py index 8c53ffc..d6ec901 100644 --- a/wildedge/hubs/registry.py +++ b/wildedge/hubs/registry.py @@ -17,11 +17,17 @@ class HubSpec: name: str required_modules: tuple[str, ...] + tracker_class: str # tracker class in the wildedge.hubs module with this name + + @property + def tracker(self) -> str: + """Import path of the tracker class, loaded only when the hub is used.""" + return f"wildedge.hubs.{self.name}:{self.tracker_class}" HUB_SPECS: tuple[HubSpec, ...] = ( - HubSpec("huggingface", ("huggingface_hub",)), - HubSpec("torchhub", ("torch",)), + HubSpec("huggingface", ("huggingface_hub",), "HuggingFaceHubTracker"), + HubSpec("torchhub", ("torch",), "TorchHubTracker"), ) HUBS_BY_NAME: dict[str, HubSpec] = {spec.name: spec for spec in HUB_SPECS} diff --git a/wildedge/instrumentation.py b/wildedge/instrumentation.py new file mode 100644 index 0000000..25dbdcb --- /dev/null +++ b/wildedge/instrumentation.py @@ -0,0 +1,108 @@ +"""Ownership of process-wide instrumentation, independent of the default client.""" + +from __future__ import annotations + +import os +import threading +import weakref +from typing import TYPE_CHECKING + +from wildedge.logging import logger + +if TYPE_CHECKING: + from wildedge.client import WildEdge + +_lock = threading.Lock() +_owner: weakref.ReferenceType[WildEdge] | None = None +_enabled: set[str] = set() +_strict_ownership = False + + +class InstrumentationOwnedError(RuntimeError): + """Another live client already owns process-wide instrumentation.""" + + +def set_strict_ownership(strict: bool) -> None: + """Raise on ownership conflicts instead of warning (`--strict-integrations`).""" + global _strict_ownership + _strict_ownership = strict + + +def strict_ownership() -> bool: + return _strict_ownership + + +class ClientRoute: + """Stable callable captured by a factory patch, resolving its current owner.""" + + def __init__(self, integration: str) -> None: + self.integration = integration + + def __call__(self) -> WildEdge | None: + with _lock: + client = _owner() if _owner is not None else None + if client is None or client.closed or self.integration not in _enabled: + return None + return client + + +class BoundClientReference: + """Existing provider clients stay with their original telemetry destination.""" + + def __init__(self, client: WildEdge) -> None: + self._ref = weakref.ref(client) + self._warned = False + + def __call__(self) -> WildEdge | None: + client = self._ref() + if client is None or client.closed: + if not self._warned: + logger.warning( + "wildedge: provider client belongs to an unavailable SDK client; " + "recreate it after initializing the replacement SDK client" + ) + self._warned = True + return None + return client + + +def bind_client_reference(reference: object, client: WildEdge) -> object: + """Snapshot a shared factory route when creating an individual provider client.""" + return ( + BoundClientReference(client) + if isinstance(reference, ClientRoute) + else reference + ) + + +def claim_instrumentation(client: WildEdge, integration: str) -> ClientRoute: + global _owner + with _lock: + if client.closed: + raise RuntimeError("Cannot instrument with a closed WildEdge client") + owner = _owner() if _owner is not None else None + if owner is not None and not owner.closed and owner is not client: + raise InstrumentationOwnedError( + "Auto-instrumentation is owned by another live WildEdge client; " + "close that client before enabling instrumentation on its replacement" + ) + if owner is not client: + _enabled.clear() + _owner = weakref.ref(client) + _enabled.add(integration) + return ClientRoute(integration) + + +def has_other_live_owner(client: WildEdge) -> bool: + with _lock: + owner = _owner() if _owner is not None else None + return owner is not None and not owner.closed and owner is not client + + +def _reset_lock_after_fork() -> None: + global _lock + _lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reset_lock_after_fork) diff --git a/wildedge/integrations/anthropic.py b/wildedge/integrations/anthropic.py index d64bb60..f21ecd3 100644 --- a/wildedge/integrations/anthropic.py +++ b/wildedge/integrations/anthropic.py @@ -2,17 +2,20 @@ from __future__ import annotations +import asyncio import functools import threading import time -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from wildedge.events.inference import ApiMeta, GenerationOutputMeta, TextInputMeta +from wildedge.instrumentation import bind_client_reference from wildedge.integrations.base import BaseExtractor from wildedge.integrations.common import ( AsyncStreamWrapper, SyncStreamWrapper, debug_failure, + stream_error_code, track_api_error, ) from wildedge.model import ModelInfo @@ -152,10 +155,20 @@ def make_anthropic_stream_callbacks( cached_tokens: list[int | None] = [None] stop_reason: list[str | None] = [None] resolved_model: list[str | None] = [None] + message_stopped = False def on_chunk(event: object) -> None: + nonlocal message_stopped event_type = type(event).__name__ - if event_type == "RawMessageStartEvent": + if ( + getattr(event, "type", None) == "message_stop" + or event_type == "RawMessageStopEvent" + ): + message_stopped = True + if ( + event_type == "RawMessageStartEvent" + or getattr(event, "type", None) == "message_start" + ): msg = getattr(event, "message", None) if msg is not None: usage = getattr(msg, "usage", None) @@ -163,7 +176,10 @@ def on_chunk(event: object) -> None: tokens_in[0] = getattr(usage, "input_tokens", None) cached_tokens[0] = getattr(usage, "cache_read_input_tokens", None) resolved_model[0] = getattr(msg, "model", None) - elif event_type == "RawMessageDeltaEvent": + elif ( + event_type == "RawMessageDeltaEvent" + or getattr(event, "type", None) == "message_delta" + ): usage = getattr(event, "usage", None) if usage is not None: tokens_out[0] = getattr(usage, "output_tokens", None) @@ -173,7 +189,11 @@ def on_chunk(event: object) -> None: if reason: stop_reason[0] = reason - def on_done(duration_ms: int, ttft_ms: int | None) -> None: + def on_done( + duration_ms: int, ttft_ms: int | None, *, outcome: str = "completed" + ) -> None: + if outcome == "completed" and not message_stopped: + outcome = "incomplete" ti, to, cr, sr = tokens_in[0], tokens_out[0], cached_tokens[0], stop_reason[0] tps = round(to / duration_ms * 1000, 1) if duration_ms > 0 and to else None api_meta = ( @@ -183,7 +203,8 @@ def on_done(duration_ms: int, ttft_ms: int | None) -> None: duration_ms=duration_ms, input_modality="text", output_modality="generation", - success=True, + success=outcome == "completed", + error_code=stream_error_code(outcome), input_meta=build_input_meta(messages, ti), output_meta=GenerationOutputMeta( task="generation", @@ -221,19 +242,18 @@ def patched_create(*args, **kwargs): handle, messages ) return SyncStreamWrapper(result, handle, t0, on_chunk, on_done) - else: - duration = elapsed_ms(t0) - usage = getattr(result, "usage", None) - tokens_in = getattr(usage, "input_tokens", None) if usage else None - handle.track_inference( - duration_ms=duration, - input_modality="text", - output_modality="generation", - success=True, - input_meta=build_input_meta(messages, tokens_in), - output_meta=build_output_meta(result, duration), - api_meta=build_api_meta(result), - ) + duration = elapsed_ms(t0) + usage = getattr(result, "usage", None) + tokens_in = getattr(usage, "input_tokens", None) if usage else None + handle.track_inference( + duration_ms=duration, + input_modality="text", + output_modality="generation", + success=True, + input_meta=build_input_meta(messages, tokens_in), + output_meta=build_output_meta(result, duration), + api_meta=build_api_meta(result), + ) return result except Exception as exc: if handle is not None: @@ -241,6 +261,7 @@ def patched_create(*args, **kwargs): raise messages_obj.create = patched_create # type: ignore[attr-defined] + wrap_native_stream(messages_obj, client_ref) def wrap_async_messages(messages_obj: object, client_ref: object) -> None: @@ -264,19 +285,18 @@ async def patched_create(*args, **kwargs): handle, messages ) return AsyncStreamWrapper(result, handle, t0, on_chunk, on_done) - else: - duration = elapsed_ms(t0) - usage = getattr(result, "usage", None) - tokens_in = getattr(usage, "input_tokens", None) if usage else None - handle.track_inference( - duration_ms=duration, - input_modality="text", - output_modality="generation", - success=True, - input_meta=build_input_meta(messages, tokens_in), - output_meta=build_output_meta(result, duration), - api_meta=build_api_meta(result), - ) + duration = elapsed_ms(t0) + usage = getattr(result, "usage", None) + tokens_in = getattr(usage, "input_tokens", None) if usage else None + handle.track_inference( + duration_ms=duration, + input_modality="text", + output_modality="generation", + success=True, + input_meta=build_input_meta(messages, tokens_in), + output_meta=build_output_meta(result, duration), + api_meta=build_api_meta(result), + ) return result except Exception as exc: if handle is not None: @@ -284,6 +304,121 @@ async def patched_create(*args, **kwargs): raise messages_obj.create = patched_create # type: ignore[attr-defined] + wrap_native_stream(messages_obj, client_ref) + + +class AnthropicStreamManager: + """Instrument the native helper's raw iterator while preserving its public API.""" + + def __init__(self, manager: Any, handle: ModelHandle, messages: list) -> None: + self._manager = manager + self._handle = handle + self._messages = messages + self._tracked = None + + def _safe(self, callback, *args, **kwargs) -> None: + try: + callback(*args, **kwargs) + except Exception as exc: + debug_anthropic_failure("native stream telemetry", exc) + + def _wrap(self, stream: Any, t0: float, asynchronous: bool) -> Any: + # Both old and current SDKs route event/text/final-message helpers + # through _iterator. Preserve the SDK parser and returned stream object. + try: + callbacks = make_anthropic_stream_callbacks(self._handle, self._messages) + wrapper = AsyncStreamWrapper if asynchronous else SyncStreamWrapper + tracked = wrapper(stream._iterator, self._handle, t0, *callbacks) + original_close = stream.close + if asynchronous: + + @functools.wraps(original_close) + async def close(): + try: + return await original_close() + finally: + tracked._finish("closed") + else: + + @functools.wraps(original_close) + def close(): + try: + return original_close() + finally: + tracked._finish("closed") + + stream._iterator = tracked + stream.close = close + self._tracked = tracked + except Exception as exc: + debug_anthropic_failure("native stream wrap", exc) + return stream + + def __enter__(self) -> Any: + t0 = time.perf_counter() + try: + stream = self._manager.__enter__() + except Exception as exc: + self._safe(track_api_error, self._handle, exc, elapsed_ms(t0)) + raise + return self._wrap(stream, t0, False) + + def __exit__(self, *args: Any) -> Any: + return self._manager.__exit__(*args) + + async def __aenter__(self) -> Any: + t0 = time.perf_counter() + try: + stream = await self._manager.__aenter__() + except asyncio.CancelledError: + self._safe( + self._handle.track_inference, + duration_ms=elapsed_ms(t0), + success=False, + error_code="STREAM_CANCELLED", + ) + raise + except Exception as exc: + self._safe(track_api_error, self._handle, exc, elapsed_ms(t0)) + raise + return self._wrap(stream, t0, True) + + async def __aexit__(self, *args: Any) -> Any: + if ( + self._tracked is not None + and args + and args[0] is not None + and issubclass(args[0], asyncio.CancelledError) + ): + self._tracked._finish("cancelled") + return await self._manager.__aexit__(*args) + + def __getattr__(self, name: str) -> Any: + return getattr(self._manager, name) + + +def wrap_native_stream(messages_obj: object, client_ref: object) -> None: + original_stream = getattr(messages_obj, "stream", None) + if original_stream is None: + return + model_handles: dict[str, ModelHandle] = {} + + @functools.wraps(original_stream) + def patched_stream(*args, **kwargs): + client = client_ref() + model_id = kwargs.get("model") + if client is None or client.closed or not model_id: + return original_stream(*args, **kwargs) + handle = resolve_handle(model_id, messages_obj, model_handles, client) + manager = original_stream(*args, **kwargs) + if handle is None: + return manager + messages = kwargs.get("messages", []) + return AnthropicStreamManager( + manager, handle, messages if isinstance(messages, list) else [] + ) + + messages_obj.stream = patched_stream class AnthropicExtractor(BaseExtractor): @@ -337,8 +472,9 @@ def patched_sync_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-d original_sync_init(self_inner, *args, **kwargs) c = client_ref() # type: ignore[call-arg] if c is not None and not c.closed: + bound_ref = bind_client_reference(client_ref, c) try: - wrap_sync_messages(self_inner.messages, client_ref) + wrap_sync_messages(self_inner.messages, bound_ref) except Exception as exc: debug_anthropic_failure("sync client wrap", exc) @@ -346,8 +482,9 @@ def patched_async_init(self_inner, *args, **kwargs): # type: ignore[no-untyped- original_async_init(self_inner, *args, **kwargs) c = client_ref() # type: ignore[call-arg] if c is not None and not c.closed: + bound_ref = bind_client_reference(client_ref, c) try: - wrap_async_messages(self_inner.messages, client_ref) + wrap_async_messages(self_inner.messages, bound_ref) except Exception as exc: debug_anthropic_failure("async client wrap", exc) @@ -359,3 +496,7 @@ def patched_async_init(self_inner, *args, **kwargs): # type: ignore[no-untyped- _anthropic.Anthropic.__init__ = patched_sync_init _anthropic.AsyncAnthropic.__init__ = patched_async_init _anthropic_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = AnthropicExtractor.install_auto_load_patch diff --git a/wildedge/integrations/base.py b/wildedge/integrations/base.py index b6e2c0e..e525bf9 100644 --- a/wildedge/integrations/base.py +++ b/wildedge/integrations/base.py @@ -20,7 +20,17 @@ def patch_instance_call_once( *, patch_name: str, make_patched_call: Callable[[Callable[..., Any]], Callable[..., Any]], + method: str = "__call__", + extra_methods: dict[str, Callable[[Callable[..., Any]], Callable[..., Any]]] + | None = None, ) -> None: + """Swap obj's class for a cached subclass whose ``method`` is wrapped. + + ``method`` defaults to ``__call__``; pass another name when the runtime's + entry points all funnel through a different method. ``extra_methods`` maps + more method names to wrapper factories, patched in the same subclass so a + repeat call with the same ``patch_name`` stays a no-op. + """ current_cls = type(obj) if getattr(current_cls, "__wildedge_patch_name__", None) == patch_name: return @@ -37,18 +47,19 @@ def patch_instance_call_once( patched_cls = per_class.get(patch_name) if patched_cls is None: - original_call = current_cls.__call__ + original_call = getattr(current_cls, method) patched_call = make_patched_call(original_call) + namespace: dict[str, Any] = { + method: patched_call, + "__wildedge_patch_name__": patch_name, + "__wildedge_original_call__": original_call, + } + for extra_name, make_extra in (extra_methods or {}).items(): + original_extra = getattr(current_cls, extra_name, None) + if original_extra is not None: + namespace[extra_name] = make_extra(original_extra) class_name = f"{current_cls.__name__}WildEdgePatched{patch_name.title()}" - patched_cls = type( - class_name, - (current_cls,), - { - "__call__": patched_call, - "__wildedge_patch_name__": patch_name, - "__wildedge_original_call__": original_call, - }, - ) + patched_cls = type(class_name, (current_cls,), namespace) per_class[patch_name] = patched_cls obj.__class__ = patched_cls # type: ignore[assignment] diff --git a/wildedge/integrations/common.py b/wildedge/integrations/common.py index 7442d43..86ffb01 100644 --- a/wildedge/integrations/common.py +++ b/wildedge/integrations/common.py @@ -2,14 +2,18 @@ from __future__ import annotations +import asyncio +import contextvars +import itertools from collections.abc import Callable from typing import TYPE_CHECKING, Any from urllib.parse import urlparse -from wildedge import constants +from wildedge import accelerators, constants, quantizations from wildedge.events.error import ErrorCode from wildedge.events.inference import TextInputMeta from wildedge.logging import logger +from wildedge.platforms import CURRENT_PLATFORM from wildedge.timing import elapsed_ms if TYPE_CHECKING: @@ -20,6 +24,69 @@ def debug_failure(framework: str, context: str, exc: BaseException) -> None: logger.debug("wildedge: %s %s failed: %s", framework, context, exc) +# Folder prefix of a repo in the Hugging Face hub cache: models--org--name. +HF_CACHE_REPO_PREFIX = "models--" + + +def hf_repo_from_cache_path(path: str) -> str | None: + """``org/name`` for a path inside the Hugging Face hub cache, else None.""" + for part in path.replace("\\", "/").split("/"): + if part.startswith(HF_CACHE_REPO_PREFIX): + return part[len(HF_CACHE_REPO_PREFIX) :].replace("--", "/", 1) + return None + + +def torch_module_accelerator(module: object) -> str | None: + """Device type of a torch module's first parameter, such as cpu, cuda, + mps or xpu. None when the module has no parameters.""" + try: + first = next(module.parameters()) # type: ignore[attr-defined] + except StopIteration: + return None + device_type = getattr(first.device, "type", None) + return str(device_type) if device_type else None + + +def tf_weights_accelerator(model: object) -> str: + """The platform's GPU when a TensorFlow/Keras model's first weight is on a + GPU, else CPU. TF names devices like ``/job:localhost/.../device:GPU:0``.""" + weights = getattr(model, "weights", None) or [] + if weights: + device = str(getattr(weights[0], "device", "") or "").upper() + if "GPU" in device or "CUDA" in device: + return CURRENT_PLATFORM.gpu_accelerator_for_offload() + return accelerators.CPU + + +def mlx_accelerator() -> str: + """MLX runs on the Apple GPU unless the app switched its default to CPU.""" + try: + import mlx.core as mx # noqa: PLC0415 -- optional runtime + + return accelerators.CPU if mx.default_device() == mx.cpu else accelerators.MPS + except Exception as exc: + debug_failure("mlx", "device detection", exc) + return accelerators.MPS + + +def mlx_layer_quantization(model: object) -> str | None: + """Quantization of an MLX model's layers: the bit width of the first + quantized layer, ``quantized`` when the width is unknown, or None when no + layer is quantized.""" + for _, layer in model.named_modules(): # type: ignore[attr-defined] + bits = getattr(layer, "bits", None) + if isinstance(bits, int): + return quantizations.int_bits(bits) + if "quantized" in type(layer).__name__.lower(): + return quantizations.QUANTIZED + return None + + +def stream_error_code(outcome: str) -> str | None: + """None for a stream that completed, else ``STREAM_``.""" + return None if outcome == "completed" else "STREAM_" + outcome.upper() + + # --------------------------------------------------------------------------- # Chat API helpers, shared by the openai integration and wildedge.llm_api # --------------------------------------------------------------------------- @@ -129,7 +196,7 @@ def track_api_error( error_code = no_response_error_code(exc) or fallback_code handle.track_error( error_code=error_code, - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), http_status=http_status, provider_error_code=provider_error_code, duration_ms=duration_ms, @@ -161,14 +228,17 @@ def build_input_meta(messages: list, tokens_in: int | None) -> TextInputMeta | N ) +# Substrings of framework dtype names (torch.float16, tensor(float16), +# mlx.core.bfloat16, torch.qint8). Checked in order, so bfloat16 must come +# before float16, which it contains. DTYPE_QUANTIZATION_MAP: dict[str, str] = { - "bfloat16": "bf16", - "float16": "f16", - "float32": "f32", - "int8": "int8", - "qint": "int8", - "quint": "int8", - "uint8": "int8", + "bfloat16": quantizations.BF16, + "float16": quantizations.F16, + "float32": quantizations.F32, + "int8": quantizations.INT8, + "qint": quantizations.INT8, + "quint": quantizations.INT8, + "uint8": quantizations.INT8, } @@ -191,7 +261,7 @@ def image_brightness_histogram(norm_flat: Any) -> tuple[float, float, list[int]] edges = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] buckets = [ int(((norm_flat >= lo) & (norm_flat < hi)).sum()) - for lo, hi in zip(edges, edges[1:]) + for lo, hi in itertools.pairwise(edges) ] buckets[-1] += int((norm_flat == 1.0).sum()) @@ -267,105 +337,241 @@ def num_classes_from_output_shape(shape: tuple) -> int: # --------------------------------------------------------------------------- # Each integration provides: # on_chunk(chunk) -> None : update mutable state from a single chunk -# on_done(duration_ms, ttft_ms) : record inference once the stream is exhausted +# on_done(duration_ms, ttft_ms, *, outcome="completed") : record one outcome # # The wrappers handle TTFT capture, error tracking, context-manager delegation, # and attribute proxying so callers get a drop-in replacement for the raw stream. -class SyncStreamWrapper: - """Wraps a sync iterable stream to capture TTFT and record inference on exhaustion.""" +def stream_chunk_has_output(chunk: object) -> bool: + """Ignore role, usage, and lifecycle chunks when measuring first output.""" + + def field(value: object, name: str) -> Any: + return ( + value.get(name) if isinstance(value, dict) else getattr(value, name, None) + ) + + choices = field(chunk, "choices") + if choices is not None: + for choice in choices: + delta = field(choice, "delta") + if field(choice, "text") or any( + field(delta, key) + for key in ("content", "refusal", "tool_calls", "function_call") + ): + return True + return False + event_type = field(chunk, "type") + if event_type is not None: + if event_type == "content_block_delta": + delta = field(chunk, "delta") + return any( + field(delta, key) + for key in ("text", "thinking", "partial_json", "data") + ) + return event_type in ( + "response.output_text.delta", + "response.refusal.delta", + "response.function_call_arguments.delta", + "response.output_audio.delta", + ) and bool(field(chunk, "delta")) + return True + + +class StreamState: + """One lifecycle per stream, with telemetry isolated from application code.""" def __init__( self, - original: object, handle: ModelHandle, t0: float, on_chunk: Callable[[object], None] | None, - on_done: Callable[[int, int | None], None], + on_done: Callable[..., None], ) -> None: - self._original = original self._handle = handle self._t0 = t0 self._on_chunk = on_chunk self._on_done = on_done + self._finished = False + self._ttft_ms = None + self._context = contextvars.copy_context() - def __iter__(self): - return self._track() + @property + def finished(self) -> bool: + return self._finished - def _track(self): - ttft_ms: int | None = None + def _safe(self, callback: Callable, *args: Any, **kwargs: Any) -> Any: try: - for chunk in self._original: # type: ignore[union-attr] - if ttft_ms is None: - ttft_ms = elapsed_ms(self._t0) - if self._on_chunk is not None: - self._on_chunk(chunk) - yield chunk + return self._context.run(callback, *args, **kwargs) except Exception as exc: - track_api_error(self._handle, exc, elapsed_ms(self._t0)) - raise + debug_failure("stream", "telemetry callback", exc) + return None + + def _chunk(self, chunk: object) -> None: + if self._ttft_ms is None and self._safe(stream_chunk_has_output, chunk): + self._ttft_ms = elapsed_ms(self._t0) + if self._on_chunk is not None: + self._safe(self._on_chunk, chunk) + + def _finish(self, outcome: str = "completed") -> None: + if self._finished: + return + self._finished = True + if outcome == "completed": + self._safe(self._on_done, elapsed_ms(self._t0), self._ttft_ms) else: - self._on_done(elapsed_ms(self._t0), ttft_ms) + self._safe( + self._on_done, elapsed_ms(self._t0), self._ttft_ms, outcome=outcome + ) + + def _error(self, exc: Exception) -> None: + if not self._finished: + self._finished = True + self._safe(track_api_error, self._handle, exc, elapsed_ms(self._t0)) + + +class SyncStreamWrapper(StreamState): + """A single-pass iterator tracking exhaustion, failure, and explicit closure.""" + + def __init__( + self, + original: Any, + handle: ModelHandle, + t0: float, + on_chunk: Callable[[object], None] | None, + on_done: Callable[..., None], + ) -> None: + super().__init__(handle, t0, on_chunk, on_done) + self._original = original + self._iterator = None + + @property + def started(self) -> bool: + """Whether iteration has begun.""" + return self._iterator is not None + + def __iter__(self) -> SyncStreamWrapper: + return self + + def __next__(self) -> Any: + if self._finished: + raise StopIteration + try: + if self._iterator is None: + self._iterator = iter(self._original) + chunk = next(self._iterator) + except StopIteration: + self._finish() + raise + except Exception as exc: + self._error(exc) + raise + self._chunk(chunk) + return chunk + + def close(self) -> Any: + try: + close = getattr(self._original, "close", None) + if close is not None: + return close() + return None + finally: + self._finish("closed") def __enter__(self) -> SyncStreamWrapper: if hasattr(self._original, "__enter__"): - self._original.__enter__() # type: ignore[union-attr] + self._original.__enter__() return self - def __exit__(self, *args: object) -> object: - if hasattr(self._original, "__exit__"): - return self._original.__exit__(*args) # type: ignore[union-attr] - return None - - def __getattr__(self, name: str) -> object: + def __exit__(self, *args: Any) -> Any: + try: + if hasattr(self._original, "__exit__"): + return self._original.__exit__(*args) + close = getattr(self._original, "close", None) + if close is not None: + close() + return None + finally: + self._finish("closed") + + def __getattr__(self, name: str) -> Any: return getattr(self._original, name) -class AsyncStreamWrapper: - """Wraps an async iterable stream to capture TTFT and record inference on exhaustion.""" +class AsyncStreamWrapper(StreamState): + """Async counterpart, including cancellation while waiting for a chunk.""" def __init__( self, - original: object, + original: Any, handle: ModelHandle, t0: float, on_chunk: Callable[[object], None] | None, - on_done: Callable[[int, int | None], None], + on_done: Callable[..., None], ) -> None: + super().__init__(handle, t0, on_chunk, on_done) self._original = original - self._handle = handle - self._t0 = t0 - self._on_chunk = on_chunk - self._on_done = on_done + self._iterator = None - def __aiter__(self): - return self._track() + def __aiter__(self) -> AsyncStreamWrapper: + return self - async def _track(self): - ttft_ms: int | None = None + async def __anext__(self) -> Any: + if self._finished: + raise StopAsyncIteration try: - async for chunk in self._original: # type: ignore[union-attr] - if ttft_ms is None: - ttft_ms = elapsed_ms(self._t0) - if self._on_chunk is not None: - self._on_chunk(chunk) - yield chunk + if self._iterator is None: + self._iterator = self._original.__aiter__() + chunk = await self._iterator.__anext__() + except StopAsyncIteration: + self._finish() + raise + except asyncio.CancelledError: + self._finish("cancelled") + raise except Exception as exc: - track_api_error(self._handle, exc, elapsed_ms(self._t0)) + self._error(exc) raise - else: - self._on_done(elapsed_ms(self._t0), ttft_ms) + self._chunk(chunk) + return chunk + + async def close(self) -> Any: + try: + close = getattr(self._original, "aclose", None) or getattr( + self._original, "close", None + ) + if close is not None: + return await close() + return None + finally: + self._finish("closed") + + async def aclose(self) -> Any: + return await self.close() async def __aenter__(self) -> AsyncStreamWrapper: if hasattr(self._original, "__aenter__"): - await self._original.__aenter__() # type: ignore[union-attr] + await self._original.__aenter__() return self - async def __aexit__(self, *args: object) -> object: - if hasattr(self._original, "__aexit__"): - return await self._original.__aexit__(*args) # type: ignore[union-attr] - return None - - def __getattr__(self, name: str) -> object: + async def __aexit__(self, *args: Any) -> Any: + try: + if hasattr(self._original, "__aexit__"): + return await self._original.__aexit__(*args) + close = getattr(self._original, "aclose", None) or getattr( + self._original, "close", None + ) + if close is not None: + await close() + return None + finally: + self._finish( + "cancelled" + if args + and args[0] is not None + and issubclass(args[0], asyncio.CancelledError) + else "closed" + ) + + def __getattr__(self, name: str) -> Any: return getattr(self._original, name) diff --git a/wildedge/integrations/faster_whisper.py b/wildedge/integrations/faster_whisper.py new file mode 100644 index 0000000..4bc2c31 --- /dev/null +++ b/wildedge/integrations/faster_whisper.py @@ -0,0 +1,192 @@ +"""Opt-in instrumentation for faster-whisper's lazy transcription API.""" + +from __future__ import annotations + +import functools +import inspect +import threading +import time + +from wildedge import accelerators, quantizations +from wildedge.instrumentation import bind_client_reference +from wildedge.integrations.common import ( + SyncStreamWrapper, + debug_failure, + track_api_error, +) +from wildedge.integrations.whisper import ( + CLIENT_ATTR, + DESCRIPTOR_ATTR, + FASTER_WHISPER, + HANDLE_ATTR, + WhisperDescriptor, + report_load, + track_transcription, +) +from wildedge.timing import elapsed_ms + +_PATCH_LOCK = threading.Lock() + +# WhisperModel.__init__'s first parameter: a size name, hub repo or local path. +MODEL_PATH_PARAM = "model_size_or_path" + +# CTranslate2 compute types, labelled by weight precision. The activation type +# after "int8_" isn't reported. Unknown types pass through as they are. +COMPUTE_TYPE_QUANTIZATION = { + "float32": quantizations.F32, + "float16": quantizations.F16, + "bfloat16": quantizations.BF16, + "int8": quantizations.INT8, + "int8_float32": quantizations.INT8, + "int8_float16": quantizations.INT8, + "int8_bfloat16": quantizations.INT8, + "int16": quantizations.int_bits(16), +} + +# Outcome of a segment loop that stopped without reading to the end. +OUTCOME_ABANDONED = "abandoned" + + +class TranscriptCollector: + """Counts words and timing without retaining transcript text.""" + + def __init__(self, t0: float) -> None: + self._t0 = t0 + self.words = 0 + self.last_segment_ms: int | None = None + + def add(self, segment) -> None: + self.last_segment_ms = elapsed_ms(self._t0) + self.words += len(segment.text.split()) + + +class SegmentStream(SyncStreamWrapper): + """Segment iterator that also reports a loop that stopped early. + + ``for segment in segments: break`` drops the generator without closing + it. The event is then sent when the iterator is garbage collected, as + ``STREAM_ABANDONED`` with the time up to the last segment received. A + stream that was never iterated, for example when only ``info.language`` + was needed, decoded nothing and sends no event. + """ + + # Class-level defaults keep __del__ safe if __init__ never finished; + # otherwise these lookups would fall through to the wrapped generator. + _finished = True + _iterator = None + + def __del__(self) -> None: + if self.started and not self.finished: + self._finish(OUTCOME_ABANDONED) + + +def tracked_handle(model): + """The model's handle while its SDK client is still open, else None.""" + client_ref = getattr(model, CLIENT_ATTR, None) + if client_ref is None or client_ref() is None: + return None + return getattr(model, HANDLE_ATTR, None) + + +def wrap_transcribe(original, model_of=lambda owner: owner): + """Wrap ``WhisperModel.transcribe`` or ``BatchedInferencePipeline.transcribe``. + + ``model_of`` finds the ``WhisperModel`` behind the object being called. + """ + + @functools.wraps(original) + def transcribe(self, audio, *args, **kwargs): + handle = tracked_handle(model_of(self)) + if handle is None: + return original(self, audio, *args, **kwargs) + t0 = time.perf_counter() + try: + segments, info = original(self, audio, *args, **kwargs) + except Exception as exc: + try: + track_api_error(handle, exc, elapsed_ms(t0)) + except Exception as telemetry_exc: + debug_failure(FASTER_WHISPER, "error tracking", telemetry_exc) + raise + collector = TranscriptCollector(t0) + + def done(duration_ms, first_segment_ms, *, outcome="completed"): + if outcome == OUTCOME_ABANDONED and collector.last_segment_ms is not None: + duration_ms = collector.last_segment_ms + track_transcription( + handle, + audio, + "", + duration_ms=duration_ms, + language=getattr(info, "language", None), + language_probability=getattr(info, "language_probability", None), + audio_duration_s=getattr(info, "duration", None), + word_count=collector.words, + first_segment_ms=first_segment_ms, + outcome=outcome, + ) + + return SegmentStream(segments, handle, t0, collector.add, done), info + + return transcribe + + +def install_auto_load_patch(client_ref): + with _PATCH_LOCK: + _install(client_ref) + + +def _install(client_ref): + import faster_whisper # noqa: PLC0415 -- optional runtime + + cls = faster_whisper.WhisperModel + original = cls.__init__ + if getattr(original, "__wildedge_patch_name__", None) == FASTER_WHISPER: + return + signature = inspect.signature(original) + + @functools.wraps(original) + def init(self, *args, **kwargs): + client = client_ref() + t0 = time.perf_counter() + original(self, *args, **kwargs) + if client is None or client.closed: + return + try: + bound = signature.bind(self, *args, **kwargs) + bound.apply_defaults() + # The ctranslate2.models.Whisper instance. Its device is the + # resolved one ("cpu" or "cuda"), already an accelerator name. + runtime = self.model + device = getattr(runtime, "device", None) or accelerators.CPU + compute = getattr( + runtime, "compute_type", bound.arguments.get("compute_type") + ) + setattr( + self, + DESCRIPTOR_ATTR, + WhisperDescriptor.from_path( + bound.arguments[MODEL_PATH_PARAM], + FASTER_WHISPER, + device, + COMPUTE_TYPE_QUANTIZATION.get(compute, compute), + ), + ) + report_load(client, self, elapsed_ms(t0)) + # Already registered by the load above; this returns its handle. + setattr(self, HANDLE_ATTR, client.register_model(self)) + # Keep this model with the client that loaded it, as provider + # clients do. After a client replacement it warns once instead. + setattr(self, CLIENT_ATTR, bind_client_reference(client_ref, client)) + except Exception as exc: + debug_failure(FASTER_WHISPER, "model registration", exc) + + init.__wildedge_patch_name__ = FASTER_WHISPER + cls.__init__ = init + cls.transcribe = wrap_transcribe(cls.transcribe) + batched = getattr(faster_whisper, "BatchedInferencePipeline", None) + if batched is not None: + batched.transcribe = wrap_transcribe( + batched.transcribe, + model_of=lambda pipeline: getattr(pipeline, "model", None), + ) diff --git a/wildedge/integrations/gguf.py b/wildedge/integrations/gguf.py index dc44360..7e4ead8 100644 --- a/wildedge/integrations/gguf.py +++ b/wildedge/integrations/gguf.py @@ -2,17 +2,28 @@ from __future__ import annotations +import contextvars import functools +import inspect import re import threading import time from pathlib import Path from typing import TYPE_CHECKING -from wildedge import constants -from wildedge.events.inference import GenerationOutputMeta, TextInputMeta +from wildedge import accelerators, quantizations +from wildedge.events.inference import ( + GenerationConfig, + GenerationOutputMeta, + TextInputMeta, +) from wildedge.integrations.base import BaseExtractor, patch_instance_call_once -from wildedge.integrations.common import SyncStreamWrapper, debug_failure +from wildedge.integrations.common import ( + SyncStreamWrapper, + build_input_meta, + debug_failure, + stream_error_code, +) from wildedge.logging import logger from wildedge.model import ModelInfo from wildedge.platforms import CURRENT_PLATFORM @@ -34,12 +45,16 @@ _llama_patched = False _LLAMA_PATCH_LOCK = threading.Lock() -GGUF_CALL_PATCH_NAME = "gguf_call" +# llama-cpp-python's __call__ and create_chat_completion (via the chat handler) +# both run create_completion, so patching it tracks every entry point once. +GGUF_COMPLETION_PATCH_NAME = "gguf_create_completion" +GGUF_PATCHED_METHOD = "create_completion" +GGUF_CHAT_METHOD = "create_chat_completion" GGUF_HANDLE_ATTR = "__wildedge_gguf_handle__" LLAMA_AUTO_LOAD_PATCH_NAME = "gguf_auto_load" -debug_gguf_failure = functools.partial(debug_failure, "gguf") +debug_gguf_failure = functools.partial(debug_failure, "llama_cpp") def n_gpu_layers(obj: object) -> int: @@ -59,17 +74,87 @@ def detect_accelerator(obj: object) -> str: return CURRENT_PLATFORM.gpu_accelerator_for_offload() except Exception as exc: debug_gguf_failure("accelerator detection", exc) - return "cpu" + return accelerators.CPU + + +# Filename suffixes that have a shared label under another name. +_SUFFIX_LABELS = {"i8": quantizations.INT8} + +# llama.cpp file types (gguf-py's LlamaFileType) as stored in the model's +# general.file_type metadata. Removed and unknown types are left out. +GGUF_FILE_TYPES = { + 0: quantizations.F32, + 1: quantizations.F16, + 2: "q4_0", + 3: "q4_1", + 7: "q8_0", + 8: "q5_0", + 9: "q5_1", + 10: "q2_k", + 11: "q3_k_s", + 12: "q3_k_m", + 13: "q3_k_l", + 14: "q4_k_s", + 15: "q4_k_m", + 16: "q5_k_s", + 17: "q5_k_m", + 18: "q6_k", + 19: "iq2_xxs", + 20: "iq2_xs", + 21: "q2_k_s", + 22: "iq3_xs", + 23: "iq3_xxs", + 24: "iq1_s", + 25: "iq4_nl", + 26: "iq3_s", + 27: "iq3_m", + 28: "iq2_s", + 29: "iq2_m", + 30: "iq4_xs", + 31: "iq1_m", + 32: quantizations.BF16, + 36: "tq1_0", + 37: "tq2_0", + 38: "mxfp4_moe", + 39: "nvfp4", + 40: "q1_0", +} def parse_quantization(filename: str) -> str | None: match = _QUANT_RE.search(filename) if match: - return match.group(1).lower() + suffix = match.group(1).lower() + return _SUFFIX_LABELS.get(suffix, suffix) return None -def make_gguf_input_meta(prompt: object, tokens_in: int | None) -> TextInputMeta | None: +def file_type_quantization(file_type: object) -> str | None: + """Label for a general.file_type metadata value, such as "15" (q4_k_m).""" + try: + return GGUF_FILE_TYPES.get(int(str(file_type))) + except ValueError: + return None + + +# Messages of the create_chat_completion call in progress, so the +# create_completion it runs can record the event as a chat turn. +_chat_messages: contextvars.ContextVar[list | None] = contextvars.ContextVar( + "wildedge_gguf_chat_messages", default=None +) + + +def make_gguf_input_meta( + prompt: object, tokens_in: int | None, messages: list | None = None +) -> TextInputMeta | None: + if messages is not None: + meta = build_input_meta(messages, tokens_in) + if meta is not None: + user_turns = sum( + 1 for m in messages if isinstance(m, dict) and m.get("role") == "user" + ) + meta.turn_index = max(user_turns - 1, 0) + return meta if not isinstance(prompt, str) or not prompt: return None return TextInputMeta( @@ -86,7 +171,7 @@ def make_gguf_output_meta( ttft_ms: int | None, duration_ms: int, ) -> GenerationOutputMeta | None: - if tokens_out is None and ttft_ms is None: + if tokens_out is None and ttft_ms is None and stop_reason is None: return None tps = ( round(tokens_out / duration_ms * 1000, 1) @@ -103,7 +188,34 @@ def make_gguf_output_meta( ) -def make_gguf_stream_callbacks(handle: object, prompt: object) -> tuple: +def make_gguf_generation_config(kwargs: dict) -> GenerationConfig | None: + stop = kwargs.get("stop") + if isinstance(stop, str): + stop_count: int | None = 1 + elif isinstance(stop, (list, tuple)): + stop_count = len(stop) or None + else: + stop_count = None + config = GenerationConfig( + temperature=kwargs.get("temperature"), + top_p=kwargs.get("top_p"), + top_k=kwargs.get("top_k"), + max_tokens=kwargs.get("max_tokens"), + repetition_penalty=kwargs.get("repeat_penalty"), + frequency_penalty=kwargs.get("frequency_penalty"), + presence_penalty=kwargs.get("presence_penalty"), + seed=kwargs.get("seed"), + stop_sequences_count=stop_count, + ) + return config if config.to_dict() else None + + +def make_gguf_stream_callbacks( + handle: object, + prompt: object, + messages: list | None = None, + generation_config: GenerationConfig | None = None, +) -> tuple: """Return (on_chunk, on_done) callbacks for a llama-cpp-python streaming response. Chunks are dicts; usage appears in the final chunk when available. @@ -120,70 +232,133 @@ def on_chunk(chunk: object) -> None: tokens_in[0] = usage.get("prompt_tokens") tokens_out[0] = usage.get("completion_tokens") choices = chunk.get("choices") or [] - if choices: + if choices and isinstance(choices[0], dict): reason = choices[0].get("finish_reason") if reason: stop_reason[0] = reason - def on_done(duration_ms: int, ttft_ms: int | None) -> None: + def on_done( + duration_ms: int, ttft_ms: int | None, *, outcome: str = "completed" + ) -> None: ti, to, sr = tokens_in[0], tokens_out[0], stop_reason[0] handle.track_inference( # type: ignore[union-attr] duration_ms=duration_ms, input_modality="text", output_modality="generation", - input_meta=make_gguf_input_meta(prompt, ti), - success=True, + input_meta=make_gguf_input_meta(prompt, ti, messages), + success=outcome == "completed", + error_code=stream_error_code(outcome), output_meta=make_gguf_output_meta(ti, to, sr, ttft_ms, duration_ms), + generation_config=generation_config, ) return on_chunk, on_done +def call_arguments(signature, self_inner, args, kwargs) -> dict: + """Arguments of the call with the method's defaults filled in. + + The chat handler and __call__ pass every sampling argument explicitly; this + gives a direct create_completion call the same generation config. + """ + if signature is not None: + try: + bound = signature.bind(self_inner, *args, **kwargs) + bound.apply_defaults() + arguments = dict(bound.arguments) + arguments.update(arguments.pop("kwargs", None) or {}) + return arguments + except Exception as exc: + debug_gguf_failure("argument binding", exc) + arguments = dict(kwargs) + if args: + arguments.setdefault("prompt", args[0]) + return arguments + + def build_patched_call(original_call): + try: + signature = inspect.signature(original_call) + except (TypeError, ValueError): + signature = None + def patched_call(self_inner, *args, **kwargs): handle = getattr(self_inner, GGUF_HANDLE_ATTR, None) if handle is None: return original_call(self_inner, *args, **kwargs) - prompt = args[0] if args else kwargs.get("prompt", "") - is_streaming: bool = bool(kwargs.get("stream", False)) + arguments = call_arguments(signature, self_inner, args, kwargs) + prompt = arguments.get("prompt", "") + messages = _chat_messages.get() + is_streaming: bool = bool(arguments.get("stream", False)) + generation_config = None + try: + generation_config = make_gguf_generation_config(arguments) + except Exception as exc: + debug_gguf_failure("generation config extraction", exc) t0 = time.perf_counter() try: result = original_call(self_inner, *args, **kwargs) if is_streaming: - on_chunk, on_done = make_gguf_stream_callbacks(handle, prompt) + on_chunk, on_done = make_gguf_stream_callbacks( + handle, prompt, messages, generation_config + ) return SyncStreamWrapper(result, handle, t0, on_chunk, on_done) duration_ms = elapsed_ms(t0) tokens_in = None tokens_out = None + stop_reason = None try: if isinstance(result, dict): - usage = result.get("usage", {}) + usage = result.get("usage") or {} tokens_in = usage.get("prompt_tokens") tokens_out = usage.get("completion_tokens") + choices = result.get("choices") or [] + if choices and isinstance(choices[0], dict): + stop_reason = choices[0].get("finish_reason") except Exception as exc: debug_gguf_failure("usage extraction", exc) handle.track_inference( duration_ms=duration_ms, input_modality="text", output_modality="generation", - input_meta=make_gguf_input_meta(prompt, tokens_in), + input_meta=make_gguf_input_meta(prompt, tokens_in, messages), success=True, output_meta=make_gguf_output_meta( - tokens_in, tokens_out, None, None, duration_ms + tokens_in, tokens_out, stop_reason, None, duration_ms ), + generation_config=generation_config, ) return result except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise return patched_call +def build_patched_chat(original_chat): + """Expose the chat messages to the create_completion call the handler makes. + + Tracking itself stays in create_completion, so each call is counted once. + """ + + def patched_chat(self_inner, *args, **kwargs): + if getattr(self_inner, GGUF_HANDLE_ATTR, None) is None: + return original_chat(self_inner, *args, **kwargs) + messages = args[0] if args else kwargs.get("messages") + token = _chat_messages.set(list(messages) if messages else None) + try: + return original_chat(self_inner, *args, **kwargs) + finally: + _chat_messages.reset(token) + + return patched_chat + + class GgufExtractor(BaseExtractor): # llama-cpp-python's Llama class is the universal GGUF loader and covers Mistral, # Phi, Gemma, etc. Model family is read from general.architecture in GGUF metadata. @@ -206,8 +381,8 @@ def extract_info( try: metadata = getattr(obj, "metadata", {}) or {} file_type = metadata.get("general.file_type") - if file_type: - quantization = str(file_type).lower() + if file_type is not None: + quantization = file_type_quantization(file_type) except Exception as exc: debug_gguf_failure("quantization metadata read", exc) if quantization is None: @@ -269,8 +444,10 @@ def install_hooks(self, obj: object, handle: ModelHandle) -> None: setattr(obj, GGUF_HANDLE_ATTR, handle) patch_instance_call_once( obj, - patch_name=GGUF_CALL_PATCH_NAME, + patch_name=GGUF_COMPLETION_PATCH_NAME, make_patched_call=build_patched_call, + method=GGUF_PATCHED_METHOD, + extra_methods={GGUF_CHAT_METHOD: build_patched_chat}, ) @classmethod @@ -322,3 +499,7 @@ def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] patched_init.__wildedge_original_call__ = original_init # type: ignore[attr-defined] _llama_cpp.Llama.__init__ = patched_init _llama_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = GgufExtractor.install_auto_load_patch diff --git a/wildedge/integrations/keras.py b/wildedge/integrations/keras.py index 51b55d5..f1f101b 100644 --- a/wildedge/integrations/keras.py +++ b/wildedge/integrations/keras.py @@ -11,16 +11,16 @@ except ImportError: _np = None # type: ignore[assignment] -from wildedge import constants +from wildedge import accelerators from wildedge.events.inference import ClassificationOutputMeta, TopKPrediction from wildedge.integrations.base import BaseExtractor, patch_instance_call_once from wildedge.integrations.common import ( debug_failure, infer_input_modality_from_layer_types, num_classes_from_output_shape, + tf_weights_accelerator, ) from wildedge.model import ModelInfo -from wildedge.platforms import CURRENT_PLATFORM from wildedge.timing import elapsed_ms if TYPE_CHECKING: @@ -99,7 +99,7 @@ def patched_call(self_inner, *args, **kwargs): except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise @@ -108,14 +108,10 @@ def patched_call(self_inner, *args, **kwargs): def detect_accelerator(obj: object) -> str: try: - weights = getattr(obj, "weights", None) or [] - if weights: - device = str(getattr(weights[0], "device", "") or "") - if any(g in device.upper() for g in ("GPU", "CUDA")): - return CURRENT_PLATFORM.gpu_accelerator_for_offload() + return tf_weights_accelerator(obj) except Exception as exc: debug_keras_failure("accelerator detection", exc) - return "cpu" + return accelerators.CPU class KerasExtractor(BaseExtractor): diff --git a/wildedge/integrations/mlx.py b/wildedge/integrations/mlx.py index 4613219..b78a5e3 100644 --- a/wildedge/integrations/mlx.py +++ b/wildedge/integrations/mlx.py @@ -3,14 +3,20 @@ from __future__ import annotations import functools +import sys import threading import time from typing import TYPE_CHECKING -from wildedge import constants +from wildedge import quantizations 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.integrations.common import ( + debug_failure, + mlx_accelerator, + mlx_layer_quantization, +) +from wildedge.loading import optional_module from wildedge.model import ModelInfo from wildedge.timing import elapsed_ms @@ -23,10 +29,9 @@ _mlx_nn = None # type: ignore[assignment] _tree_flatten = None # type: ignore[assignment] -try: - import mlx_lm as _mlx_lm -except ImportError: - _mlx_lm = None # type: ignore[assignment] +# Imported when the integration is activated, so registering an MLX model +# does not import mlx-lm. +_mlx_lm = sys.modules.get("mlx_lm") if TYPE_CHECKING: from wildedge.model import ModelHandle @@ -64,7 +69,7 @@ def extract_model_args(obj: object) -> tuple[str | None, str | None]: quant = getattr(args, "quantization", None) if quant is not None: bits = getattr(quant, "bits", None) - q_str = f"q{int(bits)}" if bits else "quantized" + q_str = quantizations.int_bits(bits) if bits else quantizations.QUANTIZED else: q_str = detect_quantization_from_layers(obj) return model_type, q_str @@ -74,15 +79,11 @@ def extract_model_args(obj: object) -> tuple[str | None, str | None]: def detect_quantization_from_layers(obj: object) -> str | None: - """Inspect layer class names for quantized linear layers.""" try: - for _, module in obj.named_modules(): # type: ignore[union-attr] - cls_name = type(module).__name__ - if "Quantized" in cls_name or "quantized" in cls_name: - return "quantized" - except Exception: - pass - return None + return mlx_layer_quantization(obj) + except Exception as exc: + debug_mlx_failure("quantization detection", exc) + return None def count_tokens(tokenizer: object, text: str) -> int | None: @@ -115,7 +116,7 @@ def patched_call(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise @@ -171,6 +172,7 @@ def memory_bytes(self, obj: object) -> int | None: return None def install_hooks(self, obj: object, handle: ModelHandle) -> None: + handle.detected_accelerator = mlx_accelerator() setattr(obj, MLX_HANDLE_ATTR, handle) patch_instance_call_once( obj, @@ -195,8 +197,11 @@ def install_auto_load_patch(cls, client_ref: object) -> None: tokens/second. The autoregressive ``model()`` loop inside generate is suppressed via a thread-local guard so it does not double-count. """ - global _mlx_patched - if _mlx_patched or _mlx_lm is None: + global _mlx_patched, _mlx_lm + if _mlx_patched: + return + _mlx_lm = _mlx_lm or optional_module("mlx_lm") + if _mlx_lm is None: return with _MLX_PATCH_LOCK: @@ -244,7 +249,7 @@ def patched_load(path_or_hf_repo, *args, **kwargs): # type: ignore[no-untyped-d _mlx_lm.load = patched_load @classmethod - def _patch_generate(cls, client_ref: object) -> None: # noqa: ARG003 + def _patch_generate(cls, client_ref: object) -> None: original_generate = _mlx_lm.generate if ( getattr(original_generate, "__wildedge_patch_name__", None) @@ -268,7 +273,7 @@ def patched_generate(model, tokenizer, prompt, *args, **kwargs): # type: ignore if handle is not None: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise finally: @@ -306,3 +311,7 @@ def patched_generate(model, tokenizer, prompt, *args, **kwargs): # type: ignore patched_generate.__wildedge_patch_name__ = MLX_GENERATE_PATCH_NAME # type: ignore[attr-defined] patched_generate.__wildedge_original_call__ = original_generate # type: ignore[attr-defined] _mlx_lm.generate = patched_generate + + +# Entry point named by the integration registry. +install_auto_load_patch = MlxExtractor.install_auto_load_patch diff --git a/wildedge/integrations/mlx_whisper.py b/wildedge/integrations/mlx_whisper.py new file mode 100644 index 0000000..df91ff8 --- /dev/null +++ b/wildedge/integrations/mlx_whisper.py @@ -0,0 +1,152 @@ +"""Opt-in instrumentation for mlx-whisper's function API and model cache.""" + +from __future__ import annotations + +import contextvars +import functools +import importlib +import inspect +import threading +import time + +from wildedge import quantizations +from wildedge.integrations.common import ( + debug_failure, + dtype_to_quantization, + mlx_accelerator, + mlx_layer_quantization, + track_api_error, +) +from wildedge.integrations.whisper import ( + DESCRIPTOR_ATTR, + MLX_WHISPER, + WhisperDescriptor, + descriptor, + report_cached_load, + report_load, + track_transcription, +) +from wildedge.timing import elapsed_ms + +_PATCH_LOCK = threading.Lock() + +# Defaults of mlx_whisper.transcribe, which picks the model and its dtype from +# path_or_hf_repo and decode_options["fp16"]. +DEFAULT_REPO = "mlx-community/whisper-tiny" +DEFAULT_FP16 = True + + +def describe( + path: object, weights: str | None, model: object = None +) -> WhisperDescriptor: + """Describe a model. ``weights`` is the weight dtype label, reported unless + the model's layers are quantized. Without ``model`` (when loading failed) + only ``weights`` is known.""" + quantization = weights + if model is not None: + try: + quantization = mlx_layer_quantization(model) or weights + except Exception as exc: + debug_failure(MLX_WHISPER, "quantization detection", exc) + return WhisperDescriptor.from_path( + path, MLX_WHISPER, mlx_accelerator(), quantization + ) + + +def requested_model(kwargs: dict) -> tuple[str, str]: + """Model path and weight dtype label that ``mlx_whisper.transcribe`` uses + for these arguments.""" + fp16 = kwargs.get("fp16", DEFAULT_FP16) + weights = quantizations.F16 if fp16 else quantizations.F32 + return kwargs.get("path_or_hf_repo", DEFAULT_REPO), weights + + +def install_auto_load_patch(client_ref): + with _PATCH_LOCK: + _install(client_ref) + + +def _install(client_ref): + root = importlib.import_module("mlx_whisper") + module = importlib.import_module("mlx_whisper.transcribe") + loaders = importlib.import_module("mlx_whisper.load_models") + original = module.transcribe + if getattr(original, "__wildedge_patch_name__", None) == MLX_WHISPER: + return + original_load = module.load_model + load_signature = inspect.signature(original_load) + # The client that owns the transcription in progress, so a load inside it + # reports to the same client even if the default changes meanwhile. + active_client = contextvars.ContextVar("wildedge_mlx_whisper_client", default=None) + + @functools.wraps(original_load) + def load(*args, **kwargs): + client = active_client.get() or client_ref() + t0 = time.perf_counter() + model = original_load(*args, **kwargs) + try: + bound = load_signature.bind(*args, **kwargs) + bound.apply_defaults() + info = describe( + bound.arguments["path_or_hf_repo"], + dtype_to_quantization(str(bound.arguments.get("dtype"))), + model, + ) + setattr(model, DESCRIPTOR_ATTR, info) + if client is not None and not client.closed: + report_load(client, model, elapsed_ms(t0)) + except Exception as exc: + debug_failure(MLX_WHISPER, "model registration", exc) + return model + + @functools.wraps(original) + def transcribe(audio, *args, **kwargs): + client = client_ref() + if client is None: + return original(audio, *args, **kwargs) + t0 = time.perf_counter() + token = active_client.set(client) + try: + result = original(audio, *args, **kwargs) + except Exception as exc: + try: + if not client.closed: + handle = client.register_model(describe(*requested_model(kwargs))) + track_api_error(handle, exc, elapsed_ms(t0)) + except Exception as telemetry_exc: + debug_failure(MLX_WHISPER, "error tracking", telemetry_exc) + raise + finally: + active_client.reset(token) + if client.closed: + return result + try: + # mlx_whisper keeps the last model in ModelHolder. It was either + # loaded by this call (and reported by load above) or reused. + model = module.ModelHolder.model + info = descriptor(model) + if info is None: + # Loaded before instrumentation started. + info = describe(*requested_model(kwargs), model) + setattr(model, DESCRIPTOR_ATTR, info) + if info.model_id not in client.auto_loaded: + report_cached_load(client, model) + handle = client.registry.handles.get( + info.model_id + ) or client.register_model(model) + track_transcription( + handle, + audio, + result.get("text", ""), + duration_ms=elapsed_ms(t0), + language=result.get("language"), + ) + except Exception as exc: + debug_failure(MLX_WHISPER, "transcription tracking", exc) + return result + + transcribe.__wildedge_patch_name__ = MLX_WHISPER + module.load_model = load + loaders.load_model = load + module.transcribe = transcribe + root.transcribe = transcribe diff --git a/wildedge/integrations/onnx.py b/wildedge/integrations/onnx.py index 0fa07b6..894aaee 100644 --- a/wildedge/integrations/onnx.py +++ b/wildedge/integrations/onnx.py @@ -8,7 +8,7 @@ import time from typing import TYPE_CHECKING, Any -from wildedge import constants +from wildedge import accelerators from wildedge.events.inference import ( ClassificationOutputMeta, HistogramSummary, @@ -20,6 +20,7 @@ from wildedge.integrations.common import ( debug_failure, dtype_to_quantization, + hf_repo_from_cache_path, image_brightness_histogram, infer_input_modality_from_names, ) @@ -43,7 +44,19 @@ _ort_patched = False _ORT_PATCH_LOCK = threading.Lock() -_GENERIC_GRAPH_NAMES = {"torch_jit", "main", "model", "network", "graph", "onnx_model"} +# Default graph names from exporters (torch.onnx.export uses torch_jit or +# main_graph), so they say nothing about which model this is. +_GENERIC_GRAPH_NAMES = { + "torch_jit", + "main_graph", + "main", + "model", + "network", + "graph", + "onnx_model", +} +# Exporters write INT64_MAX into model_version to mean "unset". +_UNSET_MODEL_VERSION = 2**63 - 1 ONNX_AUTO_LOAD_PATCH_NAME = "onnx_auto_load" @@ -57,25 +70,46 @@ def model_id_from_path(path: str) -> str | None: - HF Hub cache path (``models--org--name/...``): returns ``org/name`` - Plain file path: returns the stem if it's not a generic name """ - for part in path.replace("\\", "/").split("/"): - if part.startswith("models--"): - return part[len("models--") :].replace("--", "/", 1) - stem = os.path.splitext(os.path.basename(path))[0] + return hf_repo_from_cache_path(path) or model_file_stem(path) + + +def model_file_stem(path: object) -> str | None: + """File name without extension, unless it is a generic name like model.onnx.""" + if not isinstance(path, str) or not path: + return None + stem = os.path.splitext(os.path.basename(path.replace("\\", "/")))[0] return stem if stem and stem not in _GENERIC_GRAPH_NAMES else None +def parse_model_version(raw: object) -> str | None: + """Return the ONNX model_version as a string, or None when it is unset. + + 0, negative values and INT64_MAX all mean the exporter left it unset. + """ + if raw is None: + return None + try: + number = int(raw) + except (TypeError, ValueError): + text = str(raw).strip() + return text or None + if number <= 0 or number >= _UNSET_MODEL_VERSION: + return None + return str(number) + + def is_ort_session(obj: object) -> bool: return ort is not None and isinstance(obj, ort.InferenceSession) _PROVIDER_TO_ACCELERATOR: dict[str, str] = { - "CUDAExecutionProvider": "cuda", - "TensorrtExecutionProvider": "cuda", - "ROCMExecutionProvider": "rocm", - "MIGraphXExecutionProvider": "rocm", - "CoreMLExecutionProvider": "coreml", - "DmlExecutionProvider": "directml", - "OpenVINOExecutionProvider": "openvino", + "CUDAExecutionProvider": accelerators.CUDA, + "TensorrtExecutionProvider": accelerators.CUDA, + "ROCMExecutionProvider": accelerators.ROCM, + "MIGraphXExecutionProvider": accelerators.ROCM, + "CoreMLExecutionProvider": accelerators.COREML, + "DmlExecutionProvider": accelerators.DIRECTML, + "OpenVINOExecutionProvider": accelerators.OPENVINO, } @@ -87,7 +121,7 @@ def detect_accelerator(session: Any) -> str: return acc except Exception as exc: debug_onnx_failure("accelerator detection", exc) - return "cpu" + return accelerators.CPU def detect_quantization(session: Any) -> str | None: @@ -151,12 +185,15 @@ def extract_info( graph_name = ( raw_name if raw_name and raw_name not in _GENERIC_GRAPH_NAMES else None ) - model_id = overrides.pop("id", None) or graph_name - - raw_version = str(meta.version) if meta.version else None - version = overrides.pop("version", None) or ( - raw_version if raw_version and raw_version != "0" else None + # onnxruntime keeps the path a session was created from (None for bytes). + model_path = getattr(obj, "_model_path", None) + model_id = ( + overrides.pop("id", None) + or graph_name + or (model_id_from_path(model_path) if isinstance(model_path, str) else None) ) + + version = overrides.pop("version", None) or parse_model_version(meta.version) if version is None: logger.warning( "wildedge: ONNX model version could not be detected - sending as null" @@ -180,7 +217,10 @@ def extract_info( source = overrides.pop("source", "local") info = ModelInfo( - model_name=graph_name or model_id or "onnx-model", + model_name=graph_name + or model_file_stem(model_path) + or model_id + or "onnx-model", model_version=version, model_source=source, model_format="onnx", @@ -290,7 +330,7 @@ def patched_run(output_names, input_feed, run_options=None): except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise @@ -351,3 +391,7 @@ def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] ) ort.InferenceSession = patched_cls _ort_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = OnnxExtractor.install_auto_load_patch diff --git a/wildedge/integrations/openai.py b/wildedge/integrations/openai.py index 299f9b9..10ad19a 100644 --- a/wildedge/integrations/openai.py +++ b/wildedge/integrations/openai.py @@ -2,12 +2,19 @@ from __future__ import annotations +import base64 import functools import threading import time from typing import TYPE_CHECKING -from wildedge.events.inference import ApiMeta, GenerationOutputMeta +from wildedge.events.inference import ( + ApiMeta, + EmbeddingOutputMeta, + GenerationOutputMeta, + TextInputMeta, +) +from wildedge.instrumentation import bind_client_reference from wildedge.integrations.base import BaseExtractor from wildedge.integrations.common import ( SOURCE_BY_HOSTNAME, # noqa: F401 (re-exported; moved to common) @@ -16,6 +23,7 @@ build_input_meta, debug_failure, source_from_base_url, + stream_error_code, track_api_error, ) from wildedge.model import ModelInfo @@ -171,14 +179,19 @@ def on_chunk(chunk: object) -> None: if reason: stop_reason[0] = reason - def on_done(duration_ms: int, ttft_ms: int | None) -> None: + def on_done( + duration_ms: int, ttft_ms: int | None, *, outcome: str = "completed" + ) -> None: ti, to, sr = tokens_in[0], tokens_out[0], stop_reason[0] + if outcome == "completed" and sr is None: + outcome = "incomplete" output_meta = build_streaming_output_meta(ttft_ms, ti, to, sr, duration_ms) handle.track_inference( duration_ms=duration_ms, input_modality="text", output_modality="generation", - success=True, + success=outcome == "completed", + error_code=stream_error_code(outcome), input_meta=build_input_meta(messages, ti), output_meta=output_meta, api_meta=build_api_meta(last_chunk[0]) @@ -208,8 +221,7 @@ def patched_create(*args, **kwargs): if is_streaming: on_chunk, on_done = make_openai_stream_callbacks(handle, messages) return SyncStreamWrapper(result, handle, t0, on_chunk, on_done) - else: - record_inference(handle, result, messages, elapsed_ms(t0)) + record_inference(handle, result, messages, elapsed_ms(t0)) return result except Exception as exc: if handle is not None: @@ -240,8 +252,7 @@ async def patched_create(*args, **kwargs): if is_streaming: on_chunk, on_done = make_openai_stream_callbacks(handle, messages) return AsyncStreamWrapper(result, handle, t0, on_chunk, on_done) - else: - record_inference(handle, result, messages, elapsed_ms(t0)) + record_inference(handle, result, messages, elapsed_ms(t0)) return result except Exception as exc: if handle is not None: @@ -251,6 +262,285 @@ async def patched_create(*args, **kwargs): completions.create = patched_create # type: ignore[attr-defined] +def responses_messages(value: object) -> list: + """Normalize text inputs without consuming arbitrary input iterables.""" + if isinstance(value, str): + return [{"role": "user", "content": value}] + return value if isinstance(value, list) else [] + + +def build_responses_output_meta( + response: object, duration_ms: int +) -> GenerationOutputMeta: + usage = getattr(response, "usage", None) + tokens_out = getattr(usage, "output_tokens", None) + status = getattr(response, "status", None) + reason = getattr(getattr(response, "incomplete_details", None), "reason", None) + return GenerationOutputMeta( + tokens_in=getattr(usage, "input_tokens", None), + tokens_out=tokens_out, + cached_input_tokens=getattr( + getattr(usage, "input_tokens_details", None), "cached_tokens", None + ), + reasoning_tokens_out=getattr( + getattr(usage, "output_tokens_details", None), "reasoning_tokens", None + ), + tokens_per_second=round(tokens_out / duration_ms * 1000, 1) + if duration_ms > 0 and tokens_out + else None, + stop_reason=reason or status, + ) + + +def record_response( + handle: ModelHandle, + response: object, + messages: list, + duration: int, + ttft_ms: int | None = None, +) -> None: + if getattr(response, "status", None) in ("queued", "in_progress"): + return + output = build_responses_output_meta(response, duration) + output.time_to_first_token_ms = ttft_ms + error = getattr(response, "error", None) + handle.track_inference( + duration_ms=duration, + input_modality="text", + output_modality="generation", + success=getattr(response, "status", None) != "failed" and error is None, + error_code=getattr(error, "code", None), + input_meta=build_input_meta(messages, output.tokens_in), + output_meta=output, + api_meta=build_api_meta(response), + ) + + +def make_responses_stream_callbacks( + handle: ModelHandle, messages: list, t0: float +) -> tuple: + response = None + ttft = None + recorded = False + + def on_chunk(chunk: object) -> None: + nonlocal response, ttft, recorded + event_type = getattr(chunk, "type", "") + if ( + event_type + in ( + "response.output_text.delta", + "response.refusal.delta", + "response.function_call_arguments.delta", + "response.output_audio.delta", + ) + and getattr(chunk, "delta", None) + and ttft is None + ): + ttft = elapsed_ms(t0) + if event_type in ( + "response.completed", + "response.incomplete", + "response.failed", + ): + response = getattr(chunk, "response", None) + elif event_type == "error" and not recorded: + handle.track_error( + error_code=getattr(chunk, "code", None) or "PROVIDER_ERROR", + duration_ms=elapsed_ms(t0), + ) + recorded = True + + def on_done( + duration_ms: int, _first_event_ms: int | None, *, outcome: str = "completed" + ) -> None: + nonlocal recorded + if response is None and not recorded: + recorded = True + handle.track_inference( + duration_ms=duration_ms, + input_modality="text", + output_modality="generation", + input_meta=build_input_meta(messages, None), + output_meta=GenerationOutputMeta( + time_to_first_token_ms=ttft, + stop_reason=outcome if outcome != "completed" else "incomplete", + ), + success=False, + error_code="STREAM_INCOMPLETE" + if outcome == "completed" + else stream_error_code(outcome), + ) + if response is not None and not recorded: + recorded = True + record_response(handle, response, messages, duration_ms, ttft) + + return on_chunk, on_done + + +def wrap_sync_responses(responses: object, source: str, client_ref: object) -> None: + original_create = responses.create + model_handles: dict[str, ModelHandle] = {} + + @functools.wraps(original_create) + def patched_create(*args, **kwargs): + client = client_ref() + model_id = kwargs.get("model") + if client is None or client.closed or not model_id: + return original_create(*args, **kwargs) + handle = resolve_handle(model_id, responses, model_handles, client, source) + t0 = time.perf_counter() + try: + result = original_create(*args, **kwargs) + except Exception as exc: + if handle is not None: + track_api_error(handle, exc, elapsed_ms(t0)) + raise + if handle is not None: + try: + messages = responses_messages(kwargs.get("input")) + if kwargs.get("stream"): + callbacks = make_responses_stream_callbacks(handle, messages, t0) + return SyncStreamWrapper(result, handle, t0, *callbacks) + record_response(handle, result, messages, elapsed_ms(t0)) + except Exception as exc: + debug_openai_failure("responses tracking", exc) + return result + + responses.create = patched_create + + +def wrap_async_responses(responses: object, source: str, client_ref: object) -> None: + original_create = responses.create + model_handles: dict[str, ModelHandle] = {} + + @functools.wraps(original_create) + async def patched_create(*args, **kwargs): + client = client_ref() + model_id = kwargs.get("model") + if client is None or client.closed or not model_id: + return await original_create(*args, **kwargs) + handle = resolve_handle(model_id, responses, model_handles, client, source) + t0 = time.perf_counter() + try: + result = await original_create(*args, **kwargs) + except Exception as exc: + if handle is not None: + track_api_error(handle, exc, elapsed_ms(t0)) + raise + if handle is not None: + try: + messages = responses_messages(kwargs.get("input")) + if kwargs.get("stream"): + callbacks = make_responses_stream_callbacks(handle, messages, t0) + return AsyncStreamWrapper(result, handle, t0, *callbacks) + record_response(handle, result, messages, elapsed_ms(t0)) + except Exception as exc: + debug_openai_failure("responses tracking", exc) + return result + + responses.create = patched_create + + +def record_embedding( + handle: ModelHandle, result: object, value: object, duration: int +) -> None: + """Record metadata only; never retain input text or embedding vectors.""" + usage = getattr(result, "usage", None) + tokens = getattr(usage, "prompt_tokens", None) + batch_size = None + chars = words = None + if isinstance(value, str): + batch_size = 1 + chars, words = len(value), len(value.split()) + elif isinstance(value, (list, tuple)): + if value and isinstance(value[0], int): + batch_size = 1 + else: + batch_size = len(value) + if value and all(isinstance(item, str) for item in value): + chars = sum(len(item) for item in value) + words = sum(len(item.split()) for item in value) + data = getattr(result, "data", None) + dimensions = None + if data: + vector = getattr(data[0], "embedding", None) + if isinstance(vector, (list, tuple)): + dimensions = len(vector) + elif isinstance(vector, str): + byte_count = len(base64.b64decode(vector, validate=True)) + if byte_count % 4 == 0: + dimensions = byte_count // 4 + handle.track_inference( + duration_ms=duration, + input_modality="text", + output_modality="embedding", + success=True, + batch_size=batch_size, + input_meta=TextInputMeta( + char_count=chars, word_count=words, token_count=tokens + ), + output_meta=EmbeddingOutputMeta(dimensions=dimensions, tokens_in=tokens), + api_meta=build_api_meta(result), + ) + + +def wrap_sync_embeddings(embeddings: object, source: str, client_ref: object) -> None: + original_create = embeddings.create + model_handles: dict[str, ModelHandle] = {} + + @functools.wraps(original_create) + def patched_create(*args, **kwargs): + client = client_ref() + model_id = kwargs.get("model") + if client is None or client.closed or not model_id: + return original_create(*args, **kwargs) + handle = resolve_handle(model_id, embeddings, model_handles, client, source) + t0 = time.perf_counter() + try: + result = original_create(*args, **kwargs) + except Exception as exc: + if handle is not None: + track_api_error(handle, exc, elapsed_ms(t0)) + raise + if handle is not None: + try: + record_embedding(handle, result, kwargs.get("input"), elapsed_ms(t0)) + except Exception as exc: + debug_openai_failure("embedding tracking", exc) + return result + + embeddings.create = patched_create + + +def wrap_async_embeddings(embeddings: object, source: str, client_ref: object) -> None: + original_create = embeddings.create + model_handles: dict[str, ModelHandle] = {} + + @functools.wraps(original_create) + async def patched_create(*args, **kwargs): + client = client_ref() + model_id = kwargs.get("model") + if client is None or client.closed or not model_id: + return await original_create(*args, **kwargs) + handle = resolve_handle(model_id, embeddings, model_handles, client, source) + t0 = time.perf_counter() + try: + result = await original_create(*args, **kwargs) + except Exception as exc: + if handle is not None: + track_api_error(handle, exc, elapsed_ms(t0)) + raise + if handle is not None: + try: + record_embedding(handle, result, kwargs.get("input"), elapsed_ms(t0)) + except Exception as exc: + debug_openai_failure("embedding tracking", exc) + return result + + embeddings.create = patched_create + + class OpenAIExtractor(BaseExtractor): def can_handle(self, obj: object) -> bool: return type(obj).__name__ in ( @@ -258,6 +548,10 @@ def can_handle(self, obj: object) -> bool: "AsyncOpenAI", "Completions", "AsyncCompletions", + "Responses", + "AsyncResponses", + "Embeddings", + "AsyncEmbeddings", ) def extract_info( @@ -282,7 +576,7 @@ def install_hooks(self, obj: object, handle: ModelHandle) -> None: @classmethod def install_auto_load_patch(cls, client_ref: object) -> None: - """Patch openai.OpenAI and openai.AsyncOpenAI to wrap chat.completions.create.""" + """Patch new OpenAI clients to track Chat Completions, Responses, and embeddings.""" global _openai_patched if _openai_patched or _openai is None: return @@ -305,29 +599,57 @@ def patched_sync_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-d original_sync_init(self_inner, *args, **kwargs) c = client_ref() # type: ignore[call-arg] if c is not None and not c.closed: + bound_ref = bind_client_reference(client_ref, c) source = source_from_base_url( str(getattr(self_inner, "base_url", None) or "") ) try: wrap_sync_completions( - self_inner.chat.completions, source, client_ref + self_inner.chat.completions, source, bound_ref ) except Exception as exc: debug_openai_failure("sync client wrap", exc) + if hasattr(self_inner, "responses"): + try: + wrap_sync_responses(self_inner.responses, source, bound_ref) + except Exception as exc: + debug_openai_failure("sync responses wrap", exc) + if hasattr(self_inner, "embeddings"): + try: + wrap_sync_embeddings( + self_inner.embeddings, source, bound_ref + ) + except Exception as exc: + debug_openai_failure("sync embeddings wrap", exc) def patched_async_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] original_async_init(self_inner, *args, **kwargs) c = client_ref() # type: ignore[call-arg] if c is not None and not c.closed: + bound_ref = bind_client_reference(client_ref, c) source = source_from_base_url( str(getattr(self_inner, "base_url", None) or "") ) try: wrap_async_completions( - self_inner.chat.completions, source, client_ref + self_inner.chat.completions, source, bound_ref ) except Exception as exc: debug_openai_failure("async client wrap", exc) + if hasattr(self_inner, "responses"): + try: + wrap_async_responses( + self_inner.responses, source, bound_ref + ) + except Exception as exc: + debug_openai_failure("async responses wrap", exc) + if hasattr(self_inner, "embeddings"): + try: + wrap_async_embeddings( + self_inner.embeddings, source, bound_ref + ) + except Exception as exc: + debug_openai_failure("async embeddings wrap", exc) patched_sync_init.__wildedge_patch_name__ = OPENAI_INIT_PATCH_NAME # type: ignore[attr-defined] patched_sync_init.__wildedge_original_call__ = original_sync_init # type: ignore[attr-defined] @@ -337,3 +659,7 @@ def patched_async_init(self_inner, *args, **kwargs): # type: ignore[no-untyped- _openai.OpenAI.__init__ = patched_sync_init _openai.AsyncOpenAI.__init__ = patched_async_init _openai_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = OpenAIExtractor.install_auto_load_patch diff --git a/wildedge/integrations/pytorch.py b/wildedge/integrations/pytorch.py index c9f414b..48119be 100644 --- a/wildedge/integrations/pytorch.py +++ b/wildedge/integrations/pytorch.py @@ -3,10 +3,12 @@ from __future__ import annotations import functools +import sys import threading import time from typing import TYPE_CHECKING +from wildedge import accelerators, quantizations from wildedge.events.inference import ( ClassificationOutputMeta, HistogramSummary, @@ -20,17 +22,16 @@ dtype_to_quantization, image_brightness_histogram, infer_input_modality_from_names, + torch_module_accelerator, ) +from wildedge.loading import optional_module from wildedge.logging import logger from wildedge.model import ModelInfo from wildedge.timing import elapsed_ms -try: - import timm as _timm - from timm.data import ImageNetInfo as _ImageNetInfo -except ImportError: - _timm = None # type: ignore[assignment] - _ImageNetInfo = None # type: ignore[assignment,misc] +# Imported when the timm integration is activated, so registering a plain +# PyTorch model does not import timm. +_timm = sys.modules.get("timm") try: import torch as _torch @@ -44,6 +45,7 @@ _timm_patched = False _TIMM_PATCH_LOCK = threading.Lock() TIMM_AUTO_LOAD_PATCH_NAME = "timm_auto_load" +TIMM_MODEL_NAME_ATTR = "__wildedge_timm_model_name__" debug_pytorch_failure = functools.partial(debug_failure, "pytorch") @@ -56,22 +58,12 @@ def is_torch_module(obj: object) -> bool: ) -def parameter_device_type(obj: object) -> str | None: +def detect_accelerator(obj: object) -> str: try: - first_param = next(obj.parameters()) # type: ignore[union-attr] - except StopIteration: - return None + return torch_module_accelerator(obj) or accelerators.CPU except Exception as exc: - debug_pytorch_failure("parameter device discovery", exc) - return None - return str(getattr(first_param.device, "type", "") or None) - - -def detect_accelerator(obj: object) -> str: - device_type = parameter_device_type(obj) - if device_type: - return device_type # 'cpu', 'cuda', 'mps', 'xpu', etc. - return "cpu" + debug_pytorch_failure("accelerator detection", exc) + return accelerators.CPU def detect_quantization(obj: object) -> str | None: @@ -79,10 +71,10 @@ def detect_quantization(obj: object) -> str | None: for module in obj.modules(): # type: ignore[union-attr] cls = type(module) if "quantized" in cls.__module__ or "quant" in cls.__name__.lower(): - return "int8" + return quantizations.INT8 for p in obj.parameters(): # type: ignore[union-attr] q = dtype_to_quantization(str(p.dtype)) - if q and q != "f32": + if q and q != quantizations.F32: return q except Exception as exc: debug_pytorch_failure("quantization detection", exc) @@ -102,10 +94,7 @@ def image_input_meta(tensor: object) -> ImageInputMeta | None: t_min = float(tensor.min().item()) # type: ignore[union-attr] t_max = float(tensor.max().item()) # type: ignore[union-attr] span = t_max - t_min - if span > 0: - norm = (tensor - t_min) / span # type: ignore[operator] - else: - norm = tensor # type: ignore[assignment] + norm = (tensor - t_min) / span if span > 0 else tensor # type: ignore[operator] brightness_mean, brightness_stddev, buckets = image_brightness_histogram( norm.flatten() @@ -129,7 +118,9 @@ def image_input_meta(tensor: object) -> ImageInputMeta | None: def build_imagenet_labels() -> list[str] | None: """Return a 1000-entry description list using timm's bundled ImageNet metadata.""" try: - info = _ImageNetInfo() # type: ignore[misc] + from timm.data import ImageNetInfo # noqa: PLC0415 -- optional, 1000-class only + + info = ImageNetInfo() return [info.index_to_description(i) for i in range(1000)] except Exception as exc: debug_pytorch_failure("imagenet label build", exc) @@ -157,7 +148,7 @@ def classification_output_meta( label=labels[int(idx.item())] if labels else str(int(idx.item())), confidence=round(float(conf.item()), 4), ) - for conf, idx in zip(top_vals, top_idx) + for conf, idx in zip(top_vals, top_idx, strict=True) ] return ClassificationOutputMeta( @@ -167,6 +158,20 @@ def classification_output_meta( ) +def timm_model_name(obj: object) -> str | None: + """Prefer the full requested name, including pretrained tags or hub paths.""" + name = getattr(obj, TIMM_MODEL_NAME_ATTR, None) + if isinstance(name, str) and name: + return name + for key in ("pretrained_cfg", "default_cfg"): + config = getattr(obj, key, None) + if isinstance(config, dict): + name = config.get("architecture") + if isinstance(name, str) and name: + return name + return None + + class PytorchExtractor(BaseExtractor): def can_handle(self, obj: object) -> bool: return is_torch_module(obj) @@ -175,9 +180,15 @@ def extract_info( self, obj: object, overrides: dict ) -> tuple[str | None, ModelInfo]: class_name = type(obj).__name__ - model_id = overrides.pop("id", None) or class_name - - if class_name in _GENERIC_MODULE_NAMES: + model_name = timm_model_name(obj) or class_name + explicit_id = overrides.pop("id", None) + model_id = explicit_id or model_name + + if ( + model_name == class_name + and class_name in _GENERIC_MODULE_NAMES + and not explicit_id + ): logger.warning( "wildedge: PyTorch model_id defaults to generic class name '%s' - " "pass model_id= to register_model() for a stable unique ID", @@ -190,7 +201,7 @@ def extract_info( source = overrides.pop("source", "local") info = ModelInfo( - model_name=class_name, + model_name=model_name, model_version=version, model_source=source, model_format="pytorch", @@ -224,7 +235,7 @@ def install_hooks(self, obj: object, handle: ModelHandle) -> None: nc = getattr(obj, "num_classes", 0) if isinstance(nc, int) and nc > 1: num_classes = nc - if num_classes == 1000 and _timm is not None: + if num_classes == 1000: labels = build_imagenet_labels() except Exception as exc: debug_pytorch_failure("num_classes detection", exc) @@ -292,8 +303,11 @@ def install_timm_patch(cls, client_ref: object) -> None: the duration of the call and emitted as a model_download event. Inference tracking is handled by the existing PyTorch forward hooks. """ - global _timm_patched - if _timm_patched or _timm is None: + global _timm_patched, _timm + if _timm_patched: + return + _timm = _timm or optional_module("timm") + if _timm is None: return with _TIMM_PATCH_LOCK: @@ -318,6 +332,9 @@ def patched_create_model(*args, **kwargs): # type: ignore[no-untyped-def] load_ms = elapsed_ms(t0) logger.debug("wildedge: timm.create_model done load_ms=%d", load_ms) if c is not None and not c.closed: + name = args[0] if args else kwargs.get("model_name") + if isinstance(name, str) and name: + setattr(model, TIMM_MODEL_NAME_ATTR, name) downloads = c._diff_hub_caches(hub_before, load_ms) or None c._on_model_auto_loaded(model, load_ms=load_ms, downloads=downloads) return model diff --git a/wildedge/integrations/registry.py b/wildedge/integrations/registry.py index 8321720..94fc78c 100644 --- a/wildedge/integrations/registry.py +++ b/wildedge/integrations/registry.py @@ -3,28 +3,48 @@ Contains only ML *framework* integrations (inference tracking, load/unload timing). Model hub and repository trackers (download provenance) live in ``wildedge.hubs.registry``. They have different activation semantics. + +Adapters are listed by name and imported only when used, so importing the SDK +does not import any ML framework (see ``wildedge.loading``). """ from __future__ import annotations +import sys from dataclasses import dataclass +from functools import cache from typing import Literal +from wildedge.loading import resolve + IntegrationKind = Literal["noop", "client_patch"] +ADAPTER_PACKAGE = "wildedge.integrations" + @dataclass(frozen=True) class IntegrationSpec: - """Describes one integration's runtime behavior and import requirements.""" + """Describes one integration's runtime behavior and import requirements. + + A ``client_patch`` integration is activated by the + ``install_auto_load_patch`` function of the adapter module with its name. + """ name: str required_modules: tuple[str, ...] kind: IntegrationKind + adapter_module: str | None = None + + @property + def installer(self) -> str | None: + if self.kind != "client_patch": + return None + return f"{ADAPTER_PACKAGE}.{self.adapter_module or self.name}:install_auto_load_patch" INTEGRATION_SPECS: tuple[IntegrationSpec, ...] = ( IntegrationSpec("anthropic", ("anthropic",), "client_patch"), - IntegrationSpec("gguf", ("llama_cpp",), "client_patch"), + IntegrationSpec("llama_cpp", ("llama_cpp",), "client_patch", adapter_module="gguf"), IntegrationSpec("onnx", ("onnxruntime",), "client_patch"), IntegrationSpec("openai", ("openai",), "client_patch"), IntegrationSpec("timm", ("timm",), "client_patch"), @@ -33,6 +53,8 @@ class IntegrationSpec: IntegrationSpec("tensorflow", ("tensorflow",), "client_patch"), IntegrationSpec("ultralytics", ("ultralytics",), "client_patch"), IntegrationSpec("transformers", ("transformers",), "client_patch"), + IntegrationSpec("faster_whisper", ("faster_whisper",), "client_patch"), + IntegrationSpec("mlx_whisper", ("mlx_whisper",), "client_patch"), IntegrationSpec("mlx", ("mlx_lm",), "client_patch"), ) @@ -41,11 +63,113 @@ class IntegrationSpec: } +# Public names describe runtimes; retain the pre-0.3 name as an input alias. +INTEGRATION_ALIASES = {"gguf": "llama_cpp"} +INTEGRATIONS_BY_NAME.update( + {alias: INTEGRATIONS_BY_NAME[name] for alias, name in INTEGRATION_ALIASES.items()} +) + + +def canonical_integration(name: str) -> str: + return INTEGRATION_ALIASES.get(name, name) + + def supported_integrations() -> set[str]: """Return all supported integration names.""" - return set(INTEGRATIONS_BY_NAME.keys()) + return {spec.name for spec in INTEGRATION_SPECS} def noop_integrations() -> set[str]: """Return integrations that intentionally perform no global patching.""" return {name for name, spec in INTEGRATIONS_BY_NAME.items() if spec.kind == "noop"} + + +@dataclass(frozen=True) +class ExtractorSpec: + """An extractor, and cheap checks for when it is worth importing. + + ``could_handle`` only decides whether to import the adapter; its + ``can_handle`` still makes the final call. + """ + + module: str # adapter module in wildedge.integrations + class_name: str # extractor class in that module + # Packages whose classes the extractor handles, matched anywhere in the + # object's class hierarchy so application subclasses count too. + packages: tuple[str, ...] = () + # Exact class names, for extractors whose can_handle checks the name. + class_names: tuple[str, ...] = () + # Attribute the SDK sets on objects it has already described. + marker_attr: str | None = None + + def could_handle(self, obj: object) -> bool: + if f"{ADAPTER_PACKAGE}.{self.module}" in sys.modules: + return True # already imported, so asking it costs nothing + if type(obj).__name__ in self.class_names: + return True + if self.marker_attr and getattr(obj, self.marker_attr, None) is not None: + return True + return any( + cls.__module__.partition(".")[0] in self.packages + for cls in type(obj).__mro__ + ) + + def load(self): + return load_extractor(f"{ADAPTER_PACKAGE}.{self.module}:{self.class_name}") + + +@cache +def load_extractor(target: str): + return resolve(target)() + + +# Order matters: the first extractor that handles a model wins, so wrappers +# and specialized formats come before the generic tensor frameworks. +EXTRACTOR_SPECS: tuple[ExtractorSpec, ...] = ( + ExtractorSpec("onnx", "OnnxExtractor", ("onnxruntime",)), + ExtractorSpec("gguf", "GgufExtractor", class_names=("Llama",)), + ExtractorSpec( + "openai", + "OpenAIExtractor", + class_names=( + "OpenAI", + "AsyncOpenAI", + "Completions", + "AsyncCompletions", + "Responses", + "AsyncResponses", + "Embeddings", + "AsyncEmbeddings", + ), + ), + ExtractorSpec( + "anthropic", + "AnthropicExtractor", + class_names=("Anthropic", "AsyncAnthropic", "Messages", "AsyncMessages"), + ), + ExtractorSpec("ultralytics", "UltralyticsExtractor", class_names=("YOLO",)), + ExtractorSpec("transformers", "TransformersExtractor", ("transformers",)), + # Before MLX, which would otherwise claim MLX Whisper models. + ExtractorSpec( + "whisper", + "WhisperExtractor", + class_names=("WhisperDescriptor",), + marker_attr="__wildedge_whisper_descriptor__", + ), + ExtractorSpec("mlx", "MlxExtractor", ("mlx",)), + ExtractorSpec("pytorch", "PytorchExtractor", ("torch",)), + ExtractorSpec( + "tensorflow", "TensorflowExtractor", ("tensorflow", "keras", "tf_keras") + ), + ExtractorSpec("keras", "KerasExtractor", ("keras", "tf_keras")), +) + + +def find_extractor(obj: object): + """The first extractor that handles ``obj``, or None.""" + for spec in EXTRACTOR_SPECS: + if spec.could_handle(obj): + extractor = spec.load() + if extractor.can_handle(obj): + return extractor + return None diff --git a/wildedge/integrations/tensorflow.py b/wildedge/integrations/tensorflow.py index 03d2416..114ba4a 100644 --- a/wildedge/integrations/tensorflow.py +++ b/wildedge/integrations/tensorflow.py @@ -3,6 +3,7 @@ from __future__ import annotations import functools +import sys import threading import time from typing import TYPE_CHECKING @@ -12,25 +13,25 @@ except ImportError: _np = None # type: ignore[assignment] -from wildedge import constants +from wildedge import accelerators from wildedge.events.inference import ClassificationOutputMeta, TopKPrediction from wildedge.integrations.base import BaseExtractor, patch_instance_call_once from wildedge.integrations.common import ( debug_failure, infer_input_modality_from_layer_types, num_classes_from_output_shape, + tf_weights_accelerator, ) +from wildedge.loading import optional_module from wildedge.model import ModelInfo -from wildedge.platforms import CURRENT_PLATFORM from wildedge.timing import elapsed_ms if TYPE_CHECKING: from wildedge.model import ModelHandle -try: - import tensorflow as _tf # type: ignore[import-untyped] -except ImportError: - _tf = None # type: ignore[assignment] +# Imported when the integration is activated, so registering a Keras model +# does not import TensorFlow. +_tf = sys.modules.get("tensorflow") TENSORFLOW_CALL_PATCH_NAME = "tensorflow_call" TENSORFLOW_HANDLE_ATTR = "__wildedge_tensorflow_handle__" @@ -80,14 +81,10 @@ def is_tensorflow_model(obj: object) -> bool: def detect_accelerator(obj: object) -> str: try: - weights = getattr(obj, "weights", None) or [] - if weights: - device = str(getattr(weights[0], "device", "") or "") - if any(g in device.upper() for g in ("GPU", "CUDA")): - return CURRENT_PLATFORM.gpu_accelerator_for_offload() + return tf_weights_accelerator(obj) except Exception as exc: debug_tensorflow_failure("accelerator detection", exc) - return "cpu" + return accelerators.CPU def build_patched_call( @@ -120,7 +117,7 @@ def patched_call(self_inner, *args, **kwargs): except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise @@ -164,11 +161,11 @@ def patched_predict(*args, **kwargs): except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise - setattr(obj, "predict", patched_predict) + obj.predict = patched_predict setattr(obj, TENSORFLOW_PREDICT_PATCHED_ATTR, True) @@ -249,8 +246,11 @@ def install_hooks(self, obj: object, handle: ModelHandle) -> None: @classmethod def install_auto_load_patch(cls, client_ref: object) -> None: - global _tf_auto_patched - if _tf_auto_patched or _tf is None: + global _tf_auto_patched, _tf + if _tf_auto_patched: + return + _tf = _tf or optional_module("tensorflow") + if _tf is None: return with _TF_AUTO_PATCH_LOCK: @@ -281,3 +281,7 @@ def install_auto_load_patch(cls, client_ref: object) -> None: saved_model.load = build_load_patch(client_ref, saved_model_load) _tf_auto_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = TensorflowExtractor.install_auto_load_patch diff --git a/wildedge/integrations/timm.py b/wildedge/integrations/timm.py new file mode 100644 index 0000000..187dd8f --- /dev/null +++ b/wildedge/integrations/timm.py @@ -0,0 +1,9 @@ +"""timm integration: times ``timm.create_model`` and records its downloads. + +The instrumentation lives with the PyTorch adapter, since timm models are +PyTorch modules. This module is the entry point the integration registry names. +""" + +from wildedge.integrations.pytorch import PytorchExtractor + +install_auto_load_patch = PytorchExtractor.install_timm_patch diff --git a/wildedge/integrations/transformers.py b/wildedge/integrations/transformers.py index e5a3b2c..bdf4a27 100644 --- a/wildedge/integrations/transformers.py +++ b/wildedge/integrations/transformers.py @@ -6,9 +6,10 @@ import os import threading import time +from collections import Counter from typing import TYPE_CHECKING -from wildedge import constants +from wildedge import accelerators, quantizations from wildedge.events.inference import ( ClassificationOutputMeta, EmbeddingOutputMeta, @@ -17,7 +18,11 @@ TopKPrediction, ) from wildedge.integrations.base import BaseExtractor -from wildedge.integrations.common import debug_failure, dtype_to_quantization +from wildedge.integrations.common import ( + debug_failure, + dtype_to_quantization, + torch_module_accelerator, +) from wildedge.model import ModelInfo from wildedge.timing import elapsed_ms @@ -41,9 +46,23 @@ # 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.__init__ so the pipeline-level tracking takes precedence. -_tl = threading.local() + +class _ThreadState(threading.local): + """Per-thread state that keeps pipeline-level tracking the only record. + + inside_pipeline_init: from_pretrained calls made by Pipeline.__init__ are + not tracked on their own. + pipeline_models: id() of the model each running tracked pipeline call + wraps, counted for nesting. That model's forward hooks record nothing + while it is here; other models called during the pipeline call do. + """ + + def __init__(self) -> None: + self.inside_pipeline_init = False + self.pipeline_models: Counter[int] = Counter() + + +_tl = _ThreadState() debug_transformers_failure = functools.partial(debug_failure, "transformers") @@ -108,7 +127,7 @@ def detect_quantization(obj: object) -> str | None: if quant_type: return str(quant_type).lower() if bits: - return f"int{int(bits)}" + return quantizations.int_bits(bits) # Fall back to model dtype model = getattr(obj, "model", obj) dtype = getattr(model, "dtype", None) @@ -121,12 +140,12 @@ def detect_quantization(obj: object) -> str | None: def detect_accelerator(obj: object) -> str: try: + # A pipeline wraps the model; a bare model is its own module. model = getattr(obj, "model", obj) - first = next(model.parameters()) # type: ignore[union-attr] - return str(getattr(first.device, "type", "cpu") or "cpu") + return torch_module_accelerator(model) or accelerators.CPU except Exception as exc: debug_transformers_failure("accelerator detection", exc) - return "cpu" + return accelerators.CPU def infer_task_from_arch(arch: str | None) -> str | None: @@ -304,10 +323,16 @@ def install_hooks(self, obj: object, handle: ModelHandle) -> None: _, _, arch = extract_model_config(obj) task_hint = infer_task_from_arch(arch) + # A tracked pipeline call records its own model's forward passes, + # so both hooks skip them. def pre_hook(module, args): # type: ignore[no-untyped-def] + if _tl.pipeline_models[id(module)]: + return _local.t0 = time.perf_counter() def post_hook(module, args, output): # type: ignore[no-untyped-def] + if _tl.pipeline_models[id(module)]: + return t0 = getattr(_local, "t0", None) duration_ms = elapsed_ms(t0) if t0 is not None else 0 @@ -419,6 +444,8 @@ def patched_call(self, inputs, *args, **kwargs): # type: ignore[no-untyped-def] input_modality, output_modality = pipeline_modalities(task) t0 = time.perf_counter() + model_key = id(getattr(self, "model", None)) + _tl.pipeline_models[model_key] += 1 try: outputs = original_call(self, inputs, *args, **kwargs) duration_ms = elapsed_ms(t0) @@ -436,9 +463,13 @@ def patched_call(self, inputs, *args, **kwargs): # type: ignore[no-untyped-def] except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise + finally: + _tl.pipeline_models[model_key] -= 1 + if not _tl.pipeline_models[model_key]: + del _tl.pipeline_models[model_key] patched_call.__wildedge_patch_name__ = PIPELINE_CALL_PATCH_NAME # type: ignore[attr-defined] patched_call.__wildedge_original_call__ = original_call # type: ignore[attr-defined] @@ -457,7 +488,7 @@ def _patch_from_pretrained(cls, client_ref: object) -> None: def patched_from_pretrained(model_cls, *args, **kwargs): # type: ignore[no-untyped-def] # Don't double-track models loaded inside Pipeline.__init__ - if getattr(_tl, "inside_pipeline_init", False): + if _tl.inside_pipeline_init: return original_func(model_cls, *args, **kwargs) c = client_ref() # type: ignore[call-arg] hub_before = ( @@ -478,3 +509,7 @@ def patched_from_pretrained(model_cls, *args, **kwargs): # type: ignore[no-unty _transformers.PreTrainedModel.from_pretrained = classmethod( # type: ignore[assignment] patched_from_pretrained ) + + +# Entry point named by the integration registry. +install_auto_load_patch = TransformersExtractor.install_auto_load_patch diff --git a/wildedge/integrations/ultralytics.py b/wildedge/integrations/ultralytics.py index ccd4531..397a26e 100644 --- a/wildedge/integrations/ultralytics.py +++ b/wildedge/integrations/ultralytics.py @@ -9,7 +9,7 @@ from pathlib import Path from typing import TYPE_CHECKING -from wildedge import constants +from wildedge import accelerators from wildedge.events.inference import ( ClassificationOutputMeta, DetectionOutputMeta, @@ -22,19 +22,34 @@ debug_failure, dtype_to_quantization, image_brightness_histogram, + torch_module_accelerator, ) from wildedge.logging import logger from wildedge.model import ModelInfo from wildedge.timing import elapsed_ms +# Oldest supported release; earlier ones keep their helpers under +# ultralytics.yolo.utils instead of ultralytics.utils. +MIN_ULTRALYTICS_VERSION = "8.0.200" + try: import ultralytics as _ultralytics - from ultralytics.utils import ( - WEIGHTS_DIR as _ULTRALYTICS_WEIGHTS_DIR, # type: ignore[import-untyped] - ) except ImportError: _ultralytics = None # type: ignore[assignment] - _ULTRALYTICS_WEIGHTS_DIR = None # type: ignore[assignment] + +_ULTRALYTICS_WEIGHTS_DIR = None +# Set when ultralytics is installed but older than MIN_ULTRALYTICS_VERSION. +_unsupported_ultralytics_version: str | None = None +if _ultralytics is not None: + try: + from ultralytics.utils import ( + WEIGHTS_DIR as _ULTRALYTICS_WEIGHTS_DIR, # type: ignore[import-untyped] + ) + except ImportError: + _unsupported_ultralytics_version = getattr( + _ultralytics, "__version__", "unknown" + ) + _ultralytics = None # type: ignore[assignment] try: import numpy as np # type: ignore[import-untyped] @@ -49,7 +64,7 @@ _ULTRALYTICS_PATCH_LOCK = threading.Lock() # --- Marker names written onto patched objects and classes --- -YOLO_CALL_PATCH_NAME = "ultralytics_call" +YOLO_PREDICT_PATCH_NAME = "ultralytics_predict" YOLO_HANDLE_ATTR = "__wildedge_yolo_handle__" YOLO_AUTO_LOAD_PATCH_NAME = "ultralytics_auto_load" @@ -67,18 +82,12 @@ def is_yolo(obj: object) -> bool: def detect_accelerator(obj: object) -> str: try: - inner = getattr(obj, "model", None) + inner = getattr(obj, "model", None) # the torch module inside YOLO if inner is not None: - try: - first_param = next(inner.parameters()) - device_type = str(getattr(first_param.device, "type", "")) - if device_type: - return device_type - except StopIteration: - pass + return torch_module_accelerator(inner) or accelerators.CPU except Exception as exc: debug_ultralytics_failure("accelerator detection", exc) - return "cpu" + return accelerators.CPU def detect_quantization(obj: object) -> str | None: @@ -155,7 +164,7 @@ def detection_output_meta( total_preds += len(confs) all_confs.extend(confs) - for i, (conf, cls_id) in enumerate(zip(confs, classes)): + for i, (conf, cls_id) in enumerate(zip(confs, classes, strict=False)): if len(top_preds) >= 5: break label = (names or {}).get(int(cls_id), str(int(cls_id))) @@ -199,7 +208,7 @@ def classify_output_meta( top5conf = ( probs.top5conf.tolist() if hasattr(probs.top5conf, "tolist") else [] ) - for cls_id, conf in zip(top5, top5conf): + for cls_id, conf in zip(top5, top5conf, strict=False): label = (names or {}).get(int(cls_id), str(int(cls_id))) top_preds.append( TopKPrediction(label=label, confidence=round(float(conf), 4)) @@ -228,12 +237,10 @@ def weights_file_exists(model_arg: object) -> bool: if p.is_file(): return True # ultralytics resolves short names (e.g. "yolov8n.pt") against its assets dir - if ( + return ( _ULTRALYTICS_WEIGHTS_DIR is not None and (_ULTRALYTICS_WEIGHTS_DIR / p.name).is_file() - ): - return True - return False + ) def build_download_record(obj: object, load_ms: int) -> dict | None: @@ -316,7 +323,7 @@ def patched_call(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] except Exception as exc: handle.track_error( error_code="UNKNOWN", - error_message=str(exc)[: constants.ERROR_MSG_MAX_LEN], + error_message=str(exc), ) raise @@ -379,9 +386,11 @@ def memory_bytes(self, obj: object) -> int | None: def install_hooks(self, obj: object, handle: ModelHandle) -> None: handle.detected_accelerator = detect_accelerator(obj) setattr(obj, YOLO_HANDLE_ATTR, handle) + # YOLO.__call__ delegates to predict; patch their common entry point. patch_instance_call_once( obj, - patch_name=YOLO_CALL_PATCH_NAME, + patch_name=YOLO_PREDICT_PATCH_NAME, + method="predict", make_patched_call=build_patched_call, ) @@ -395,6 +404,14 @@ def install_auto_load_patch(cls, client_ref: object) -> None: is emitted with the file size and the total load duration. """ global _ultralytics_patched + if _unsupported_ultralytics_version is not None: + logger.warning( + "wildedge: ultralytics %s is older than %s, the oldest supported " + "release; the ultralytics integration is disabled", + _unsupported_ultralytics_version, + MIN_ULTRALYTICS_VERSION, + ) + return if _ultralytics_patched or _ultralytics is None: return @@ -433,3 +450,7 @@ def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] patched_init.__wildedge_original_call__ = original_init # type: ignore[attr-defined] _ultralytics.YOLO.__init__ = patched_init _ultralytics_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = UltralyticsExtractor.install_auto_load_patch diff --git a/wildedge/integrations/whisper.py b/wildedge/integrations/whisper.py new file mode 100644 index 0000000..55d5ba2 --- /dev/null +++ b/wildedge/integrations/whisper.py @@ -0,0 +1,191 @@ +"""Shared metadata tracking for local Whisper runtimes.""" + +from __future__ import annotations + +import wave +from dataclasses import dataclass +from pathlib import Path + +from wildedge.events.inference import AudioInputMeta +from wildedge.integrations.base import BaseExtractor +from wildedge.integrations.common import ( + hf_repo_from_cache_path, + stream_error_code, +) +from wildedge.model import ModelHandle, ModelInfo + +# Backends, used in model IDs and as the integration names. +FASTER_WHISPER = "faster_whisper" +MLX_WHISPER = "mlx_whisper" +MODEL_FORMATS = {FASTER_WHISPER: "ctranslate2", MLX_WHISPER: "mlx"} + +SOURCE_LOCAL = "local" +SOURCE_HUGGINGFACE = "huggingface" + +DESCRIPTOR_ATTR = "__wildedge_whisper_descriptor__" +HANDLE_ATTR = "__wildedge_whisper_handle__" +CLIENT_ATTR = "__wildedge_whisper_client__" + +# Both libraries expect 1-D sample arrays at this rate. +SAMPLE_RATE = 16000 + +# Languages written without spaces between words, where splitting on +# whitespace does not give a word count. +UNSPACED_LANGUAGES = frozenset({"bo", "ja", "km", "lo", "my", "th", "yue", "zh"}) + + +@dataclass +class WhisperDescriptor: + """Model identity, worked out once when the model is loaded.""" + + name: str + source: str + backend: str + accelerator: str + quantization: str | None = None + + @classmethod + def from_path( + cls, + path: object, + backend: str, + accelerator: str, + quantization: str | None = None, + ) -> WhisperDescriptor: + name, source = model_name_and_source(str(path)) + return cls(name, source, backend, accelerator, quantization) + + @property + def model_id(self) -> str: + return f"{self.backend}:{self.name}" + + @property + def model_format(self) -> str: + return MODEL_FORMATS[self.backend] + + +def model_name_and_source(path: str) -> tuple[str, str]: + """Name a model by its hub repo, its local folder, or the ID as given. + + A Hugging Face cache path (``.../models--org--name/snapshots/``) is + named ``org/name`` rather than after the snapshot hash. + """ + repo = hf_repo_from_cache_path(path) + if repo is not None: + return repo, SOURCE_HUGGINGFACE + local = Path(path).expanduser() + if local.exists(): + return local.resolve().name, SOURCE_LOCAL + return path, SOURCE_HUGGINGFACE + + +def descriptor(obj: object) -> WhisperDescriptor | None: + return ( + obj + if isinstance(obj, WhisperDescriptor) + else getattr(obj, DESCRIPTOR_ATTR, None) + ) + + +def report_load(client, model: object, load_ms: int) -> None: + """Record a model load. Whisper loads always read the weights again, so a + model ID seen before is still a new load, for example when MLX switches + back to an earlier model.""" + client._on_model_auto_loaded(model, load_ms=load_ms, reload=True) + + +def report_cached_load(client, model: object) -> None: + """Record a model that was already in memory before this client existed.""" + client._on_model_auto_loaded( + model, load_ms=0, load_kwargs={"context": {"model_cache_hit": True}} + ) + + +class WhisperExtractor(BaseExtractor): + """Must come before MlxExtractor in ``EXTRACTOR_SPECS``, which + would otherwise claim MLX Whisper models as plain MLX models.""" + + def can_handle(self, obj: object) -> bool: + return isinstance(descriptor(obj), WhisperDescriptor) + + def extract_info(self, obj: object, overrides: dict) -> tuple[str, ModelInfo]: + info = descriptor(obj) + assert info is not None + return overrides.pop("id", None) or info.model_id, ModelInfo( + model_name=info.name, + model_version=overrides.pop("version", "unknown"), + model_source=overrides.pop("source", info.source), + model_format=overrides.pop("format", info.model_format), + model_family=overrides.pop("family", "whisper"), + quantization=overrides.pop("quantization", None) or info.quantization, + ) + + def install_hooks(self, obj: object, handle: ModelHandle) -> None: + info = descriptor(obj) + assert info is not None + handle.detected_accelerator = info.accelerator + + +def is_sample_array(audio: object) -> bool: + return getattr(audio, "ndim", None) == 1 + + +def audio_meta(audio: object, audio_duration_s: float | None = None) -> AudioInputMeta: + meta = AudioInputMeta() + if isinstance(audio, (str, Path)): + try: + with wave.open(str(audio), "rb") as wav: + meta.duration_ms = round(wav.getnframes() / wav.getframerate() * 1000) + meta.sample_rate = wav.getframerate() + meta.channels = wav.getnchannels() + meta.bit_depth = wav.getsampwidth() * 8 + meta.format = "wav" + except (OSError, EOFError, wave.Error): + pass + elif is_sample_array(audio): + meta.duration_ms = round(audio.shape[0] / SAMPLE_RATE * 1000) + meta.sample_rate = SAMPLE_RATE + meta.channels = 1 + if audio_duration_s is not None: + meta.duration_ms = round(audio_duration_s * 1000) + return meta + + +def track_transcription( + handle: ModelHandle, + audio: object, + text: str, + *, + duration_ms: int, + language: str | None = None, + language_probability: float | None = None, + audio_duration_s: float | None = None, + word_count: int | None = None, + first_segment_ms: int | None = None, + outcome: str = "completed", +) -> None: + """Record one transcription. + + ``duration_ms`` is how long the transcription took; ``audio_duration_s`` + is the length of the audio when the runtime reports it. + """ + context: dict = {} + if language not in UNSPACED_LANGUAGES: + context["transcript_word_count"] = ( + len(text.split()) if word_count is None else word_count + ) + if language is not None: + context["language"] = language + if language_probability is not None: + context["language_confidence"] = language_probability + if first_segment_ms is not None: + context["time_to_first_segment_ms"] = first_segment_ms + handle.track_inference( + duration_ms=duration_ms, + input_modality="audio", + output_modality="text", + input_meta=audio_meta(audio, audio_duration_s), + context=context, + success=outcome == "completed", + error_code=stream_error_code(outcome), + ) diff --git a/wildedge/llm_api.py b/wildedge/llm_api.py index 47ff17f..7b0ca24 100644 --- a/wildedge/llm_api.py +++ b/wildedge/llm_api.py @@ -22,6 +22,7 @@ import time from typing import TYPE_CHECKING, Any +from wildedge.defaults import get_client from wildedge.events.inference import ApiMeta, GenerationOutputMeta, TextInputMeta from wildedge.integrations.common import ( build_input_meta, @@ -179,7 +180,6 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool: return self.__exit__(exc_type, exc_val, exc_tb) def _handle(self) -> ModelHandle | None: - from wildedge.defaults import get_client # noqa: PLC0415 (import cycle) try: return get_client().register_model( diff --git a/wildedge/loading.py b/wildedge/loading.py new file mode 100644 index 0000000..e2328aa --- /dev/null +++ b/wildedge/loading.py @@ -0,0 +1,37 @@ +"""Import optional parts of the SDK only when a feature uses them. + +Importing ``wildedge`` must not import ML frameworks or API clients. The +registries name each adapter by import path instead, and these helpers load it +on first use. +""" + +from __future__ import annotations + +from functools import cache +from importlib import import_module +from types import ModuleType +from typing import Any + + +@cache +def resolve(target: str) -> Any: + """Import ``package.module:name`` and return that attribute.""" + module, name = target.split(":", 1) + return getattr(import_module(module), name) + + +def lazy_callable(target: str): + """A function that imports ``target`` on its first call and forwards to it.""" + + def call(*args, **kwargs): + return resolve(target)(*args, **kwargs) + + return call + + +def optional_module(name: str) -> ModuleType | None: + """Import an optional third-party runtime, or None if it is not installed.""" + try: + return import_module(name) + except ImportError: + return None diff --git a/wildedge/model.py b/wildedge/model.py index bfbcab8..1c86fd6 100644 --- a/wildedge/model.py +++ b/wildedge/model.py @@ -7,6 +7,7 @@ from pathlib import Path from typing import Any +from wildedge import accelerators from wildedge.attachments import Attachment from wildedge.events import ( ApiMeta, @@ -68,7 +69,7 @@ def __init__( self.info = info self.publish = publish self.capture_attachments = capture_attachments - self.detected_accelerator: str = "cpu" + self.detected_accelerator: str = accelerators.CPU self.last_inference_id: str | None = None def track_load( @@ -405,17 +406,22 @@ def register( capture_attachments: Callable[[list[Attachment], str, datetime], list[dict]] | None = None, ) -> tuple[ModelHandle, bool]: - """Return (handle, is_new). is_new=False means already registered; skip install_hooks.""" + """Return (handle, is_new). is_new=False means already registered; skip install_hooks. + + The first registration in a process sets the model's info, replacing any + info persisted by an earlier process; the file only covers models that + queued events still reference but that are not registered again. + """ if model_id in self.handles: return self.handles[model_id], False - if model_id not in self.models: + if self.models.get(model_id) != info: self.models[model_id] = info self.save_to_disk() handle = ModelHandle( model_id=model_id, - info=self.models[model_id], + info=info, publish=publish, capture_attachments=capture_attachments, ) diff --git a/wildedge/platforms/base.py b/wildedge/platforms/base.py index 62df959..7ebefd5 100644 --- a/wildedge/platforms/base.py +++ b/wildedge/platforms/base.py @@ -11,7 +11,7 @@ from datetime import datetime from pathlib import Path -from wildedge import constants +from wildedge import accelerators, constants from wildedge.logging import logger from wildedge.platforms.device_info import DeviceInfo from wildedge.platforms.hardware import HardwareContext, ThermalContext @@ -113,7 +113,7 @@ def detect_device( cpu_cores=os.cpu_count(), ram_total_bytes=self.ram_bytes(), disk_total_bytes=self.disk_bytes(), - accelerators=["cpu", *gpu_accs], + accelerators=[accelerators.CPU, *gpu_accs], gpu_name=gpu_name, ) if overrides: diff --git a/wildedge/platforms/linux.py b/wildedge/platforms/linux.py index c189e4a..13e70ba 100644 --- a/wildedge/platforms/linux.py +++ b/wildedge/platforms/linux.py @@ -5,6 +5,7 @@ import shutil from pathlib import Path +from wildedge import accelerators from wildedge.platforms.base import ( Platform, cuda_device_count, @@ -91,18 +92,18 @@ def disk_bytes(self) -> int | None: return None def gpu_accelerators(self) -> tuple[list[str], str | None]: - accelerators = [] + found = [] gpu_name = None if cuda_device_count("libcuda.so.1") > 0: - accelerators.append("cuda") + found.append(accelerators.CUDA) gpu_name = nvml_gpu_name("libnvidia-ml.so.1") if hip_device_count("libamdhip64.so") > 0: - accelerators.append("rocm") - return accelerators, gpu_name + found.append(accelerators.ROCM) + return found, gpu_name def gpu_accelerator_for_offload(self) -> str: accs, _ = self.gpu_accelerators() - return accs[0] if accs else "cpu" + return accs[0] if accs else accelerators.CPU def cpu_freq(self) -> tuple[int | None, int | None]: """Read current and max CPU frequency (MHz) from cpufreq for cpu0.""" diff --git a/wildedge/platforms/macos.py b/wildedge/platforms/macos.py index dc86852..2321356 100644 --- a/wildedge/platforms/macos.py +++ b/wildedge/platforms/macos.py @@ -5,6 +5,7 @@ import shutil from pathlib import Path +from wildedge import accelerators from wildedge.platforms.base import Platform, debug_detection_failure from wildedge.platforms.hardware import ThermalContext @@ -21,7 +22,7 @@ def _get_libc() -> ctypes.CDLL: CF_STRING_ENCODING_UTF8 = 0x08000100 CF_NUMBER_SINT32_TYPE = 3 -# NSProcessInfoThermalState values (0–3) +# NSProcessInfoThermalState values (0-3) MACOS_THERMAL_STATES: dict[int, tuple[str, str]] = { 0: ("nominal", "nominal"), 1: ("fair", "fair"), @@ -111,11 +112,11 @@ def os_version(self) -> str | None: def gpu_accelerators(self) -> tuple[list[str], str | None]: if platform.machine() == "arm64": - return ["mps"], None + return [accelerators.MPS], None return [], None def gpu_accelerator_for_offload(self) -> str: - return "mps" if platform.machine() == "arm64" else "cpu" + return accelerators.MPS if platform.machine() == "arm64" else accelerators.CPU def cpu_freq(self) -> tuple[int | None, int | None]: """Read CPU frequency via sysctl. diff --git a/wildedge/platforms/unknown.py b/wildedge/platforms/unknown.py index 973dd4e..f89aecc 100644 --- a/wildedge/platforms/unknown.py +++ b/wildedge/platforms/unknown.py @@ -5,6 +5,7 @@ import sys from pathlib import Path +from wildedge import accelerators from wildedge.platforms.base import Platform, debug_detection_failure @@ -41,4 +42,4 @@ def gpu_accelerators(self) -> tuple[list[str], str | None]: return [], None def gpu_accelerator_for_offload(self) -> str: - return "cpu" + return accelerators.CPU diff --git a/wildedge/platforms/windows.py b/wildedge/platforms/windows.py index 7748816..7aca452 100644 --- a/wildedge/platforms/windows.py +++ b/wildedge/platforms/windows.py @@ -6,6 +6,7 @@ import shutil from pathlib import Path +from wildedge import accelerators from wildedge.platforms.base import ( Platform, cuda_device_count, @@ -91,7 +92,7 @@ def meminfo(self) -> tuple[int | None, int | None]: def disk_bytes(self) -> int | None: try: - drive = os.environ.get("SystemDrive", "C:\\") + drive = os.environ.get("SYSTEMDRIVE", "C:\\") return shutil.disk_usage(drive).total except Exception as exc: debug_detection_failure("windows disk_bytes", exc) @@ -99,12 +100,12 @@ def disk_bytes(self) -> int | None: def gpu_accelerators(self) -> tuple[list[str], str | None]: if cuda_device_count("nvcuda.dll") > 0: - return ["cuda"], nvml_gpu_name("nvml.dll") + return [accelerators.CUDA], nvml_gpu_name("nvml.dll") return [], None def gpu_accelerator_for_offload(self) -> str: accs, _ = self.gpu_accelerators() - return accs[0] if accs else "cpu" + return accs[0] if accs else accelerators.CPU def battery(self) -> tuple[float | None, bool | None]: """Read battery level and charging state via GetSystemPowerStatus.""" diff --git a/wildedge/quantizations.py b/wildedge/quantizations.py new file mode 100644 index 0000000..1d3cf90 --- /dev/null +++ b/wildedge/quantizations.py @@ -0,0 +1,17 @@ +"""Quantization labels reported in model metadata. + +Common weight precisions use the names below. Format-specific schemes pass +through as their runtime names them, such as GGUF's ``q4_k_m`` or +bitsandbytes' ``nf4``, so this list is not closed. +""" + +F32 = "f32" +F16 = "f16" +BF16 = "bf16" +INT8 = "int8" +QUANTIZED = "quantized" # quantized, with an unknown bit width + + +def int_bits(bits: int) -> str: + """Label for integer weights of the given width, such as ``int4``.""" + return f"int{int(bits)}" diff --git a/wildedge/runtime/bootstrap.py b/wildedge/runtime/bootstrap.py index 63e0ee4..308b9b6 100644 --- a/wildedge/runtime/bootstrap.py +++ b/wildedge/runtime/bootstrap.py @@ -16,6 +16,7 @@ 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.instrumentation import set_strict_ownership from wildedge.integrations.registry import INTEGRATIONS_BY_NAME, supported_integrations from wildedge.settings import read_runtime_env @@ -208,6 +209,10 @@ def install_runtime(*, install_signal_handlers: bool = True) -> RuntimeContext: file=os.sys.stderr, ) + # The app may call init() with its own DSN later. Under + # --strict-integrations an ownership conflict then raises instead of warning. + set_strict_ownership(env.strict_integrations) + if env.strict_integrations: failures = [row for row in statuses if row["status"] in STRICT_FAILURE_STATUSES] if failures: diff --git a/wildedge/settings.py b/wildedge/settings.py index 6c445b8..855e9e1 100644 --- a/wildedge/settings.py +++ b/wildedge/settings.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from wildedge import constants +from wildedge.integrations.registry import canonical_integration TRUE_VALUES = {"1", "true", "yes", "on"} @@ -46,7 +47,13 @@ def parse_integration_list(value: str | None, all_values: list[str]) -> list[str """Parse ``--integrations`` value; defaults to all when unset.""" if not value or value == "all": return sorted(all_values) - return [item.strip() for item in value.split(",") if item.strip()] + return list( + dict.fromkeys( + canonical_integration(item.strip()) + for item in value.split(",") + if item.strip() + ) + ) def parse_hub_list(value: str | None, all_values: list[str]) -> list[str]: diff --git a/wildedge/transmitter.py b/wildedge/transmitter.py index f31c88d..de74bed 100644 --- a/wildedge/transmitter.py +++ b/wildedge/transmitter.py @@ -15,7 +15,7 @@ class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): """Raise immediately on any redirect instead of following it.""" - def redirect_request(self, req, fp, code, msg, headers, newurl): # noqa: PLR0913 + def redirect_request(self, req, fp, code, msg, headers, newurl): raise urllib.error.HTTPError(newurl, code, msg, headers, fp) @@ -78,7 +78,7 @@ def send(self, batch: dict) -> IngestResponse: server_time=data.get("server_time"), rejected=data.get("rejected"), ) - elif status_code == 400: + if status_code == 400: logger.warning( "wildedge: batch rejected (400) - discarding: %s", raw[: constants.ERROR_MSG_MAX_LEN], @@ -89,7 +89,7 @@ def send(self, batch: dict) -> IngestResponse: events_accepted=0, events_rejected=len(batch.get("events", [])), ) - elif status_code == 401: + if status_code == 401: logger.error("wildedge: authentication failed (401) - check your API key") return IngestResponse( status="unauthorized", @@ -97,7 +97,7 @@ def send(self, batch: dict) -> IngestResponse: events_accepted=0, events_rejected=len(batch.get("events", [])), ) - elif 300 <= status_code < 400: + if 300 <= status_code < 400: # Redirects should never occur (we disable redirect following). # Treat as a permanent config error so we don't loop. logger.error( @@ -111,7 +111,7 @@ def send(self, batch: dict) -> IngestResponse: events_accepted=0, events_rejected=len(batch.get("events", [])), ) - elif status_code == 404: + if status_code == 404: logger.error( "wildedge: endpoint not found (404) at %s; check WILDEDGE_DSN", url ) @@ -121,11 +121,11 @@ def send(self, batch: dict) -> IngestResponse: events_accepted=0, events_rejected=len(batch.get("events", [])), ) - elif status_code == 429 or status_code >= 500: + if status_code == 429 or status_code >= 500: raise TransmitError( f"HTTP {status_code}: {raw[: constants.ERROR_MSG_MAX_LEN]!r}" ) - elif 400 <= status_code < 500: + if 400 <= status_code < 500: # Other 4xx (e.g. 422 Unprocessable) are permanent client errors; discard. logger.warning( "wildedge: batch rejected (%d) - discarding: %s", @@ -138,11 +138,10 @@ def send(self, batch: dict) -> IngestResponse: events_accepted=0, events_rejected=len(batch.get("events", [])), ) - else: - # Guard against unexpected status codes that would otherwise return None. - raise TransmitError( - f"Unexpected HTTP {status_code}: {raw[: constants.ERROR_MSG_MAX_LEN]!r}" - ) + # Guard against unexpected status codes that would otherwise return None. + raise TransmitError( + f"Unexpected HTTP {status_code}: {raw[: constants.ERROR_MSG_MAX_LEN]!r}" + ) def close(self) -> None: pass