# GPU Pipelining - Phase 2: CPU Preprocessing Pipeline

**Status:** 🔮 Future / Not Implemented

## Goal

For UDFs with heavy CPU preprocessing (image decoding, tokenization), separate CPU work from GPU work to maximize both CPU and GPU utilization.

## Problem

Many GPU UDFs spend significant time on CPU-bound preprocessing:

```python
@udf(num_gpus=1)
class ImageCaptionUDF:
    def __call__(self, image_bytes: bytes) -> str:
        # CPU-bound preprocessing (runs on GPU actor, wastes GPU time)
        image = decode_image(image_bytes)     # ~10ms
        tensor = preprocess(image)             # ~5ms
        tensor = tensor.to("cuda")             # ~2ms

        # GPU inference (actual GPU work)
        with torch.no_grad():
            output = self.model(tensor)        # ~20ms

        return decode_output(output)
```

The GPU actor spends ~17ms on CPU work per batch, leaving GPU idle.

## Proposed Architecture

```
DataLoaderActor ──→ CPUPreprocessorActor ──→ TensorQueue ──→ GPUApplierActor
     (I/O)              (decode, prep)         (staging)       (inference)
      ×N                     ×M                                    ×P
```

### Three-Stage Pipeline

1. **DataLoaderActors** (existing from Phase 1): Read raw data from Lance
2. **CPUPreprocessorActors** (new): Decode images, tokenize text, prepare tensors
3. **GPUApplierActors** (modified): Only run model inference

## Proposed UDF Interface Extension

### Option A: Separate Methods

```python
@udf(num_gpus=1)
class PipelinedImageCaption:
    """UDF with explicit preprocessing step."""

    def preprocess(self, batch: pa.RecordBatch) -> PreprocessedBatch:
        """CPU-bound preprocessing (runs on CPUPreprocessorActor)."""
        images = [decode_image(b) for b in batch["image"]]
        tensors = torch.stack([to_tensor(img) for img in images])
        return PreprocessedBatch(tensors=tensors, metadata=batch)

    def __call__(self, preprocessed: PreprocessedBatch) -> pa.Array:
        """GPU inference only (runs on GPUApplierActor)."""
        tensors = preprocessed.tensors.to("cuda")
        with torch.no_grad():
            outputs = self.model(tensors)
        return pa.array(decode_outputs(outputs))
```

### Option B: Decorator-based

```python
@udf(num_gpus=1, cpu_preprocessing=True)
class PipelinedImageCaption:
    @cpu_stage
    def decode_and_prepare(self, batch: pa.RecordBatch) -> TensorBatch:
        """Runs on CPU actors."""
        images = [decode_image(b) for b in batch["image"]]
        return TensorBatch(torch.stack([to_tensor(img) for img in images]))

    @gpu_stage
    def inference(self, tensors: TensorBatch) -> pa.Array:
        """Runs on GPU actors."""
        with torch.no_grad():
            return self.model(tensors.to("cuda"))
```

## Key Components to Implement

### 1. PreprocessedBatch

Data structure for passing preprocessed data between stages:

```python
@attrs.define
class PreprocessedBatch:
    """Container for preprocessed tensors and metadata."""
    tensors: torch.Tensor  # Preprocessed tensors (on CPU)
    metadata: pa.RecordBatch  # Original batch metadata
    original_task: ReadTask  # For checkpointing
```

### 2. CPUPreprocessorActor

Actor that runs CPU preprocessing:

```python
@ray.remote(num_cpus=2)
class CPUPreprocessorActor:
    """Runs CPU preprocessing stage of pipelined UDFs."""

    def __init__(self, udf: UDF):
        self.udf = udf
        self.udf.setup()  # Initialize on CPU

    def preprocess(
        self, read_task: ReadTask, batch_ref: ObjectRef
    ) -> tuple[ReadTask, ObjectRef]:
        """Run preprocessing and store result in object store."""
        batch = ray.get(batch_ref)
        preprocessed = self.udf.preprocess(batch)
        result_ref = ray.put(preprocessed)
        return read_task, result_ref
```

### 3. PreprocessorPool

Pool manager for CPU preprocessors:

```python
@attrs.define
class PreprocessorPool:
    """Manages CPUPreprocessorActor pool."""

    num_preprocessors: int = 4
    queue_depth: int = 8

    def preprocess_iter(
        self, preloaded_tasks: Iterator[PreloadedReadTask]
    ) -> Iterator[PreprocessedTask]:
        """Iterate over preprocessed tasks."""
        ...
```

### 4. Modified GPUApplierActor

GPU actors that only run inference:

```python
@ray.remote(num_gpus=1)
class GPUInferenceActor:
    """Runs GPU inference on preprocessed batches."""

    def __init__(self, udf: UDF):
        self.udf = udf
        self.udf.setup()  # Load model to GPU

    def inference(self, preprocessed_ref: ObjectRef) -> ObjectRef:
        """Run inference on preprocessed batch."""
        preprocessed = ray.get(preprocessed_ref)
        result = self.udf(preprocessed)  # GPU inference only
        return ray.put(result)
```

## Pipeline Flow

```
                    ┌─────────────────────┐
ReadTask ─→ Queue ─→│ DataLoaderActor     │─→ ObjectRef (raw batch)
                    │ (read from Lance)    │
                    └─────────────────────┘
                              ↓
                    ┌─────────────────────┐
          Queue ───→│ CPUPreprocessorActor│─→ ObjectRef (preprocessed)
                    │ (decode, tokenize)  │
                    └─────────────────────┘
                              ↓
                    ┌─────────────────────┐
          Queue ───→│ GPUInferenceActor   │─→ Result
                    │ (model inference)   │
                    └─────────────────────┘
```

## Configuration

```python
# Proposed config additions
class JobConfig:
    # Phase 1 (existing)
    enable_gpu_pipelining: bool = False
    num_loader_actors: int = 4
    prefetch_depth: int = 8

    # Phase 2 (new)
    enable_cpu_preprocessing: bool = False
    num_preprocessor_actors: int = 4
    preprocessor_queue_depth: int = 8
```

## Detection of Preprocessing UDFs

```python
def _should_use_cpu_preprocessing(udf: UDF) -> bool:
    """Check if UDF has preprocessing stage."""
    return hasattr(udf, 'preprocess') and callable(udf.preprocess)
```

## Backward Compatibility

- UDFs without `preprocess()` method work as before
- Preprocessing pipeline only activated for UDFs that define it
- Phase 1 pipelining still works independently

## Benefits

| Metric | Without Preprocessing Pipeline | With Preprocessing Pipeline |
|--------|-------------------------------|----------------------------|
| GPU Utilization | ~50% (CPU work on GPU actor) | ~95% (GPU only inference) |
| CPU Utilization | Low (only on GPU nodes) | High (parallel preprocessing) |
| Throughput | Limited by preprocessing | Limited by slowest stage |

## Challenges

1. **UDF interface change**: Authors must split preprocessing from inference
2. **Memory pressure**: Preprocessed tensors staged in object store
3. **Ordering**: Must maintain batch ordering through pipeline
4. **Error handling**: Failures at any stage need proper handling

## Implementation Steps

1. Define `PreprocessedBatch` data structure
2. Add `preprocess()` method detection to UDF
3. Implement `CPUPreprocessorActor`
4. Implement `PreprocessorPool`
5. Modify `GPUApplierActor` for inference-only mode
6. Update `ColumnAddPipelineJob` to orchestrate 3-stage pipeline
7. Add configuration options
8. Update documentation and examples

## Open Questions

1. How to handle UDFs that need GPU for preprocessing (e.g., GPU-accelerated image decode)?
2. Should preprocessing be async within the actor (overlapping batches)?
3. How to handle different tensor types (PyTorch, NumPy, Arrow)?
