# SageMaker + LanceDB Enterprise + Geneva: Managed GPU Inference for the Multimodal Lakehouse

SageMaker provides managed GPU compute — AWS-integrated, auto-scaling endpoints,
spot instance support. LanceDB Enterprise + Geneva provides data lifecycle
management — checkpointing, incremental processing, versioning, orchestration.
Together they fill each other's gaps: SageMaker handles GPU infrastructure and
model serving within AWS, Geneva handles the data pipeline that feeds it.

| Concern | SageMaker provides | Geneva provides |
|---------|-------------------|-----------------|
| GPU provisioning | Managed (endpoint or processing job) | — |
| Auto-scaling inference | Yes (endpoint auto-scaling policies) | — |
| Spot instance support | Yes (managed spot training/processing) | — |
| Model registry | Yes (SageMaker Model Registry) | — |
| AWS ecosystem integration | Native (S3, IAM, CloudWatch, ECR) | — |
| Checkpointing | — | Automatic (fragment + batch) |
| Incremental processing | — | Automatic (only new rows) |
| Data versioning | — | Automatic (Lance format) |
| Retry / error handling | — | Configurable (tenacity-based) |
| Distributed orchestration | — | Ray Core |
| Vector search / serving | — | LanceDB Enterprise |

---

## Why Combine Them?

Each tool has a gap the other fills.

| Problem | SageMaker alone | Geneva alone | SageMaker + Geneva |
|---------|----------------|--------------|---------------------|
| GPU provisioning | Managed | Manual cluster | Managed |
| AWS-native integration | Yes | Limited | Yes |
| Any model (not just APIs) | Yes | Yes (needs GPU cluster) | Yes (no cluster) |
| Container management | You build + push to ECR | Not needed (pyproject.toml) | SageMaker handles serving |
| Checkpointing | Build yourself (~40 lines) | Automatic | Automatic |
| Incremental processing | Build yourself | Automatic | Automatic |
| Data versioning | Build yourself | Automatic | Automatic |
| Spot instance savings | Yes (up to 90% off) | Manual | Yes |

**The key insight from `cloud_managed_vs_geneva.md`:**

> **SageMaker asks**: "Give me managed infrastructure for my ML code"
> **Geneva asks**: "Give me a managed feature column"

When combined, SageMaker manages the GPU inference service and Geneva manages
the data lifecycle around it. The ~300 lines of glue code from the DIY approach
drops to ~20 lines of pipeline code.

---

## Architecture

```mermaid
flowchart TB
    A["<b>Your Code</b><br/>UDF definitions · table.backfill()"]
    A --> LH

    subgraph LH["Multimodal Lakehouse (CPU-only cluster)"]
        direction TB

        subgraph G["Geneva (compute layer)"]
            direction TB
            G1["Checkpointing"] ~~~ G2["Incremental Processing"] ~~~ G3["Retry / Error Handling"]
            G4["UDF Versions"] ~~~ G5["Adaptive Batching"] ~~~ G6["Orchestration"]
        end

        subgraph R["Ray Workers (CPU)"]
            W1["Worker 1<br/>(UDF)"]
            W2["Worker 2<br/>(UDF)"]
            W3["Worker 3<br/>(UDF)"]
        end

        subgraph L["LanceDB Enterprise (storage + serving)"]
            direction TB
            L1["Lance Tables"] ~~~ L2["Vector Search"] ~~~ L3["Data Versions"]
        end

        G --> R
        R --> L
    end

    subgraph SM["AWS SageMaker (managed GPU platform)"]
        direction TB
        subgraph EP["SageMaker Endpoint"]
            EP1["GPU Instance<br/>(Model)<br/>ml.g5.xlarge"]
            EP2["GPU Instance<br/>(Model)<br/>ml.g5.xlarge"]
            EP3["GPU Instance<br/>(Model)<br/>ml.g5.xlarge"]
        end
        SMC["ECR container images<br/>Auto-scaling policies · CloudWatch"]
    end

    W1 -->|"HTTPS (boto3)"| SM
    W2 -->|"HTTPS (boto3)"| SM
    W3 -->|"HTTPS (boto3)"| SM
```

### What Runs Where

| Concern | Owner | Notes |
|---------|-------|-------|
| Data storage + versioning | LanceDB Enterprise | Lance format on S3, vector indices |
| Checkpointing + resume | Geneva | Fragment + batch level |
| Incremental processing | Geneva | Only new/changed rows |
| Retry / error handling | Geneva | Configurable tenacity-based |
| Distributed orchestration | Geneva + Ray | CPU-only cluster |
| GPU inference | SageMaker | Managed endpoint, auto-scaled |
| Model container images | SageMaker + ECR | Built once, versioned in registry |
| Scaling GPU workers | SageMaker | Auto-scaling policies |
| Model registry | SageMaker | Model versioning, A/B testing |
| Billing | SageMaker (GPU) + cluster (CPU) | GPU is per-instance-hour (spot available) |

**Key point:** The Geneva cluster is **CPU-only** — cheap to run. All GPU cost is
via SageMaker endpoints, which auto-scale and support spot instances. No GPU
management in the Geneva cluster.

---

## Two Integration Modes

### Mode 1: Call Existing SageMaker Endpoint

For teams that already have a SageMaker endpoint deployed (common in enterprise AWS shops):

```python
# Using the SageMakerEmbeddingUDF sketched below in the E2E section
udf = SageMakerEmbeddingUDF(
    endpoint_name="clip-embedding-endpoint",  # Existing endpoint
    region="us-east-1",
)

table.add_columns({"embedding": udf})
table.backfill("embedding")
```

This is the most common pattern for enterprise teams — the ML platform team
deploys and manages SageMaker endpoints, and data teams consume them via Geneva.

### Mode 2: Self-Provisioning

The UDF creates a SageMaker endpoint in `setup()` and tears it down after:

```python
# A provisioning variant would extend SageMakerEndpointUDF to create
# the endpoint in setup() and delete it in teardown()
udf = SageMakerProvisionedEmbeddingUDF(
    model="sentence-transformers/all-MiniLM-L6-v2",
    instance_type="ml.g5.xlarge",
    teardown_on_complete=True,  # Delete endpoint after backfill
)

table.add_columns({"embedding": udf})
table.backfill("embedding")
```

**What happens under the hood:**

1. **First `setup()` call** — UDF creates SageMaker Model + Endpoint Configuration + Endpoint
2. **Endpoint becomes InService** — ~5-10 minutes for cold start (container pull + model load)
3. **Subsequent calls** — Reuses the live endpoint
4. **Teardown** — Optional cleanup after job completes (avoids idle endpoint costs)

### Mode Comparison

| Aspect | Mode 1: Existing Endpoint | Mode 2: Self-Provisioning |
|--------|--------------------------|--------------------------|
| Pre-deployment needed | Yes (ML platform team deploys) | No (Geneva handles it) |
| Who manages endpoint | ML platform / MLOps team | Geneva UDF |
| Container + model config | Team controls via SageMaker console/CDK | Specified in UDF params |
| Best for | Shared inference services, enterprise | Per-pipeline ephemeral GPU |
| Cold start | Warm (always running) | ~5-10 min (endpoint creation) |
| Idle cost | ~$0.70-1.40/hr per instance | None if teardown enabled |
| IAM required | Invoke endpoint permission | Full SageMaker create/delete permissions |

---

## End-to-End Example

Image embedding + captioning via SageMaker, with Geneva handling checkpointing,
incremental processing, and versioning:

### The SageMaker UDF (sketch)

Unlike Modal (where Geneva can call `Function.lookup()` directly), SageMaker
requires a UDF that wraps `boto3` endpoint invocation. This would be built by
the Geneva team or the user — it's a generic wrapper, not per-model code.

```python
import geneva
import json
from typing import Any


@geneva.udf(version="0.1")
class SageMakerEndpointUDF:
    """
    Generic Geneva UDF that calls a SageMaker real-time endpoint.

    No GPU needed on the Geneva side — the GPU runs behind the
    SageMaker endpoint. This UDF runs on CPU Ray workers and makes
    HTTP calls via boto3.
    """

    def __init__(
        self,
        endpoint_name: str,
        region: str = "us-east-1",
        content_type: str = "application/json",
        accept: str = "application/json",
    ):
        self.endpoint_name = endpoint_name
        self.region = region
        self.content_type = content_type
        self.accept = accept

    def setup(self):
        import boto3

        self.client = boto3.client(
            "sagemaker-runtime", region_name=self.region
        )

    def __call__(self, row: dict[str, Any]) -> dict[str, Any]:
        """
        Send a row to the SageMaker endpoint, return the response.

        The payload format depends on the model behind the endpoint.
        Users subclass or configure this for their specific model.
        """
        response = self.client.invoke_endpoint(
            EndpointName=self.endpoint_name,
            ContentType=self.content_type,
            Accept=self.accept,
            Body=json.dumps(row),
        )
        return json.loads(response["Body"].read())


# Concrete UDFs built on top of the generic one:

@geneva.udf(version="0.1")
class SageMakerEmbeddingUDF(SageMakerEndpointUDF):
    """CLIP embedding via a SageMaker endpoint."""

    def __call__(self, image: bytes) -> list[float]:
        import base64

        payload = {"image": base64.b64encode(image).decode()}
        response = self.client.invoke_endpoint(
            EndpointName=self.endpoint_name,
            ContentType=self.content_type,
            Accept=self.accept,
            Body=json.dumps(payload),
        )
        result = json.loads(response["Body"].read())
        return result["embedding"]


@geneva.udf(version="0.1")
class SageMakerCaptionUDF(SageMakerEndpointUDF):
    """BLIP captioning via a SageMaker endpoint."""

    def __call__(self, image: bytes) -> str:
        import base64

        payload = {"image": base64.b64encode(image).decode()}
        response = self.client.invoke_endpoint(
            EndpointName=self.endpoint_name,
            ContentType=self.content_type,
            Accept=self.accept,
            Body=json.dumps(payload),
        )
        result = json.loads(response["Body"].read())
        return result["caption"]
```

**Key points:**
- **No `num_gpus`** — the UDF runs on CPU. The GPU is behind the SageMaker endpoint.
- **`boto3` is the only dependency** — no ML frameworks needed on the Geneva cluster.
- **Generic base class** — `SageMakerEndpointUDF` handles the boto3 plumbing;
  subclasses define the payload format and return type for schema inference.
- **Batch support** — could be extended to use `invoke_endpoint()` with batched
  payloads for throughput, or use SageMaker Batch Transform for very large jobs.

### Pipeline Code

With the UDFs above, the pipeline is straightforward:

```python
import geneva

conn = geneva.connect("s3://my-bucket/workspace")
table = conn.open_table("images")

# Embedding via SageMaker — CLIP on ml.g5.xlarge
table.add_columns({
    "embedding": SageMakerEmbeddingUDF(
        endpoint_name="clip-embedding-prod",  # Pre-deployed by ML platform
        region="us-east-1",
    ),
})

# Captioning via SageMaker — BLIP on ml.g5.xlarge
table.add_columns({
    "caption": SageMakerCaptionUDF(
        endpoint_name="blip-captioning-prod",  # Pre-deployed by ML platform
        region="us-east-1",
    ),
})

# Run — Geneva handles everything else
table.backfill("embedding", batch_size=100)
table.backfill("caption", batch_size=50)

# Later: add new images → only new rows processed
table.add(new_images)
table.backfill("embedding")
table.backfill("caption")
```

**What Geneva handles:**
- Checkpointing (fragment + batch level, resume from exact failure point)
- Incremental processing (only rows where column IS NULL)
- Data versioning (every backfill creates a new version)
- Retry with exponential backoff on transient failures
- Adaptive batch sizing based on execution time
- Distributed orchestration across Ray workers

**What SageMaker handles:**
- GPU provisioning (managed instances behind endpoint)
- Model serving (container with model weights, health checks)
- Auto-scaling (CloudWatch-based scaling policies)
- Model registry (versioned models, approval workflows)
- CloudWatch logging and monitoring
- Spot instance support (for processing jobs)

---

## Comparison: SageMaker + Geneva vs DIY SageMaker

The `cloud_managed_vs_geneva.md` doc shows ~280 lines + Terraform config for a
DIY SageMaker pipeline. With the SageMaker + Geneva pattern:

| Component | DIY SageMaker | SageMaker + Geneva |
|-----------|--------------|---------------------|
| Inference code | ~80 lines (inference.py in container) | Endpoint already deployed |
| Pipeline/orchestration | ~60 lines (run_job.py) | ~15 lines (table.backfill) |
| Dockerfile | ~20 lines | 0 (endpoint handles this) |
| Checkpoint logic | ~40 lines | 0 (Geneva automatic) |
| IAM/permissions | ~50 lines (Terraform) | Invoke-only IAM |
| Compute setup | ~30 lines (Terraform) | 0 (endpoint pre-exists) |
| **TOTAL** | **~280 lines + config** | **~15 lines** |

The container, model serving, and GPU management still exist — they're just
owned by the ML platform team via SageMaker, not reimplemented in every pipeline.

---

## Comparison with Other Geneva Inference Approaches

| Factor | Third-Party API (Gemini/OpenAI UDFs) | Local GPU (`num_gpus=1`) | Modal + Geneva | SageMaker + Geneva |
|--------|--------------------------------------|--------------------------|----------------|---------------------|
| **Setup complexity** | Low | High | Medium | Medium-High |
| **GPU management** | None | Full responsibility | None | SageMaker manages |
| **Model flexibility** | Provider models only | Any model | Any model | Any model |
| **AWS integration** | None | Manual | None | Native |
| **Data privacy** | Low (data leaves infra) | High (stays local) | Medium (crosses network) | High (stays in AWS VPC) |
| **Latency** | High (50-500ms) | Low (<10ms) | Medium (10-100ms) | Medium (10-50ms) |
| **Scaling** | Automatic | Manual / complex | Automatic | Automatic (policy-based) |
| **Cold start** | None | 10-60s | 3-30s | 5-10 min (endpoint) |
| **Idle cost** | None | ~$1-3/hr per GPU | None | ~$0.70-1.40/hr per instance |
| **Spot support** | N/A | Manual | N/A | Built-in |
| **Governance** | Limited | Manual | Limited | Full (IAM, VPC, CloudTrail) |

### Cost Deep Dive (1M embeddings, 384-dim)

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

| Approach | Compute Cost | Notes |
|----------|-------------|-------|
| OpenAI API (text-embedding-3-small) | $10 | Pay-per-token |
| Local A100 (1hr job) | $4 | Assumes full utilization |
| Modal A10G | $2.50 | ~15,000 sec total inference, pay-per-second |
| SageMaker ml.g5.xlarge (on-demand) | $5.50 | ~1hr runtime + endpoint overhead |
| SageMaker ml.g5.xlarge (spot) | $1.65 | 70% discount, but can be interrupted |

**SageMaker + Geneva is the enterprise choice** for teams that need AWS-native
governance, VPC isolation, and compliance — with Geneva eliminating the data
pipeline glue code.

---

## When to Use This Pattern

### Decision Criteria

**Use SageMaker + Geneva when:**
- Your organization is standardized on AWS
- You need VPC isolation / data residency within AWS
- ML platform team already manages SageMaker endpoints
- Compliance requires IAM, CloudTrail, and governance controls
- You want managed GPU infrastructure without building pipelines
- Workloads are steady enough to justify always-on endpoints (or use spot for batch)

**Use Modal + Geneva when:**
- You want simplest possible GPU setup (no containers, no IAM)
- Workloads are bursty (pay-per-second matters)
- No cloud provider lock-in requirement
- Startup / small team without ML platform engineers

**Use Third-Party APIs (Gemini/OpenAI UDFs) when:**
- Provider models are sufficient
- Simplest possible setup
- Dataset is < 10M rows

**Use Local GPU when:**
- Sustained 24/7 workloads
- Latency-critical (< 10ms)
- Strict data isolation (no network boundary crossing)
- Already have GPU infrastructure

### Decision Flowchart

```mermaid
flowchart TB
    Start["Need GPU inference<br/>in Geneva pipeline?"]
    Start --> Q1{"Provider model<br/>sufficient?<br/>(GPT, Gemini, Claude)"}

    Q1 -->|Yes| Q2{"Dataset < 10M rows?"}
    Q1 -->|"No — need custom<br/>or open-source model"| Q3{"Standardized<br/>on AWS?"}

    Q2 -->|Yes| API["<b>Third-Party API</b><br/>(Gemini/OpenAI UDFs)<br/>Simplest setup"]
    Q2 -->|"No — cost matters<br/>at scale"| Q3

    Q3 -->|"Yes — enterprise<br/>AWS shop"| Q4{"Have GPU<br/>infrastructure<br/>already?"}
    Q3 -->|"No — or<br/>multi-cloud"| MODAL["<b>Modal + Geneva</b><br/>Serverless GPU<br/>Pay-per-second"]

    Q4 -->|"Yes — sustained<br/>24/7 workload"| LOCAL["<b>Local GPU</b><br/>(num_gpus=1)<br/>Lowest cost at scale"]
    Q4 -->|"No — let SageMaker<br/>manage it"| SM["<b>SageMaker + Geneva</b><br/>Managed GPU<br/>AWS-native governance"]
```

---

## SageMaker Capabilities Reference

How SageMaker features map to Geneva integration points:

| SageMaker Feature | What It Does | Geneva Integration |
|-------------------|-------------|-------------------|
| Real-time Endpoints | Low-latency model serving | Mode 1: call via `invoke_endpoint()` |
| Batch Transform | Batch inference on S3 data | Alternative to endpoint for large jobs |
| Processing Jobs | Run arbitrary containers on managed GPU | Could wrap Geneva UDF as processing job |
| Model Registry | Version and approve models | Track model version in UDF `version` param |
| Auto-scaling | Scale endpoint instances on metrics | Complements Geneva's distributed workers |
| Spot instances | Up to 90% cost savings | Use for processing jobs, not endpoints |
| Multi-model endpoints | Serve multiple models on one endpoint | One endpoint per UDF, or shared |
| VPC configuration | Network isolation | Data stays within AWS VPC |
| CloudWatch | Monitoring and alerting | Monitor inference latency/errors |
| IAM | Fine-grained access control | Invoke-only permissions for Geneva |
| SageMaker Pipelines | ML workflow orchestration | Geneva replaces this for data pipelines |

---

## Implementation Status

| Phase | Description | Status |
|-------|-------------|--------|
| Phase 1 | API provider UDFs (Gemini/OpenAI embedding + generative) | **Implemented** |
| Phase 2 | Local GPU support (`num_gpus=1` on Ray workers) | **Implemented** |
| Phase 3 | Serverless GPU integration (Modal) | **Planned** |
| Phase 3b | Managed GPU integration (SageMaker) | **Planned** |
| Phase 4 | Unified `InferenceBackend` abstraction | **Planned** |

### Phase 3b Scope (SageMaker Integration)

- `SageMakerInferenceUDF` base class
- `SageMakerEndpointUDF` for calling existing endpoints (Mode 1)
- `SageMakerProvisionedUDF` for auto-provisioning (Mode 2)
- Batch Transform mode for large offline jobs
- Request batching to maximize throughput per endpoint
- IAM credential handling (instance role or explicit)

### Open Questions

1. **Endpoint lifecycle management?** Should Geneva create/delete endpoints, or always expect pre-existing ones?
2. **Batch Transform vs Endpoint?** When should Geneva use Batch Transform (S3-to-S3) vs real-time endpoint calls?
3. **Multi-model endpoints?** Can multiple UDFs share a single SageMaker endpoint?
4. **Cost attribution?** SageMaker endpoint costs are per-hour — how to attribute across Geneva jobs?
5. **Spot interruption handling?** If a spot processing job is interrupted, Geneva's checkpointing resumes — but does the endpoint need re-provisioning?

---

## Related Documents

- [Cloud Managed vs Geneva](cloud_managed_vs_geneva.md) — Competitive comparison (SageMaker/Vertex/Azure as standalone alternatives)
- [Modal + Geneva](modal_plus_geneva.md) — Serverless GPU partner integration (similar pattern, different trade-offs)
- [Comparison Overview](comparison_overview.md) — Category map of all competitors
