Metadata-Version: 2.4
Name: gen-worker
Version: 0.7.0
Summary: A library used to build custom functions in Cozy Creator's serverless function platform.
Project-URL: Homepage, https://github.com/cozy-creator/python-gen-worker
Project-URL: Repository, https://github.com/cozy-creator/python-gen-worker
Project-URL: Issues, https://github.com/cozy-creator/python-gen-worker/issues
Author-email: Paul Fidika <paul@fidika.com>
License-Expression: MIT
License-File: LICENSE
Keywords: ai,cozy,inference,ml,serverless
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Requires-Dist: aiohttp>=3.11.14
Requires-Dist: backoff>=2.2.1
Requires-Dist: blake3>=1.0.0
Requires-Dist: grpcio>=1.71.0
Requires-Dist: huggingface-hub>=0.26.0
Requires-Dist: msgspec>=0.18.6
Requires-Dist: numpy>=1.24.0
Requires-Dist: pillow>=9.0.0
Requires-Dist: protobuf>=6.30.0
Requires-Dist: psutil>=7.0.0
Requires-Dist: pyjwt[crypto]>=2.8.0
Requires-Dist: pyyaml>=6.0.0
Requires-Dist: requests>=2.32.0
Requires-Dist: tomli-w>=1.2.0
Requires-Dist: tqdm>=4.66.0
Provides-Extra: audio
Requires-Dist: numpy>=1.24; extra == 'audio'
Requires-Dist: soundfile>=0.12; extra == 'audio'
Provides-Extra: dev
Requires-Dist: grpcio-tools>=1.71.0; extra == 'dev'
Requires-Dist: mypy>=1.10.0; extra == 'dev'
Requires-Dist: pytest>=9.0.0; extra == 'dev'
Requires-Dist: types-pyyaml>=6.0.12.20250915; extra == 'dev'
Requires-Dist: types-requests>=2.32.4.20250913; extra == 'dev'
Provides-Extra: images
Requires-Dist: pillow>=10.0; extra == 'images'
Provides-Extra: torch
Requires-Dist: flashpack>=0.2.1; extra == 'torch'
Requires-Dist: safetensors>=0.7.0; extra == 'torch'
Requires-Dist: torch>=2.11.0; extra == 'torch'
Requires-Dist: torchaudio>=2.11.0; extra == 'torch'
Requires-Dist: torchvision>=0.25.0; extra == 'torch'
Provides-Extra: trainer
Requires-Dist: pyarrow>=17.0.0; extra == 'trainer'
Description-Content-Type: text/markdown

# gen-worker

Python SDK for writing **endpoints** that run on Cozy's worker pool. You write
a decorated Python function; the SDK handles discovery, scheduling, model
loading, cancellation, file I/O, and reporting to the control plane.

Three endpoint kinds:

- **Inference** — request/response (optionally streaming).
- **Training** — long-running, stateful, can publish checkpoints back to a repo.
- **Conversion** — produces weight artifacts on a destination repo.

## Install

```bash
uv add gen-worker          # core
uv add gen-worker[torch]   # with PyTorch
```

## Quick start

```python
import msgspec
from gen_worker import RequestContext, inference_function

class Input(msgspec.Struct):
    prompt: str

class Output(msgspec.Struct):
    text: str

@inference_function
def hello(ctx: RequestContext, payload: Input) -> Output:
    return Output(text=f"Hello, {payload.prompt}!")
```

Pair it with an `endpoint.toml`:

```toml
schema_version = 1
name = "hello"
main = "my_pkg.main"   # import path that contains your @inference_function

[resources]
ram_gb = 2
cpu_cores = 1
```

…and a `Dockerfile`:

```dockerfile
FROM python:3.12-slim
WORKDIR /app
COPY . /app
RUN pip install uv && uv sync --frozen
RUN mkdir -p /app/.tensorhub && \
    uv run python -m gen_worker.discovery > /app/.tensorhub/endpoint.lock
ENTRYPOINT ["uv", "run", "python", "-m", "gen_worker.entrypoint"]
```

Publish with `cozyctl endpoint deploy` (or via the platform UI). The control
plane reads `/app/.tensorhub/endpoint.lock` from the image and routes invocations.

## Reference

See [`docs/endpoint-authoring.md`](docs/endpoint-authoring.md) for the full
authoring guide: model injection, streaming, file uploads, training/conversion
contracts, error types, and local testing.

## Public surface

The top-level `gen_worker` module exports only what endpoint authors need:

- Decorators: `inference_function`, `Resources`
- Bindings: `Repo`, `Dispatch`, `dispatch`
- Context: `RequestContext` (inference; the base), `ConversionContext` (transform / conversion endpoints), `DatasetContext` (dataset-generation), `TrainingContext` (trainer-class)
- Types: `Asset`, `Tensors`, `Compute`, `LoraSpec`
- Errors: `ValidationError`, `RetryableError`, `FatalError`, `ResourceError`,
  `AuthError`, `CanceledError`, `OutputTooLargeError`, `WorkerError`
- Helpers: `Clamp`, `iter_transformers_text_deltas`, `load_loras`,
  `apply_low_vram_config`, `with_oom_retry`

Training and conversion live in their own submodules: `gen_worker.trainer`,
`gen_worker.conversion`, `gen_worker.clone`.

## Migrating 0.6.x → 0.7.0

The 0.7.0 cut replaces the `Annotated[T, ModelRef(...)]` injection pattern and
the `endpoint.toml [models]` table with a single `models={...}` kwarg on
`@inference_function`. `ResourceRequirements` and `ScalingHints` merged into
one `Resources` struct (declared **per function**). The `require_vram` /
`require_compute_capability` / `require_cuda_library` runtime helpers are
gone — the worker now boot-checks each function's `Resources` envelope
against host hardware and self-advertises only runnable functions.

```python
# 0.6.x:
from gen_worker import ModelRef, ResourceRequirements, ScalingHints, inference_function
from gen_worker.capability import require_vram

@inference_function(
    resources=ResourceRequirements(min_vram_gb=4.0),
    scaling_hints=ScalingHints(vram_scales_with=("width", "height")),
)
def generate(
    ctx,
    pipe: Annotated[FluxPipeline, ModelRef(Src.FIXED, ref="acme/flux", flavor="bf16")],
    payload: Input,
) -> Output:
    require_vram(22 * 1024**3)
    ...

# 0.7.0:
from gen_worker import Repo, Resources, inference_function

flux = Repo("acme/flux")

@inference_function(
    resources=Resources(
        requires_gpu=True,
        min_vram_gb=22.0,
        vram_scales_with=("width", "height"),
    ),
    models={"pipe": flux.flavor("bf16")},
)
def generate(ctx, pipe: FluxPipeline, payload: Input) -> Output:
    ...
```

Bare imports of the removed symbols (`ModelRef`, `ModelRefSource`, `Src`,
`ResourceRequirements`, `ScalingHints`, `require_vram`, etc.) raise
`ImportError` with a one-line migration pointer.
