# Traditional vs Geneva: Image Embedding & Captioning Pipeline

This file demonstrates the dramatic difference between building a production
image embedding pipeline the traditional way vs using Geneva.

## Summary

| Approach | Lines of Code | Time to Production |
|----------|---------------|-------------------|
| Traditional | 250+ lines across 5 files | 2-4 weeks |
| Geneva | 60 lines across 2 files | Hours |

---

## Traditional Approach

### Files Needed
1. `infrastructure/Dockerfile`
2. `infrastructure/ray-cluster.yaml`
3. `src/checkpoint_store.py`
4. `src/ray_actors.py`
5. `src/pipeline.py`

### Teams Involved
- **Platform/Infra team**: Dockerfile, ray-cluster.yaml, deployment
- **ML Engineer**: checkpoint_store.py, ray_actors.py, pipeline.py

### Timeline
2-4 weeks for initial pipeline, ongoing maintenance

---

### infrastructure/Dockerfile

```dockerfile
FROM rayproject/ray-ml:2.44.0-py310-gpu

# Install dependencies - hope there are no conflicts!
RUN pip install \
    torch==2.1.0 \
    torchvision==0.16.0 \
    open-clip-torch==2.24.0 \
    transformers==4.36.0 \
    Pillow==10.1.0 \
    pyarrow==14.0.1

# Copy application code
COPY src/ /app/src/
WORKDIR /app

# Build and push to registry
# docker build -t gcr.io/my-project/embedding-pipeline:v1 .
# docker push gcr.io/my-project/embedding-pipeline:v1
```

### infrastructure/ray-cluster.yaml

```yaml
apiVersion: ray.io/v1
kind: RayCluster
metadata:
  name: embedding-cluster
spec:
  headGroupSpec:
    rayStartParams:
      dashboard-host: '0.0.0.0'
    template:
      spec:
        containers:
        - name: ray-head
          image: gcr.io/my-project/embedding-pipeline:v1
          resources:
            limits:
              cpu: "4"
              memory: "8Gi"
  workerGroupSpecs:
  - groupName: gpu-workers
    replicas: 4
    minReplicas: 1
    maxReplicas: 10
    rayStartParams: {}
    template:
      spec:
        containers:
        - name: ray-worker
          image: gcr.io/my-project/embedding-pipeline:v1
          resources:
            limits:
              cpu: "4"
              memory: "16Gi"
              nvidia.com/gpu: "1"
```

### src/checkpoint_store.py

```python
import hashlib
import json
from pathlib import Path
from typing import Optional
import threading


class CheckpointStore:
    """
    Track which images have been processed to enable resume after failures.

    This is a simplified version. Production needs:
    - Distributed locking for concurrent writers
    - Cloud storage backend (S3/GCS)
    - Atomic updates
    - Garbage collection for old checkpoints
    """

    def __init__(self, checkpoint_dir: str):
        self.checkpoint_dir = Path(checkpoint_dir)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        self._lock = threading.Lock()
        self._processed: set[str] = self._load_state()
        self._dirty = False

    def _load_state(self) -> set[str]:
        state_file = self.checkpoint_dir / "processed.json"
        if state_file.exists():
            try:
                return set(json.loads(state_file.read_text()))
            except json.JSONDecodeError:
                # Corrupted checkpoint file - start fresh or fail?
                return set()
        return set()

    def _save_state(self) -> None:
        state_file = self.checkpoint_dir / "processed.json"
        temp_file = self.checkpoint_dir / "processed.json.tmp"
        # Atomic write
        temp_file.write_text(json.dumps(list(self._processed)))
        temp_file.rename(state_file)
        self._dirty = False

    def mark_processed(self, image_id: str) -> None:
        with self._lock:
            self._processed.add(image_id)
            self._dirty = True
            # Periodic save to avoid losing too much progress
            if len(self._processed) % 100 == 0:
                self._save_state()

    def is_processed(self, image_id: str) -> bool:
        return image_id in self._processed

    def flush(self) -> None:
        with self._lock:
            if self._dirty:
                self._save_state()

    def get_stats(self) -> dict:
        return {
            "processed_count": len(self._processed),
            "checkpoint_dir": str(self.checkpoint_dir),
        }
```

### src/ray_actors.py

```python
import ray
from ray.util.actor_pool import ActorPool
from typing import Optional
import logging

logger = logging.getLogger(__name__)


@ray.remote(num_gpus=1, max_restarts=3)
class EmbeddingActor:
    """
    Ray actor for GPU-based CLIP embedding computation.

    Handles:
    - Lazy model loading (load once per actor)
    - GPU memory management
    - Error isolation (one bad image doesn't crash the actor)
    """

    def __init__(self, model_name: str = 'ViT-B-32', pretrained: str = 'laion2b_s34b_b79k'):
        self.model_name = model_name
        self.pretrained = pretrained
        self.model = None
        self.preprocess = None
        self._initialized = False

    def _load_model(self) -> None:
        if self._initialized:
            return

        import open_clip
        import torch

        logger.info(f"Loading CLIP model {self.model_name}...")
        self.model, _, self.preprocess = open_clip.create_model_and_transforms(
            self.model_name,
            pretrained=self.pretrained
        )
        self.model = self.model.cuda().eval()
        self._initialized = True
        logger.info("CLIP model loaded successfully")

    def compute_embedding(self, image_bytes: bytes, image_id: str) -> dict:
        """Compute embedding for a single image."""
        self._load_model()

        from PIL import Image
        import torch
        import io

        try:
            # Decode image
            image = Image.open(io.BytesIO(image_bytes)).convert('RGB')

            # Preprocess and compute embedding
            image_tensor = self.preprocess(image).unsqueeze(0).cuda()
            with torch.no_grad():
                embedding = self.model.encode_image(image_tensor)
                embedding = embedding.cpu().numpy().flatten().tolist()

            return {
                "id": image_id,
                "embedding": embedding,
                "success": True,
                "error": None,
            }

        except Exception as e:
            logger.error(f"Error processing image {image_id}: {e}")
            return {
                "id": image_id,
                "embedding": None,
                "success": False,
                "error": str(e),
            }

    def health_check(self) -> bool:
        """Verify actor is healthy."""
        return self._initialized


@ray.remote(num_gpus=1, max_restarts=3)
class CaptionActor:
    """
    Ray actor for GPU-based BLIP captioning.

    Note: This uses different dependencies than EmbeddingActor.
    In a real setup, you might need separate Docker images or
    careful dependency management to avoid conflicts.
    """

    def __init__(self, model_name: str = "Salesforce/blip-image-captioning-base"):
        self.model_name = model_name
        self.model = None
        self.processor = None
        self._initialized = False

    def _load_model(self) -> None:
        if self._initialized:
            return

        from transformers import BlipProcessor, BlipForConditionalGeneration
        import torch

        logger.info(f"Loading BLIP model {self.model_name}...")
        self.processor = BlipProcessor.from_pretrained(self.model_name)
        self.model = BlipForConditionalGeneration.from_pretrained(self.model_name)
        self.model = self.model.cuda().eval()
        self._initialized = True
        logger.info("BLIP model loaded successfully")

    def generate_caption(self, image_bytes: bytes, image_id: str) -> dict:
        """Generate caption for a single image."""
        self._load_model()

        from PIL import Image
        import torch
        import io

        try:
            # Decode image
            image = Image.open(io.BytesIO(image_bytes)).convert('RGB')

            # Generate caption
            inputs = self.processor(image, return_tensors="pt")
            inputs = {k: v.to("cuda") for k, v in inputs.items()}

            with torch.no_grad():
                output = self.model.generate(**inputs, max_new_tokens=50)

            caption = self.processor.decode(output[0], skip_special_tokens=True)

            return {
                "id": image_id,
                "caption": caption,
                "success": True,
                "error": None,
            }

        except Exception as e:
            logger.error(f"Error generating caption for {image_id}: {e}")
            return {
                "id": image_id,
                "caption": None,
                "success": False,
                "error": str(e),
            }
```

### src/pipeline.py

```python
import pyarrow as pa
import pyarrow.parquet as pq
import ray
from ray.util.actor_pool import ActorPool
import time
import logging
from typing import Optional
from dataclasses import dataclass

# Assuming these are importable
# from checkpoint_store import CheckpointStore
# from ray_actors import EmbeddingActor, CaptionActor

logger = logging.getLogger(__name__)


@dataclass
class PipelineConfig:
    """Configuration for the embedding pipeline."""
    input_path: str
    output_path: str
    checkpoint_dir: str
    batch_size: int = 100
    num_embedding_actors: int = 4
    num_caption_actors: int = 4
    max_retries: int = 3
    retry_delay_base: float = 2.0
    commit_interval: int = 1000


class EmbeddingPipeline:
    """
    Main pipeline orchestration for computing embeddings and captions.

    This class handles:
    - Ray cluster connection
    - Actor pool management
    - Batch processing with retries
    - Checkpoint management
    - Progress tracking
    - Result writing

    Missing (would need in production):
    - Metrics/monitoring integration
    - Dead letter queue for failed items
    - Graceful shutdown handling
    - Schema versioning
    - Incremental refresh logic
    - Concurrent write handling
    """

    def __init__(self, config: PipelineConfig):
        self.config = config
        self.checkpoint: Optional[CheckpointStore] = None
        self.embedding_pool: Optional[ActorPool] = None
        self.caption_pool: Optional[ActorPool] = None
        self._results: list[dict] = []
        self._failed: list[str] = []

    def _init_ray(self) -> None:
        """Connect to Ray cluster."""
        if not ray.is_initialized():
            ray.init(address="auto")
        logger.info(f"Connected to Ray cluster: {ray.cluster_resources()}")

    def _init_actors(self) -> None:
        """Create actor pools."""
        logger.info(f"Creating {self.config.num_embedding_actors} embedding actors...")
        embedding_actors = [
            EmbeddingActor.remote()
            for _ in range(self.config.num_embedding_actors)
        ]
        self.embedding_pool = ActorPool(embedding_actors)

        logger.info(f"Creating {self.config.num_caption_actors} caption actors...")
        caption_actors = [
            CaptionActor.remote()
            for _ in range(self.config.num_caption_actors)
        ]
        self.caption_pool = ActorPool(caption_actors)

    def _load_pending_images(self) -> list[dict]:
        """Load images that haven't been processed yet."""
        logger.info(f"Loading data from {self.config.input_path}...")
        table = pq.read_table(self.config.input_path)

        pending = []
        skipped = 0

        for i in range(len(table)):
            row_id = str(table['id'][i].as_py())
            if self.checkpoint.is_processed(row_id):
                skipped += 1
                continue
            pending.append({
                'id': row_id,
                'image': table['image'][i].as_py(),
            })

        logger.info(f"Found {len(pending)} pending images ({skipped} already processed)")
        return pending

    def _process_batch_embeddings(self, batch: list[dict]) -> list[dict]:
        """Process a batch of images through embedding actors."""
        results = []

        # Submit all tasks
        for row in batch:
            self.embedding_pool.submit(
                lambda actor, r: actor.compute_embedding.remote(r['image'], r['id']),
                row
            )

        # Collect results with retries
        while self.embedding_pool.has_next():
            try:
                result = self.embedding_pool.get_next(timeout=60)

                if result['success']:
                    results.append(result)
                else:
                    # Retry logic
                    row = next(r for r in batch if r['id'] == result['id'])
                    retry_result = self._retry_with_backoff(
                        lambda: ray.get(
                            self.embedding_pool._idle_actors[0].compute_embedding.remote(
                                row['image'], row['id']
                            ),
                            timeout=60
                        )
                    )
                    if retry_result and retry_result['success']:
                        results.append(retry_result)
                    else:
                        self._failed.append(result['id'])

            except Exception as e:
                logger.error(f"Batch processing error: {e}")

        return results

    def _process_batch_captions(self, batch: list[dict], embeddings: list[dict]) -> None:
        """Add captions to embedding results."""
        # Create lookup for results
        result_map = {r['id']: r for r in embeddings}

        # Submit caption tasks
        for row in batch:
            if row['id'] in result_map:  # Only caption successful embeddings
                self.caption_pool.submit(
                    lambda actor, r: actor.generate_caption.remote(r['image'], r['id']),
                    row
                )

        # Collect results
        while self.caption_pool.has_next():
            try:
                result = self.caption_pool.get_next(timeout=120)
                if result['success'] and result['id'] in result_map:
                    result_map[result['id']]['caption'] = result['caption']
            except Exception as e:
                logger.error(f"Caption error: {e}")

    def _retry_with_backoff(self, func, max_attempts: Optional[int] = None) -> Optional[dict]:
        """Retry a function with exponential backoff."""
        max_attempts = max_attempts or self.config.max_retries

        for attempt in range(max_attempts):
            try:
                return func()
            except Exception as e:
                if attempt == max_attempts - 1:
                    logger.error(f"Failed after {max_attempts} attempts: {e}")
                    return None
                delay = self.config.retry_delay_base ** attempt
                logger.warning(f"Attempt {attempt + 1} failed, retrying in {delay}s...")
                time.sleep(delay)

        return None

    def _write_results(self, results: list[dict]) -> None:
        """Write results to output path."""
        if not results:
            return

        schema = pa.schema([
            ('id', pa.string()),
            ('embedding', pa.list_(pa.float32(), 512)),
            ('caption', pa.string()),
        ])

        table = pa.table({
            'id': pa.array([r['id'] for r in results]),
            'embedding': pa.array([r.get('embedding') for r in results]),
            'caption': pa.array([r.get('caption') for r in results]),
        }, schema=schema)

        # NOTE: This overwrites! Production needs append logic with versioning
        pq.write_table(table, self.config.output_path)
        logger.info(f"Wrote {len(results)} results to {self.config.output_path}")

    def run(self) -> dict:
        """Execute the full pipeline."""
        start_time = time.time()

        # Initialize
        self._init_ray()
        self.checkpoint = CheckpointStore(self.config.checkpoint_dir)
        self._init_actors()

        # Load pending work
        pending = self._load_pending_images()
        total = len(pending)

        if total == 0:
            logger.info("No pending images to process")
            return {"processed": 0, "failed": 0, "duration": 0}

        # Process in batches
        for batch_start in range(0, total, self.config.batch_size):
            batch_end = min(batch_start + self.config.batch_size, total)
            batch = pending[batch_start:batch_end]

            # Compute embeddings
            embeddings = self._process_batch_embeddings(batch)

            # Compute captions
            self._process_batch_captions(batch, embeddings)

            # Checkpoint and accumulate results
            for result in embeddings:
                self.checkpoint.mark_processed(result['id'])
                self._results.append(result)

            # Periodic write
            if len(self._results) >= self.config.commit_interval:
                self._write_results(self._results)
                self._results = []

            # Progress
            processed = batch_end
            pct = 100 * processed / total
            logger.info(f"Progress: {processed}/{total} ({pct:.1f}%)")

        # Final write
        if self._results:
            self._write_results(self._results)

        # Cleanup
        self.checkpoint.flush()

        duration = time.time() - start_time

        return {
            "processed": total - len(self._failed),
            "failed": len(self._failed),
            "failed_ids": self._failed,
            "duration": duration,
        }


def main():
    """Entry point for traditional pipeline."""
    config = PipelineConfig(
        input_path="s3://my-bucket/images.parquet",
        output_path="s3://my-bucket/embeddings.parquet",
        checkpoint_dir="/tmp/checkpoints",
        batch_size=100,
        num_embedding_actors=4,
        num_caption_actors=4,
    )

    pipeline = EmbeddingPipeline(config)
    result = pipeline.run()

    print(f"Pipeline complete: {result}")
```

---

## Geneva Approach

### Files Needed
1. `udfs/pyproject.toml` (admin creates once)
2. `udfs/embeddings.py` (researcher writes)
3. `pipeline.py` (researcher writes)

### Roles
- **Admin (one-time)**: Cluster setup, connection config
- **Researcher**: UDF logic and pipeline execution

### Timeline
Hours to production

---

### Admin Setup (one-time)

#### admin/cluster_config.py

Admin creates cluster configuration once. Researchers reuse it.

```python
import geneva
from geneva.cluster import GenevaClusterBuilder, WorkerGroupBuilder

def setup_workspace():
    """One-time workspace and cluster setup by platform admin."""

    # Define GPU cluster for ML workloads
    gpu_cluster = (
        GenevaClusterBuilder.create("ml-gpu-cluster")
        .head_group(
            image="rayproject/ray-ml:2.44.0-py310-gpu",
            cpus=4,
            memory="8Gi"
        )
        .add_worker_group(
            WorkerGroupBuilder.gpu_worker(cpus=4, memory="16Gi", gpus=1)
            .image("rayproject/ray-ml:2.44.0-py310-gpu")
            .instances(min=2, max=10)  # Autoscaling
        )
        .build()
    )

    # Register cluster in workspace
    conn = geneva.connect("s3://company-data/geneva-workspace")
    conn.define_cluster("ml-gpu-cluster", gpu_cluster)

    print("Cluster 'ml-gpu-cluster' registered. Researchers can now use it.")
```

---

### Researcher Workflow

#### udfs/pyproject.toml

```toml
[project]
name = "image-features"
version = "0.1.0"
dependencies = [
    "open-clip-torch>=2.24.0",
    "transformers>=4.36.0",
    "torch>=2.1.0",
    "Pillow>=10.0.0",
]

# That's it. Geneva handles:
# - Dependency resolution
# - Environment packaging
# - Shipping to workers
# - Isolation between UDFs
```

#### udfs/embeddings.py

Researcher writes ONLY the ML logic. No infrastructure code.

```python
import geneva


@geneva.udf(num_gpus=1, version="0.1")
class ComputeEmbedding:
    """
    CLIP embedding computation.

    What the researcher writes:
    - Model loading in setup()
    - Inference logic in __call__()
    - Type annotations for schema

    What Geneva handles automatically:
    - Ray actor creation
    - GPU scheduling
    - Model caching across batches
    - Error handling and retries
    - Checkpointing
    - Progress tracking
    """

    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 GenerateCaption:
    """
    BLIP image captioning.

    Note: Different dependencies than CLIP - Geneva isolates automatically.
    No Docker image conflicts to resolve.
    """

    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.to("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.py

Researcher's complete pipeline script.

```python
import geneva
from udfs.embeddings import ComputeEmbedding, GenerateCaption


def run_pipeline():
    """
    Complete pipeline in 15 lines.

    Compare to 250+ lines in traditional approach.
    """
    # Connect to workspace (cluster already configured by admin)
    conn = geneva.connect("s3://company-data/geneva-workspace")

    # Open existing table with images
    table = conn.open_table("product_images")

    # Add computed columns (schema inferred from type annotations)
    table.add_columns({
        "embedding": ComputeEmbedding,
        "caption": GenerateCaption,
    })

    # Run on GPU cluster - everything else is automatic
    with conn.context(cluster="ml-gpu-cluster"):
        # Embeddings first
        table.backfill("embedding", batch_size=100, commit_granularity=10)

        # Then captions
        table.backfill("caption", batch_size=50, commit_granularity=10)

    # Done. Results are:
    # - Checkpointed (resume from failure automatically)
    # - Versioned (time-travel to any point)
    # - Queryable immediately

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


def run_pipeline_async():
    """
    Async version with progress monitoring.
    """
    conn = geneva.connect("s3://company-data/geneva-workspace")
    table = conn.open_table("product_images")

    table.add_columns({
        "embedding": ComputeEmbedding,
        "caption": GenerateCaption,
    })

    with conn.context(cluster="ml-gpu-cluster"):
        # Start async backfill
        embedding_job = table.backfill_async(
            "embedding",
            batch_size=100,
            commit_granularity=10
        )

        # Monitor progress
        while not embedding_job.done():
            table.checkout_latest()
            computed = table.count_rows("embedding IS NOT NULL")
            total = table.count_rows()
            print(f"Embedding progress: {computed}/{total}")
            time.sleep(10)

        # Check result
        result = embedding_job.result()
        print(f"Embedding complete: {result}")

        # Now captions
        caption_job = table.backfill_async("caption", batch_size=50)
        caption_job.result()  # Wait for completion

    print("Pipeline complete!")
```

---

## Comparison Summary

### Traditional Approach

```
├── infrastructure/
│   ├── Dockerfile (20 lines)
│   └── ray-cluster.yaml (40 lines)
├── src/
│   ├── checkpoint_store.py (60 lines)
│   ├── ray_actors.py (100 lines)
│   └── pipeline.py (150 lines)
└── TOTAL: 370+ lines, 5 files, 2+ weeks
```

**Still missing:**
- Monitoring/alerting
- Dead letter queue
- Schema versioning
- Incremental refresh
- Concurrent write handling
- Graceful shutdown

### Geneva Approach

```
├── admin/ (one-time setup)
│   └── cluster_config.py (25 lines)
├── udfs/
│   ├── pyproject.toml (10 lines)
│   └── embeddings.py (45 lines)
└── pipeline.py (20 lines)
└── TOTAL: 100 lines, 4 files, hours
```

**Automatically included:**
- Checkpointing & resume
- Retry logic
- Progress tracking
- Schema management
- Versioning
- Incremental refresh
- GPU scheduling
- Dependency isolation
- Error handling

---

## Role Separation

### Traditional

- **Platform Team**: Dockerfile, K8s manifests, Ray cluster, monitoring
- **ML Engineer**: Everything else (200+ lines of infrastructure code)
- **Ratio**: 80% infrastructure, 20% ML logic

### Geneva

- **Admin (one-time)**: 25 lines cluster config
- **Researcher**: 65 lines (45 UDF + 20 pipeline)
- **Ratio**: 20% infrastructure, 80% ML logic
