Metadata-Version: 2.3
Name: atdata-lambda
Version: 0.1.1a2
Summary: Modal serverless processing backend for atdata datasets
Author: Maxine Levesque
Author-email: Maxine Levesque <170461181+maxinelevesque@users.noreply.github.com>
Requires-Dist: atdata>=0.3.0b1
Requires-Dist: modal
Requires-Dist: boto3
Requires-Dist: webdataset
Requires-Python: >=3.12
Description-Content-Type: text/markdown

# atdata-lambda

Modal serverless processing backend for [atdata](https://github.com/foundation-ac/atdata) datasets.

Fan out sample-level transformations across sharded WebDataset tar files using
[Modal](https://modal.com) for compute and [Cloudflare R2](https://developers.cloudflare.com/r2/)
for storage. One shard = one worker — no job queue required.

## Installation

```bash
pip install atdata-lambda
```

Requires Python 3.12+.

## Quickstart

```python
import modal
from atdata_lambda import processor, Dispatcher

app = modal.App("my-pipeline")

@processor(app, gpu="T4", memory=8192)
def embed_images(sample: ImageSample) -> EmbeddingSample:
    embedding = model.encode(sample.image)
    return EmbeddingSample(embedding=embedding, label=sample.label)

dispatcher = Dispatcher(embed_images)
result = dispatcher.run(
    input_prefix="data/raw/train/",
    output_prefix="data/embedded/train/",
)
print(result.summary())
# Job a1b2c3d4e5f6: SUCCESS — 128/128 shards completed, 0 failed
```

### R2 Configuration

Credentials are read from environment variables (typically injected via
[Modal Secrets](https://modal.com/docs/guide/secrets)):

| Variable | Description |
|---|---|
| `R2_ENDPOINT_URL` | Cloudflare R2 S3-compatible endpoint |
| `R2_ACCESS_KEY_ID` | R2 access key |
| `R2_SECRET_ACCESS_KEY` | R2 secret key |
| `R2_BUCKET` | Target bucket name |

Or pass an `R2Config` explicitly:

```python
from atdata_lambda import R2Config

config = R2Config(
    endpoint_url="https://acct.r2.cloudflarestorage.com",
    access_key_id="...",
    secret_access_key="...",
    bucket="my-bucket",
)
dispatcher = Dispatcher(embed_images, r2_config=config)
```

### HTTP Endpoint

Expose a Modal web endpoint for event-driven dispatch (e.g. R2 bucket notifications):

```python
from atdata_lambda import create_dispatch_endpoint

create_dispatch_endpoint(app, dispatcher)
```

This registers a POST endpoint accepting `{"input_prefix": "...", "output_prefix": "..."}`.

## Architecture

```
Client (notebook / CLI / R2 event)
     │
     ▼
  Dispatcher ── list_shards(prefix) ──→ R2
     │
     │  .map()
     ▼
  Modal workers (1 per shard)
     │  read shard → apply @processor fn → write output shard
     │  write .manifest.json
     ▼
  R2 (output shards + manifests)
```

### Modules

| Module | Responsibility |
|---|---|
| `_processor.py` | `@processor` decorator — wraps user functions, extracts Packable type info, registers with Modal |
| `_dispatcher.py` | `Dispatcher` — enumerates shards, fans out via `.map()`, collects `JobResult` |
| `_shard_worker.py` | `process_shard` — runs inside each Modal worker, reads/writes tar shards |
| `_manifest.py` | `ManifestBuilder` — per-shard `.manifest.json` generation alongside output |
| `_config.py` | `R2Config` — R2 credentials, `list_shards` for shard enumeration via boto3 |

### Public API

```python
from atdata_lambda import (
    processor,           # decorator
    Dispatcher,          # fan-out orchestrator
    JobResult,           # aggregate result
    R2Config,            # R2 connection config
    ShardIndex,          # shard listing result
    list_shards,         # enumerate .tar files at prefix
    ManifestBuilder,     # incremental manifest construction
    ShardManifest,       # manifest dataclass
    write_manifest,      # upload manifest to R2
    ShardResult,         # per-shard processing result
    ProcessorError,      # invalid processor configuration
    create_dispatch_endpoint,  # HTTP trigger registration
)
```

## Development

```bash
uv sync                          # install dependencies
uv run pytest                    # run tests
uv run ruff check .              # lint
uv run ruff format .             # format
```

## License

See [LICENSE](LICENSE).
