# Ray Data vs Geneva: Image Embedding & Captioning Pipeline

This file compares Ray Data (Ray's native data processing library) with Geneva
for building production image embedding pipelines.

## Key Insight

Ray Data is the foundation Geneva builds on. Geneva adds the production
lifecycle features that Ray Data intentionally doesn't include.

Think of it as:
- **Ray Data** = Powerful distributed compute engine
- **Geneva** = Feature engineering platform (built on Ray Data + Lance)

## Summary

| Approach | Lines of Code |
|----------|---------------|
| Ray Data | ~100 lines for computation + ~150 lines for production ops (custom) |
| Geneva | ~80 lines total, production ops included automatically |

---

## Ray Data Approach

### Ray Data Excels At

- Streaming data between CPU and GPU stages
- High GPU utilization (90%+ reported)
- Heterogeneous cluster support
- Integration with ML frameworks (PyTorch, HuggingFace, vLLM)

### You'll Need to Add

- Checkpointing for fault tolerance (only in Anyscale paid tier)
- Incremental processing logic
- Schema management and versioning
- Dependency packaging

---

### UDF Definitions

```python
import ray
from ray.data import ActorPoolStrategy
import numpy as np
from typing import Dict
import torch


class ComputeEmbeddingRayData:
    """
    CLIP embedding using Ray Data's class-based UDF pattern.

    - __init__ for model loading (called once per actor)
    - __call__ receives batch as Dict[str, np.ndarray]
    - Must handle batching and error recovery yourself
    """

    def __init__(self):
        import open_clip
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            'ViT-B-32', pretrained='laion2b_s34b_b79k'
        )
        self.model = self.model.cuda().eval()

    def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        from PIL import Image
        import io

        embeddings = []
        for img_bytes in batch["image"]:
            if img_bytes is None:
                embeddings.append(np.zeros(512, dtype=np.float32))
                continue
            try:
                img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
                tensor = self.preprocess(img).unsqueeze(0).cuda()
                with torch.no_grad():
                    embedding = self.model.encode_image(tensor)
                embeddings.append(embedding.cpu().numpy().flatten())
            except Exception:
                embeddings.append(np.zeros(512, dtype=np.float32))

        batch["embedding"] = np.array(embeddings)
        return batch


class GenerateCaptionRayData:
    """BLIP captioning using Ray Data's class-based UDF pattern."""

    def __init__(self):
        from transformers import BlipProcessor, BlipForConditionalGeneration
        self.processor = BlipProcessor.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        )
        self.model = BlipForConditionalGeneration.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        ).cuda().eval()

    def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        from PIL import Image
        import io

        captions = []
        for img_bytes in batch["image"]:
            if img_bytes is None:
                captions.append("")
                continue
            try:
                img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
                inputs = self.processor(img, return_tensors="pt")
                inputs = {k: v.cuda() for k, v in inputs.items()}
                with torch.no_grad():
                    output = self.model.generate(**inputs, max_new_tokens=50)
                captions.append(self.processor.decode(output[0], skip_special_tokens=True))
            except Exception:
                captions.append("")

        batch["caption"] = np.array(captions)
        return batch
```

### Basic Pipeline (no production ops)

```python
def run_raydata_pipeline_basic():
    """
    Basic Ray Data pipeline - clean and efficient for one-shot jobs.

    This is where Ray Data shines: streaming execution, high GPU utilization.
    """
    ray.init(address="auto")

    # Read data
    ds = ray.data.read_parquet("s3://my-bucket/images/")

    # Compute embeddings with GPU actors
    ds = ds.map_batches(
        ComputeEmbeddingRayData,
        compute=ActorPoolStrategy(size=4),
        num_gpus=1,
        batch_size=32,
    )

    # Compute captions with GPU actors
    ds = ds.map_batches(
        GenerateCaptionRayData,
        compute=ActorPoolStrategy(size=4),
        num_gpus=1,
        batch_size=16,
    )

    # Write results
    ds.write_parquet("s3://my-bucket/embeddings/")
```

### Production Pipeline (with custom ops)

For production with open-source Ray Data, you need to build:
1. **Checkpointing** - Ray Data OSS doesn't checkpoint data pipelines
2. **Incremental processing** - Process only new data
3. **Progress tracking** - Know where you are

> **Note**: Anyscale (paid) offers job-level checkpointing, but with limitations:
> - Only for map-based operators (no joins, aggregations)
> - Requires id_column for row identification
> - Known bug: may produce duplicate outputs on retry

```python
import json
from pathlib import Path
from datetime import datetime


class RayDataCheckpointManager:
    """
    Custom checkpoint manager for Ray Data pipelines.

    You build this yourself - Ray Data OSS doesn't include it.
    """

    def __init__(self, checkpoint_dir: str):
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self.state_file = self.checkpoint_dir / "state.json"
        self.state = self._load_state()

    def _load_state(self) -> dict:
        if self.state_file.exists():
            return json.loads(self.state_file.read_text())
        return {"processed_files": [], "rows_processed": 0}

    def _save_state(self):
        self.state["last_checkpoint"] = datetime.now().isoformat()
        self.state_file.write_text(json.dumps(self.state, indent=2))

    def mark_file_processed(self, file_path: str, row_count: int):
        self.state["processed_files"].append(file_path)
        self.state["rows_processed"] += row_count
        self._save_state()

    def is_file_processed(self, file_path: str) -> bool:
        return file_path in self.state["processed_files"]

    def get_stats(self) -> dict:
        return self.state.copy()


class RayDataIncrementalProcessor:
    """
    Incremental processing wrapper for Ray Data.

    Implements file-level checkpointing since Ray Data doesn't
    support row-level or block-level checkpointing in OSS.
    """

    def __init__(
        self,
        input_path: str,
        output_path: str,
        checkpoint_dir: str,
        num_gpu_actors: int = 4,
        batch_size: int = 32,
    ):
        self.input_path = input_path
        self.output_path = output_path
        self.num_gpu_actors = num_gpu_actors
        self.batch_size = batch_size
        self.checkpoint = RayDataCheckpointManager(checkpoint_dir)

    def _list_input_files(self) -> list[str]:
        """List all parquet files in input path."""
        import pyarrow.fs as pafs
        fs, path = pafs.FileSystem.from_uri(self.input_path)
        file_info = fs.get_file_info(pafs.FileSelector(path, recursive=True))
        return [f.path for f in file_info if f.path.endswith('.parquet')]

    def _get_pending_files(self) -> list[str]:
        all_files = self._list_input_files()
        return [f for f in all_files if not self.checkpoint.is_file_processed(f)]

    def run(self):
        """Run incremental processing with checkpointing."""
        ray.init(address="auto", ignore_reinit_error=True)

        pending_files = self._get_pending_files()
        print(f"Files to process: {len(pending_files)}")

        if not pending_files:
            print("Nothing to process!")
            return

        # Process file by file for checkpointing granularity
        for i, file_path in enumerate(pending_files):
            print(f"Processing [{i+1}/{len(pending_files)}]: {file_path}")

            try:
                ds = ray.data.read_parquet(file_path)

                ds = ds.map_batches(
                    ComputeEmbeddingRayData,
                    compute=ActorPoolStrategy(size=self.num_gpu_actors),
                    num_gpus=1,
                    batch_size=self.batch_size,
                )

                ds = ds.map_batches(
                    GenerateCaptionRayData,
                    compute=ActorPoolStrategy(size=self.num_gpu_actors),
                    num_gpus=1,
                    batch_size=self.batch_size // 2,
                )

                row_count = ds.count()
                output_file = f"{self.output_path}/batch_{i:06d}.parquet"
                ds.write_parquet(output_file)

                self.checkpoint.mark_file_processed(file_path, row_count)
                print(f"  Completed: {row_count} rows")

            except Exception as e:
                print(f"  FAILED: {e}")
                print("  Re-run to resume from this file.")
                raise

        print(f"Complete! Total rows: {self.checkpoint.get_stats()['rows_processed']}")


def run_raydata_pipeline_production():
    """Production Ray Data pipeline with custom checkpointing."""
    processor = RayDataIncrementalProcessor(
        input_path="s3://my-bucket/images",
        output_path="s3://my-bucket/embeddings",
        checkpoint_dir="/tmp/raydata_checkpoints",
    )
    processor.run()
```

### Dependency Management

```python
def configure_raydata_with_dependencies():
    """
    Configure Ray cluster with required dependencies.

    You manage this yourself - Ray Data doesn't auto-package deps.
    """
    runtime_env = {
        "pip": [
            "torch>=2.1.0",
            "open-clip-torch>=2.24.0",
            "transformers>=4.36.0",
            "Pillow>=10.0.0",
        ],
    }
    ray.init(address="auto", runtime_env=runtime_env)
```

---

## Geneva Approach

Geneva is built on Ray (uses Ray Data internally) but adds:
- Fragment-level AND batch-level checkpointing
- Automatic incremental processing
- Per-UDF dependency isolation
- Table versioning and time travel

### UDF Definitions

```python
import geneva


@geneva.udf(num_gpus=1, version="0.1")
class ComputeEmbeddingGeneva:
    """
    CLIP embedding UDF for Geneva.

    Key differences from Ray Data:
    - setup() instead of __init__()
    - Process single row, Geneva handles batching
    - Type annotation -> schema (no manual dtype)
    """

    def setup(self):
        import open_clip
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            'ViT-B-32', pretrained='laion2b_s34b_b79k'
        )
        self.model = self.model.cuda().eval()

    def __call__(self, image: bytes) -> list[float]:
        from PIL import Image
        import torch
        import io

        img = Image.open(io.BytesIO(image)).convert('RGB')
        tensor = self.preprocess(img).unsqueeze(0).cuda()
        with torch.no_grad():
            embedding = self.model.encode_image(tensor)
        return embedding.cpu().numpy().flatten().tolist()


@geneva.udf(num_gpus=1, version="0.1")
class GenerateCaptionGeneva:
    """BLIP captioning UDF for Geneva."""

    def setup(self):
        from transformers import BlipProcessor, BlipForConditionalGeneration
        self.processor = BlipProcessor.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        )
        self.model = BlipForConditionalGeneration.from_pretrained(
            "Salesforce/blip-image-captioning-base"
        ).cuda()

    def __call__(self, image: bytes) -> str:
        from PIL import Image
        import torch
        import io

        img = Image.open(io.BytesIO(image)).convert('RGB')
        inputs = self.processor(img, return_tensors="pt")
        inputs = {k: v.cuda() for k, v in inputs.items()}
        with torch.no_grad():
            output = self.model.generate(**inputs, max_new_tokens=50)
        return self.processor.decode(output[0], skip_special_tokens=True)
```

### Pipeline

```python
def run_geneva_pipeline():
    """
    Geneva pipeline - built on Ray, production-ready by default.

    No custom checkpoint code needed.
    No incremental processing code needed.
    No dependency management code needed.
    """
    conn = geneva.connect("s3://my-bucket/geneva-workspace")
    table = conn.open_table("images")

    table.add_columns({
        "embedding": ComputeEmbeddingGeneva,
        "caption": GenerateCaptionGeneva,
    })

    # Automatic checkpointing at fragment AND batch level
    table.backfill("embedding", batch_size=100, commit_granularity=10)
    table.backfill("caption", batch_size=50, commit_granularity=10)

    print(f"Processed {table.count_rows()} images")
    print(f"Table version: {table.version}")


def run_geneva_incremental():
    """
    Incremental processing in Geneva - no custom code needed.
    """
    conn = geneva.connect("s3://my-bucket/geneva-workspace")
    table = conn.open_table("images")

    # Add new data
    new_images = load_new_images()
    table.add(new_images)

    # Only processes rows where embedding IS NULL
    table.backfill("embedding")


def run_geneva_with_failure_recovery():
    """
    Failure recovery in Geneva.

    Ray Data OSS: Job fails -> start over
    Geneva: Job fails -> resume from exact batch
    """
    conn = geneva.connect("s3://my-bucket/geneva-workspace")
    table = conn.open_table("images")

    table.add_columns({"embedding": ComputeEmbeddingGeneva})

    # If this fails at batch 5000 of 10000...
    try:
        table.backfill("embedding", batch_size=100, commit_granularity=10)
    except Exception:
        pass

    # ...just run again. Geneva resumes from batch 5000.
    table.backfill("embedding", batch_size=100, commit_granularity=10)
```

---

## Side-by-Side Comparison

### Feature Comparison (3-way)

| Feature | Ray Data (OSS) | Ray Data (Anyscale) | Geneva |
|---------|----------------|---------------------|--------|
| Streaming execution | Excellent | Excellent | Uses Ray |
| GPU utilization | 90%+ | 90%+ | Same |
| map_batches API | Native | Native | Abstracted |
| ActorPoolStrategy | Configurable | Configurable | Automatic |
| Framework integration | First-class | First-class | First-class |
| **Checkpointing** | Not included | Limited [1] | Fragment+batch |
| **Incremental processing** | Build yourself | Build yourself | Automatic |
| **Resume from failure** | Start over | Resume [2] | Exact batch |
| **Duplicate prevention** | N/A | Known bug [3] | Guaranteed |
| **Dependency isolation** | runtime_env | runtime_env | Per-UDF auto |
| **Table versioning** | Not included | Not included | Automatic |
| **Operator support** | All operators | Map-only [4] | All operators |
| Pricing | Free (OSS) | Paid | Paid |

**Anyscale Checkpoint Limitations:**
1. Only for read → map ops → write patterns
2. Resumes from checkpoint, but see [3]
3. Known bug: may produce duplicate outputs on retry
4. No joins, aggregations, or shuffles supported

### Code Comparison (3-way)

| Component | Ray Data (OSS) | Ray Data (Anyscale) | Geneva |
|-----------|----------------|---------------------|--------|
| UDF definitions | ~55 lines | ~55 lines | ~45 lines |
| Basic pipeline | ~20 lines | ~20 lines | ~15 lines |
| Checkpoint config | ~40 (custom) | ~10 lines | 0 (included) |
| Incremental processor | ~60 (custom) | ~60 (custom) | 0 (included) |
| Dependency config | ~15 lines | ~15 lines | pyproject.toml |
| id_column requirement | N/A | Required | Not needed |
| **TOTAL (production)** | **~190 lines** | **~160 lines** | **~60 lines** |

### Checkpoint Granularity Comparison

**Scenario: Job fails at row 500,000 of 1,000,000**

| System | Recovery Behavior |
|--------|-------------------|
| **Ray Data (OSS)** | Restart from row 0, reprocess all 500,000 rows, no checkpoint storage needed |
| **Ray Data (Anyscale)** | Restart from last checkpoint, skip processed rows (by id_column), may produce duplicates (known bug), only works with map-based operators, requires unique id_column per row |
| **Geneva** | Restart from fragment 50, batch 15, no reprocessing of completed batches, no duplicates (guaranteed), works with all operators, no id_column required |

---

## Deep Dive: Ray Data LLM Batch Inference (`ray.data.llm`)

Anyscale has invested heavily in [LLM batch inference with `ray.data.llm`](https://docs.anyscale.com/llm/batch-inference/ray-data-llm)
(requires `ray[data,llm]>=2.49.1`). This is Ray Data's strongest story —
a purpose-built API for running vLLM-backed inference at scale.

### What Ray Data LLM Does Well

```python
from ray.data.llm import vLLMEngineProcessorConfig, build_processor

config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen2.5-7B-Instruct",
    batch_size=16,
    engine_kwargs={"max_model_len": 8192},
    accelerator_type="L4",
    concurrency=1,
)

processor = build_processor(
    config,
    preprocess=lambda row: {
        "messages": [{"role": "user", "content": row["text"]}],
        "sampling_params": {"max_tokens": 500},
    },
    postprocess=lambda row: {
        "input": row["text"],
        "output": row["generated_text"],
    },
)

ds = ray.data.read_parquet("s3://my-bucket/data/")
results = processor(ds)
results.write_parquet("s3://my-bucket/output/")
```

**Strong points:**
- **3-stage pipeline**: preprocess → inference → postprocess, each scaling independently
- **vLLM-native**: continuous batching, PagedAttention, optimized throughput
- **Model parallelism**: tensor parallel + pipeline parallel for 200B+ models
- **GPU memory tuning**: batch_size, max_model_len, gpu_memory_utilization knobs
- **Heterogeneous scheduling**: minimize CPU/GPU idle time across stages
- **Model format support**: GPTQ, GGUF, LoRA from HuggingFace Hub

### What Ray Data LLM Still Doesn't Handle

Even with the polished `ray.data.llm` API, the same gaps remain:

| Concern | Ray Data LLM | Geneva |
|---------|-------------|--------|
| Run inference on a dataset | Excellent | Excellent (uses Ray) |
| Checkpoint mid-job | Not included (Anyscale: limited) | Fragment + batch |
| Resume from failure | Restart from scratch | Exact batch resume |
| Process only new rows | Build yourself | Automatic |
| Track which model version produced which output | Build yourself | UDF version tracking |
| Store outputs with versioning | Write to Parquet (no versioning) | Lance (versioned) |
| Re-run when model changes | Manual, full recompute | Automatic (version change detected) |
| Prevent duplicates on retry | Not guaranteed | Guaranteed |

### The Core Difference

Ray Data LLM optimizes for **throughput on a single run**: how fast can I
push data through vLLM on N GPUs? It's excellent at this.

Geneva optimizes for **feature maintenance over time**: what happens on day 2
when new data arrives? Day 30 when the model changes? Day 90 when the job
fails at 3am?

```
Ray Data LLM:
  Day 1: Run inference on 1M rows → write Parquet      ✓ Great
  Day 2: 10K new rows arrive → ???                      Build yourself
  Day 30: Upgrade model → recompute all 1.01M rows      Manual
  Day 90: Job fails at row 500K → restart from row 0    No checkpoints

Geneva:
  Day 1: table.backfill("output") → 1M rows computed    ✓ Same Ray engine
  Day 2: table.backfill("output") → only 10K new rows   ✓ Automatic
  Day 30: Bump UDF version → recompute triggered         ✓ Automatic
  Day 90: Re-run backfill() → resumes from row 500K     ✓ Checkpointed
```

### When Ray Data LLM Is the Right Choice

- **Multi-stage inference pipelines** where you need fine-grained control over
  each stage's scaling (preprocess on CPU, inference on GPU, postprocess on CPU)
- **Custom vLLM tuning** — you want direct access to engine_kwargs, sampling_params,
  tensor/pipeline parallelism configuration
- **One-shot batch jobs** — run once, write output, done. No ongoing maintenance.
- **Integration with Ray Train/Serve** — training → batch inference → serving in
  one Ray ecosystem

### When Geneva Is the Right Choice

- **Ongoing feature columns** — embeddings, captions, classifications that need
  to stay current as data grows
- **Fault tolerance matters** — long-running jobs where restart-from-scratch is
  expensive (days of GPU time)
- **Multiple inference columns** — embedding + caption + classification on the same
  table, each with different models and schedules
- **Team collaboration** — versioned tables, data lineage, reproducibility

### Note on Compatibility

`ray.data.llm` is **not compatible with Geneva** — it owns the full data pipeline
(read → preprocess → inference → postprocess → write), which conflicts with
Geneva's own checkpoint-aware pipeline. However, a Geneva UDF could use **vLLM
directly** (not via `ray.data.llm`) as its inference engine, getting optimized
LLM inference while Geneva handles checkpointing, incremental processing, and
data lifecycle.

---

## When to Use Each

### Ray Data is great for:
- Multi-stage inference pipelines you control end-to-end
- Streaming ML pipelines (train → inference → serve)
- When you need fine-grained control over execution and scaling
- Integration with Ray Train, Ray Serve
- LLM batch inference with vLLM (via `ray.data.llm`)
- Ad-hoc data processing where restart is acceptable

### Ray Data (Anyscale) adds value when:
- You need basic checkpointing for map-only pipelines
- Your data has natural unique IDs
- You're already on Anyscale platform
- Duplicate risk is acceptable or handled downstream

### Geneva is great for:
- Production feature engineering pipelines
- Ongoing incremental updates (daily/hourly)
- When fault tolerance is critical (no duplicates)
- Complex pipelines with joins/aggregations
- Multi-team environments
- Long-running jobs that may fail

---

## The Key Insight

**Geneva is built ON Ray Data, not LIMITED BY it.**

We use Ray Data for distributed execution—just as Kubernetes uses Linux
containers. The value is in the orchestration layer we own and control.

**Geneva (Feature Engineering Platform)**

| What Geneva Adds |
|------------------|
| Fragment & batch checkpointing |
| Automatic incremental processing |
| Per-UDF dependency isolation |
| Table versioning & time travel |
| Guaranteed no duplicates |
| Schema from type hints |

**Our Scheduling Roadmap** (what Ray Data won't prioritize):
- Job-level resource quotas (not just task-level)
- Cost-aware scheduling (spot vs on-demand based on priority)
- Multi-tenant fairness (team A's backfill doesn't starve team B)
- Preemption for high-priority incremental updates

↓ *built on*

**Ray Data (Distributed Compute Engine)**

| Capabilities |
|--------------|
| Streaming execution |
| 90%+ GPU utilization |
| map_batches API |
| ActorPoolStrategy |
| HuggingFace/PyTorch integration |
| Heterogeneous clusters |

**The analogy**: Kubernetes is "just a wrapper" around Linux containers. But the
orchestration layer is where the value lives. Ray Data is our container runtime;
Geneva is the orchestrator purpose-built for feature engineering workloads.

---

## Why "Checkpointing" Matters More Than You Think

**Critics say**: "Checkpointing is overhead for simple jobs."

**Reality**: Checkpoints aren't just for failure recovery—they enable incremental
processing, which is the core value for production pipelines.

```
Day 1: 1M images → compute 1M embeddings
Day 2: 10K new images added

Without Geneva:
  Option A: Recompute all 1.01M (wasteful)
  Option B: Build custom "what's new?" logic

With Geneva:
  table.backfill("embedding")  # Only computes 10K new rows
```

**Ray Data asks**: "How do I efficiently process this data?"
**Geneva asks**: "How do I maintain this feature over time?"

Both use the same distributed execution. Geneva also handles:
- What if the job fails at 3am? → Resume from exact batch
- What if new data arrives tomorrow? → Incremental processing
- What if I need to recompute with a new model? → UDF version tracking
- What if another team needs these features? → Table versioning

---

## Architecture

### Ray Data Architecture (you build the gray boxes)

**Your Code** (you build these):

| Checkpoint Manager | Incremental Processor | Schema Versioning | Dependency Config |
|--------------------|----------------------|-------------------|-------------------|
| (custom) | (custom) | (custom) | (manual) |

↓

**Ray Data**

| Streaming Execution | map_batches API | ActorPool Strategy | GPU Scheduling |
|---------------------|-----------------|-------------------|----------------|

### Geneva Architecture (everything included)

**Your Code**

| UDFs (ML logic) | backfill commands |
|-----------------|-------------------|

↓

**Geneva**

| Fragment Checkpoint | Batch Checkpoint | Incremental Detection | Dependency Isolation | Table Version |
|--------------------|------------------|----------------------|---------------------|---------------|

↓

**Ray Data + Lance Storage**
