caption-flow 0.2.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

caption_flow/cli.py

@@ -120,13 +120,19 @@ class ConfigManager:
 
 
 def setup_logging(verbose: bool = False):
-    """Configure logging with rich handler."""
+    """Configure logging with rich handler, including timestamp."""
     level = logging.DEBUG if verbose else logging.INFO
     logging.basicConfig(
         level=level,
-        format="%(message)s",
+        format="%(asctime)s %(message)s",
+        datefmt="[%Y-%m-%d %H:%M:%S]",
         handlers=[
-            RichHandler(console=console, rich_tracebacks=True, show_path=False, show_time=False)
+            RichHandler(
+                console=console,
+                rich_tracebacks=True,
+                show_path=False,
+                show_time=True,  # Enables timestamp in RichHandler output
+            )
         ],
     )
 

caption_flow/monitor.py

@@ -107,7 +107,7 @@ class Monitor:
         """Main display update loop."""
         layout = self._create_layout()
 
-        with Live(layout, console=self.console, refresh_per_second=4, screen=True) as live:
+        with Live(layout, console=self.console, refresh_per_second=1, screen=True) as live:
             while self.running:
                 self._update_layout(layout)
                 await asyncio.sleep(0.25)

caption_flow/orchestrator.py

@@ -363,6 +363,7 @@ class Orchestrator:
         self.ssl_context = self._setup_ssl()
 
         # Statistics
+        self.is_generating_stats = False
         self.stats = {
             "total_chunks": 0,
             "completed_chunks": 0,
@@ -1409,28 +1410,36 @@ class Orchestrator:
         finally:
             del self.data_workers[worker_id]
 
-    async def _handle_monitor(self, websocket: WebSocketServerProtocol):
-        """Handle monitor connection."""
-        self.monitors.add(websocket)
-        logger.info("Monitor connected")
-
+    async def _send_leaderboard_to_monitor(self, websocket: WebSocketServerProtocol):
+        """Send leaderboard data to a specific monitor."""
+        total_start = time.time()
         try:
-            # Send initial stats
-            await websocket.send(safe_json_dumps({"type": "stats", "data": self.stats}))
-
-            # Send chunk stats
-            chunk_stats = self.chunk_manager.get_stats()
-            await websocket.send(safe_json_dumps({"type": "chunk_stats", "data": chunk_stats}))
+            if websocket not in self.monitors:
+                return
 
-            # Send contributor leaderboard with active worker counts
+            # Get contributors asynchronously
+            contributors_start = time.time()
             contributors = await self.storage.get_top_contributors(10)
+            logger.debug(
+                f"Contributors retrieved in {(time.time() - contributors_start)*1000:.1f}ms"
+            )
 
-            # Enhance contributor data with active worker counts
-            enhanced_contributors = []
-            worker_counts = (
-                self.get_workers_by_user_stats() if hasattr(self, "workers_by_user") else {}
+            # Get worker counts in thread pool
+            worker_counts_start = time.time()
+            loop = asyncio.get_event_loop()
+            worker_counts = await loop.run_in_executor(
+                None,
+                lambda: (
+                    self.get_workers_by_user_stats() if hasattr(self, "workers_by_user") else {}
+                ),
+            )
+            logger.debug(
+                f"Worker counts retrieved in {(time.time() - worker_counts_start)*1000:.1f}ms"
             )
 
+            # Build enhanced contributors list
+            build_start = time.time()
+            enhanced_contributors = []
             for contributor in contributors:
                 contrib_dict = {
                     "contributor_id": contributor.contributor_id,
@@ -1442,40 +1451,157 @@ class Orchestrator:
                     ),
                 }
                 enhanced_contributors.append(contrib_dict)
+            logger.debug(f"Enhanced contributors built in {(time.time() - build_start)*1000:.1f}ms")
 
-            await websocket.send(
-                safe_json_dumps({"type": "leaderboard", "data": enhanced_contributors})
+            # Cache for future monitors
+            self._cached_leaderboard = enhanced_contributors
+
+            # Send if still connected
+            if websocket in self.monitors:
+                send_start = time.time()
+                await websocket.send(
+                    safe_json_dumps({"type": "leaderboard", "data": enhanced_contributors})
+                )
+                logger.debug(
+                    f"Leaderboard sent to monitor in {(time.time() - send_start)*1000:.1f}ms"
+                )
+
+            logger.debug(
+                f"Leaderboard send to monitor completed in {(time.time() - total_start)*1000:.1f}ms"
+            )
+
+        except websockets.exceptions.ConnectionClosed:
+            logger.debug("Monitor disconnected during leaderboard send")
+        except Exception as e:
+            logger.error(f"Error sending leaderboard to monitor: {e}")
+
+    async def _send_initial_monitor_data(self, websocket: WebSocketServerProtocol):
+        """Send initial data to monitor in a separate task to avoid blocking."""
+        total_start = time.time()
+        try:
+            # Check if websocket is still in monitors set
+            if websocket not in self.monitors:
+                logger.debug("Monitor disconnected before initial data send")
+                return
+
+            # Send current stats (already in memory)
+            stats_start = time.time()
+            await websocket.send(safe_json_dumps({"type": "stats", "data": self.stats}))
+            logger.debug(f"Monitor stats sent in {(time.time() - stats_start)*1000:.1f}ms")
+
+            # Get chunk stats asynchronously
+            chunk_stats_start = time.time()
+            loop = asyncio.get_event_loop()
+            chunk_stats = await loop.run_in_executor(None, self.chunk_manager.get_stats)
+            logger.debug(f"Chunk stats retrieved in {(time.time() - chunk_stats_start)*1000:.1f}ms")
+
+            if websocket not in self.monitors:
+                return
+
+            chunk_send_start = time.time()
+            await websocket.send(safe_json_dumps({"type": "chunk_stats", "data": chunk_stats}))
+            logger.debug(f"Chunk stats sent in {(time.time() - chunk_send_start)*1000:.1f}ms")
+
+            # For leaderboard, check if we have a cached version first
+            if hasattr(self, "_cached_leaderboard") and self._cached_leaderboard:
+                # Use cached leaderboard if available
+                cache_send_start = time.time()
+                await websocket.send(
+                    safe_json_dumps({"type": "leaderboard", "data": self._cached_leaderboard})
+                )
+                logger.debug(
+                    f"Cached leaderboard sent in {(time.time() - cache_send_start)*1000:.1f}ms"
+                )
+            else:
+                # Schedule leaderboard update separately
+                leaderboard_task_start = time.time()
+                asyncio.create_task(self._send_leaderboard_to_monitor(websocket))
+                logger.debug(
+                    f"Leaderboard task created in {(time.time() - leaderboard_task_start)*1000:.1f}ms"
+                )
+
+            logger.debug(
+                f"Monitor initial data send completed in {(time.time() - total_start)*1000:.1f}ms"
+            )
+
+        except websockets.exceptions.ConnectionClosed:
+            logger.debug("Monitor disconnected during initial data send")
+        except Exception as e:
+            logger.error(f"Error sending initial monitor data: {e}")
+
+    async def _handle_monitor(self, websocket: WebSocketServerProtocol):
+        """Handle monitor connection - truly non-blocking version."""
+        monitor_start = time.time()
+        self.monitors.add(websocket)
+        logger.info(f"Monitor connected (total monitors: {len(self.monitors)})")
+
+        try:
+            # Send welcome message immediately
+            welcome_start = time.time()
+            await websocket.send(safe_json_dumps({"type": "welcome", "role": "monitor"}))
+            logger.debug(f"Monitor welcome sent in {(time.time() - welcome_start)*1000:.1f}ms")
+
+            # Schedule initial data send as a separate task to avoid blocking
+            task_create_start = time.time()
+            asyncio.create_task(self._send_initial_monitor_data(websocket))
+            logger.debug(
+                f"Monitor initial data task created in {(time.time() - task_create_start)*1000:.1f}ms"
             )
 
-            # Keep connection alive
-            async for _ in websocket:
-                pass
+            # Just keep the connection alive - no blocking work here
+            try:
+                async for message in websocket:
+                    # Handle any incoming messages from monitor if needed
+                    # For now, just ignore them
+                    pass
+            except websockets.exceptions.ConnectionClosed:
+                pass  # Normal disconnection
 
         except websockets.exceptions.ConnectionClosed:
             logger.info("Monitor disconnected")
+        except Exception as e:
+            logger.error(f"Error in monitor handler: {e}")
         finally:
             self.monitors.discard(websocket)
+            logger.debug(f"Monitor handler completed in {(time.time() - monitor_start)*1000:.1f}ms")
 
     async def _broadcast_stats(self):
-        """Broadcast statistics to all monitors - enhanced for multi-stage."""
+        """Broadcast statistics to all monitors - truly non-blocking version."""
         if not self.monitors:
             return
+        if self.is_generating_stats:
+            return  # Already generating stats, skip this call
+        self.is_generating_stats = True
+        total_start = time.time()
+
+        # Prepare all the data first
+        data_prep_start = time.time()
+        loop = asyncio.get_event_loop()
 
-        # Get storage stats
+        # Get storage stats (already async)
+        storage_stats_start = time.time()
         storage_stats = await self.storage.get_storage_stats()
+        logger.debug(f"Storage stats retrieved in {(time.time() - storage_stats_start)*1000:.1f}ms")
+
+        caption_stats_start = time.time()
         caption_stats = await self.storage.get_caption_stats()
+        logger.debug(f"Caption stats retrieved in {(time.time() - caption_stats_start)*1000:.1f}ms")
 
-        # Include chunk stats
-        chunk_stats = self.chunk_manager.get_stats()
-        self.stats.update({f"chunks_{k}": v for k, v in chunk_stats.items()})
+        # Get chunk stats in thread pool
+        chunk_stats_start = time.time()
+        chunk_stats = await loop.run_in_executor(None, self.chunk_manager.get_stats)
+        logger.debug(f"Chunk stats retrieved in {(time.time() - chunk_stats_start)*1000:.1f}ms")
 
-        # Merge storage stats
-        self.stats.update(storage_stats)
-        self.stats["field_breakdown"] = caption_stats.get("field_stats", {})
-        self.stats["output_fields_list"] = caption_stats.get("output_fields", [])
+        # Build stats dict
+        build_stats_start = time.time()
+        stats_update = self.stats.copy()
+        stats_update.update({f"chunks_{k}": v for k, v in chunk_stats.items()})
+        stats_update.update(storage_stats)
+        stats_update["field_breakdown"] = caption_stats.get("field_stats", {})
+        stats_update["output_fields_list"] = caption_stats.get("output_fields", [])
 
         # Add rate information
-        self.stats.update(
+        stats_update.update(
             {
                 "current_rate": self.rate_tracker["current_rate"],
                 "average_rate": self.rate_tracker["average_rate"],
@@ -1483,41 +1609,106 @@ class Orchestrator:
             }
         )
 
-        # Add vLLM info - now includes stage count
-        self.stats["vllm_model"] = self.vllm_config.get("model", "unknown")
-        self.stats["vllm_batch_size"] = self.vllm_config.get("batch_size", 0)
+        # Add vLLM info
+        stats_update["vllm_model"] = self.vllm_config.get("model", "unknown")
+        stats_update["vllm_batch_size"] = self.vllm_config.get("batch_size", 0)
 
-        # NEW: Add stage information
+        # Add stage information
         stages = self.vllm_config.get("stages", [])
         if stages:
-            self.stats["stage_count"] = len(stages)
-            self.stats["stage_names"] = [s.get("name", "unnamed") for s in stages]
+            stats_update["stage_count"] = len(stages)
+            stats_update["stage_names"] = [s.get("name", "unnamed") for s in stages]
         else:
-            self.stats["stage_count"] = 1  # Backward compatibility
-            self.stats["stage_names"] = ["default"]
+            stats_update["stage_count"] = 1
+            stats_update["stage_names"] = ["default"]
 
+        # Get field stats
+        field_stats_start = time.time()
         field_stats = await self.storage.get_output_field_stats()
-        self.stats["output_fields"] = field_stats
+        stats_update["output_fields"] = field_stats
+        logger.debug(f"Field stats retrieved in {(time.time() - field_stats_start)*1000:.1f}ms")
 
-        message = safe_json_dumps({"type": "stats", "data": self.stats})
+        # Update our internal stats
+        self.stats = stats_update
+        logger.debug(f"Stats prepared in {(time.time() - build_stats_start)*1000:.1f}ms")
 
-        # Send to all monitors
-        disconnected = set()
-        _monitors = self.monitors.copy()
-        for monitor in _monitors:
+        logger.debug(f"Total data preparation took {(time.time() - data_prep_start)*1000:.1f}ms")
+
+        # Create message once
+        message_create_start = time.time()
+        stats_message = safe_json_dumps({"type": "stats", "data": self.stats})
+        logger.debug(f"Stats message created in {(time.time() - message_create_start)*1000:.1f}ms")
+
+        # Send to all monitors asynchronously in parallel
+        send_start = time.time()
+
+        async def send_to_monitor(monitor):
             try:
-                await monitor.send(message)
+                await monitor.send(stats_message)
             except websockets.exceptions.ConnectionClosed:
-                disconnected.add(monitor)
+                return monitor  # Return for removal
+            except Exception as e:
+                logger.debug(f"Error sending stats to monitor: {e}")
+                return monitor  # Return for removal
+            return None
 
-        # send updated leaderboard
+        # Send to all monitors in parallel
+        monitors_copy = self.monitors.copy()
+        results = await asyncio.gather(
+            *[send_to_monitor(m) for m in monitors_copy], return_exceptions=True
+        )
+
+        # Remove disconnected monitors
+        disconnected = {
+            m
+            for m, r in zip(monitors_copy, results)
+            if r is not None and not isinstance(r, Exception)
+        }
+        self.monitors -= disconnected
+
+        logger.debug(
+            f"Stats sent to {len(monitors_copy)} monitors in {(time.time() - send_start)*1000:.1f}ms"
+        )
+
+        # Send leaderboard update in a separate task to avoid blocking
+        leaderboard_task_start = time.time()
+        asyncio.create_task(self._broadcast_leaderboard())
+        self.is_generating_stats = False
+        logger.debug(
+            f"Leaderboard broadcast task created in {(time.time() - leaderboard_task_start)*1000:.1f}ms"
+        )
+        logger.debug(f"Stats broadcast completed in {(time.time() - total_start)*1000:.1f}ms")
+
+    async def _broadcast_leaderboard(self):
+        """Send leaderboard updates to monitors - separate from stats to avoid blocking."""
+        if not self.monitors:
+            return
+
+        total_start = time.time()
         try:
+            # Get contributors
+            contributors_start = time.time()
             contributors = await self.storage.get_top_contributors(10)
-            enhanced_contributors = []
-            worker_counts = (
-                self.get_workers_by_user_stats() if hasattr(self, "workers_by_user") else {}
+            logger.debug(
+                f"Contributors retrieved for broadcast in {(time.time() - contributors_start)*1000:.1f}ms"
+            )
+
+            # Get worker counts
+            worker_counts_start = time.time()
+            loop = asyncio.get_event_loop()
+            worker_counts = await loop.run_in_executor(
+                None,
+                lambda: (
+                    self.get_workers_by_user_stats() if hasattr(self, "workers_by_user") else {}
+                ),
+            )
+            logger.debug(
+                f"Worker counts retrieved for broadcast in {(time.time() - worker_counts_start)*1000:.1f}ms"
             )
 
+            # Build enhanced contributors list
+            build_start = time.time()
+            enhanced_contributors = []
             for contributor in contributors:
                 contrib_dict = {
                     "contributor_id": contributor.contributor_id,
@@ -1529,26 +1720,64 @@ class Orchestrator:
                     ),
                 }
                 enhanced_contributors.append(contrib_dict)
+            logger.debug(
+                f"Enhanced contributors built for broadcast in {(time.time() - build_start)*1000:.1f}ms"
+            )
+
+            # Cache it
+            self._cached_leaderboard = enhanced_contributors
 
+            # Create message once
+            message_create_start = time.time()
             leaderboard_message = safe_json_dumps(
                 {"type": "leaderboard", "data": enhanced_contributors}
             )
+            logger.debug(
+                f"Leaderboard message created in {(time.time() - message_create_start)*1000:.1f}ms"
+            )
 
-            # Send to all monitors
-            disconnected = set()
-            for monitor in self.monitors.copy():
+            # Send to all monitors in parallel
+            send_start = time.time()
+
+            async def send_leaderboard(monitor):
                 try:
                     await monitor.send(leaderboard_message)
-                except websockets.exceptions.ConnectionClosed:
-                    disconnected.add(monitor)
+                except:
+                    return monitor  # Mark for removal
+                return None
+
+            monitors_copy = self.monitors.copy()
+            results = await asyncio.gather(
+                *[send_leaderboard(m) for m in monitors_copy], return_exceptions=True
+            )
 
+            # Remove disconnected
+            disconnected = {
+                m
+                for m, r in zip(monitors_copy, results)
+                if r is not None and not isinstance(r, Exception)
+            }
             self.monitors -= disconnected
 
-        except Exception as e:
-            logger.error(f"Error sending leaderboard update: {e}")
+            logger.debug(
+                f"Leaderboard sent to {len(monitors_copy)} monitors in {(time.time() - send_start)*1000:.1f}ms"
+            )
+            logger.debug(
+                f"Leaderboard broadcast completed in {(time.time() - total_start)*1000:.1f}ms"
+            )
 
-        # Clean up disconnected monitors
-        self.monitors -= disconnected
+        except Exception as e:
+            logger.error(f"Error broadcasting leaderboard: {e}")
+
+    def _get_queue_stats(self) -> Dict[str, int]:
+        """Get queue statistics - synchronous helper for thread pool."""
+        with self.chunk_manager.lock:
+            return {
+                "pending_chunks": len(self.chunk_manager.pending_chunks),
+                "assigned_chunks": sum(
+                    len(chunks) for chunks in self.chunk_manager.assigned_chunks.values()
+                ),
+            }
 
     async def _flush_processed_items(self):
         """Flush batched processed items to chunk tracker."""
@@ -1591,12 +1820,14 @@ class Orchestrator:
             self.last_item_batch_flush = time.time()
 
     def get_workers_by_user_stats(self) -> Dict[str, Any]:
-        """Get statistics about workers grouped by user/token."""
+        """Get statistics about workers grouped by user/token - thread-safe version."""
         if not hasattr(self, "workers_by_user"):
             return {}
 
+        # Create a copy to avoid issues with concurrent modification
         stats = {}
-        for user, worker_ids in self.workers_by_user.items():
+        workers_snapshot = dict(self.workers_by_user)
+        for user, worker_ids in workers_snapshot.items():
             stats[user] = {"worker_count": len(worker_ids), "worker_ids": list(worker_ids)}
         return stats
 
@@ -1621,21 +1852,63 @@ class Orchestrator:
     async def _heartbeat_loop(self):
         """Send periodic heartbeats to maintain connections."""
         while True:
-            await asyncio.sleep(30)
+            try:
+                await asyncio.sleep(30)
 
-            # Ping workers
-            disconnected = []
-            for worker_id, ws in self.workers.items():
-                try:
-                    await ws.ping()
-                except:
-                    disconnected.append(worker_id)
+                # Create a copy of worker items to avoid modification during iteration
+                worker_items = list(self.workers.items())
+                disconnected = []
 
-            # Clean up disconnected workers
-            for worker_id in disconnected:
-                if worker_id in self.workers:
-                    del self.workers[worker_id]
-                    self.chunk_manager.release_worker_chunks(worker_id)
+                for worker_id, ws in worker_items:
+                    try:
+                        # Check if worker still exists before pinging
+                        if worker_id not in self.workers:
+                            continue
+
+                        # Send ping with timeout
+                        pong_waiter = await ws.ping()
+                        try:
+                            await asyncio.wait_for(pong_waiter, timeout=10)
+                        except asyncio.TimeoutError:
+                            logger.warning(f"Worker {worker_id} failed to respond to ping")
+                            disconnected.append(worker_id)
+                    except websockets.exceptions.ConnectionClosed:
+                        logger.info(f"Worker {worker_id} connection already closed")
+                        disconnected.append(worker_id)
+                    except Exception as e:
+                        logger.error(f"Error pinging worker {worker_id}: {e}")
+                        disconnected.append(worker_id)
+
+                # Clean up disconnected workers
+                for worker_id in disconnected:
+                    if worker_id in self.workers:
+                        logger.info(f"Removing unresponsive worker {worker_id}")
+                        del self.workers[worker_id]
+                        self.chunk_manager.release_worker_chunks(worker_id)
+
+                        # Update stats
+                        self.stats["connected_workers"] = len(self.workers)
+
+                        # Also clean up from workers_by_user if it exists
+                        if hasattr(self, "workers_by_user"):
+                            worker_user = (
+                                worker_id.rsplit("_", 1)[0] if "_" in worker_id else worker_id
+                            )
+                            if worker_user in self.workers_by_user:
+                                self.workers_by_user[worker_user].discard(worker_id)
+                                if not self.workers_by_user[worker_user]:
+                                    del self.workers_by_user[worker_user]
+
+                        # Notify monitors
+                        await self._broadcast_stats()
+                        await self._send_activity(
+                            f"Worker {worker_id} removed due to heartbeat timeout"
+                        )
+
+            except Exception as e:
+                logger.error(f"Error in heartbeat loop: {e}", exc_info=True)
+                # Continue the loop even if there's an error
+                await asyncio.sleep(5)
 
     async def _checkpoint_loop(self):
         """Periodically checkpoint storage."""
@@ -1663,7 +1936,10 @@ class Orchestrator:
                 )
 
     async def _stats_update_loop(self):
-        """Periodically update and broadcast stats."""
+        """Periodically update and broadcast stats - non-blocking version."""
+        # Get the event loop for running blocking operations
+        loop = asyncio.get_event_loop()
+
         # Track session start values
         storage_stats = await self.storage.get_storage_stats()
         session_start_outputs = storage_stats["total_captions"]  # This now counts ALL outputs
@@ -1675,8 +1951,8 @@ class Orchestrator:
         while True:
             await asyncio.sleep(10)
 
-            # Update chunk stats
-            chunk_stats = self.chunk_manager.get_stats()
+            # Update chunk stats in thread pool to avoid blocking
+            chunk_stats = await loop.run_in_executor(None, self.chunk_manager.get_stats)
             storage_stats = await self.storage.get_storage_stats()
             current_total_outputs = storage_stats["total_captions"]  # ALL outputs
             if self.chunk_tracker:
@@ -1690,12 +1966,9 @@ class Orchestrator:
             self.stats["total_outputs"] = current_total_outputs
             self.stats["total_captions"] = current_total_outputs  # Keep for backward compatibility
 
-            # Add queue information
-            with self.chunk_manager.lock:
-                self.stats["pending_chunks"] = len(self.chunk_manager.pending_chunks)
-                self.stats["assigned_chunks"] = sum(
-                    len(chunks) for chunks in self.chunk_manager.assigned_chunks.values()
-                )
+            # Get queue stats in thread pool to avoid blocking
+            queue_stats = await loop.run_in_executor(None, self._get_queue_stats)
+            self.stats.update(queue_stats)
 
             # Calculate if we need more chunks
             worker_count = self.stats.get("connected_workers", 0)

caption_flow/utils/dataset_loader.py

@@ -234,6 +234,54 @@ class DatasetLoader:
             for key, url, image_data in ds:
                 yield key, url, image_data
 
+    def _create_dataset_at_position(self, dataset_path: str, split: str, start_idx: int):
+        """Create a dataset iterator positioned at start_idx using state_dict if available."""
+        try:
+            # Load dataset in streaming mode
+            dataset = load_dataset(
+                dataset_path,
+                split=split,
+                streaming=True,
+                token=self.token,
+            )
+
+            # Check if the dataset supports state_dict (newer versions of datasets library)
+            if hasattr(dataset, "load_state_dict") and hasattr(dataset, "state_dict"):
+                # Try to use the dataset's native state management
+                try:
+                    # Get current state
+                    state = dataset.state_dict()
+
+                    # Modify the state to skip to start_idx
+                    if "epoch" in state:
+                        state["epoch"] = 0
+                    if "num_examples_since_previous_state" in state:
+                        state["num_examples_since_previous_state"] = start_idx
+
+                    # For newer datasets with examples_iterable state
+                    if "examples_iterable" in state:
+                        if isinstance(state["examples_iterable"], dict):
+                            if "shard_example_idx" in state["examples_iterable"]:
+                                state["examples_iterable"]["shard_example_idx"] = start_idx
+
+                    # Load the modified state
+                    dataset.load_state_dict(state)
+                    logger.info(f"Positioned dataset at index {start_idx} using state_dict")
+                    return dataset
+                except Exception as e:
+                    logger.debug(f"Could not use state_dict approach: {e}")
+
+            # Fall back to skip() for large skips
+            if start_idx > 0:
+                logger.info(f"Using skip() to position dataset at index {start_idx}")
+                dataset = dataset.skip(start_idx)
+
+            return dataset
+
+        except Exception as e:
+            logger.warning(f"Error creating positioned dataset: {e}")
+            return None
+
     def _iterate_hf_dataset_shard_with_metadata(
         self, shard_url: str, processed_keys: Optional[set] = None
     ) -> Generator[Tuple[str, str, bytes, Dict[str, Any]], None, None]:
@@ -248,7 +296,65 @@ class DatasetLoader:
         )
 
         try:
-            # Load dataset in streaming mode
+            # Try optimized approach for large skips
+            if start_idx > 100:
+                dataset = self._create_dataset_at_position(dataset_path, self.split, start_idx)
+                if dataset:
+                    items_processed = 0
+
+                    for item in dataset:
+                        # Stop after processing chunk_size items
+                        if items_processed >= chunk_size:
+                            break
+
+                        # Generate a unique key for this item
+                        key = f"{dataset_path.replace('/', '_')}_{start_idx + items_processed:08d}"
+
+                        if key in processed_keys:
+                            items_processed += 1
+                            continue
+
+                        try:
+                            # Extract image data
+                            if self.image_column in item:
+                                img_data = item[self.image_column]
+
+                                # Process image to bytes
+                                image_bytes = ImageProcessor.process_image_data(img_data)
+
+                                if image_bytes:
+                                    # Extract all metadata (excluding the image column)
+                                    metadata = {
+                                        k: v for k, v in item.items() if k != self.image_column
+                                    }
+
+                                    # URL is virtual for HF datasets
+                                    url = f"hf://{dataset_path}#{start_idx + items_processed}"
+                                    items_processed += 1
+                                    yield key, url, image_bytes, metadata
+                                else:
+                                    logger.warning(
+                                        f"Failed to process image for item at index {start_idx + items_processed}"
+                                    )
+                                    items_processed += 1
+                                    continue
+                            else:
+                                logger.warning(
+                                    f"No image column '{self.image_column}' found in item at index {start_idx + items_processed}. "
+                                    f"Available columns: {list(item.keys())}"
+                                )
+                                items_processed += 1
+
+                        except Exception as e:
+                            logger.error(
+                                f"Error processing item at index {start_idx + items_processed}: {e}"
+                            )
+                            items_processed += 1
+                            continue
+
+                    return
+
+            # Fall back to regular approach for small skips or if StatefulDataLoader not available
             dataset = load_dataset(
                 dataset_path,
                 split=self.split,
@@ -256,7 +362,7 @@ class DatasetLoader:
                 token=self.token,
             )
 
-            # Skip to start index if needed - CONSISTENT WITH OTHER METHOD
+            # Skip to start index if needed
             if start_idx > 0:
                 dataset = dataset.skip(start_idx)
 
@@ -267,7 +373,7 @@ class DatasetLoader:
                 if items_processed >= chunk_size:
                     break
 
-                # Generate a unique key for this item - CONSISTENT FORMAT
+                # Generate a unique key for this item
                 key = f"{dataset_path.replace('/', '_')}_{start_idx + items_processed:08d}"
 
                 if key in processed_keys:
@@ -337,7 +443,76 @@ class DatasetLoader:
         )
 
         try:
-            # Load dataset in streaming mode
+            # Try optimized approach for large skips
+            if start_idx > 100:
+                dataset = self._create_dataset_at_position(dataset_path, self.split, start_idx)
+                if dataset:
+                    items_processed = 0
+
+                    for item in dataset:
+                        # Stop after processing chunk_size items
+                        if items_processed >= chunk_size:
+                            logger.info(f"Completed chunk: processed {items_processed} items")
+                            break
+
+                        # Also stop if we've reached the dataset end
+                        if (
+                            self._hf_total_items
+                            and (start_idx + items_processed) >= self._hf_total_items
+                        ):
+                            logger.info(
+                                f"Reached dataset end at item {start_idx + items_processed} "
+                                f"(total: {self._hf_total_items})"
+                            )
+                            break
+
+                        # Generate a unique key for this item
+                        key = f"{dataset_path.replace('/', '_')}_{start_idx + items_processed:08d}"
+
+                        if key in processed_keys:
+                            items_processed += 1
+                            continue
+
+                        try:
+                            # Extract image data
+                            if self.image_column in item:
+                                img_data = item[self.image_column]
+
+                                # Delegate image processing to ImageProcessor
+                                image_bytes = ImageProcessor.process_image_data(img_data)
+
+                                if image_bytes:
+                                    # URL is virtual for HF datasets
+                                    url = f"hf://{dataset_path}#{start_idx + items_processed}"
+                                    items_processed += 1
+                                    yield key, url, image_bytes
+                                else:
+                                    logger.warning(
+                                        f"Failed to process image for item at index {start_idx + items_processed}"
+                                    )
+                                    items_processed += 1
+                                    continue
+                            else:
+                                logger.warning(
+                                    f"No image column '{self.image_column}' found in item at index {start_idx + items_processed}. "
+                                    f"Available columns: {list(item.keys())}"
+                                )
+                                items_processed += 1
+
+                        except Exception as e:
+                            logger.error(
+                                f"Error processing item at index {start_idx + items_processed}: {e}"
+                            )
+                            items_processed += 1
+                            continue
+
+                    logger.info(
+                        f"Virtual shard complete: processed {items_processed} items "
+                        f"(start_idx: {start_idx})"
+                    )
+                    return
+
+            # Fall back to regular approach for small skips or if StatefulDataLoader not available
             dataset = load_dataset(
                 dataset_path,
                 split=self.split,

{caption_flow-0.2.0.dist-info → caption_flow-0.2.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.2.0
+Version: 0.2.1
 Summary: Self-contained distributed community captioning system
 Author-email: bghira <bghira@users.github.com>
 License: MIT
@@ -32,6 +32,7 @@ Requires-Dist: pandas<3.0.0,>=2.3.1
 Requires-Dist: arrow<2.0.0,>=1.3.0
 Requires-Dist: datasets<5.0.0,>=4.0.0
 Requires-Dist: boto3<2.0.0,>=1.40.11
+Requires-Dist: torchdata<0.12.0,>=0.11.0
 Provides-Extra: dev
 Requires-Dist: pytest>=7.4.0; extra == "dev"
 Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"

{caption_flow-0.2.0.dist-info → caption_flow-0.2.1.dist-info}/RECORD

@@ -1,8 +1,8 @@
 caption_flow/__init__.py,sha256=NLPJ25lRN7xHqncXweINDNwbt0q8lgjZ30G21zlPdRs,303
-caption_flow/cli.py,sha256=DVVN4e4uL0jL0gRTaIC5BL0DBU2IU_2yUOi4lg6-lEw,28639
+caption_flow/cli.py,sha256=bHxx66CPsCmSieaH3pw8NZBojIIbniRTdU9mEBHMmWA,28832
 caption_flow/models.py,sha256=qo6lQiO10UISbaBVr6Cs-fSW_pmjwE6kmiTmmU_l3Wk,2140
-caption_flow/monitor.py,sha256=MltOwBqcFwni1XEPWu5dIO-os5NKDbH_LInOBXUWHAY,7870
-caption_flow/orchestrator.py,sha256=vLW_w5KuRn9Asy_343DxZDRxiUs0xYgbfuuNGgqIf7k,76403
+caption_flow/monitor.py,sha256=ZZCSasYLKJ-UzA3-RoAtytv-tbNA-m3h5YjlZg_vukg,7870
+caption_flow/orchestrator.py,sha256=bZ8NnGdqoXSmu7Nq-_7cOSH1DLHkBT88cne0uDyPeNY,89112
 caption_flow/storage.py,sha256=hC6ZHT_PHFoUVjqD5JUwy3_79oAD1e1H30neA_xsz7s,40748
 caption_flow/utils/__init__.py,sha256=F1BChVoCsj9zn1GJRBOLHET1kLW6xrAmsbzcR7hHy6Y,202
 caption_flow/utils/auth.py,sha256=UrxX2n8OEEcfMD1Ey27TxGfrJFmUCpC59x-SCrQJoVE,2253
@@ -10,7 +10,7 @@ caption_flow/utils/caption_utils.py,sha256=esUMAdcCkNjRroZ0Bhxv0_yKlLtMf0XeDCTt-
 caption_flow/utils/certificates.py,sha256=eu4blQZEkL9NRaY1ynQWg1asvDorRYhGRZea7STonJE,4635
 caption_flow/utils/checkpoint_tracker.py,sha256=8tsTFF-HcygitK92YcS-QWzeg-qRm9AuCpQoQRfC8M0,3335
 caption_flow/utils/chunk_tracker.py,sha256=hKn8CN6ubErc9kuCWZMj12ZCZKxVlqXqAEocbzjfa-k,17296
-caption_flow/utils/dataset_loader.py,sha256=qjoRuPnCv_2nGPfrdqf45AgBXlthw1HwqZ1IqwIXzH4,20792
+caption_flow/utils/dataset_loader.py,sha256=ZplJv655ZMyUbaZC4BBiL5II18sBy4JSJhxGZtK_VmA,29107
 caption_flow/utils/image_processor.py,sha256=Zl8TAv9gYPdAYat3UiTuuNdIb2fXNfZ35AxsxuovJTs,5650
 caption_flow/utils/job_queue.py,sha256=itdfXcrkvGjmXn4qtpgMF63k1ufRBaejDe4V6WcxzgU,1104
 caption_flow/utils/json_utils.py,sha256=IiZYn8uCM-3pYmyIbX2fmaOIyutArn67SqAyp0ggNpU,5396
@@ -21,9 +21,9 @@ caption_flow/utils/vllm_config.py,sha256=TC7Rmjk0zRKbBXbWUXrFL4Z58hzax_-4L0pXZn0
 caption_flow/workers/base.py,sha256=jPm_Xw4Lxd0cnrPs-biBqKRQKkTOJLvHLolmp0Gb1CI,7530
 caption_flow/workers/caption.py,sha256=NZ9kTjk2uOoNwyyNSkB_arYk213vLr5mowHN-OjiFkk,54631
 caption_flow/workers/data.py,sha256=0Tg8NE0wdONeMlivYQ4nvbcfWdLuU51O7vR8_YSnJgo,14813
-caption_flow-0.2.0.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-caption_flow-0.2.0.dist-info/METADATA,sha256=6qwt05U0S23Omjz1yR6VzLq_wRHbRx_xl3YzhwHyDLc,11900
-caption_flow-0.2.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-caption_flow-0.2.0.dist-info/entry_points.txt,sha256=KnVlyrGKZj6p2zNyuEnCx4Y6jvJ4V-mcfN0lddPKTlQ,55
-caption_flow-0.2.0.dist-info/top_level.txt,sha256=_bXpKRutqded0FQ80dCChIz26ETV7tL4d4e2E_Y1FXs,13
-caption_flow-0.2.0.dist-info/RECORD,,
+caption_flow-0.2.1.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+caption_flow-0.2.1.dist-info/METADATA,sha256=fxNfSOqkCklb96aq3ZFU7SvRuXEBUQ11xbjkQn7Yzuo,11941
+caption_flow-0.2.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+caption_flow-0.2.1.dist-info/entry_points.txt,sha256=KnVlyrGKZj6p2zNyuEnCx4Y6jvJ4V-mcfN0lddPKTlQ,55
+caption_flow-0.2.1.dist-info/top_level.txt,sha256=_bXpKRutqded0FQ80dCChIz26ETV7tL4d4e2E_Y1FXs,13
+caption_flow-0.2.1.dist-info/RECORD,,