caption-flow 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

caption_flow/__init__.py

@@ -1,6 +1,6 @@
 """CaptionFlow - Distributed community captioning system."""
 
-__version__ = "0.3.2"
+__version__ = "0.3.4"
 
 from .orchestrator import Orchestrator
 from .workers.data import DataWorker

caption_flow/cli.py

@@ -124,7 +124,7 @@ def setup_logging(verbose: bool = False):
     level = logging.DEBUG if verbose else logging.INFO
     logging.basicConfig(
         level=level,
-        format="%(message)s",
+        format="%(name)s: %(message)s",
         datefmt="[%Y-%m-%d %H:%M:%S]",
         handlers=[
             RichHandler(
@@ -490,7 +490,9 @@ def reload_config(
 
     async def send_reload():
         try:
-            async with websockets.connect(server, ssl=ssl_context) as websocket:
+            async with websockets.connect(
+                server, ssl=ssl_context, ping_interval=20, ping_timeout=60, close_timeout=10
+            ) as websocket:
                 # Authenticate as admin
                 await websocket.send(json.dumps({"token": token, "role": "admin"}))
 

caption_flow/monitor.py

@@ -73,6 +73,9 @@ class Monitor:
                 async with websockets.connect(
                     self.server_url,
                     ssl=self.ssl_context if self.server_url.startswith("wss://") else None,
+                    ping_interval=20,
+                    ping_timeout=60,
+                    close_timeout=10,
                 ) as websocket:
                     # Authenticate
                     await websocket.send(json.dumps({"token": self.token}))

caption_flow/orchestrator.py

@@ -124,13 +124,14 @@ class Orchestrator:
 
         # Initialize storage
         await self.storage.initialize()
-        await self.update_unprocessed_ranges()
 
         # Start background tasks
         asyncio.create_task(self._heartbeat_loop())
         asyncio.create_task(self._checkpoint_loop())
         asyncio.create_task(self._stats_update_loop())
 
+        await self.update_unprocessed_ranges()
+
         # Start WebSocket server
         websocket_logger = logging.getLogger("websockets")
         websocket_logger.setLevel(logging.WARNING)
@@ -376,16 +377,17 @@ class Orchestrator:
         """Process results submission from worker."""
         # Extract user from worker_id
         worker_user = worker_id.rsplit("_", 1)[0] if "_" in worker_id else worker_id
+
         # Create work result
         _job_id = data.get("job_id")
         job_id = JobId.from_str(_job_id)
-        shard_name = job_id.shard_id  # >data-0000<
-        chunk_name = job_id.chunk_id  # data-0000:chunk:>0<
-        # logger.debug(f"({job_id}) Worker result: {data}")
+        shard_name = job_id.shard_id
+        chunk_name = job_id.chunk_id
+
         result = WorkResult(
             unit_id=data["unit_id"],
             source_id=shard_name,
-            chunk_id=job_id.get_chunk_str(),  # we want the full string here
+            chunk_id=job_id.get_chunk_str(),
             sample_id=data["sample_id"],
             dataset=data["dataset"],
             outputs=data["outputs"],
@@ -393,7 +395,9 @@ class Orchestrator:
             processing_time_ms=data.get("processing_time_ms", 0),
         )
 
-        # Let processor handle any custom processing
+        # Let processor handle any custom processing - this updates chunk tracker
+        # IMPORTANT: Call this BEFORE saving to storage so chunk tracker is updated
+        # regardless of whether the item is a duplicate
         processed = self.processor.handle_result(result)
 
         # Create caption record for storage
@@ -411,6 +415,7 @@ class Orchestrator:
         for key in to_delete_metadata_keys:
             if key in result.metadata:
                 del result.metadata[key]
+
         caption = Caption(
             job_id=job_id,
             dataset=result.dataset,
@@ -432,14 +437,15 @@ class Orchestrator:
             image_format=image_format,
         )
 
-        # Save to storage
-        await self.storage.save_caption(caption)
+        # Save to storage (might skip if duplicate)
+        saved = await self.storage.save_caption(caption)
 
-        # Update contributor stats
-        contributor = await self.storage.get_contributor(worker_user)
-        if contributor:
-            contributor.total_captions += total_outputs
-            await self.storage.save_contributor(contributor)
+        # Update contributor stats only if actually saved
+        if saved:
+            contributor = await self.storage.get_contributor(worker_user)
+            if contributor:
+                contributor.total_captions += total_outputs
+                await self.storage.save_contributor(contributor)
 
     async def _handle_monitor(self, websocket: WebSocketServerProtocol):
         """Handle monitor connection."""
@@ -839,39 +845,55 @@ class Orchestrator:
         self.monitors -= disconnected
 
     async def _heartbeat_loop(self):
-        """Send periodic heartbeats to maintain connections."""
+        """Collect and log worker status periodically."""
         while True:
             await asyncio.sleep(30)
 
-            disconnected = []
+            # Just collect status - no ping/pong
+            active_workers = []
             for worker_id, ws in list(self.workers.items()):
-                try:
-                    pong_waiter = await ws.ping()
-                    await asyncio.wait_for(pong_waiter, timeout=10)
-                except:
-                    disconnected.append(worker_id)
-
-            # Clean up disconnected workers
-            for worker_id in disconnected:
-                logger.warning(f"Worker {worker_id} did not respond to ping, disconnecting")
-                if worker_id in self.workers:
+                # Check if WebSocket is still open (don't ping)
+                if ws.state == websockets.protocol.State.OPEN:
+                    active_workers.append(worker_id)
+                else:
+                    # Clean up closed connections
+                    logger.info(f"Worker {worker_id} connection closed")
                     del self.workers[worker_id]
-                    logger.warning(
-                        f"Releasing assignments for worker {worker_id} because it did not respond to ping"
-                    )
                     self.processor.release_assignments(worker_id)
-                    self.stats["connected_workers"] = len(self.workers)
+
+            # Log status
+            if active_workers:
+                logger.debug(
+                    f"Active workers: {len(active_workers)} - {', '.join(active_workers[:5])}"
+                )
+                logger.debug(f"Inactive workers: {len(self.workers) - len(active_workers)}")
+            # add to self.stats
+            self.stats["active_workers"] = len(active_workers)
+            self.stats["inactive_workers"] = len(self.workers) - len(active_workers)
 
     async def _checkpoint_loop(self):
-        """Periodically checkpoint storage."""
+        """Periodically checkpoint storage and chunk tracker."""
         interval = self.config.get("storage", {}).get("checkpoint_interval", 60)
 
         while True:
             await asyncio.sleep(interval)
 
-            await self.storage.checkpoint()
-            self.stats["last_checkpoint"] = datetime.utcnow().isoformat()
-            logger.info("Storage checkpoint complete")
+            try:
+                # Checkpoint storage
+                await self.storage.checkpoint()
+
+                # Also checkpoint the chunk tracker if using webdataset processor
+                if hasattr(self.processor, "chunk_tracker") and self.processor.chunk_tracker:
+                    # Save checkpoint in thread pool to avoid blocking
+                    await asyncio.get_event_loop().run_in_executor(
+                        None, self.processor.chunk_tracker.save
+                    )
+                    logger.debug("Saved chunk tracker checkpoint")
+
+                self.stats["last_checkpoint"] = datetime.utcnow().isoformat()
+                logger.info("Storage and chunk tracker checkpoint complete")
+            except Exception as e:
+                logger.error(f"Error during checkpoint: {e}", exc_info=True)
 
     async def _stats_update_loop(self):
         """Periodically update and broadcast stats."""

caption_flow/processors/huggingface.py

@@ -551,7 +551,7 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
 
             # Force checkpoint save if needed
             if self.chunk_tracker:
-                self.chunk_tracker.save_checkpoint()
+                self.chunk_tracker.save()
 
     def get_work_units(self, count: int, worker_id: str) -> List[WorkUnit]:
         """Get available work units for a worker."""
@@ -717,7 +717,7 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
 
         # Save final state
         if self.chunk_tracker:
-            self.chunk_tracker.save_checkpoint()
+            self.chunk_tracker.save()
 
 
 class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):

caption_flow/processors/webdataset.py

@@ -110,58 +110,86 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
         return self.shard_info_cache[shard_idx]
 
     def _restore_state(self, storage: StorageManager) -> None:
-        """Restore state from chunk tracker."""
-        logger.debug("Restoring state from chunk tracker")
+        """Restore state from chunk tracker and synchronize with storage."""
+        logger.info("Restoring state from chunk tracker and synchronizing with storage")
         if not self.chunk_tracker:
             return
 
+        # First, update chunk tracker from storage
+        processed_job_ids = storage.get_all_processed_job_ids()
+        if processed_job_ids:
+            logger.info(
+                f"Synchronizing chunk tracker with {len(processed_job_ids)} processed items from storage"
+            )
+            self.update_from_storage(processed_job_ids)
+
+        # Then restore work units from chunk tracker
         shards_summary = self.chunk_tracker.get_shards_summary()
-        logger.debug(f"Restoring state: {shards_summary}")
+        logger.info(f"Restoring work units from chunk tracker: {len(shards_summary)} shards")
 
         with self.lock:
+            restored_count = 0
             for shard_name, shard_info in shards_summary.items():
                 chunks = shard_info.get("chunks", [])
-                logger.debug(f"Existing job ids: {storage.get_all_processed_job_ids()}")
                 for chunk_state in chunks:
                     # Only add incomplete chunks
-                    if chunk_state.status != "completed":
-                        logger.debug(f"Restoring incomplete chunk {chunk_state}")
+                    if chunk_state.status == "completed":
+                        logger.debug(f"Skipping completed chunk {chunk_state.chunk_id}")
+                        continue
 
-                        # Get unprocessed ranges
-                        unprocessed_ranges = chunk_state.get_unprocessed_ranges()
+                    # Get unprocessed ranges
+                    unprocessed_ranges = chunk_state.get_unprocessed_ranges()
+                    if not unprocessed_ranges:
                         logger.debug(
-                            f"Chunk {chunk_state.chunk_id} unprocessed ranges: {unprocessed_ranges}"
+                            f"Chunk {chunk_state.chunk_id} has no unprocessed ranges, marking as completed"
                         )
-                        if not unprocessed_ranges:
-                            continue
+                        self.chunk_tracker.mark_completed(chunk_state.chunk_id)
+                        continue
 
-                        # Convert relative ranges to absolute file indices
-                        absolute_ranges = []
-                        for start, end in unprocessed_ranges:
-                            abs_start = chunk_state.start_index + start
-                            abs_end = chunk_state.start_index + end
-                            absolute_ranges.append((abs_start, abs_end))
+                    logger.info(
+                        f"Restoring chunk {chunk_state.chunk_id} with unprocessed ranges: {unprocessed_ranges}"
+                    )
 
-                        unit = WorkUnit(
-                            unit_id=chunk_state.chunk_id,
-                            chunk_id=chunk_state.chunk_id,
-                            source_id=shard_name,
-                            unit_size=chunk_state.chunk_size,
-                            data={
-                                "shard_url": chunk_state.shard_url,
-                                "shard_name": shard_name,
-                                "start_index": chunk_state.start_index,
-                                "chunk_size": chunk_state.chunk_size,
-                                "unprocessed_ranges": absolute_ranges,
-                            },
-                            metadata={
-                                "shard_name": shard_name,
-                                "chunk_index": chunk_state.start_index // self.chunk_size,
-                            },
-                        )
+                    # Convert relative ranges to absolute file indices
+                    absolute_ranges = []
+                    for start, end in unprocessed_ranges:
+                        abs_start = chunk_state.start_index + start
+                        abs_end = chunk_state.start_index + end
+                        absolute_ranges.append((abs_start, abs_end))
+
+                    # Get shard index if available
+                    shard_idx = None
+                    if self.dataset:
+                        for idx in range(self.dataset.num_shards):
+                            shard_info = self._get_shard_info_cached(idx)
+                            if shard_info and shard_info["name"] == shard_name:
+                                shard_idx = idx
+                                break
 
-                        self.work_units[unit.unit_id] = unit
-                        self.pending_units.append(unit.unit_id)
+                    unit = WorkUnit(
+                        unit_id=chunk_state.chunk_id,
+                        chunk_id=chunk_state.chunk_id,
+                        source_id=shard_name,
+                        unit_size=chunk_state.chunk_size,
+                        data={
+                            "shard_url": chunk_state.shard_url,
+                            "shard_name": shard_name,
+                            "shard_idx": shard_idx,
+                            "start_index": chunk_state.start_index,
+                            "chunk_size": chunk_state.chunk_size,
+                            "unprocessed_ranges": absolute_ranges,
+                        },
+                        metadata={
+                            "shard_name": shard_name,
+                            "chunk_index": chunk_state.start_index // self.chunk_size,
+                        },
+                    )
+
+                    self.work_units[unit.unit_id] = unit
+                    self.pending_units.append(unit.unit_id)
+                    restored_count += 1
+
+            logger.info(f"Restored {restored_count} incomplete work units")
 
     def _create_units_background(self) -> None:
         """Background thread to create work units on demand."""
@@ -278,8 +306,15 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
         assigned = []
 
         with self.lock:
-            while len(assigned) < count and self.pending_units:
+            units_checked = 0
+            max_units_to_check = len(self.pending_units)
+
+            while len(assigned) < count and units_checked < max_units_to_check:
+                if not self.pending_units:
+                    break
+
                 unit_id = self.pending_units.popleft()
+                units_checked += 1
                 unit = self.work_units.get(unit_id)
 
                 if unit:
@@ -288,6 +323,16 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
                         chunk_state = self.chunk_tracker.chunks[unit_id]
                         relative_unprocessed = chunk_state.get_unprocessed_ranges()
 
+                        # If no unprocessed ranges, mark as completed and skip
+                        if not relative_unprocessed:
+                            logger.info(
+                                f"Chunk {unit_id} has no unprocessed ranges, marking as completed"
+                            )
+                            self.chunk_tracker.mark_completed(unit_id)
+                            # Remove from work units
+                            del self.work_units[unit_id]
+                            continue
+
                         # Convert relative to absolute indices
                         absolute_ranges = []
                         for start, end in relative_unprocessed:
@@ -307,6 +352,9 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
 
                     if self.chunk_tracker:
                         self.chunk_tracker.mark_assigned(unit_id, worker_id)
+                else:
+                    # Put it back if we couldn't get the unit
+                    self.pending_units.append(unit_id)
 
         logger.debug(f"Assigned {len(assigned)} units to worker {worker_id}")
         return assigned
@@ -366,8 +414,20 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
         logger.info(f"Released {len(unit_ids)} assignments from {worker_id}")
 
     def handle_result(self, result: WorkResult) -> Dict[str, Any]:
-        """Handle result from worker."""
-        # Track processed items if we have chunk tracker
+        """Handle result from worker and update chunk tracker."""
+        # Extract the actual item index from the metadata
+        item_index = result.metadata.get("_item_index", None)
+
+        # If we have an item index, mark it as processed in the chunk tracker
+        if self.chunk_tracker and item_index is not None and result.chunk_id:
+            try:
+                # Mark single item as processed
+                self.chunk_tracker.mark_items_processed(result.chunk_id, item_index, item_index)
+                # logger.debug(f"Marked item {item_index} as processed in chunk {result.chunk_id}")
+            except Exception as e:
+                logger.error(f"Error marking item {item_index} as processed: {e}")
+
+        # Also handle batch results if present (backward compatibility)
         if self.chunk_tracker and "item_indices" in result.metadata:
             indices = result.metadata["item_indices"]
 
@@ -391,6 +451,9 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
                 # Mark ranges as processed
                 for start_idx, end_idx in ranges:
                     self.chunk_tracker.mark_items_processed(result.chunk_id, start_idx, end_idx)
+                    logger.debug(
+                        f"Marked range {start_idx}-{end_idx} as processed in chunk {result.chunk_id}"
+                    )
 
         return {
             "source_id": result.source_id,
@@ -407,22 +470,46 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
             # Group by chunk
             processed_by_chunk = defaultdict(set)
 
-            for job_id in processed_job_ids:
-                # Parse job_id to extract chunk and index
-                # Expected format: "shard:chunk:X:idx:Y"
-                parts = job_id.split(":")
-                if len(parts) >= 5 and parts[3] == "idx":
-                    chunk_id = ":".join(parts[:3])  # "shard:chunk:X"
-                    try:
-                        idx = int(parts[4])
-                        processed_by_chunk[chunk_id].add(idx)
-                    except ValueError:
-                        continue
+            for job_id_str in processed_job_ids:
+                try:
+                    # Use JobId to parse the job ID string
+                    job_id = JobId.from_str(job_id_str)
+                    chunk_id = job_id.get_chunk_str()
+                    sample_idx = int(job_id.sample_id)
+                    processed_by_chunk[chunk_id].add(sample_idx)
+                except ValueError as e:
+                    logger.warning(f"Invalid job ID format: {job_id_str} - {e}")
+                    continue
 
             # Update chunk tracker with processed items
             if self.chunk_tracker:
                 for chunk_id, indices in processed_by_chunk.items():
                     if indices:
+                        # Get or create chunk state
+                        chunk_state = self.chunk_tracker.chunks.get(chunk_id)
+                        if not chunk_state:
+                            # Parse chunk_id using JobId to get shard info
+                            try:
+                                # chunk_id format: "shard_id:chunk:chunk_idx"
+                                parts = chunk_id.split(":")
+                                if len(parts) >= 3:
+                                    shard_name = parts[0]
+                                    chunk_idx = int(parts[2])
+                                    # Infer start index from chunk index and size
+                                    start_index = chunk_idx * self.chunk_size
+                                    # Create chunk state
+                                    self.chunk_tracker.add_chunk(
+                                        chunk_id,
+                                        shard_name,
+                                        f"{shard_name}.tar",
+                                        start_index,
+                                        self.chunk_size,
+                                    )
+                                    logger.info(f"Created missing chunk state for {chunk_id}")
+                            except (ValueError, IndexError) as e:
+                                logger.error(f"Failed to create chunk state for {chunk_id}: {e}")
+                                continue
+
                         # Sort indices and convert to ranges
                         sorted_indices = sorted(indices)
                         if not sorted_indices:
@@ -443,10 +530,13 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
                         ranges.append((start_range, end_range))
 
                         # Mark each contiguous range as processed
-                        logger.debug(f"Marking ranges {ranges} as processed in chunk {chunk_id}")
+                        logger.info(f"Marking ranges {ranges} as processed in chunk {chunk_id}")
                         for start_idx, end_idx in ranges:
                             self.chunk_tracker.mark_items_processed(chunk_id, start_idx, end_idx)
 
+                # Save checkpoint after updating
+                self.chunk_tracker.save()
+
     def get_stats(self) -> Dict[str, Any]:
         """Get processor statistics."""
         with self.lock:
@@ -484,7 +574,7 @@ class WebDatasetOrchestratorProcessor(OrchestratorProcessor):
 
         # Save checkpoint
         if self.chunk_tracker:
-            self.chunk_tracker.save_checkpoint()
+            self.chunk_tracker.save()
 
 
 class WebDatasetWorkerProcessor(WorkerProcessor):
@@ -555,7 +645,11 @@ class WebDatasetWorkerProcessor(WorkerProcessor):
             # Generate mock results for unprocessed ranges
             for start_idx, end_idx in unprocessed_ranges:
                 for idx in range(start_idx, end_idx + 1):
-                    job_id = f"{shard_name}:chunk:{chunk_index}:idx:{idx}"
+                    # Use JobId to create consistent job ID
+                    job_id = JobId.from_values(
+                        shard_id=shard_name, chunk_id=str(chunk_index), sample_id=str(idx)
+                    )
+                    job_id_str = job_id.get_sample_str()
 
                     yield {
                         "image": self._create_mock_image(idx),
@@ -565,11 +659,11 @@ class WebDatasetWorkerProcessor(WorkerProcessor):
                         "metadata": {
                             "_item_index": idx,
                             "_chunk_relative_index": idx - unit.data["start_index"],
-                            "_job_id": job_id,
+                            "_job_id": job_id_str,
                             "_mock": True,
                             "_processed_indices": processed_indices,
                         },
-                        "job_id": job_id,
+                        "job_id": job_id_str,
                     }
 
                     processed_indices.append(idx)
@@ -614,8 +708,11 @@ class WebDatasetWorkerProcessor(WorkerProcessor):
                                         f"Error decoding image {entry.path} with cv2: {img_e}"
                                     )
 
-                            # Generate job ID compatible with chunk tracker
-                            job_id = f"{shard_name}:chunk:{chunk_index}:idx:{idx}"
+                            # Generate job ID using JobId class
+                            job_id = JobId.from_values(
+                                shard_id=shard_name, chunk_id=str(chunk_index), sample_id=str(idx)
+                            )
+                            job_id_str = job_id.get_sample_str()
 
                             yield {
                                 "image": image,
@@ -625,12 +722,12 @@ class WebDatasetWorkerProcessor(WorkerProcessor):
                                 "metadata": {
                                     "_item_index": idx,
                                     "_chunk_relative_index": idx - unit.data["start_index"],
-                                    "_job_id": job_id,
+                                    "_job_id": job_id_str,
                                     "_filename": entry.path,
                                     "_file_size": entry.size,
                                     "_processed_indices": processed_indices,
                                 },
-                                "job_id": job_id,
+                                "job_id": job_id_str,
                             }
 
                             processed_indices.append(idx)

caption_flow/utils/checkpoint_tracker.py

@@ -6,6 +6,7 @@ from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Dict, Any, Optional
 from datetime import datetime
+from concurrent.futures import ThreadPoolExecutor
 
 logger = logging.getLogger(__name__)
 
@@ -52,35 +53,54 @@ class CheckpointTracker(ABC):
 
     def save(self) -> None:
         """Save checkpoint to disk atomically."""
-        try:
-            # Prepare data with metadata
-            data = self._serialize_state()
-            data["updated_at"] = datetime.utcnow().isoformat()
+        with self.lock:
+            try:
+                # Prepare data with metadata
+                data = self._serialize_state()
+                data["updated_at"] = datetime.utcnow().isoformat()
+
+                # Write atomically using temp file
+                tmp_file = self.checkpoint_path.with_suffix(".tmp")
+                # If a save is already in progress, let it finish.
+                # This prevents race conditions if save() is called rapidly.
+                if (
+                    hasattr(self, "_save_future")
+                    and self._save_future
+                    and not self._save_future.done()
+                ):
+                    self._save_future.result()  # Wait for the previous save to complete
+
+                # Use an executor to run the save operation in a background thread.
+                # This makes the save call non-blocking.
+                with ThreadPoolExecutor(max_workers=1) as executor:
+                    data_to_save = data.copy()
+                    self._save_future = executor.submit(self._write_to_disk, data_to_save, tmp_file)
+            except Exception as e:
+                logger.error(f"Failed to submit save task: {e}", exc_info=True)
 
-            # Write atomically using temp file
-            tmp_file = self.checkpoint_path.with_suffix(".tmp")
+    def _write_to_disk(self, data: Dict[str, Any]) -> None:
+        """Write checkpoint data to disk atomically."""
+        # Create a temporary file in the same directory as the checkpoint
+        tmp_file = self.checkpoint_path.with_suffix(".tmp")
+
+        try:
+            # Ensure the parent directory exists
+            self.checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
 
             with open(tmp_file, "w") as f:
                 json.dump(data, f, indent=2)
 
-            # Ensure temp file was created
-            if not tmp_file.exists():
-                raise IOError(f"Failed to create temporary file: {tmp_file}")
-
-            # Move atomically
+            # Atomically replace the checkpoint file
             tmp_file.replace(self.checkpoint_path)
-
             logger.debug(f"Saved checkpoint to {self.checkpoint_path}")
-
         except Exception as e:
-            # logger.error(f"Error saving checkpoint: {e}", exc_info=True)
-            # Try direct write as fallback
-            try:
-                with open(self.checkpoint_path, "w") as f:
-                    json.dump(data, f, indent=2)
-                # logger.info("Saved checkpoint using fallback direct write")
-            except Exception as fallback_error:
-                logger.error(f"Fallback save also failed: {fallback_error}")
+            logger.error(f"Failed to save checkpoint atomically: {e}", exc_info=True)
+            # Try to clean up the temp file if it exists
+            if tmp_file.exists():
+                try:
+                    tmp_file.unlink()
+                except:
+                    pass
 
     def get_stats(self) -> Dict[str, Any]:
         """Get statistics about tracked items. Override for custom stats."""

caption_flow/utils/chunk_tracker.py

@@ -8,6 +8,7 @@ from datetime import datetime, timedelta
 from dataclasses import dataclass, asdict, field
 
 from .checkpoint_tracker import CheckpointTracker
+from threading import Lock
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -60,12 +61,12 @@ class ChunkState:
         self.status = "completed"
         self.completed_at = datetime.utcnow()
         # Clear processed_ranges since we don't need them after completion
-        self.processed_ranges = []
-        self.assigned_to = None
-        self.assigned_at = None
+        # self.processed_ranges = []
+        # self.assigned_to = None
+        # self.assigned_at = None
 
     def get_unprocessed_ranges(self) -> List[Tuple[int, int]]:
-        """Get ranges that haven't been processed yet."""
+        """Get ranges of unprocessed items within the chunk (relative indices)."""
         if self.status == "completed":
             return []
 
@@ -73,22 +74,57 @@ class ChunkState:
             logger.info(f"Chunk {self.chunk_id} has no processed ranges, returning full range")
             return [(0, self.chunk_size - 1)]
 
+        # Merge ranges first to ensure no overlaps
+        merged_ranges = self._merge_ranges(self.processed_ranges)
+
         unprocessed = []
-        current = 0
+        current_pos = 0
 
-        logger.info(
-            f"Processing {len(self.processed_ranges)} processed ranges for chunk {self.chunk_id}"
-        )
-        for start, end in self.processed_ranges:
-            if current < start:
-                unprocessed.append((current, start - 1))
-            current = max(current, end + 1)
+        for start, end in merged_ranges:
+            if current_pos < start:
+                unprocessed.append((current_pos, start - 1))
+            current_pos = max(current_pos, end + 1)
 
-        if current < self.chunk_size:
-            unprocessed.append((current, self.chunk_size - 1))
+        # Add any remaining range
+        if current_pos < self.chunk_size:
+            unprocessed.append((current_pos, self.chunk_size - 1))
+
+        # Log for debugging
+        if not unprocessed:
+            logger.info(
+                f"Chunk {self.chunk_id} has processed ranges {merged_ranges} covering entire chunk size {self.chunk_size}"
+            )
+        else:
+            total_processed = sum(end - start + 1 for start, end in merged_ranges)
+            total_unprocessed = sum(end - start + 1 for start, end in unprocessed)
+            logger.debug(
+                f"Chunk {self.chunk_id}: {total_processed} processed, {total_unprocessed} unprocessed"
+            )
 
         return unprocessed
 
+    def _merge_ranges(self, ranges: List[Tuple[int, int]]) -> List[Tuple[int, int]]:
+        """Merge overlapping or adjacent ranges."""
+        if not ranges:
+            return []
+
+        # Sort ranges by start index, ensuring all are tuples
+        sorted_ranges = sorted([tuple(r) for r in ranges])
+        merged = [sorted_ranges[0]]
+
+        for current_start, current_end in sorted_ranges[1:]:
+            last_start, last_end = merged[-1]
+
+            # Check if ranges overlap or are adjacent
+            if current_start <= last_end + 1:
+                # Merge the ranges
+                merged[-1] = (last_start, max(last_end, current_end))
+            else:
+                # Add as new range
+                merged.append((current_start, current_end))
+
+        return merged
+
     def to_dict(self):
         """Convert to dictionary for JSON serialization."""
         d = asdict(self)
@@ -124,6 +160,7 @@ class ChunkTracker(CheckpointTracker):
         self.max_completed_chunks_in_memory = max_completed_chunks_in_memory
         self.archive_after_hours = archive_after_hours
         self._completed_count = 0  # Track count without storing all IDs
+        self.lock = Lock()
         super().__init__(checkpoint_file)
 
     def _get_default_state(self) -> Dict[str, Any]:
@@ -132,16 +169,17 @@ class ChunkTracker(CheckpointTracker):
 
     def _deserialize_state(self, data: Dict[str, Any]) -> None:
         """Deserialize loaded data into instance state."""
-        self.chunks = {}
-        self._completed_count = data.get("completed_count", 0)
-
-        # Load chunk states
-        completed_chunks = 0
-        for chunk_id, chunk_data in data.get("chunks", {}).items():
-            chunk_state = ChunkState.from_dict(chunk_data)
-            self.chunks[chunk_id] = chunk_state
-            if chunk_state.status == "completed":
-                completed_chunks += 1
+        with self.lock:
+            self.chunks = {}
+            self._completed_count = data.get("completed_count", 0)
+
+            # Load chunk states
+            completed_chunks = 0
+            for chunk_id, chunk_data in data.get("chunks", {}).items():
+                chunk_state = ChunkState.from_dict(chunk_data)
+                self.chunks[chunk_id] = chunk_state
+                if chunk_state.status == "completed":
+                    completed_chunks += 1
 
         logger.info(
             f"Loaded {len(self.chunks)} chunks from checkpoint, "
@@ -494,40 +532,40 @@ class ChunkTracker(CheckpointTracker):
                 for start_idx, end_idx in ranges:
                     chunk.add_processed_range(start_idx, end_idx)
 
-    def mark_items_processed(self, chunk_id: str, start_idx: int, end_idx: int):
-        """Mark a range of items as processed within a chunk (expects ABSOLUTE indices)."""
+    def mark_items_processed(self, chunk_id: str, start_idx: int, end_idx: int) -> None:
+        """Mark a range of items as processed within a chunk."""
         if chunk_id not in self.chunks:
-            logger.error(f"Unknown chunk: {chunk_id}")
+            logger.warning(f"Chunk {chunk_id} not found in tracker")
             return
 
-        chunk = self.chunks[chunk_id]
+        chunk_state = self.chunks[chunk_id]
 
-        # Convert absolute indices to chunk-relative
-        relative_start = start_idx - chunk.start_index
-        relative_end = end_idx - chunk.start_index
+        # Convert absolute indices to chunk-relative indices
+        relative_start = start_idx - chunk_state.start_index
+        relative_end = end_idx - chunk_state.start_index
 
-        # Validate boundaries
-        if relative_start < 0 or relative_end >= chunk.chunk_size:
-            logger.error(
-                f"Invalid indices for chunk {chunk_id}: "
-                f"absolute {start_idx}-{end_idx} (relative {relative_start}-{relative_end}) "
-                f"outside chunk bounds [{chunk.start_index}, {chunk.start_index + chunk.chunk_size - 1}]"
-            )
-            return
+        # Ensure indices are within chunk bounds
+        relative_start = max(0, relative_start)
+        relative_end = min(chunk_state.chunk_size - 1, relative_end)
 
-        # Add the relative range
-        chunk.add_processed_range(relative_start, relative_end)
+        # Add to processed ranges
+        chunk_state.processed_ranges.append((relative_start, relative_end))
 
-        # If chunk is now complete, increment counter
-        if chunk.status == "completed":
-            self._completed_count += 1
+        # Merge overlapping ranges
+        chunk_state.processed_ranges = chunk_state._merge_ranges(chunk_state.processed_ranges)
 
-        self.save()
         logger.debug(
-            f"Marked items {start_idx}-{end_idx} as processed in chunk {chunk_id} "
-            f"(relative indices: {relative_start}-{relative_end})"
+            f"Marked items {start_idx}-{end_idx} as processed in chunk {chunk_id} (relative indices: {relative_start}-{relative_end})"
         )
 
+        # Check if chunk is now complete
+        if chunk_state.get_unprocessed_ranges() == []:
+            logger.info(f"Chunk {chunk_id} is now complete")
+            chunk_state.status = "completed"
+
+        # Save checkpoint after updating
+        self.save()
+
     def get_chunk_with_unprocessed_items(self, chunk_id: str) -> Optional[Dict[str, Any]]:
         """Get chunk info with unprocessed item ranges."""
         chunk_state = self.chunks.get(chunk_id)

caption_flow/workers/base.py

@@ -89,8 +89,13 @@ class BaseWorker(ABC):
     async def _connect_and_run(self):
         """Connect to orchestrator and run main loop."""
         logger.info(f"Connecting to {self.server_url}")
-
-        async with websockets.connect(self.server_url, ssl=self.ssl_context) as websocket:
+        async with websockets.connect(
+            self.server_url,
+            ssl=self.ssl_context,
+            ping_interval=20,
+            ping_timeout=60,
+            close_timeout=10,
+        ) as websocket:
             self.websocket = websocket
             self.connected.set()
 

caption_flow/workers/caption.py

@@ -248,7 +248,13 @@ class CaptionWorker(BaseWorker):
     async def _initial_connect_for_config(self):
         """Connect initially just to get configuration."""
         logger.info(f"Connecting to {self.server_url}")
-        async with websockets.connect(self.server_url, ssl=self.ssl_context) as websocket:
+        async with websockets.connect(
+            self.server_url,
+            ssl=self.ssl_context,
+            ping_interval=20,
+            ping_timeout=60,
+            close_timeout=10,
+        ) as websocket:
             await websocket.send(json.dumps(self._get_auth_data()))
 
             welcome = await websocket.recv()

{caption_flow-0.3.2.dist-info → caption_flow-0.3.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.3.2
+Version: 0.3.4
 Summary: Self-contained distributed community captioning system
 Author-email: bghira <bghira@users.github.com>
 License: MIT

{caption_flow-0.3.2.dist-info → caption_flow-0.3.4.dist-info}/RECORD

@@ -1,14 +1,14 @@
-caption_flow/__init__.py,sha256=09Vyr0RqKrKe1caUhXq9beficJkmclryjT6BNiASUxQ,303
-caption_flow/cli.py,sha256=t_cYCxJE7f5UtB3br2Es51JjO5KPsWM1JTdDXAxM_Lw,41371
+caption_flow/__init__.py,sha256=2M1VLvkVjUmTHXuJFMLnZKqVYni5A0HJfxcnjz53K7c,303
+caption_flow/cli.py,sha256=K3lML3WIYjD7OluGltHGP4N98S5w-KyhDUlQZudDQXE,41464
 caption_flow/models.py,sha256=2n6iphTEL62xK2FFcJM6axMsaE8KwsUv5Ak_cCF-TdQ,5652
-caption_flow/monitor.py,sha256=bAt9EJqfPgT_KdbknGdCxwBRH002pRDgyUmYIj6Dyso,7885
-caption_flow/orchestrator.py,sha256=34gZvaW14YZ7a7LagYOO3VKKwlbuS4aw0yoP1L8gwf0,36192
+caption_flow/monitor.py,sha256=z2HakZSG799HvTJgjgG7u_MHvhq9-JL1LXzxBwP3WQc,7998
+caption_flow/orchestrator.py,sha256=3XKZXFE1Aw1kCqb_Vw9loYpkmJ5LTLyZZf9pj4k6ldA,37175
 caption_flow/viewer.py,sha256=HxO98eHR1xtivG0dEdYC2U9T_RgeRfJqqTK-37u9bNM,20471
 caption_flow/processors/__init__.py,sha256=hvq-OuAJWQe6hFglKe7QmkS8473k20FmxZDSxfXpCrg,423
 caption_flow/processors/base.py,sha256=IAEr0pqHRuSkXunvDWk1vf2IKeYQ-2YERqej9iSQm94,6931
-caption_flow/processors/huggingface.py,sha256=w0j7PRosXYyJXZ0A0Y-J6_n-aHCGVW8tbt8lcvguO_Y,41237
+caption_flow/processors/huggingface.py,sha256=t_dklhmNRAyk2jISu4FqmNecjg9hfY47omOiRVkbhvA,41215
 caption_flow/processors/local_filesystem.py,sha256=OuNNDemy0sdtpBBC_5GbI-c1vMqp8OIz983Cq85gdb8,27964
-caption_flow/processors/webdataset.py,sha256=TkC6xZO6m2FcwiBQGJsSQcrshBKcLdr4edFVtnBOd3U,28999
+caption_flow/processors/webdataset.py,sha256=tUBCUKunqooHibTWtQ1wljuRI55Wc6M1WrI2hOZgt7g,33858
 caption_flow/storage/__init__.py,sha256=IVnzcSCPpPuyp-QLlgJirRZ9Sb3tR0F4sfuF5u2cNMk,36
 caption_flow/storage/exporter.py,sha256=mFJqMDQ61cP-qcXe118_-oL1TUqULdQZ8LdjSTym44I,19697
 caption_flow/storage/manager.py,sha256=KPExcKPuFVQSsBnfCBdne5PO4PwN4NTfd-EJQk13OY0,47459
@@ -16,18 +16,18 @@ caption_flow/utils/__init__.py,sha256=bDcO5uR455TKCQ2hX-_XcdTnRXDBaT8Yn4jWqWzfFs
 caption_flow/utils/auth.py,sha256=UrxX2n8OEEcfMD1Ey27TxGfrJFmUCpC59x-SCrQJoVE,2253
 caption_flow/utils/caption_utils.py,sha256=esUMAdcCkNjRroZ0Bhxv0_yKlLtMf0XeDCTt-5k6bik,5309
 caption_flow/utils/certificates.py,sha256=eu4blQZEkL9NRaY1ynQWg1asvDorRYhGRZea7STonJE,4635
-caption_flow/utils/checkpoint_tracker.py,sha256=-nN5gLvXyMdKOCT2SNNL2Km6UYm2Hii9wuXeezWhwx4,3339
-caption_flow/utils/chunk_tracker.py,sha256=HntWeINTbJmIERsW21p4q4FK8D9-4xKbZQUsj24DIqo,19975
+caption_flow/utils/checkpoint_tracker.py,sha256=nOZIIGsXTRUj09tFSnWtRgj_zoa8Og_-rutkr2GFz8Y,4417
+caption_flow/utils/chunk_tracker.py,sha256=JZIFvaHS5AYaVOzsSJKrnNlS4E3BdzV64cRkQa_65g0,21508
 caption_flow/utils/image_processor.py,sha256=wmOExkVfM7OeuLfX3AwMefsH-TxL8TNcn22gp0NmJKY,1541
 caption_flow/utils/json_utils.py,sha256=IiZYn8uCM-3pYmyIbX2fmaOIyutArn67SqAyp0ggNpU,5396
 caption_flow/utils/prompt_template.py,sha256=AKp0diSZqNBMwZkpiTNjw8-bbQwHStr7QZTOJ7o1dC4,4345
 caption_flow/utils/vllm_config.py,sha256=TC7Rmjk0zRKbBXbWUXrFL4Z58hzax_-4L0pXZn09hdM,6019
-caption_flow/workers/base.py,sha256=2AGWERC5hbmO-0V_A1MUbgRVvRNN3blqGPyDokvvzmM,7575
-caption_flow/workers/caption.py,sha256=X4BEmb6C1c73hvgJDMsHtgCUlCuECtnloWSVolVpa4s,39353
+caption_flow/workers/base.py,sha256=nEWohozFZ0Bw3_8U8xirnKLeZsGR5k69rSu4j-oDitc,7698
+caption_flow/workers/caption.py,sha256=swE4pYg4ZYAAtMxvyvlETa3wv4yKWUPXXulCAwPhPiQ,39477
 caption_flow/workers/data.py,sha256=0Tg8NE0wdONeMlivYQ4nvbcfWdLuU51O7vR8_YSnJgo,14813
-caption_flow-0.3.2.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-caption_flow-0.3.2.dist-info/METADATA,sha256=8bHECzNi4R6_FlbHWSHMx9TDo4uTVKWWgVbqAe5cCIs,9708
-caption_flow-0.3.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-caption_flow-0.3.2.dist-info/entry_points.txt,sha256=KnVlyrGKZj6p2zNyuEnCx4Y6jvJ4V-mcfN0lddPKTlQ,55
-caption_flow-0.3.2.dist-info/top_level.txt,sha256=_bXpKRutqded0FQ80dCChIz26ETV7tL4d4e2E_Y1FXs,13
-caption_flow-0.3.2.dist-info/RECORD,,
+caption_flow-0.3.4.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+caption_flow-0.3.4.dist-info/METADATA,sha256=dfB40EF_Zgz2Ux8qvdBbfLdhzY85_MUFRX-904I-qb4,9708
+caption_flow-0.3.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+caption_flow-0.3.4.dist-info/entry_points.txt,sha256=KnVlyrGKZj6p2zNyuEnCx4Y6jvJ4V-mcfN0lddPKTlQ,55
+caption_flow-0.3.4.dist-info/top_level.txt,sha256=_bXpKRutqded0FQ80dCChIz26ETV7tL4d4e2E_Y1FXs,13
+caption_flow-0.3.4.dist-info/RECORD,,