caption-flow 0.5.2__tar.gz → 0.5.3__tar.gz

{caption_flow-0.5.2/src/caption_flow.egg-info → caption_flow-0.5.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.5.2
+Version: 0.5.3
 Summary: Self-contained distributed community captioning system
 Author-email: bghira <bghira@users.github.com>
 License: MIT
@@ -76,7 +76,7 @@ a fast websocket-based orchestrator paired with lightweight gpu workers achieves
 ```bash
 python -m venv .venv
 source .venv/bin/activate  # windows: .venv\Scripts\activate
-pip install caption-flow
+pip install "caption-flow[vllm]"
 ```
 
 ## quickstart (single box)
@@ -240,7 +240,7 @@ PRs welcome. keep it simple and fast.
 
 To contribute compute to a cluster:
 
-1. Install caption-flow: `pip install caption-flow`
+1. Install caption-flow: `pip install "caption-flow[vllm]"`
 2. Get a worker token from the project maintainer
 3. Run: `caption-flow worker --server wss://project.domain.com:8765 --token YOUR_TOKEN`
 

{caption_flow-0.5.2 → caption_flow-0.5.3}/README.md

@@ -21,7 +21,7 @@ a fast websocket-based orchestrator paired with lightweight gpu workers achieves
 ```bash
 python -m venv .venv
 source .venv/bin/activate  # windows: .venv\Scripts\activate
-pip install caption-flow
+pip install "caption-flow[vllm]"
 ```
 
 ## quickstart (single box)
@@ -185,7 +185,7 @@ PRs welcome. keep it simple and fast.
 
 To contribute compute to a cluster:
 
-1. Install caption-flow: `pip install caption-flow`
+1. Install caption-flow: `pip install "caption-flow[vllm]"`
 2. Get a worker token from the project maintainer
 3. Run: `caption-flow worker --server wss://project.domain.com:8765 --token YOUR_TOKEN`
 

{caption_flow-0.5.2 → caption_flow-0.5.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "caption-flow"
-version = "0.5.2"
+version = "0.5.3"
 description = "Self-contained distributed community captioning system"
 readme = "README.md"
 requires-python = ">=3.11,<3.14"

{caption_flow-0.5.2 → caption_flow-0.5.3}/src/caption_flow/__init__.py

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

{caption_flow-0.5.2 → caption_flow-0.5.3}/src/caption_flow/processors/huggingface.py

@@ -32,6 +32,8 @@ from .base import OrchestratorProcessor, ProcessorConfig, WorkerProcessor, WorkR
 logger = logging.getLogger(__name__)
 logger.setLevel(os.environ.get("CAPTIONFLOW_LOG_LEVEL", "INFO").upper())
 
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif", ".tiff", ".tif"}
+
 
 def log_memory(location: str):
     """Log memory usage at specific location."""
@@ -244,6 +246,10 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
             return match.group(1)
         return url.split("/")[-1]
 
+    def _is_image_file(self, filename: str) -> bool:
+        """Return whether a Hugging Face dataset file is a supported raw image."""
+        return Path(filename).suffix.lower() in IMAGE_EXTENSIONS
+
     def _get_data_files_from_builder(self) -> List[str]:
         """Get data files using dataset builder with minimal memory usage."""
         # Load builder to get correct file structure
@@ -317,21 +323,29 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
                     token=self.token,
                 )
 
-                # Read only metadata
-                metadata = pq.read_metadata(local_path)
-                size = metadata.num_rows
+                file_type = "parquet"
+                try:
+                    # Read only metadata for parquet-backed datasets.
+                    metadata = pq.read_metadata(local_path)
+                    size = metadata.num_rows
+                except Exception:
+                    if not self._is_image_file(filename):
+                        raise
+                    file_type = "image"
+                    size = 1
 
                 self.shard_info[i] = {
                     "shard_id": i,
                     "file_url": file_url,
                     "filename": filename,
+                    "file_type": file_type,
                     "start_offset": cumulative_offset,
                     "size": size,
                     "end_offset": cumulative_offset + size - 1,
                 }
 
                 cumulative_offset += size
-                logger.info(f"Shard {i} ({filename}): {size} rows")
+                logger.info(f"Shard {i} ({filename}): {size} {file_type} item(s)")
 
             except Exception as e:
                 logger.error(f"Failed to discover shard {i}: {e}")
@@ -371,6 +385,14 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
                 return shard_id, local_index
         raise ValueError(f"Global index {global_index} not found in any shard")
 
+    def _get_shards_for_range(self, start_index: int, end_index: int) -> List[int]:
+        """Get all shard IDs intersecting an absolute index range."""
+        shard_ids = []
+        for shard_id, sinfo in self.shard_info.items():
+            if sinfo["start_offset"] <= end_index and sinfo["end_offset"] >= start_index:
+                shard_ids.append(shard_id)
+        return shard_ids
+
     def _restore_state(self, storage: StorageManager) -> None:
         """Restore state from chunk tracker and synchronize with storage."""
         logger.debug("Restoring state from chunk tracker and synchronizing with storage")
@@ -468,9 +490,10 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
         current_index = chunk_index * self.chunk_size
         chunk_size = min(self.chunk_size, self.total_items - current_index)
 
-        # Find shard for this chunk
+        # Find all shards touched by this chunk/range set
         shard_id, local_idx = self._get_shard_for_index(current_index)
         shard_name = Path(self.shard_info[shard_id]["filename"]).stem
+        shard_ids = self._get_shards_for_range(ranges[0][0], ranges[-1][1])
 
         # Create unique unit ID that includes range info
         range_suffix = f"r{len(ranges)}"  # r2 = 2 ranges, etc.
@@ -497,8 +520,9 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
                 "unprocessed_ranges": ranges,
                 "range_based": True,
                 "is_split_unit": True,  # Flag to indicate this is a split from larger chunk
-                "shard_ids": [shard_id],
+                "shard_ids": shard_ids,
                 "data_files": self.data_files,
+                "shard_info": self.shard_info,
             },
             metadata={
                 "dataset": self.dataset_name,
@@ -519,9 +543,10 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
 
         chunk_size = min(self.chunk_size, self.total_items - current_index)
 
-        # Find shard for this chunk
+        # Find all shards touched by this chunk
         shard_id, local_idx = self._get_shard_for_index(current_index)
         shard_name = Path(self.shard_info[shard_id]["filename"]).stem
+        shard_ids = self._get_shards_for_range(current_index, current_index + chunk_size - 1)
 
         # Calculate RELATIVE chunk index within the shard
         job_id_obj = JobId(
@@ -556,13 +581,6 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
         # Calculate actual unprocessed items and total work to be assigned
         unprocessed_items = sum(end - start + 1 for start, end in unprocessed_ranges)
 
-        # Skip assignment if there are very few unprocessed items (< 10 items)
-        if unprocessed_items < 10:
-            logger.debug(
-                f"Chunk {unit_id} has only {unprocessed_items} unprocessed items, skipping assignment"
-            )
-            return None
-
         # Create work unit that represents ONLY the unprocessed ranges
         # This is the key fix: don't assign the full chunk, assign only unprocessed parts
         unit = WorkUnit(
@@ -579,8 +597,9 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
                 "actual_work_size": unprocessed_items,  # NEW: actual work to be done
                 "unprocessed_ranges": unprocessed_ranges,  # The specific ranges to process
                 "range_based": True,  # NEW: flag to indicate this is range-based
-                "shard_ids": [shard_id],
+                "shard_ids": shard_ids,
                 "data_files": self.data_files,
+                "shard_info": self.shard_info,
             },
             metadata={
                 "dataset": self.dataset_name,
@@ -629,7 +648,7 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
                     )
                     break
                 # Get shard info for proper unit_id
-                current_index = self.current_chunk_index
+                current_index = self.current_chunk_index * self.chunk_size
                 if current_index < self.total_items:
                     shard_id, _ = self._get_shard_for_index(current_index)
                     shard_name = Path(self.shard_info[shard_id]["filename"]).stem
@@ -980,6 +999,9 @@ class HuggingFaceDatasetOrchestratorProcessor(OrchestratorProcessor):
 
                     for start_idx, end_idx in ranges:
                         self.chunk_tracker.mark_items_processed(result.chunk_id, start_idx, end_idx)
+            elif "_item_index" in result.metadata and result.metadata["_item_index"] is not None:
+                item_index = int(result.metadata["_item_index"])
+                self.chunk_tracker.mark_items_processed(result.chunk_id, item_index, item_index)
 
         return base_result
 
@@ -1070,6 +1092,10 @@ class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):
             return match.group(1)
         return url.split("/")[-1]
 
+    def _is_image_file(self, filename: str) -> bool:
+        """Return whether a Hugging Face dataset file is a supported raw image."""
+        return Path(filename).suffix.lower() in IMAGE_EXTENSIONS
+
     def _create_dummy_image(self, index: int, metadata: Dict[str, Any]) -> Image.Image:
         """Create a dummy image."""
         color = (0, 0, 0)
@@ -1091,13 +1117,41 @@ class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):
         )
         shard_ids = unit.data.get("shard_ids", [])
         data_files = unit.data.get("data_files", [])
+        unit_shard_info = unit.data.get("shard_info", {})
+        unit_shard_info = {int(k): v for k, v in unit_shard_info.items()}
 
         logger.info(f"Processing unit {unit.unit_id} with ranges: {unprocessed_ranges}")
 
         # Build shard info from provided data files (no dataset builder needed)
         shard_info = {}
 
-        if data_files:
+        if unit_shard_info:
+            for shard_id in shard_ids:
+                shard_id = int(shard_id)
+                if shard_id not in unit_shard_info:
+                    continue
+                metadata = unit_shard_info[shard_id]
+                filename = metadata["filename"]
+                shard_path = self._get_shard_path(dataset_name, filename)
+                file_type = metadata.get("file_type") or (
+                    "image" if self._is_image_file(filename) else "parquet"
+                )
+
+                shard_info[shard_id] = {
+                    "path": shard_path,
+                    "filename": filename,
+                    "file_type": file_type,
+                    "start_offset": metadata["start_offset"],
+                    "end_offset": metadata["end_offset"],
+                    "size": metadata["size"],
+                    "metadata": None,
+                }
+
+                if file_type == "parquet":
+                    shard_info[shard_id]["metadata"] = pq.read_metadata(shard_path)
+
+        elif data_files:
+            shard_ids = {int(shard_id) for shard_id in shard_ids}
             # Use provided data files
             for i, file_url in enumerate(data_files):
                 if i in shard_ids:
@@ -1110,6 +1164,8 @@ class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):
 
                     shard_info[i] = {
                         "path": shard_path,
+                        "filename": filename,
+                        "file_type": "parquet",
                         "start_offset": 0,  # Will be set below
                         "end_offset": 0,  # Will be set below
                         "size": size,
@@ -1152,6 +1208,49 @@ class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):
         # Process items shard by shard
         for shard_id, idx_pairs in indices_by_shard.items():
             shard_path = shard_info[shard_id]["path"]
+            file_type = shard_info[shard_id].get("file_type", "parquet")
+            shard_name = Path(shard_info[shard_id]["filename"]).stem
+
+            if file_type == "image":
+                for global_idx, local_idx in idx_pairs:
+                    try:
+                        with Image.open(shard_path) as img:
+                            image = img.copy()
+
+                        chunk_index = unit.metadata["chunk_index"]
+                        job_id_obj = JobId(
+                            shard_id=shard_name,
+                            chunk_id=str(chunk_index),
+                            sample_id=str(global_idx),
+                        )
+                        job_id = job_id_obj.get_sample_str()
+
+                        clean_metadata = {
+                            "_item_index": global_idx,
+                            "_chunk_relative_index": global_idx - start_index,
+                            "_job_id": job_id,
+                            "_shard_id": shard_id,
+                            "_local_index": local_idx,
+                            "_filename": shard_info[shard_id]["filename"],
+                            "_url": None,
+                            "_mock": self.mock_results,
+                        }
+
+                        yield {
+                            "image": image,
+                            "item_key": str(global_idx),
+                            "item_index": global_idx,
+                            "metadata": clean_metadata,
+                            "job_id": job_id,
+                            "_processed_indices": processed_indices,
+                        }
+
+                        processed_indices.append(global_idx)
+
+                    except Exception as e:
+                        logger.error(f"Error processing image file at index {global_idx}: {e}")
+
+                continue
 
             # Process in batches to avoid loading entire table
             batch_size = 100
@@ -1325,11 +1424,10 @@ class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):
 
                             # Build job ID
                             chunk_index = unit.metadata["chunk_index"]
-                            shard_name = unit.metadata["shard_name"]
                             job_id_obj = JobId(
                                 shard_id=shard_name,
                                 chunk_id=str(chunk_index),
-                                sample_id=str(local_idx),
+                                sample_id=str(global_idx),
                             )
                             job_id = job_id_obj.get_sample_str()
 
@@ -1362,7 +1460,7 @@ class HuggingFaceDatasetWorkerProcessor(WorkerProcessor):
                                 "_processed_indices": processed_indices,
                             }
 
-                            processed_indices.append(local_idx)
+                            processed_indices.append(global_idx)
 
                         except Exception as e:
                             logger.error(f"Error processing item at index {global_idx}: {e}")

{caption_flow-0.5.2 → caption_flow-0.5.3}/src/caption_flow/utils/vllm_config.py

@@ -31,6 +31,9 @@ class VLLMConfigManager:
         "enforce_eager",
         "limit_mm_per_prompt",
         "disable_mm_preprocessor_cache",
+        "mm_processor_cache_gb",
+        "mm_processor_cache_type",
+        "mm_shm_cache_max_object_size_mb",
     }
 
     # Fields that can be updated without reload
@@ -115,7 +118,7 @@ class VLLMConfigManager:
 
     def get_vllm_init_params(self, vllm_config: Dict[str, Any]) -> Dict[str, Any]:
         """Extract vLLM initialization parameters from config."""
-        return {
+        params = {
             "model": vllm_config["model"],
             "trust_remote_code": True,
             "tensor_parallel_size": vllm_config.get("tensor_parallel_size", 1),
@@ -126,6 +129,15 @@ class VLLMConfigManager:
             "limit_mm_per_prompt": vllm_config.get("limit_mm_per_prompt", {"image": 1}),
             "disable_mm_preprocessor_cache": vllm_config.get("disable_mm_preprocessor_cache", True),
         }
+        for optional_key in (
+            "mm_processor_cache_gb",
+            "mm_processor_cache_type",
+            "mm_shm_cache_max_object_size_mb",
+        ):
+            if vllm_config.get(optional_key) is not None:
+                params[optional_key] = vllm_config[optional_key]
+
+        return params
 
     def requires_tokenizer_reload(
         self, old_config: Optional[Dict[str, Any]], new_config: Dict[str, Any]

{caption_flow-0.5.2 → caption_flow-0.5.3}/src/caption_flow/workers/caption.py

@@ -5,6 +5,7 @@ import os
 os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
 
 import asyncio
+import inspect
 import json
 import logging
 import time
@@ -66,6 +67,39 @@ class MultiStageVLLMManager:
         self.tokenizers: Dict[str, Any] = {}  # model_name -> tokenizer
         self.sampling_params: Dict[str, Any] = {}  # stage_name -> SamplingParams
 
+    @staticmethod
+    def _filter_engine_args(params: Dict[str, Any], engine_args_cls: Any) -> Dict[str, Any]:
+        """Filter vLLM init kwargs to those supported by the installed EngineArgs."""
+        signature = inspect.signature(engine_args_cls)
+        supported_params = set(signature.parameters)
+        filtered = params.copy()
+
+        disable_mm_cache = filtered.pop("disable_mm_preprocessor_cache", None)
+        if (
+            disable_mm_cache is True
+            and "disable_mm_preprocessor_cache" not in supported_params
+            and "mm_processor_cache_gb" in supported_params
+            and "mm_processor_cache_gb" not in filtered
+        ):
+            filtered["mm_processor_cache_gb"] = 0
+            logger.info(
+                "Mapped disable_mm_preprocessor_cache=True to mm_processor_cache_gb=0 "
+                "for this vLLM version"
+            )
+        elif "disable_mm_preprocessor_cache" in supported_params and disable_mm_cache is not None:
+            filtered["disable_mm_preprocessor_cache"] = disable_mm_cache
+
+        unsupported = sorted(key for key in filtered if key not in supported_params)
+        if unsupported:
+            logger.warning(
+                "Dropping unsupported vLLM initialization parameter(s): %s",
+                ", ".join(unsupported),
+            )
+            for key in unsupported:
+                filtered.pop(key, None)
+
+        return filtered
+
     def load_model(self, model_name: str, stage: ProcessingStage, base_config: Dict[str, Any]):
         """Load a model if not already loaded."""
         if model_name in self.models:
@@ -74,6 +108,7 @@ class MultiStageVLLMManager:
 
         from transformers import AutoProcessor, AutoTokenizer
         from vllm import LLM
+        from vllm.engine.arg_utils import EngineArgs
 
         logger.info(f"Loading model {model_name} for stage {stage.name}")
 
@@ -113,8 +148,15 @@ class MultiStageVLLMManager:
                 "disable_mm_preprocessor_cache", True
             ),
         }
-
-        self.models[model_name] = LLM(**vllm_params)
+        for optional_key in (
+            "mm_processor_cache_gb",
+            "mm_processor_cache_type",
+            "mm_shm_cache_max_object_size_mb",
+        ):
+            if model_config.get(optional_key) is not None:
+                vllm_params[optional_key] = model_config[optional_key]
+
+        self.models[model_name] = LLM(**self._filter_engine_args(vllm_params, EngineArgs))
         logger.info(f"Model {model_name} loaded successfully")
 
     def create_sampling_params(self, stage: ProcessingStage, base_sampling: Dict[str, Any]):
@@ -544,6 +586,11 @@ class CaptionWorker(BaseWorker):
                 "disable_mm_preprocessor_cache", True
             ),
             "limit_mm_per_prompt": self.vllm_config.get("limit_mm_per_prompt", {"image": 1}),
+            "mm_processor_cache_gb": self.vllm_config.get("mm_processor_cache_gb"),
+            "mm_processor_cache_type": self.vllm_config.get("mm_processor_cache_type"),
+            "mm_shm_cache_max_object_size_mb": self.vllm_config.get(
+                "mm_shm_cache_max_object_size_mb"
+            ),
         }
 
         base_sampling = self.vllm_config.get("sampling", {})

{caption_flow-0.5.2 → caption_flow-0.5.3/src/caption_flow.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: caption-flow
-Version: 0.5.2
+Version: 0.5.3
 Summary: Self-contained distributed community captioning system
 Author-email: bghira <bghira@users.github.com>
 License: MIT
@@ -76,7 +76,7 @@ a fast websocket-based orchestrator paired with lightweight gpu workers achieves
 ```bash
 python -m venv .venv
 source .venv/bin/activate  # windows: .venv\Scripts\activate
-pip install caption-flow
+pip install "caption-flow[vllm]"
 ```
 
 ## quickstart (single box)
@@ -240,7 +240,7 @@ PRs welcome. keep it simple and fast.
 
 To contribute compute to a cluster:
 
-1. Install caption-flow: `pip install caption-flow`
+1. Install caption-flow: `pip install "caption-flow[vllm]"`
 2. Get a worker token from the project maintainer
 3. Run: `caption-flow worker --server wss://project.domain.com:8765 --token YOUR_TOKEN`
 

{caption_flow-0.5.2 → caption_flow-0.5.3}/tests/test_processors.py

@@ -455,6 +455,42 @@ class TestHuggingFaceDatasetProcessors(ProcessorTestBase):
                 finally:
                     orchestrator.cleanup()
 
+    def test_raw_image_file_discovery(self, orchestrator_config, storage_manager, temp_dir):
+        """Test Hugging Face repos that expose raw image files instead of parquet shards."""
+        image_paths = []
+        data_files = []
+        for i in range(3):
+            image_path = temp_dir / f"image-{i}.jpg"
+            Image.new("RGB", (10, 10), color=(i, i, i)).save(image_path)
+            image_paths.append(str(image_path))
+            data_files.append(f"image-{i}.jpg")
+
+        with patch(
+            "caption_flow.processors.huggingface.HuggingFaceDatasetOrchestratorProcessor._get_data_files_from_builder",
+            return_value=data_files,
+        ):
+            with patch(
+                "caption_flow.processors.huggingface.hf_hub_download",
+                side_effect=image_paths,
+            ):
+                orchestrator = HuggingFaceDatasetOrchestratorProcessor()
+                try:
+                    orchestrator.initialize(orchestrator_config, storage_manager)
+
+                    assert orchestrator.total_items == 3
+                    assert len(orchestrator.shard_info) == 3
+                    assert all(
+                        shard["file_type"] == "image"
+                        for shard in orchestrator.shard_info.values()
+                    )
+
+                    unit = orchestrator._create_work_unit(0)
+                    assert unit is not None
+                    assert unit.unit_size == 3
+                    assert unit.data["shard_ids"] == [0, 1, 2]
+                finally:
+                    orchestrator.cleanup()
+
     # For test_worker_processing_with_ranges, replace with:
     def test_worker_processing_with_ranges(self, worker_config, temp_dir):
         """Test worker processing with specific unprocessed ranges."""
@@ -517,6 +553,68 @@ class TestHuggingFaceDatasetProcessors(ProcessorTestBase):
             processed_indices = context.get("_processed_indices", [])
             assert len(processed_indices) == 20
 
+    def test_worker_processing_raw_image_files(self, worker_config, temp_dir):
+        """Test worker processing for raw image files from Hugging Face repos."""
+        worker = HuggingFaceDatasetWorkerProcessor()
+        worker.gpu_id = 0
+
+        image_paths = []
+        for i in range(2):
+            image_path = temp_dir / f"raw-{i}.webp"
+            Image.new("RGB", (10, 10), color=(i, i, i)).save(image_path)
+            image_paths.append(str(image_path))
+
+        with patch(
+            "caption_flow.processors.huggingface.hf_hub_download",
+            side_effect=image_paths,
+        ):
+            worker.initialize(worker_config)
+
+            unit = WorkUnit(
+                unit_id="raw-0:chunk:0",
+                chunk_id="raw-0:chunk:0",
+                source_id="raw-0",
+                unit_size=2,
+                data={
+                    "dataset_name": "test/dataset",
+                    "config": "default",
+                    "split": "train",
+                    "start_index": 0,
+                    "chunk_size": 2,
+                    "unprocessed_ranges": [(0, 1)],
+                    "shard_ids": [0, 1],
+                    "data_files": ["raw-0.webp", "raw-1.webp"],
+                    "shard_info": {
+                        0: {
+                            "filename": "raw-0.webp",
+                            "file_type": "image",
+                            "start_offset": 0,
+                            "end_offset": 0,
+                            "size": 1,
+                        },
+                        1: {
+                            "filename": "raw-1.webp",
+                            "file_type": "image",
+                            "start_offset": 1,
+                            "end_offset": 1,
+                            "size": 1,
+                        },
+                    },
+                },
+                metadata={"chunk_index": 0, "shard_name": "raw-0"},
+            )
+
+            context = {}
+            items = list(worker.process_unit(unit, context))
+
+            assert len(items) == 2
+            assert [item["item_index"] for item in items] == [0, 1]
+            assert [item["job_id"] for item in items] == [
+                "raw-0:chunk:0:idx:0",
+                "raw-1:chunk:0:idx:1",
+            ]
+            assert context["_processed_indices"] == [0, 1]
+
     @pytest.mark.asyncio
     async def test_storage_update_flow(self, orchestrator_config, storage_manager, temp_dir):
         """Test updating chunk tracker from storage."""

{caption_flow-0.5.2 → caption_flow-0.5.3}/tests/test_vllm_config.py

@@ -268,6 +268,22 @@ class TestVLLMConfigManagerInitParams:
 
         assert params == expected
 
+    def test_get_vllm_init_params_new_mm_cache_config(self):
+        """Test getting vLLM init params with current multimodal cache config."""
+        manager = VLLMConfigManager()
+        vllm_config = {
+            "model": "custom-model",
+            "mm_processor_cache_gb": 8,
+            "mm_processor_cache_type": "shm",
+            "mm_shm_cache_max_object_size_mb": 256,
+        }
+
+        params = manager.get_vllm_init_params(vllm_config)
+
+        assert params["mm_processor_cache_gb"] == 8
+        assert params["mm_processor_cache_type"] == "shm"
+        assert params["mm_shm_cache_max_object_size_mb"] == 256
+
 
 class TestVLLMConfigManagerTokenizerReload:
     """Test VLLMConfigManager tokenizer reload checking."""

{caption_flow-0.5.2 → caption_flow-0.5.3}/tests/test_worker_caption.py

@@ -19,6 +19,7 @@ from caption_flow.storage import StorageManager
 # Import the modules to test
 from caption_flow.workers.caption import (
     CaptionWorker,
+    MultiStageVLLMManager,
     ProcessingItem,
 )
 
@@ -37,6 +38,96 @@ def create_fast_caption_worker(worker_config):
                         return worker
 
 
+class TestMultiStageVLLMManager:
+    """Test vLLM manager compatibility helpers."""
+
+    def test_filter_engine_args_maps_removed_mm_cache_flag(self):
+        """Test newer vLLM EngineArgs compatibility for removed cache flag."""
+
+        class EngineArgs:
+            def __init__(
+                self,
+                model,
+                trust_remote_code=False,
+                limit_mm_per_prompt=None,
+                mm_processor_cache_gb=4,
+            ):
+                pass
+
+        params = {
+            "model": "test-model",
+            "trust_remote_code": True,
+            "limit_mm_per_prompt": {"image": 1},
+            "disable_mm_preprocessor_cache": True,
+            "unsupported": "drop-me",
+        }
+
+        filtered = MultiStageVLLMManager._filter_engine_args(params, EngineArgs)
+
+        assert filtered == {
+            "model": "test-model",
+            "trust_remote_code": True,
+            "limit_mm_per_prompt": {"image": 1},
+            "mm_processor_cache_gb": 0,
+        }
+
+    def test_filter_engine_args_preserves_explicit_mm_cache_size(self):
+        """Test explicit newer vLLM cache config takes precedence over legacy alias."""
+
+        class EngineArgs:
+            def __init__(
+                self,
+                model,
+                trust_remote_code=False,
+                mm_processor_cache_gb=4,
+                mm_processor_cache_type="lru",
+            ):
+                pass
+
+        params = {
+            "model": "test-model",
+            "trust_remote_code": True,
+            "disable_mm_preprocessor_cache": True,
+            "mm_processor_cache_gb": 8,
+            "mm_processor_cache_type": "shm",
+        }
+
+        filtered = MultiStageVLLMManager._filter_engine_args(params, EngineArgs)
+
+        assert filtered == {
+            "model": "test-model",
+            "trust_remote_code": True,
+            "mm_processor_cache_gb": 8,
+            "mm_processor_cache_type": "shm",
+        }
+
+    def test_filter_engine_args_preserves_supported_mm_cache_flag(self):
+        """Test older vLLM EngineArgs compatibility for the legacy cache flag."""
+
+        class EngineArgs:
+            def __init__(
+                self,
+                model,
+                trust_remote_code=False,
+                disable_mm_preprocessor_cache=True,
+            ):
+                pass
+
+        params = {
+            "model": "test-model",
+            "trust_remote_code": True,
+            "disable_mm_preprocessor_cache": False,
+        }
+
+        filtered = MultiStageVLLMManager._filter_engine_args(params, EngineArgs)
+
+        assert filtered == {
+            "model": "test-model",
+            "trust_remote_code": True,
+            "disable_mm_preprocessor_cache": False,
+        }
+
+
 class TestCaptionWorker:
     """Test suite for CaptionWorker."""