caption-flow 0.5.6__tar.gz → 0.5.7__tar.gz

{caption_flow-0.5.6 → caption_flow-0.5.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.5.6
+Version: 0.5.7
 Summary: Self-contained distributed community captioning system
 Author: bghira
 Author-email: bghira <bghira@users.github.com>
@@ -31,7 +31,7 @@ Requires-Dist: trainingsample<0.4.0,>=0.3.2
 Requires-Dist: pandas<4.0.0,>=3.0.0
 Requires-Dist: datasets<6.0.0,>=5.0.0
 Requires-Dist: boto3<2.0.0,>=1.43.0
-Requires-Dist: webshart<0.6.0,>=0.5.2
+Requires-Dist: webshart<0.7.0,>=0.6.0
 Requires-Dist: pylance<9.0.0,>=8.0.0
 Requires-Dist: duckdb<2.0.0,>=1.5.0
 Requires-Dist: aiohttp<4.0.0,>=3.13.3
@@ -43,35 +43,35 @@ Requires-Dist: requests<3.0.0,>=2.32.0
 Requires-Dist: tqdm<5.0.0,>=4.67.0
 Requires-Dist: urwid<5.0.0,>=3.0.2
 Provides-Extra: vllm
-Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "vllm"
-Requires-Dist: torch==2.11.0; extra == "vllm"
-Requires-Dist: torchvision==0.26.0; extra == "vllm"
+Requires-Dist: vllm<0.32.0,>=0.31.0; extra == "vllm"
+Requires-Dist: torch==2.13.0; extra == "vllm"
+Requires-Dist: torchvision==0.28.0; extra == "vllm"
 Requires-Dist: torchaudio==2.11.0; extra == "vllm"
-Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "vllm"
+Requires-Dist: transformers<5.18.0,>=5.10.4; extra == "vllm"
 Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "vllm"
 Provides-Extra: captioning
-Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "captioning"
-Requires-Dist: torch==2.11.0; extra == "captioning"
-Requires-Dist: torchvision==0.26.0; extra == "captioning"
+Requires-Dist: vllm<0.32.0,>=0.31.0; extra == "captioning"
+Requires-Dist: torch==2.13.0; extra == "captioning"
+Requires-Dist: torchvision==0.28.0; extra == "captioning"
 Requires-Dist: torchaudio==2.11.0; extra == "captioning"
-Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "captioning"
+Requires-Dist: transformers<5.18.0,>=5.10.4; extra == "captioning"
 Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "captioning"
 Provides-Extra: openai
 Provides-Extra: cpu
-Requires-Dist: torch>=2.11.0; extra == "cpu"
-Requires-Dist: torchvision>=0.26.0; extra == "cpu"
+Requires-Dist: torch>=2.13.0; extra == "cpu"
+Requires-Dist: torchvision>=0.28.0; extra == "cpu"
 Requires-Dist: torchaudio>=2.11.0; extra == "cpu"
 Provides-Extra: cuda
-Requires-Dist: torch>=2.11.0; extra == "cuda"
-Requires-Dist: torchvision>=0.26.0; extra == "cuda"
+Requires-Dist: torch>=2.13.0; extra == "cuda"
+Requires-Dist: torchvision>=0.28.0; extra == "cuda"
 Requires-Dist: torchaudio>=2.11.0; extra == "cuda"
 Provides-Extra: cuda13
-Requires-Dist: torch>=2.11.0; extra == "cuda13"
-Requires-Dist: torchvision>=0.26.0; extra == "cuda13"
+Requires-Dist: torch>=2.13.0; extra == "cuda13"
+Requires-Dist: torchvision>=0.28.0; extra == "cuda13"
 Requires-Dist: torchaudio>=2.11.0; extra == "cuda13"
 Provides-Extra: rocm
-Requires-Dist: torch>=2.11.0; extra == "rocm"
-Requires-Dist: torchvision>=0.26.0; extra == "rocm"
+Requires-Dist: torch>=2.13.0; extra == "rocm"
+Requires-Dist: torchvision>=0.28.0; extra == "rocm"
 Requires-Dist: torchaudio>=2.11.0; extra == "rocm"
 Provides-Extra: apple
 Requires-Dist: torch>=2.13.0; extra == "apple"
@@ -82,16 +82,16 @@ Provides-Extra: dev
 Requires-Dist: pytest<9.0.0,>=8.0.0; extra == "dev"
 Requires-Dist: pytest-asyncio<2.0.0,>=1.1.0; extra == "dev"
 Requires-Dist: pytest-cov<7.0.0,>=6.0.0; extra == "dev"
-Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "dev"
+Requires-Dist: transformers<5.18.0,>=5.10.4; extra == "dev"
 Requires-Dist: black<26.0.0,>=25.0.0; extra == "dev"
 Requires-Dist: ruff<1.0.0,>=0.12.0; extra == "dev"
 Requires-Dist: mypy<2.0.0,>=1.17.0; extra == "dev"
 Provides-Extra: all
-Requires-Dist: vllm<0.26.0,>=0.25.1; extra == "all"
-Requires-Dist: torch==2.11.0; extra == "all"
-Requires-Dist: torchvision==0.26.0; extra == "all"
+Requires-Dist: vllm<0.32.0,>=0.31.0; extra == "all"
+Requires-Dist: torch==2.13.0; extra == "all"
+Requires-Dist: torchvision==0.28.0; extra == "all"
 Requires-Dist: torchaudio==2.11.0; extra == "all"
-Requires-Dist: transformers<6.0.0,>=5.5.3; extra == "all"
+Requires-Dist: transformers<5.18.0,>=5.10.4; extra == "all"
 Requires-Dist: qwen-vl-utils<0.1.0,>=0.0.14; extra == "all"
 Dynamic: author
 Dynamic: license-file
@@ -128,6 +128,11 @@ pip install --upgrade pip
 pip install "caption-flow[vllm]"
 ```
 
+The Linux GPU worker extra targets vLLM 0.31.x with PyTorch 2.13.0,
+torchvision 0.28.0, and torchaudio 2.11.0 (as required by vLLM), plus
+Transformers 5.10.4–5.17.x. To upgrade an existing worker from a checkout,
+run `pip install --upgrade -e ".[vllm]"` in its environment and restart it.
+
 An OpenAI-compatible API worker does not need the GPU dependencies:
 
 ```bash
@@ -220,7 +225,7 @@ Options:
 * **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata
 * **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets
 * **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions
-* **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. for this format, pass `--output` as the path to the existing shard metadata JSON file when exporting one shard. if you export multiple shards, pass `--output` as a directory containing one existing `{shard_name}.json` file per shard.
+* **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. Strings, native JSON objects, and mixed lists are supported. JSON objects stored as text in Lance output columns are decoded during export, preserving nested values instead of writing Python dictionary representations. Native objects supplied through `Caption.outputs` are encoded as JSON only for Lance persistence. Requires webshart 0.6.0 or newer for native JSON caption support (included in the base dependencies). For this format, pass `--output` as the path to the existing shard metadata JSON file when exporting one shard. If you export multiple shards, pass `--output` as a directory containing one existing `{shard_name}.json` file per shard.
 * **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary
 * **all**: creates the directory/file-generating export formats in a specified `--output` directory. prefer a directory here; `webshart` is a special case that expects existing per-shard metadata `.json` files rather than creating new metadata files.
 
@@ -279,6 +284,13 @@ prompts, sampling, and output fields. Existing `vllm:` configurations are also
 accepted by API workers for backward compatibility; endpoint-local `model`
 values take precedence over the broadcast model name.
 
+Direct vLLM and API workers share JSON validation, output transformations,
+refusal detection, and semantic retries. Configure those under
+`inference.output_processing`; use `response_format` and `retry_response_format`
+for constrained decoding on either backend. See the
+[structured JSON recipe](recipes/STRUCTURED_JSON.md) for examples and native
+vLLM tuning details.
+
 ---
 
 ## dataset formats
@@ -288,6 +300,15 @@ values take precedence over the broadcast model name.
 * webdatasets shards containing full image data; also can be hosted on the hub
 * local folder filled with images; orchestrator will serve the data to workers
 
+For indexed Webshart reads, set `orchestrator.webshart_prefetch_shards: 1` to
+download the next shard in the background while captioning the current one.
+The default is `0` (disabled). This works with both direct vLLM and API workers,
+does not move the sample cursor, and is ignored for explicit remote-range reads.
+Size `shard_cache_gb` to hold the current shard plus the lookahead. Prefetch
+retries are spaced by `webshart_cache_retry_interval` (30 seconds by default);
+a prefetch failure never marks samples processed or prevents reading the
+current cached shard.
+
 ## configuration path
 
 ### config discovery order

{caption_flow-0.5.6 → caption_flow-0.5.7}/README.md

@@ -27,6 +27,11 @@ pip install --upgrade pip
 pip install "caption-flow[vllm]"
 ```
 
+The Linux GPU worker extra targets vLLM 0.31.x with PyTorch 2.13.0,
+torchvision 0.28.0, and torchaudio 2.11.0 (as required by vLLM), plus
+Transformers 5.10.4–5.17.x. To upgrade an existing worker from a checkout,
+run `pip install --upgrade -e ".[vllm]"` in its environment and restart it.
+
 An OpenAI-compatible API worker does not need the GPU dependencies:
 
 ```bash
@@ -119,7 +124,7 @@ Options:
 * **csv**: exports CSV-compatible data columns to the `--output` path containing incomplete metadata
 * **json**: creates a `.json` file for each sample inside the `--output` subdirectory containing **complete** metadata; useful for webdatasets
 * **txt**: creates `.txt` file for each sample inside the `--output` subdirectory containing ONLY captions
-* **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. for this format, pass `--output` as the path to the existing shard metadata JSON file when exporting one shard. if you export multiple shards, pass `--output` as a directory containing one existing `{shard_name}.json` file per shard.
+* **webshart**: updates an **existing per-shard metadata `.json` file** by writing captions under the plural `captions` key. Strings, native JSON objects, and mixed lists are supported. JSON objects stored as text in Lance output columns are decoded during export, preserving nested values instead of writing Python dictionary representations. Native objects supplied through `Caption.outputs` are encoded as JSON only for Lance persistence. Requires webshart 0.6.0 or newer for native JSON caption support (included in the base dependencies). For this format, pass `--output` as the path to the existing shard metadata JSON file when exporting one shard. If you export multiple shards, pass `--output` as a directory containing one existing `{shard_name}.json` file per shard.
 * **huggingface_hub**: creates a dataset on Hugging Face Hub, possibly `--private` and `--nsfw` where necessary
 * **all**: creates the directory/file-generating export formats in a specified `--output` directory. prefer a directory here; `webshart` is a special case that expects existing per-shard metadata `.json` files rather than creating new metadata files.
 
@@ -178,6 +183,13 @@ prompts, sampling, and output fields. Existing `vllm:` configurations are also
 accepted by API workers for backward compatibility; endpoint-local `model`
 values take precedence over the broadcast model name.
 
+Direct vLLM and API workers share JSON validation, output transformations,
+refusal detection, and semantic retries. Configure those under
+`inference.output_processing`; use `response_format` and `retry_response_format`
+for constrained decoding on either backend. See the
+[structured JSON recipe](recipes/STRUCTURED_JSON.md) for examples and native
+vLLM tuning details.
+
 ---
 
 ## dataset formats
@@ -187,6 +199,15 @@ values take precedence over the broadcast model name.
 * webdatasets shards containing full image data; also can be hosted on the hub
 * local folder filled with images; orchestrator will serve the data to workers
 
+For indexed Webshart reads, set `orchestrator.webshart_prefetch_shards: 1` to
+download the next shard in the background while captioning the current one.
+The default is `0` (disabled). This works with both direct vLLM and API workers,
+does not move the sample cursor, and is ignored for explicit remote-range reads.
+Size `shard_cache_gb` to hold the current shard plus the lookahead. Prefetch
+retries are spaced by `webshart_cache_retry_interval` (30 seconds by default);
+a prefetch failure never marks samples processed or prevents reading the
+current cached shard.
+
 ## configuration path
 
 ### config discovery order

{caption_flow-0.5.6 → caption_flow-0.5.7}/setup.py

@@ -38,7 +38,7 @@ base_deps = [
     "pandas>=3.0.0,<4.0.0",
     "datasets>=5.0.0,<6.0.0",
     "boto3>=1.43.0,<2.0.0",
-    "webshart>=0.5.2,<0.6.0",
+    "webshart>=0.6.0,<0.7.0",
     "pylance>=8.0.0,<9.0.0",
     "duckdb>=1.5.0,<2.0.0",
     "aiohttp>=3.13.3,<4.0.0",
@@ -51,21 +51,21 @@ base_deps = [
     "urwid>=3.0.2,<5.0.0",
 ]
 
-# vLLM 0.25.1 currently selects the PyTorch 2.11.0 family. Keep these
-# explicit so a worker install cannot silently combine incompatible torch
-# packages when another dependency broadens its requirements.
+# Match vLLM 0.31's published requirements, including torchaudio 2.11.0.
+# Keep these explicit so a worker install cannot silently combine incompatible
+# torch packages when another dependency broadens its requirements.
 vllm_deps = [
-    "vllm>=0.25.1,<0.26.0",
-    "torch==2.11.0",
-    "torchvision==0.26.0",
+    "vllm>=0.31.0,<0.32.0",
+    "torch==2.13.0",
+    "torchvision==0.28.0",
     "torchaudio==2.11.0",
-    "transformers>=5.5.3,<6.0.0",
+    "transformers>=5.10.4,<5.18.0",
     "qwen-vl-utils>=0.0.14,<0.1.0",
 ]
 
 PYTORCH_DEPENDENCIES = [
-    "torch>=2.11.0",
-    "torchvision>=0.26.0",
+    "torch>=2.13.0",
+    "torchvision>=0.28.0",
     "torchaudio>=2.11.0",
 ]
 
@@ -92,7 +92,7 @@ extras_require = {
         "pytest>=8.0.0,<9.0.0",
         "pytest-asyncio>=1.1.0,<2.0.0",
         "pytest-cov>=6.0.0,<7.0.0",
-        "transformers>=5.5.3,<6.0.0",
+        "transformers>=5.10.4,<5.18.0",
         "black>=25.0.0,<26.0.0",
         "ruff>=0.12.0,<1.0.0",
         "mypy>=1.17.0,<2.0.0",

{caption_flow-0.5.6 → caption_flow-0.5.7}/src/caption_flow/__init__.py

@@ -1,6 +1,6 @@
 """CaptionFlow - Distributed community captioning system."""
 
-__version__ = "0.5.6"
+__version__ = "0.5.7"
 
 from .monitor import Monitor
 from .orchestrator import Orchestrator

{caption_flow-0.5.6 → caption_flow-0.5.7}/src/caption_flow/models.py

@@ -6,7 +6,7 @@ import os
 from dataclasses import dataclass, field
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from PIL import Image
 
@@ -128,7 +128,9 @@ class Caption:
     caption_count: int = 1  # Number of captions generated for this item
     caption: Optional[str] = None
     captions: Optional[List[str]] = None
-    outputs: Dict[str, List[str]] = field(default_factory=dict)
+    outputs: Dict[str, Union[str, Dict[str, Any], List[Union[str, Dict[str, Any]]]]] = field(
+        default_factory=dict
+    )
     quality_score: Optional[float] = None
     quality_scores: Optional[List[float]] = None
 
@@ -179,12 +181,13 @@ class ProcessingStage:
     dtype: Optional[str] = None
     gpu_memory_utilization: Optional[float] = None
 
-    # Optional semantic fallback for empty/refused caption responses. API
-    # workers can omit the image on the fallback request when an upstream
-    # provider rejects the image before inference.
+    # Caption policy and recovery apply to every inference backend.
     retry_prompt: Optional[str] = None
     retry_without_image: bool = False
     retry_sampling: Optional[Dict[str, Any]] = None
+    output_processing: Optional[Dict[str, Any]] = None
+    response_format: Optional[Dict[str, Any]] = None
+    retry_response_format: Optional[Dict[str, Any]] = None
 
 
 @dataclass

{caption_flow-0.5.6 → caption_flow-0.5.7}/src/caption_flow/orchestrator.py

@@ -51,7 +51,7 @@ class Orchestrator:
         assert (
             processor_type is not None
         ), "You must supply processor_type in your orchestrator dataset configuration."
-        processor_config = ProcessorConfig(processor_type=processor_type, config=config)
+        self.processor_config = ProcessorConfig(processor_type=processor_type, config=config)
 
         # Initialize processor
         if processor_type == "webdataset":
@@ -70,7 +70,6 @@ class Orchestrator:
             caption_buffer_size=storage_config.get("caption_buffer_size", 1000),
         )
         self.auth = AuthManager(config.get("auth", {}))
-        self.processor.initialize(processor_config, self.storage)
 
         # Processing configuration
         self.chunks_per_request = config.get("chunks_per_request", 2)
@@ -132,6 +131,8 @@ class Orchestrator:
 
         # Initialize storage
         await self.storage.initialize()
+        # Checkpoints can include results that had not reached a durable flush.
+        self.processor.initialize(self.processor_config, self.storage)
 
         # Start background tasks
         asyncio.create_task(self._heartbeat_loop())
@@ -410,7 +411,9 @@ class Orchestrator:
 
         elif msg_type == "work_complete":
             unit_id = data["unit_id"]
-            self.processor.mark_completed(unit_id, worker_id)
+            if self.processor.mark_completed(unit_id, worker_id) is False:
+                logger.warning(f"Work unit {unit_id} reported complete before all results arrived")
+                return
             logger.debug(f"Work unit {unit_id} completed by worker {worker_id}")
             await self._send_activity(f"Chunk completed: {unit_id}")
 
@@ -428,9 +431,8 @@ class Orchestrator:
             logger.debug(f"Heartbeat from {worker_id}: {data}")
 
     async def _handle_results_submission(self, worker_id: str, data: Dict):
-        """Process results submission from worker - fires off async task and returns immediately."""
-        # Fire and forget - process in background
-        asyncio.create_task(self._process_result_async(worker_id, data))
+        """Accept results in connection order before handling a completion message."""
+        await self._process_result_async(worker_id, data)
 
     async def _process_result_async(self, worker_id: str, data: Dict):
         """Actually process the result in background."""
@@ -444,6 +446,9 @@ class Orchestrator:
             shard_name = job_id.shard_id
             chunk_name = job_id.chunk_id
 
+            if data.get("error") or not data.get("outputs"):
+                raise ValueError("A failed result cannot be recorded as completed work")
+
             result = WorkResult(
                 unit_id=data["unit_id"],
                 source_id=shard_name,
@@ -455,11 +460,6 @@ class Orchestrator:
                 processing_time_ms=data.get("processing_time_ms", 0),
             )
 
-            # 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
             total_outputs = sum(len(v) for v in result.outputs.values())
 
@@ -480,7 +480,7 @@ class Orchestrator:
             caption = Caption(
                 job_id=job_id,
                 dataset=result.dataset,
-                shard=processed["source_id"],
+                shard=result.source_id,
                 chunk_id=chunk_name,
                 item_key=item_key,
                 captions=result.outputs.get("captions", []),
@@ -501,6 +501,11 @@ class Orchestrator:
             # Save to storage (might skip if duplicate)
             saved = await self.storage.save_caption(caption)
 
+            # Progress may only follow storage acceptance. A duplicate is also
+            # accepted: its caption is already present. Startup reconstructs
+            # progress from durable storage if buffered results are lost.
+            self.processor.handle_result(result)
+
             # Update contributor stats only if actually saved
             if saved:
                 contributor = await self.storage.get_contributor(worker_user)
@@ -509,8 +514,11 @@ class Orchestrator:
                     await self.storage.save_contributor(contributor)
 
         except Exception as e:
+            if data.get("unit_id"):
+                self.processor.mark_failed(data["unit_id"], worker_id, "Result acceptance failed")
             logger.error(
-                f"Error processing result from {worker_id} for unit {data.get('unit_id', 'unknown')}: {e}",
+                f"Error processing result from {worker_id} "
+                f"for unit {data.get('unit_id', 'unknown')}: {e}",
                 exc_info=True,
             )