# Dask vs Geneva: Image Embedding & Captioning Pipeline

This file compares Dask (parallel computing library) with Geneva
for building production image embedding pipelines.

## Key Insight

Dask scales pandas/numpy workflows to clusters. It's a general-purpose
parallel computing library, not a feature engineering platform.

Think of it as:
- **Dask** = Parallel pandas/numpy for clusters
- **Geneva** = Feature engineering platform (built on Ray + Lance)

## Summary

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

---

## Dask Approach

### Dask Excels At

- Scaling pandas/numpy workflows to clusters
- Lazy execution with task graphs (DAG)
- Familiar pandas-like API
- Integration with scikit-learn, XGBoost
- GPU support via cuDF (RAPIDS)

### You'll Need to Add

- Checkpointing for fault tolerance (Dask has task retry, not checkpointing)
- Incremental processing logic
- GPU scheduling (manual or via RAPIDS)
- Dependency packaging

---

### UDF Definitions

```python
import dask
import dask.dataframe as dd
from dask.distributed import Client, LocalCluster
import pandas as pd
import numpy as np
from typing import Any


def compute_embedding_dask(partition: pd.DataFrame) -> pd.DataFrame:
    """
    CLIP embedding using Dask's map_partitions pattern.

    Key differences from Geneva:
    - Receives entire partition (DataFrame), not single row
    - Must handle GPU device management yourself
    - Model loading happens per-partition (inefficient) or needs workaround
    - Returns DataFrame with new column
    """
    import open_clip
    import torch
    from PIL import Image
    import io

    # Load model - THIS HAPPENS EVERY PARTITION (inefficient)
    # Dask doesn't have stateful actors like Ray
    model, _, preprocess = open_clip.create_model_and_transforms(
        'ViT-B-32', pretrained='laion2b_s34b_b79k'
    )
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = model.to(device).eval()

    embeddings = []
    for img_bytes in partition["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 = preprocess(img).unsqueeze(0).to(device)
            with torch.no_grad():
                embedding = model.encode_image(tensor)
            embeddings.append(embedding.cpu().numpy().flatten())
        except Exception:
            embeddings.append(np.zeros(512, dtype=np.float32))

    partition = partition.copy()
    partition["embedding"] = embeddings
    return partition


def compute_caption_dask(partition: pd.DataFrame) -> pd.DataFrame:
    """BLIP captioning using Dask's map_partitions pattern."""
    from transformers import BlipProcessor, BlipForConditionalGeneration
    import torch
    from PIL import Image
    import io

    # Load model - AGAIN, every partition
    device = "cuda" if torch.cuda.is_available() else "cpu"
    processor = BlipProcessor.from_pretrained(
        "Salesforce/blip-image-captioning-base"
    )
    model = BlipForConditionalGeneration.from_pretrained(
        "Salesforce/blip-image-captioning-base"
    ).to(device).eval()

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

    partition = partition.copy()
    partition["caption"] = captions
    return partition
```

### Stateful Actor Workaround

```python
class DaskModelWorker:
    """
    Workaround for Dask's lack of stateful actors.

    You can use dask.delayed with a class, but it's awkward.
    Or use dask-actor (experimental) or dask-cuda workers.
    """

    def __init__(self):
        self.model = None
        self.preprocess = None

    def load_model(self):
        if self.model is None:
            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, partition: pd.DataFrame) -> pd.DataFrame:
        self.load_model()
        # ... process partition
        return partition
```

### Basic Pipeline (no production ops)

```python
def run_dask_pipeline_basic():
    """
    Basic Dask pipeline - familiar pandas-like API.

    Note: GPU scheduling is manual. Model reloading per partition
    is inefficient without custom workarounds.
    """
    # Start Dask cluster
    cluster = LocalCluster(n_workers=4)
    client = Client(cluster)

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

    # Compute embeddings (model loads per partition - inefficient)
    meta = ddf._meta.copy()
    meta["embedding"] = [np.zeros(512, dtype=np.float32)]

    ddf = ddf.map_partitions(
        compute_embedding_dask,
        meta=meta,
    )

    # Compute captions
    meta["caption"] = ""
    ddf = ddf.map_partitions(
        compute_caption_dask,
        meta=meta,
    )

    # Write results (triggers computation)
    ddf.to_parquet("s3://my-bucket/embeddings/")

    client.close()
    cluster.close()
```

### GPU Pipeline with RAPIDS

```python
def run_dask_pipeline_gpu():
    """
    Dask with GPU via RAPIDS/cuDF.

    Requires dask-cuda and specific cluster setup.
    cuDF doesn't support all pandas operations.
    """
    from dask_cuda import LocalCUDACluster
    from dask.distributed import Client

    # Create GPU cluster
    cluster = LocalCUDACluster()
    client = Client(cluster)

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

    # Note: cuDF for data loading, but model inference
    # still needs custom GPU management
    ddf = ddf.map_partitions(
        compute_embedding_dask,  # Still loads model per partition
        meta={"image": bytes, "embedding": object},
    )

    ddf.to_parquet("s3://my-bucket/embeddings/")

    client.close()
    cluster.close()
```

### Production Pipeline (with custom ops)

For production with Dask, you need to build:
1. **Checkpointing** - Dask has task retry, NOT pipeline checkpointing
2. **Incremental processing** - Process only new data
3. **Model caching** - Avoid reloading per partition
4. **Progress tracking**

**Dask resilience model:**
- If worker fails, tasks are re-executed on other workers
- Scattered data is NOT replicated by default (lost on failure)
- No checkpoint persistence to resume jobs after scheduler restart

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


class DaskCheckpointManager:
    """
    Custom checkpoint manager for Dask pipelines.

    You build this yourself - Dask doesn't include it.

    Dask's built-in resilience:
    - Task retry on worker failure (in-memory only)
    - No persistence across scheduler restarts
    - Scattered data can be lost
    """

    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_partitions": [], "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_partition_processed(self, partition_id: str, row_count: int):
        self.state["processed_partitions"].append(partition_id)
        self.state["rows_processed"] += row_count
        self._save_state()

    def is_partition_processed(self, partition_id: str) -> bool:
        return partition_id in self.state["processed_partitions"]

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


class DaskIncrementalProcessor:
    """
    Incremental processing wrapper for Dask.

    Implements partition-level checkpointing since Dask doesn't
    support row-level checkpointing.
    """

    def __init__(
        self,
        input_path: str,
        output_path: str,
        checkpoint_dir: str,
        n_workers: int = 4,
    ):
        self.input_path = input_path
        self.output_path = output_path
        self.n_workers = n_workers
        self.checkpoint = DaskCheckpointManager(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_partition_processed(f)]

    def run(self):
        """Run incremental processing with checkpointing."""
        cluster = LocalCluster(n_workers=self.n_workers)
        client = Client(cluster)

        try:
            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
            for i, file_path in enumerate(pending_files):
                print(f"Processing [{i+1}/{len(pending_files)}]: {file_path}")

                try:
                    # Read single file
                    ddf = dd.read_parquet(file_path)

                    # Apply transformations
                    meta = ddf._meta.copy()
                    meta["embedding"] = [np.zeros(512)]
                    meta["caption"] = ""

                    ddf = ddf.map_partitions(compute_embedding_dask, meta=meta)
                    ddf = ddf.map_partitions(compute_caption_dask, meta=meta)

                    # Compute and write
                    result = ddf.compute()
                    row_count = len(result)

                    output_file = f"{self.output_path}/batch_{i:06d}.parquet"
                    result.to_parquet(output_file)

                    self.checkpoint.mark_partition_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']}")

        finally:
            client.close()
            cluster.close()


def run_dask_pipeline_production():
    """Production Dask pipeline with custom checkpointing."""
    processor = DaskIncrementalProcessor(
        input_path="s3://my-bucket/images",
        output_path="s3://my-bucket/embeddings",
        checkpoint_dir="/tmp/dask_checkpoints",
    )
    processor.run()
```

### Dependency Management

```python
def configure_dask_with_dependencies():
    """
    Configure Dask cluster with required dependencies.

    Options:
    1. Pre-install on all workers (conda/pip)
    2. Use dask-cloudprovider with custom images
    3. Scatter dependencies (brittle)

    You manage this yourself - Dask doesn't auto-package deps.
    """
    from dask.distributed import Client

    # Workers must have same packages pre-installed
    # No automatic dependency packaging like Geneva
    client = Client("scheduler-address:8786")

    # Check worker packages match client
    def check_versions():
        import torch
        import transformers
        return {
            "torch": torch.__version__,
            "transformers": transformers.__version__,
        }

    versions = client.run(check_versions)
    print(f"Worker versions: {versions}")
```

---

## Geneva Approach

Geneva provides:
- Stateful UDFs (model loads once per worker, not per partition)
- 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 Dask:
    - setup() loads model ONCE per worker (not per partition)
    - Process single row, Geneva handles batching
    - Type annotation -> schema (no meta dict)
    - GPU scheduling automatic
    """

    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 - stateful UDFs, production-ready by default.

    No model reloading per partition.
    No custom checkpoint code needed.
    No incremental processing 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")
```

---

## Side-by-Side Comparison

### Feature Comparison

| Feature | Dask | Geneva |
|---------|------|--------|
| Pandas-like API | Native | Table-centric |
| Lazy execution (DAG) | Yes | Eager backfill |
| GPU support | Via RAPIDS/cuDF | num_gpus=1 |
| Scikit-learn integration | dask-ml | Via UDFs |
| Distributed scheduler | dask.distributed | Ray |
| **Stateful UDFs** | Model reloads [1] | setup() once |
| **Checkpointing** | Task retry only | Fragment + batch |
| **Incremental processing** | Build yourself | Automatic |
| **Resume from failure** | Restart job [2] | Exact batch |
| **Dependency isolation** | Pre-install | Per-UDF automatic |
| **Table versioning** | Not included | Automatic |
| **Scattered data resilience** | Lost on failure | Checkpointed |

**Dask limitations:**
1. map_partitions reloads model per partition (no actors)
2. Scheduler restart loses all progress

### Code Comparison

| Component | Dask | Geneva |
|-----------|------|--------|
| UDF definitions | ~65 lines | ~45 lines |
| Basic pipeline | ~25 lines | ~15 lines |
| Checkpoint manager | ~45 lines (custom) | 0 (included) |
| Incremental processor | ~70 lines (custom) | 0 (included) |
| Model caching workaround | ~30 lines | 0 (built-in) |
| Dependency config | Pre-install req'd | pyproject.toml |
| **TOTAL (production)** | **~235 lines** | **~60 lines** |

### Model Loading Comparison

**Dask (map_partitions):**
- 10 partitions × 4 workers = model loads 10 times
- Each partition: Load model → process → unload
- GPU memory thrashing, slow startup per partition

**Dask (workaround with delayed):**
- Complex, requires manual actor management
- Not officially supported pattern
- Debugging is difficult

**Geneva (stateful UDFs):**
- 4 workers = model loads 4 times (once per worker)
- Worker setup() → process all batches → teardown
- Efficient GPU utilization

### Resilience Comparison

**Dask:**
- Worker fails → tasks re-scheduled to other workers
- Scheduler fails → ALL progress lost, restart from scratch
- Scattered data not replicated → data lost on worker failure
- No checkpoint persistence

**Geneva:**
- Worker fails → tasks re-scheduled (same as Dask)
- Scheduler fails → resume from last checkpoint
- All data checkpointed to Lance storage
- Fragment + batch level checkpointing

---

## When to Use Each

### Dask is great for:
- Scaling pandas/numpy workflows to clusters
- Interactive data exploration
- Scikit-learn and XGBoost distributed training
- When you have existing pandas code to scale
- Short-running jobs where restart is acceptable

### Geneva is great for:
- Production feature engineering pipelines
- GPU-heavy model inference at scale
- Ongoing incremental updates (daily/hourly)
- When fault tolerance is critical
- Long-running jobs that may fail
- Multi-team environments

---

## The Key Insight

**Dask scales your pandas code. Geneva manages your features.**

**Dask asks**: "How do I run this pandas code on a cluster?"
**Geneva asks**: "How do I maintain this feature over time?"

Dask is a parallel execution engine.
Geneva is a feature engineering platform.

---

## Architecture

### Dask Architecture (you build the gray boxes)

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

| Checkpoint Manager | Incremental Processor | Model Caching | Dependency Pre-install |
|--------------------|----------------------|---------------|------------------------|
| (custom) | (custom) | (workaround) | (manual) |

↓

**Dask**

| DataFrame API | map_partitions | Task Scheduler | Distributed Workers |
|---------------|----------------|----------------|---------------------|

### Geneva Architecture (everything included)

**Your Code**

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

↓

**Geneva**

| Stateful UDFs | Fragment Checkpoint | Incremental Detection | Dependency Isolation | Table Version |
|---------------|--------------------|-----------------------|---------------------|---------------|

↓

**Ray + Lance Storage**
