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

caption_flow/utils/checkpoint_tracker.py

@@ -73,12 +73,12 @@ class CheckpointTracker(ABC):
             logger.debug(f"Saved checkpoint to {self.checkpoint_path}")
 
         except Exception as e:
-            logger.error(f"Error saving checkpoint: {e}", exc_info=True)
+            # 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")
+                # logger.info("Saved checkpoint using fallback direct write")
             except Exception as fallback_error:
                 logger.error(f"Fallback save also failed: {fallback_error}")
 

caption_flow/utils/chunk_tracker.py

@@ -10,6 +10,7 @@ from dataclasses import dataclass, asdict, field
 from .checkpoint_tracker import CheckpointTracker
 
 logger = logging.getLogger(__name__)
+# logger.setLevel(logging.DEBUG)
 
 
 @dataclass
@@ -58,11 +59,15 @@ class ChunkState:
     def get_unprocessed_ranges(self) -> List[Tuple[int, int]]:
         """Get ranges that haven't been processed yet."""
         if not self.processed_ranges:
+            logger.info(f"Chunk {self.chunk_id} has no processed ranges, returning full range")
             return [(0, self.chunk_size - 1)]
 
         unprocessed = []
         current = 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))
@@ -132,6 +137,11 @@ class ChunkTracker(CheckpointTracker):
         self, chunk_id: str, shard_name: str, shard_url: str, start_index: int, chunk_size: int
     ) -> bool:
         """Add a new chunk. Returns False if chunk already exists and is completed."""
+        if chunk_id in self.chunks:
+            logger.debug(
+                f"Chunk {chunk_id} already exists with status: {self.chunks[chunk_id].status}, not creating"
+            )
+            return False
         if chunk_id in self.completed_chunks:
             logger.debug(f"Chunk {chunk_id} already completed, skipping")
             return False
@@ -166,7 +176,7 @@ class ChunkTracker(CheckpointTracker):
             chunk.completed_at = datetime.utcnow()
             self.completed_chunks.add(chunk_id)
             self.save()
-            logger.info(f"Chunk {chunk_id} marked as completed")
+            logger.debug(f"Chunk {chunk_id} marked as completed")
 
     def mark_failed(self, chunk_id: str):
         """Mark chunk as failed."""
@@ -207,6 +217,49 @@ class ChunkTracker(CheckpointTracker):
                     pending.append(chunk_id)
         return pending
 
+    def get_processed_indices_for_chunk(
+        self, chunk_id: str, processed_job_ids: Set[str]
+    ) -> List[Tuple[int, int]]:
+        """Convert processed job_ids back to ranges for a chunk."""
+        # Extract indices from job_ids like "data-0000:chunk:0:idx:42"
+        processed_indices = []
+        # this will be slow as shit, but it's simple for now, Proof of Concept.
+        for job_id in processed_job_ids:
+            test_chunk_id = chunk_id.replace("_", ":")
+            if test_chunk_id in job_id:
+                parts = job_id.split(":")
+                logger.debug(
+                    f"Found matching job_id {job_id} for chunk {chunk_id} with {len(parts)=} and {parts[3]=}"
+                )
+                if len(parts) >= 5 and parts[3] == "idx":
+                    idx = int(parts[4])
+                    processed_indices.append(idx)
+
+        # Convert to ranges
+        if not processed_indices:
+            # logger.warning(
+            #     f"Chunk {chunk_id} had no pre-processed ranges discovered, will process all elements"
+            # )
+            return []
+        else:
+            logger.debug(f"Chunk {chunk_id} has {len(processed_indices)} pre-processed indices")
+
+        processed_indices.sort()
+        ranges = []
+        start = processed_indices[0]
+        end = processed_indices[0]
+
+        for idx in processed_indices[1:]:
+            if idx == end + 1:
+                end = idx
+            else:
+                ranges.append((start, end))
+                start = idx
+                end = idx
+
+        ranges.append((start, end))
+        return ranges
+
     def is_shard_complete(self, shard_name: str) -> bool:
         """Check if all chunks for a shard are complete."""
         shard_chunks = [chunk for chunk in self.chunks.values() if chunk.shard_name == shard_name]
@@ -236,20 +289,8 @@ class ChunkTracker(CheckpointTracker):
 
         for chunk_id, chunk_state in self.chunks.items():
             shard_name = chunk_state.shard_name
-
-            # For virtual HF dataset shards, normalize the shard name
-            if shard_name.startswith("hf_dataset:"):
-                parts = shard_name.split(":")
-                if len(parts) >= 4 and parts[2] == "chunk":
-                    # Use just the dataset identifier as the shard name
-                    normalized_shard_name = ":".join(parts[:2])
-                else:
-                    normalized_shard_name = shard_name
-            else:
-                normalized_shard_name = shard_name
-
-            if normalized_shard_name not in shards:
-                shards[normalized_shard_name] = {
+            if shard_name not in shards:
+                shards[shard_name] = {
                     "total_chunks": 0,
                     "completed_chunks": 0,
                     "pending_chunks": 0,
@@ -259,20 +300,20 @@ class ChunkTracker(CheckpointTracker):
                     "chunks": [],
                 }
 
-            shards[normalized_shard_name]["chunks"].append(chunk_state)
-            shards[normalized_shard_name]["total_chunks"] += 1
+            shards[shard_name]["chunks"].append(chunk_state)
+            shards[shard_name]["total_chunks"] += 1
 
             if chunk_state.status == "completed":
-                shards[normalized_shard_name]["completed_chunks"] += 1
+                shards[shard_name]["completed_chunks"] += 1
             elif chunk_state.status == "pending":
-                shards[normalized_shard_name]["pending_chunks"] += 1
-                shards[normalized_shard_name]["is_complete"] = False
+                shards[shard_name]["pending_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
             elif chunk_state.status == "assigned":
-                shards[normalized_shard_name]["assigned_chunks"] += 1
-                shards[normalized_shard_name]["is_complete"] = False
+                shards[shard_name]["assigned_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
             elif chunk_state.status == "failed":
-                shards[normalized_shard_name]["failed_chunks"] += 1
-                shards[normalized_shard_name]["is_complete"] = False
+                shards[shard_name]["failed_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
 
         return shards
 
@@ -322,13 +363,7 @@ class ChunkTracker(CheckpointTracker):
                         continue
 
                     # Infer shard URL and create chunk with default size
-                    if shard_name.replace("_", "/") in chunk_id or "_" in shard_name:
-                        # HF dataset
-                        dataset_path = shard_name.replace("_", "/")
-                        shard_url = f"hf_dataset:{dataset_path}:chunk:{start_idx}"
-                    else:
-                        # WebDataset
-                        shard_url = f"unknown://{shard_name}.tar"
+                    shard_url = f"unknown://{shard_name}.tar"
 
                     self.chunks[chunk_id] = ChunkState(
                         chunk_id=chunk_id,
@@ -410,6 +445,7 @@ class ChunkTracker(CheckpointTracker):
         """Mark a range of items as processed within a chunk (expects ABSOLUTE indices)."""
         if chunk_id not in self.chunks:
             logger.error(f"Unknown chunk: {chunk_id}")
+            logger.debug(f"Known chunks: {list(self.chunks.keys())}")
             return
 
         chunk = self.chunks[chunk_id]
@@ -441,9 +477,32 @@ class ChunkTracker(CheckpointTracker):
         )
 
     def get_chunk_with_unprocessed_items(self, chunk_id: str) -> Optional[Dict[str, Any]]:
-        """Get chunk info including unprocessed ranges."""
-        if chunk_id not in self.chunks:
+        """Get chunk info with unprocessed item ranges."""
+        chunk_state = self.chunks.get(chunk_id)
+        if not chunk_state:
             return None
 
-        chunk = self.chunks[chunk_id]
-        return {"chunk": chunk.to_dict(), "unprocessed_ranges": chunk.get_unprocessed_ranges()}
+        # During startup or if no worker is assigned, treat all unprocessed as available
+        if not hasattr(self, "_startup_complete"):
+            self._startup_complete = False
+
+        if not self._startup_complete or (
+            not chunk_state.assigned_to or chunk_state.completed_at is None
+        ):
+            # Return all unprocessed ranges
+            logger.debug(
+                f"Returning all unprocessed ranges. Status {self._startup_complete=} {chunk_state=}"
+            )
+            return {
+                "chunk_id": chunk_id,
+                "unprocessed_ranges": chunk_state.get_unprocessed_ranges(),
+                "status": chunk_state.status,
+            }
+
+        # Normal operation - only return ranges not being worked on
+        # This would need more complex tracking of which ranges each worker is processing
+        return {
+            "chunk_id": chunk_id,
+            "unprocessed_ranges": chunk_state.get_unprocessed_ranges(),
+            "status": chunk_state.status,
+        }