# Inference UDF Internal Design Doc

**Status:** Draft
**Author:** [Author]
**Date:** 2026-02-05

## Overview

This document compares three architectural approaches for running ML inference within Geneva UDFs:

1. **Third-Party Inference as a Service** - API calls to hosted providers (OpenAI, Gemini, TwelveLabs, etc.)
2. **Local GPU in UDF Job** - GPU resources allocated directly to Ray workers (`num_gpus=1.0`)
3. **Hybrid GPU-as-a-Service** - UDFs call out to GPU-rich platforms (Modal, Nebius, etc.)

## Current State

Geneva already provides key building blocks for inference UDFs:

### Stateful UDFs

Class-based UDFs support one-time initialization of clients, models, or connections per Ray worker. By convention, a `setup()` method is used for lazy initialization (not a formal API):

```python
@udf(data_type=pa.list_(pa.float32(), 384), num_gpus=1.0)
class EmbeddingUDF:
    def setup(self):
        # Convention: called lazily on first __call__ - initialize clients, load models
        self.model = load_model().to("cuda")
        # Or for API-based inference:
        # self.client = openai.OpenAI(api_key=os.environ["OPENAI_API_KEY"])

    def __call__(self, batch: pa.RecordBatch) -> pa.Array:
        return self.model.encode(batch)
```

### GPU Allocation

GPU resources flow through: `UDF.num_gpus` → `MapTask` → `Ray actor.options(num_gpus=X)`.

### Built-in Retry Policies

Geneva provides out-of-the-box retry handling via `ErrorHandlingConfig` (tenacity-based):

```python
from geneva.debug.error_store import ErrorHandlingConfig, FaultIsolation, RetryConfig
from tenacity import stop_after_attempt, wait_exponential, retry_if_exception_type

@udf(
    data_type=pa.list_(pa.float32(), 384),
    error_handling=ErrorHandlingConfig(
        fault_isolation=FaultIsolation.SKIP_ROWS,  # or FAIL_BATCH
        retry_config=RetryConfig(
            stop=stop_after_attempt(3),
            wait=wait_exponential(multiplier=1, max=60),
            retry=retry_if_exception_type((RateLimitError, TimeoutError)),
        ),
        log_errors=True,
    ),
)
class APIInferenceUDF:
    def setup(self):
        self.client = openai.OpenAI()

    def __call__(self, texts: pa.StringArray) -> pa.Array:
        # Retries handled automatically by framework
        return self.client.embeddings.create(...)
```

### Adaptive Batch Sizing

`AdaptiveCheckpointSizer` dynamically adjusts batch sizes based on execution time, optimizing for both throughput and checkpoint frequency.

### Gap: Users Must Write Their Own UDFs

While Geneva provides the building blocks (stateful UDFs, retry policies, batch sizing), users currently must write their own inference UDFs from examples. There are no out-of-the-box UDFs for common third-party inference services.

---

## Approach 1: Third-Party Inference as a Service

### Description

UDFs make HTTP/gRPC calls to hosted inference APIs (OpenAI, Gemini, Anthropic, TwelveLabs, Replicate, etc.). No local GPU required.

### Proposal: Out-of-the-Box UDFs

The key value-add for this approach is providing **pre-built UDFs** with simple, consistent configuration:

- **Provider selection** - Choose from OpenAI, Gemini, Anthropic, Cohere, TwelveLabs, etc.
- **Model selection** - Easy parameter to specify model (e.g., `model="text-embedding-3-small"`)
- **API key conventions** - Standard environment variable names or secrets integration
- **Prompt/input configuration** - Consistent patterns for passing prompts, system messages, input columns

#### Embeddings

```python
from geneva.udfs.inference import openai_embedding_udf

# Simple one-liner with sensible defaults
udf = openai_embedding_udf(model="text-embedding-3-small", column="text")

# Or with more options
udf = openai_embedding_udf(
    model="text-embedding-3-large",
    column="text",
    api_key_env="OPENAI_API_KEY",  # Convention: defaults to this
    dimensions=1536,
)

table.backfill("embeddings", udf=udf)
```

#### LLM Completions

```python
from geneva.udfs.inference import openai_llm_udf, gemini_llm_udf

# Apply a prompt to each row
udf = openai_llm_udf(
    model="gpt-5.2-instant",
    prompt="Summarize the following article in 2 sentences: {article}",
    column="article",  # Column to substitute into {article}
)

table.backfill("summary", udf=udf)

# Multi-column input with system message
udf = openai_llm_udf(
    model="gpt-5.2",
    system_prompt="You are a sentiment analysis expert.",
    prompt="Analyze the sentiment of this review for {product}: {review_text}",
    columns=["product", "review_text"],  # Multiple columns
    temperature=0.0,  # Deterministic output
)

table.backfill("sentiment", udf=udf)
```

### Architecture

```mermaid
flowchart TB
    subgraph Geneva["Geneva Cluster (CPU-only)"]
        W1["Ray Worker<br/>(UDF)"]
        W2["Ray Worker<br/>(UDF)"]
        W3["Ray Worker<br/>(UDF)"]
    end

    subgraph APIs["Third-Party APIs"]
        OpenAI["OpenAI"]
        Gemini["Gemini"]
        TwelveLabs["TwelveLabs"]
    end

    W1 -->|HTTP/gRPC| APIs
    W2 -->|HTTP/gRPC| APIs
    W3 -->|HTTP/gRPC| APIs
```

### Pros

| Aspect | Benefit |
|--------|---------|
| **Zero GPU infrastructure** | No need to provision/manage GPU nodes |
| **Model variety** | Access to latest models (GPT-4V, Gemini 1.5, etc.) without deployment |
| **Automatic scaling** | Provider handles scaling to demand |
| **Simplicity** | Minimal code, no model loading/versioning |
| **Cost predictability** | Pay-per-token pricing, no idle GPU costs |
| **Maintenance-free** | No model updates, driver compatibility, or CUDA issues |

### Cons

| Aspect | Concern |
|--------|---------|
| **Latency** | Network round-trip adds 50-500ms per request |
| **Cost at scale** | $0.0001/1K tokens adds up: 1B tokens = $100K+ |
| **Rate limits** | Provider throttling (e.g., OpenAI 10K RPM) |
| **Data egress** | Sensitive data leaves your infrastructure |
| **Vendor lock-in** | API changes, pricing changes, deprecations |
| **Availability** | Dependent on provider uptime |
| **Limited models** | Restricted to what providers offer |

### Cost Estimation

| Provider | Model | Price | 1M rows (avg 500 tokens) | 1B rows |
|----------|-------|-------|--------------------------|---------|
| OpenAI | text-embedding-3-small | $0.02/1M tokens | $10 | $10,000 |
| OpenAI | text-embedding-3-large | $0.13/1M tokens | $65 | $65,000 |
| Gemini | text-embedding-004 | $0.00025/1K chars | ~$125 | ~$125,000 |
| Cohere | embed-v3 | $0.10/1M tokens | $50 | $50,000 |

### Best For

- Quick prototyping and experimentation
- Access to proprietary models (GPT-4, Claude, Gemini)
- Small-to-medium datasets (< 10M rows)
- Teams without GPU infrastructure expertise
- Variable/unpredictable workloads

---

## Approach 2: Local GPU in UDF Job

### Description

GPU resources allocated directly to Ray workers. Models loaded locally, inference runs on cluster GPUs.

### Architecture

```mermaid
flowchart TB
    subgraph Geneva["Geneva Cluster (GPU-enabled)"]
        subgraph W1["Ray Worker"]
            subgraph UDF1["UDF"]
                M1["Model"]
            end
            G1[("GPU 0")]
        end

        subgraph W2["Ray Worker"]
            subgraph UDF2["UDF"]
                M2["Model"]
            end
            G2[("GPU 1")]
        end

        Storage[("Lance Storage<br/>(local or object store)")]
    end

    UDF1 --> G1
    UDF2 --> G2
    W1 <--> Storage
    W2 <--> Storage
```

### Example Implementation

```python
@udf(
    data_type=pa.list_(pa.float32(), 384),
    num_gpus=1.0,
    batch_size=256,
    memory="16GB"
)
class LocalEmbeddingUDF:
    def setup(self):
        from sentence_transformers import SentenceTransformer
        self.model = SentenceTransformer("all-MiniLM-L6-v2", device="cuda")

    def __call__(self, texts: pa.StringArray) -> pa.Array:
        with torch.inference_mode():
            embeddings = self.model.encode(
                texts.to_pylist(),
                convert_to_numpy=True,
                batch_size=64
            )
        return pa.array(embeddings.tolist(), type=pa.list_(pa.float32(), 384))
```

### Fractional GPU Sharing

```python
# Run 4 UDF instances sharing one GPU
@udf(data_type=pa.float32(), num_gpus=0.25)
class SharedGPUUDF:
    def setup(self):
        # Limit GPU memory per instance
        torch.cuda.set_per_process_memory_fraction(0.25)
        self.model = load_small_model().to("cuda")
```

### Pros

| Aspect | Benefit |
|--------|---------|
| **Lowest latency** | No network hop, data stays local |
| **Cost efficient at scale** | Amortized GPU cost over large datasets |
| **Data privacy** | Data never leaves your infrastructure |
| **Model flexibility** | Run any model (custom, fine-tuned, open-source) |
| **Deterministic** | Same model version, reproducible results |
| **Batching control** | Optimize batch size for your workload |

### Cons

| Aspect | Concern |
|--------|---------|
| **GPU provisioning** | Must provision and manage GPU nodes |
| **Idle cost** | GPUs cost ~$1-3/hr even when idle |
| **Scaling complexity** | Auto-scaling GPU clusters is hard |
| **Model management** | Must version, deploy, and update models |
| **CUDA/driver issues** | Compatibility matrix hell |
| **Cold start** | Model loading takes 10-60s per worker |
| **Memory limits** | Large models may not fit on single GPU |

### GPU Cluster Options

| Platform | GPU Type | Price/hr | Notes |
|----------|----------|----------|-------|
| AWS (p4d.24xlarge) | 8x A100 | ~$32/hr | Best for large jobs |
| GCP (a2-highgpu-1g) | 1x A100 | ~$4/hr | Good balance |
| Azure (NC6s_v3) | 1x V100 | ~$3/hr | Budget option |
| Lambda Labs | 1x A100 | ~$1.10/hr | Cheapest A100 |

### Cold Start Mitigation

```python
# Pre-warm actors before job starts
pipeline = RayPipeline(...)
pipeline.prewarm_actors(num_actors=8, timeout=120)

# Or use Ray's actor pool with min replicas
@ray.remote(num_gpus=1, max_restarts=3)
class PersistentModelActor:
    def __init__(self):
        self.model = load_model()  # Loaded once
```

### Best For

- Large-scale production workloads (100M+ rows)
- Latency-sensitive applications
- Custom or fine-tuned models
- Data that cannot leave your infrastructure
- Predictable, sustained workloads

---

## Approach 3: Hybrid GPU-as-a-Service

### Description

UDFs run on CPU-only nodes but call out to serverless GPU platforms (Modal, Nebius, RunPod, Banana, Replicate) for inference.

### Proposal: OOTB UDFs with Two Modes

Similar to Approach 1, we provide **pre-built UDFs**. These support two modes:

#### Mode 1: Connect to Existing Endpoint (like Approach 1)

For persistent microservices that are always running:

```python
from geneva.udfs.inference import modal_embedding_udf

# Connect to an already-deployed Modal endpoint
udf = modal_embedding_udf(
    endpoint="geneva-inference/Embedder",  # Existing deployment
    column="text",
)

table.backfill("embeddings", udf=udf)
```

#### Mode 2: Self-Provisioning

The `setup()` method provisions the GPU microservice automatically:

```python
# Auto-provision Modal endpoint
udf = modal_embedding_udf(
    model="sentence-transformers/all-MiniLM-L6-v2",
    column="text",
    gpu="A10G",  # or "A100", "H100"
    provision=True,  # Explicitly opt into provisioning
)

table.backfill("embeddings", udf=udf)
```

**What happens under the hood:**

1. **First `setup()` call** - UDF provisions/deploys the inference microservice on Modal
2. **Subsequent calls** - Reuses the deployed endpoint
3. **Warm pool** - Platform maintains warm containers for low latency
4. **Teardown** - Optional cleanup after job completes (or keep warm for future jobs)

```python
# Full options
udf = modal_embedding_udf(
    model="BAAI/bge-large-en-v1.5",
    column="text",
    gpu="A100",
    provision=True,
    container_idle_timeout=300,  # Keep warm for 5 min
    allow_concurrent_inputs=100,  # Batching on Modal side
    api_key_env="MODAL_TOKEN_ID",  # Convention for credentials
)
```

### Architecture

```mermaid
flowchart TB
    subgraph Geneva["Geneva Cluster (CPU-only)"]
        W1["Ray Worker<br/>(UDF)"]
        W2["Ray Worker<br/>(UDF)"]
        W3["Ray Worker<br/>(UDF)"]
    end

    subgraph Platform["Serverless GPU Platform (Modal / Nebius / RunPod)"]
        direction TB
        subgraph Workers["Auto-scaled GPU Workers"]
            GW1["GPU Worker<br/>(Model)<br/>A100"]
            GW2["GPU Worker<br/>(Model)<br/>A100"]
            GW3["GPU Worker<br/>(Model)<br/>A100"]
        end
        Cache["Model cached in container image<br/>Warm pool maintained"]
    end

    W1 -->|gRPC/HTTP| Platform
    W2 -->|gRPC/HTTP| Platform
    W3 -->|gRPC/HTTP| Platform
```

### Pros

| Aspect | Benefit |
|--------|---------|
| **No GPU management** | Platform handles provisioning, scaling, drivers |
| **Pay-per-use** | Only pay for actual inference time |
| **Auto-scaling** | Automatic scale-up/down based on demand |
| **Model flexibility** | Deploy any model you want (unlike API providers) |
| **Warm pools** | Platforms maintain warm containers, reducing cold starts |
| **Separation of concerns** | Geneva handles orchestration, platform handles inference |
| **Cost optimization** | Use spot/preemptible GPUs for 60-80% savings |

### Cons

| Aspect | Concern |
|--------|---------|
| **Network latency** | 10-100ms per request (better than public APIs) |
| **Platform dependency** | Tied to specific platform's API/SDK |
| **Data transfer** | Data crosses network boundary |
| **Debugging complexity** | Failures span two systems |
| **Cost at very high scale** | May exceed local GPU cost for sustained loads |
| **Cold starts** | First request may take 10-30s (mitigated by warm pools) |

### Platform Comparison

| Platform | GPU Options | Pricing | Cold Start | Concurrency | Notes |
|----------|-------------|---------|------------|-------------|-------|
| **Modal** | T4, A10G, A100, H100 | $0.000164/sec (A10G) | ~3s (warm) | 100+ per container | Best DevX |
| **Nebius** | H100, H200, L40S | €2.60/hr (H100) | ~10s | High | Europe-focused |
| **RunPod** | 3090, A100, H100 | $0.74/hr (A100) | ~15s | Variable | Cheapest H100 |
| **Banana** | A100 | ~$0.0015/sec | ~5s | Limited | Simple API |
| **Replicate** | T4, A40, A100 | $0.00055/sec (A40) | ~10s | Auto | Model marketplace |

### Cost Comparison (1M embeddings, 384-dim)

Assumptions: 500 tokens/row avg, 50ms inference time per batch of 32

| Approach | Compute Cost | Notes |
|----------|--------------|-------|
| OpenAI API | $10 | text-embedding-3-small |
| Local A100 (1hr job) | $4 | Assumes full utilization |
| Modal A10G | $2.50 | ~15,000 sec total inference |
| Modal A100 | $4.20 | Faster but more expensive |

### Best For

- Variable/bursty workloads
- Teams wanting model flexibility without GPU ops
- Multi-cloud or hybrid deployments
- Cost optimization via spot instances
- Rapid experimentation with different models

---

## Comparison Matrix

| Factor | Third-Party API | Local GPU | Hybrid (Modal/Nebius) |
|--------|-----------------|-----------|----------------------|
| **Setup complexity** | Low | High | Medium |
| **Latency** | High (50-500ms) | Low (<10ms) | Medium (10-100ms) |
| **Cost at 1M rows** | $10-65 | $4-10 | $2.50-5 |
| **Cost at 1B rows** | $10K-65K | $1-2K | $2-5K |
| **Scaling** | Automatic | Manual/Complex | Automatic |
| **Model flexibility** | Low (provider models only) | High (any model) | High (any model) |
| **Data privacy** | Low (data leaves infra) | High (data stays local) | Medium (crosses network) |
| **GPU management** | None | Full responsibility | None |
| **Cold start** | None | 10-60s | 3-30s (warm pools) |
| **Availability** | Provider-dependent | Self-managed | Platform-dependent |

---

## Recommendation Framework

### Use Third-Party APIs when:
- [x] Prototyping or experimenting quickly
- [x] Need access to proprietary models (GPT-4, Claude)
- [x] Dataset < 10M rows
- [x] No GPU infrastructure or expertise
- [x] Unpredictable, sporadic workloads

### Use Local GPU when:
- [x] Large-scale production (100M+ rows)
- [x] Latency is critical (<10ms required)
- [x] Running custom/fine-tuned models
- [x] Strict data residency requirements
- [x] Sustained, predictable workloads
- [x] Already have GPU infrastructure

### Use Hybrid GPU-as-a-Service when:
- [x] Want model flexibility without GPU ops
- [x] Variable/bursty workloads
- [x] Cost optimization is priority
- [x] Multi-model experimentation
- [x] Bridge between prototyping and production

---

## Implementation Roadmap

### Phase 1: Enhanced API Provider Support (Approach 1)
- [ ] Generic `APIInferenceUDF` base class with retry/rate-limiting
- [ ] Pre-built UDFs for OpenAI, Gemini, Anthropic, Cohere
- [ ] Credential management via environment/secrets
- [ ] Cost tracking and budget alerts

### Phase 2: Improved Local GPU Support (Approach 2)
- [ ] GPU memory monitoring and OOM handling
- [ ] Automatic batch size tuning based on GPU memory
- [ ] Multi-GPU model parallelism support
- [ ] Pre-warming actor pools before job start
- [ ] Model caching across jobs (Ray object store)

### Phase 3: Serverless GPU Integration (Approach 3)
- [ ] `ModalInferenceUDF` base class
- [ ] `NebiusInferenceUDF` base class
- [ ] Connection pooling and request batching
- [ ] Automatic failover between providers
- [ ] Cost optimization (spot instances, region selection)

### Phase 4: Unified Inference Abstraction
- [ ] Common `InferenceBackend` interface
- [ ] Automatic backend selection based on workload
- [ ] Seamless fallback between approaches
- [ ] Unified monitoring and cost tracking

---

## Open Questions

1. **Should Geneva manage inference endpoints directly?** Or leave deployment to users?

2. **How to handle model versioning?** Same UDF code should produce consistent results.

3. **Connection pooling strategy?** Per-actor vs shared across actors?

4. **Cost attribution?** How to track inference costs per job/table?

---

## Appendix: Code Examples

> **Note:** Retry logic and batch sizing are already handled by Geneva's framework:
> - **Retry:** Use `ErrorHandlingConfig` with `retry_config` (tenacity-based) - see `geneva/apply/error_handling.py`
> - **Batching:** Use `checkpoint_size` parameter; `AdaptiveCheckpointSizer` handles dynamic sizing - see `geneva/apply/adaptive.py`

### Model Registry Pattern

```python
class ModelRegistry:
    """Central registry for inference models."""

    _models: dict[str, Any] = {}

    @classmethod
    def get(cls, model_id: str, loader: Callable) -> Any:
        if model_id not in cls._models:
            cls._models[model_id] = loader()
        return cls._models[model_id]

    @classmethod
    def clear(cls):
        cls._models.clear()

# Usage in UDF
class EmbeddingUDF:
    def setup(self):
        self.model = ModelRegistry.get(
            "sentence-transformers/all-MiniLM-L6-v2",
            lambda: SentenceTransformer("all-MiniLM-L6-v2", device="cuda")
        )
```
