Handle missing streamed token metrics
This commit is contained in:
@@ -310,7 +310,7 @@ def stream_chat_completion(
|
|||||||
return StreamMetrics(
|
return StreamMetrics(
|
||||||
ttft_ms=ttft_ms,
|
ttft_ms=ttft_ms,
|
||||||
tpot_ms=tpot_ms,
|
tpot_ms=tpot_ms,
|
||||||
completion_tokens=used_tokens if used_tokens > 0 else None,
|
completion_tokens=used_tokens if used_tokens is not None and used_tokens > 0 else None,
|
||||||
completion_tokens_source=completion_tokens_source,
|
completion_tokens_source=completion_tokens_source,
|
||||||
streamed_chunk_count=chunk_token_count,
|
streamed_chunk_count=chunk_token_count,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -8,6 +8,7 @@ import statistics
|
|||||||
import subprocess
|
import subprocess
|
||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
|
import traceback
|
||||||
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
|
from concurrent.futures import FIRST_COMPLETED, Future, ThreadPoolExecutor, wait
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
@@ -596,6 +597,7 @@ def run_trial(trial_spec_path: Path) -> dict[str, Any]:
|
|||||||
"best_request_count": None,
|
"best_request_count": None,
|
||||||
"failure_stage": failure_stage,
|
"failure_stage": failure_stage,
|
||||||
"failure_reason": str(exc),
|
"failure_reason": str(exc),
|
||||||
|
"failure_traceback": traceback.format_exc(),
|
||||||
"probes": probe_history,
|
"probes": probe_history,
|
||||||
}
|
}
|
||||||
StudyStore.write_json(Path(trial.result_path), result)
|
StudyStore.write_json(Path(trial.result_path), result)
|
||||||
|
|||||||
@@ -11,7 +11,13 @@ from unittest import mock
|
|||||||
from aituner.cli import main as cli_main
|
from aituner.cli import main as cli_main
|
||||||
from aituner.compare import _aggregate_summary, load_compare_spec, run_compare
|
from aituner.compare import _aggregate_summary, load_compare_spec, run_compare
|
||||||
from aituner.engine import build_launch_recipe
|
from aituner.engine import build_launch_recipe
|
||||||
from aituner.http_client import StreamMetrics, _auth_headers, _openai_url, _should_bypass_proxy
|
from aituner.http_client import (
|
||||||
|
StreamMetrics,
|
||||||
|
_auth_headers,
|
||||||
|
_openai_url,
|
||||||
|
_should_bypass_proxy,
|
||||||
|
stream_chat_completion,
|
||||||
|
)
|
||||||
from aituner.job import append_job, build_trial_job
|
from aituner.job import append_job, build_trial_job
|
||||||
from aituner.harness import (
|
from aituner.harness import (
|
||||||
build_harness_context,
|
build_harness_context,
|
||||||
@@ -3767,6 +3773,29 @@ class CoreFlowTests(unittest.TestCase):
|
|||||||
"http://example.com/v1/chat/completions",
|
"http://example.com/v1/chat/completions",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_stream_chat_completion_handles_missing_usage_and_chunks(self) -> None:
|
||||||
|
class FakeResponse:
|
||||||
|
def __enter__(self):
|
||||||
|
return self
|
||||||
|
|
||||||
|
def __exit__(self, exc_type, exc, traceback):
|
||||||
|
return False
|
||||||
|
|
||||||
|
def __iter__(self):
|
||||||
|
return iter([b"data: {\"choices\": []}\n", b"data: [DONE]\n"])
|
||||||
|
|
||||||
|
with mock.patch("aituner.http_client._urlopen", return_value=FakeResponse()):
|
||||||
|
metrics = stream_chat_completion(
|
||||||
|
base_url="http://127.0.0.1:8000",
|
||||||
|
body={"model": "m", "messages": [{"role": "user", "content": "x"}]},
|
||||||
|
timeout_s=1.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertIsNone(metrics.ttft_ms)
|
||||||
|
self.assertIsNone(metrics.tpot_ms)
|
||||||
|
self.assertIsNone(metrics.completion_tokens)
|
||||||
|
self.assertEqual(metrics.completion_tokens_source, "none")
|
||||||
|
|
||||||
def test_loopback_urls_bypass_proxy(self) -> None:
|
def test_loopback_urls_bypass_proxy(self) -> None:
|
||||||
self.assertTrue(_should_bypass_proxy("http://127.0.0.1:8000/v1/models"))
|
self.assertTrue(_should_bypass_proxy("http://127.0.0.1:8000/v1/models"))
|
||||||
self.assertTrue(_should_bypass_proxy("http://localhost:8000/health"))
|
self.assertTrue(_should_bypass_proxy("http://localhost:8000/health"))
|
||||||
|
|||||||
Reference in New Issue
Block a user