fix(weave): retry chunked file reads when replication is lagging

Author: gtarpenning

tests/trace_server/test_clickhouse_trace_server_batched.py

@@ -1244,6 +1244,8 @@ def test_file_content_read_retries_eventual_consistency():
         assert mock_read_once.call_count == 2
 
     # Real miss: after exhausting retries, the original NotFoundError still surfaces.
+    # `_file_content_read_with_retry` defaults to 5 attempts so chunked-file reads
+    # can absorb replication lag across multiple shards/replicas.
     with patch.object(
         server,
         "_file_content_read_once",
@@ -1255,7 +1257,7 @@ def test_file_content_read_retries_eventual_consistency():
         ):
             server.file_content_read(req)
 
-        assert mock_read_once.call_count == 2
+        assert mock_read_once.call_count == 5
 
 
 @pytest.mark.disable_logging_error_check

weave/trace_server/clickhouse_trace_server_batched.py

@@ -3446,9 +3446,7 @@ def annotator_queue_items_progress_update(
                 annotation_state=req.annotation_state,
                 pb=update_pb,
                 cluster_name=self.clickhouse_cluster_name,
-                table_name=self._mutation_table_name(
-                    "annotator_queue_items_progress"
-                ),
+                table_name=self._mutation_table_name("annotator_queue_items_progress"),
             )
             self._command(
                 update_query,
@@ -5557,9 +5555,15 @@ def _obj_read_with_retry(
         name="clickhouse_trace_server_batched._file_content_read_with_retry"
     )
     def _file_content_read_with_retry(
-        self, req: tsi.FileContentReadReq, max_attempts: int = 2
+        self, req: tsi.FileContentReadReq, max_attempts: int = 5
     ) -> tsi.FileContentReadRes:
-        """Read file content with retry for ClickHouse eventual consistency."""
+        """Read file content with retry for ClickHouse eventual consistency.
+
+        Higher attempt count than `_obj_read_with_retry` because chunked-file
+        reads need all chunks visible across all shards/replicas: in
+        replicated/distributed mode each missing chunk on any replica trips
+        the retry, and lag can briefly exceed the 50ms-1s window.
+        """
         return self._read_with_retry(
             lambda: self._file_content_read_once(req), max_attempts=max_attempts
         )
@@ -5995,7 +5999,12 @@ def _file_content_read_once(
         result_rows = list(query_result.result_rows)
 
         if len(result_rows) < n_chunks:
-            raise ValueError("Missing chunks")
+            # Treat as not-found so the tenacity retry in `_read_with_retry`
+            # picks it up. Replicated/distributed reads can transiently see
+            # fewer rows than `n_chunks` while replication catches up.
+            raise NotFoundError(
+                f"File with digest {req.digest} has {len(result_rows)}/{n_chunks} chunks visible"
+            )
         elif len(result_rows) > n_chunks:
             # The general case where this can occur is when there are multiple
             # writes of the same digest AND the effective `FILE_CHUNK_SIZE`

weave/trace_server/clickhouse_trace_server_migrator.py

@@ -175,6 +175,12 @@ def _is_transient_ch_error(exc: BaseException) -> bool:
     # "versions for agent" queries have the same locality as the agent row.
     "agents": "project_id, agent_name",
     "agent_versions": "project_id, agent_name",
+    # Files are chunked: `_file_content_read_once` selects all rows for a
+    # (project_id, digest) and checks the count against `n_chunks`. With
+    # rand() sharding chunks land on different shards, so any per-shard
+    # replication lag manifests as "Missing chunks". Co-locate chunks of
+    # one file on one shard so the read sees an atomic set.
+    "files": "project_id, digest",
 }