streaming: episodes_per_epoch worker lifetime cap (dataloader workers leaked to OOM-kill on multi-hour runs twice)

tinyvla_b200/src/tinyvla/data/streaming_hub.py

@@ -67,6 +67,9 @@ class HubEpisodeStream(IterableDataset):
         robot_prompts: dict | None = None,      # morph_key -> str
         support_cache_datasets: int = 8,
         support_cache_size: int = 48,
+        episodes_per_epoch: int = 200,  # воркер завершает поток после N эпизодов:
+                                        # ограничивает жизнь процесса и любые
+                                        # медленные утечки (OOM-kill на длинных прогонах)
     ):
         super().__init__()
         self.specs = specs
@@ -75,6 +78,7 @@ class HubEpisodeStream(IterableDataset):
         self.robot_prompts = robot_prompts or {}
         self.support_cache_datasets = support_cache_datasets
         self.support_cache_size = support_cache_size
+        self.episodes_per_epoch = episodes_per_epoch
         self.token = token
         self.chunk = chunk
         self.image_size = image_size
@@ -325,7 +329,9 @@ class HubEpisodeStream(IterableDataset):
             item["support_actions"] = torch.stack(acts)
             return item
 
-        while True:  # infinite stream; the training loop counts steps, not epochs
+        ep_count = 0
+        while self.episodes_per_epoch <= 0 or ep_count < self.episodes_per_epoch:
+            ep_count += 1
             si = rng.choices(range(len(self.specs)), weights=weights, k=1)[0]
             try:
                 meta = self._load_meta(si)
@@ -351,3 +357,5 @@ class HubEpisodeStream(IterableDataset):
                 import logging
 
                 logging.warning(f"stream: skipping {self.specs[si]['repo_id']} ep after {type(e).__name__}: {e}")
+        rng.shuffle(buf)
+        yield from buf  # эпоха воркера кончилась: дослать буфер и умереть чистым