feat(templates): add vlm distillation catalog enrichment template + test

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

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

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

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<a href=\"https://console.anyscale.com/register/ha?render_flow=ray&utm_source=ray_docs&utm_medium=docs&utm_campaign=vlm-distillation-catalog-enrichment&redirectTo=/v2/template-preview/vlm-distillation-catalog-enrichment\">\n<img src=\"https://raw.githubusercontent.com/ray-project/ray/c34b74c22a9390aa89baf80815ede59397786d2e/doc/source/_static/img/run-on-anyscale.svg\" alt=\"Run on Anyscale\">\n</a>\n<br></br>\n<div align=\"left\">\n<a href=\"https://github.com/anyscale/templates\" role=\"button\"><img src=\"https://img.shields.io/static/v1?label=&amp;message=View%20On%20GitHub&amp;color=586069&amp;logo=github&amp;labelColor=2f363d\"></a>&nbsp;\n</div>\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
+}