[bugfix] Kafka: batch + retry offsets_for_times to survive broker timeouts (#547)

tiankongdeguiji · claude · web-flow · commit 516644068892 · 2026-06-23T09:48:01.000+08:00
Co-authored-by: Claude Opus 4.8 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/docs/source/feature/data.md b/docs/source/feature/data.md
@@ -106,8 +106,13 @@ data_config {
   - `kafka://broker:9092/topic?group.id=consumer_group&auto.offset.reset=earliest`
   - 需以`&`分隔符来分隔kafka的参数，`group.id`是必选参数，其余参数参考[Kafka配置文档](https://github.com/confluentinc/librdkafka/blob/master/CONFIGURATION.md)
   - `enable.auto.commit`默认设置为`false`，KafkaDataset使用自身的checkpoint机制管理消费位点，不依赖Kafka broker端的offset提交。如需启用自动提交，可在URI中显式设置`enable.auto.commit=true`
-  - 支持`start.timestamp.ms`参数，指定从某个时间戳（毫秒）开始消费，消费者会从各分区中时间戳 >= 该值的最早offset开始读取。当同时存在checkpoint时，checkpoint优先级更高。示例:
-    - `kafka://broker:9092/topic?group.id=consumer_group&auto.offset.reset=earliest&start.timestamp.ms=1711929600000`
+  - 支持`start.timestamp.ms`参数，指定从某个时间戳（毫秒）开始消费，消费者会从各分区中时间戳 >= 该值的最早offset开始读取。当同时存在checkpoint时，checkpoint优先级更高
+    - 示例:
+      - `kafka://broker:9092/topic?group.id=consumer_group&auto.offset.reset=earliest&start.timestamp.ms=1711929600000`
+    - 使用`start.timestamp.ms`时，时间戳到offset的解析（`offsets_for_times`）会分批请求并带重试，可通过如下环境变量调优，避免分区较多时单次请求超时：
+      - `TZREC_KAFKA_OFFSETS_FOR_TIMES_BATCH_SIZE`: 每次`offsets_for_times`请求解析的分区数，默认`32`
+      - `TZREC_KAFKA_OFFSETS_FOR_TIMES_TIMEOUT`: 每次`offsets_for_times`请求的超时时间（秒），默认`30.0`
+      - `TZREC_KAFKA_OFFSETS_FOR_TIMES_RETRIES`: 每批请求失败后的重试次数，默认`2`
 
 - 注意:
 
diff --git a/tzrec/datasets/kafka_dataset.py b/tzrec/datasets/kafka_dataset.py
@@ -14,11 +14,16 @@
 import os
 import threading
 import time
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
 from urllib.parse import parse_qsl, urlparse
 
 import pyarrow as pa
-from confluent_kafka import OFFSET_INVALID, Consumer, TopicPartition
+from confluent_kafka import (
+    OFFSET_INVALID,
+    Consumer,
+    KafkaException,
+    TopicPartition,
+)
 
 from tzrec.datasets.dataset import BaseDataset, BaseReader
 from tzrec.datasets.utils import (
@@ -91,6 +96,72 @@ def _parse_kafka_uri(uri: str) -> Tuple[str, Dict[str, Any], Optional[int]]:
     return topic, params, start_timestamp_ms
 
 
+def _offsets_for_times_batched(
+    consumer: Consumer,
+    ts_partitions: List[TopicPartition],
+    start_timestamp_ms: int,
+    batch_size: int,
+    timeout: float,
+    max_retries: int,
+    sleep_fn: Callable[[float], None] = time.sleep,
+) -> Dict[Tuple[str, int], int]:
+    """Resolve a start timestamp to offsets in chunks, with retry.
+
+    ``consumer.offsets_for_times`` is split into chunks of ``batch_size``
+    partitions to keep each broker request small, and every chunk is retried
+    up to ``max_retries`` times on ``KafkaException`` (e.g. request timeout).
+    A persistent failure re-raises rather than silently falling back, because
+    this reader is shared with training, where falling back to
+    ``auto.offset.reset`` would replay the whole topic.
+
+    Args:
+        consumer (Consumer): the kafka consumer.
+        ts_partitions (list): partitions to resolve.
+        start_timestamp_ms (int): target timestamp in milliseconds.
+        batch_size (int): number of partitions per ``offsets_for_times`` call.
+        timeout (float): per-call request timeout in seconds.
+        max_retries (int): retries per chunk after the first attempt.
+        sleep_fn (callable): sleep between retries (injectable for testing).
+
+    Returns:
+        dict: map of (topic, partition) to resolved offset, including only
+            partitions the broker resolved to a real (>= 0) offset.
+
+    Raises:
+        KafkaException: if a chunk still fails after exhausting retries.
+    """
+    resolved_map: Dict[Tuple[str, int], int] = {}
+    for i in range(0, len(ts_partitions), batch_size):
+        chunk = ts_partitions[i : i + batch_size]
+        # offsets_for_times reads the target timestamp from the offset field.
+        for tp in chunk:
+            tp.offset = start_timestamp_ms
+        retry_cnt = 0
+        while True:
+            try:
+                resolved = consumer.offsets_for_times(chunk, timeout=timeout)
+            except KafkaException as e:
+                if retry_cnt >= max_retries:
+                    logger.error(
+                        f"offsets_for_times failed after {max_retries} retries "
+                        f"for {len(chunk)} partition(s) (timeout={timeout}s); "
+                        f"tune TZREC_KAFKA_OFFSETS_FOR_TIMES_* env vars"
+                    )
+                    raise
+                retry_cnt += 1
+                logger.warning(
+                    f"offsets_for_times attempt {retry_cnt}/{max_retries} failed "
+                    f"for {len(chunk)} partition(s), retrying: {e}"
+                )
+                sleep_fn(min(2.0 * retry_cnt, 10.0))
+                continue
+            break
+        for r in resolved:
+            if r.offset is not None and r.offset >= 0:
+                resolved_map[(r.topic, r.partition)] = r.offset
+    return resolved_map
+
+
 class KafkaDataset(BaseDataset):
     """Dataset for reading data from Kafka (each message is a serialized Arrow batch).
 
@@ -260,6 +331,14 @@ def _reader(
         """
         topic, config, start_timestamp_ms = _parse_kafka_uri(self._input_path)
         max_poll_interval_ms = int(config.get("max.poll.interval.ms", "300000"))
+        # offsets_for_times tunables (chunk size / per-call timeout / retries).
+        offt_batch_size = max(
+            1, int(os.environ.get("TZREC_KAFKA_OFFSETS_FOR_TIMES_BATCH_SIZE", "32"))
+        )
+        offt_timeout = float(
+            os.environ.get("TZREC_KAFKA_OFFSETS_FOR_TIMES_TIMEOUT", "30.0")
+        )
+        offt_retries = int(os.environ.get("TZREC_KAFKA_OFFSETS_FOR_TIMES_RETRIES", "2"))
         consumer = Consumer(config)
 
         stop_event = threading.Event()
@@ -283,10 +362,14 @@ def on_assign(consumer: Consumer, partitions: List[TopicPartition]) -> None:
                     tp.offset = OFFSET_INVALID
 
             if ts_partitions:
-                for tp in ts_partitions:
-                    tp.offset = start_timestamp_ms
-                resolved = consumer.offsets_for_times(ts_partitions, timeout=30.0)
-                resolved_map = {(r.topic, r.partition): r.offset for r in resolved}
+                resolved_map = _offsets_for_times_batched(
+                    consumer,
+                    ts_partitions,
+                    start_timestamp_ms,
+                    offt_batch_size,
+                    offt_timeout,
+                    offt_retries,
+                )
                 for tp in ts_partitions:
                     res_offset = resolved_map.get((tp.topic, tp.partition))
                     if res_offset is not None and res_offset >= 0:
diff --git a/tzrec/datasets/kafka_dataset_test.py b/tzrec/datasets/kafka_dataset_test.py
@@ -14,18 +14,23 @@
 import os
 import time
 import unittest
+from unittest import mock
 
 import pyarrow as pa
 from alibabacloud_alikafka20190916 import models as alikafka_models
 from alibabacloud_alikafka20190916.client import Client as AliKafkaClient
 from alibabacloud_credentials.client import Client as CredClient
 from alibabacloud_tea_openapi import models as openapi_models
 from alibabacloud_tea_util import models as util_models
-from confluent_kafka import Producer
+from confluent_kafka import KafkaException, Producer, TopicPartition
 from parameterized import parameterized
 from torch.utils.data import DataLoader
 
-from tzrec.datasets.kafka_dataset import KafkaDataset, _parse_kafka_uri
+from tzrec.datasets.kafka_dataset import (
+    KafkaDataset,
+    _offsets_for_times_batched,
+    _parse_kafka_uri,
+)
 from tzrec.features.feature import FgMode, create_features
 from tzrec.protos import data_pb2, feature_pb2
 from tzrec.utils.checkpoint_util import update_dataloder_state
@@ -69,6 +74,122 @@ def test_parse_with_non_integer_timestamp_raises(self):
             _parse_kafka_uri(uri)
 
 
+def _resolved(partition, offset, topic="t"):
+    """Build a TopicPartition as returned by offsets_for_times."""
+    tp = TopicPartition(topic, partition)
+    tp.offset = offset
+    return tp
+
+
+class OffsetsForTimesBatchedTest(unittest.TestCase):
+    """Unit tests for _offsets_for_times_batched (broker-free, mocked)."""
+
+    START_TS = 1711929600000
+
+    def test_empty_partitions_returns_empty(self):
+        consumer = mock.MagicMock()
+        result = _offsets_for_times_batched(
+            consumer, [], self.START_TS, 30, 30.0, 3, sleep_fn=mock.MagicMock()
+        )
+        self.assertEqual(result, {})
+        consumer.offsets_for_times.assert_not_called()
+
+    def test_single_chunk_all_resolved(self):
+        seen_offsets = []
+
+        def _fake(chunk, timeout):
+            # offsets_for_times receives the target timestamp in tp.offset.
+            seen_offsets.extend(tp.offset for tp in chunk)
+            return [_resolved(tp.partition, tp.partition + 100) for tp in chunk]
+
+        consumer = mock.MagicMock()
+        consumer.offsets_for_times.side_effect = _fake
+        parts = [TopicPartition("t", p) for p in range(5)]
+        result = _offsets_for_times_batched(
+            consumer, parts, self.START_TS, 30, 30.0, 3, sleep_fn=mock.MagicMock()
+        )
+        self.assertEqual(consumer.offsets_for_times.call_count, 1)
+        self.assertEqual(result, {("t", p): p + 100 for p in range(5)})
+        # target timestamp was set on every partition before the call.
+        self.assertEqual(seen_offsets, [self.START_TS] * 5)
+
+    def test_chunking_boundary(self):
+        consumer = mock.MagicMock()
+        consumer.offsets_for_times.side_effect = lambda chunk, timeout: [
+            _resolved(tp.partition, tp.partition + 100) for tp in chunk
+        ]
+        parts = [TopicPartition("t", p) for p in range(65)]
+        result = _offsets_for_times_batched(
+            consumer, parts, self.START_TS, 30, 30.0, 3, sleep_fn=mock.MagicMock()
+        )
+        chunk_sizes = [
+            len(call.args[0]) for call in consumer.offsets_for_times.call_args_list
+        ]
+        self.assertEqual(chunk_sizes, [30, 30, 5])
+        self.assertEqual(len(result), 65)
+
+    def test_partial_resolution_filters_negative(self):
+        # broker returns -1 when the timestamp is past the last message.
+        consumer = mock.MagicMock()
+        consumer.offsets_for_times.return_value = [
+            _resolved(0, 100),
+            _resolved(1, -1),
+            _resolved(2, 102),
+        ]
+        parts = [TopicPartition("t", p) for p in range(3)]
+        result = _offsets_for_times_batched(
+            consumer, parts, self.START_TS, 30, 30.0, 3, sleep_fn=mock.MagicMock()
+        )
+        self.assertEqual(result, {("t", 0): 100, ("t", 2): 102})
+        self.assertNotIn(("t", 1), result)
+
+    def test_retry_then_success(self):
+        sleep_fn = mock.MagicMock()
+        consumer = mock.MagicMock()
+        consumer.offsets_for_times.side_effect = [
+            KafkaException("transient timeout"),
+            [_resolved(0, 100)],
+        ]
+        parts = [TopicPartition("t", 0)]
+        result = _offsets_for_times_batched(
+            consumer, parts, self.START_TS, 30, 30.0, 3, sleep_fn=sleep_fn
+        )
+        self.assertEqual(result, {("t", 0): 100})
+        self.assertEqual(consumer.offsets_for_times.call_count, 2)
+        self.assertEqual(sleep_fn.call_count, 1)
+
+    def test_retry_exhausted_raises(self):
+        sleep_fn = mock.MagicMock()
+        consumer = mock.MagicMock()
+        consumer.offsets_for_times.side_effect = KafkaException("persistent timeout")
+        parts = [TopicPartition("t", 0)]
+        with self.assertRaises(KafkaException):
+            _offsets_for_times_batched(
+                consumer, parts, self.START_TS, 30, 30.0, 3, sleep_fn=sleep_fn
+            )
+        # 1 initial attempt + 3 retries = 4 calls; sleep before each retry = 3.
+        self.assertEqual(consumer.offsets_for_times.call_count, 4)
+        self.assertEqual(sleep_fn.call_count, 3)
+
+    def test_retry_isolated_per_chunk(self):
+        sleep_fn = mock.MagicMock()
+        # chunk 0 fails once then succeeds; chunk 1 succeeds first try.
+        outcomes = [
+            KafkaException("blip"),
+            [_resolved(0, 100)],
+            [_resolved(1, 101)],
+        ]
+        consumer = mock.MagicMock()
+        consumer.offsets_for_times.side_effect = outcomes
+        parts = [TopicPartition("t", 0), TopicPartition("t", 1)]
+        result = _offsets_for_times_batched(
+            consumer, parts, self.START_TS, 1, 30.0, 3, sleep_fn=sleep_fn
+        )
+        self.assertEqual(result, {("t", 0): 100, ("t", 1): 101})
+        self.assertEqual(consumer.offsets_for_times.call_count, 3)
+        self.assertEqual(sleep_fn.call_count, 1)
+
+
 class KafkaDatasetTest(unittest.TestCase):
     def setUp(self):
         credential = CredClient()
