From 48d3c769b2c0e1119da70a70315da82cc998640e Mon Sep 17 00:00:00 2001 From: bghira Date: Sun, 6 Sep 2026 20:01:00 -0600 Subject: [PATCH 1/3] fix: harden remote captioning workflows --- recipes/GLM_5.3_FLASH.md | 214 ++++++++++++++++++ src/caption_flow/models.py | 6 + src/caption_flow/processors/webdataset.py | 47 +++- src/caption_flow/utils/chunk_tracker.py | 33 +++ src/caption_flow/utils/vllm_config.py | 2 + src/caption_flow/workers/caption.py | 9 + src/caption_flow/workers/openai_compatible.py | 179 +++++++++++++-- tests/test_openai_compatible_worker.py | 162 +++++++++++++ tests/test_vllm_config.py | 2 + tests/test_webdataset_ranges.py | 59 +++++ tests/test_worker_caption.py | 10 +- 11 files changed, 688 insertions(+), 35 deletions(-) create mode 100644 recipes/GLM_5.3_FLASH.md diff --git a/recipes/GLM_5.3_FLASH.md b/recipes/GLM_5.3_FLASH.md new file mode 100644 index 0000000..78b0293 --- /dev/null +++ b/recipes/GLM_5.3_FLASH.md @@ -0,0 +1,214 @@ +# GLM-5.3-Flash with vLLM + +This recipe serves GLM-5.3-Flash from local weights with vLLM and connects it +to CaptionFlow through the OpenAI-compatible worker. No GLM-specific +CaptionFlow code is required once the vLLM Chat Completions endpoint is +available. + +## Tested configuration + +- 4 NVIDIA H100 80 GB GPUs +- `wtdcode/GLM-5.3-Flash-AWQ-W4A16`, revision + `abd7b07719111f137e1de8a0c1b7e01c11b74d1a` +- `vllm/vllm-openai:glm53-flash` +- vLLM `0.1.dev20051+g487ecf187` +- tensor parallelism across all four GPUs +- an 8,192-token model context +- up to 64 concurrent vLLM sequences + +The tested checkpoint occupies approximately 177.7 GiB on disk and about +41.7 GiB of model memory per GPU with tensor parallel size 4. It quantizes the +routed MoE experts to W4A16 while retaining the vision encoder and several +sensitive components in BF16. + +The official FP8 checkpoint did not fit this 4x80 GB setup: its estimated +runtime requirement was approximately 386 GiB, exceeding the 320 GB of total +VRAM. Use the W4A16 checkpoint for this hardware configuration. This recipe +does not claim that W4A16 and FP8 have identical output quality. + +## Download the checkpoint + +Allow at least 180 GiB of free local storage: + +```bash +export GLM_MODEL_DIR=/models/GLM-5.3-Flash-AWQ-W4A16 + +hf download wtdcode/GLM-5.3-Flash-AWQ-W4A16 \ + --revision abd7b07719111f137e1de8a0c1b7e01c11b74d1a \ + --local-dir "$GLM_MODEL_DIR" +``` + +Pinning the revision prevents a later checkpoint update from silently changing +the deployment. + +## Start vLLM + +Use a vLLM build that recognizes `Glm5NextForConditionalGeneration`. The +tested image reports vLLM version `0.1.dev20051+g487ecf187`; an older stable +image may reject the architecture. + +Inside that image or environment, start the server with: + +```bash +vllm serve "$GLM_MODEL_DIR" \ + --host 127.0.0.1 \ + --port 8000 \ + --served-model-name glm-5.3-flash \ + --tensor-parallel-size 4 \ + --max-model-len 8192 \ + --max-num-seqs 64 \ + --gpu-memory-utilization 0.94 \ + --reasoning-parser glm45 \ + --no-enable-flashinfer-autotune +``` + +Bind to `0.0.0.0` only when the endpoint must be reachable from another host, +and protect it with appropriate network controls or authentication. + +The first startup can take roughly five minutes while vLLM loads the weights, +compiles kernels, profiles memory, and captures CUDA graphs. The tested run +loaded the weights in about 80 seconds and completed the full engine startup in +about five minutes. + +`--no-enable-flashinfer-autotune` skips an additional startup autotuning phase. +It is part of the known-good invocation, but it is not a CaptionFlow +requirement. Benchmark before removing it or changing the vLLM kernel setup. + +Confirm that the server is ready: + +```bash +curl --fail http://127.0.0.1:8000/health +curl --fail http://127.0.0.1:8000/v1/models +``` + +The model list should contain `glm-5.3-flash`. + +## Configure the CaptionFlow worker + +Install CaptionFlow's lightweight API worker dependencies in the environment +that will run the worker: + +```bash +pip install "caption-flow[openai]" +``` + +Create a worker configuration such as `worker.glm-5.3-flash.yaml`: + +```yaml +worker: + server: "ws://orchestrator.example:8765" + token: "replace-with-captionflow-worker-token" + name: "vllm-glm-5.3-flash" + when_finished: "stay_connected" + batch_image_processing: false + + openai_compatible: + batch_size: 64 + include_image: true + image_detail: "auto" + image_format: "jpeg" + image_quality: 90 + max_image_dimension: 1024 + system_prompt: >- + You are a precise, minimalist image captioner. Write short factual + descriptions using only visually supported details. Return only the + requested caption. + + endpoints: + - name: "local-vllm-glm-5.3-flash" + base_url: "http://127.0.0.1:8000/v1" + api_key_env: "VLLM_API_KEY" + model: "glm-5.3-flash" + initial_concurrency: 64 + max_concurrency: 64 + probe_after_successes: 8 + timeout_seconds: 300 + max_retries: 3 + extra_body: + chat_template_kwargs: + reasoning_effort: "low" +``` + +The OpenAI client requires a non-empty API key even when the local vLLM server +does not enforce authentication. Give it a non-secret placeholder unless the +server was started with a real API key: + +```bash +export VLLM_API_KEY=local-vllm +caption-flow worker \ + --config worker.glm-5.3-flash.yaml \ + --openai-compatible +``` + +The worker token must match one of the orchestrator's `auth.worker_tokens`. +When vLLM and the CaptionFlow worker run in different containers, replace +`127.0.0.1` with a private, reachable vLLM address. + +## Configure caption generation + +The orchestrator owns the prompt and sampling configuration. This is the +configuration used for concise captions: + +```yaml +orchestrator: + inference: + model: "glm-5.3-flash" + batch_size: 8 + sampling: + temperature: 0.3 + top_p: 0.9 + max_tokens: 3000 + inference_prompts: + - >- + Write one concise, minimalist standalone caption for this photograph in + one or two short sentences, preferably 20 to 45 words. Identify only + the main subject, action, and essential setting or visual context. Omit + exhaustive detail, interpretation, and speculation. Do not use labels + or bullet points. Return only the caption. +``` + +The high token ceiling prevents reasoning from being truncated. The prompt and +`reasoning_effort: low` still keep the returned caption short, while +`--reasoning-parser glm45` separates model reasoning from the final content. + +Hosted providers may reject an image in a safety prefilter before GLM sees the +captioning prompt. When the source dataset contains a usable text description, +the OpenAI-compatible worker can make one text-only fallback attempt: + +```yaml +orchestrator: + inference: + retry_prompt: >- + Rewrite the following source description as one concise, neutral, + standalone caption. Preserve only concrete visual details, omit + speculation and sensitive details, and return only the caption. + Source description: {column:captions} + retry_without_image: true +``` + +The fallback runs only when every primary output for an item is empty, a known +refusal, or a provider content-policy error. Ordinary authentication, network, +and malformed-request failures remain visible instead of being converted into +captions. Replace `captions` in `{column:captions}` with the source metadata +column available in the dataset. + +## Operational notes + +- Keep the served model name and the endpoint model name identical. The tested + alias is exactly `glm-5.3-flash`. +- `batch_size: 64` and `initial_concurrency: 64` keep the local endpoint full. + On the tested 4x H100 system, 128 matched images took 20.1 seconds at this + setting versus 71.4 seconds when the endpoint began at concurrency 8. Use a + conservative initial value for rate-limited hosted providers. +- `max_image_dimension: 1024` bounds vision-token use and request size. Without + normalization, 40 of 128 high-resolution test images exceeded the shared + 8,192-token context window. +- Reduce `initial_concurrency`, `max_concurrency`, or `max_num_seqs` if the + workload encounters memory pressure. Increase them only after measuring + throughput and latency on representative images. +- A warning that no MLA prefill backend supports the model was non-fatal in the + tested vLLM build. Sparse MLA used its supported top-k path and captioning + continued normally. +- Do not judge quantization quality from unmatched worker outputs. A useful + comparison requires the same images, prompt, sampling settings, and decoding + path on W4A16 and a higher-precision checkpoint. diff --git a/src/caption_flow/models.py b/src/caption_flow/models.py index b30be60..b0bf00d 100644 --- a/src/caption_flow/models.py +++ b/src/caption_flow/models.py @@ -179,6 +179,12 @@ 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. + retry_prompt: Optional[str] = None + retry_without_image: bool = False + @dataclass class StageResult: diff --git a/src/caption_flow/processors/webdataset.py b/src/caption_flow/processors/webdataset.py index 59c44af..f84187f 100644 --- a/src/caption_flow/processors/webdataset.py +++ b/src/caption_flow/processors/webdataset.py @@ -130,6 +130,25 @@ def _restore_state(self, storage: StorageManager) -> None: shards_summary = self.chunk_tracker.get_shards_summary() logger.info(f"Restoring work units from chunk tracker: {len(shards_summary)} shards") + # Resolve each current shard once. A checkpoint may describe an older + # version of a mutable dataset whose final chunk was larger. + incomplete_shards = { + shard_name + for shard_name, summary in shards_summary.items() + if any(chunk.status != "completed" for chunk in summary.get("chunks", [])) + } + current_shards = {} + if self.dataset and incomplete_shards: + for shard_idx in range(self.dataset.num_shards): + current_info = self._get_shard_info_cached(shard_idx) + if not current_info or current_info.get("name") not in incomplete_shards: + continue + sample_count = current_info.get("num_samples", current_info.get("num_files")) + if sample_count is not None: + current_shards[current_info["name"]] = (shard_idx, sample_count) + if len(current_shards) == len(incomplete_shards): + break + with self.lock: restored_count = 0 for shard_name, shard_info in shards_summary.items(): @@ -140,6 +159,23 @@ def _restore_state(self, storage: StorageManager) -> None: logger.debug(f"Skipping completed chunk {chunk_state.chunk_id}") continue + current_shard = current_shards.get(shard_name) + if current_shard: + _shard_idx, sample_count = current_shard + current_size = max( + 0, + min(chunk_state.chunk_size, sample_count - chunk_state.start_index), + ) + if self.chunk_tracker.shrink_chunk(chunk_state.chunk_id, current_size): + logger.warning( + "Clamped restored chunk %s to %d samples " + "using current shard bounds", + chunk_state.chunk_id, + current_size, + ) + if chunk_state.status == "completed": + continue + # Get unprocessed ranges unprocessed_ranges = chunk_state.get_unprocessed_ranges() if not unprocessed_ranges: @@ -161,19 +197,13 @@ def _restore_state(self, storage: StorageManager) -> None: 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 + shard_idx = current_shard[0] if current_shard else None unit = WorkUnit( unit_id=chunk_state.chunk_id, chunk_id=chunk_state.chunk_id, source_id=shard_name, - unit_size=chunk_state.chunk_size, + unit_size=sum(end - start + 1 for start, end in absolute_ranges), data={ "shard_url": chunk_state.shard_url, "shard_name": shard_name, @@ -345,6 +375,7 @@ def get_work_units(self, count: int, worker_id: str) -> List[WorkUnit]: # Update the work unit's unprocessed ranges unit.data["unprocessed_ranges"] = absolute_ranges + unit.unit_size = sum(end - start + 1 for start, end in absolute_ranges) logger.debug( f"Updated unit {unit_id} with unprocessed ranges: {absolute_ranges}" diff --git a/src/caption_flow/utils/chunk_tracker.py b/src/caption_flow/utils/chunk_tracker.py index ed4d045..4b435cf 100644 --- a/src/caption_flow/utils/chunk_tracker.py +++ b/src/caption_flow/utils/chunk_tracker.py @@ -369,6 +369,39 @@ def mark_completed(self, chunk_id: str): if self._completed_count % 50 == 0: self._limit_completed_chunks_in_memory() + def shrink_chunk(self, chunk_id: str, chunk_size: int) -> bool: + """Shrink a chunk to the source's current bounds. + + Checkpoints can outlive a mutable remote dataset. Keep already + processed indices that still exist and complete the chunk when the + remaining range consisted only of indices past the end of the shard. + """ + if chunk_size < 0: + raise ValueError("chunk_size must be non-negative") + + chunk = self.chunks.get(chunk_id) + if not chunk or chunk_size >= chunk.chunk_size: + return False + + clamped_ranges = [ + (max(0, start), min(end, chunk_size - 1)) + for start, end in chunk.processed_ranges + if start < chunk_size and end >= 0 + ] + chunk.chunk_size = chunk_size + chunk.processed_ranges = chunk._merge_ranges(clamped_ranges) + chunk.processed_count = sum(end - start + 1 for start, end in chunk.processed_ranges) + chunk._invalidate_cache() + + if chunk.processed_count >= chunk_size: + was_completed = chunk.status == "completed" + chunk.mark_completed() + if not was_completed: + self._completed_count += 1 + + self._mark_dirty() + return True + def mark_failed(self, chunk_id: str): """Mark chunk as failed.""" if chunk_id in self.chunks: diff --git a/src/caption_flow/utils/vllm_config.py b/src/caption_flow/utils/vllm_config.py index 0d1b834..1a28459 100644 --- a/src/caption_flow/utils/vllm_config.py +++ b/src/caption_flow/utils/vllm_config.py @@ -41,6 +41,8 @@ class VLLMConfigManager: "batch_size", "sampling", "inference_prompts", + "retry_prompt", + "retry_without_image", } def __init__(self): diff --git a/src/caption_flow/workers/caption.py b/src/caption_flow/workers/caption.py index 6460184..9d19791 100644 --- a/src/caption_flow/workers/caption.py +++ b/src/caption_flow/workers/caption.py @@ -515,6 +515,8 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt prompts=vllm_config.get("inference_prompts", ["describe this image"]), output_field="captions", requires=[], + retry_prompt=vllm_config.get("retry_prompt"), + retry_without_image=bool(vllm_config.get("retry_without_image", False)), ) ] @@ -532,6 +534,13 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt max_model_len=stage_cfg.get("max_model_len"), dtype=stage_cfg.get("dtype"), gpu_memory_utilization=stage_cfg.get("gpu_memory_utilization"), + retry_prompt=stage_cfg.get("retry_prompt", vllm_config.get("retry_prompt")), + retry_without_image=bool( + stage_cfg.get( + "retry_without_image", + vllm_config.get("retry_without_image", False), + ) + ), ) stages.append(stage) diff --git a/src/caption_flow/workers/openai_compatible.py b/src/caption_flow/workers/openai_compatible.py index 7b41d2c..ff92f2f 100644 --- a/src/caption_flow/workers/openai_compatible.py +++ b/src/caption_flow/workers/openai_compatible.py @@ -565,6 +565,31 @@ class OpenAICompatibleWorker(CaptionWorker): "presence_penalty", "seed", } + _REFUSAL_MARKERS = ( + "i'm sorry", + "i’m sorry", + "i cannot", + "i can't", + "i can’t", + "unable to provide", + "unable to describe", + "cannot provide", + "can't provide", + "can’t provide", + "cannot assist", + "can't assist", + "can’t assist", + ) + _POLICY_ERROR_MARKERS = ( + "contentfilter", + "content filter", + "content policy", + "content_policy", + "moderation", + "policy violation", + "safety system", + "unsafe or sensitive", + ) def __init__(self, config: Dict[str, Any]): super().__init__(config) @@ -593,6 +618,8 @@ def __init__(self, config: Dict[str, Any]): self.image_detail = raw_config.get("image_detail", "auto") self.image_format = str(raw_config.get("image_format", "jpeg")).lower() self.image_quality = int(raw_config.get("image_quality", 90)) + configured_dimension = int(raw_config.get("max_image_dimension", 0) or 0) + self.max_image_dimension = configured_dimension if configured_dimension > 0 else None async def _pre_start(self): """Fetch shared stage settings, then start the API processing thread.""" @@ -671,13 +698,7 @@ def _process_batch_multi_stage( # noqa: C901 sampling = self._sampling_for_stage(stage) for item in active_batch: - context = dict(item.metadata) - for previous_name, result in item.stage_results.items(): - for index, output in enumerate(result.outputs): - context[f"{previous_name}_output_{index}"] = output - context[result.output_field] = ( - result.outputs[0] if len(result.outputs) == 1 else result.outputs - ) + context = self._stage_context(item) for prompt in PromptTemplateManager(stage.prompts).format_all(context): requests.append( @@ -692,8 +713,10 @@ def _process_batch_multi_stage( # noqa: C901 ) owners.append(item) + api_started = time.monotonic() responses = self.api_loop.run_until_complete(self.endpoint_pool.run_many(requests)) outputs_by_item: Dict[int, List[str]] = defaultdict(list) + retryable_item_ids = set() for owner, response in zip(owners, responses, strict=True): if isinstance(response, Exception): logger.error( @@ -702,10 +725,76 @@ def _process_batch_multi_stage( # noqa: C901 stage_name, response, ) + if self._is_semantic_caption_failure(response): + retryable_item_ids.add(id(owner)) continue cleaned = self._clean_output(response) - if cleaned: + if cleaned and not self._is_refusal_text(cleaned): outputs_by_item[id(owner)].append(cleaned) + else: + retryable_item_ids.add(id(owner)) + + retry_items = [ + item + for item in active_batch + if id(item) in retryable_item_ids + and not outputs_by_item.get(id(item)) + and stage.retry_prompt + ] + if retry_items: + retry_requests = [] + for item in retry_items: + context = self._stage_context(item) + retry_prompt = PromptTemplateManager([stage.retry_prompt]).format_all(context)[ + 0 + ] + retry_requests.append( + ChatRequest( + prompt=retry_prompt, + image_data_url=( + None if stage.retry_without_image else image_urls.get(id(item)) + ), + requested_model=stage.model, + parameters=sampling, + system_prompt=self.system_prompt, + image_detail=self.image_detail, + ) + ) + + logger.info( + "Retrying %d empty or refused caption(s) in stage %s%s", + len(retry_items), + stage_name, + " without images" if stage.retry_without_image else "", + ) + retry_responses = self.api_loop.run_until_complete( + self.endpoint_pool.run_many(retry_requests) + ) + for item, response in zip(retry_items, retry_responses, strict=True): + if isinstance(response, Exception): + logger.error( + "Fallback API request failed for item %s in stage %s: %s", + item.item_key, + stage_name, + response, + ) + continue + cleaned = self._clean_output(response) + if cleaned and not self._is_refusal_text(cleaned): + outputs_by_item[id(item)].append(cleaned) + else: + logger.error( + "Fallback returned no usable output for %s in stage %s", + item.item_key, + stage_name, + ) + + logger.info( + "API stage %s completed %d request(s) in %.3f seconds", + stage_name, + len(requests) + len(retry_items), + time.monotonic() - api_started, + ) next_batch = [] for item in active_batch: @@ -733,6 +822,32 @@ def _process_batch_multi_stage( # noqa: C901 self.items_processed += 1 return results + @classmethod + def _is_refusal_text(cls, text: str) -> bool: + normalized = text.strip().lower() + return any(marker in normalized for marker in cls._REFUSAL_MARKERS) + + @classmethod + def _is_semantic_caption_failure(cls, error: Exception) -> bool: + if isinstance(error, EndpointRequestError): + message = error.message.lower() + return any(marker in message for marker in cls._POLICY_ERROR_MARKERS) + if isinstance(error, ValueError): + message = str(error).lower() + return "message content" in message + return False + + @staticmethod + def _stage_context(item: ProcessingItem) -> Dict[str, Any]: + context = dict(item.metadata) + for previous_name, result in item.stage_results.items(): + for index, output in enumerate(result.outputs): + context[f"{previous_name}_output_{index}"] = output + context[result.output_field] = ( + result.outputs[0] if len(result.outputs) == 1 else result.outputs + ) + return context + def _sampling_for_stage(self, stage: ProcessingStage) -> Dict[str, Any]: sampling = dict(self.vllm_config.get("sampling", {})) if stage.sampling: @@ -743,44 +858,60 @@ def _sampling_for_stage(self, stage: ProcessingStage) -> Dict[str, Any]: if key in self._SUPPORTED_SAMPLING_KEYS and value is not None } - def _image_data_url(self, item: ProcessingItem) -> str: - image_format = ( + def _image_data_url(self, item: ProcessingItem) -> str: # noqa: C901 + source_format = ( (item.image.format if item.image else None) or item.metadata.get("image_format") or item.metadata.get("_image_format") ) - if not image_format and item.image_data: + if not source_format and item.image_data: try: with Image.open(io.BytesIO(item.image_data)) as source: - image_format = source.format + source_format = source.format except (OSError, ValueError): pass - image_format = image_format or self.image_format - image_format = str(image_format).lower() - mime_format = "jpeg" if image_format in {"jpg", "jpeg"} else image_format - supported_passthrough = mime_format in {"jpeg", "png", "gif", "webp"} - if item.image_data and supported_passthrough: + target_format = "jpeg" if self.image_format in {"jpg", "jpeg"} else self.image_format + if target_format not in {"jpeg", "png", "gif", "webp"}: + target_format = "jpeg" + normalized_source = str(source_format or "").lower() + if normalized_source in {"jpg", "jpeg"}: + normalized_source = "jpeg" + + image = item.image + needs_resize = False + if self.max_image_dimension: + if image is not None: + needs_resize = max(image.size) > self.max_image_dimension + elif item.image_data: + with Image.open(io.BytesIO(item.image_data)) as source: + needs_resize = max(source.size) > self.max_image_dimension + + if item.image_data and normalized_source == target_format and not needs_resize: data = item.image_data - elif item.image is not None: + elif image is not None or item.image_data: output = io.BytesIO() - save_format = ( - "JPEG" if self.image_format in {"jpg", "jpeg"} else self.image_format.upper() - ) - image = item.image + save_format = "JPEG" if target_format == "jpeg" else target_format.upper() + if image is None: + image = Image.open(io.BytesIO(item.image_data)) + if self.max_image_dimension and max(image.size) > self.max_image_dimension: + image = image.copy() + image.thumbnail( + (self.max_image_dimension, self.max_image_dimension), + Image.Resampling.LANCZOS, + ) save_kwargs: Dict[str, Any] = {} if save_format == "JPEG": if image.mode not in {"RGB", "L"}: image = image.convert("RGB") save_kwargs["quality"] = self.image_quality - mime_format = "jpeg" image.save(output, format=save_format, **save_kwargs) data = output.getvalue() else: raise ValueError(f"Item {item.item_key} has no image data") encoded = base64.b64encode(data).decode("ascii") - return f"data:image/{mime_format};base64,{encoded}" + return f"data:image/{target_format};base64,{encoded}" def _get_heartbeat_data(self) -> Dict[str, Any]: data = super()._get_heartbeat_data() diff --git a/tests/test_openai_compatible_worker.py b/tests/test_openai_compatible_worker.py index f2ada01..1a91e9e 100644 --- a/tests/test_openai_compatible_worker.py +++ b/tests/test_openai_compatible_worker.py @@ -1,6 +1,7 @@ """Tests for the BYOK OpenAI-compatible caption worker.""" import asyncio +import base64 import io import os from unittest.mock import AsyncMock, patch @@ -421,7 +422,11 @@ def test_worker_image_encoding_handles_passthrough_conversion_and_missing_data() image_data=png_buffer.getvalue(), metadata={}, ) + assert worker._image_data_url(passthrough).startswith("data:image/jpeg;base64,") + + worker.image_format = "png" assert worker._image_data_url(passthrough).startswith("data:image/png;base64,") + worker.image_format = "jpeg" converted = ProcessingItem( unit_id="unit", @@ -440,6 +445,54 @@ def test_worker_image_encoding_handles_passthrough_conversion_and_missing_data() worker._image_data_url(converted) +def test_worker_resizes_large_images_before_encoding(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "max_image_dimension": 1024, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + jpeg_buffer = io.BytesIO() + Image.new("RGB", (2400, 1600), color="green").save(jpeg_buffer, format="JPEG") + item = ProcessingItem( + unit_id="unit", + job_id="job", + chunk_id="chunk", + item_key="large.jpg", + item_index=0, + image=None, + image_data=jpeg_buffer.getvalue(), + metadata={}, + ) + + encoded = worker._image_data_url(item) + encoded_bytes = base64.b64decode(encoded.split(",", 1)[1]) + with Image.open(io.BytesIO(encoded_bytes)) as resized: + assert resized.size == (1024, 683) + + +def test_worker_disables_image_resize_for_non_positive_dimension(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "max_image_dimension": -1, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + assert worker.max_image_dimension is None + + def test_worker_validates_config_and_applies_shared_updates(): base = {"server": "ws://localhost:8765", "token": "orchestrator-token"} with pytest.raises(ValueError, match="must be a mapping"): @@ -523,6 +576,115 @@ def test_worker_runs_shared_caption_stage_through_endpoint_pool(): assert request.image_data_url.startswith("data:image/jpeg;base64,") +def test_worker_retries_policy_rejection_with_metadata_and_without_image(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "batch_image_processing": False, + "openai_compatible": { + "endpoints": [ + { + "base_url": "https://provider.example/v1", + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "provider-vision-model", + } + ] + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + worker.vllm_config = { + "model": "shared-model", + "inference_prompts": ["Describe this image"], + "retry_prompt": "Rewrite neutrally: {column:captions}", + "retry_without_image": True, + } + worker.stages = worker._parse_stages_config(worker.vllm_config) + worker.stage_order = worker._topological_sort_stages(worker.stages) + policy_error = EndpointRequestError( + "provider", + 400, + '{"contentFilter": [{"level": 2}], "error": "unsafe or sensitive content"}', + {}, + ) + worker.endpoint_pool.run_many = AsyncMock( + side_effect=[[policy_error], ["A woman sits among fallen leaves in a forest."]] + ) + worker.api_loop = asyncio.new_event_loop() + item = ProcessingItem( + unit_id="unit", + job_id="shard:chunk:0:idx:6", + chunk_id="chunk", + item_key="image.png", + item_index=6, + image=Image.new("RGB", (2, 2), color="red"), + image_data=b"", + metadata={"captions": "A woman sits in a forest."}, + ) + + try: + results = worker._process_batch_multi_stage([item]) + finally: + worker.api_loop.close() + + assert results == [(item, {"captions": ["A woman sits among fallen leaves in a forest."]})] + assert worker.endpoint_pool.run_many.await_count == 2 + first_request = worker.endpoint_pool.run_many.await_args_list[0].args[0][0] + retry_request = worker.endpoint_pool.run_many.await_args_list[1].args[0][0] + assert first_request.image_data_url.startswith("data:image/jpeg;base64,") + assert retry_request.prompt == "Rewrite neutrally: A woman sits in a forest." + assert retry_request.image_data_url is None + + +def test_worker_retries_empty_refusal_but_not_unrelated_request_error(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + worker.vllm_config = { + "model": "vision", + "inference_prompts": ["Describe"], + "retry_prompt": "Try a neutral description", + } + worker.stages = worker._parse_stages_config(worker.vllm_config) + worker.stage_order = worker._topological_sort_stages(worker.stages) + worker.endpoint_pool.run_many = AsyncMock( + side_effect=[["I'm sorry, I cannot describe this image."], ["A neutral caption."]] + ) + worker.api_loop = asyncio.new_event_loop() + item = ProcessingItem( + unit_id="unit", + job_id="job", + chunk_id="chunk", + item_key="image.png", + item_index=0, + image=Image.new("RGB", (2, 2), color="red"), + image_data=b"", + metadata={}, + ) + + try: + results = worker._process_batch_multi_stage([item]) + finally: + worker.api_loop.close() + + assert results == [(item, {"captions": ["A neutral caption."]})] + retry_request = worker.endpoint_pool.run_many.await_args_list[1].args[0][0] + assert retry_request.image_data_url.startswith("data:image/jpeg;base64,") + assert worker._is_semantic_caption_failure(ValueError("returned no message content")) + assert not worker._is_semantic_caption_failure( + EndpointRequestError("provider", 400, "invalid request", {}) + ) + + def test_cli_selects_openai_compatible_worker(): from caption_flow.cli import main diff --git a/tests/test_vllm_config.py b/tests/test_vllm_config.py index 227cfae..edaea01 100644 --- a/tests/test_vllm_config.py +++ b/tests/test_vllm_config.py @@ -55,6 +55,8 @@ def test_class_constants(self): assert "tensor_parallel_size" in VLLMConfigManager.RELOAD_REQUIRED_FIELDS assert "batch_size" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS assert "sampling" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS + assert "retry_prompt" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS + assert "retry_without_image" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS class TestVLLMConfigManagerAnalyzeChange: diff --git a/tests/test_webdataset_ranges.py b/tests/test_webdataset_ranges.py index e8521c9..ef1671b 100644 --- a/tests/test_webdataset_ranges.py +++ b/tests/test_webdataset_ranges.py @@ -163,6 +163,7 @@ def test_restore_state_with_incomplete_chunks(self, orchestrator_processor, mock # Should have the unprocessed ranges as absolute indices expected_ranges = [(10, 20), (50, 99)] assert unit.data["unprocessed_ranges"] == expected_ranges + assert unit.unit_size == 61 def test_restore_state_skips_completed_chunks(self, orchestrator_processor, mock_storage): """Test that completed chunks are skipped during restoration.""" @@ -194,6 +195,63 @@ def test_restore_state_skips_completed_chunks(self, orchestrator_processor, mock assert len(orchestrator_processor.work_units) == 0 assert len(orchestrator_processor.pending_units) == 0 + def test_restore_state_clamps_stale_final_chunk(self, orchestrator_processor, mock_storage): + """Restore only missing samples that still exist in the current shard.""" + mock_storage.get_all_processed_job_ids.return_value = set() + processor = orchestrator_processor + processor.chunk_tracker.chunks.clear() + processor.work_units.clear() + processor.pending_units.clear() + processor.shard_info_cache[0] = { + "name": "shard_0", + "path": "shard_0.tar", + "num_samples": 75, + "num_files": 96, + } + + chunk_id = "shard_0:chunk:2" + processor.chunk_tracker.add_chunk(chunk_id, "shard_0", "shard_0.tar", 64, 32) + processor.chunk_tracker.mark_items_processed(chunk_id, 64, 68) + processor.chunk_tracker.mark_items_processed(chunk_id, 70, 74) + + processor._restore_state(mock_storage) + + state = processor.chunk_tracker.chunks[chunk_id] + assert state.chunk_size == 11 + assert state.processed_ranges == [(0, 4), (6, 10)] + assert state.get_unprocessed_ranges() == [(5, 5)] + assert processor.work_units[chunk_id].unit_size == 1 + assert processor.work_units[chunk_id].data["chunk_size"] == 11 + assert processor.work_units[chunk_id].data["unprocessed_ranges"] == [(69, 69)] + + def test_restore_state_completes_stale_out_of_bounds_tail( + self, orchestrator_processor, mock_storage + ): + """A nonexistent tail must not keep an otherwise complete chunk pending.""" + mock_storage.get_all_processed_job_ids.return_value = set() + processor = orchestrator_processor + processor.chunk_tracker.chunks.clear() + processor.work_units.clear() + processor.pending_units.clear() + processor.shard_info_cache[0] = { + "name": "shard_0", + "path": "shard_0.tar", + "num_samples": 75, + "num_files": 96, + } + + chunk_id = "shard_0:chunk:2" + processor.chunk_tracker.add_chunk(chunk_id, "shard_0", "shard_0.tar", 64, 32) + processor.chunk_tracker.mark_items_processed(chunk_id, 64, 74) + + processor._restore_state(mock_storage) + + state = processor.chunk_tracker.chunks[chunk_id] + assert state.chunk_size == 11 + assert state.status == "completed" + assert chunk_id not in processor.work_units + assert chunk_id not in processor.pending_units + def test_work_unit_creation_basic(self, orchestrator_processor): """Test basic work unit creation logic by directly testing unit creation.""" # Clear any existing state @@ -277,6 +335,7 @@ def test_work_unit_assignment_updates_ranges(self, orchestrator_processor, mock_ # Should have updated unprocessed ranges to absolute indices expected_ranges = [(5, 15), (30, 40)] # The gaps we left assert assigned_unit.data["unprocessed_ranges"] == expected_ranges + assert assigned_unit.unit_size == 22 def test_work_unit_assignment_skips_completed_chunks(self, orchestrator_processor): """Test that work unit assignment skips chunks with no unprocessed ranges.""" diff --git a/tests/test_worker_caption.py b/tests/test_worker_caption.py index f82b628..16dec60 100644 --- a/tests/test_worker_caption.py +++ b/tests/test_worker_caption.py @@ -255,6 +255,8 @@ def test_parse_stages_backward_compatibility(self, caption_worker): old_config = { "model": "single-model", "inference_prompts": ["Test prompt"], + "retry_prompt": "Fallback prompt", + "retry_without_image": True, } stages = caption_worker._parse_stages_config(old_config) @@ -263,6 +265,8 @@ def test_parse_stages_backward_compatibility(self, caption_worker): assert stages[0].name == "default" assert stages[0].model == "single-model" assert stages[0].output_field == "captions" + assert stages[0].retry_prompt == "Fallback prompt" + assert stages[0].retry_without_image is True def test_topological_sort_stages(self, caption_worker): """Test dependency sorting of stages.""" @@ -950,9 +954,9 @@ def test_unit_incomplete_reporting_to_orchestrator(self, caption_worker, mock_vl work_failed_call = sent_data break - assert work_failed_call is not None, ( - "work_failed message should have been sent for incomplete unit" - ) + assert ( + work_failed_call is not None + ), "work_failed message should have been sent for incomplete unit" assert work_failed_call["unit_id"] == "unit1" assert "Processing incomplete" in work_failed_call["error"] assert "1/3 items processed" in work_failed_call["error"] From 65500b10466b73d92ea092033ea56431bd77018c Mon Sep 17 00:00:00 2001 From: bghira Date: Sun, 6 Sep 2026 22:06:16 -0600 Subject: [PATCH 2/3] perf: batch image preprocessing in Rust --- recipes/GLM_5.3_FLASH.md | 14 ++++ setup.py | 2 +- src/caption_flow/processors/webdataset.py | 29 ++----- src/caption_flow/utils/image_processor.py | 76 ++++++++++++++++- src/caption_flow/workers/caption.py | 2 +- src/caption_flow/workers/openai_compatible.py | 67 +++++++++++++-- tests/test_image_processor.py | 84 +++++++++++++++++++ tests/test_openai_compatible_worker.py | 78 +++++++++++++++++ tests/test_webdataset_ranges.py | 66 ++++++++++----- 9 files changed, 361 insertions(+), 57 deletions(-) create mode 100644 tests/test_image_processor.py diff --git a/recipes/GLM_5.3_FLASH.md b/recipes/GLM_5.3_FLASH.md index 78b0293..e8eb812 100644 --- a/recipes/GLM_5.3_FLASH.md +++ b/recipes/GLM_5.3_FLASH.md @@ -144,6 +144,20 @@ The worker token must match one of the orchestrator's `auth.worker_tokens`. When vLLM and the CaptionFlow worker run in different containers, replace `127.0.0.1` with a private, reachable vLLM address. +For WebDataset runs, keep encoded samples intact until the worker has assembled +a complete request batch: + +```yaml +orchestrator: + dataset: + decode_images: false +``` + +The API worker will then use `trainingsample`'s Rust-backed fused batch +decode/resize pipeline. This avoids decoding each sample once in the dataset +reader and again while normalizing it for the endpoint. Pillow remains a +per-image fallback for malformed or unsupported inputs. + ## Configure caption generation The orchestrator owns the prompt and sampling configuration. This is the diff --git a/setup.py b/setup.py index 140042a..b34a266 100644 --- a/setup.py +++ b/setup.py @@ -34,6 +34,7 @@ def get_version() -> str: "certbot>=5.0.0,<6.0.0", "numpy>=2.2.0,<3.0.0", "Pillow>=11.3.0,<13.0.0", + "trainingsample>=0.3.2,<0.4.0", "pandas>=3.0.0,<4.0.0", "datasets>=5.0.0,<6.0.0", "boto3>=1.43.0,<2.0.0", @@ -44,7 +45,6 @@ def get_version() -> str: "fastapi>=0.133.0,<0.137.0", "uvicorn>=0.35.0,<1.0.0", "huggingface-hub>=1.5.0,<2.0.0", - "opencv-python-headless>=4.13.0,<6.0.0", "psutil>=7.0.0,<8.0.0", "requests>=2.32.0,<3.0.0", "tqdm>=4.67.0,<5.0.0", diff --git a/src/caption_flow/processors/webdataset.py b/src/caption_flow/processors/webdataset.py index f84187f..cdf78f6 100644 --- a/src/caption_flow/processors/webdataset.py +++ b/src/caption_flow/processors/webdataset.py @@ -1,7 +1,6 @@ """WebDataset processor implementation using webshart TarDataLoader.""" import gc -import io import logging import os import threading @@ -11,8 +10,6 @@ from types import SimpleNamespace from typing import Any, Deque, Dict, Iterator, List, Optional, Set -import cv2 -import numpy as np import requests import webshart from PIL import Image @@ -21,6 +18,7 @@ from caption_flow.storage import StorageManager from ..utils import ChunkTracker +from ..utils.image_processor import ImageProcessor from .base import OrchestratorProcessor, ProcessorConfig, WorkerProcessor, WorkResult, WorkUnit logger = logging.getLogger(__name__) @@ -625,6 +623,7 @@ def __init__(self): self.remote_range_reads = False self.remote_range_timeout = 120.0 self.remote_range_retries = 3 + self.decode_images = True self._remote_shard_layouts: Dict[int, Dict[str, Any]] = {} self.http_session: Optional[requests.Session] = None @@ -639,6 +638,7 @@ def initialize(self, config: ProcessorConfig) -> None: self.remote_range_reads = bool(dataset_cfg.get("remote_range_reads", False)) self.remote_range_timeout = float(dataset_cfg.get("remote_range_timeout", 120)) self.remote_range_retries = max(1, int(dataset_cfg.get("remote_range_retries", 3))) + self.decode_images = bool(dataset_cfg.get("decode_images", True)) split_worker_cache = dataset_cfg.get( "split_worker_cache", True ) # multiple workers get their own cache by default @@ -835,28 +835,11 @@ def process_unit(self, unit: WorkUnit, context: Dict[str, Any]) -> Iterator[Dict # Decode image image = None - if entry.data: + if entry.data and self.decode_images: try: - # Use cv2 to decode from memory - nparr = np.frombuffer(entry.data, np.uint8) - img_np = cv2.imdecode(nparr, cv2.IMREAD_COLOR) - - if img_np is not None: - # Convert from BGR (OpenCV default) to RGB (PIL default) - img_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB) - image = Image.fromarray(img_rgb) - else: - logger.warning(f"cv2.imdecode failed for {entry.path}") - - except ImportError: - logger.warning( - "cv2 or numpy not installed, falling back to PIL" - ) - image = Image.open(io.BytesIO(entry.data)) + image = ImageProcessor.decode_image_data(entry.data) except Exception as img_e: - logger.error( - f"Error decoding image {entry.path} with cv2: {img_e}" - ) + logger.error(f"Error decoding image {entry.path}: {img_e}") # Generate job ID using JobId class job_id = JobId.from_values( diff --git a/src/caption_flow/utils/image_processor.py b/src/caption_flow/utils/image_processor.py index 661c58e..e42d801 100644 --- a/src/caption_flow/utils/image_processor.py +++ b/src/caption_flow/utils/image_processor.py @@ -4,7 +4,10 @@ import os from concurrent.futures import ProcessPoolExecutor from io import BytesIO +from typing import Iterable, Sequence +import numpy as np +import trainingsample as tsr from PIL import Image from ..models import ProcessingItem @@ -43,13 +46,84 @@ def prepare_for_inference(item: ProcessingItem) -> Image.Image: return image item.image = None - image = Image.open(BytesIO(item.image_data)) + image = ImageProcessor.decode_image_data(item.image_data) item.image_data = b"" item.metadata["image_format"] = image.format or "unknown" item.metadata["image_width"], item.metadata["image_height"] = image.size return image + @staticmethod + def decode_image_data(image_data: bytes) -> Image.Image: + """Decode encoded image bytes through the Rust-backed fast path.""" + try: + decoded = tsr.imdecode_py(image_data, 1) + return Image.fromarray(np.asarray(decoded)) + except Exception: + logger.debug("trainingsample decode failed; falling back to Pillow", exc_info=True) + with Image.open(BytesIO(image_data)) as source: + return source.convert("RGB") + + @staticmethod + def constrained_size(size: tuple[int, int], max_dimension: int) -> tuple[int, int]: + """Preserve aspect ratio while bounding the longest image edge.""" + width, height = size + longest = max(width, height) + if longest <= max_dimension: + return width, height + scale = max_dimension / longest + return max(1, round(width * scale)), max(1, round(height * scale)) + + @staticmethod + def preprocess_encoded_batch( + image_buffers: Sequence[bytes], target_sizes: Sequence[tuple[int, int]] + ) -> list[Image.Image]: + """Decode and resize an encoded-image batch in the Rust extension.""" + if not image_buffers: + return [] + if len(image_buffers) != len(target_sizes): + raise ValueError("image_buffers and target_sizes must have the same length") + + processor = tsr.PyBatchProcessor.with_config(True, min(32, len(image_buffers))) + arrays = processor.batch_preprocess_pipeline( + list(image_buffers), + list(target_sizes), + None, + 1, + tsr.INTER_LANCZOS4, + ) + if len(arrays) != len(image_buffers): + raise RuntimeError( + f"trainingsample returned {len(arrays)} images for {len(image_buffers)} inputs" + ) + return [Image.fromarray(np.asarray(array)) for array in arrays] + + @staticmethod + def resize_images( + images: Iterable[Image.Image], target_sizes: Sequence[tuple[int, int]] + ) -> list[Image.Image]: + """Resize PIL images as one Rust-backed batch, with a Pillow fallback.""" + image_list = list(images) + if len(image_list) != len(target_sizes): + raise ValueError("images and target_sizes must have the same length") + if not image_list: + return [] + + try: + arrays = [np.asarray(image.convert("RGB")) for image in image_list] + resized = tsr.batch_resize_images(arrays, list(target_sizes)) + if len(resized) != len(image_list): + raise RuntimeError( + f"trainingsample returned {len(resized)} images for {len(image_list)} inputs" + ) + return [Image.fromarray(np.asarray(array)) for array in resized] + except Exception: + logger.debug("trainingsample resize failed; falling back to Pillow", exc_info=True) + return [ + image.resize(size, Image.Resampling.LANCZOS) + for image, size in zip(image_list, target_sizes, strict=True) + ] + def shutdown(self): """Shutdown the executor.""" self.executor.shutdown(wait=True) diff --git a/src/caption_flow/workers/caption.py b/src/caption_flow/workers/caption.py index 9d19791..d027ff7 100644 --- a/src/caption_flow/workers/caption.py +++ b/src/caption_flow/workers/caption.py @@ -1042,7 +1042,7 @@ def _resize_image_for_tokens( new_height = int(item.image.height * target_ratio) # Resize image - resized_image = item.image.resize((new_width, new_height), Image.Resampling.LANCZOS) + resized_image = ImageProcessor.resize_images([item.image], [(new_width, new_height)])[0] # Create new item with resized image new_item = ProcessingItem( diff --git a/src/caption_flow/workers/openai_compatible.py b/src/caption_flow/workers/openai_compatible.py index ff92f2f..83b6b08 100644 --- a/src/caption_flow/workers/openai_compatible.py +++ b/src/caption_flow/workers/openai_compatible.py @@ -25,6 +25,7 @@ from .. import __version__ from ..models import ProcessingStage, StageResult +from ..utils.image_processor import ImageProcessor from ..utils.prompt_template import PromptTemplateManager from .caption import CaptionWorker, ProcessingItem @@ -688,8 +689,7 @@ def _process_batch_multi_stage( # noqa: C901 active_batch = list(batch) image_urls: Dict[int, Optional[str]] = {} if self.include_image: - for item in active_batch: - image_urls[id(item)] = self._image_data_url(item) + image_urls = self._image_data_urls(active_batch) for stage_name in self.stage_order: stage = next(stage for stage in self.stages if stage.name == stage_name) @@ -893,13 +893,10 @@ def _image_data_url(self, item: ProcessingItem) -> str: # noqa: C901 output = io.BytesIO() save_format = "JPEG" if target_format == "jpeg" else target_format.upper() if image is None: - image = Image.open(io.BytesIO(item.image_data)) + image = ImageProcessor.decode_image_data(item.image_data) if self.max_image_dimension and max(image.size) > self.max_image_dimension: - image = image.copy() - image.thumbnail( - (self.max_image_dimension, self.max_image_dimension), - Image.Resampling.LANCZOS, - ) + target_size = ImageProcessor.constrained_size(image.size, self.max_image_dimension) + image = ImageProcessor.resize_images([image], [target_size])[0] save_kwargs: Dict[str, Any] = {} if save_format == "JPEG": if image.mode not in {"RGB", "L"}: @@ -913,6 +910,60 @@ def _image_data_url(self, item: ProcessingItem) -> str: # noqa: C901 encoded = base64.b64encode(data).decode("ascii") return f"data:image/{target_format};base64,{encoded}" + def _image_data_urls(self, items: Iterable[ProcessingItem]) -> Dict[int, Optional[str]]: + """Encode a request batch, fusing byte decode and resize when possible.""" + item_list = list(items) + results: Dict[int, Optional[str]] = {} + target_format = "jpeg" if self.image_format in {"jpg", "jpeg"} else self.image_format + fast_items: List[ProcessingItem] = [] + fast_sizes: List[tuple[int, int]] = [] + + if target_format == "jpeg": + for item in item_list: + if item.image is not None or not item.image_data: + continue + try: + with Image.open(io.BytesIO(item.image_data)) as source: + source_format = str(source.format or "").lower() + source_size = source.size + normalized_source = ( + "jpeg" if source_format in {"jpg", "jpeg"} else source_format + ) + target_size = source_size + if self.max_image_dimension: + target_size = ImageProcessor.constrained_size( + source_size, self.max_image_dimension + ) + if normalized_source == "jpeg" and target_size == source_size: + encoded = base64.b64encode(item.image_data).decode("ascii") + results[id(item)] = f"data:image/jpeg;base64,{encoded}" + else: + fast_items.append(item) + fast_sizes.append(target_size) + except (OSError, ValueError): + continue + + if fast_items: + try: + images = ImageProcessor.preprocess_encoded_batch( + [item.image_data for item in fast_items], fast_sizes + ) + for item, image in zip(fast_items, images, strict=True): + output = io.BytesIO() + image.save(output, format="JPEG", quality=self.image_quality) + encoded = base64.b64encode(output.getvalue()).decode("ascii") + results[id(item)] = f"data:image/jpeg;base64,{encoded}" + except Exception: + logger.warning( + "Rust batch image preprocessing failed; falling back to individual images", + exc_info=True, + ) + + for item in item_list: + if id(item) not in results: + results[id(item)] = self._image_data_url(item) + return results + def _get_heartbeat_data(self) -> Dict[str, Any]: data = super()._get_heartbeat_data() data.update( diff --git a/tests/test_image_processor.py b/tests/test_image_processor.py new file mode 100644 index 0000000..0953590 --- /dev/null +++ b/tests/test_image_processor.py @@ -0,0 +1,84 @@ +"""Tests for Rust-backed image preprocessing.""" + +import io +from unittest.mock import Mock, patch + +import numpy as np +import pytest +from PIL import Image + +from caption_flow.utils.image_processor import ImageProcessor + + +def encoded_image(size=(20, 10), image_format="PNG") -> bytes: + buffer = io.BytesIO() + Image.new("RGB", size, color="red").save(buffer, format=image_format) + return buffer.getvalue() + + +def test_decode_image_data_uses_trainingsample_and_falls_back_to_pillow(): + data = encoded_image() + decoded = ImageProcessor.decode_image_data(data) + assert decoded.mode == "RGB" + assert decoded.size == (20, 10) + + with patch( + "caption_flow.utils.image_processor.tsr.imdecode_py", + side_effect=RuntimeError("unsupported"), + ): + fallback = ImageProcessor.decode_image_data(data) + assert fallback.mode == "RGB" + assert fallback.size == (20, 10) + + +def test_constrained_size_preserves_aspect_ratio(): + assert ImageProcessor.constrained_size((80, 40), 100) == (80, 40) + assert ImageProcessor.constrained_size((200, 100), 100) == (100, 50) + assert ImageProcessor.constrained_size((1, 1000), 1) == (1, 1) + + +def test_preprocess_encoded_batch_validates_and_converts_results(): + assert ImageProcessor.preprocess_encoded_batch([], []) == [] + with pytest.raises(ValueError, match="same length"): + ImageProcessor.preprocess_encoded_batch([b"one"], []) + + processor = Mock() + processor.batch_preprocess_pipeline.return_value = [np.zeros((5, 10, 3), dtype=np.uint8)] + with patch( + "caption_flow.utils.image_processor.tsr.PyBatchProcessor.with_config", + return_value=processor, + ) as factory: + [image] = ImageProcessor.preprocess_encoded_batch([b"one"], [(10, 5)]) + factory.assert_called_once_with(True, 1) + assert image.size == (10, 5) + + processor.batch_preprocess_pipeline.return_value = [] + with ( + patch( + "caption_flow.utils.image_processor.tsr.PyBatchProcessor.with_config", + return_value=processor, + ), + pytest.raises(RuntimeError, match="returned 0 images"), + ): + ImageProcessor.preprocess_encoded_batch([b"one"], [(10, 5)]) + + +def test_resize_images_uses_batch_and_has_pillow_fallback(): + image = Image.new("RGB", (20, 10), color="green") + assert ImageProcessor.resize_images([], []) == [] + with pytest.raises(ValueError, match="same length"): + ImageProcessor.resize_images([image], []) + + resized = ImageProcessor.resize_images([image], [(10, 5)]) + assert resized[0].size == (10, 5) + + with patch( + "caption_flow.utils.image_processor.tsr.batch_resize_images", + side_effect=RuntimeError("unsupported"), + ): + fallback = ImageProcessor.resize_images([image], [(8, 4)]) + assert fallback[0].size == (8, 4) + + with patch("caption_flow.utils.image_processor.tsr.batch_resize_images", return_value=[]): + mismatch_fallback = ImageProcessor.resize_images([image], [(6, 3)]) + assert mismatch_fallback[0].size == (6, 3) diff --git a/tests/test_openai_compatible_worker.py b/tests/test_openai_compatible_worker.py index 1a91e9e..bd1f669 100644 --- a/tests/test_openai_compatible_worker.py +++ b/tests/test_openai_compatible_worker.py @@ -11,6 +11,7 @@ from PIL import Image from caption_flow.workers.caption import ProcessingItem +from caption_flow.utils.image_processor import ImageProcessor from caption_flow.workers.openai_compatible import ( AdaptiveEndpointPool, ChatRequest, @@ -493,6 +494,83 @@ def test_worker_disables_image_resize_for_non_positive_dimension(): assert worker.max_image_dimension is None +def test_worker_fuses_encoded_image_batch_preprocessing(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "max_image_dimension": 100, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + items = [] + for index, size in enumerate(((200, 100), (100, 200))): + buffer = io.BytesIO() + Image.new("RGB", size, color="blue").save(buffer, format="PNG") + items.append( + ProcessingItem( + unit_id="unit", + job_id=f"job-{index}", + chunk_id="chunk", + item_key=f"image-{index}.png", + item_index=index, + image=None, + image_data=buffer.getvalue(), + metadata={}, + ) + ) + + with patch.object( + ImageProcessor, + "preprocess_encoded_batch", + wraps=ImageProcessor.preprocess_encoded_batch, + ) as preprocess: + urls = worker._image_data_urls(items) + + preprocess.assert_called_once() + assert preprocess.call_args.args[1] == [(100, 50), (50, 100)] + for item, expected_size in zip(items, ((100, 50), (50, 100)), strict=True): + encoded = base64.b64decode(urls[id(item)].split(",", 1)[1]) + with Image.open(io.BytesIO(encoded)) as image: + assert image.size == expected_size + + +def test_worker_batch_preprocessing_falls_back_per_item(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + buffer = io.BytesIO() + Image.new("RGB", (8, 8), color="blue").save(buffer, format="PNG") + item = ProcessingItem( + unit_id="unit", + job_id="job", + chunk_id="chunk", + item_key="image.png", + item_index=0, + image=None, + image_data=buffer.getvalue(), + metadata={}, + ) + with ( + patch.object(ImageProcessor, "preprocess_encoded_batch", side_effect=RuntimeError), + patch.object(worker, "_image_data_url", return_value="fallback") as fallback, + ): + assert worker._image_data_urls([item]) == {id(item): "fallback"} + fallback.assert_called_once_with(item) + + def test_worker_validates_config_and_applies_shared_updates(): base = {"server": "ws://localhost:8765", "token": "orchestrator-token"} with pytest.raises(ValueError, match="must be a mapping"): diff --git a/tests/test_webdataset_ranges.py b/tests/test_webdataset_ranges.py index ef1671b..b4857dd 100644 --- a/tests/test_webdataset_ranges.py +++ b/tests/test_webdataset_ranges.py @@ -933,20 +933,14 @@ def test_process_unit_real_mode_with_mock_loader(self, worker_processor_real): metadata={"chunk_index": 0}, ) - # Mock the image decoding + # Mock the Rust-backed image decoding test_image = Image.new("RGB", (100, 100), color="red") - with patch("caption_flow.processors.webdataset.cv2.imdecode") as mock_decode: - with patch("caption_flow.processors.webdataset.cv2.cvtColor") as mock_convert: - with patch( - "caption_flow.processors.webdataset.Image.fromarray", - return_value=test_image, - ): - # Mock cv2 processing chain - mock_decode.return_value = "fake_cv2_image" - mock_convert.return_value = "fake_rgb_array" - - results = list(worker_processor_real.process_unit(unit, {})) + with patch( + "caption_flow.processors.webdataset.ImageProcessor.decode_image_data", + return_value=test_image, + ) as mock_decode: + results = list(worker_processor_real.process_unit(unit, {})) assert len(results) == 3 @@ -965,6 +959,7 @@ def test_process_unit_real_mode_with_mock_loader(self, worker_processor_real): assert "json_path" not in result["metadata"] assert result["metadata"]["_filename"] != "wrong.jpg" assert not result["metadata"].get("_mock", False) # Should not have mock flag + assert mock_decode.call_count == 3 # Verify loader was called correctly worker_processor_real.loader.load_sample.assert_any_call(0, 5) @@ -1160,17 +1155,17 @@ def test_process_unit_real_mode_shard_by_name(self, worker_processor_real): test_image = Image.new("RGB", (50, 50)) - with patch( - "caption_flow.processors.webdataset.webshart.next_with_cache_wait", - return_value=mock_entry, + with ( + patch( + "caption_flow.processors.webdataset.webshart.next_with_cache_wait", + return_value=mock_entry, + ), + patch( + "caption_flow.processors.webdataset.ImageProcessor.decode_image_data", + return_value=test_image, + ), ): - with patch("caption_flow.processors.webdataset.Image.open", return_value=test_image): - # Simulate cv2 import error to test PIL fallback - with patch( - "caption_flow.processors.webdataset.cv2.imdecode", - side_effect=ImportError("cv2 not available"), - ): - results = list(worker_processor_real.process_unit(unit, {})) + results = list(worker_processor_real.process_unit(unit, {})) assert len(results) == 1 @@ -1179,10 +1174,35 @@ def test_process_unit_real_mode_shard_by_name(self, worker_processor_real): filename="shard_unknown", cursor_idx=10 ) - # Should have used PIL fallback result = results[0] assert result["image"] == test_image + def test_process_unit_can_defer_image_decode(self, worker_processor_real): + """API workers can preserve encoded bytes for fused batch preprocessing.""" + worker_processor_real.decode_images = False + mock_entry = Mock(data=b"encoded", path="test.jpg", size=7, metadata={}) + worker_processor_real.loader.load_sample = Mock(return_value=mock_entry) + unit = WorkUnit( + unit_id="shard_0:chunk:0", + chunk_id="shard_0:chunk:0", + source_id="shard_0", + unit_size=1, + data={ + "shard_name": "shard_0", + "shard_idx": 0, + "start_index": 0, + "unprocessed_ranges": [(0, 0)], + }, + metadata={"chunk_index": 0}, + ) + + with patch("caption_flow.processors.webdataset.ImageProcessor.decode_image_data") as decode: + [result] = list(worker_processor_real.process_unit(unit, {})) + + decode.assert_not_called() + assert result["image"] is None + assert result["image_data"] == b"encoded" + def test_get_dataset_info_mock_mode(self, worker_processor): """Test dataset info in mock mode.""" info = worker_processor.get_dataset_info() From 631c71a80d8112dce448f02dc2749134c4a6c1f7 Mon Sep 17 00:00:00 2001 From: bghira Date: Sun, 6 Sep 2026 22:23:55 -0600 Subject: [PATCH 3/3] perf: expose WebShart loader tuning --- src/caption_flow/processors/webdataset.py | 6 ++++ tests/test_webdataset_ranges.py | 39 +++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src/caption_flow/processors/webdataset.py b/src/caption_flow/processors/webdataset.py index cdf78f6..b06f1e2 100644 --- a/src/caption_flow/processors/webdataset.py +++ b/src/caption_flow/processors/webdataset.py @@ -624,6 +624,8 @@ def __init__(self): self.remote_range_timeout = 120.0 self.remote_range_retries = 3 self.decode_images = True + self.webshart_parallel_downloads = 4 + self.webshart_chunk_size_mb = 10 self._remote_shard_layouts: Dict[int, Dict[str, Any]] = {} self.http_session: Optional[requests.Session] = None @@ -639,6 +641,8 @@ def initialize(self, config: ProcessorConfig) -> None: self.remote_range_timeout = float(dataset_cfg.get("remote_range_timeout", 120)) self.remote_range_retries = max(1, int(dataset_cfg.get("remote_range_retries", 3))) self.decode_images = bool(dataset_cfg.get("decode_images", True)) + self.webshart_parallel_downloads = max(1, int(cfg.get("webshart_parallel_downloads", 4))) + self.webshart_chunk_size_mb = max(1, int(cfg.get("webshart_chunk_size_mb", 10))) split_worker_cache = dataset_cfg.get( "split_worker_cache", True ) # multiple workers get their own cache by default @@ -663,6 +667,7 @@ def initialize(self, config: ProcessorConfig) -> None: else str(cache_dir / "shard_cache") ), cache_limit_gb=cfg.get("shard_cache_gb", 10.0), + parallel_downloads=self.webshart_parallel_downloads, ) # Create loader @@ -671,6 +676,7 @@ def initialize(self, config: ProcessorConfig) -> None: buffer_size=cfg.get("buffer_size", 10), max_file_size=cfg.get("max_file_size", 100 * 1024 * 1024), load_file_data=True, + chunk_size_mb=self.webshart_chunk_size_mb, ) if self.remote_range_reads: diff --git a/tests/test_webdataset_ranges.py b/tests/test_webdataset_ranges.py index b4857dd..50bde30 100644 --- a/tests/test_webdataset_ranges.py +++ b/tests/test_webdataset_ranges.py @@ -786,6 +786,45 @@ def test_initialization_real_mode(self, worker_processor_real): assert worker_processor_real.dataset is not None assert worker_processor_real.loader is not None + def test_initialization_forwards_webshart_performance_settings(self, worker_config, temp_dir): + config_dict = { + **worker_config.config, + "dataset": { + **worker_config.config["dataset"], + "mock_results": False, + }, + "cache_dir": str(temp_dir / "performance-cache"), + "webshart_parallel_downloads": 8, + "webshart_chunk_size_mb": 32, + } + dataset = Mock() + + with ( + patch( + "caption_flow.processors.webdataset.webshart.discover_dataset", + return_value=dataset, + ), + patch("caption_flow.processors.webdataset.webshart.TarDataLoader") as loader, + ): + processor = WebDatasetWorkerProcessor() + processor.gpu_id = 0 + processor.initialize(ProcessorConfig(processor_type="webdataset", config=config_dict)) + + dataset.enable_shard_cache.assert_called_once_with( + location=str(temp_dir / "performance-cache" / "shard_cache" / "0"), + cache_limit_gb=1.0, + parallel_downloads=8, + ) + loader.assert_called_once_with( + dataset, + buffer_size=10, + max_file_size=100 * 1024 * 1024, + load_file_data=True, + chunk_size_mb=32, + ) + assert processor.webshart_parallel_downloads == 8 + assert processor.webshart_chunk_size_mb == 32 + def test_mock_image_creation(self, worker_processor): """Test mock image creation produces different images.""" img1 = worker_processor._create_mock_image(0)