fix: handle truncated batch stream with sequential GET fallback

- catch StopIteration during batch result iteration and fall back
  to individual GET requests instead of crashing the DataLoader worker
- use batch_stream_failed flag to skip dead iterator for remaining objects
- reuse existing _get_object_from_moss_in() retry path for recovery
- update test to verify fallback behavior instead of expecting crash

Signed-off-by: Abhishek Gaikwad <gaikwadabhishek1997@gmail.com>

Author: gaikwadabhishek

lhotse/ais/batch_loader.py

@@ -133,6 +133,7 @@ def __call__(self, cuts: CutSet) -> CutSet:
             logger.warning(
                 f"AIStore batch.get() failed: {e}. Falling back to sequential GET requests."
             )
+
             # Fallback: make sequential GET requests for each object in the batch
             # Use a generator to maintain consistency with batch.get() which returns an iterator
             def sequential_get():
@@ -153,33 +154,49 @@ def sequential_get():
 
         # Apply the received data back into each manifest that had a URL
         request_idx = 0
+        batch_stream_failed = False
         for manifest, has_url in manifest_list:
             if has_url:
                 info = None
                 content = None
 
-                try:
-                    info, content = next(batch_result)
-                except StopIteration:
-                    raise AISBatchLoaderError(
-                        "Batch result iterator exhausted prematurely. "
-                        f"Expected more objects for manifests with URLs."
-                    )
-                except TimeoutError as e:
-                    # Timeout occurred - recover the request info from saved_requests_list
-                    logger.warning(
-                        f"Timeout while fetching batch result at index {request_idx}: {e}. "
-                        f"Falling back to direct AIStore API call."
-                    )
-
-                    if request_idx < len(saved_requests_list):
-                        info = saved_requests_list[request_idx]
-                        content = b""  # Mark as empty to trigger retry
-                    else:
-                        raise AISBatchLoaderError(
-                            f"Timeout at request index {request_idx}, but cannot recover: "
-                            f"index out of range for saved_requests_list (len={len(saved_requests_list)})"
-                        ) from e
+                if batch_stream_failed:
+                    # Batch stream already broke — go straight to individual GET
+                    info = saved_requests_list[request_idx]
+                    content = b""  # trigger retry below
+                else:
+                    try:
+                        info, content = next(batch_result)
+                    except StopIteration:
+                        # Batch stream was truncated (e.g., connection reset mid-tar).
+                        # Fall back to individual GET for this and all remaining objects.
+                        batch_stream_failed = True
+                        logger.warning(
+                            f"Batch stream truncated at index {request_idx}/{len(saved_requests_list)}. "
+                            f"Falling back to direct AIStore API calls for remaining objects."
+                        )
+                        if request_idx < len(saved_requests_list):
+                            info = saved_requests_list[request_idx]
+                            content = b""  # trigger retry below
+                        else:
+                            raise AISBatchLoaderError(
+                                f"Batch stream truncated at index {request_idx}, but cannot recover: "
+                                f"index out of range for saved_requests_list (len={len(saved_requests_list)})"
+                            )
+                    except TimeoutError as e:
+                        # Timeout occurred - recover the request info from saved_requests_list
+                        logger.warning(
+                            f"Timeout while fetching batch result at index {request_idx}: {e}. "
+                            f"Falling back to direct AIStore API call."
+                        )
+                        if request_idx < len(saved_requests_list):
+                            info = saved_requests_list[request_idx]
+                            content = b""  # Mark as empty to trigger retry
+                        else:
+                            raise AISBatchLoaderError(
+                                f"Timeout at request index {request_idx}, but cannot recover: "
+                                f"index out of range for saved_requests_list (len={len(saved_requests_list)})"
+                            ) from e
 
                 # Retry with direct API call if content is empty (from timeout or actual empty response)
                 if content == b"":

test/cut/test_ais_batch_loader.py

@@ -4,6 +4,7 @@
 These tests use mocking to simulate AIStore client behavior,
 allowing them to run in CI environments without AIStore infrastructure.
 """
+
 from unittest.mock import MagicMock, patch
 
 import numpy as np
@@ -968,35 +969,45 @@ def test_fallback_failure_raises_error(
             loader(cuts)
 
     @patch("lhotse.ais.batch_loader.get_aistore_client")
-    def test_iterator_exhausted_raises_error(
+    def test_iterator_exhausted_falls_back_to_sequential(
         self, mock_get_client, cut_with_url_recording
     ):
-        """Test that iterator exhaustion raises AISBatchLoaderError."""
+        """Test that iterator exhaustion falls back to individual GET requests."""
         client = MagicMock()
         batch = MagicMock()
 
         # Track add() calls
         add_count = []
         batch.add.side_effect = lambda *args, **kwargs: add_count.append(1)
 
-        # Mock batch.get() to return fewer items than expected
+        # Mock batch.get() to return fewer items than expected (empty iterator)
         def mock_batch_get():
-            # Return nothing even though we expect 1 item
             return iter([])
 
         batch.get.side_effect = lambda: mock_batch_get()
+        batch.requests_list = [
+            MagicMock(
+                obj_name="test.wav", bck="test-bucket", provider="ais", archpath=""
+            )
+        ]
         client.batch.return_value = batch
-        client.bucket.return_value = MagicMock()
+        mock_bucket = MagicMock()
+        mock_obj = MagicMock()
+        mock_reader = MagicMock()
+        mock_reader.read_all.return_value = b"\x00" * 16000
+        mock_obj.get_reader.return_value = mock_reader
+        mock_bucket.object.return_value = mock_obj
+        client.bucket.return_value = mock_bucket
         mock_get_client.return_value = (client, None)
 
         loader = AISBatchLoader()
         cuts = CutSet.from_cuts([cut_with_url_recording])
 
-        # Should raise AISBatchLoaderError when iterator is exhausted
-        with pytest.raises(
-            AISBatchLoaderError, match="Batch result iterator exhausted prematurely"
-        ):
-            loader(cuts)
+        # Should NOT raise — falls back to individual GET
+        result = loader(cuts)
+        assert result is not None
+        # Verify the fallback GET was called
+        mock_obj.get_reader.assert_called()
 
     @patch("lhotse.ais.batch_loader.get_aistore_client")
     def test_multiple_cuts_with_fallback(self, mock_get_client, cut_with_url_recording):