diff --git a/composer/diagnostics/timing.py b/composer/diagnostics/timing.py index dd3c7704..a523c9cc 100644 --- a/composer/diagnostics/timing.py +++ b/composer/diagnostics/timing.py @@ -16,16 +16,18 @@ from typing import AsyncIterator, Iterable, Protocol import uuid -from graphcore.utils import TokenUsageDict +from graphcore.utils import NormalizedTokenUsage @dataclass(frozen=True) class TokenTotals: - """Raw LLM token counts accumulated across one or more calls. + """LLM token counts accumulated across one or more calls. ``input`` is fresh + (uncached) input only; the two cache buckets are counted separately. - ``from_dict`` builds one from a ``graphcore.utils.TokenUsageDict`` (its - ``input_tokens`` / ``output_tokens`` / ``cache_read_input_tokens`` / - ``cache_creation_input_tokens`` keys). + ``from_normalized`` builds one from a ``graphcore.utils.NormalizedTokenUsage``, + whose ``total_input_tokens`` *includes* both cache buckets — they are subtracted + back out so the stored fields (and everything serialized from them) keep the + same meaning they always had. """ input: int = 0 output: int = 0 @@ -45,12 +47,14 @@ def __bool__(self) -> bool: return self.input > 0 or self.output > 0 or self.cache_read > 0 or self.cache_write > 0 @classmethod - def from_dict(cls, u: TokenUsageDict) -> "TokenTotals": + def from_normalized(cls, u: NormalizedTokenUsage) -> "TokenTotals": + cache_read = u["cache_read_tokens"] + cache_write = u["cache_write_tokens"] return TokenTotals( - input=u["input_tokens"], - output=u["output_tokens"], - cache_read=u["cache_read_input_tokens"], - cache_write=u["cache_creation_input_tokens"], + input=max(0, u["total_input_tokens"] - cache_read - cache_write), + output=u["total_output_tokens"], + cache_read=cache_read, + cache_write=cache_write, ) def as_dict(self) -> dict[str, int]: @@ -145,12 +149,12 @@ def record_prover_runtime(self, ms: int, *, task_id: str | None = None) -> None: self._active_prover_reported_by_task.get(task_id, 0) + ms ) - def record_token_usage(self, usage: TokenUsageDict, *, task_id: str | None = None) -> None: + def record_token_usage(self, usage: NormalizedTokenUsage, *, task_id: str | None = None) -> None: """Accumulate one LLM call's token counts into the run-wide per-model totals and (if a task is active) into that task's in-flight bucket, later folded into its ``PhaseRecord`` by ``record_phase``. Defaults attribution to the active task.""" model = usage.get("model_name") or "unknown" - update = TokenTotals.from_dict(usage) + update = TokenTotals.from_normalized(usage) self.token_usage_by_model[model] = self.token_usage_by_model.get(model, TokenTotals()) + update if (task_id := task_id or get_current_task_id()) is not None: bucket = self._active_tokens_by_task.setdefault(task_id, {}) diff --git a/composer/diagnostics/usage_callback.py b/composer/diagnostics/usage_callback.py index c03a1514..deb5c1c2 100644 --- a/composer/diagnostics/usage_callback.py +++ b/composer/diagnostics/usage_callback.py @@ -20,7 +20,7 @@ from langchain_core.messages import AIMessage from langchain_core.outputs import ChatGeneration, LLMResult -from graphcore.utils import get_token_usage +from graphcore.utils import get_normalized_token_usage from composer.diagnostics.timing import get_run_summary @@ -43,6 +43,9 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: return msg = generation.message if isinstance(msg, AIMessage): - # get_run_summary() returns an inert throwaway outside a run, so this is - # a no-op when no autoprove run is active (e.g. ad-hoc model use). - get_run_summary().record_token_usage(get_token_usage(msg)) + # The normalized usage_metadata is the one source every transport and + # provider fills in — the raw response_metadata["usage"] dict exists only + # on non-streamed Anthropic responses. get_run_summary() returns an inert + # throwaway outside a run, so this is a no-op when no autoprove run is + # active (e.g. ad-hoc model use). + get_run_summary().record_token_usage(get_normalized_token_usage(msg)) diff --git a/composer/llm/anthropic.py b/composer/llm/anthropic.py index b1101fcf..7b07bbfb 100644 --- a/composer/llm/anthropic.py +++ b/composer/llm/anthropic.py @@ -293,7 +293,16 @@ def builder_for( return ChatAnthropic( model_name=self.model_name, max_tokens_to_sample=opts.tokens, - timeout=None, + # An explicit None DISABLES the SDK's timeouts (None != not-given), so a + # socket that dies silently mid-stream hangs the session forever. A float + # is a per-phase httpx timeout — for a streamed response, the max silence + # between chunks, not a cap on the whole turn. + timeout=300.0, + # Stream every request: a long authoring turn (Opus + thinking on a large + # prompt) can exceed the SDK's 600s non-streaming ceiling, and a silent + # 10-minute wait is long enough for NAT/idle killers to drop the socket + # (surfaces as APIConnectionError mid-run). Streaming keeps bytes flowing. + streaming=True, max_retries=8, stop=None, betas=betas, diff --git a/composer/rag/db.py b/composer/rag/db.py index 0d1549fd..f691f2ec 100644 --- a/composer/rag/db.py +++ b/composer/rag/db.py @@ -3,6 +3,7 @@ from dataclasses import dataclass import asyncio import logging +import threading from abc import ABC, abstractmethod import os @@ -20,6 +21,7 @@ from composer.rag.types import ManualRef, BlockChunk, ManualSectionHit from composer.rag.text import code_ref_tag +from composer.rag.models import ENCODE_LOCK import sqlite3 @@ -96,19 +98,29 @@ async def get_manual_section(self, headers: list[str]) -> str | None: ... + # Encodes race when concurrent (see ENCODE_LOCK in composer.rag.models — the lock + # is shared with the sync DefaultEmbedder so neither path can race the other). + def _encode_query_locked(self, query: str) -> ndarray: + with ENCODE_LOCK: + return cast(ndarray, self.tr.encode_query(query, show_progress_bar=False)) + + def _encode_docs_locked(self, docs: list[str]) -> list[ndarray]: + with ENCODE_LOCK: + return cast(list[ndarray], self.tr.encode_document(docs, show_progress_bar=False)) + async def embed_query( self, query: str ) -> ndarray: - return cast(ndarray, await asyncio.to_thread( - self.tr.encode_query, f"search_query: {query}", show_progress_bar=False - )) + return await asyncio.to_thread( + self._encode_query_locked, f"search_query: {query}" + ) async def embed_docs( self, doc: list[BlockChunk] ) -> list[ndarray]: - return cast(list[ndarray], await asyncio.to_thread( - self.tr.encode_document, [f"search_document: {d.chunk}" for d in doc], show_progress_bar=False - )) + return await asyncio.to_thread( + self._encode_docs_locked, [f"search_document: {d.chunk}" for d in doc] + ) type RagConnection = str | AsyncConnectionPool[AsyncConnection[TupleRow]] diff --git a/composer/rag/models.py b/composer/rag/models.py index e5844ad4..e4a974ce 100644 --- a/composer/rag/models.py +++ b/composer/rag/models.py @@ -1,7 +1,16 @@ +import os +import threading from typing import TYPE_CHECKING, override from langchain_core.embeddings import Embeddings +# One process-wide gate for every sentence-transformer encode: the model's remote +# code caches positional tensors per sequence length, so concurrent encodes race +# (shape mismatches on CPU, SIGSEGV in torch's MPS shader cache). Shared between +# the sync Embeddings API here and the async ComposerRAGDB wrappers so the two +# paths cannot race each other on the same model instance either. +ENCODE_LOCK = threading.Lock() + # claim we always import ST if TYPE_CHECKING: from sentence_transformers import SentenceTransformer @@ -12,7 +21,14 @@ def get_model() -> SentenceTransformer: from sentence_transformers import SentenceTransformer #type: ignore def get_model() -> SentenceTransformer: - return SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True) + # COMPOSER_EMBED_DEVICE overrides the auto-picked device. torch's MPS shader + # cache is not thread-safe under concurrent encodes (SIGSEGV on Apple + # Silicon), so mac hosts should set it to "cpu". + return SentenceTransformer( + 'nomic-ai/nomic-embed-text-v1.5', + trust_remote_code=True, + device=os.environ.get("COMPOSER_EMBED_DEVICE"), + ) except ImportError: # for tests (no ST dependency) def get_model() -> "SentenceTransformer": @@ -25,12 +41,14 @@ def __init__(self, model: "SentenceTransformer | None" = None): @override def embed_documents(self, texts: list[str]) -> list[list[float]]: - return self.model.encode_document( - texts - ).tolist() #type: ignore + with ENCODE_LOCK: + return self.model.encode_document( + texts + ).tolist() #type: ignore @override def embed_query(self, text: str) -> list[float]: - return self.model.encode_query( - [text] - ).tolist()[0] #type: ignore + with ENCODE_LOCK: + return self.model.encode_query( + [text] + ).tolist()[0] #type: ignore diff --git a/composer/spec/source/autosetup.py b/composer/spec/source/autosetup.py index 338c064b..e3f4e2b7 100644 --- a/composer/spec/source/autosetup.py +++ b/composer/spec/source/autosetup.py @@ -18,7 +18,7 @@ import asyncio from composer.prover.core import ProverOptions -from graphcore.utils import TokenUsageDict +from graphcore.utils import NormalizedTokenUsage from composer.io.context import emit_custom_event # Locators for autosetup's on-disk usage files (certora_autosetup owns that layout). from certora_autosetup.utils.paths import ( @@ -218,22 +218,27 @@ def log_complete(self, returncode: int): # owns that on-disk layout and exposes resolve_autosetup_{llm,prover}_usage_file() to locate them. -def _to_token_usage(model: str, bucket: dict) -> TokenUsageDict: - """Build a graphcore ``TokenUsageDict`` from one AutoSetup rollup bucket, - keeping only the four token fields composer tracks (AutoSetup's ``calls`` - count has no slot in ``TokenTotals`` and is dropped).""" +def _to_token_usage(model: str, bucket: dict) -> NormalizedTokenUsage: + """Build a graphcore ``NormalizedTokenUsage`` from one AutoSetup rollup bucket. + The bucket keeps the raw Anthropic convention (input excludes the cache + buckets) while the normalized form totals them, hence the sum. AutoSetup's + ``calls`` count has no slot in ``TokenTotals`` and is dropped; it reports no + thinking-token split, so ``thinking_tokens`` is 0.""" + cache_read = int(bucket.get("cache_read_input_tokens", 0)) + cache_write = int(bucket.get("cache_creation_input_tokens", 0)) return { "model_name": model, - "input_tokens": int(bucket.get("input_tokens", 0)), - "output_tokens": int(bucket.get("output_tokens", 0)), - "cache_read_input_tokens": int(bucket.get("cache_read_input_tokens", 0)), - "cache_creation_input_tokens": int(bucket.get("cache_creation_input_tokens", 0)), + "total_input_tokens": int(bucket.get("input_tokens", 0)) + cache_read + cache_write, + "total_output_tokens": int(bucket.get("output_tokens", 0)), + "cache_read_tokens": cache_read, + "cache_write_tokens": cache_write, + "thinking_tokens": 0, } -def read_autosetup_usage(project_root: Path) -> list[TokenUsageDict]: +def read_autosetup_usage(project_root: Path) -> list[NormalizedTokenUsage]: """Return AutoSetup's per-model token usage for the most recent run, one - ``TokenUsageDict`` per model — ready to feed straight into + ``NormalizedTokenUsage`` per model — ready to feed straight into ``RunSummary.record_token_usage``. Returns ``[]`` on any failure (file absent — autosetup skipped, cache hit, diff --git a/tests/test_token_usage.py b/tests/test_token_usage.py index 04596223..ea438072 100644 --- a/tests/test_token_usage.py +++ b/tests/test_token_usage.py @@ -28,15 +28,17 @@ ) from composer.diagnostics.usage_callback import UsageCallback from composer.spec.source.autosetup import read_autosetup_usage -from graphcore.utils import TokenUsageDict +from graphcore.utils import NormalizedTokenUsage -def _usage(model: str, i: int, o: int, cr: int, cw: int) -> TokenUsageDict: +def _usage(model: str, i: int, o: int, cr: int, cw: int) -> NormalizedTokenUsage: + """Normalized usage for FRESH input ``i`` — total input includes the caches.""" return { - "input_tokens": i, - "output_tokens": o, - "cache_read_input_tokens": cr, - "cache_creation_input_tokens": cw, + "total_input_tokens": i + cr + cw, + "total_output_tokens": o, + "cache_read_tokens": cr, + "cache_write_tokens": cw, + "thinking_tokens": 0, "model_name": model, } @@ -130,16 +132,17 @@ async def test_token_usage_persisted_to_run_meta_tags(): # --------------------------------------------------------------------------- # def _fake_model(callbacks): + # usage_metadata is the normalized field every transport fills in (a streamed + # response carries no raw response_metadata["usage"] dict at all). Normalized + # input_tokens is the total INCLUDING the cache buckets: 100 fresh + 5 + 2. resp = AIMessage( content="ok", - response_metadata={ - "model_name": "claude-test", - "usage": { - "input_tokens": 100, - "output_tokens": 10, - "cache_read_input_tokens": 5, - "cache_creation_input_tokens": 2, - }, + response_metadata={"model_name": "claude-test"}, + usage_metadata={ + "input_tokens": 107, + "output_tokens": 10, + "total_tokens": 117, + "input_token_details": {"cache_read": 5, "cache_creation": 2}, }, ) return FakeMessagesListChatModel(responses=[resp, resp], callbacks=callbacks) @@ -222,21 +225,24 @@ def _write_autosetup_usage( ) -def test_read_autosetup_usage_returns_token_usage_dicts(tmp_path): +def test_read_autosetup_usage_returns_normalized_usage(tmp_path): _write_autosetup_usage(tmp_path, { "claude-sonnet-4-6": _autosetup_bucket(100, 10, 5, 2), "claude-opus-4": _autosetup_bucket(50, 5, 0, 1), }) by_model = {u["model_name"]: u for u in read_autosetup_usage(tmp_path)} + # The disk bucket keeps the raw convention (input excludes caches); the + # normalized form totals them: 100 fresh + 5 read + 2 write. assert by_model["claude-sonnet-4-6"] == { "model_name": "claude-sonnet-4-6", - "input_tokens": 100, - "output_tokens": 10, - "cache_read_input_tokens": 5, - "cache_creation_input_tokens": 2, + "total_input_tokens": 107, + "total_output_tokens": 10, + "cache_read_tokens": 5, + "cache_write_tokens": 2, + "thinking_tokens": 0, } - assert by_model["claude-opus-4"]["input_tokens"] == 50 + assert by_model["claude-opus-4"]["total_input_tokens"] == 51 assert "calls" not in by_model["claude-sonnet-4-6"] # AutoSetup-only field dropped @@ -285,4 +291,4 @@ def test_autosetup_usage_fallback_newest_dir(tmp_path): timestamp="20260102_000000", write_result=False) usage = read_autosetup_usage(tmp_path) assert [u["model_name"] for u in usage] == ["new"] - assert usage[0]["input_tokens"] == 9 + assert usage[0]["total_input_tokens"] == 9