Fix issues that made long runs fail or hang

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, {})

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))

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,

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]]
 

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

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,

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