Fix issues that made long runs fail or hang

composer/diagnostics/usage_callback.py

@@ -20,10 +20,41 @@
 from langchain_core.messages import AIMessage
 from langchain_core.outputs import ChatGeneration, LLMResult
 
-from graphcore.utils import get_token_usage
+from graphcore.utils import TokenUsageDict, get_token_usage
 from composer.diagnostics.timing import get_run_summary
 
 
+def _usage_of(msg: AIMessage) -> TokenUsageDict:
+    """Token usage of a response, whichever transport produced it.
+
+    ``get_token_usage`` reads the raw Anthropic ``response_metadata["usage"]`` dict,
+    which only a non-streamed response carries; a streamed response reports usage
+    solely through the provider-normalized ``usage_metadata``. Fall back to that,
+    translating back to the raw shape: normalized ``input_tokens`` is the total
+    including both cache buckets, where the raw count excludes them."""
+    usage = get_token_usage(msg)
+    if any(
+        usage[k]
+        for k in (
+            "input_tokens", "output_tokens",
+            "cache_read_input_tokens", "cache_creation_input_tokens",
+        )
+    ):
+        return usage
+    if (um := msg.usage_metadata) is None:
+        return usage
+    details = um.get("input_token_details") or {}
+    cache_read = details.get("cache_read", 0)
+    cache_creation = details.get("cache_creation", 0)
+    return {
+        "model_name": usage["model_name"],
+        "input_tokens": max(0, um["input_tokens"] - cache_read - cache_creation),
+        "output_tokens": um["output_tokens"],
+        "cache_read_input_tokens": cache_read,
+        "cache_creation_input_tokens": cache_creation,
+    }
+
+
 class UsageCallback(BaseCallbackHandler):
     """Records each LLM response's token usage into the active ``RunSummary``."""
 
@@ -45,4 +76,4 @@ def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
         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))
+            get_run_summary().record_token_usage(_usage_of(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