# GPU Pipelining - Phase 1: DataLoader Actors

**Status:** 📋 Proposed / Planned

## Goal

Decouple data loading from GPU execution so GPU actors always have work ready. This is a "quick win" that provides immediate benefits without requiring UDF changes.

## Architecture

```
┌─────────────────────┐                              ┌─────────────────────┐
│  DataLoaderActor    │ ──→ Ray ObjectStore ───────→ │  GPUApplierActor    │
│  (CPU-only)         │     (prefetched batches)     │  (num_gpus=1)       │
│  - Lance I/O        │                              │  - UDF execution    │
│  - Decompression    │     ObjectRef queue          │  - No I/O waits     │
└─────────────────────┘     (bounded depth)          └─────────────────────┘
         ×N loaders                                           ×M GPUs
```

## Key Components

### 1. PreloadedReadTask

A wrapper that implements the `ReadTask` interface but returns pre-loaded batches from Ray's object store instead of reading from Lance.

```python
@attrs.define
class PreloadedReadTask(ReadTask):
    """A ReadTask where batches have been pre-loaded into Ray object store."""

    original_task: ReadTask
    batch_refs: list[ObjectRef]

    def to_batches(self, *, batch_size=DEFAULT_CHECKPOINT_ROWS) -> Iterator[pa.RecordBatch]:
        """Yield pre-loaded batches from object store."""
        for ref in self.batch_refs:
            yield ray.get(ref)  # Fast - already in object store

    def checkpoint_key(self) -> str:
        return self.original_task.checkpoint_key()

    # ... other delegated methods
```

**Why this design?**
- Implements existing `ReadTask` interface
- No changes needed to `ApplierActor` or `CheckpointingApplier`
- Transparent drop-in replacement

### 2. DataLoaderActor

A Ray actor that runs on CPU and performs I/O operations (reading from Lance, decompression).

```python
@ray.remote(num_cpus=1)
class DataLoaderActor:
    """Prefetches batches from Lance into Ray object store."""

    def load(self, read_task: ReadTask, batch_size: int) -> tuple[ReadTask, list[ObjectRef]]:
        """Load all batches from a ReadTask into Ray object store."""
        batch_refs = []
        for batch in read_task.to_batches(batch_size=batch_size):
            ref = ray.put(batch)  # Store in object store
            batch_refs.append(ref)
        return read_task, batch_refs
```

**Design decisions:**
- `num_cpus=1`: Lightweight, I/O-bound work
- Returns original task + ObjectRefs for creating `PreloadedReadTask`
- Uses `ray.put()` for efficient zero-copy sharing

### 3. LoaderPool

Manages a pool of DataLoaderActors with a bounded prefetch queue.

```python
@attrs.define
class LoaderPool:
    num_loaders: int = 4
    prefetch_depth: int = 8

    def prefetch_iter(self, read_tasks: Iterator[ReadTask], batch_size: int) -> Iterator[PreloadedReadTask]:
        """Iterate over prefetched ReadTasks."""
        # Maintains prefetch queue by submitting new load tasks
        # as PreloadedReadTasks are consumed
```

**Key behaviors:**
- Primes prefetch queue on initialization
- Maintains `prefetch_depth` in-flight loads
- Handles actor failures with automatic restart
- Provides iterator interface for easy integration

## Integration with Pipeline

In `ColumnAddPipelineJob.run()`:

```python
def run(self) -> None:
    # Setup GPU pipelining if enabled
    loader_pool: LoaderPool | None = None
    if self._should_use_gpu_pipelining():
        loader_pool = LoaderPool(
            num_loaders=self.config.num_loader_actors,
            prefetch_depth=self.config.prefetch_depth,
        )

    # Determine input tasks
    if loader_pool is not None:
        # GPU pipelining: prefetch batches and wrap in PreloadedReadTask
        input_tasks = loader_pool.prefetch_iter(plans, batch_size)
    else:
        # Standard path: ReadTasks loaded on-demand
        input_tasks = plans

    # Rest of pipeline unchanged
    applier_iter = pool.map_unordered(
        lambda actor, value: actor.run.remote(value),
        input_tasks,
    )
```

## Configuration

```python
# In JobConfig
enable_gpu_pipelining: bool = False  # Off by default
num_loader_actors: int = 4           # Loader actor count
prefetch_depth: int = 8              # Batches to prefetch

# Via backfill API
table.backfill(
    "column_name",
    enable_gpu_pipelining=True,
    num_loader_actors=4,
    prefetch_depth=8,
)
```

## Files Modified

| File | Changes |
|------|---------|
| `src/geneva/runners/ray/loader.py` | **NEW** - DataLoaderActor, LoaderPool, PreloadedReadTask |
| `src/geneva/runners/ray/pipeline.py` | Added pipelining logic to ColumnAddPipelineJob |
| `src/geneva/jobs/config.py` | Added configuration options |
| `src/geneva/table.py` | Updated docstrings for backfill() |

## How to Validate

1. **Test with BLIP caption UDF:**
   ```python
   table.backfill("caption", enable_gpu_pipelining=True)
   ```

2. **Monitor GPU utilization:**
   ```bash
   watch -n 1 nvidia-smi
   ```

3. **Compare throughput:**
   - Run with `enable_gpu_pipelining=False` (baseline)
   - Run with `enable_gpu_pipelining=True` (pipelined)
   - Measure rows/second improvement

## Tuning Guidelines

| Parameter | Recommendation |
|-----------|---------------|
| `num_loader_actors` | Start with 2-4, increase if GPU still starved |
| `prefetch_depth` | 2-3× number of GPU actors |
| Memory impact | `prefetch_depth × batch_size × row_size` |

## Limitations

1. **Memory overhead**: Pre-loaded batches consume Ray object store memory
2. **No CPU preprocessing**: Phase 1 doesn't separate CPU preprocessing from GPU inference
3. **No CUDA stream overlap**: Data transfer and compute not overlapped within GPU

## Future Improvements (Phase 2, 3)

- Add CPU preprocessing pipeline for image decode, tokenization
- Add CUDA stream optimization for overlapping data transfer with compute
