diff --git a/BUILD.yaml b/BUILD.yaml index 5df92fee7..b59c71052 100644 --- a/BUILD.yaml +++ b/BUILD.yaml @@ -593,6 +593,19 @@ command: bash tests.sh timeout_in_sec: 1800 +# owner: @geoff-counihan +- name: vlm-distillation-catalog-enrichment + dir: templates/vlm-distillation-catalog-enrichment + cluster_env: + image_uri: anyscale/ray-llm:2.55.1-py311-cu128 + compute_config: + AWS: configs/vlm-distillation-catalog-enrichment/aws.yaml + GCP: configs/vlm-distillation-catalog-enrichment/gce.yaml + test: + tests_path: tests/vlm-distillation-catalog-enrichment/ + command: bash tests.sh + timeout_in_sec: 3600 + # owner: @christian-stano - name: ecommerce_multi_model_serving dir: templates/ecommerce_multi_model_serving diff --git a/configs/vlm-distillation-catalog-enrichment/aws.yaml b/configs/vlm-distillation-catalog-enrichment/aws.yaml new file mode 100644 index 000000000..d5f29b093 --- /dev/null +++ b/configs/vlm-distillation-catalog-enrichment/aws.yaml @@ -0,0 +1,15 @@ +# Head node (CPU only) +head_node_type: + name: head + instance_type: m5.2xlarge + resources: + cpu: 0 + gpu: 0 + +# Worker nodes with L4 GPUs +worker_node_types: + - name: gpu-worker-4xL4 + instance_type: g6.12xlarge # 4x L4 GPUs + min_workers: 0 + max_workers: 2 + use_spot: false diff --git a/configs/vlm-distillation-catalog-enrichment/gce.yaml b/configs/vlm-distillation-catalog-enrichment/gce.yaml new file mode 100644 index 000000000..aa5ad479d --- /dev/null +++ b/configs/vlm-distillation-catalog-enrichment/gce.yaml @@ -0,0 +1,15 @@ +# Head node (CPU only) +head_node_type: + name: head + instance_type: n2-standard-8 + resources: + cpu: 0 + gpu: 0 + +# Worker nodes with L4 GPUs +worker_node_types: + - name: gpu-worker-4xL4 + instance_type: g2-standard-48 # 4x L4 GPUs + min_workers: 0 + max_workers: 2 + use_spot: false diff --git a/templates/vlm-distillation-catalog-enrichment/README.ipynb b/templates/vlm-distillation-catalog-enrichment/README.ipynb new file mode 100644 index 000000000..3654430e2 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/README.ipynb @@ -0,0 +1,133 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": "# VLM teacher-student distillation for ecommerce catalogs\n\n\n\"Run\n\n

\n
\n \n
\n\n**⏱️ Time to complete**: 45 min\n\nThis tutorial distills a Qwen2.5-VL-7B teacher into a Qwen2.5-VL-3B student so a product catalog can be enriched and embedded at 3B inference cost without sacrificing teacher-quality structured outputs. The full pipeline runs on Ray and Anyscale across three stages, all on the same 4× L4 GPU node:\n\n1. **Teacher batch labeling** — Run Qwen2.5-VL-7B over a catalog subset with [`ray.data.llm`](https://docs.ray.io/en/latest/data/working-with-llms.html) to produce `{category, attributes, search_tags, description}` JSON per product.\n2. **Distillation SFT** — Fine-tune Qwen2.5-VL-3B on the teacher labels with Ray Train + FSDP + LoRA. Only the language model gets adapters; the vision tower stays frozen.\n3. **Enrichment and embeddings** — Run the LoRA-adapted student to emit catalog JSON *and* SigLIP-2 image and text embeddings in a single streaming Ray Data graph.\n\nRay is particularly powerful for this workload because it:\n- **Schedules CPU and GPU stages in one cluster** so image fetching, VLM inference, and embedding extraction share a single resource pool with no idle hardware between stages.\n- **Streams data between stages** without intermediate disk writes — SigLIP embeddings begin computing as soon as the first VLM-enriched batch is ready.\n- **Hot-swaps LoRA adapters** at inference via vLLM's [multi-LoRA support](https://docs.vllm.ai/en/stable/features/multi_lora.html), so the distilled student loads from the same base model as the un-tuned variant.\n- **Promotes any notebook to a scheduled production job** with `anyscale job submit` — no rewrite, same cluster shape." + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Architecture\n\n```\nStage 1 Stage 2 Stage 3\n───────────── ────────────── ─────────────────\n7B teacher ───► 3B student SFT ───► 3B distilled enrich\nbatch labeling (FSDP + LoRA) + SigLIP embeddings\n (one Ray Data graph)\n\nray.data.llm Ray Train + FSDP ray.data.llm + Ray Data\n ▼ ▼ ▼\nteacher.parquet LoRA adapter enriched_with_embeddings.parquet\n(JSON labels) (~100 MB) (JSON + 1152-dim vectors)\n```\n\nStages 1 and 2 each take 15–30 min on the smoke configuration (N=20 rows); Stage 3 takes ~5 min. Production runs at N=10,000 take roughly 1.5 hours on the same cluster.\n\nThe deep-dive notebooks under [`notebooks/`](notebooks/) walk through each stage cell by cell — useful when you want to swap models, change the catalog, or tune the SFT hyperparameters." + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Set up\n\nThe pipeline pulls model weights and the [Amazon-Reviews-2023](https://huggingface.co/datasets/McAuley-Lab/Amazon-Reviews-2023) dataset from Hugging Face. Make sure `HF_TOKEN` is available in your environment before running the cells below." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\nimport subprocess\nimport sys\n\n# Cache HF downloads on cluster storage so the three stages share weights.\nos.environ.setdefault(\"HF_HOME\", \"/mnt/cluster_storage/hf_cache\")\nos.environ.setdefault(\"HF_HUB_ENABLE_HF_TRANSFER\", \"1\")\n\n# Smoke knobs — overridden by tests.sh during CI.\n# Set these higher for a production run (N_ROWS=10000 is the default in each script).\nN_ROWS = int(os.environ.get(\"N_ROWS\", \"20\"))\nTEACHER_N_ROWS = int(os.environ.get(\"TEACHER_N_ROWS\", str(N_ROWS)))\nCATEGORY = os.environ.get(\"CATEGORY\", \"Electronics\")\n\nos.environ[\"N_ROWS\"] = str(N_ROWS)\nos.environ[\"TEACHER_N_ROWS\"] = str(TEACHER_N_ROWS)\nos.environ[\"CATEGORY\"] = CATEGORY\n\nprint(f\"Running with N_ROWS={N_ROWS} TEACHER_N_ROWS={TEACHER_N_ROWS} CATEGORY={CATEGORY}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install -q -r requirements.txt" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Stage 1 — Teacher batch labeling\n\nQwen2.5-VL-7B reads a product image and merchant-supplied title and returns a four-key JSON object (`category`, `attributes`, `search_tags`, `description`). [`ray.data.llm`](https://docs.ray.io/en/latest/data/working-with-llms.html#multimodal) wraps the vLLM engine — preprocessing, batching, and postprocessing all happen as `map_batches` stages over a Ray Dataset. With `concurrency=4`, one 7B replica runs on each of the four L4 GPUs.\n\nThe output parquet becomes Stage 2's supervised training data.\n\nThe deep dive: [`notebooks/01_teacher_batch_label.ipynb`](notebooks/01_teacher_batch_label.ipynb)." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!python scripts/run_teacher_batch_label.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Inspect a sample row of the teacher output.\nimport pyarrow.parquet as pq\n\nteacher_path = f\"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/teacher_7b_enriched_{TEACHER_N_ROWS}.parquet\"\ntbl = pq.read_table(teacher_path)\nprint(f\"rows: {tbl.num_rows}\")\nprint(f\"schema: {tbl.schema.names}\")\nrow = tbl.slice(0, 1).to_pylist()[0]\nprint(f\"\\nsample title: {row['title'][:100]}\")\nprint(f\"sample teacher output (raw_output):\\n{row['raw_output']}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Stage 2 — Distill the 3B student with FSDP + LoRA\n\nRay Train + FSDP fine-tunes Qwen2.5-VL-3B on the teacher parquet using LoRA adapters. The vision tower stays frozen (standard [LLaVA recipe](https://arxiv.org/abs/2304.08485)); only the language model gets adapters, so trainable parameters drop to ~1% of the model. FSDP `FULL_SHARD` distributes the optimizer state across the four GPUs at bf16 mixed precision.\n\nThe output is a small (~100 MB) adapter directory that drops into Stage 3 as a model swap — same base weights, fine-tuned head.\n\nThe deep dive: [`notebooks/02_distill_student_lora.ipynb`](notebooks/02_distill_student_lora.ipynb)." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!python scripts/run_distill_student_lora.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Inspect the adapter.\nfrom pathlib import Path\n\nadapter_dir = Path(f\"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/qwen25vl_3b_enrichment_lora_{N_ROWS}\")\nprint(f\"adapter dir: {adapter_dir}\")\nfor p in sorted(adapter_dir.rglob(\"*\"))[:10]:\n print(f\" {p.relative_to(adapter_dir)}\")\n\n# Point Stage 3 at the freshly trained adapter.\nos.environ[\"QWEN_LORA_ADAPTER_DIR\"] = str(adapter_dir)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Stage 3 — Enrichment and embeddings in one Ray Data graph\n\nThe third stage chains together the LoRA-adapted 3B student (via `ray.data.llm`) and the SigLIP-2 dual-tower encoder (via a Ray Data actor pool) into a single streaming pipeline. No intermediate disk writes between the VLM and the embedding stages — Ray Data hands batches directly from one `map_batches` stage to the next, so the SigLIP encoders start working as soon as the first VLM-enriched batch is ready.\n\nEach row of the output parquet carries the structured catalog JSON *and* two 1152-dimensional embeddings (image + text), ready to load into a vector store.\n\nThe deep dive: [`notebooks/03_enrich_and_embed.ipynb`](notebooks/03_enrich_and_embed.ipynb)." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!python scripts/run_enrich_and_embed.py" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Inspect a sample row of the final output.\nimport pyarrow.parquet as pq\n\nout_path = f\"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/enc_vlm_enriched_with_embeddings_{N_ROWS}.parquet\"\ntbl = pq.read_table(out_path)\nprint(f\"rows: {tbl.num_rows}\")\nprint(f\"schema: {tbl.schema.names}\")\nrow = tbl.slice(0, 1).to_pylist()[0]\nprint(f\"\\ntitle: {row['title'][:100]}\")\nprint(f\"raw_output (3B student): {row['raw_output']}\")\nprint(f\"image_embedding dim: {len(row['image_embedding'])}\")\nprint(f\"text_embedding dim: {len(row['text_embedding'])}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Run as a scheduled Anyscale Job\n\nThe same scripts run unchanged as Anyscale Jobs — no rewrite, no second codebase. [`job_config.yaml`](job_config.yaml) submits Stage 3 on the same 4× L4 cluster used for the workspace runs, which is the daily / weekly cadence most teams use to refresh their catalog:\n\n```bash\nanyscale job submit --config-file job_config.yaml --env HF_TOKEN=$HF_TOKEN\n```\n\nSwap the `entrypoint` field to `scripts/run_teacher_batch_label.py` or `scripts/run_distill_student_lora.py` to submit Stages 1 or 2 as jobs." + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Clean up\n\nRemove the cached parquets, adapters, and intermediate checkpoints from cluster storage." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import shutil\n\nshutil.rmtree(\"/mnt/cluster_storage/vlm-distillation-catalog-enrichment\", ignore_errors=True)\nprint(\"cluster storage cleared\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/templates/vlm-distillation-catalog-enrichment/README.md b/templates/vlm-distillation-catalog-enrichment/README.md new file mode 100644 index 000000000..2ea2fe60b --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/README.md @@ -0,0 +1,179 @@ +# VLM teacher-student distillation for ecommerce catalogs + + +Run on Anyscale + +

+
+  +
+ +**⏱️ Time to complete**: 45 min + +This tutorial distills a Qwen2.5-VL-7B teacher into a Qwen2.5-VL-3B student so a product catalog can be enriched and embedded at 3B inference cost without sacrificing teacher-quality structured outputs. The full pipeline runs on Ray and Anyscale across three stages, all on the same 4× L4 GPU node: + +1. **Teacher batch labeling** — Run Qwen2.5-VL-7B over a catalog subset with [`ray.data.llm`](https://docs.ray.io/en/latest/data/working-with-llms.html) to produce `{category, attributes, search_tags, description}` JSON per product. +2. **Distillation SFT** — Fine-tune Qwen2.5-VL-3B on the teacher labels with Ray Train + FSDP + LoRA. Only the language model gets adapters; the vision tower stays frozen. +3. **Enrichment and embeddings** — Run the LoRA-adapted student to emit catalog JSON *and* SigLIP-2 image and text embeddings in a single streaming Ray Data graph. + +Ray is particularly powerful for this workload because it: +- **Schedules CPU and GPU stages in one cluster** so image fetching, VLM inference, and embedding extraction share a single resource pool with no idle hardware between stages. +- **Streams data between stages** without intermediate disk writes — SigLIP embeddings begin computing as soon as the first VLM-enriched batch is ready. +- **Hot-swaps LoRA adapters** at inference via vLLM's [multi-LoRA support](https://docs.vllm.ai/en/stable/features/multi_lora.html), so the distilled student loads from the same base model as the un-tuned variant. +- **Promotes any notebook to a scheduled production job** with `anyscale job submit` — no rewrite, same cluster shape. + +## Architecture + +``` +Stage 1 Stage 2 Stage 3 +───────────── ────────────── ───────────────── +7B teacher ───► 3B student SFT ───► 3B distilled enrich +batch labeling (FSDP + LoRA) + SigLIP embeddings + (one Ray Data graph) + +ray.data.llm Ray Train + FSDP ray.data.llm + Ray Data + ▼ ▼ ▼ +teacher.parquet LoRA adapter enriched_with_embeddings.parquet +(JSON labels) (~100 MB) (JSON + 1152-dim vectors) +``` + +Stages 1 and 2 each take 15–30 min on the smoke configuration (N=20 rows); Stage 3 takes ~5 min. Production runs at N=10,000 take roughly 1.5 hours on the same cluster. + +The deep-dive notebooks under [`notebooks/`](notebooks/) walk through each stage cell by cell — useful when you want to swap models, change the catalog, or tune the SFT hyperparameters. + +## Set up + +The pipeline pulls model weights and the [Amazon-Reviews-2023](https://huggingface.co/datasets/McAuley-Lab/Amazon-Reviews-2023) dataset from Hugging Face. Make sure `HF_TOKEN` is available in your environment before running the cells below. + + +```python +import os +import subprocess +import sys + +# Cache HF downloads on cluster storage so the three stages share weights. +os.environ.setdefault("HF_HOME", "/mnt/cluster_storage/hf_cache") +os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") + +# Smoke knobs — overridden by tests.sh during CI. +# Set these higher for a production run (N_ROWS=10000 is the default in each script). +N_ROWS = int(os.environ.get("N_ROWS", "20")) +TEACHER_N_ROWS = int(os.environ.get("TEACHER_N_ROWS", str(N_ROWS))) +CATEGORY = os.environ.get("CATEGORY", "Electronics") + +os.environ["N_ROWS"] = str(N_ROWS) +os.environ["TEACHER_N_ROWS"] = str(TEACHER_N_ROWS) +os.environ["CATEGORY"] = CATEGORY + +print(f"Running with N_ROWS={N_ROWS} TEACHER_N_ROWS={TEACHER_N_ROWS} CATEGORY={CATEGORY}") +``` + + +```python +!pip install -q -r requirements.txt +``` + +## Stage 1 — Teacher batch labeling + +Qwen2.5-VL-7B reads a product image and merchant-supplied title and returns a four-key JSON object (`category`, `attributes`, `search_tags`, `description`). [`ray.data.llm`](https://docs.ray.io/en/latest/data/working-with-llms.html#multimodal) wraps the vLLM engine — preprocessing, batching, and postprocessing all happen as `map_batches` stages over a Ray Dataset. With `concurrency=4`, one 7B replica runs on each of the four L4 GPUs. + +The output parquet becomes Stage 2's supervised training data. + +The deep dive: [`notebooks/01_teacher_batch_label.ipynb`](notebooks/01_teacher_batch_label.ipynb). + + +```python +!python scripts/run_teacher_batch_label.py +``` + + +```python +# Inspect a sample row of the teacher output. +import pyarrow.parquet as pq + +teacher_path = f"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/teacher_7b_enriched_{TEACHER_N_ROWS}.parquet" +tbl = pq.read_table(teacher_path) +print(f"rows: {tbl.num_rows}") +print(f"schema: {tbl.schema.names}") +row = tbl.slice(0, 1).to_pylist()[0] +print(f"\nsample title: {row['title'][:100]}") +print(f"sample teacher output (raw_output):\n{row['raw_output']}") +``` + +## Stage 2 — Distill the 3B student with FSDP + LoRA + +Ray Train + FSDP fine-tunes Qwen2.5-VL-3B on the teacher parquet using LoRA adapters. The vision tower stays frozen (standard [LLaVA recipe](https://arxiv.org/abs/2304.08485)); only the language model gets adapters, so trainable parameters drop to ~1% of the model. FSDP `FULL_SHARD` distributes the optimizer state across the four GPUs at bf16 mixed precision. + +The output is a small (~100 MB) adapter directory that drops into Stage 3 as a model swap — same base weights, fine-tuned head. + +The deep dive: [`notebooks/02_distill_student_lora.ipynb`](notebooks/02_distill_student_lora.ipynb). + + +```python +!python scripts/run_distill_student_lora.py +``` + + +```python +# Inspect the adapter. +from pathlib import Path + +adapter_dir = Path(f"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/qwen25vl_3b_enrichment_lora_{N_ROWS}") +print(f"adapter dir: {adapter_dir}") +for p in sorted(adapter_dir.rglob("*"))[:10]: + print(f" {p.relative_to(adapter_dir)}") + +# Point Stage 3 at the freshly trained adapter. +os.environ["QWEN_LORA_ADAPTER_DIR"] = str(adapter_dir) +``` + +## Stage 3 — Enrichment and embeddings in one Ray Data graph + +The third stage chains together the LoRA-adapted 3B student (via `ray.data.llm`) and the SigLIP-2 dual-tower encoder (via a Ray Data actor pool) into a single streaming pipeline. No intermediate disk writes between the VLM and the embedding stages — Ray Data hands batches directly from one `map_batches` stage to the next, so the SigLIP encoders start working as soon as the first VLM-enriched batch is ready. + +Each row of the output parquet carries the structured catalog JSON *and* two 1152-dimensional embeddings (image + text), ready to load into a vector store. + +The deep dive: [`notebooks/03_enrich_and_embed.ipynb`](notebooks/03_enrich_and_embed.ipynb). + + +```python +!python scripts/run_enrich_and_embed.py +``` + + +```python +# Inspect a sample row of the final output. +import pyarrow.parquet as pq + +out_path = f"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/enc_vlm_enriched_with_embeddings_{N_ROWS}.parquet" +tbl = pq.read_table(out_path) +print(f"rows: {tbl.num_rows}") +print(f"schema: {tbl.schema.names}") +row = tbl.slice(0, 1).to_pylist()[0] +print(f"\ntitle: {row['title'][:100]}") +print(f"raw_output (3B student): {row['raw_output']}") +print(f"image_embedding dim: {len(row['image_embedding'])}") +print(f"text_embedding dim: {len(row['text_embedding'])}") +``` + +## Run as a scheduled Anyscale Job + +The same scripts run unchanged as Anyscale Jobs — no rewrite, no second codebase. [`job_config.yaml`](job_config.yaml) submits Stage 3 on the same 4× L4 cluster used for the workspace runs, which is the daily / weekly cadence most teams use to refresh their catalog: + +```bash +anyscale job submit --config-file job_config.yaml --env HF_TOKEN=$HF_TOKEN +``` + +Swap the `entrypoint` field to `scripts/run_teacher_batch_label.py` or `scripts/run_distill_student_lora.py` to submit Stages 1 or 2 as jobs. + +## Clean up + +Remove the cached parquets, adapters, and intermediate checkpoints from cluster storage. + + +```python +import shutil + +shutil.rmtree("/mnt/cluster_storage/vlm-distillation-catalog-enrichment", ignore_errors=True) +print("cluster storage cleared") +``` diff --git a/templates/vlm-distillation-catalog-enrichment/job_config.yaml b/templates/vlm-distillation-catalog-enrichment/job_config.yaml new file mode 100644 index 000000000..e93fe7ac6 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/job_config.yaml @@ -0,0 +1,34 @@ +# Anyscale Job — VLM Batch Enrichment (Qwen2.5-VL-7B, TP=2 PP=2, 10k rows) +name: ecommerce-enc-vlm-batch-embeddings-enrichment-3b +entrypoint: python scripts/run_enc_vlm_batch_emb_enrich_3b.py + +# Pin to the working image from another day's run used (Ray 2.55.0). Without this, +# Anyscale Jobs inherit a default image that may ship an older Ray, which +# triggers "Changing the ray version is not allowed" because requirements.txt +# pins ray[data]>=2.55. +image_uri: anyscale/ray:2.55.0-py311-cu121 + +working_dir: . +requirements: requirements.txt + +compute_config: + cloud: aws-public-us-west-2 + head_node: + instance_type: m5.2xlarge + worker_nodes: + # CPU pool kept declared but locked at 0 — the g6.12xlarge worker has 48 + # vCPUs, more than enough for image fetch + decode at observed ~5/48 use. + # max_nodes:0 prevents Ray from queueing placement groups against this + # pool (we saw 145 pending PGs on the previous run with max_nodes:8). + - name: cpu-workers + instance_type: m5.4xlarge + min_nodes: 0 + max_nodes: 4 + # 4x L4 per node — one full TP=2 PP=2 replica fits per node. + # concurrency=1 in the script → exactly 1 replica → exactly 1 node. + - name: gpu-workers + instance_type: g6.12xlarge + min_nodes: 1 + max_nodes: 1 + +max_retries: 2 diff --git a/templates/vlm-distillation-catalog-enrichment/notebooks/01_teacher_batch_label.ipynb b/templates/vlm-distillation-catalog-enrichment/notebooks/01_teacher_batch_label.ipynb new file mode 100644 index 000000000..9be708596 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/notebooks/01_teacher_batch_label.ipynb @@ -0,0 +1,89 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": "# Stage 1 — Teacher batch labeling with Qwen2.5-VL-7B\n\nThis notebook runs the same labeling pipeline as [`scripts/run_teacher_batch_label.py`](../scripts/run_teacher_batch_label.py), cell by cell. It loads a slice of [Amazon-Reviews-2023](https://huggingface.co/datasets/McAuley-Lab/Amazon-Reviews-2023), prompts a Qwen2.5-VL-7B model with the product image and title, and writes a parquet of `{category, attributes, search_tags, description}` JSON per product.\n\nThe output parquet is the labeled corpus that Stage 2 distills into a 3B student. For the end-to-end story across all three stages, see the template's [`README.ipynb`](../README.ipynb).", + "id": "gen-0" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "49379612", + "metadata": {}, + "outputs": [], + "source": "import os, sys, json\nimport ray\nsys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), \"..\")))\n\nray.init(\n ignore_reinit_error=True,\n runtime_env={\n \"working_dir\": \".\",\n # vLLM 0.20+ moved TokensPrompt up from vllm.inputs.data to vllm.inputs.\n # Ray 2.55's batch LLM stage still imports via the old path, so every\n # worker patches it on startup. Must be set here on ray.init — adding it\n # to a per-actor runtime_env later is silently ignored by Ray's setup-hook\n # registry.\n \"worker_process_setup_hook\": \"src._vllm_compat.patch\",\n },\n)\nprint(\"Cluster resources:\", json.dumps(ray.cluster_resources(), indent=2))" + }, + { + "cell_type": "markdown", + "id": "ba6bc0aa", + "metadata": {}, + "source": "## Load and preprocess the catalog\n\n`ray.data.read_parquet` streams the Amazon-Reviews-2023 metadata directly from the Hugging Face Hub. The dataset arrives row-per-product with a struct of image URLs per row; the helpers below flatten that into the `{id, product_id, title, description, image_url, source}` schema the VLM stage expects.\n\nA deterministic `id` (SHA-1 of title + image URL) lets Ray Data's `CheckpointConfig` resume mid-run if the cluster blips. UUIDs would regenerate on retry and the checkpoint would match zero rows." + }, + { + "cell_type": "markdown", + "id": "2895e5b9", + "source": "### Sample row\n\nPrint the preprocessed schema, take one row, and display its image to confirm the load worked end-to-end before lighting up the GPUs.", + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8ec6cd5f", + "metadata": {}, + "outputs": [], + "source": "# load data\n\nimport os\nimport ray\nfrom huggingface_hub import HfFileSystem\nfrom typing import Optional\nfrom src.load_catalog import load_amazon_reviews_2023\nfrom PIL import Image\nimport hashlib\nimport io, requests\n\n# data source params\ncategory = \"Electronics\"\nhf_path = f\"hf://datasets/McAuley-Lab/Amazon-Reviews-2023/raw_meta_{category}\"\n\n# preprocess params\nn_rows = 50\nseed = 42\n\n# paths\nbase_dir = \"/mnt/cluster_storage/vlm-distillation-catalog-enrichment\"\ncache_path = f\"{base_dir}/catalog_demo_{n_rows}.parquet\"\ncheckpoint_path = f\"{base_dir}/catalog_demo_{n_rows}_checkpoint\"\n\n\ndef _extract_image_url(images_field) -> Optional[str]:\n \"\"\"Pick the first 'large' image URL (or hi_res/thumb fallback) from Amazon's image struct.\n\n HF parquet stores `images` as a struct of parallel lists,\n e.g. ``{\"large\": [url1, url2], \"hi_res\": [...], \"thumb\": [...], \"variant\": [...]}``,\n not as a list of per-image dicts. We take the first URL of the best\n available size.\n \"\"\"\n if not images_field or not isinstance(images_field, dict):\n return None\n for key in (\"large\", \"hi_res\", \"thumb\"):\n urls = images_field.get(key)\n if urls is not None and len(urls) > 0:\n return urls[0]\n return None\n\n\ndef _has_title_and_image(row: dict) -> bool:\n \"\"\"Drop rows missing either a title or any usable image URL.\"\"\"\n title = row.get(\"title\")\n if not (title and title.strip()):\n return False\n return _extract_image_url(row.get(\"images\")) is not None\n\n\ndef _coerce_description(desc_field) -> str:\n if isinstance(desc_field, list):\n return \" \".join(str(x) for x in desc_field if x).strip()\n if isinstance(desc_field, str):\n return desc_field.strip()\n return \"\"\n\n\ndef _row_id(title: str, image_url: str) -> str:\n \"\"\"Stable per-(product, image) ID: SHA-1 of title + image_url, truncated.\n\n Deterministic across runs so Ray Data's CheckpointConfig can resume —\n `uuid.uuid4()` regenerated fresh IDs each submission, making the\n checkpoint match zero rows on retry. Hashing on (title, image_url) keeps\n the ID unique under future multi-image fan-out (same product, different\n image → different ID).\n \"\"\"\n return hashlib.sha1(f\"{title}|{image_url}\".encode(\"utf-8\")).hexdigest()[:16]\n\n\ndef _normalize_amazon_row_to_image(row: dict) -> dict:\n \"\"\"One catalog row per product, using the single best image URL (use with `.map`).\"\"\"\n title = row[\"title\"].strip()[:512]\n image_url = _extract_image_url(row[\"images\"])\n return {\n \"id\": _row_id(title, image_url),\n \"product_id\": row.get(\"parent_asin\") or row.get(\"asin\") or \"\",\n \"title\": title,\n \"description\": _coerce_description(row.get(\"description\"))[:1024],\n \"image_url\": image_url,\n \"source\": \"amazon-reviews-2023\",\n }\n\ndef _normalize_amazon_row_to_images(row: dict, max_per_product: int = 8) -> list[dict]:\n \"\"\"One catalog row per *image* — explodes a product into N rows (use with `.flat_map`).\n\n Takes 'large' URLs (falling back to hi_res / thumb), capped at\n ``max_per_product``. Each output row carries an ``image_idx`` so\n ``(product_id, image_idx)`` is a stable row key.\n \"\"\"\n images = row.get(\"images\") or {}\n if not isinstance(images, dict):\n return []\n urls = (\n images.get(\"large\")\n or images.get(\"hi_res\")\n or images.get(\"thumb\")\n or []\n )[:max_per_product]\n if not urls:\n return []\n\n product_id = row.get(\"parent_asin\") or row.get(\"asin\") or \"\"\n title = row[\"title\"].strip()[:512]\n description = _coerce_description(row.get(\"description\"))[:1024]\n\n return [\n {\n \"id\": _row_id(title, url),\n \"product_id\": product_id,\n \"image_idx\": i,\n \"title\": title,\n \"description\": description,\n \"image_url\": url,\n \"source\": \"amazon-reviews-2023\",\n }\n for i, url in enumerate(urls)\n ]\n\nif True: #not os.path.exists(cache_path):\n # load from hf\n print(f\"[load hf] Loading data from huggingface hub: {hf_path}...\")\n ds = ray.data.read_parquet(\n hf_path,\n file_extensions=[\"parquet\"],\n filesystem=HfFileSystem(),\n )\n print(\"[load hf] count:\", ds.count())\n print(\"[load hf] original schema:\", ds.schema())\n\n # reduce and filter\n ds = ds.limit(n_rows)\n ds = ds.filter(_has_title_and_image)\n\n # naively assign 1:1 image per product, in reality would explode images across same product for 1:many\n ds = ds.map(_normalize_amazon_row_to_image)\n # ds = ds.map(_normalize_amazon_row_to_images) # untested to his is the idea\n ds = ds.random_shuffle(seed=seed)\n\n # cache preprocessed dataset\n ds.write_parquet(cache_path)\nelse:\n # cache exists already\n ds = ray.data.read_parquet(cache_path)\n" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "480b21fe", + "metadata": {}, + "outputs": [], + "source": "# show example\nprint()\nprint(\"[load] preprocessed schema:\", ds.schema())\nprint()\n\nrow = ds.take(1)[0]\nprint()\nprint(\"[load] sample row:\", json.dumps(row, indent=2))\nprint()\n\n\nresp = requests.get(row[\"image_url\"], timeout=5,\n headers={\"User-Agent\": \"vlm-distillation-catalog-enrichment/1.0\"})\nimg = Image.open(io.BytesIO(resp.content))\nimg.thumbnail((256, 256))\nimg" + }, + { + "cell_type": "markdown", + "id": "74444c5e", + "metadata": {}, + "source": "## Run the 7B teacher with `ray.data.llm`\n\n[`ray.data.llm`](https://docs.ray.io/en/latest/data/working-with-llms.html#multimodal) wraps the vLLM engine as a Ray Data stage. `vLLMEngineProcessorConfig` declares the model, the per-replica engine kwargs, and the multimodal preparation stage; `build_processor` then takes a `preprocess`/`postprocess` pair and returns a callable that maps over the Dataset.\n\nFor the 7B teacher: one replica per L4 (`tensor_parallel_size=1`), `concurrency=4` replicas, `batch_size=16`. The prompt asks for a four-key JSON object; the `vlm_postprocess` step lifts `generated_text` into a `raw_output` column alongside the original row fields, then `write_parquet` materializes the result." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0fa6d414", + "metadata": {}, + "outputs": [], + "source": "import io, json, re\nfrom ray.data.llm import vLLMEngineProcessorConfig, build_processor\nfrom ray.data.checkpoint import CheckpointConfig\n\n# Resumable batch inference: the row `id` (deterministic SHA-1 above) keys the\n# checkpoint, so a re-submission picks up where the prior run left off.\nctx = ray.data.DataContext.get_current()\nctx.checkpoint_config = CheckpointConfig(\n id_column=\"id\",\n checkpoint_path=checkpoint_path,\n delete_checkpoint_on_success=False,\n)\n\nconfig = vLLMEngineProcessorConfig(\n model_source=\"Qwen/Qwen2.5-VL-7B-Instruct\",\n engine_kwargs={\n \"tensor_parallel_size\": 1, # 7B fits on a single L4 at bf16\n \"pipeline_parallel_size\": 1,\n \"max_model_len\": 4096,\n \"trust_remote_code\": True,\n \"limit_mm_per_prompt\": {\"image\": 1},\n },\n batch_size=16,\n concurrency=4, # one replica per GPU on a g6.12xlarge\n prepare_multimodal_stage={\"enabled\": True},\n runtime_env={\"worker_process_setup_hook\": \"src._vllm_compat.patch\"},\n)\n\nprompt = \"\"\"\\\nYou are a product catalog enrichment assistant. Given a product image and \\\nthe merchant-supplied title, output a JSON object with exactly these keys:\n\n category: one short string (e.g. \"Wireless Earbuds\")\n attributes: a list of 3 short attribute strings\n search_tags: a list of 5 short search keywords\n description: a single sentence (<= 30 words)\n\nTitle: {title}\n\nReturn ONLY the JSON object, no commentary.\\\n\"\"\"\n\n\ndef build_messages_url(url, title):\n \"\"\"Image + text in OpenAI chat-completion format.\"\"\"\n return [\n {\n \"role\": \"user\",\n \"content\": [\n {\"type\": \"image_url\", \"image_url\": {\"url\": url}},\n {\"type\": \"text\", \"text\": prompt.format(title=title)},\n ],\n }\n ]\n\n\ndef vlm_preprocess(row):\n return {\n \"id\": row[\"id\"],\n \"messages\": build_messages_url(row[\"image_url\"], row[\"title\"]),\n \"sampling_params\": {\"max_tokens\": 256, \"temperature\": 0.0},\n }\n\n\ndef vlm_postprocess(row):\n return {\n \"id\": row[\"id\"], # required for checkpoint id_column\n \"product_id\": row[\"product_id\"],\n \"title\": row[\"title\"],\n \"image_url\": row[\"image_url\"],\n \"source\": row[\"source\"],\n \"raw_output\": row[\"generated_text\"],\n }\n\n\nvlm_processor = build_processor(\n config,\n preprocess=vlm_preprocess,\n postprocess=vlm_postprocess,\n)\n\nds = vlm_processor(ds)\n\n# Show the first two enriched rows.\nfor output in ds.take(limit=2):\n print(json.dumps(output, indent=2))\n\nds.write_parquet(f\"{base_dir}/teacher_7b_enriched_{n_rows}.parquet\")" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0a392c48", + "metadata": {}, + "outputs": [], + "source": "print(ds.stats())" + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/templates/vlm-distillation-catalog-enrichment/notebooks/02_distill_student_lora.ipynb b/templates/vlm-distillation-catalog-enrichment/notebooks/02_distill_student_lora.ipynb new file mode 100644 index 000000000..e91f6530d --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/notebooks/02_distill_student_lora.ipynb @@ -0,0 +1,129 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro-md", + "metadata": {}, + "source": "# Stage 2 — Distill the 3B student with FSDP + LoRA\n\nThis notebook fine-tunes Qwen2.5-VL-3B on the Stage 1 teacher parquet using Ray Train + FSDP + LoRA. The vision tower stays frozen (standard [LLaVA recipe](https://arxiv.org/abs/2304.08485)); only the language model gets LoRA adapters, so trainable parameters drop to ~1% of the model. FSDP `FULL_SHARD` distributes optimizer state across the four GPUs at bf16 mixed precision.\n\nThe output is a ~100 MB adapter directory that drops into Stage 3 as a model swap. The notebook scales every knob (rows, workers, epochs) down so the loop finishes in minutes on a workspace cluster; [`scripts/run_distill_student_lora.py`](../scripts/run_distill_student_lora.py) runs the same pipeline at production scale. For the end-to-end story across all three stages, see the template's [`README.ipynb`](../README.ipynb)." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-init", + "metadata": {}, + "outputs": [], + "source": "import os, sys, json\nimport ray\n\nsys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), \"..\")))\n\n# If the cluster is already up with stale runtime_env (e.g. missing peft on\n# workers), tear it down so the new pip list below actually takes effect.\nif ray.is_initialized():\n ray.shutdown()\n\nray.init(\n runtime_env={\n \"working_dir\": os.path.abspath(os.path.join(os.getcwd(), \"..\")),\n # peft + accelerate aren't in the workspace base image; install on\n # workers so train_loop_per_worker can `from peft import ...`.\n \"pip\": [\"peft>=0.12\", \"accelerate>=0.34\"],\n },\n)\nprint(\"Cluster resources:\", json.dumps(ray.cluster_resources(), indent=2))" + }, + { + "cell_type": "markdown", + "id": "config-md", + "metadata": {}, + "source": "## Configure the run\n\nThe data pipeline and training loop import from [`scripts/run_distill_student_lora.py`](../scripts/run_distill_student_lora.py), so the notebook exercises the same code paths as a production job. Only row count, worker count, and epoch count are scaled down here." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-config", + "metadata": {}, + "outputs": [], + "source": "# Production helpers — same code paths as the submitted job\nfrom scripts.run_distill_student_lora import (\n BuildSFTExample, train_loop_per_worker,\n fetch_and_resize_images, _split_bucket, _strip_code_fence,\n STUDENT_MODEL_ID, MAX_SEQ_LEN, MIN_PIXELS, MAX_PIXELS,\n LORA_R, LORA_ALPHA, LORA_DROPOUT, LORA_TARGET_MODULES,\n)\nfrom ray.train import ScalingConfig, RunConfig, CheckpointConfig\nfrom ray.train.torch import TorchTrainer\n\n# Smoke-scale knobs — small enough to run end-to-end in ~10 min on 2× L4.\nN_SMOKE = 100\nNUM_WORKERS_SMOKE = 2\nNUM_EPOCHS_SMOKE = 1\nPER_DEVICE_BS_SMOKE = 1\nGRAD_ACCUM_SMOKE = 4 # effective batch = 1 × 2 × 4 = 8\nLR_SMOKE = 1e-4\nWARMUP_RATIO_SMOKE = 0.0 # tiny run, skip warmup\n\nBASE_DIR = \"/mnt/cluster_storage/vlm-distillation-catalog-enrichment\"\nTEACHER_PARQUET = f\"{BASE_DIR}/teacher_7b_enriched_10000.parquet\"\nSMOKE_CACHE_PATH = f\"{BASE_DIR}/sft_cache_smoke_{N_SMOKE}.parquet\"\nSMOKE_ADAPTER_DIR = f\"{BASE_DIR}/qwen25vl_3b_enrichment_lora_smoke\"\nSMOKE_RUN_DIR = f\"{BASE_DIR}/qwen25vl_3b_enrichment_runs_smoke\"\n\nif not os.path.exists(TEACHER_PARQUET):\n raise FileNotFoundError(\n f\"Teacher parquet not found at {TEACHER_PARQUET}. \"\n f\"Run scripts/run_teacher_batch_label.py first.\"\n )\n\nprint(f\"smoke config: N={N_SMOKE} workers={NUM_WORKERS_SMOKE} \"\n f\"epochs={NUM_EPOCHS_SMOKE} effective_batch={PER_DEVICE_BS_SMOKE * NUM_WORKERS_SMOKE * GRAD_ACCUM_SMOKE}\")" + }, + { + "cell_type": "markdown", + "id": "cache-md", + "metadata": {}, + "source": "## Build the SFT cache\n\nRead the teacher parquet, drop rows whose teacher output isn't valid JSON, assign a deterministic train/val/test split, fetch each image once, and write a self-contained training parquet. The image bytes live in this cache so epochs 2..N never go back to HTTP." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-cache", + "metadata": {}, + "outputs": [], + "source": "def _teacher_output_is_valid(row):\n try:\n obj = json.loads(_strip_code_fence(row[\"raw_output\"]))\n return all(k in obj for k in (\"category\", \"attributes\", \"search_tags\", \"description\"))\n except Exception:\n return False\n\ndef _attach_split(row):\n row[\"split\"] = _split_bucket(row[\"id\"])\n return row\n\nif os.path.exists(SMOKE_CACHE_PATH):\n cached = ray.data.read_parquet(SMOKE_CACHE_PATH)\n print(f\"[cache] reusing {SMOKE_CACHE_PATH}\")\nelse:\n ds = ray.data.read_parquet(TEACHER_PARQUET)\n ds = ds.filter(_teacher_output_is_valid).limit(N_SMOKE).map(_attach_split)\n ds = ds.map_batches(\n fetch_and_resize_images,\n batch_size=16,\n concurrency=8,\n batch_format=\"numpy\",\n )\n ds.write_parquet(SMOKE_CACHE_PATH)\n cached = ray.data.read_parquet(SMOKE_CACHE_PATH)\n\nn_train = cached.filter(lambda r: r[\"split\"] == \"train\").count()\nn_val = cached.filter(lambda r: r[\"split\"] == \"val\").count()\nn_test = cached.filter(lambda r: r[\"split\"] == \"test\").count()\nprint(f\"[split] train={n_train} val={n_val} test={n_test}\")" + }, + { + "cell_type": "markdown", + "id": "sample-md", + "metadata": {}, + "source": "## Sanity check one row\n\nConfirm the teacher target parses, the image decodes, and the split assignment looks right. Any data issue surfaces here before the GPUs spin up." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-sample", + "metadata": {}, + "outputs": [], + "source": "from PIL import Image\nimport io as _io\n\nrow = cached.take(1)[0]\nprint(\"id: \", row[\"id\"])\nprint(\"split: \", row[\"split\"])\nprint(\"title: \", row[\"title\"][:90])\nprint(\"teacher: \", _strip_code_fence(row[\"raw_output\"])[:240])\nimg = Image.open(_io.BytesIO(row[\"image_bytes\"]))\nimg.thumbnail((256, 256))\nimg" + }, + { + "cell_type": "markdown", + "id": "examples-md", + "metadata": {}, + "source": "## Build SFT examples\n\n`BuildSFTExample` applies the Qwen2.5-VL chat template and processor to each row, producing fixed-shape numpy columns (`input_ids`, `labels`, `attention_mask`, `pixel_values`, `image_grid_thw`) that Ray Data streams into the trainer. Loss masking on the user prefix means only the assistant's JSON tokens contribute to the loss — image and prompt tokens are `-100`." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-examples", + "metadata": {}, + "outputs": [], + "source": "build_kwargs = dict(\n fn_constructor_kwargs={\n \"model_id\": STUDENT_MODEL_ID,\n \"max_length\": MAX_SEQ_LEN,\n \"min_pixels\": MIN_PIXELS,\n \"max_pixels\": MAX_PIXELS,\n },\n num_cpus=2,\n concurrency=4,\n)\n\n# Repartition before BuildSFTExample so the OutputSplitter at the trainer\n# end has multiple blocks to balance across ranks. The smoke cache parquet\n# is a single block; without this, OutputSplitter[split(2, equal=True)] can\n# stall one rank waiting for the splitter to materialize, which on FSDP\n# manifests as a NCCL collective hang at the first all-gather.\ntrain_ds = (cached.filter(lambda r: r[\"split\"] == \"train\")\n .repartition(NUM_WORKERS_SMOKE * 2)\n .map(BuildSFTExample, **build_kwargs))\nval_ds = (cached.filter(lambda r: r[\"split\"] == \"val\")\n .repartition(NUM_WORKERS_SMOKE * 2)\n .map(BuildSFTExample, **build_kwargs))\n\n# Inspect one tokenized example. The shapes here drive every batching\n# decision in the trainer; if anything's variable across rows, training will\n# fail on the first stack() call.\nexample = train_ds.take(1)[0]\nfor k, v in example.items():\n if hasattr(v, \"shape\"):\n print(f\" {k:18} shape={tuple(v.shape)} dtype={v.dtype}\")\n else:\n print(f\" {k:18} = {v!r}\")\n\n# Loss-mask sanity check: assistant tokens are non-(-100), prefix tokens are\n# all -100. The fraction non-masked ≈ length(target_json) / MAX_SEQ_LEN.\nimport numpy as np\nlabels = example[\"labels\"]\nn_loss_tokens = int((labels != -100).sum())\nprint(f\"\\nloss-mask: {n_loss_tokens}/{len(labels)} tokens contribute to loss \"\n f\"({100 * n_loss_tokens / len(labels):.1f}%)\")" + }, + { + "cell_type": "markdown", + "id": "train-md", + "metadata": {}, + "source": "## Run training\n\n`TorchTrainer` runs `train_loop_per_worker` across the GPU workers. Each worker loads the 3B base model, applies LoRA, FSDP-wraps the language model, and streams the SFT shards through. On 2× L4 at smoke scale this typically finishes in 8–12 minutes — most of which is the first-step compile and model load.\n\nIf the run OOMs, drop `MAX_SEQ_LEN` to 1024 or set `PER_DEVICE_BS_SMOKE = 1` and `GRAD_ACCUM_SMOKE = 8`." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-train", + "metadata": {}, + "outputs": [], + "source": "trainer = TorchTrainer(\n train_loop_per_worker=train_loop_per_worker,\n train_loop_config={\n \"model_id\": STUDENT_MODEL_ID,\n \"lr\": LR_SMOKE,\n \"num_epochs\": NUM_EPOCHS_SMOKE,\n \"per_device_bs\": PER_DEVICE_BS_SMOKE,\n \"grad_accum\": GRAD_ACCUM_SMOKE,\n \"warmup_ratio\": WARMUP_RATIO_SMOKE,\n \"weight_decay\": 0.01,\n \"lora_r\": LORA_R,\n \"lora_alpha\": LORA_ALPHA,\n \"lora_dropout\": LORA_DROPOUT,\n \"lora_target_modules\": LORA_TARGET_MODULES,\n \"train_size\": n_train,\n \"adapter_dir\": SMOKE_ADAPTER_DIR,\n # Production uses 50; smoke has only ~10 grad steps total, so drop\n # to 5 to demonstrate that step-level checkpointing actually fires.\n # Expect ckpts at step 5 and step 10 (the latter coincides with the\n # end-of-epoch ckpt, so retention picks whichever has lower val_loss).\n \"save_every_n_steps\": 5,\n # Eval bs > train bs: no backward, no grad checkpoint tax. Smoke val\n # is ~8 rows, sharded across 2 workers → 4 rows each → exactly 1\n # batch per rank at eval_bs=4. Production uses 4 for the same reason.\n \"eval_per_device_bs\": 4,\n # Reproducibility: torch + cuda seeds offset per-rank, random/numpy\n # uniform. Without this, dropout and any python-side randomness drift\n # run-to-run even with identical data and config.\n \"seed\": 42,\n },\n scaling_config=ScalingConfig(\n num_workers=NUM_WORKERS_SMOKE,\n use_gpu=True,\n accelerator_type=\"L4\",\n resources_per_worker={\"GPU\": 1, \"CPU\": 4},\n # PACK both workers onto the same node so FSDP all-gathers /\n # reduce-scatters use intra-node shared memory instead of crossing\n # the cluster network. Without this, a workspace with only\n # single-GPU L4 nodes available will land the 2 workers on\n # separate nodes; the first NCCL collective then either crawls\n # (10x+ slower per step) or hangs at all-gather. Requires a node\n # with >= num_workers GPUs (e.g. g6.12xlarge with 4x L4).\n # placement_strategy=\"PACK\",\n ),\n run_config=RunConfig(\n storage_path=SMOKE_RUN_DIR,\n # num_to_keep=3 so all of (step5, step10, epoch0) can co-exist for\n # the demo — easy to inspect each in result.best_checkpoints.\n checkpoint_config=CheckpointConfig(\n num_to_keep=3,\n checkpoint_score_attribute=\"val_loss\",\n checkpoint_score_order=\"min\",\n ),\n ),\n datasets={\"train\": train_ds, \"val\": val_ds},\n)\nresult = trainer.fit()\nprint(f\"\\n[done] metrics: {result.metrics}\")\nprint(f\"[done] best checkpoint: {result.checkpoint}\")\nprint(f\"[done] adapter: {SMOKE_ADAPTER_DIR}\")" + }, + { + "cell_type": "markdown", + "id": "verify-md", + "metadata": {}, + "source": "## Verify the adapter\n\nConfirm the saved directory has the expected files (`adapter_config.json` plus the LoRA weights) and that the config records the LoRA rank and target modules used during training. The full inference test belongs in Stage 3, which loads this adapter into vLLM at serve time." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-verify", + "metadata": {}, + "outputs": [], + "source": "import json as _json\n\nfiles = sorted(os.listdir(SMOKE_ADAPTER_DIR))\nprint(\"adapter directory contents:\")\nfor f in files:\n p = os.path.join(SMOKE_ADAPTER_DIR, f)\n size_mb = os.path.getsize(p) / 1e6\n print(f\" {f:40} {size_mb:>8.2f} MB\")\n\nwith open(os.path.join(SMOKE_ADAPTER_DIR, \"adapter_config.json\")) as fh:\n adapter_cfg = _json.load(fh)\nprint(\"\\nadapter_config.json (key fields):\")\nfor k in (\"r\", \"lora_alpha\", \"lora_dropout\", \"target_modules\",\n \"task_type\", \"base_model_name_or_path\"):\n print(f\" {k:28} = {adapter_cfg.get(k)!r}\")" + }, + { + "cell_type": "markdown", + "id": "next-md", + "metadata": {}, + "source": "## Submit the full run as an Anyscale Job\n\nOnce the smoke run looks healthy, submit the full 10k-row training as a job — same compute config as the workspace, no code changes:\n\n```bash\nanyscale job submit --config-file ../job_config.yaml \\\n --entrypoint \"python scripts/run_distill_student_lora.py\" \\\n --env HF_TOKEN=$HF_TOKEN\n```\n\nThen load the trained adapter into the Stage 3 inference pipeline by setting `QWEN_LORA_ADAPTER_DIR` to the adapter directory before running [`scripts/run_enrich_and_embed.py`](../scripts/run_enrich_and_embed.py)." + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/templates/vlm-distillation-catalog-enrichment/notebooks/03_enrich_and_embed.ipynb b/templates/vlm-distillation-catalog-enrichment/notebooks/03_enrich_and_embed.ipynb new file mode 100644 index 000000000..b782bc63c --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/notebooks/03_enrich_and_embed.ipynb @@ -0,0 +1,95 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "intro-md", + "metadata": {}, + "source": "# Stage 3 — Enrichment and embeddings in one Ray Data graph\n\nThis notebook runs the LoRA-adapted 3B student over the catalog and emits both structured enrichment JSON *and* SigLIP-2 (so400m, 1152-dim) image and text embeddings in a single streaming Ray Data pipeline. No intermediate disk writes between the VLM and the embedding stages — Ray Data hands batches directly from one `map_batches` stage to the next.\n\nIf the Stage 2 adapter directory exists on cluster storage, the VLM stage routes through the fine-tuned weights via vLLM multi-LoRA; otherwise it falls back to the base Qwen2.5-VL-3B model. [`scripts/run_enrich_and_embed.py`](../scripts/run_enrich_and_embed.py) runs the same pipeline at production scale. For the end-to-end story across all three stages, see the template's [`README.ipynb`](../README.ipynb)." + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-init", + "metadata": {}, + "outputs": [], + "source": "import os, sys, json\n\nsys.path.insert(0, os.path.abspath(os.path.join(os.getcwd(), \"..\")))\n\n# MUST be set before importing the script — the LoRA detection runs at\n# module-import time. If the dir doesn't exist the script falls back to base 3B.\nos.environ.setdefault(\n \"QWEN_LORA_ADAPTER_DIR\",\n \"/mnt/cluster_storage/vlm-distillation-catalog-enrichment/qwen25vl_3b_enrichment_lora_smoke\",\n)\n\nimport ray\nif ray.is_initialized():\n ray.shutdown()\n\nray.init(\n runtime_env={\n \"working_dir\": os.path.abspath(os.path.join(os.getcwd(), \"..\")),\n },\n)\nprint(\"Cluster resources:\", json.dumps(ray.cluster_resources(), indent=2))" + }, + { + "cell_type": "markdown", + "id": "db422898", + "source": "## Configure the run\n\nThe pipeline imports from [`scripts/run_enrich_and_embed.py`](../scripts/run_enrich_and_embed.py) and rebinds a few module attributes so `build_catalog` and `build_vlm_processor` pick up the smaller smoke knobs at call time.", + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-config", + "metadata": {}, + "outputs": [], + "source": "from scripts import run_enrich_and_embed as enc\n\n# Smoke overrides — same code paths as the production job, just smaller. We\n# rebind module attributes so the script's helpers (build_catalog,\n# build_vlm_processor) pick them up at call time via the module globals.\nN_SMOKE = 50\nenc.N_ROWS = N_SMOKE\nenc.VLM_CONCURRENCY = 1\nenc.EMB_CONCURRENCY = 1\nenc.PROCESS_CONCURRENCY = 4\nenc.FETCH_CONCURRENCY = 8\nenc.CACHE_PATH = f\"{enc.BASE_DIR}/catalog_{enc.CATEGORY}_{N_SMOKE}.parquet\"\nenc.CHECKPOINT_PATH = f\"{enc.BASE_DIR}/enc_vlm_emb_enrich_smoke_{N_SMOKE}_checkpoint\"\nenc.OUTPUT_PATH = f\"{enc.BASE_DIR}/enc_vlm_enriched_with_embeddings_smoke_{N_SMOKE}.parquet\"\n\nvlm_mode = (\n f\"LoRA: {enc.LORA_ADAPTER_NAME} (loaded from {enc.LORA_REMOTE_PREFIX})\"\n if enc.LORA_ADAPTER_NAME\n else f\"base 3B ({enc.VLM_MODEL_SOURCE})\"\n)\nprint(f\"smoke config: N={N_SMOKE} VLM_CONCURRENCY={enc.VLM_CONCURRENCY} EMB_CONCURRENCY={enc.EMB_CONCURRENCY}\")\nprint(f\"VLM mode: {vlm_mode}\")\nprint(f\"output: {enc.OUTPUT_PATH}\")" + }, + { + "cell_type": "markdown", + "id": "5e6f2b6a", + "source": "## Load the catalog\n\n`build_catalog` reads Amazon-Reviews-2023 metadata from the Hugging Face Hub, filters to rows with usable images, normalizes the schema, and caches the result as parquet on cluster storage. First run does the HF read + URL HEAD checks; later runs reuse the cache.", + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-input", + "metadata": {}, + "outputs": [], + "source": "import io as _io\nimport requests as _r\nfrom PIL import Image as _Image\n\n# build_catalog() caches at CACHE_PATH; first run does the HF read + URL HEAD\n# checks, repeat runs reuse the parquet. Returns a Ray Dataset.\nds = enc.build_catalog()\nprint(f\"\\n[catalog] {ds.count()} rows after filtering\")\n\nrow = ds.take(1)[0]\nprint(\"\\n[input row]\")\nfor k, v in row.items():\n print(f\" {k:14} = {repr(v)[:100]}\")\n\nresp = _r.get(row[\"image_url\"], timeout=5, headers={\"User-Agent\": \"vlm-distillation-catalog-enrichment/1.0\"})\nimg = _Image.open(_io.BytesIO(resp.content))\nimg.thumbnail((256, 256))\nimg" + }, + { + "cell_type": "markdown", + "id": "c07de5cf", + "source": "## Build the streaming pipeline\n\nFour `map_batches` stages chain together: CPU image fetch → VLM enrichment (`ray.data.llm`, optionally LoRA-routed) → CPU SigLIP processor → GPU SigLIP embed. `CheckpointConfig` runs over the row `id` so a re-submission picks up from the last successful batch.\n\n`write_parquet` triggers execution; rows stream through the four stages and write out as soon as each batch finishes embedding.", + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-pipeline", + "metadata": {}, + "outputs": [], + "source": "from ray.data.checkpoint import CheckpointConfig\n\n# CPU fetch + decode + resize\nds = ds.map_batches(\n enc.fetch_and_resize,\n batch_size=16,\n concurrency=enc.FETCH_CONCURRENCY,\n batch_format=\"numpy\",\n)\n\n# VLM enrichment via ray.data.llm (vLLM under the hood; LoRA-routed if attached)\nds = enc.build_vlm_processor()(ds)\n\n# SigLIP CPU process — image processor + text tokenizer over title + VLM description\nds = ds.map(\n enc.ProcessSigLIP,\n fn_constructor_kwargs={\"model_id\": enc.EMB_MODEL_SOURCE},\n num_cpus=1,\n concurrency=enc.PROCESS_CONCURRENCY,\n)\n\n# SigLIP GPU embed — pure forward pass on pre-tensorized inputs\nds = ds.map_batches(\n enc.EmbedSigLIP,\n fn_constructor_kwargs={\"model_id\": enc.EMB_MODEL_SOURCE},\n batch_size=enc.EMB_BATCH_SIZE,\n num_gpus=1,\n concurrency=enc.EMB_CONCURRENCY,\n batch_format=\"numpy\",\n)\n\n# Set CheckpointConfig AFTER build_catalog so the catalog's parquet write\n# doesn't poison the inference checkpoint with already-seen IDs (would skip\n# ~99% of rows on the next run).\nctx = ray.data.DataContext.get_current()\nctx.checkpoint_config = CheckpointConfig(\n id_column=\"id\",\n checkpoint_path=enc.CHECKPOINT_PATH,\n delete_checkpoint_on_success=False,\n)\n\nds.write_parquet(enc.OUTPUT_PATH)\nprint(f\"[done] wrote {enc.OUTPUT_PATH}\")" + }, + { + "cell_type": "markdown", + "id": "a219b392", + "source": "## Inspect the output\n\nRead back a couple of rows of the output parquet and print the structured catalog JSON next to the first six dimensions of each embedding. Each row in the output carries the original product fields, the VLM's `raw_output` JSON, an `image_embedding`, and a `text_embedding`.", + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-output", + "metadata": {}, + "outputs": [], + "source": "sample = ray.data.read_parquet(enc.OUTPUT_PATH).take(2)\nprint(f\"[output] read back {len(sample)} sample rows from {enc.OUTPUT_PATH}\\n\")\n\nfor r in sample:\n print(json.dumps({\n \"title\": r[\"title\"],\n \"raw_output\": r[\"raw_output\"],\n \"image_embedding[:6]\": [round(float(x), 3) for x in r[\"image_embedding\"][:6]],\n \"text_embedding[:6]\": [round(float(x), 3) for x in r[\"text_embedding\"][:6]],\n }, indent=2))\n print()" + } + ], + "metadata": { + "kernelspec": { + "display_name": "base", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/templates/vlm-distillation-catalog-enrichment/requirements.txt b/templates/vlm-distillation-catalog-enrichment/requirements.txt new file mode 100644 index 000000000..21849a1dc --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/requirements.txt @@ -0,0 +1,14 @@ +ray[data,train]>=2.55 +vllm>=0.6.5 +transformers>=4.45 +torch>=2.4 +peft>=0.12 +accelerate>=0.34 +pillow>=10.0 +requests>=2.31 +datasets>=2.20 +huggingface_hub>=1.1.6 +pyarrow>=14.0 +pandas>=2.0 +numpy>=1.24 +scipy>=1.13 diff --git a/templates/vlm-distillation-catalog-enrichment/requirements_ft.txt b/templates/vlm-distillation-catalog-enrichment/requirements_ft.txt new file mode 100644 index 000000000..127c3422b --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/requirements_ft.txt @@ -0,0 +1,2 @@ +peft>=0.12 +accelerate>=0.34 diff --git a/templates/vlm-distillation-catalog-enrichment/scripts/run_distill_student_lora.py b/templates/vlm-distillation-catalog-enrichment/scripts/run_distill_student_lora.py new file mode 100644 index 000000000..7372a4129 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/scripts/run_distill_student_lora.py @@ -0,0 +1,799 @@ +""" +Stage 2 — Qwen2.5-VL-3B SFT (FSDP + LoRA) on Stage 1's teacher labels. + +Distills the 7B teacher's enrichment JSON into the 3B student so you ship +teacher-grade structured catalog output at 3B inference cost. Reads the +parquet that scripts/run_teacher_batch_label.py wrote and uses each row's +`raw_output` (the teacher's JSON string) as the SFT target. + +The output is a LoRA adapter that drops into Stage 3 (run_enrich_and_embed.py) +as a model swap — same Qwen2.5-VL-3B base, fine-tuned weights. + +Pipeline: + LOAD teacher parquet → CPU FETCH images + SPLIT → BUILD SFT EXAMPLES (CPU) + → SFT TRAIN (Ray Train + FSDP + LoRA on 4× L4) + → WRITE adapter + +Run on the workspace cluster (after the Stage 1 teacher parquet exists): + python scripts/run_distill_student_lora.py + +Or as an Anyscale Job (same 4× L4 fleet as Stage 1): + anyscale job submit --config-file job_config.yaml \\ + --entrypoint "python scripts/run_distill_student_lora.py" \\ + --env HF_TOKEN=$HF_TOKEN + +────────────────────────────────────────────────────────── +Why distill the 7B teacher into a 3B student +────────────────────────────────────────────────────────── +Same prompt, same 4-key JSON schema, but a fraction of the inference cost. +The student is bounded by teacher quality — it will not exceed the 7B teacher +— but closes most of the gap on a focused task (structured JSON over a narrow +product category) at 3B model size. The category-specialization is a feature: +the teacher parquet is produced on one CATEGORY at a time, and you ship a +fleet of per-category adapters served from one base model checkpoint. + +────────────────────────────────────────────────────────── +What's optimized for L4 24GB and the 4× L4 node shape +────────────────────────────────────────────────────────── + 1. LoRA on the LLM, vision tower frozen. Trainable params drop to ~1% of + the model. The vision encoder is already strong for product photography + out-of-the-box and freezing it sidesteps multimodal training instability. + 2. FSDP FULL_SHARD with bf16 mixed precision. 3B in bf16 is ~6 GB; sharded + across 4 GPUs that's ~1.5 GB/GPU for params. Plenty of headroom for + activations on L4 24 GB. + 3. Gradient checkpointing on the LLM. Halves activation memory for the + long image-token prefix at the cost of one extra forward per step — + standard VLM SFT recipe. + 4. Per-device batch size of 1 with grad accumulation. Image-token prefixes + (~512 tokens at our pixel cap) make per-device batches expensive in + activation memory; bs=1 + grad_accum 16 → effective batch of 64. + 5. Cached training parquet with image_bytes pre-fetched. Epochs 2..N read + from cache and never go back to HTTP. Same trick as the inference + scripts, applied to the SFT path. + 6. Loss masking on the user prefix. Standard SFT — only the assistant's + JSON tokens contribute to the loss, image and prompt tokens are -100. + 7. Pixel cap matches the inference-time cap (max_pixels = 512·28²) so + the fine-tuned weights see exactly the same visual feature distribution + they will see at deployment. +""" + +import os, sys, json, hashlib, io + +import numpy as np +import ray +import requests +from ray.train import ScalingConfig, RunConfig, CheckpointConfig +from ray.train.torch import TorchTrainer + + +# Repo root — so `src._vllm_compat` resolves on the driver and Ray workers. +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + + +# ────────────────────────────────────────────────────────── +# Knobs — the only things you usually tune +# ────────────────────────────────────────────────────────── +# Teacher parquet always reads from the 10k file (the source of truth produced +# by stage 1). Training subset is N_ROWS — env-overridable so the same script +# + job config can submit a small smoke job (N_ROWS=1000, ~45 min) or the +# full run (N_ROWS=10000, ~6 hr) without code edits. +TEACHER_N_ROWS = int(os.environ.get("TEACHER_N_ROWS", 10_000)) +N_ROWS = int(os.environ.get("N_ROWS", TEACHER_N_ROWS)) +SEED = int(os.environ.get("SEED", 42)) + +BASE_DIR = "/mnt/cluster_storage/vlm-distillation-catalog-enrichment" +TEACHER_PARQUET = f"{BASE_DIR}/teacher_7b_enriched_{TEACHER_N_ROWS}.parquet" +SFT_CACHE_PATH = f"{BASE_DIR}/sft_cache_{N_ROWS}.parquet" +ADAPTER_OUTPUT_DIR = f"{BASE_DIR}/qwen25vl_3b_enrichment_lora_{N_ROWS}" +TRAIN_RUN_DIR = f"{BASE_DIR}/qwen25vl_3b_enrichment_runs_{N_ROWS}" + +STUDENT_MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct" + +# One Ray Train worker per L4 — matches the g6.12xlarge shape (4× L4 24GB). +NUM_WORKERS = 4 + +# SFT hyperparameters. LoRA tolerates a higher LR than full FT because only +# adapter params move; 1e-4 is the standard safe default for r=16. +LEARNING_RATE = 1e-4 +NUM_EPOCHS = 2 +PER_DEVICE_BATCH_SIZE = 1 # bs=1 + grad accum keeps activation memory in check +# Eval batch can be larger than train: no backward → no activation tape, no +# grad checkpointing tax, no optimizer state. With seq=2048 + bf16 model on +# L4 24GB, bs=4 fits with room. Cuts val pass wall-clock by ~4x. +EVAL_PER_DEVICE_BATCH_SIZE = 4 +GRAD_ACCUM_STEPS = 16 # effective batch = 1 × 4 workers × 16 = 64 +WARMUP_RATIO = 0.03 +WEIGHT_DECAY = 0.01 +MAX_SEQ_LEN = 1024 # ~512 visual + ~200 prompt + ≤160 target ≈ 870 + # tokens worst-case; 1024 is a safe ceiling. Lower + # than teacher-batch-label's 2048 max_model_len to + # halve the cross-entropy logits allocation + # (batch × seq × ~152K vocab × bf16) which OOMed + # the loss step on 4× L4. Inference (online + # search) still uses max_model_len=2048 — the + # model's positional encoding is unchanged, only + # the SFT context window is shorter. + +# Mid-epoch checkpoint cadence. 0 = disabled (epoch-only checkpoints). +# Each step checkpoint costs one FSDP FULL_STATE_DICT gather + ~74MB write, +# roughly 5–10s of pause. At 50, that's ~3 step-ckpts per epoch on the 10k-row +# config (156 grad steps), <1% wall-clock overhead, and lets you resume +# closer to a crash than epoch boundaries allow. +SAVE_EVERY_N_STEPS = 50 + +# LoRA config. r=16 is the sweet spot for 3B-scale models on a focused task; +# bump to 32 for harder tasks, drop to 8 for compute-tight runs. +LORA_R = 16 +LORA_ALPHA = 32 +LORA_DROPOUT = 0.05 +LORA_TARGET_MODULES = [ + "q_proj", "k_proj", "v_proj", "o_proj", # attention + "gate_proj", "up_proj", "down_proj", # MLP +] + +# Image processing — same caps the inference-time scripts use, so the model +# trains on exactly the visual feature distribution it will see at serve time. +MIN_PIXELS = 256 * 28 * 28 +MAX_PIXELS = 512 * 28 * 28 +IMAGE_RESIZE = 512 # pre-resize to a fixed square for stable batch shape + +# Deterministic train/val/test split by SHA-1(id) — stable across re-runs. +TRAIN_FRAC = 0.80 +VAL_FRAC = 0.10 +# remainder is test + +# CPU FETCH (HTTP + PIL decode + resize). Network IO bound. +FETCH_TIMEOUT_S = 5.0 +FETCH_CONCURRENCY = 16 + +# CPU SFT EXAMPLE BUILDER (chat template + processor). CPU bound. +BUILD_CONCURRENCY = 8 + + +# ────────────────────────────────────────────────────────── +# STAGE 1 — BUILD SFT TRAINING CACHE (CPU) +# ────────────────────────────────────────────────────────── +# Reads the 7B teacher parquet, drops rows whose teacher output isn't valid +# JSON (bad supervision), assigns a deterministic train/val/test split, +# fetches every image once, and writes a self-contained training parquet at +# SFT_CACHE_PATH. Once cached, every epoch reads from this file — no HTTP. +# https://docs.ray.io/en/latest/data/loading-data.html + +def _split_bucket(row_id: str) -> str: + """Deterministic split by hashing the stable row id.""" + h = int(hashlib.sha1(row_id.encode("utf-8")).hexdigest()[:8], 16) / 0xFFFFFFFF + if h < TRAIN_FRAC: + return "train" + if h < TRAIN_FRAC + VAL_FRAC: + return "val" + return "test" + + +def _strip_code_fence(text: str) -> str: + """The 7B teacher sometimes wraps JSON in ```json ... ``` fences. Strip + them so json.loads succeeds. Same recovery the serve app does.""" + text = (text or "").strip() + if text.startswith("```"): + text = text.split("```", 2)[1] + if text.startswith("json"): + text = text[4:].lstrip("\n") + return text.strip() + + +def _teacher_output_is_valid(row): + """Drop rows where the teacher output isn't a complete JSON object with + all four expected keys. Bad targets pollute the loss.""" + try: + obj = json.loads(_strip_code_fence(row["raw_output"])) + return all(k in obj for k in ("category", "attributes", "search_tags", "description")) + except Exception: + return False + + +def fetch_and_resize_images(batch): + """One HTTP round-trip per image. Resize to a fixed square so all rows + produce identically-shaped processor outputs downstream — that matters + because Ray Data's `iter_batches(numpy)` stacks rows into a single array + per column.""" + from PIL import Image + keep = ("id", "product_id", "title", "image_url", "raw_output", "split") + out = {k: [] for k in keep} + out["image_bytes"] = [] + for i in range(len(batch["id"])): + try: + r = requests.get( + batch["image_url"][i], + timeout=FETCH_TIMEOUT_S, + headers={"User-Agent": "anyscale-finetune/1.0"}, + ) + if r.status_code != 200: + continue + img = Image.open(io.BytesIO(r.content)).convert("RGB").resize( + (IMAGE_RESIZE, IMAGE_RESIZE) + ) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=88) + jpeg = buf.getvalue() + except Exception: + continue + for k in keep: + out[k].append(batch[k][i]) + out["image_bytes"].append(jpeg) + return out + + +def build_sft_cache(): + if os.path.exists(SFT_CACHE_PATH): + cached = ray.data.read_parquet(SFT_CACHE_PATH) + print(f"[cache] reusing {SFT_CACHE_PATH}") + return cached + + if not os.path.exists(TEACHER_PARQUET): + raise FileNotFoundError( + f"Teacher parquet not found at {TEACHER_PARQUET!r}. Run " + f"scripts/run_teacher_batch_label.py first to produce it." + ) + + print(f"[cache] reading teacher parquet {TEACHER_PARQUET}") + ds = ray.data.read_parquet(TEACHER_PARQUET) + ds = ds.filter(_teacher_output_is_valid) + if N_ROWS < TEACHER_N_ROWS: + # Subset for smoke/quick runs. .limit() before image fetch so we don't + # do 10× more network IO than we need. + ds = ds.limit(N_ROWS) + print(f"[cache] limiting to first {N_ROWS} rows (subset of {TEACHER_N_ROWS})") + + def _attach_split(row): + row["split"] = _split_bucket(row["id"]) + return row + + ds = ds.map(_attach_split) + ds = ds.map_batches( + fetch_and_resize_images, + batch_size=16, + concurrency=FETCH_CONCURRENCY, + batch_format="numpy", + ) + ds.write_parquet(SFT_CACHE_PATH) + + cached = ray.data.read_parquet(SFT_CACHE_PATH) + print(f"[cache] wrote {cached.count()} rows → {SFT_CACHE_PATH}") + return cached + + +# ────────────────────────────────────────────────────────── +# STAGE 2 — BUILD SFT EXAMPLES (CPU, autoscaled actor pool) +# ────────────────────────────────────────────────────────── +# One AutoProcessor instance per CPU actor. Each call returns a fully +# tokenized, padded SFT example with loss masked on the user prefix so only +# the assistant's JSON tokens contribute to the loss. Output shapes are +# deterministic — fixed image size + fixed max_length — which is what lets +# Ray Data's `iter_batches(numpy)` stack rows for the trainer. + +VLM_PROMPT = """\ +You are a product catalog enrichment assistant. Given a product image and \ +the merchant-supplied title, output a JSON object with exactly these keys: + + category: one short string (e.g. "Wireless Earbuds") + attributes: a list of 3 short attribute strings + search_tags: a list of 5 short search keywords + description: a single sentence (<= 30 words) + +Title: {title} + +Return ONLY the JSON object, no commentary.\ +""" + + +class BuildSFTExample: + """Tokenize one (image, title, teacher_json) triple into an SFT example. + + Output columns: input_ids, labels, attention_mask, pixel_values, + image_grid_thw, split + + Loss masking: tokens up to the end of the user prefix get label=-100 so + they don't contribute to cross-entropy. Pad tokens are also masked. + """ + + def __init__(self, model_id: str, max_length: int, + min_pixels: int, max_pixels: int): + from transformers import AutoProcessor + self.processor = AutoProcessor.from_pretrained( + model_id, min_pixels=min_pixels, max_pixels=max_pixels + ) + self.max_length = max_length + # Qwen tokenizer has eos but not always pad — fall back to eos for + # padding (with the attention mask doing the right thing). + self.pad_id = ( + self.processor.tokenizer.pad_token_id + or self.processor.tokenizer.eos_token_id + ) + + def __call__(self, row: dict) -> dict: + from PIL import Image + + img = Image.open(io.BytesIO(row["image_bytes"])).convert("RGB") + + user_msg = { + "role": "user", + "content": [ + {"type": "image", "image": img}, + {"type": "text", "text": VLM_PROMPT.format(title=row["title"])}, + ], + } + # Strip the teacher's code fences before training on the JSON. Means + # the student learns to emit clean JSON, not fenced JSON — matches the + # robust-parser the serve app wraps the output in anyway. + target = _strip_code_fence(row["raw_output"]) + asst_msg = { + "role": "assistant", + "content": [{"type": "text", "text": target}], + } + + # User-only template + generation prompt so we know where the + # assistant tokens begin (everything before that gets -100). + user_text = self.processor.apply_chat_template( + [user_msg], tokenize=False, add_generation_prompt=True + ) + full_text = self.processor.apply_chat_template( + [user_msg, asst_msg], tokenize=False + ) + + user_inputs = self.processor( + text=[user_text], images=[img], + padding="max_length", truncation=True, + max_length=self.max_length, return_tensors="pt", + ) + full_inputs = self.processor( + text=[full_text], images=[img], + padding="max_length", truncation=True, + max_length=self.max_length, return_tensors="pt", + ) + + # Boundary = number of non-pad tokens in the user-only prompt. + user_len = int((user_inputs["input_ids"][0] != self.pad_id).sum().item()) + + input_ids = full_inputs["input_ids"][0] + attention_mask = full_inputs["attention_mask"][0] + labels = input_ids.clone() + labels[:user_len] = -100 # mask user prefix + labels[input_ids == self.pad_id] = -100 # mask padding + + return { + "input_ids": input_ids.numpy(), + "labels": labels.numpy(), + "attention_mask": attention_mask.numpy(), + "pixel_values": full_inputs["pixel_values"].numpy(), + "image_grid_thw": full_inputs["image_grid_thw"][0].numpy(), + "split": row["split"], + } + + +# ────────────────────────────────────────────────────────── +# STAGE 3 — TRAIN LOOP (Ray Train worker, FSDP + LoRA) +# ────────────────────────────────────────────────────────── +# One process per GPU. Each worker loads the full Qwen2.5-VL-3B in bf16, +# applies LoRA via PEFT, FSDP-wraps the result, and runs SGD on the data +# shard Ray Train hands it. After training, rank 0 gathers the full state +# and saves the LoRA adapter (config + weights only) to ADAPTER_OUTPUT_DIR. +# https://docs.ray.io/en/latest/train/getting-started-pytorch.html +# https://huggingface.co/docs/peft/main/en/accelerate/fsdp + +def train_loop_per_worker(config: dict): + import functools + import random + import tempfile + import torch + from torch.distributed.fsdp import ( + FullyShardedDataParallel as FSDP, + MixedPrecision, ShardingStrategy, + StateDictType, FullStateDictConfig, + ) + from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy + from transformers import ( + Qwen2_5_VLForConditionalGeneration, + get_cosine_schedule_with_warmup, + ) + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLDecoderLayer, + ) + from peft import LoraConfig, get_peft_model + import ray.train as train + + # ── Distributed context ── + rank = train.get_context().get_world_rank() + world_size = train.get_context().get_world_size() + local_rank = train.get_context().get_local_rank() + device = torch.device(f"cuda:{local_rank}") + + # ── Seed propagation ── + # torch + cuda get a per-rank offset so dropout differs across ranks + # (otherwise FSDP's regularization is undermined — every rank dropping + # the same units is just a uniform scale-down). random/numpy stay + # rank-uniform; they're consumed by CPU-side ops that should be + # deterministic across ranks for the same logical sample. + seed = int(config.get("seed", 42)) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed + rank) + torch.cuda.manual_seed_all(seed + rank) + + if rank == 0: + print(f"[train] world_size={world_size}, model={config['model_id']}, " + f"effective_batch={config['per_device_bs'] * world_size * config['grad_accum']}, " + f"seed={seed}") + + # ── Load student + apply LoRA (BEFORE FSDP wrap) ── + model = Qwen2_5_VLForConditionalGeneration.from_pretrained( + config["model_id"], torch_dtype=torch.bfloat16, + ) + + # Freeze the vision encoder. Vision-tower weights are already strong on + # product photography out-of-the-box; full-multimodal SFT is unstable for + # short runs. The LLM-only LoRA path is the standard VLM SFT recipe. + for p in model.visual.parameters(): + p.requires_grad = False + + lora_cfg = LoraConfig( + r=config["lora_r"], + lora_alpha=config["lora_alpha"], + lora_dropout=config["lora_dropout"], + target_modules=config["lora_target_modules"], + task_type="CAUSAL_LM", + bias="none", + ) + model = get_peft_model(model, lora_cfg) + if rank == 0: + model.print_trainable_parameters() + + # FSDP's FlatParameter requires uniform dtype within each wrap unit, but + # PEFT creates LoRA weights (lora_A, lora_B) in fp32 while the base is + # bf16. Cast trainable adapter params to bf16 so each decoder block has + # uniform dtype. Stable for r=16 SFT at this scale; if you need fp32 + # LoRA, switch FSDP to HSDP/no_shard or use a transformer_auto_wrap_policy + # that pulls adapters into their own wrap unit. + for p in model.parameters(): + if p.requires_grad and p.dtype == torch.float32: + p.data = p.data.to(torch.bfloat16) + + # Required when using gradient checkpointing on a PEFT model: input grads + # need to flow back through the frozen base layers. + model.enable_input_require_grads() + model.gradient_checkpointing_enable({"use_reentrant": False}) + + # ── FSDP wrap ── + # transformer_auto_wrap_policy keyed on Qwen2_5_VLDecoderLayer is the + # PEFT-compatible policy: each decoder layer is its own FSDP unit while + # lm_head, embed_tokens, and the (frozen) vision tower stay replicated. + # Why not size_based: PEFT's tuner wrapper calls the base model via + # `self.model.forward(...)` (peft/tuners/tuners_utils.py:330), bypassing + # __call__ and therefore any FSDP pre-forward unshard hook installed on + # an inner-LLM-level wrap. The symptom is a `size mismatch ... vec (N/2)` + # error inside lm_head where N is the lm_head's full param count. + # Wrapping at the decoder-layer level avoids that path entirely. + auto_wrap_policy = functools.partial( + transformer_auto_wrap_policy, + transformer_layer_cls={Qwen2_5_VLDecoderLayer}, + ) + fsdp_model = FSDP( + model, + auto_wrap_policy=auto_wrap_policy, + mixed_precision=MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + buffer_dtype=torch.bfloat16, + ), + sharding_strategy=ShardingStrategy.FULL_SHARD, + device_id=device, + use_orig_params=True, # required for PEFT param-name preservation + ) + + # ── Optimizer + cosine LR schedule ── + optimizer = torch.optim.AdamW( + [p for p in fsdp_model.parameters() if p.requires_grad], + lr=config["lr"], weight_decay=config["weight_decay"], + ) + + steps_per_epoch = max(1, config["train_size"] // ( + config["per_device_bs"] * world_size * config["grad_accum"] + )) + total_steps = steps_per_epoch * config["num_epochs"] + warmup_steps = max(1, int(total_steps * config["warmup_ratio"])) + scheduler = get_cosine_schedule_with_warmup( + optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps, + ) + + # ── Data shards (Ray Data autoshards across workers) ── + train_shard = train.get_dataset_shard("train") + val_shard = train.get_dataset_shard("val") + + def _to_device(batch: dict) -> dict: + # Each numpy column is shape (bs, ...). Squeeze pixel_values back to + # (bs * num_patches_per_image, patch_dim) since the processor returned + # (1, N, D) per row and we stacked bs of them. + pv = np.stack(batch["pixel_values"]) + if pv.ndim == 4: + pv = pv.reshape(-1, pv.shape[-1]) + return { + "input_ids": torch.tensor(np.stack(batch["input_ids"])).to(device), + "labels": torch.tensor(np.stack(batch["labels"])).to(device), + "attention_mask": torch.tensor(np.stack(batch["attention_mask"])).to(device), + "pixel_values": torch.tensor(pv).to(device, dtype=torch.bfloat16), + "image_grid_thw": torch.tensor(np.stack(batch["image_grid_thw"])).to(device), + } + + # ── Helper: gather FSDP shards on rank 0 and report a Ray Train ckpt ── + # Used both at mid-epoch step cadence and at end-of-epoch. Pulled out so + # the two callsites stay identical (same gather + same report shape). + def _save_checkpoint_and_report(metrics: dict, prefix: str): + ckpt_dir = tempfile.mkdtemp(prefix=prefix) + save_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_cfg): + full_state = fsdp_model.state_dict() + if rank == 0: + fsdp_model.module.save_pretrained( + ckpt_dir, state_dict=full_state, safe_serialization=True, + ) + ckpt = train.Checkpoint.from_directory(ckpt_dir) if rank == 0 else None + train.report(metrics=metrics, checkpoint=ckpt) + + # ── Train ── + fsdp_model.train() + global_step = 0 + save_every = int(config.get("save_every_n_steps", 0)) + # Carry-forward last val_loss into step-checkpoint metrics so retention + # by score_attribute="val_loss" can rank step ckpts vs. epoch ckpts. + # Epoch 0 step ckpts get inf — they'll lose to any post-validation ckpt, + # which is the right preference (no val signal yet = lower confidence). + last_val_loss = float("inf") + for epoch in range(config["num_epochs"]): + optimizer.zero_grad(set_to_none=True) + for step, batch in enumerate(train_shard.iter_batches( + batch_size=config["per_device_bs"], batch_format="numpy", + )): + inputs = _to_device(batch) + outputs = fsdp_model(**inputs) + loss = outputs.loss / config["grad_accum"] + loss.backward() + + if (step + 1) % config["grad_accum"] == 0: + optimizer.step() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + global_step += 1 + if rank == 0 and global_step % 10 == 0: + print(f" step {global_step:>5}/{total_steps} " + f"loss={loss.item() * config['grad_accum']:.4f} " + f"lr={scheduler.get_last_lr()[0]:.2e}") + + # ── Mid-epoch step checkpoint ── + # All ranks must enter the FSDP gather collective together; + # `global_step` is identical on all ranks because Ray Data + # iter_batches is in lockstep, so this branch fires uniformly. + if save_every and global_step % save_every == 0: + step_train_loss = float(loss.item() * config["grad_accum"]) + _save_checkpoint_and_report( + metrics={ + "step": global_step, + "epoch": epoch, + "train_loss": step_train_loss, + "val_loss": last_val_loss, + }, + prefix=f"adapter_step{global_step}_", + ) + if rank == 0: + print(f" [ckpt] step={global_step} " + f"train_loss={step_train_loss:.4f} " + f"(carrying val_loss={last_val_loss:.4f})") + + # ── Validation pass ── + # Eval batch can be larger than train (no backward, no grad + # checkpointing tax). val_loss accumulation is sample-weighted, not + # batch-weighted, so the result is correct regardless of whether the + # last batch is full — important once eval_bs > 1. + fsdp_model.eval() + val_loss_sum, val_count = 0.0, 0 + with torch.no_grad(): + for batch in val_shard.iter_batches( + batch_size=int(config.get("eval_per_device_bs", config["per_device_bs"])), + batch_format="numpy", + ): + inputs = _to_device(batch) + outputs = fsdp_model(**inputs) + bs = inputs["input_ids"].shape[0] + val_loss_sum += outputs.loss.item() * bs # un-mean per-batch + val_count += bs # count samples + + # All-reduce so the reported val_loss is a true global mean across + # ranks, not a per-rank shard mean. Without this, rank 0 and rank 1 + # report different val_loss values for the same epoch. + val_stats = torch.tensor([val_loss_sum, float(val_count)], device=device) + torch.distributed.all_reduce(val_stats, op=torch.distributed.ReduceOp.SUM) + avg_val = (val_stats[0] / torch.clamp(val_stats[1], min=1.0)).item() + last_val_loss = avg_val # carry forward into the next epoch's step ckpts + + # Per-epoch checkpoint via the same helper used by the step path — + # populates result.metrics + result.checkpoint and lets + # CheckpointConfig(score_attribute="val_loss") keep the best epoch. + _save_checkpoint_and_report( + metrics={"epoch": epoch, "step": global_step, "val_loss": avg_val}, + prefix=f"adapter_epoch{epoch}_", + ) + + if rank == 0: + print(f"[epoch {epoch}] val_loss = {avg_val:.4f}") + fsdp_model.train() + + # ── Save LoRA adapter (rank 0 only) ── + # FullStateDictConfig gathers the sharded params on rank 0 only — every + # other rank gets an empty dict, so we only call save_pretrained there. + save_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(fsdp_model, StateDictType.FULL_STATE_DICT, save_cfg): + full_state = fsdp_model.state_dict() + if rank == 0: + peft_model = fsdp_model.module + peft_model.save_pretrained( + config["adapter_dir"], + state_dict=full_state, # PEFT auto-filters to LoRA weights + safe_serialization=True, + ) + print(f"[save] adapter written to {config['adapter_dir']}") + + +# ────────────────────────────────────────────────────────── +# Driver +# ────────────────────────────────────────────────────────── + +def main(): + ray.init( + ignore_reinit_error=True, + runtime_env={ + "worker_process_setup_hook": "src._vllm_compat.patch", + }, + ) + print(f"[run] N_ROWS={N_ROWS} (teacher source has {TEACHER_N_ROWS}) " + f"SEED={SEED} NUM_WORKERS={NUM_WORKERS} " + f"effective_batch={PER_DEVICE_BATCH_SIZE * NUM_WORKERS * GRAD_ACCUM_STEPS}") + print("Cluster resources:", json.dumps(ray.cluster_resources(), indent=2)) + + # STAGE 1 — build / reuse the SFT cache (image bytes + split labels) + cached = build_sft_cache() + + # Materialized counts so the schedule math in train_loop_per_worker is + # right — we avoid calling .count() on a lazy view inside the trainer. + train_size = cached.filter(lambda r: r["split"] == "train").count() + val_size = cached.filter(lambda r: r["split"] == "val").count() + test_size = cached.filter(lambda r: r["split"] == "test").count() + print(f"[split] train={train_size} val={val_size} test={test_size}") + + # STAGE 2 — SFT example builder (one processor instance per CPU actor) + common_kwargs = dict( + fn_constructor_kwargs={ + "model_id": STUDENT_MODEL_ID, + "max_length": MAX_SEQ_LEN, + "min_pixels": MIN_PIXELS, + "max_pixels": MAX_PIXELS, + }, + num_cpus=2, + concurrency=BUILD_CONCURRENCY, + ) + train_ds = cached.filter(lambda r: r["split"] == "train").map(BuildSFTExample, **common_kwargs) + val_ds = cached.filter(lambda r: r["split"] == "val").map(BuildSFTExample, **common_kwargs) + + # STAGE 3 — Ray Train + FSDP + LoRA + trainer = TorchTrainer( + train_loop_per_worker=train_loop_per_worker, + train_loop_config={ + "model_id": STUDENT_MODEL_ID, + "lr": LEARNING_RATE, + "num_epochs": NUM_EPOCHS, + "per_device_bs": PER_DEVICE_BATCH_SIZE, + "grad_accum": GRAD_ACCUM_STEPS, + "warmup_ratio": WARMUP_RATIO, + "weight_decay": WEIGHT_DECAY, + "lora_r": LORA_R, + "lora_alpha": LORA_ALPHA, + "lora_dropout": LORA_DROPOUT, + "lora_target_modules": LORA_TARGET_MODULES, + "train_size": train_size, + "adapter_dir": ADAPTER_OUTPUT_DIR, + "save_every_n_steps": SAVE_EVERY_N_STEPS, + "eval_per_device_bs": EVAL_PER_DEVICE_BATCH_SIZE, + "seed": SEED, + }, + scaling_config=ScalingConfig( + num_workers=NUM_WORKERS, + use_gpu=True, + accelerator_type="L4", + resources_per_worker={"GPU": 1, "CPU": 4}, + ), + run_config=RunConfig( + storage_path=TRAIN_RUN_DIR, + # Keep top-3 by val_loss across both epoch and step checkpoints. + # With NUM_EPOCHS=2 + SAVE_EVERY_N_STEPS=50 + ~156 steps/epoch + # → ~6 step ckpts + 2 epoch ckpts = 8 total candidates; we retain + # the 3 best. Drop num_to_keep if you're storage-tight. + checkpoint_config=CheckpointConfig( + num_to_keep=3, + checkpoint_score_attribute="val_loss", + checkpoint_score_order="min", + ), + ), + datasets={"train": train_ds, "val": val_ds}, + ) + result = trainer.fit() + print(f"\n[done] last metrics: {result.metrics}") + print(f"[done] best checkpoint: {result.checkpoint}") + print(f"[done] LoRA adapter at: {ADAPTER_OUTPUT_DIR}") + + # ── PREVIEW — qualitative side-by-side on the held-out test split ── + # The student outputs aren't generated here (that's an inference-time + # concern); we just print what the teacher said for 3 test products so + # you have a reference to compare against after loading the adapter into + # run_enrich_and_embed.py. + print("\n[preview] held-out test set — teacher targets:") + for r in cached.filter(lambda r: r["split"] == "test").take(3): + print(f"\n title: {r['title'][:90]}") + print(f" teacher: {_strip_code_fence(r['raw_output'])[:240]}") + + print("\n[next] swap the adapter into the inference scripts:") + print(" from peft import PeftModel") + print(" base = Qwen2_5_VLForConditionalGeneration.from_pretrained(STUDENT_MODEL_ID)") + print(f" model = PeftModel.from_pretrained(base, '{ADAPTER_OUTPUT_DIR}')") + + +if __name__ == "__main__": + main() + + +# ────────────────────────────────────────────────────────── +# CHEAT SHEET — Where each stage runs +# ────────────────────────────────────────────────────────── +# +# Stage Runs on Scales via Bottleneck +# ──────────────────── ───────────── ─────────────────────────────── ────────── +# LOAD teacher parquet CPU Ray Data block parallelism cheap +# FILTER bad JSON CPU block parallelism cheap +# CPU FETCH images CPU pool FETCH_CONCURRENCY net IO +# WRITE sft cache CPU block parallelism disk +# BUILD SFT examples CPU pool BUILD_CONCURRENCY processor (PIL + tok) +# TRAIN LOOP GPU pool NUM_WORKERS (one Ray Train L4 FLOPs (heavy) +# worker per L4) +# SAVE adapter GPU rank 0 — cheap +# +# GPU footprint at defaults (g6.12xlarge, 4× L4 24GB): +# NUM_WORKERS=4 × num_gpus=1 → 4 L4s (one full node) +# +# Memory math (per L4): +# Qwen2.5-VL-3B params bf16, FULL_SHARD across 4 GPUs ~1.5 GB +# LoRA trainable params (~50M @ bf16) ~0.1 GB +# AdamW state for LoRA (fp32) ~0.4 GB +# Activations (bs=1, seq=2048, grad checkpoint) ~6–8 GB +# Vision encoder (frozen, full-replicated bf16) ~1.2 GB +# Total ~10 GB / 24 GB +# +# Multi-node / multi-category scale-up: +# - Bump NUM_WORKERS for more GPUs (one ScalingConfig knob). +# - Per-category training: re-run with a different teacher parquet +# (point TEACHER_PARQUET at teacher_7b_enriched_.parquet) and a +# different ADAPTER_OUTPUT_DIR. One base model, one LoRA per category. +# +# Output: +# ADAPTER_OUTPUT_DIR/ +# adapter_config.json ← LoraConfig (r, alpha, target_modules) +# adapter_model.safetensors ← LoRA weights only (~100 MB for r=16) +# +# Loading the adapter at inference time: +# from peft import PeftModel +# base = Qwen2_5_VLForConditionalGeneration.from_pretrained(STUDENT_MODEL_ID) +# model = PeftModel.from_pretrained(base, ADAPTER_OUTPUT_DIR) +# +# Or, with vLLM in the existing batch / serve pipelines, attach the LoRA +# via engine_kwargs: +# LLMConfig( +# ..., +# engine_kwargs={..., "enable_lora": True, +# "lora_modules": [{"name": "enrich", +# "path": ADAPTER_OUTPUT_DIR}]}, +# ) diff --git a/templates/vlm-distillation-catalog-enrichment/scripts/run_enrich_and_embed.py b/templates/vlm-distillation-catalog-enrichment/scripts/run_enrich_and_embed.py new file mode 100644 index 000000000..9f7259bea --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/scripts/run_enrich_and_embed.py @@ -0,0 +1,740 @@ +""" +Ray Data — Multimodal Catalog Pipeline: Qwen2.5-VL Enrichment + SigLIP 2 Embeddings + +One Ray Data graph that produces, for every product, BOTH: + - the structured-attribute JSON a generative VLM emits (category / attributes / + tags / description), and + - dense image AND text vectors in a shared multimodal space (1152-dim, + L2-normalized) ready to push into FAISS / OpenSearch k-NN / Vespa. + +This is the shape of a real production-scale catalog-understanding pipeline: + - Generative VLM enrichment fills the search index (categories, tags, descriptions). + - Dual-tower image+text encoder fills the vector index (visual search, + out-of-stock substitution, listing dedup, "more like this", multimodal recall). + Same images, two parallel outputs, one Ray Data run. + +Pipeline shape: + LOAD → PREPROCESS → URL_CHECK → CPU_FETCH + → VLM_ENRICH (GPU pool, ray.data.llm) ← Qwen2.5-VL-3B + → CPU_PROCESS (CPU pool, SigLIP processor + tok) ← consumes image_bytes + raw_output + → GPU_EMBED (GPU pool, pure inference) ← SigLIP image+text features + → WRITE + +Each comment block below says WHERE the stage runs (CPU or GPU) and links to +the relevant docs. + +Run directly on the workspace cluster: + python scripts/run_enrich_and_embed.py + +Or as an Anyscale Job (re-use job_config.yaml — same g6.12xlarge fits; +flip cpu-workers max_nodes back to a positive number — this pipeline WANTS the +CPU pool): + anyscale job submit --config-file job_config.yaml \\ + --entrypoint "python scripts/run_enrich_and_embed.py" \\ + --env HF_TOKEN=$HF_TOKEN + +NOTE: this is a batch Ray Data pipeline (datasets in → parquet out). It is +NOT a Ray Serve app and cannot be run via `serve run`. To expose both models +as HTTP endpoints — Qwen on /v1/chat/completions, SigLIP on /embed — use a +separate Serve file (not built here). See scripts/run_vlm_online_enrich_3b.py +for the OpenAI-compatible Serve idiom. + +────────────────────────────────────────────────────────── +Why two GPU stages, not one — and why the CPU split matters +────────────────────────────────────────────────────────── +Qwen2.5-VL is a *generative* VLM (autoregressive token decode, hidden states +not contrastively aligned for retrieval, hundreds of ms/image). SigLIP 2 is +a contrastive dual-tower encoder (image + text into the same 1152-dim space, +~thousands of items/sec on one L4). Same images, complementary outputs. + +The CPU/GPU split for the SigLIP path mirrors the Anyscale cross-modal-search +pattern: a CPU `map` actor (Process) does the image processor + text +tokenizer; a GPU `map_batches` actor (Embed) does PURE inference on +pre-tensorized inputs. The L4 spends every cycle on forward-pass FLOPs, never +on PIL or tokenization. Each stage has its own concurrency knob. + + Blog: https://www.anyscale.com/blog/cross-modal-search-for-e-commerce-building-and-scaling-a-cross-modal-image-retrieval-app + Code: https://github.com/anyscale/cross-modal-search-ecommerce-project + SigLIP 2: https://huggingface.co/google/siglip2-so400m-patch14-384 + +Domain-tuned alternative (better for ecommerce-specific recall): + Marqo/marqo-ecommerce-embeddings-L — same dual-tower shape, drop-in swap. + +────────────────────────────────────────────────────────── +What gets optimized for IO and GPU at multi-node scale +────────────────────────────────────────────────────────── + 1. Single image fetch, two consumers. CPU_FETCH decodes, resizes to + 512×512, attaches raw JPEG bytes. The VLM stage (via base64 data URL — + no HTTP refetch) and the SigLIP CPU_PROCESS stage (direct PIL decode) + both read those same bytes. + 2. Resize to a size both models like. 512×512 sits below Qwen2.5-VL's + max_pixels cap (~633×633) so its dynamic-resolution processor uses the + image as-is, and SigLIP's processor downscales to 384 internally. + 3. Cap Qwen2.5-VL vision tokens at ~512/image via mm_processor_kwargs. + Halves prefill work — the dominant cost on the VLM stage. + 4. SigLIP text tower embeds the *enriched* description, not just the + title. We pull "description" out of the VLM's raw_output JSON in + ProcessSigLIP and concat with the title before tokenization. Sparse + merchant titles → richer text vectors → better retrieval recall. + 5. Process and Embed are separate Ray Data stages. Each has its own + concurrency knob (CPU pool vs GPU pool); the GPU actor never touches + PIL, tokenizers, or Python transforms. + 6. Streaming executor pipelines all stages. SigLIP_GPU runs downstream of + VLM (the heavy bottleneck) — when VLM is busy, SigLIP idles between + batches. Both GPU pools autoscale concurrency independently. + 7. Stable SHA-1(title|image_url) row IDs + CheckpointConfig. Job-level + resumability across worker death or cancel-resubmit. Random UUIDs + would regenerate per submission and the checkpoint would never match. +""" + +import os, sys, json, io, base64, hashlib +from typing import Optional + +import numpy as np +import ray +import requests +from huggingface_hub import HfFileSystem +from ray.data.llm import vLLMEngineProcessorConfig, build_processor +from ray.data.checkpoint import CheckpointConfig + + +# Repo root — so `src._vllm_compat` resolves on the driver and inside Ray workers. +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + + +# ────────────────────────────────────────────────────────── +# Knobs — the only things you usually tune +# ────────────────────────────────────────────────────────── +CATEGORY = os.environ.get("CATEGORY", "Electronics") +N_ROWS = int(os.environ.get("N_ROWS", 10_000)) +SEED = int(os.environ.get("SEED", 42)) + +# VLM (generative) — Qwen2.5-VL-3B fits on a single L4 at bf16. +VLM_MODEL_SOURCE = "Qwen/Qwen2.5-VL-3B-Instruct" +VLM_TENSOR_PARALLEL_SIZE = 1 +VLM_PIPELINE_PARALLEL_SIZE = 1 +VLM_MAX_MODEL_LEN = 2048 # image (~512 tok) + prompt (~150) + gen (≤160) ≈ 822 +VLM_BATCH_SIZE = 48 +VLM_CONCURRENCY = 2 # set (min, max) for autoscaling, e.g. (1, 4) + +# Embedding (contrastive dual-tower) — SigLIP 2 so400m, 1152-dim shared space. +EMB_MODEL_SOURCE = "google/siglip2-so400m-patch14-384" +EMB_BATCH_SIZE = 32 # bs=64 fits comfortably on L4 24GB at bf16; 32 is conservative. +EMB_CONCURRENCY = 2 # one actor per L4. Same autoscale story as VLM_CONCURRENCY. + +# CPU FETCH (HTTP + PIL decode + resize). Network IO bound. +IMAGE_RESIZE = 512 +FETCH_TIMEOUT_S = 5.0 +FETCH_CONCURRENCY = 16 + +# CPU PROCESS (SigLIP image processor + text tokenizer). CPU bound. +PROCESS_CONCURRENCY = 8 + +BASE_DIR = "/mnt/cluster_storage/vlm-distillation-catalog-enrichment" +HF_PATH = f"hf://datasets/McAuley-Lab/Amazon-Reviews-2023/raw_meta_{CATEGORY}" +CACHE_PATH = f"{BASE_DIR}/catalog_{CATEGORY}_{N_ROWS}.parquet" +CHECKPOINT_PATH = f"{BASE_DIR}/enc_vlm_emb_enrich_{N_ROWS}_checkpoint" +OUTPUT_PATH = f"{BASE_DIR}/enc_vlm_enriched_with_embeddings_{N_ROWS}.parquet" + + +# ────────────────────────────────────────────────────────── +# Optional LoRA adapter (mirrors run_vlm_online_search_3b.py) +# ────────────────────────────────────────────────────────── +# If the adapter dir exists, route VLM enrichment through the fine-tuned +# adapter via vLLM's LoRA multiplexing; otherwise fall back to the base 3B +# model. Override the location with QWEN_LORA_ADAPTER_DIR. +# +# ray.data.llm batch wiring: +# - dynamic_lora_loading_path (top-level on vLLMEngineProcessorConfig) is the +# S3 prefix that *contains* adapter subfolders. +# - Per-row "model" = "" triggers a LoRARequest; vLLM then +# pulls "//" on first use and +# caches the loaded adapter for the rest of the run. +# - dynamic_lora_loading_path requires a cloud URI (s3://, gs://, ...). When +# the adapter dir is local (under /mnt), we sync it to the workspace's +# ANYSCALE_ARTIFACT_STORAGE bucket once and use that S3 prefix. +LORA_ADAPTER_DIR = os.environ.get( + "QWEN_LORA_ADAPTER_DIR", + None, +) +LORA_MAX_RANK = 16 # must be ≥ LORA_R used during training (run_distill_student_lora.py) + + +def _ensure_lora_remote(local_dir: str) -> Optional[str]: + """Sync a local LoRA dir to the workspace artifact bucket and return the + parent S3 prefix. If `local_dir` is already a cloud URI, return its + parent. Returns None if we can't make a usable cloud prefix. + """ + if local_dir.startswith(("s3://", "gs://", "abfss://", "azure://")): + return os.path.dirname(local_dir.rstrip("/")) + + artifact_root = os.environ.get("ANYSCALE_ARTIFACT_STORAGE") + if not artifact_root: + print( + f"[vlm] LoRA dir {local_dir!r} is local and " + "ANYSCALE_ARTIFACT_STORAGE is not set; can't auto-sync to cloud. " + "Either set QWEN_LORA_ADAPTER_DIR to an s3:// path or run on an " + "Anyscale workspace with artifact storage configured." + ) + return None + + adapter_name = os.path.basename(local_dir.rstrip("/")) + s3_prefix = f"{artifact_root.rstrip('/')}/loras" + s3_dest = f"{s3_prefix}/{adapter_name}" + import subprocess + print(f"[vlm] syncing LoRA {local_dir} → {s3_dest}") + res = subprocess.run( + ["aws", "s3", "sync", local_dir, s3_dest], + capture_output=True, text=True, + ) + if res.returncode != 0: + print(f"[vlm] aws s3 sync failed:\n{res.stderr}") + return None + return s3_prefix + + +_lora_dir = LORA_ADAPTER_DIR.rstrip("/") +_lora_ready = ( + os.path.isdir(_lora_dir) + and os.path.exists(os.path.join(_lora_dir, "adapter_config.json")) +) or _lora_dir.startswith(("s3://", "gs://", "abfss://", "azure://")) + +LORA_REMOTE_PREFIX: Optional[str] = None +LORA_ADAPTER_NAME: Optional[str] = None +if _lora_ready: + LORA_REMOTE_PREFIX = _ensure_lora_remote(_lora_dir) + _lora_ready = LORA_REMOTE_PREFIX is not None + +if _lora_ready: + LORA_ADAPTER_NAME = os.path.basename(_lora_dir.rstrip("/")) + print(f"[vlm] LoRA adapter ready at {LORA_REMOTE_PREFIX}/{LORA_ADAPTER_NAME}; " + f"routing batch enrichment requests with model='{LORA_ADAPTER_NAME}'") +else: + print(f"[vlm] no usable LoRA adapter; falling back to base model='{VLM_MODEL_SOURCE}'") + + +# ────────────────────────────────────────────────────────── +# STAGE 1 — LOAD + STAGE 2 — PREPROCESS (CPU) +# ────────────────────────────────────────────────────────── +# Pulls Amazon-Reviews-2023 metadata from HF, drops rows missing title/image, +# normalizes one row per product, HEAD-checks the URL, caches as parquet. +# Runs entirely on CPU; Ray Data block parallelism does the work. +# https://docs.ray.io/en/latest/data/loading-data.html + +def _extract_image_url(images_field): + if not images_field or not isinstance(images_field, dict): + return None + for key in ("large", "hi_res", "thumb"): + urls = images_field.get(key) + if urls is not None and len(urls) > 0: + return urls[0] + return None + + +def _has_title_and_image(row): + title = row.get("title") + if not (title and title.strip()): + return False + return _extract_image_url(row.get("images")) is not None + + +def _coerce_description(desc_field): + if isinstance(desc_field, list): + return " ".join(str(x) for x in desc_field if x).strip() + if isinstance(desc_field, str): + return desc_field.strip() + return "" + + +def _url_is_reachable(row): + try: + return requests.head(row["image_url"], timeout=FETCH_TIMEOUT_S, allow_redirects=True).ok + except Exception: + return False + + +def _normalize_amazon_row(row): + title = row["title"].strip()[:512] + image_url = _extract_image_url(row["images"]) + row_id = hashlib.sha1(f"{title}|{image_url}".encode("utf-8")).hexdigest()[:16] + return { + "id": row_id, + "product_id": row.get("parent_asin") or row.get("asin") or "", + "title": title, + "description": _coerce_description(row.get("description"))[:1024], + "image_url": image_url, + "source": "amazon-reviews-2023", + } + + +def build_catalog(): + # Cache check: /mnt/cluster_storage/vlm-distillation-catalog-enrichment is per-user persistent, so once we've + # built the catalog at this (CATEGORY, N_ROWS) we never need to redo the + # HF read + URL HEAD checks. Saves ~10 min/run for 10K rows. + if os.path.exists(CACHE_PATH): + cached = ray.data.read_parquet(CACHE_PATH) + print(f"[load hf] reusing cache at {CACHE_PATH}") + return cached + + print(f"[load hf] Loading {HF_PATH} ...") + ds = ray.data.read_parquet( + HF_PATH, + file_extensions=["parquet"], + filesystem=HfFileSystem(), + ) + ds = ds.limit(N_ROWS) + ds = ds.filter(_has_title_and_image) + ds = ds.map(_normalize_amazon_row) + ds = ds.filter(_url_is_reachable) + ds = ds.random_shuffle(seed=SEED) + # mode="overwrite": write_parquet defaults to APPEND (writes new files into + # the directory without removing existing ones). If a stale partial run + # left files behind, the next run would silently mix old + new rows on read + # — exactly the 19,991-row "two runs accumulated" bug we hit before the + # cache_check above was added. + ds.write_parquet(CACHE_PATH, mode="overwrite") + # CRITICAL: count from the materialized parquet, NOT the lazy `ds`. Calling + # ds.count() here would re-execute the entire upstream plan (HF read + + # filters + URL HEAD checks + shuffle) from scratch — wasted ~10 min on + # the first 10K-row run before this fix. read_parquet().count() reads + # parquet row-group metadata, near-instant. + cached = ray.data.read_parquet(CACHE_PATH) + print(f"[load hf] cached {cached.count()} rows → {CACHE_PATH}") + return cached + + +# ────────────────────────────────────────────────────────── +# STAGE 3 — CPU FETCH (CPU, autoscaled actor pool) +# ────────────────────────────────────────────────────────── +# One HTTP round-trip per image. Bytes attached to the row so BOTH downstream +# GPU stages can consume them without re-fetching: +# - VLM stage gets a base64 data URL (built in vlm_preprocess) — vLLM's +# prepare_multimodal_stage decodes the base64 in-place, no HTTP. +# - SigLIP CPU_PROCESS decodes the bytes directly via PIL. + +def fetch_and_resize(batch): + from PIL import Image + out_id, out_pid, out_title, out_desc, out_url, out_src, out_bytes = ( + [], [], [], [], [], [], [], + ) + for i in range(len(batch["id"])): + try: + r = requests.get( + batch["image_url"][i], + timeout=FETCH_TIMEOUT_S, + headers={"User-Agent": "vlm-distillation-catalog-enrichment/1.0"}, + ) + if r.status_code != 200: + continue + img = Image.open(io.BytesIO(r.content)).convert("RGB").resize( + (IMAGE_RESIZE, IMAGE_RESIZE) + ) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=88) + jpeg_bytes = buf.getvalue() + except Exception: + continue + out_id.append(batch["id"][i]) + out_pid.append(batch["product_id"][i]) + out_title.append(batch["title"][i]) + out_desc.append(batch["description"][i]) + out_url.append(batch["image_url"][i]) + out_src.append(batch["source"][i]) + out_bytes.append(jpeg_bytes) + return { + "id": out_id, + "product_id": out_pid, + "title": out_title, + "description": out_desc, + "image_url": out_url, + "source": out_src, + "image_bytes": out_bytes, + } + + +# ────────────────────────────────────────────────────────── +# STAGE 4 — VLM ENRICH (GPU, ray.data.llm processor) +# ────────────────────────────────────────────────────────── +# build_processor builds a multi-stage sub-pipeline under the hood: +# +# PREPROCESS (CPU) → PREPARE_IMAGES (CPU) → ChatTemplate (CPU) +# → Tokenize (CPU) → vLLM Engine (GPU) → Detokenize (CPU) → POSTPROCESS (CPU) +# +# Because we already have image_bytes on the row, we feed vLLM a base64 data +# URL — prepare_multimodal_stage decodes the base64 in-place instead of doing +# an HTTP fetch. One fetch upstream, zero refetches. +# +# CRITICAL: vlm_preprocess and vlm_postprocess must round-trip image_bytes +# through this stage so the downstream SigLIP CPU_PROCESS can read them. +# vLLM only forwards columns we name; everything else is dropped after detokenize. +# https://docs.ray.io/en/latest/data/working-with-llms.html +# https://docs.ray.io/en/latest/data/working-with-llms.html#multimodal + +VLM_PROMPT = """\ +You are a product catalog enrichment assistant. Given a product image and \ +the merchant-supplied title, output a JSON object with exactly these keys: + + category: one short string (e.g. "Wireless Earbuds") + attributes: a list of 3 short attribute strings + search_tags: a list of 5 short search keywords + description: a single sentence (<= 30 words) + +Title: {title} + +Return ONLY the JSON object, no commentary.\ +""" + + +def vlm_preprocess(row): + b64 = base64.b64encode(row["image_bytes"]).decode("ascii") + data_url = f"data:image/jpeg;base64,{b64}" + out = { + "id": row["id"], + # Round-trip these so the SigLIP CPU_PROCESS stage downstream can pick + # them up. vLLM's processor only forwards columns we name in postprocess; + # everything else is dropped after detokenize. + "image_bytes": row["image_bytes"], + "product_id": row["product_id"], + "title": row["title"], + "image_url": row["image_url"], + "source": row["source"], + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_url}}, + {"type": "text", "text": VLM_PROMPT.format(title=row["title"])}, + ], + } + ], + # max_tokens=160: 4-key JSON is ~120 tok typical; 160 is a safe ceiling + # and shrinks per-seq KV reservation vs 256. + "sampling_params": {"max_tokens": 160, "temperature": 0.0}, + } + # When a LoRA adapter is loaded, ray.data.llm routes the request through + # it iff the row's "model" field differs from the processor's + # `model_source` (the base). Set it to the adapter name so vLLM downloads + # `//` once and reuses it. + if LORA_ADAPTER_NAME is not None: + out["model"] = LORA_ADAPTER_NAME + return out + + +def vlm_postprocess(row): + return { + "id": row["id"], + "product_id": row["product_id"], + "title": row["title"], + "image_url": row["image_url"], + "image_bytes": row["image_bytes"], + "source": row["source"], + "raw_output": row["generated_text"], + } + + +def build_vlm_processor(): + engine_kwargs = { + "tensor_parallel_size": VLM_TENSOR_PARALLEL_SIZE, + "pipeline_parallel_size": VLM_PIPELINE_PARALLEL_SIZE, + "max_model_len": VLM_MAX_MODEL_LEN, + "trust_remote_code": True, + # Chunked prefill: image-token prefills are huge (~500 tok/image + # after the cap below). Chunking lets prefill interleave with + # ongoing decode steps in the same forward pass — big throughput + # win for multimodal. + "enable_chunked_prefill": True, + "max_num_batched_tokens": 8192, + "limit_mm_per_prompt": {"image": 1}, + # Cap Qwen2.5-VL's dynamic-resolution processor. Token count + # ≈ pixels / (28×28). 512×(28²) = 401,408 px (~633×633) → ~512 + # vision tokens / image. Halves prefill vs uncapped. + "mm_processor_kwargs": { + "min_pixels": 256 * 28 * 28, + "max_pixels": 512 * 28 * 28, + }, + "gpu_memory_utilization": 0.85, + } + extra_kwargs = {} + if _lora_ready: + engine_kwargs.update({ + "enable_lora": True, + "max_lora_rank": LORA_MAX_RANK, + "max_loras": 1, + }) + # Top-level vLLMEngineProcessorConfig field — the cloud prefix that + # contains adapter subfolders. Per-row "model" picks which subfolder. + extra_kwargs["dynamic_lora_loading_path"] = LORA_REMOTE_PREFIX + + config = vLLMEngineProcessorConfig( + model_source=VLM_MODEL_SOURCE, + engine_kwargs=engine_kwargs, + batch_size=VLM_BATCH_SIZE, + concurrency=VLM_CONCURRENCY, + prepare_multimodal_stage={"enabled": True}, + # vLLM 0.20+ moved TokensPrompt; Ray 2.55's batch LLM stage still + # imports the old path. Patch lives in src/_vllm_compat.py. + runtime_env={"worker_process_setup_hook": "src._vllm_compat.patch"}, + **extra_kwargs, + ) + return build_processor(config, preprocess=vlm_preprocess, postprocess=vlm_postprocess) + + +# ────────────────────────────────────────────────────────── +# STAGE 5 — CPU PROCESS (CPU, autoscaled actor pool) +# ────────────────────────────────────────────────────────── +# SigLIP image processor (PIL → pixel_values) + SigLIP tokenizer (enriched +# text → input_ids). One CPU actor since both operate on the same row. +# Decoupled from the GPU stage on purpose: the GPU actor receives ready-to-go +# tensors and spends all its cycles on forward-pass FLOPs. +# +# One-pipeline payoff: we tokenize "title + VLM description" instead of just +# title. Sparse merchant titles get enriched by the VLM upstream → richer +# text vectors → better retrieval recall on descriptive queries. +# https://huggingface.co/google/siglip2-so400m-patch14-384 + +class ProcessSigLIP: + """CPU-side image+text preprocessing. + + Input columns: image_bytes (jpeg), title, raw_output (VLM JSON string) + Output columns: pixel_values, input_ids + (image_bytes dropped — no need to ship raw bytes to GPU; + raw_output kept so it lands in the final parquet) + """ + + def __init__(self, model_id: str): + from transformers import AutoProcessor + + self.processor = AutoProcessor.from_pretrained(model_id) + + @staticmethod + def _compose_text(title: str, raw_output: str) -> str: + # Embed title + VLM description in the text tower. The VLM enrichment + # makes the embedding stage's text side richer, which improves recall + # for descriptive queries ("noise-cancelling over-ear") that wouldn't + # match a sparse merchant title. This is the one-pipeline payoff. + try: + obj = json.loads(raw_output) + desc = obj.get("description") or "" + return f"{title}. {desc}".strip(". ").strip() if desc else title + except Exception: + return title + + def __call__(self, row: dict) -> dict: + from PIL import Image + + img = Image.open(io.BytesIO(row["image_bytes"])).convert("RGB") + img_inputs = self.processor(images=img, return_tensors="pt") + # SigLIP 2 was trained with fixed 64-token text inputs; padding="max_length" + # is what the model expects, NOT padding=True. Match it or quality drops. + # NOTE: SigLIP's tokenizer always pads to max_length and the model treats + # every token as valid — the processor does NOT return an attention_mask, + # and get_text_features doesn't need one. Don't try to grab one. + text = self._compose_text(row["title"], row.get("raw_output") or "") + txt_inputs = self.processor( + text=text, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + + row["pixel_values"] = img_inputs["pixel_values"][0].numpy() # (3, 384, 384) + row["input_ids"] = txt_inputs["input_ids"][0].numpy() # (64,) + # Drop raw bytes — the GPU stage doesn't need them, and shipping them + # across the streaming executor is just memory pressure. + row.pop("image_bytes", None) + return row + + +# ────────────────────────────────────────────────────────── +# STAGE 6 — GPU EMBED (GPU, actor pool, num_gpus=1) +# ────────────────────────────────────────────────────────── +# Pure inference. No PIL, no tokenizer, no Python transforms — just two +# forward passes on pre-tensorized inputs. Loads model ONCE per actor on init. +# +# - Image tower: pixel_values → get_image_features → 1152-dim vector +# - Text tower: input_ids → get_text_features → 1152-dim vector +# (SigLIP doesn't take an attention_mask — fixed-length padded inputs) +# +# Both vectors L2-normalized so cosine similarity = dot product downstream. + +class EmbedSigLIP: + """GPU forward-pass actor. Pure inference.""" + + def __init__(self, model_id: str): + import torch + from transformers import AutoModel + + self._torch = torch + self.device = "cuda" if torch.cuda.is_available() else "cpu" + # bf16: SigLIP 2 trains and ships in bf16. ~1.8 GB weights. + self.model = AutoModel.from_pretrained( + model_id, torch_dtype=torch.bfloat16 + ).to(self.device).eval() + print(f" [embed] loaded {model_id} on {self.device} (bf16)") + + def __call__(self, batch: dict) -> dict: + torch = self._torch + + pixel_values = torch.tensor( + np.stack(batch["pixel_values"]), + device=self.device, + dtype=torch.bfloat16, + ) + input_ids = torch.tensor(np.stack(batch["input_ids"]), device=self.device) + + with torch.inference_mode(): + img_out = self.model.get_image_features(pixel_values=pixel_values) + txt_out = self.model.get_text_features(input_ids=input_ids) + + # SigLIP2 returns BaseModelOutputWithPooling on some transformers + # versions; older SigLIP returned a tensor. Unwrap defensively. + if hasattr(img_out, "pooler_output"): + img_out = img_out.pooler_output + if hasattr(txt_out, "pooler_output"): + txt_out = txt_out.pooler_output + + img_feat = torch.nn.functional.normalize(img_out, dim=-1).float().cpu().numpy() + txt_feat = torch.nn.functional.normalize(txt_out, dim=-1).float().cpu().numpy() + + batch["image_embedding"] = list(img_feat) + batch["text_embedding"] = list(txt_feat) + # Drop intermediate tensors — we don't want them in the output parquet. + for k in ("pixel_values", "input_ids"): + batch.pop(k, None) + return batch + + +# ────────────────────────────────────────────────────────── +# Driver +# ────────────────────────────────────────────────────────── + +def main(): + ray.init( + ignore_reinit_error=True, + runtime_env={ + "worker_process_setup_hook": "src._vllm_compat.patch", + }, + ) + print("Cluster resources:", json.dumps(ray.cluster_resources(), indent=2)) + + # STAGE 1+2 — load + preprocess + # NB: CheckpointConfig is set AFTER build_catalog (further down) on + # purpose. Setting it before would cause build_catalog's write_parquet to + # populate the checkpoint with all catalog row IDs — and then the + # downstream inference pipeline would think those IDs are already done + # and skip 99%+ of the rows. (Caught a 1-row-out-of-9998 run this way.) + # The catalog has its own resume mechanism: the cache parquet itself. + ds = build_catalog() + + # Job-level checkpointing for the inference pipeline ONLY — resumes after + # worker death / cancel-resubmit. https://docs.anyscale.com/runtime/data + ctx = ray.data.DataContext.get_current() + ctx.checkpoint_config = CheckpointConfig( + id_column="id", + checkpoint_path=CHECKPOINT_PATH, + delete_checkpoint_on_success=False, + ) + + # STAGE 3 — CPU fetch + decode + resize + print(f"\n[fetch+resize] CPU pool, concurrency={FETCH_CONCURRENCY}, target {IMAGE_RESIZE}px") + ds = ds.map_batches( + fetch_and_resize, + batch_size=16, + concurrency=FETCH_CONCURRENCY, + batch_format="numpy", + ) + + # STAGE 4 — VLM enrichment via ray.data.llm + print(f"\n[vlm enrich] GPU pool, concurrency={VLM_CONCURRENCY}, model={VLM_MODEL_SOURCE}") + ds = build_vlm_processor()(ds) + + # STAGE 5 — CPU process (SigLIP image processor + text tokenizer) + print(f"\n[siglip process] CPU pool, concurrency={PROCESS_CONCURRENCY}, model={EMB_MODEL_SOURCE}") + ds = ds.map( + ProcessSigLIP, + fn_constructor_kwargs={"model_id": EMB_MODEL_SOURCE}, + num_cpus=1, + concurrency=PROCESS_CONCURRENCY, + ) + + # STAGE 6 — GPU embed (pure inference) + print(f"\n[siglip embed] GPU pool, concurrency={EMB_CONCURRENCY}, batch={EMB_BATCH_SIZE}") + ds = ds.map_batches( + EmbedSigLIP, + fn_constructor_kwargs={"model_id": EMB_MODEL_SOURCE}, + batch_size=EMB_BATCH_SIZE, + num_gpus=1, + concurrency=EMB_CONCURRENCY, + batch_format="numpy", + ) + + # STAGE 7 — write parquet (the sink that triggers checkpointing) + # materialize() pins the result blocks in the object store so both + # write_parquet and the preview .take() below consume the same in-memory + # dataset — no disk round-trip, no re-execution of the GPU pipeline. + # https://docs.ray.io/en/latest/data/saving-data.html + ds = ds.materialize() + ds.write_parquet(OUTPUT_PATH) + print(f"\n[done] wrote enriched + embedded catalog to {OUTPUT_PATH}") + + # ── PREVIEW — sample rows + in-driver retrieval demo ── + # Stand-in for what FAISS / OpenSearch k-NN / Vespa do at scale. + # + # SigLIP raw-cosine regime (the model applies its own learned + # logit_scale ≈ 100 / bias ≈ -10 internally, so raw cos clusters near 0): + # image → image : relevant pairs at ~0.6+ + # text → image : relevant pairs at ~0.10–0.20; <0.05 = noise + print(f"\n[schema]\n{ds.schema()}") + print(f"\n[count] {ds.count()} rows") + print(f"\n[row 0]\n{ds.take(1)[0]}") + print(f"\n[checkpoint] manifest at {CHECKPOINT_PATH}") + + +if __name__ == "__main__": + main() + + +# ────────────────────────────────────────────────────────── +# CHEAT SHEET — Where each stage runs +# ────────────────────────────────────────────────────────── +# +# Stage Runs on Scales via Bottleneck +# ──────────────────── ───────── ────────────────────────────── ────────── +# LOAD (read_parquet) CPU Ray Data block parallelism HF CDN +# PREPROCESS CPU block parallelism cheap +# URL_CHECK CPU block parallelism net IO +# CPU_FETCH CPU pool FETCH_CONCURRENCY net IO +# VLM_ENRICH GPU pool VLM_CONCURRENCY (TP×PP=1 each) L4 FLOPs (heavy) +# CPU_PROCESS CPU pool PROCESS_CONCURRENCY CPU +# GPU_EMBED GPU pool EMB_CONCURRENCY (1 GPU each) L4 FLOPs (light) +# WRITE (parquet) CPU block parallelism disk +# +# GPU footprint at defaults (g6.12xlarge, 4× L4 24 GB): +# VLM_CONCURRENCY=2 × (TP=1, PP=1) → 2 L4s +# EMB_CONCURRENCY=2 × 1 GPU → 2 L4s +# Total = 4 L4s ✓ (one node) +# +# Multi-node scale-up: +# - Bump VLM_CONCURRENCY / EMB_CONCURRENCY to (min, max) tuples for autoscale. +# - Each GPU replica is independent; Ray Data shards blocks across them. +# - GPU_EMBED is the cheaper stage (~thousands of items/sec/L4); usually +# idles between VLM batches in the streaming executor — that's the +# pipeline doing the right thing, not a bug. Don't optimize a stage +# that's already idle. +# +# Output schema (one row per product): +# id, product_id, title, image_url, source, +# raw_output ← VLM enrichment JSON string +# image_embedding[1152] ← SigLIP image tower (L2-normalized bf16→fp32) +# text_embedding[1152] ← SigLIP text tower (L2-normalized bf16→fp32) +# +# Drop-in alternatives you can swap with one line: +# EMB_MODEL_SOURCE = "Marqo/marqo-ecommerce-embeddings-L" # ecommerce-tuned +# VLM_MODEL_SOURCE = "Qwen/Qwen2.5-VL-7B-Instruct" # better quality, more GPU +# +# This is a BATCH pipeline. To expose both models as HTTP endpoints: +# - Build a separate Serve app (e.g. scripts/run_vlm_online_enrich_emb_dual.py) +# - Use ray.serve.llm.LLMConfig + build_openai_app for Qwen +# - Add a Serve deployment for SigLIP that takes (image_url|image_bytes, text) +# and returns {"image_embedding": [...], "text_embedding": [...]} diff --git a/templates/vlm-distillation-catalog-enrichment/scripts/run_teacher_batch_label.py b/templates/vlm-distillation-catalog-enrichment/scripts/run_teacher_batch_label.py new file mode 100644 index 000000000..c2dce13cd --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/scripts/run_teacher_batch_label.py @@ -0,0 +1,390 @@ +""" +Stage 1 — Teacher batch enrichment with Qwen2.5-VL-7B. + +Reads an Amazon-Reviews-2023 product subset, prompts Qwen2.5-VL-7B with the +image + title, and writes a parquet of {category, attributes, search_tags, +description} JSON per row. The output is the labeled corpus that Stage 2 +(run_distill_student_lora.py) distills into a 3B student. + +Pipeline: + LOAD → PREPROCESS → PREPARE_IMAGES → INFER (vLLM) → POSTPROCESS → WRITE + +Each replica runs on a single L4 GPU (TP=1). With CONCURRENCY=4 you saturate +all four GPUs on a g6.12xlarge node, which gives roughly 4× the throughput of +a single TP=4 replica without any all-reduce overhead. + +Run directly on the workspace cluster: + python scripts/run_teacher_batch_label.py + +Or as an Anyscale Job: + anyscale job submit --config-file job_config.yaml --env HF_TOKEN=$HF_TOKEN +""" + +import os, sys, json, hashlib +import ray +import requests +from huggingface_hub import HfFileSystem +from ray.data.llm import vLLMEngineProcessorConfig, build_processor +from ray.data.checkpoint import CheckpointConfig + + +# Repo root — so `src._vllm_compat` resolves on the driver and inside Ray workers. +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, REPO_ROOT) + + +# ────────────────────────────────────────────────────────── +# Knobs — the only things you usually tune +# ────────────────────────────────────────────────────────── +CATEGORY = os.environ.get("CATEGORY", "Electronics") +N_ROWS = int(os.environ.get("N_ROWS", 10_000)) +SEED = int(os.environ.get("SEED", 42)) + +# Qwen2.5-VL ships in 3B / 7B / 32B / 72B. 7B in bf16 (~14GB) fits on a +# single L4 (24GB), so we shard via replicas — one model copy per GPU — +# rather than tensor parallelism. +MODEL_SOURCE = "Qwen/Qwen2.5-VL-7B-Instruct" + +# tensor_parallel_size: shard weights across N GPUs on ONE node (intra-node). +# pipeline_parallel_size: split layers across M nodes (inter-node). +# Each replica claims TP × PP GPUs total. The 7B fits on one L4, so TP=1 + +# CONCURRENCY=4 beats TP=4 + CONCURRENCY=1: every GPU runs an independent +# replica every step, with no all-reduce overhead between them. +# https://docs.vllm.ai/en/stable/serving/distributed_serving.html +TENSOR_PARALLEL_SIZE = 1 +PIPELINE_PARALLEL_SIZE = 1 + +MAX_MODEL_LEN = 2048 # prompt + image tokens + generation; raise if truncating. + # With max_pixels capped (see mm_processor_kwargs below), + # image ≈ 512 tok + prompt ≈ 150 tok + gen ≤ 160 tok ≈ 822. + # 2048 leaves headroom; halves per-seq KV reservation vs 4096. +BATCH_SIZE = 16 # rows fed to each replica per step; smaller for VLMs + # because image tokens explode the per-row token count. +CONCURRENCY = 4 # one replica per L4 on a g6.12xlarge (4× L4) node. + # Set (min, max) for autoscaling, e.g. (1, 4). + +BASE_DIR = "/mnt/cluster_storage/vlm-distillation-catalog-enrichment" # per-user persistent NFS — survives cluster + # teardown AND is private to the submitting user. + # Use /mnt/cluster_storage/vlm-distillation-catalog-enrichment instead if teammates + # in the same project need to read these outputs. +HF_PATH = f"hf://datasets/McAuley-Lab/Amazon-Reviews-2023/raw_meta_{CATEGORY}" +CACHE_PATH = f"{BASE_DIR}/catalog_{CATEGORY}_{N_ROWS}.parquet" +CHECKPOINT_PATH = f"{BASE_DIR}/teacher_7b_{N_ROWS}_checkpoint" +OUTPUT_PATH = f"{BASE_DIR}/teacher_7b_enriched_{N_ROWS}.parquet" + + +# ────────────────────────────────────────────────────────── +# STAGE 1 — LOAD + STAGE 2 — PREPROCESS (CPU) +# ────────────────────────────────────────────────────────── +# Pulls Amazon-Reviews-2023 metadata from HF, drops rows missing title/image, +# normalizes one row per product, and caches as parquet on cluster storage. +# Runs entirely on CPU; Ray Data block parallelism does the work. +# https://docs.ray.io/en/latest/data/loading-data.html + +def _extract_image_url(images_field): + if not images_field or not isinstance(images_field, dict): + return None + for key in ("large", "hi_res", "thumb"): + urls = images_field.get(key) + if urls is not None and len(urls) > 0: + return urls[0] + return None + + +def _has_title_and_image(row): + title = row.get("title") + if not (title and title.strip()): + return False + return _extract_image_url(row.get("images")) is not None + + +def _coerce_description(desc_field): + if isinstance(desc_field, list): + return " ".join(str(x) for x in desc_field if x).strip() + if isinstance(desc_field, str): + return desc_field.strip() + return "" + + +def _url_is_reachable(row): + # Pre-fetch HEAD check — Amazon delists product images periodically, and + # one 404 inside prepare_multimodal_stage aborts the whole run (the stage + # has no row-level fault tolerance). Cheaper to drop here than retry-fight + # vLLM. Parallelized via Ray Data block parallelism on the CPU pool. + try: + return requests.head(row["image_url"], timeout=5, allow_redirects=True).ok + except Exception: + return False + + +def _normalize_amazon_row_to_image(row): + # Stable per-(product, image) ID: SHA-1 of title + image_url. Deterministic + # across runs so the Ray Data CheckpointConfig can actually resume — + # uuid.uuid4() regenerated fresh IDs each submission, making the checkpoint + # match zero rows. Hashing on (title, image_url) instead of product_id keeps + # the ID unique if we ever fan a single product out into one row per image. + title = row["title"].strip()[:512] + image_url = _extract_image_url(row["images"]) + row_id = hashlib.sha1(f"{title}|{image_url}".encode("utf-8")).hexdigest()[:16] + return { + "id": row_id, + "product_id": row.get("parent_asin") or row.get("asin") or "", + "title": title, + "description": _coerce_description(row.get("description"))[:1024], + "image_url": image_url, + "source": "amazon-reviews-2023", + } + + +def build_catalog(): + print(f"[load hf] Loading data from huggingface hub: {HF_PATH}...") + ds = ray.data.read_parquet( + HF_PATH, + file_extensions=["parquet"], + filesystem=HfFileSystem(), + ) + print("[load hf] count:", ds.count()) + print("[load hf] original schema:", ds.schema()) + + ds = ds.limit(N_ROWS) + ds = ds.filter(_has_title_and_image) + ds = ds.map(_normalize_amazon_row_to_image) + ds = ds.filter(_url_is_reachable) + ds = ds.random_shuffle(seed=SEED) + ds.write_parquet(CACHE_PATH) + return ray.data.read_parquet(CACHE_PATH) + + +# ────────────────────────────────────────────────────────── +# STAGE 3 — VLM INFERENCE pipeline +# ────────────────────────────────────────────────────────── +# build_processor builds a multi-stage sub-pipeline under the hood: +# +# PREPROCESS (CPU) → PREPARE_IMAGES (CPU) → ChatTemplate (CPU) +# → Tokenize (CPU) → vLLM Engine (GPU) → Detokenize (CPU) → POSTPROCESS (CPU) +# +# Each inner stage is its own auto-scaled actor pool. The PREPARE_IMAGES +# stage is the multimodal-specific one — it fetches every image_url and +# decodes to a PIL.Image so the GPU stage doesn't pay download latency. +# https://docs.ray.io/en/latest/data/working-with-llms.html +# https://docs.ray.io/en/latest/data/working-with-llms.html#multimodal + +PROMPT = """\ +You are a product catalog enrichment assistant. Given a product image and \ +the merchant-supplied title, output a JSON object with exactly these keys: + + category: one short string (e.g. "Wireless Earbuds") + attributes: a list of 3 short attribute strings + search_tags: a list of 5 short search keywords + description: a single sentence (<= 30 words) + +Title: {title} + +Return ONLY the JSON object, no commentary.\ +""" + + +def build_messages_url(url, title): + """OpenAI-spec multimodal message: image_url block + text block. + + Ray's prepare_multimodal_stage will fetch the URL on a CPU worker and + replace the content block with the decoded image before it hits the GPU. + """ + return [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": url}}, + {"type": "text", "text": PROMPT.format(title=title)}, + ], + } + ] + + +# ── PREPROCESS — (CPU, autoscaled actor pool) ── +# Builds the OpenAI-style messages + per-row sampling params. Keep this lean: +# anything you return travels through every downstream stage. +# https://docs.ray.io/en/latest/data/working-with-llms.html#preprocess +def vlm_preprocess(row): + return { + "id": row["id"], + "messages": build_messages_url(row["image_url"], row["title"]), + # max_tokens=160: output is a 4-key JSON (~120 tok typical, ~150 worst- + # case). 160 is a safe ceiling and shrinks per-seq KV reservation vs 256. + "sampling_params": {"max_tokens": 160, "temperature": 0.0}, + } + + +# ── POSTPROCESS — (CPU, autoscaled actor pool) ── +# Project to just the columns we want in the output parquet. Don't `**row` +# here — it'd carry vLLM-internal columns (prompt_token_ids, MediaWithBytes, +# generated_tokens, …) into the sink and bloat the file. Keep `id` so the +# checkpoint id_column lookup still works. +# https://docs.ray.io/en/latest/data/working-with-llms.html#postprocess +def vlm_postprocess(row): + return { + "id": row["id"], + "product_id": row["product_id"], + "title": row["title"], + "image_url": row["image_url"], + "source": row["source"], + "raw_output": row["generated_text"], + } + + +def run_inference(ds): + # ── Job-level checkpointing (Anyscale-only feature) ── + # Lets the job resume from where it left off if a worker dies or you + # cancel + resubmit. id_column must be a stable string per row. + # delete_checkpoint_on_success=False keeps the checkpoint after a clean + # finish — handy during development; flip to True in production to free + # the storage automatically. + # https://docs.anyscale.com/runtime/data + ctx = ray.data.DataContext.get_current() + ctx.checkpoint_config = CheckpointConfig( + id_column="id", + checkpoint_path=CHECKPOINT_PATH, + delete_checkpoint_on_success=False, + ) + + config = vLLMEngineProcessorConfig( + model_source=MODEL_SOURCE, + engine_kwargs={ + # TP shards the model across GPUs on ONE node (intra-node, NCCL). + # PP splits LAYERS across nodes (inter-node, slower link OK). + # Total GPUs per replica = TP × PP. + "tensor_parallel_size": TENSOR_PARALLEL_SIZE, + "pipeline_parallel_size": PIPELINE_PARALLEL_SIZE, + "max_model_len": MAX_MODEL_LEN, + "trust_remote_code": True, + # Chunked prefill: break long prefills (image tokens are huge — + # ~1k+ per image for Qwen2.5-VL) into chunks that interleave + # with ongoing decode steps in the same forward pass. Without + # this, one image's prefill stalls every other request in the + # running batch. Big throughput win for multimodal. + "enable_chunked_prefill": True, + # Per-step compute budget — total tokens vLLM processes in one + # forward pass (prefill chunks + decodes combined). When chunked + # prefill is on, this also caps the chunk size. 8192 is a good + # middle ground for image-heavy prompts; raise for higher + # throughput / lower for tighter per-step latency. + "max_num_batched_tokens": 8192, + # NB: continuous batching is the default vLLM scheduler — you're + # already getting it (admit-on-arrival, finished-requests free + # their slot mid-step). No flag needed. + # Cap images per prompt so vLLM's mm scheduler can size buffers. + # We send exactly 1 image per row in this pipeline. + "limit_mm_per_prompt": {"image": 1}, + # Cap Qwen2.5-VL's dynamic-resolution image processor. Token count + # ≈ pixels / (28×28). At Amazon's `large` image size (~1024²) the + # processor emits ~1,300 vision tokens per row; capping at + # 512×(28²) = 401,408 px (~633×633) bounds it to ~512 tokens → + # roughly halves prefill work with no visible quality loss for + # centered product photos. Floor of 256×(28²) avoids tiny + # thumbnails getting upsampled into wasted tokens. + "mm_processor_kwargs": { + "min_pixels": 256 * 28 * 28, + "max_pixels": 512 * 28 * 28, + }, + # "should_continue_on_error": True, + # ↑ row-level fault tolerance: skip a bad image / corrupt + # prompt instead of failing the whole batch. Off here so dev + # bugs surface loudly; turn on for prod. + }, + batch_size=BATCH_SIZE, + concurrency=CONCURRENCY, + # Multimodal-specific stage. Without this, the engine receives raw + # image_url strings and the GPU actor would be doing HTTP fetches — + # huge GPU-idle time. With it, fetch + decode is offloaded to CPU + # workers and pipelined alongside inference. + prepare_multimodal_stage={"enabled": True}, + # Setup hook re-run on every engine actor: vLLM 0.20+ moved + # TokensPrompt; Ray 2.55's batch LLM stage still imports the old + # path. Patch lives in src/_vllm_compat.py. + runtime_env={"worker_process_setup_hook": "src._vllm_compat.patch"}, + ) + + vlm_processor = build_processor( + config, + preprocess=vlm_preprocess, + postprocess=vlm_postprocess, + ) + return vlm_processor(ds) + + +# ────────────────────────────────────────────────────────── +# Driver +# ────────────────────────────────────────────────────────── + +def main(): + # ray.init connects to the cluster Anyscale provisioned. We only add the + # vllm setup hook here — the Job's runtime_env already supplies working_dir + # (auto-uploaded from the yaml's `working_dir: .`) and env_vars (HF_TOKEN + # from yaml). Passing those again here triggers a "Failed to merge runtime + # env" conflict, since Ray won't merge overlapping keys. + # + # HF_TOKEN: NOT threaded through here. huggingface_hub / transformers / + # datasets auto-pick from os.environ. Same pattern as the official + # Megatron example: + # https://github.com/anyscale/examples/blob/main/megatron_training/llm_sft_ray_train_megatron.py + ray.init( + ignore_reinit_error=True, + runtime_env={ + "worker_process_setup_hook": "src._vllm_compat.patch", + }, + ) + print("Cluster resources:", json.dumps(ray.cluster_resources(), indent=2)) + + ds = build_catalog() # STAGE 1+2 — load + preprocess + ds = run_inference(ds) # STAGE 3 — multimodal inference + + # ── STAGE 4 — WRITE (CPU, the "sink" that triggers checkpointing) ── + # https://docs.ray.io/en/latest/data/saving-data.html + ds.write_parquet(OUTPUT_PATH) + print(f"[done] wrote enriched output to {OUTPUT_PATH}") + + # ── PREVIEW — print a few enriched rows to the job log ── + # Reads from the just-written parquet (NOT the lazy `ds`) so we don't + # re-execute the GPU pipeline for the preview. If the cluster's + # /mnt/cluster_storage/vlm-distillation-catalog-enrichment gets blown away after the job, at least the + # sample rows live forever in the captured job logs. + print("\n[preview] sample enriched rows:") + for row in ray.data.read_parquet(OUTPUT_PATH).take(limit=4): + print(json.dumps(row, indent=2, default=str)) + + print(ds.stats()) + + +if __name__ == "__main__": + main() + + +# ────────────────────────────────────────────────────────── +# CHEAT SHEET — Where each stage runs +# ────────────────────────────────────────────────────────── +# +# Stage Runs on Scales via +# ──────────────────── ───────── ────────────────────────── +# LOAD (read_parquet) CPU Ray Data block parallelism +# PREPROCESS CPU Autoscaled actor pool +# PREPARE_IMAGES CPU ┐ Multimodal-only: HTTP fetch + PIL decode. +# Set prepare_multimodal_stage={"enabled":True}. +# ChatTemplate CPU │ +# Tokenize CPU ├── Built-in stages inside build_processor; +# vLLM Engine GPU │ each is its own actor pool. +# Detokenize CPU │ +# POSTPROCESS CPU ┘ +# WRITE (parquet) CPU Ray Data block parallelism +# +# Sizing the engine replicas: +# GPUs per replica = tensor_parallel_size × pipeline_parallel_size +# total replicas = concurrency (or autoscales between (min, max)) +# total GPUs used = (TP × PP) × concurrency +# +# This script: TP=1, PP=1, concurrency=4 → 4 GPUs total → fits on one +# g6.12xlarge (4× L4), one 7B replica per GPU. +# +# Common scale-ups: +# 7B-8B on 1 node: TP=1 PP=1, concurrency=N (one replica per GPU) +# 32B-70B on 1 node: TP=4 PP=1, concurrency=1 (or (1, num_nodes)) +# 70B+ multi-node: TP=4 PP=2, concurrency=1 (one replica spans 2 nodes) diff --git a/templates/vlm-distillation-catalog-enrichment/src/__init__.py b/templates/vlm-distillation-catalog-enrichment/src/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/templates/vlm-distillation-catalog-enrichment/src/_vllm_compat.py b/templates/vlm-distillation-catalog-enrichment/src/_vllm_compat.py new file mode 100644 index 000000000..cfba73736 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/_vllm_compat.py @@ -0,0 +1,17 @@ +"""Compat shim for vLLM ↔ Ray LLM batch stage. + +vLLM 0.20+ moved ``TokensPrompt`` and ``TextPrompt`` up from +``vllm.inputs.data`` to ``vllm.inputs`` directly. Ray 2.55's +``vllm_engine_stage.py`` still references the old path, so we alias +``vllm.inputs.data`` back to ``vllm.inputs`` if missing. + +Wired in via ``vLLMEngineProcessorConfig.runtime_env={"worker_process_setup_hook": "src._vllm_compat.patch"}`` +so it runs once per LLM worker on startup, before the stage UDF runs. +""" + + +def patch() -> None: + import vllm.inputs + + if not hasattr(vllm.inputs, "data"): + vllm.inputs.data = vllm.inputs diff --git a/templates/vlm-distillation-catalog-enrichment/src/enrich.py b/templates/vlm-distillation-catalog-enrichment/src/enrich.py new file mode 100644 index 000000000..9f71307d5 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/enrich.py @@ -0,0 +1,177 @@ +""" +GPU enrichment stage — Qwen2.5-VL-3B-Instruct via vLLM. + +Two paths, mirroring `notebooks/demo_walkthrough.ipynb`: + + 1. NaiveVLMEnricher (callable class) + Single-actor approach: each GPU actor does HTTP fetch + decode + VLM + generate inline. The GPU sits idle while images download. The "before". + + 2. build_heterogeneous_processor (ray.data.llm) + Decoupled approach: messages carry the image as an OpenAI-spec base64 + data URL (Arrow-native, no PIL → pickle fallback across stages). + CPU stage upstream produces ``image_bytes``; this stage runs only the + VLM via ray.data.llm + vLLM. The GPU stays saturated. The "after". +""" +import base64 +import io + +from PIL import Image + + +MODEL_ID = "Qwen/Qwen2.5-VL-3B-Instruct" + +ENRICHMENT_PROMPT = """\ +You are a product catalog enrichment assistant. Given a product image and \ +the merchant-supplied title, output a JSON object with exactly these keys: + + category: one short string (e.g. "Wireless Earbuds") + attributes: a list of 3 short attribute strings + search_tags: a list of 5 short search keywords + description: a single sentence (<= 30 words) + +Title: {title} + +Return ONLY the JSON object, no commentary.\ +""" + + +def build_messages(image: Image.Image, title: str) -> list: + """Naive variant: PIL.Image inline. Used inside a single actor only — + the message dict never crosses Ray Data block boundaries, so PIL is fine.""" + return [ + { + "role": "user", + "content": [ + {"type": "image", "image": image}, + {"type": "text", "text": ENRICHMENT_PROMPT.format(title=title)}, + ], + } + ] + + +def build_messages_url(data_url: str, title: str) -> list: + """Hetero variant: image as an OpenAI-spec base64 data URL string. + + Strings serialize zero-copy through Arrow across operator boundaries — + no PIL → pickle fallback warnings, no per-stage memory blow-up. + """ + return [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_url}}, + {"type": "text", "text": ENRICHMENT_PROMPT.format(title=title)}, + ], + } + ] + + +# --------------------------------------------------------------------------- +# Path 1 — NAIVE: one actor does fetch + decode + generate +# --------------------------------------------------------------------------- + +class NaiveVLMEnricher: + """One GPU actor: HTTP fetch + PIL decode + vLLM generate, all inline.""" + + def __init__(self, model_id: str = MODEL_ID): + from vllm import LLM, SamplingParams + + self.llm = LLM( + model=model_id, + trust_remote_code=True, + max_model_len=4096, + limit_mm_per_prompt={"image": 1}, + dtype="float16", + gpu_memory_utilization=0.85, + ) + self.sampling = SamplingParams(max_tokens=256, temperature=0.0) + self.tokenizer = self.llm.get_tokenizer() + + def __call__(self, batch: dict) -> dict: + import requests + + prompts, kept = [], [] + for i in range(len(batch["product_id"])): + try: + resp = requests.get( + batch["image_url"][i], + timeout=5.0, + headers={"User-Agent": "vlm-distillation-catalog-enrichment/1.0"}, + ) + if resp.status_code != 200: + continue + img = ( + Image.open(io.BytesIO(resp.content)) + .convert("RGB") + .resize((384, 384)) + ) + except Exception: + continue + + prompt = self.tokenizer.apply_chat_template( + build_messages(img, batch["title"][i]), + tokenize=False, + add_generation_prompt=True, + ) + prompts.append({"prompt": prompt, "multi_modal_data": {"image": img}}) + kept.append(i) + + if not prompts: + return {k: [] for k in ("product_id", "title", "raw_output")} + + outputs = self.llm.generate(prompts, self.sampling) + return { + "product_id": [batch["product_id"][i] for i in kept], + "title": [batch["title"][i] for i in kept], + "raw_output": [o.outputs[0].text for o in outputs], + } + + +# --------------------------------------------------------------------------- +# Path 2 — HETEROGENEOUS: ray.data.llm processor, CPU pool feeds it bytes +# --------------------------------------------------------------------------- + +def build_heterogeneous_processor(num_gpus: int = 2, batch_size: int = 8, model_id: str = MODEL_ID): + """Returns a ray.data.llm processor that consumes ``image_bytes`` from the + upstream CPU stage (see ``src.preprocess.fetch_and_decode``). + + Expected input columns: product_id, title, image_bytes + Output columns: product_id, title, raw_output + """ + from ray.data.llm import build_llm_processor, vLLMEngineProcessorConfig + + config = vLLMEngineProcessorConfig( + model_source=model_id, + engine_kwargs={ + "trust_remote_code": True, + "max_model_len": 4096, + "limit_mm_per_prompt": {"image": 1}, + "dtype": "float16", + "gpu_memory_utilization": 0.85, + }, + concurrency=num_gpus, + batch_size=batch_size, + # vLLM 0.20+ moved TokensPrompt up from vllm.inputs.data to vllm.inputs. + # Ray 2.55's batch LLM stage still imports the old path, so each worker + # patches it on startup. Job-level ray.init also sets this hook for + # belt-and-suspenders. + runtime_env={"worker_process_setup_hook": "src._vllm_compat.patch"}, + ) + + def preprocess(row): + b64 = base64.b64encode(row["image_bytes"]).decode("ascii") + data_url = f"data:image/jpeg;base64,{b64}" + return dict( + messages=build_messages_url(data_url, row["title"]), + sampling_params=dict(max_tokens=256, temperature=0.0), + ) + + def postprocess(row): + return { + "product_id": row["product_id"], + "title": row["title"], + "raw_output": row.get("generated_text", "") or "", + } + + return build_llm_processor(config, preprocess=preprocess, postprocess=postprocess) diff --git a/templates/vlm-distillation-catalog-enrichment/src/load_catalog.py b/templates/vlm-distillation-catalog-enrichment/src/load_catalog.py new file mode 100644 index 000000000..fda247018 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/load_catalog.py @@ -0,0 +1,271 @@ +""" +Catalog loaders for the VLM batch enrichment demo. + +Primary source: McAuley-Lab/Amazon-Reviews-2023 (image URLs, served from Amazon CDN) +Backup source: shunk031/amazon-berkeley-objects (images bundled in the HF dataset) + +Both write a parquet with a uniform schema so the rest of the pipeline is +source-agnostic: + + product_id : str + title : str + description: str + image_url : str (HTTP for Amazon Reviews; "abo://" for ABO local cache) + source : str ("amazon-reviews-2023" | "amazon-berkeley-objects") +""" +import os +import random +from typing import Optional + +import pandas as pd + + +AMAZON_REVIEWS_DATASET = "McAuley-Lab/Amazon-Reviews-2023" +ABO_DATASET = "shunk031/amazon-berkeley-objects" + +# Categories mirror Amazon Reviews 2023 metadata splits +AMAZON_CATEGORIES = [ + "Electronics", + "Home_and_Kitchen", + "Clothing_Shoes_and_Jewelry", + "Beauty_and_Personal_Care", + "Sports_and_Outdoors", +] + + +def _extract_image_url(images_field) -> Optional[str]: + """Pick the first 'large' image URL (or hi_res/thumb fallback) from Amazon's image struct. + + HF parquet stores `images` as a struct of parallel lists, + e.g. ``{"large": [url1, url2], "hi_res": [...], "thumb": [...], "variant": [...]}``, + not as a list of per-image dicts. We take the first URL of the best + available size. + """ + if not images_field or not isinstance(images_field, dict): + return None + for key in ("large", "hi_res", "thumb"): + urls = images_field.get(key) + if urls is not None and len(urls) > 0: + return urls[0] + return None + + +def _has_title_and_image(row: dict) -> bool: + """Drop rows missing either a title or any usable image URL.""" + title = row.get("title") + if not (title and title.strip()): + return False + return _extract_image_url(row.get("images")) is not None + + +def _coerce_description(desc_field) -> str: + if isinstance(desc_field, list): + return " ".join(str(x) for x in desc_field if x).strip() + if isinstance(desc_field, str): + return desc_field.strip() + return "" + + +def _normalize_amazon_row_to_image(row: dict) -> dict: + """One catalog row per product, using the single best image URL (use with `.map`).""" + return { + "product_id": row.get("parent_asin") or row.get("asin") or "", + "title": row["title"].strip()[:512], + "description": _coerce_description(row.get("description"))[:1024], + "image_url": _extract_image_url(row["images"]), + "source": "amazon-reviews-2023", + } + + +def _normalize_amazon_row_to_images(row: dict, max_per_product: int = 8) -> list[dict]: + """One catalog row per *image* — explodes a product into N rows (use with `.flat_map`). + + Takes 'large' URLs (falling back to hi_res / thumb), capped at + ``max_per_product``. Each output row carries an ``image_idx`` so + ``(product_id, image_idx)`` is a stable row key. + """ + images = row.get("images") or {} + if not isinstance(images, dict): + return [] + urls = ( + images.get("large") + or images.get("hi_res") + or images.get("thumb") + or [] + )[:max_per_product] + if not urls: + return [] + + product_id = row.get("parent_asin") or row.get("asin") or "" + title = row["title"].strip()[:512] + description = _coerce_description(row.get("description"))[:1024] + + return [ + { + "product_id": product_id, + "image_idx": i, + "title": title, + "description": description, + "image_url": url, + "source": "amazon-reviews-2023", + } + for i, url in enumerate(urls) + ] + + +def load_amazon_reviews_2023( + category: str, + n_rows: int, + seed: int = 42, +): + """ + Lazy Ray Dataset of N normalized product rows from Amazon Reviews 2023. + + Reads parquet directly from the HF Hub via ray.data.read_parquet — + `datasets.load_dataset` no longer supports the script-based loader + this dataset originally shipped (Amazon-Reviews-2023.py). + + Output columns: ``product_id, title, description, image_url, source``. + Image bytes are NOT fetched here — the demo's naive vs heterogeneous + pipelines diverge on *where* the fetch happens (GPU actor vs dedicated + CPU stage), so the loader leaves URLs as URLs. + """ + import ray + from huggingface_hub import HfFileSystem + + if category not in AMAZON_CATEGORIES: + raise ValueError(f"Unknown category {category!r}; pick one of {AMAZON_CATEGORIES}") + + hf_path = f"hf://datasets/{AMAZON_REVIEWS_DATASET}/raw_meta_{category}" + print(f"[load] Building pipeline from {hf_path}...") + + return ( + ray.data.read_parquet( + hf_path, + file_extensions=["parquet"], + filesystem=HfFileSystem(), + ) + .limit(n_rows) + .filter(_has_title_and_image) + .map(_normalize_amazon_row_to_image) + .random_shuffle(seed=seed) + ) + + +def load_amazon_berkeley_objects( + n_rows: int, + output_path: str, + seed: int = 42, + image_cache_dir: str = "/mnt/cluster_storage/vlm-distillation-catalog-enrichment/abo-images", +) -> str: + """ + Backup loader: Amazon Berkeley Objects. + + ABO has images bundled in the HF dataset (PIL.Image), so we materialize + them to a local cache and reference them via 'abo://' URLs + that src/preprocess.py knows how to resolve. + + Use this if Amazon Reviews 2023 image URLs are unreachable from the + cluster (e.g., demo wifi blocking Amazon CDN, or rate-limiting). + """ + from datasets import load_dataset + from PIL import Image + + print(f"[load] Loading {ABO_DATASET} (images bundled, this may take a few minutes)...") + ds = load_dataset(ABO_DATASET, "all", split="train", trust_remote_code=True) + + rng = random.Random(seed) + indices = list(range(len(ds))) + rng.shuffle(indices) + indices = indices[:n_rows * 2] # buffer for skipped rows + + os.makedirs(image_cache_dir, exist_ok=True) + records = [] + for idx in indices: + row = ds[idx] + item_id = row.get("item_id") or f"abo-{idx}" + # ABO 'item_name' is a list of language-tagged dicts; take the en_US value. + name_list = row.get("item_name") or [] + title = "" + for entry in name_list: + if isinstance(entry, dict) and entry.get("language_tag", "").startswith("en"): + title = entry.get("value", "") + break + if not title: + continue + + image: Optional[Image.Image] = row.get("image") or row.get("main_image") + if image is None: + continue + + cache_path = os.path.join(image_cache_dir, f"{item_id}.jpg") + if not os.path.exists(cache_path): + image.convert("RGB").save(cache_path, "JPEG", quality=85) + + records.append({ + "product_id": item_id, + "title": title[:512], + "description": (row.get("product_description") or "")[:1024], + "image_url": f"abo://{item_id}.jpg", # preprocess resolves to image_cache_dir + "source": "amazon-berkeley-objects", + }) + if len(records) >= n_rows: + break + + df = pd.DataFrame(records) + os.makedirs(os.path.dirname(output_path), exist_ok=True) + df.to_parquet(output_path, index=False) + print(f" [load] Wrote {len(df):,} products → {output_path}") + print(f" [load] Image cache: {image_cache_dir}") + return output_path + + +def load_catalog( + source: str = "amazon-reviews-2023", + category: str = "Electronics", + n_rows: int = 10_000, + output_path: str = "/mnt/cluster_storage/vlm-distillation-catalog-enrichment/vlm-demo/catalog.parquet", +): + """ + Unified loader. `source` is one of: + - "amazon-reviews-2023" (default, returns a lazy Ray Dataset of image URLs) + - "amazon-berkeley-objects" (backup, materializes images, returns parquet path) + """ + if source == "amazon-reviews-2023": + return load_amazon_reviews_2023(category, n_rows) + elif source == "amazon-berkeley-objects": + return load_amazon_berkeley_objects(n_rows, output_path) + else: + raise ValueError(f"Unknown source {source!r}") + + +def shard_catalog_to_parquet( + category: str, + n_rows: int, + num_shards: int, + output_dir: str, + seed: int = 42, +) -> int: + """Build a normalized + sharded catalog at ``output_dir/shard_NNNN.parquet``. + + Each shard becomes one unit of resumable work for the sharded pipeline + (see ``scripts/run_pipeline_sharded.py`` and ``src.pipeline.run_with_checkpoints``). + + Returns the number of shards actually written (chunks with zero rows are skipped). + """ + import numpy as np + + print(f"[shard] Loading {n_rows:,} rows from {AMAZON_REVIEWS_DATASET}/{category}...") + df = load_amazon_reviews_2023(category, n_rows, seed=seed).to_pandas() + print(f"[shard] Got {len(df):,} rows after filtering — splitting into {num_shards} shards") + + os.makedirs(output_dir, exist_ok=True) + written = 0 + for shard_id, chunk in enumerate(np.array_split(df, num_shards)): + if len(chunk) == 0: + continue + path = os.path.join(output_dir, f"shard_{shard_id:04d}.parquet") + chunk.to_parquet(path, index=False) + written += 1 + print(f"[shard] Wrote {written} shard files → {output_dir}") + return written diff --git a/templates/vlm-distillation-catalog-enrichment/src/observability.py b/templates/vlm-distillation-catalog-enrichment/src/observability.py new file mode 100644 index 000000000..01fa60fea --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/observability.py @@ -0,0 +1,158 @@ +""" +Observability for the sharded VLM enrichment job. + +Three tiers, all wired through one set of helpers so the per-shard loop +only calls ``record_shard_complete()``: + +1. **Structured JSON logs** (always-on) — one line per event to stdout, + picked up by Anyscale's log aggregation. Searchable by ``event`` field. +2. **Ray metrics** (Counter / Gauge / Histogram) — surfaced via Ray's + Prometheus exporter and visible in the Anyscale Grafana dashboard. +3. **W&B** (optional, gated on ``WANDB_API_KEY``) — per-shard step charts. +""" +import json +import logging +import os +import sys +import time +from typing import Optional + + +# --------------------------------------------------------------------------- +# Structured logging (always-on) +# --------------------------------------------------------------------------- + +_log_handler = logging.StreamHandler(sys.stdout) +_log_handler.setFormatter(logging.Formatter("%(message)s")) # raw JSON lines +logger = logging.getLogger("vlm-enrichment") +logger.setLevel(logging.INFO) +logger.addHandler(_log_handler) +logger.propagate = False + + +def log_event(event: str, **fields) -> None: + """Emit a single JSON line for downstream log aggregation.""" + payload = {"ts": time.time(), "event": event, **fields} + logger.info(json.dumps(payload, default=str)) + + +# --------------------------------------------------------------------------- +# Ray metrics → Anyscale Prometheus → Grafana +# --------------------------------------------------------------------------- + +_ray_metrics_initialized = False +_PRODUCTS_PROCESSED = None +_SHARDS_COMPLETED = None +_SHARD_LATENCY = None + + +def init_ray_metrics() -> None: + """Lazy-init Ray metric handles. Must be called AFTER ray.init(). + + Safe to call multiple times. + """ + global _ray_metrics_initialized, _PRODUCTS_PROCESSED, _SHARDS_COMPLETED, _SHARD_LATENCY + if _ray_metrics_initialized: + return + try: + from ray.util.metrics import Counter, Gauge, Histogram + + _PRODUCTS_PROCESSED = Counter( + "vlm_enrichment_products_processed_total", + description="Total products enriched across all shards in this run", + tag_keys=("shard",), + ) + _SHARDS_COMPLETED = Gauge( + "vlm_enrichment_shards_completed", + description="Number of shards committed in this run", + ) + _SHARD_LATENCY = Histogram( + "vlm_enrichment_shard_seconds", + description="Wall time per shard", + boundaries=[10, 30, 60, 120, 300, 600, 1200, 3600], + ) + _ray_metrics_initialized = True + except Exception as e: + log_event("ray_metrics_init_skipped", reason=str(e)) + + +def record_shard_metrics( + shard_id: int, num_products: int, elapsed_seconds: float, total_completed: int +) -> None: + if not _ray_metrics_initialized: + return + _PRODUCTS_PROCESSED.inc(num_products, tags={"shard": str(shard_id)}) + _SHARDS_COMPLETED.set(total_completed) + _SHARD_LATENCY.observe(elapsed_seconds) + + +# --------------------------------------------------------------------------- +# W&B (optional) +# --------------------------------------------------------------------------- + +_wandb_run = None + + +def maybe_init_wandb(run_name: str, total_shards: int, project: Optional[str] = None) -> bool: + """Initialize a W&B run if WANDB_API_KEY is set. No-op otherwise.""" + global _wandb_run + if not os.environ.get("WANDB_API_KEY"): + log_event("wandb_skipped", reason="WANDB_API_KEY not set") + return False + try: + import wandb + except ImportError: + log_event("wandb_skipped", reason="wandb not installed") + return False + + _wandb_run = wandb.init( + project=project or os.environ.get("WANDB_PROJECT", "anyscale-vlm-enrichment-demo"), + name=run_name, + resume="allow", + config={"total_shards": total_shards}, + ) + log_event("wandb_initialized", run_name=run_name, project=_wandb_run.project) + return True + + +def log_to_wandb(**fields) -> None: + if _wandb_run is None: + return + _wandb_run.log(fields) + + +def finish_wandb() -> None: + global _wandb_run + if _wandb_run is not None: + _wandb_run.finish() + _wandb_run = None + + +# --------------------------------------------------------------------------- +# Composite helper — call after each committed shard. +# --------------------------------------------------------------------------- + +def record_shard_complete( + shard_id: int, + num_products: int, + elapsed_seconds: float, + total_completed: int, + total_shards: int, +) -> None: + throughput = num_products / elapsed_seconds if elapsed_seconds > 0 else 0 + log_event( + "shard_complete", + shard_id=shard_id, + num_products=num_products, + elapsed_seconds=round(elapsed_seconds, 2), + throughput_per_sec=round(throughput, 2), + progress=f"{total_completed}/{total_shards}", + ) + record_shard_metrics(shard_id, num_products, elapsed_seconds, total_completed) + log_to_wandb( + shard_id=shard_id, + shard_throughput=throughput, + shard_wall_seconds=elapsed_seconds, + shards_completed=total_completed, + progress_pct=100.0 * total_completed / max(total_shards, 1), + ) diff --git a/templates/vlm-distillation-catalog-enrichment/src/pipeline.py b/templates/vlm-distillation-catalog-enrichment/src/pipeline.py new file mode 100644 index 000000000..7c2b984a8 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/pipeline.py @@ -0,0 +1,296 @@ +""" +Ray Data VLM batch enrichment pipeline. + +Three modes: + - run_pipeline_naive: Single GPU stage does fetch + decode + infer. + GPU is blocked on I/O. The "before" picture. + - run_pipeline_heterogeneous: CPU pool fetches + decodes; GPU pool only + runs the VLM (via ray.data.llm). The "after". + - run_with_checkpoints: Same heterogeneous pipeline, sharded with + atomic per-shard commits → resumable across + cluster restarts. The production-grade path. + +Catalog for the simple modes is cached at ``catalog_path``; for the sharded +mode it's pre-split into ``input_dir/shard_NNNN.parquet`` files (see +``src.load_catalog.shard_catalog_to_parquet``). +""" +import os +import time +import uuid + +import ray +import ray.data +from ray.data import ActorPoolStrategy + +from src.enrich import NaiveVLMEnricher, build_heterogeneous_processor +from src.load_catalog import load_amazon_reviews_2023 +from src.observability import ( + finish_wandb, + init_ray_metrics, + log_event, + maybe_init_wandb, + record_shard_complete, +) +from src.preprocess import fetch_and_decode +from src.shard import ( + cleanup_stale_tmp, + commit_shard, + list_completed_shards, + list_input_shards, + shard_input_path, + shard_tmp_path, +) +from src.utils import ( + calc_throughput, + estimate_job_cost, + estimate_single_node_time, + format_number, + print_metrics_table, +) + + +def ensure_catalog(category: str, n_rows: int, catalog_path: str) -> None: + """Build the normalized parquet cache on first run; reuse on subsequent runs.""" + if os.path.exists(catalog_path): + print(f"[load] Reusing existing catalog at {catalog_path}") + return + print(f"[load] Building catalog at {catalog_path} ({n_rows:,} rows from {category})") + load_amazon_reviews_2023(category, n_rows=n_rows).write_parquet(catalog_path) + + +def run_pipeline_naive(catalog_path: str, output_path: str, num_gpus: int = 2) -> dict: + """Single GPU stage does HTTP fetch + image decode + VLM generate.""" + pipeline_start = time.time() + + print(f"\n[1/3] Reading catalog from {catalog_path}") + ds = ray.data.read_parquet(catalog_path, override_num_blocks=num_gpus * 8) + total = ds.count() + print(f" Products loaded: {format_number(total)}") + + print(f"\n[2/3] NAIVE enrichment — {num_gpus} GPU actor(s) doing fetch + decode + generate") + print(f" (GPU will sit idle while each image downloads — expect low utilization)") + ds = ds.map_batches( + NaiveVLMEnricher, + batch_size=8, + num_gpus=1, + compute=ActorPoolStrategy(size=num_gpus), + batch_format="numpy", + ) + + print(f"\n[3/3] Writing enriched catalog to {output_path}") + ds.write_parquet(output_path) + + wall_time = time.time() - pipeline_start + throughput = calc_throughput(total, wall_time) + + # Real per-operator timings live in ds.stats() — Ray Data's internal + # streaming executor accumulates them as it runs. Manual time.time() + # around .map_batches() / .write_parquet() is misleading because Ray + # Data is lazy; the work all attributes to the terminal operator. + print("\n--- Ray Data per-operator stats ---") + print(ds.stats()) + + print("\n--- Sample enriched rows ---") + print(ray.data.read_parquet(output_path).take(3)) + + metrics = { + "Mode": "NAIVE (GPU does fetch + decode + infer)", + "Total products processed": format_number(total), + "Wall time": f"{wall_time:.1f}s ({wall_time / 60:.1f} min)", + "Throughput": f"{throughput:,.2f} products/sec", + "GPU workers": str(num_gpus), + "Model": "Qwen2.5-VL-3B-Instruct", + "Output path": output_path, + "Est. single-node time": estimate_single_node_time(total, throughput), + "Est. Anyscale job cost": estimate_job_cost(wall_time, num_cpu_workers=0, num_gpu_workers=num_gpus), + } + print_metrics_table(metrics) + return metrics + + +def run_pipeline_heterogeneous( + catalog_path: str, + output_path: str, + num_gpus: int = 2, + cpu_concurrency: int = 8, +) -> dict: + """CPU pool fetches + decodes; ray.data.llm processor runs only the VLM.""" + pipeline_start = time.time() + + print(f"\n[1/4] Reading catalog from {catalog_path}") + ds = ray.data.read_parquet(catalog_path, override_num_blocks=cpu_concurrency * 4) + total = ds.count() + print(f" Products loaded: {format_number(total)}") + + print(f"\n[2/4] CPU image fetch + decode — {cpu_concurrency} CPU workers") + ds = ds.map_batches( + fetch_and_decode, + batch_size=16, + concurrency=cpu_concurrency, + batch_format="numpy", + ) + + print(f"\n[3/4] VLM enrichment via ray.data.llm — {num_gpus} GPU worker(s)") + print(f" (GPU only runs inference; CPU pool keeps it fed)") + processor = build_heterogeneous_processor(num_gpus=num_gpus, batch_size=8) + ds = processor(ds) + + print(f"\n[4/4] Writing enriched catalog to {output_path}") + ds.write_parquet(output_path) + + wall_time = time.time() - pipeline_start + throughput = calc_throughput(total, wall_time) + + # Real per-operator timings live in ds.stats() — Ray Data's streaming + # executor accumulates them as it runs. Manual time.time() around lazy + # operators is misleading because the work attributes to the terminal + # write_parquet rather than the upstream stages it actually came from. + print("\n--- Ray Data per-operator stats ---") + print(ds.stats()) + + print("\n--- Sample enriched rows ---") + print(ray.data.read_parquet(output_path).take(3)) + + metrics = { + "Mode": "HETEROGENEOUS (CPU pool fetches; GPU only infers)", + "Total products processed": format_number(total), + "Wall time": f"{wall_time:.1f}s ({wall_time / 60:.1f} min)", + "Throughput": f"{throughput:,.2f} products/sec", + "CPU workers": str(cpu_concurrency), + "GPU workers": str(num_gpus), + "Model": "Qwen2.5-VL-3B-Instruct", + "Output path": output_path, + "Est. single-node time": estimate_single_node_time(total, throughput), + "Est. Anyscale job cost": estimate_job_cost( + wall_time, num_cpu_workers=cpu_concurrency, num_gpu_workers=num_gpus + ), + } + print_metrics_table(metrics) + return metrics + + +def run_with_checkpoints( + input_dir: str, + output_dir: str, + num_gpus: int = 2, + batch_size: int = 8, + cpu_concurrency: int = 8, + run_name: str = None, + max_shards_this_run: int = None, +) -> dict: + """Process input shards with atomic per-shard commits. + + - Skips already-committed shards (resume semantics). + - Cleans up half-written .tmp/ dirs from a prior crashed run. + - Builds the vLLM processor once so engine actors stay warm across shards. + - On SIGTERM (caller raises KeyboardInterrupt), the in-flight shard's + .tmp/ stays behind; the next run cleans it and reprocesses. + """ + run_name = run_name or f"run-{uuid.uuid4().hex[:8]}" + pipeline_start = time.time() + + os.makedirs(output_dir, exist_ok=True) + + stale = cleanup_stale_tmp(output_dir) + if stale: + log_event("cleanup_stale_tmp", removed=stale) + + all_inputs = list_input_shards(input_dir) + completed = list_completed_shards(output_dir) + remaining = sorted(set(all_inputs) - completed) + + log_event( + "checkpoint_scan", + input_shards=len(all_inputs), + already_completed=len(completed), + remaining=len(remaining), + run_name=run_name, + ) + + if not remaining: + print(f"All {len(all_inputs)} shards already completed at {output_dir} — nothing to do.") + return {"status": "no-op", "completed": len(completed), "total": len(all_inputs)} + + if max_shards_this_run is not None and max_shards_this_run < len(remaining): + log_event("max_shards_capped", remaining_total=len(remaining), processing=max_shards_this_run) + remaining = remaining[:max_shards_this_run] + + print(f"\n[checkpoint] Resuming run: {len(completed)}/{len(all_inputs)} done, " + f"{len(remaining)} shards to process this invocation\n") + + init_ray_metrics() + maybe_init_wandb(run_name=run_name, total_shards=len(all_inputs)) + + # Build the vLLM processor ONCE — actor pool persists across shards so + # the model loads only on the first shard. + print(f"[engine] Building Qwen2.5-VL-3B vLLM processor ({num_gpus} GPU workers)...") + processor = build_heterogeneous_processor(num_gpus=num_gpus, batch_size=batch_size) + + total_products = 0 + shard_times = [] + + for shard_id in remaining: + shard_start = time.time() + in_path = shard_input_path(input_dir, shard_id) + tmp_dir = shard_tmp_path(output_dir, shard_id) + + print(f"[shard {shard_id:04d}] Reading {in_path}") + ds = ray.data.read_parquet(in_path) + n_in = ds.count() + + ds = ds.map_batches( + fetch_and_decode, + batch_size=16, + concurrency=cpu_concurrency, + batch_format="numpy", + ) + ds = processor(ds) + ds.write_parquet(tmp_dir) + + # Atomic commit. If the process dies before this rename, the next run + # cleans up the .tmp/ and reprocesses this shard from scratch. + commit_shard(output_dir, shard_id) + + elapsed = time.time() - shard_start + total_products += n_in + shard_times.append(elapsed) + + completed_now = len(completed) + len(shard_times) + record_shard_complete( + shard_id=shard_id, + num_products=n_in, + elapsed_seconds=elapsed, + total_completed=completed_now, + total_shards=len(all_inputs), + ) + print(f"[shard {shard_id:04d}] ✓ {n_in} products in {elapsed:.1f}s " + f"({completed_now}/{len(all_inputs)} shards done)") + + finish_wandb() + + print("\n--- Sample enriched rows ---") + print(ray.data.read_parquet(output_dir).take(3)) + + wall_time = time.time() - pipeline_start + avg_shard = sum(shard_times) / len(shard_times) if shard_times else 0 + throughput = calc_throughput(total_products, wall_time) + + metrics = { + "Run name": run_name, + "Mode": "Sharded + checkpointed (ray.data.llm + Qwen2.5-VL-3B)", + "Shards completed this run": format_number(len(shard_times)), + "Shards committed total": format_number(len(completed) + len(shard_times)), + "Shards remaining": format_number(len(all_inputs) - len(completed) - len(shard_times)), + "Products processed this run": format_number(total_products), + "Wall time": f"{wall_time:.1f}s ({wall_time / 60:.1f} min)", + "Avg time per shard": f"{avg_shard:.1f}s", + "Throughput": f"{throughput:,.2f} products/sec", + "CPU concurrency": str(cpu_concurrency), + "GPU workers": str(num_gpus), + "Output dir": output_dir, + "Est. Anyscale cost (this run)": estimate_job_cost( + wall_time, num_cpu_workers=cpu_concurrency, num_gpu_workers=num_gpus + ), + } + print_metrics_table(metrics, title="VLM ENRICHMENT JOB — RUN COMPLETE") + return metrics diff --git a/templates/vlm-distillation-catalog-enrichment/src/preprocess.py b/templates/vlm-distillation-catalog-enrichment/src/preprocess.py new file mode 100644 index 000000000..5bd956b84 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/preprocess.py @@ -0,0 +1,46 @@ +""" +CPU image fetch + decode stage. + +Runs on Ray Data CPU workers via map_batches. Pulls image bytes from HTTP, +decodes to PIL.Image, resizes to 384x384 (Qwen2.5-VL input sweet spot), +re-encodes as JPEG, and emits raw bytes ready for the VLM. + +Mirrors the inline helper in `notebooks/demo_walkthrough.ipynb`. +""" +import io + +from PIL import Image + + +TARGET_SIZE = (384, 384) + + +def fetch_and_decode(batch: dict) -> dict: + """CPU stage: HTTP fetch → PIL decode → 384x384 JPEG bytes ready for vLLM. + + Input columns: product_id, title, image_url, ... + Output columns: product_id, title, image_url, image_bytes + (rows with failed fetch/decode are dropped) + """ + import requests + + out = {k: [] for k in ("product_id", "title", "image_url", "image_bytes")} + for i in range(len(batch["product_id"])): + try: + resp = requests.get( + batch["image_url"][i], + timeout=5.0, + headers={"User-Agent": "vlm-distillation-catalog-enrichment/1.0"}, + ) + if resp.status_code != 200: + continue + img = Image.open(io.BytesIO(resp.content)).convert("RGB").resize(TARGET_SIZE) + buf = io.BytesIO() + img.save(buf, "JPEG", quality=90) + except Exception: + continue + out["product_id"].append(batch["product_id"][i]) + out["title"].append(batch["title"][i]) + out["image_url"].append(batch["image_url"][i]) + out["image_bytes"].append(buf.getvalue()) + return out diff --git a/templates/vlm-distillation-catalog-enrichment/src/shard.py b/templates/vlm-distillation-catalog-enrichment/src/shard.py new file mode 100644 index 000000000..a5986c349 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/shard.py @@ -0,0 +1,102 @@ +""" +Shard checkpointing — the resumability primitive for the sharded VLM job. + +Convention: + input_dir/shard_NNNN.parquet — pre-sharded catalog input (one shard per file) + output_dir/shard_NNNN/ — committed enriched shard (atomic write) + output_dir/.shard_NNNN.tmp/ — in-flight write (gets renamed on commit) + +A shard is "completed" iff a directory named ``shard_NNNN`` exists at +``output_dir`` with at least one parquet inside. Commits use POSIX +``os.rename``, atomic for directories on the same filesystem (works on +NFS too, which is what ``/mnt/cluster_storage/vlm-distillation-catalog-enrichment`` is). +""" +import os +import re +import shutil + + +SHARD_RE = re.compile(r"shard_(\d{4})$") +TMP_RE = re.compile(r"\.shard_(\d{4})\.tmp$") + + +def shard_input_path(input_dir: str, shard_id: int) -> str: + return os.path.join(input_dir, f"shard_{shard_id:04d}.parquet") + + +def shard_output_path(output_dir: str, shard_id: int) -> str: + return os.path.join(output_dir, f"shard_{shard_id:04d}") + + +def shard_tmp_path(output_dir: str, shard_id: int) -> str: + return os.path.join(output_dir, f".shard_{shard_id:04d}.tmp") + + +def list_input_shards(input_dir: str) -> list[int]: + """Return sorted shard IDs present in input_dir.""" + if not os.path.exists(input_dir): + return [] + ids = [] + for name in os.listdir(input_dir): + m = re.match(r"shard_(\d{4})\.parquet$", name) + if m: + ids.append(int(m.group(1))) + return sorted(ids) + + +def list_completed_shards(output_dir: str) -> set[int]: + """Return set of shard IDs that have a fully-committed output dir.""" + if not os.path.exists(output_dir): + return set() + completed = set() + for name in os.listdir(output_dir): + m = SHARD_RE.match(name) + if not m: + continue + shard_path = os.path.join(output_dir, name) + # Sanity: a committed shard has at least one parquet inside. + if os.path.isdir(shard_path) and any( + f.endswith(".parquet") for f in os.listdir(shard_path) + ): + completed.add(int(m.group(1))) + return completed + + +def cleanup_stale_tmp(output_dir: str) -> int: + """Remove leftover .shard_NNNN.tmp/ dirs from a previous crashed run. + + Returns the count cleaned. Safe to call before every run. + """ + if not os.path.exists(output_dir): + return 0 + cleaned = 0 + for name in os.listdir(output_dir): + if TMP_RE.match(name): + shutil.rmtree(os.path.join(output_dir, name), ignore_errors=True) + cleaned += 1 + return cleaned + + +def commit_shard(output_dir: str, shard_id: int) -> None: + """Atomically promote a .shard_NNNN.tmp/ dir to shard_NNNN/. + + POSIX rename is atomic for directories on the same filesystem, so an + observer (e.g. another process scanning for completed shards) will + always see either the old name or the new name — never a half-state. + """ + tmp = shard_tmp_path(output_dir, shard_id) + final = shard_output_path(output_dir, shard_id) + if not os.path.exists(tmp): + raise FileNotFoundError(f"No tmp dir to commit at {tmp}") + if os.path.exists(final): + # Idempotent: a previous run already committed this shard. + shutil.rmtree(tmp, ignore_errors=True) + return + os.rename(tmp, final) + + +def remaining_shards(input_dir: str, output_dir: str) -> list[int]: + """Input shards minus already-committed output shards, sorted.""" + inputs = set(list_input_shards(input_dir)) + completed = list_completed_shards(output_dir) + return sorted(inputs - completed) diff --git a/templates/vlm-distillation-catalog-enrichment/src/utils.py b/templates/vlm-distillation-catalog-enrichment/src/utils.py new file mode 100644 index 000000000..69a32c467 --- /dev/null +++ b/templates/vlm-distillation-catalog-enrichment/src/utils.py @@ -0,0 +1,59 @@ +""" +Timing, metrics, and display utilities for the VLM enrichment pipeline. +""" +import time +from contextlib import contextmanager + + +@contextmanager +def timer(label: str): + print(f"\n⏱ {label}...") + start = time.time() + yield + elapsed = time.time() - start + print(f" ✓ {label} complete — {elapsed:.1f}s") + + +def calc_throughput(num_items: int, elapsed_seconds: float) -> float: + return num_items / elapsed_seconds if elapsed_seconds > 0 else 0.0 + + +def format_number(n: int) -> str: + return f"{n:,}" + + +def print_metrics_table(metrics: dict, title: str = "PIPELINE METRICS SUMMARY"): + width = 60 + print("\n" + "=" * width) + print(f" {title}") + print("=" * width) + for key, value in metrics.items(): + print(f" {key:<32} {value}") + print("=" * width + "\n") + + +def estimate_single_node_time(num_items: int, throughput_per_sec: float = 0.5) -> str: + """ + Rough single-node estimate for Qwen2.5-VL-3B on a single A10G with no + CPU offload of image fetch/decode. Empirically ~0.5 items/sec end-to-end + when the GPU also fetches and decodes images. + """ + seconds = num_items / max(throughput_per_sec, 1e-6) + if seconds < 60: + return f"{seconds:.0f}s" + elif seconds < 3600: + return f"{seconds / 60:.1f} min" + else: + return f"{seconds / 3600:.1f} hours" + + +def estimate_job_cost( + wall_time_seconds: float, + num_cpu_workers: int = 4, + num_gpu_workers: int = 2, + cpu_hourly: float = 0.768, # m5.4xlarge on-demand us-west-2 + gpu_hourly: float = 1.006, # g5.xlarge A10G on-demand us-west-2 +) -> str: + hours = wall_time_seconds / 3600 + total = num_cpu_workers * cpu_hourly * hours + num_gpu_workers * gpu_hourly * hours + return f"~${total:.2f} (est. on-demand, m5.4xlarge CPU + g5.xlarge A10G)" diff --git a/tests/vlm-distillation-catalog-enrichment/nb2py.py b/tests/vlm-distillation-catalog-enrichment/nb2py.py new file mode 100644 index 000000000..3c7f38322 --- /dev/null +++ b/tests/vlm-distillation-catalog-enrichment/nb2py.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +import argparse +import nbformat + + +def convert_notebook(input_path: str, output_path: str) -> None: + """ + Read a Jupyter notebook and write a Python script, converting all %%bash + cells and IPython "!" commands into subprocess.run calls that raise on error. + Cells that load or autoreload extensions are ignored. + """ + nb = nbformat.read(input_path, as_version=4) + with open(output_path, "w") as out: + for cell in nb.cells: + # Only process code cells + if cell.cell_type != "code": + continue + + lines = cell.source.splitlines() + # Skip cells that load or autoreload extensions + if any( + l.strip().startswith("%load_ext autoreload") + or l.strip().startswith("%autoreload all") + for l in lines + ): + continue + + # Detect a %%bash cell + if lines and lines[0].strip().startswith("%%bash"): + bash_script = "\n".join(lines[1:]).rstrip() + out.write("import subprocess\n") + out.write( + f"subprocess.run(r'''{bash_script}''',\n" + " shell=True,\n" + " check=True,\n" + " executable='/bin/bash')\n\n" + ) + else: + # Detect any IPython '!' shell commands in code lines + has_bang = any(line.lstrip().startswith("!") for line in lines) + if has_bang: + out.write("import subprocess\n") + for line in lines: + stripped = line.lstrip() + if stripped.startswith("!"): + cmd = stripped[1:].lstrip() + out.write( + f"subprocess.run(r'''{cmd}''',\n" + " shell=True,\n" + " check=True,\n" + " executable='/bin/bash')\n" + ) + else: + out.write(line.rstrip() + "\n") + out.write("\n") + else: + # Regular Python cell: dump as-is + out.write(cell.source.rstrip() + "\n\n") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Convert a Jupyter notebook to a Python script, preserving bash cells and '!' commands as subprocess calls." + ) + parser.add_argument("input_nb", help="Path to the input .ipynb file") + parser.add_argument("output_py", help="Path for the output .py script") + args = parser.parse_args() + convert_notebook(args.input_nb, args.output_py) + + +if __name__ == "__main__": + main() diff --git a/tests/vlm-distillation-catalog-enrichment/tests.sh b/tests/vlm-distillation-catalog-enrichment/tests.sh new file mode 100755 index 000000000..f1423899f --- /dev/null +++ b/tests/vlm-distillation-catalog-enrichment/tests.sh @@ -0,0 +1,18 @@ +#!/bin/bash + +set -euxo pipefail + +# Smoke run: N=20 rows through the full 3-stage distillation flow. +# The orchestrator README.ipynb at the template root runs end-to-end. +export N_ROWS=20 +export TEACHER_N_ROWS=20 +export NUM_EPOCHS=1 +export TRAIN_FRAC=0.5 + +cd /home/ray/default/templates/vlm-distillation-catalog-enrichment + +python /home/ray/default/tests/vlm-distillation-catalog-enrichment/nb2py.py \ + README.ipynb README.py + +python README.py +rm README.py