unsloth/studio/backend/tests/test_api_monitor.py
alkinun 14188d6f45
Add API server monitor in Studio (#5558)
* Add Studio API activity monitor

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Mark disconnected monitor streams cancelled

* Fix monitor lifecycle, parsing, and clock issues for PR #5558

Backend:
- Finalize the monitor entry as error on the chat-completions validation
  reject paths (addresses the codex P2 comment around line 2337). Adds a
  small _reject helper next to the api_monitor.start so each early-raise
  in the GGUF tool-passthrough, non-empty-messages, GGUF vision/audio,
  PIL image decode, and text-only model branches no longer leaves the
  entry stuck running until eviction.
- _monitor_openai_chunk is now defensive about malformed shapes
  (non-list choices, non-dict choice/delta) so a misbehaving provider
  chunk does not raise into the streaming generator and abort the
  user's response.
- _monitor_openai_sse_line accepts both data:value and data: value per
  SSE spec; previously the single-space-only prefix silently dropped
  tokens from compact emitters.
- openai_completions stream now keeps a residual buffer between
  aiter_bytes chunks so a data: line whose newline lands in the next
  TCP frame is reconstructed before parsing, instead of being dropped
  by the per-chunk splitlines call.

api_monitor:
- duration_ms is derived from time.monotonic anchors and clamped at 0,
  so NTP / manual clock steps no longer produce negative durations.
- finish() and fail() are idempotent: a second call (for example [DONE]
  arriving after the generator's finally block already ran) no longer
  moves finished_at or finished_monotonic.
- set_usage only derives total_tokens when no authoritative total has
  been recorded, so a later partial-usage chunk does not clobber a
  provider-reported total from an earlier chunk.
- _trim guards limit < 3 so the slice cannot underflow.

Tests:
- Adds regression coverage for the new idempotency, total preservation,
  monotonic clock, and _trim guard behaviour.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Mark monitored stream failures correctly

* Fix completions monitor failures

* Fix responses stream monitor cleanup

* Fix audio input API monitor lifecycle

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix API monitor review follow-ups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix completions monitor usage accounting

* Fix API monitor cancellation and usage gaps

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Move API monitor into settings

* Fix API monitor review followups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Lazy load API monitor details

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix API monitor review issues

* Finalize passthrough monitor on clean EOF

* Finalize monitor on chat validation rejects

* Fix API monitor review follow-ups

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix API monitor follow-up review issues

* Monitor embeddings and tool-call replies

* Cover tool-call monitor replies and cancellations

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Scope API monitor entries by subject

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Finalize monitor entries on cancellation

* fix: address PR 5558 CI failures

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: update provider proxy test stub

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: imagineer99 <samleejackson0@gmail.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
2026-06-16 20:26:13 +01:00

222 lines
6.5 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from core.inference.api_monitor import ApiMonitor, _trim
def test_api_monitor_tracks_reply_usage_and_context():
monitor = ApiMonitor(max_entries = 3)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "local-model",
prompt = "user: hello",
context_length = 100,
)
monitor.append_reply(entry_id, "hi")
monitor.append_reply(entry_id, " there")
monitor.set_usage(
entry_id,
prompt_tokens = 4,
completion_tokens = 6,
)
monitor.finish(entry_id)
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
assert entry["reply"] == "hi there"
assert entry["total_tokens"] == 10
assert entry["context_usage"] == 0.1
assert entry["duration_ms"] is not None
def test_api_monitor_summary_omits_full_prompt_and_reply():
monitor = ApiMonitor(max_entries = 3)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "local-model",
prompt = "p" * 500,
)
monitor.set_reply(entry_id, "r" * 500)
[summary] = monitor.snapshot(include_details = False)
assert "prompt" not in summary
assert "reply" not in summary
assert summary["prompt_preview"].endswith("...")
assert summary["reply_preview"].endswith("...")
assert summary["prompt_truncated"] is True
assert summary["reply_truncated"] is True
detail = monitor.get(entry_id)
assert detail is not None
assert detail["prompt"] == "p" * 500
assert detail["reply"] == "r" * 500
def test_api_monitor_filters_entries_by_subject():
monitor = ApiMonitor(max_entries = 3)
alice = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "alice prompt",
subject = "alice",
)
bob = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "bob prompt",
subject = "bob",
)
monitor.finish(bob)
alice_entries = monitor.snapshot(subject = "alice")
assert [entry["id"] for entry in alice_entries] == [alice]
assert monitor.get(bob, subject = "alice") is None
assert monitor.get(bob, subject = "bob")["id"] == bob
assert monitor.active_count(subject = "alice") == 1
assert monitor.active_count(subject = "bob") == 0
def test_api_monitor_keeps_bounded_recent_history():
monitor = ApiMonitor(max_entries = 2)
first = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "first",
)
second = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "second",
)
third = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "third",
)
monitor.finish(first)
monitor.finish(second)
monitor.finish(third)
entries = monitor.snapshot()
ids = [entry["id"] for entry in entries]
assert ids[0] == third
assert [entry["prompt"] for entry in entries] == ["third", "second"]
assert first not in ids
assert monitor.active_count() == 0
def test_api_monitor_keeps_running_entries_beyond_history_limit():
monitor = ApiMonitor(max_entries = 1)
running = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "running",
)
for prompt in ("done-1", "done-2", "done-3"):
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = prompt,
)
monitor.finish(entry_id)
entries = monitor.snapshot()
ids = [entry["id"] for entry in entries]
assert running in ids
assert monitor.active_count() == 1
monitor.finish(running)
[entry] = monitor.snapshot()
assert entry["id"] == running
assert entry["status"] == "completed"
assert monitor.active_count() == 0
def test_api_monitor_finish_is_idempotent():
monitor = ApiMonitor(max_entries = 2)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
)
monitor.finish(entry_id)
first = monitor.snapshot()[0]
monitor.finish(entry_id)
second = monitor.snapshot()[0]
assert first["finished_at"] == second["finished_at"]
assert first["duration_ms"] == second["duration_ms"]
def test_api_monitor_preserves_authoritative_total_tokens():
monitor = ApiMonitor(max_entries = 2)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
)
monitor.set_usage(
entry_id,
prompt_tokens = 10,
completion_tokens = 20,
total_tokens = 33,
)
# A later partial chunk omitting `total_tokens` must not clobber 33.
monitor.set_usage(entry_id, prompt_tokens = 11)
assert monitor.snapshot()[0]["total_tokens"] == 33
def test_api_monitor_recomputes_derived_total_tokens():
monitor = ApiMonitor(max_entries = 2)
entry_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "m",
prompt = "hi",
)
monitor.set_usage(entry_id, prompt_tokens = 10)
assert monitor.snapshot()[0]["total_tokens"] == 10
monitor.set_usage(entry_id, completion_tokens = 20)
entry = monitor.snapshot()[0]
assert entry["prompt_tokens"] == 10
assert entry["completion_tokens"] == 20
assert entry["total_tokens"] == 30
def test_api_monitor_duration_non_negative_under_clock_step(monkeypatch):
import core.inference.api_monitor as m
fake_now = [1000.0]
monkeypatch.setattr(m.time, "time", lambda: fake_now[0])
monitor = ApiMonitor(max_entries = 1)
entry_id = monitor.start(
endpoint = "/x",
method = "POST",
model = "m",
prompt = "hi",
)
fake_now[0] = 500.0
monitor.finish(entry_id)
assert monitor.snapshot()[0]["duration_ms"] >= 0
def test_api_monitor_trim_guards_tiny_limit():
assert _trim("abcdefgh", 2) == ".."
assert _trim("abcdefgh", 0) == ""
assert _trim("abcdefgh", 3) == "..."
assert _trim("abcdefgh", 4) == "a..."
assert _trim("abcdefgh", 100) == "abcdefgh"