caption-flow 0.1.0__py3-none-any.whl

caption_flow/utils/chunk_tracker.py

@@ -0,0 +1,365 @@
+"""Chunk tracking for persistent state across restarts."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Set, Dict, List, Optional, Any
+from datetime import datetime
+from dataclasses import dataclass, asdict
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ChunkState:
+    """State of a chunk."""
+
+    chunk_id: str
+    shard_name: str
+    start_index: int
+    chunk_size: int
+    status: str  # pending, assigned, completed, failed
+    completed_at: Optional[datetime] = None
+    assigned_to: Optional[str] = None
+    assigned_at: Optional[datetime] = None
+
+    def to_dict(self):
+        """Convert to dictionary for JSON serialization."""
+        d = asdict(self)
+        # Convert datetime objects to ISO format strings
+        if d["completed_at"]:
+            d["completed_at"] = d["completed_at"].isoformat()
+        if d["assigned_at"]:
+            d["assigned_at"] = d["assigned_at"].isoformat()
+        return d
+
+    @classmethod
+    def from_dict(cls, d: Dict):
+        """Create from dictionary."""
+        # Convert ISO format strings back to datetime objects
+        if d.get("completed_at"):
+            d["completed_at"] = datetime.fromisoformat(d["completed_at"])
+        if d.get("assigned_at"):
+            d["assigned_at"] = datetime.fromisoformat(d["assigned_at"])
+        return cls(**d)
+
+
+class ChunkTracker:
+    """Tracks chunk processing state persistently."""
+
+    def __init__(self, checkpoint_file: Path):
+        self.checkpoint_file = checkpoint_file
+        self.chunks: Dict[str, ChunkState] = {}
+        self.completed_chunks: Set[str] = set()
+        self._load_checkpoint()
+
+    def _load_checkpoint(self):
+        """Load checkpoint from disk."""
+        if self.checkpoint_file.exists():
+            try:
+                with open(self.checkpoint_file, "r") as f:
+                    data = json.load(f)
+
+                # Load chunk states
+                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":
+                        self.completed_chunks.add(chunk_id)
+
+                logger.info(
+                    f"Loaded {len(self.chunks)} chunks from checkpoint, "
+                    f"{len(self.completed_chunks)} completed"
+                )
+
+            except Exception as e:
+                logger.error(f"Error loading chunk checkpoint: {e}")
+                self.chunks = {}
+                self.completed_chunks = set()
+
+    def save_checkpoint(self):
+        """Save checkpoint to disk."""
+        try:
+            # Ensure parent directory exists
+            self.checkpoint_file.parent.mkdir(parents=True, exist_ok=True)
+
+            # Convert chunks to serializable format
+            data = {
+                "chunks": {chunk_id: chunk.to_dict() for chunk_id, chunk in self.chunks.items()},
+                "updated_at": datetime.utcnow().isoformat(),
+            }
+
+            # Write atomically with absolute paths
+            tmp_file = self.checkpoint_file.parent / f"{self.checkpoint_file.name}.tmp"
+
+            # Write to temp file
+            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 (use rename for same filesystem)
+            import shutil
+
+            shutil.move(str(tmp_file), str(self.checkpoint_file))
+
+            logger.debug(
+                f"Saved chunk checkpoint with {len(self.chunks)} chunks to {self.checkpoint_file}"
+            )
+
+        except Exception as e:
+            logger.error(f"Error saving chunk checkpoint: {e}", exc_info=True)
+            # Try direct write as fallback
+            try:
+                with open(self.checkpoint_file, "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}")
+
+    def add_chunk(self, chunk_id: str, shard_name: 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.completed_chunks:
+            logger.debug(f"Chunk {chunk_id} already completed, skipping")
+            return False
+
+        if chunk_id not in self.chunks:
+            self.chunks[chunk_id] = ChunkState(
+                chunk_id=chunk_id,
+                shard_name=shard_name,
+                start_index=start_index,
+                chunk_size=chunk_size,
+                status="pending",
+            )
+            self.save_checkpoint()
+
+        return True
+
+    def mark_assigned(self, chunk_id: str, worker_id: str):
+        """Mark chunk as assigned."""
+        if chunk_id in self.chunks:
+            chunk = self.chunks[chunk_id]
+            chunk.status = "assigned"
+            chunk.assigned_to = worker_id
+            chunk.assigned_at = datetime.utcnow()
+            self.save_checkpoint()
+
+    def mark_completed(self, chunk_id: str):
+        """Mark chunk as completed."""
+        if chunk_id in self.chunks:
+            chunk = self.chunks[chunk_id]
+            chunk.status = "completed"
+            chunk.completed_at = datetime.utcnow()
+            self.completed_chunks.add(chunk_id)
+            self.save_checkpoint()
+            logger.info(f"Chunk {chunk_id} marked as completed")
+
+    def mark_failed(self, chunk_id: str):
+        """Mark chunk as failed."""
+        if chunk_id in self.chunks:
+            chunk = self.chunks[chunk_id]
+            chunk.status = "pending"  # Reset to pending for retry
+            chunk.assigned_to = None
+            chunk.assigned_at = None
+            self.save_checkpoint()
+
+    def mark_pending(self, chunk_id: str):
+        """Mark chunk as pending (for manual reset)."""
+        if chunk_id in self.chunks:
+            chunk = self.chunks[chunk_id]
+            chunk.status = "pending"
+            chunk.assigned_to = None
+            chunk.assigned_at = None
+            self.save_checkpoint()
+
+    def release_worker_chunks(self, worker_id: str):
+        """Release all chunks assigned to a worker."""
+        released_chunks = []
+        for chunk_id, chunk in self.chunks.items():
+            if chunk.assigned_to == worker_id and chunk.status == "assigned":
+                chunk.status = "pending"
+                chunk.assigned_to = None
+                chunk.assigned_at = None
+                released_chunks.append(chunk_id)
+        self.save_checkpoint()
+        return released_chunks
+
+    def get_pending_chunks(self, shard_name: Optional[str] = None) -> List[str]:
+        """Get list of pending chunk IDs, optionally filtered by shard."""
+        pending = []
+        for chunk_id, chunk in self.chunks.items():
+            if chunk.status == "pending":
+                if shard_name is None or chunk.shard_name == shard_name:
+                    pending.append(chunk_id)
+        return pending
+
+    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]
+
+        if not shard_chunks:
+            return False
+
+        return all(chunk.status == "completed" for chunk in shard_chunks)
+
+    def get_stats(self) -> Dict[str, int]:
+        """Get chunk statistics."""
+        stats = {
+            "total": len(self.chunks),
+            "pending": sum(1 for c in self.chunks.values() if c.status == "pending"),
+            "assigned": sum(1 for c in self.chunks.values() if c.status == "assigned"),
+            "completed": len(self.completed_chunks),
+            "failed": sum(1 for c in self.chunks.values() if c.status == "failed"),
+        }
+        return stats
+
+    def get_shards_summary(self) -> Dict[str, Dict[str, Any]]:
+        """Get summary of all shards and their chunk status."""
+        shards = {}
+
+        for chunk_id, chunk_state in self.chunks.items():
+            shard_name = chunk_state.shard_name
+            if shard_name not in shards:
+                shards[shard_name] = {
+                    "total_chunks": 0,
+                    "completed_chunks": 0,
+                    "pending_chunks": 0,
+                    "assigned_chunks": 0,
+                    "failed_chunks": 0,
+                    "is_complete": True,
+                    "chunks": [],
+                }
+
+            shards[shard_name]["chunks"].append(chunk_state)
+            shards[shard_name]["total_chunks"] += 1
+
+            if chunk_state.status == "completed":
+                shards[shard_name]["completed_chunks"] += 1
+            elif chunk_state.status == "pending":
+                shards[shard_name]["pending_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
+            elif chunk_state.status == "assigned":
+                shards[shard_name]["assigned_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
+            elif chunk_state.status == "failed":
+                shards[shard_name]["failed_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
+
+        return shards
+
+    def get_incomplete_shards(self) -> Set[str]:
+        """Get set of shard names that have incomplete chunks."""
+        incomplete = set()
+        for chunk_id, chunk_state in self.chunks.items():
+            if chunk_state.status != "completed":
+                incomplete.add(chunk_state.shard_name)
+        return incomplete
+
+    def get_shards_summary(self) -> Dict[str, Dict[str, Any]]:
+        """Get summary of all shards and their chunk status."""
+        shards = {}
+
+        for chunk_id, chunk_state in self.chunks.items():
+            shard_name = chunk_state.shard_name
+            if shard_name not in shards:
+                shards[shard_name] = {
+                    "total_chunks": 0,
+                    "completed_chunks": 0,
+                    "pending_chunks": 0,
+                    "assigned_chunks": 0,
+                    "failed_chunks": 0,
+                    "is_complete": True,
+                    "chunks": [],
+                }
+
+            shards[shard_name]["chunks"].append(chunk_state)
+            shards[shard_name]["total_chunks"] += 1
+
+            if chunk_state.status == "completed":
+                shards[shard_name]["completed_chunks"] += 1
+            elif chunk_state.status == "pending":
+                shards[shard_name]["pending_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
+            elif chunk_state.status == "assigned":
+                shards[shard_name]["assigned_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
+            elif chunk_state.status == "failed":
+                shards[shard_name]["failed_chunks"] += 1
+                shards[shard_name]["is_complete"] = False
+
+        return shards
+
+    async def sync_with_storage(self, storage_manager):
+        """Sync chunk state with storage to detect already-processed chunks."""
+        logger.info("Syncing chunk state with storage...")
+
+        # Get all existing captions from storage
+        if storage_manager.captions_path.exists():
+            import pyarrow.parquet as pq
+
+            # Read just the job_id column
+            table = pq.read_table(storage_manager.captions_path, columns=["job_id", "chunk_id"])
+            existing_job_ids = set(table["job_id"].to_pylist())
+
+            # Also get chunk_ids if available
+            if "chunk_id" in table.column_names:
+                existing_chunk_ids = set(
+                    cid for cid in table["chunk_id"].to_pylist() if cid is not None
+                )
+
+                # Mark existing chunks as completed
+                for chunk_id in existing_chunk_ids:
+                    if chunk_id in self.chunks:
+                        self.mark_completed(chunk_id)
+                    else:
+                        # Create chunk entry for already-processed chunks
+                        # Extract shard name from chunk_id (format: shard_chunk_index)
+                        parts = chunk_id.rsplit("_chunk_", 1)
+                        if len(parts) == 2:
+                            shard_name = parts[0]
+                            try:
+                                start_idx = int(parts[1])
+                                # We don't know exact chunk size, but mark it as completed
+                                self.chunks[chunk_id] = ChunkState(
+                                    chunk_id=chunk_id,
+                                    shard_name=shard_name,
+                                    start_index=start_idx,
+                                    chunk_size=1000,  # Default chunk size
+                                    status="completed",
+                                    completed_at=datetime.utcnow(),
+                                )
+                                self.completed_chunks.add(chunk_id)
+                            except ValueError:
+                                logger.warning(f"Could not parse chunk_id: {chunk_id}")
+
+                logger.info(f"Found {len(existing_chunk_ids)} completed chunks in storage")
+
+            # Also check by job_id pattern if chunk_id column doesn't exist
+            else:
+                for job_id in existing_job_ids:
+                    # Extract chunk_id from job_id (format: chunk_id_item_key)
+                    if "_chunk_" in job_id:
+                        parts = job_id.split("_")
+                        # Find the chunk part
+                        for i, part in enumerate(parts):
+                            if part == "chunk" and i + 1 < len(parts):
+                                try:
+                                    # Reconstruct chunk_id
+                                    chunk_idx = int(parts[i + 1])
+                                    shard_parts = parts[:i]
+                                    chunk_id = f"{'_'.join(shard_parts)}_chunk_{chunk_idx}"
+
+                                    if chunk_id not in self.completed_chunks:
+                                        self.completed_chunks.add(chunk_id)
+                                        logger.debug(
+                                            f"Marked chunk {chunk_id} as completed from job_id"
+                                        )
+                                    break
+                                except ValueError:
+                                    continue
+
+                logger.info(f"Inferred {len(self.completed_chunks)} completed chunks from job_ids")
+
+            self.save_checkpoint()

caption_flow/utils/dataset_loader.py

@@ -0,0 +1,186 @@
+"""Dataset loading utilities for WebDataset and HuggingFace."""
+
+import asyncio
+import shlex
+import logging
+from pathlib import Path
+from typing import List, Dict, Any, Generator, Optional, Tuple
+import json
+
+import webdataset as wds
+from huggingface_hub import HfFileSystem, get_token, hf_hub_url
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetLoader:
+    """Handles loading datasets from various sources."""
+
+    def __init__(self, dataset_path: str, dataset_type: str = "huggingface"):
+        """
+        Initialize dataset loader.
+
+        Args:
+            dataset_path: Path to dataset (HF repo, local dir, etc.)
+            dataset_type: Type of dataset ("huggingface", "webdataset", "local")
+        """
+        self.dataset_path = dataset_path
+        self.dataset_type = dataset_type
+        self.token = get_token()
+
+        if not self.token and dataset_type == "huggingface":
+            logger.warning("No HuggingFace token found; run `huggingface-cli login`")
+
+    def get_shard_list(self) -> List[str]:
+        """Get list of all shards in the dataset."""
+        if self.dataset_type == "huggingface":
+            return self._get_hf_shards()
+        elif self.dataset_type == "local":
+            return self._get_local_shards()
+        else:
+            raise ValueError(f"Unknown dataset type: {self.dataset_type}")
+
+    def _get_hf_shards(self) -> List[str]:
+        """Get shard URLs from HuggingFace dataset."""
+        logger.info(f"Getting shard list from HuggingFace: {self.dataset_path}")
+
+        fs = HfFileSystem()
+        files = [fs.resolve_path(p) for p in fs.glob(f"hf://datasets/{self.dataset_path}/**/*.tar")]
+
+        urls = [hf_hub_url(f.repo_id, f.path_in_repo, repo_type="dataset") for f in files]
+
+        logger.info(f"Found {len(urls)} shards")
+        return sorted(urls)
+
+    def _get_local_shards(self) -> List[str]:
+        """Get shard files from local directory."""
+        path = Path(self.dataset_path)
+        if not path.exists():
+            raise ValueError(f"Local dataset path does not exist: {path}")
+
+        shards = list(path.glob("*.tar"))
+        logger.info(f"Found {len(shards)} local shards")
+        return [str(s) for s in sorted(shards)]
+
+    def load_shard(self, shard_url: str, processed_keys: Optional[set] = None) -> wds.DataPipeline:
+        """
+        Load a single shard as a WebDataset pipeline.
+
+        Args:
+            shard_url: URL or path to the shard
+            processed_keys: Set of already processed keys to skip
+        """
+        if processed_keys is None:
+            processed_keys = set()
+
+        if self.dataset_type == "huggingface":
+            # Use curl with auth token for HuggingFace
+            url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.token)}' {shlex.quote(shard_url)} || true"
+            ds = wds.DataPipeline(
+                wds.SimpleShardList(url_cmd),
+                wds.tarfile_to_samples(),
+                wds.to_tuple("__key__", "__url__", "jpg;png;jpeg;webp;jxl"),
+                wds.select(lambda x: x[0] not in processed_keys),
+            )
+        else:
+            # Local file access
+            ds = wds.DataPipeline(
+                wds.SimpleShardList(shard_url),
+                wds.tarfile_to_samples(),
+                wds.to_tuple("__key__", "__url__", "jpg;png;jpeg;webp;jxl"),
+                wds.select(lambda x: x[0] not in processed_keys),
+            )
+
+        return ds
+
+    def iterate_shard(
+        self, shard_url: str, processed_keys: Optional[set] = None
+    ) -> Generator[Tuple[str, str, bytes], None, None]:
+        """
+        Iterate over items in a shard.
+
+        Yields:
+            Tuple of (key, url, image_bytes)
+        """
+        ds = self.load_shard(shard_url, processed_keys)
+
+        for key, url, image_data in ds:
+            yield key, url, image_data
+
+    def count_shard_items(self, shard_url: str, processed_keys: Optional[set] = None) -> int:
+        """Count items in a shard (can be slow for large shards)."""
+        count = 0
+        try:
+            for _ in self.iterate_shard(shard_url, processed_keys):
+                count += 1
+        except Exception as e:
+            logger.error(f"Error counting shard {shard_url}: {e}")
+        return count
+
+
+class ShardTracker:
+    """Tracks shard processing progress."""
+
+    def __init__(self, checkpoint_path: Path):
+        """Initialize shard tracker with checkpoint file."""
+        self.checkpoint_path = checkpoint_path
+        self.checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self.completed_shards: set = set()
+        self.partial_shards: Dict[str, Dict[str, Any]] = {}
+        self.load()
+
+    def load(self):
+        """Load checkpoint from disk."""
+        if self.checkpoint_path.exists():
+            try:
+                data = json.loads(self.checkpoint_path.read_text())
+                self.completed_shards = set(data.get("completed_shards", []))
+                self.partial_shards = data.get("partial_shards", {})
+                logger.info(
+                    f"Loaded checkpoint: {len(self.completed_shards)} completed, "
+                    f"{len(self.partial_shards)} partial shards"
+                )
+            except Exception as e:
+                logger.error(f"Failed to load checkpoint: {e}")
+
+    def save(self):
+        """Save checkpoint to disk."""
+        data = {
+            "completed_shards": list(self.completed_shards),
+            "partial_shards": self.partial_shards,
+        }
+
+        tmp = self.checkpoint_path.with_suffix(".tmp")
+        tmp.write_text(json.dumps(data, indent=2))
+        tmp.replace(self.checkpoint_path)
+
+    def mark_complete(self, shard_name: str):
+        """Mark a shard as complete."""
+        self.completed_shards.add(shard_name)
+        if shard_name in self.partial_shards:
+            del self.partial_shards[shard_name]
+        self.save()
+
+    def update_partial(self, shard_name: str, processed_keys: List[str]):
+        """Update partial progress for a shard."""
+        self.partial_shards[shard_name] = {"keys": processed_keys, "count": len(processed_keys)}
+        self.save()
+
+    def get_processed_keys(self, shard_name: str) -> set:
+        """Get set of processed keys for a shard."""
+        if shard_name in self.completed_shards:
+            return set()  # All done
+
+        if shard_name in self.partial_shards:
+            return set(self.partial_shards[shard_name].get("keys", []))
+
+        return set()
+
+    def is_complete(self, shard_name: str) -> bool:
+        """Check if a shard is complete."""
+        return shard_name in self.completed_shards
+
+    def get_remaining_shards(self, all_shards: List[str]) -> List[str]:
+        """Get list of shards that still need processing."""
+        return [s for s in all_shards if Path(s).stem not in self.completed_shards]

caption_flow/utils/image_processor.py

@@ -0,0 +1,51 @@
+"""Image preprocessing utilities."""
+
+import asyncio
+from concurrent.futures import ProcessPoolExecutor
+from pathlib import Path
+from typing import List, Any
+
+import numpy as np
+from PIL import Image
+
+
+class ImageProcessor:
+    """Handles image loading and preprocessing."""
+
+    def __init__(self, num_workers: int = 4):
+        self.executor = ProcessPoolExecutor(max_workers=num_workers)
+
+    async def process_batch(self, image_paths: List[Path]) -> List[np.ndarray]:
+        """Process a batch of images in parallel."""
+        loop = asyncio.get_event_loop()
+
+        tasks = []
+        for path in image_paths:
+            task = loop.run_in_executor(self.executor, self._process_image, path)
+            tasks.append(task)
+
+        return await asyncio.gather(*tasks)
+
+    @staticmethod
+    def _process_image(path: Path) -> np.ndarray:
+        """Process a single image."""
+        img = Image.open(path)
+
+        # Resize to standard size
+        img = img.resize((224, 224), Image.Resampling.LANCZOS)
+
+        # Convert to RGB if needed
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+
+        # Convert to numpy array
+        arr = np.array(img, dtype=np.float32)
+
+        # Normalize
+        arr = arr / 255.0
+
+        return arr
+
+    def shutdown(self):
+        """Shutdown the executor."""
+        self.executor.shutdown(wait=True)

caption_flow/utils/job_queue.py

@@ -0,0 +1,41 @@
+"""Job queue management."""
+
+import asyncio
+from typing import Optional
+from collections import deque
+
+from ..models import Job
+
+
+class JobQueue:
+    """Priority job queue with backpressure."""
+
+    def __init__(self):
+        self.queue = deque()
+        self.processing = set()
+        self.lock = asyncio.Lock()
+
+    async def add(self, job: Job):
+        """Add job to queue."""
+        async with self.lock:
+            self.queue.append(job)
+
+    async def get_next(self) -> Optional[Job]:
+        """Get next available job."""
+        async with self.lock:
+            if self.queue:
+                job = self.queue.popleft()
+                self.processing.add(job.job_id)
+                return job
+        return None
+
+    async def complete(self, job_id: str):
+        """Mark job as complete."""
+        async with self.lock:
+            self.processing.discard(job_id)
+
+    async def requeue(self, job: Job):
+        """Requeue a job (for failures)."""
+        async with self.lock:
+            self.processing.discard(job.job_id)
+            self.queue.appendleft(job)  # Priority requeue