Metadata-Version: 2.5
Name: gen-worker
Version: 0.114.2
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.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.12
Requires-Dist: blake3>=1.0.0
Requires-Dist: c2pa-python>=0.36
Requires-Dist: gguf>=0.10.0
Requires-Dist: grpcio>=1.82.1
Requires-Dist: huggingface-hub>=0.26.0
Requires-Dist: msgspec>=0.18.6
Requires-Dist: protobuf>=7.35.0
Requires-Dist: psutil>=7.0.0
Requires-Dist: pyyaml>=6.0.0
Requires-Dist: requests>=2.32.0
Requires-Dist: tomli-w>=1.0.0
Provides-Extra: audio
Requires-Dist: numpy>=1.24; extra == 'audio'
Requires-Dist: soundfile>=0.12; extra == 'audio'
Provides-Extra: dev
Requires-Dist: accelerate>=1.9; extra == 'dev'
Requires-Dist: av>=12; extra == 'dev'
Requires-Dist: c2pa-python>=0.36; extra == 'dev'
Requires-Dist: diffusers>=0.39.0; extra == 'dev'
Requires-Dist: grpcio-tools>=1.82.1; extra == 'dev'
Requires-Dist: hypothesis>=6.100; extra == 'dev'
Requires-Dist: mypy>=1.10.0; extra == 'dev'
Requires-Dist: pillow>=10.0; extra == 'dev'
Requires-Dist: pyarrow>=17.0.0; extra == 'dev'
Requires-Dist: pytest-xdist>=3.8.0; extra == 'dev'
Requires-Dist: pytest>=9.0.0; extra == 'dev'
Requires-Dist: ruff>=0.6.0; extra == 'dev'
Requires-Dist: transformers>=5.13; 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: signing
Requires-Dist: c2pa-python>=0.36; extra == 'signing'
Provides-Extra: torch
Requires-Dist: accelerate>=1.9; extra == 'torch'
Requires-Dist: bitsandbytes>=0.45.0; (sys_platform == 'linux' or sys_platform == 'win32') and extra == 'torch'
Requires-Dist: safetensors>=0.8.0; extra == 'torch'
Requires-Dist: torch>=2.13.0; extra == 'torch'
Provides-Extra: video
Requires-Dist: av>=12; extra == 'video'
Requires-Dist: numpy>=1.24; extra == 'video'
Provides-Extra: vision
Requires-Dist: torchvision>=0.28.0; extra == 'vision'
Description-Content-Type: text/markdown

# gen-worker

Python SDK for writing **endpoints** that run on Cozy's worker pool. You write
one decorated function or class; the SDK handles discovery, scheduling, model
download + placement, cancellation, file I/O, streaming, and reporting back to
the control plane.

## Install

```bash
pip install gen-worker[torch]   # for PyTorch inference/training
pip install gen-worker          # plain Python (e.g. API-proxy endpoints)
```

Optional extras: `[images]` / `[audio]` / `[video]` for media I/O,
`[vision]` for torchvision.

## Hello world

**`pyproject.toml`** — the one config value:

```toml
[tool.gen_worker]
main = "myendpoint.main"
```

**`main.py`**:

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

class Input(msgspec.Struct):
    prompt: str

class Output(msgspec.Struct):
    text: str

@endpoint
def echo(ctx: RequestContext, payload: Input) -> Output:
    return Output(text=f"got: {payload.prompt}")
```

Run it locally, no orchestrator:

```bash
gen-worker run --payload '{"prompt": "hello"}'
```

`cozyctl build` / `cozyctl deploy` take it from here — the full path to a
deployed, billed endpoint is [tensorhub docs/writing-endpoints.md](https://github.com/cozy-creator/tensorhub/blob/master/docs/writing-endpoints.md).

## Adding a model

Hold state in a class: `setup()` runs once, every public method is one
routable function. The worker downloads the binding, constructs the pipeline
from the `setup()` annotation, and owns device placement + low-VRAM offload —
endpoint code never touches `.to("cuda")` or offload config.

```python
from diffusers import StableDiffusionXLPipeline
from gen_worker import HF, RequestContext, Resources, endpoint

@endpoint(
    model=HF("stabilityai/stable-diffusion-xl-base-1.0", dtype="bf16"),
    resources=Resources(gpu=True),
)
class Generate:
    def setup(self, pipeline: StableDiffusionXLPipeline) -> None:
        self.pipeline = pipeline

    def generate(self, ctx: RequestContext, payload: Input) -> Output:
        view = ctx.for_request(self.pipeline, seed=42)
        image = view(payload.prompt, generator=view.generator).images[0]
        return Output(text=ctx.save_image(image).ref)
```

`Resources` declares only what the endpoint CANNOT run without (`gpu`,
`gpu_count`, `libraries`, `strict_vram`, `vcpus`); VRAM requirements are
MEASURED by the platform's profiling gate, not declared (`vram_gb_hint` is
an optional first-build placement hint only). Handlers are exactly
`(self, ctx, payload)`; per-request state (sampler, seed, scheduler) lives
in a `ctx.for_request` view over shared weights — never assigned onto the
instance.

Bindings: `HF(id, revision=, dtype=, subfolder=, files=, storage_dtype=)`,
`Hub(ref, tag=, storage_dtype=)`, `Civitai(id, version=)`, `ModelScope(id, ...)`.
The slot name comes from the `models={}` key or the `setup()` parameter —
never a constructor argument. `storage_dtype="fp8"` keeps denoiser weights in
fp8-E4M3 storage with per-layer upcast to the compute `dtype` (half the VRAM
on any card); fp8-stored artifacts get the same treatment automatically.
Quantization itself is ahead-of-time only — a conversion endpoint produces the
artifact, never `setup()` (th#1803) — and the `flavor` axis is DELETED
(§1.32(d), pgw#1148): selection within a tag group is tensor-layout-contract
compatibility, declared per slot as `Slot(layouts=…)` (§1.33, pgw#1143). See
[docs/endpoint-authoring.md](docs/endpoint-authoring.md).

Curated checkpoint selection is a runtime payload argument: a handler declares
`model: SomeModelChoice` (a `ModelChoice` enum of `Model` rows, each carrying a
`ModelRef` binding + typed per-model defaults) and reads `payload.model.defaults`
typed — one `generate(model=)` replaces N near-identical functions. `model:
SomeModelChoice | ModelRef` opens BYOM. Streaming = an async-generator handler.
Engine-hosted endpoints declare `runtime="vllm"` and get a booted,
health-checked server subprocess injected into `setup()`.

`Slot(pipeline_cls, selected_by=, family=, default_checkpoint=)` is the hub-resolved
alternative to `ModelChoice`: the model SET lives in platform config, not
code, and the COMPONENT TREE (`pipeline.unet`, `pipeline.vae`, ...) is
derived from the pipeline class and published to the hub — parts are never
declared as sibling slots. The per-model config SCHEMA derives from the
handler's context annotation (`ctx: RequestContext[SdxlDefaults]`, a
`gen_worker.families.GenerationDefaults` vocabulary); the catalog owns the
VALUES and `ctx.defaults` hands the resolved recipe to the handler typed.
`family="sdxl"` explicitly keeps a non-root/defaultless model lane inside
that architecture's binding gate; `family=""` explicitly marks a shared
auxiliary as family-agnostic. Omitting it preserves the compatibility
inference.

Full reference: [docs/endpoint-authoring.md](docs/endpoint-authoring.md).

## Public surface

- The decorator + bindings: `endpoint`, `Resources`, `Compile`, `HF`, `Hub`,
  `Civitai`, `ModelScope`, `ModelRef`
- Model selection: `Model`, `ModelChoice`, `ModelDefaults`, `Slot`,
  `ResolvedSlot`, `gen_worker.families.GenerationDefaults`
- Compile envelope (the declared serving region — resolutions, text lengths,
  guidance, batch): `Compile`, `CompileAxis`, `AxisClass`, `DynamicDim`,
  `pad_text_sequence`; per-request views: `ctx.for_request` / `gen_worker.view`
- Contexts: `RequestContext` (≤15 members), `ConversionContext`,
  `DatasetContext`, `TrainingContext`
- Errors: `ValidationError`, `RetryableError`, `CanceledError`, `FatalError`
- Streaming: `BatchItemDelta`, `IncrementalTokenDelta`, `Done`, `Error`
- Value types: `Asset`, `ImageAsset`, `AudioAsset`, `VideoAsset`
- I/O codecs: `gen_worker.io`

The conversion ETL (hub ingest, dtype cast / quant, clone, Tensorhub
publish) is `gen_worker.convert` (see [docs/convert.md](docs/convert.md)).

## Local development

```bash
gen-worker run --payload '{"prompt": "hello"}'  # one-shot in-process
gen-worker run --list                            # describe functions (JSON)
gen-worker serve                                 # warm local server
gen-worker invoke <fn> prompt=hello              # client for serve
gen-worker prefetch                              # weights only, no GPU
```

stdout for results, stderr for events; exit 0 / 1 / 2 / 3 / 130 for success /
user-exception / usage / model-resolution / SIGINT. Details:
[docs/local-dev.md](docs/local-dev.md); host contract:
[docs/host-integration.md](docs/host-integration.md).

### Running tests

```bash
uv run --extra dev pytest
```

Plain `uv run pytest` would fall through to a global launcher — always pass
`--extra dev`. **Never `pip install` gen-worker globally:** a stale
`~/.local` install silently shadows the working tree (`tests/conftest.py`
hard-fails if `gen_worker` resolves outside `src/`).

## Documentation

- [docs/endpoint-authoring.md](docs/endpoint-authoring.md) — the `@endpoint`
  reference: bindings, variants, Resources, contexts, streaming, runtimes.
- [docs/local-dev.md](docs/local-dev.md) — the CLI: `run`/`serve`/`invoke`/
  `prefetch`, `field=value` grammar, `--offline`, exit codes.
- [docs/dockerfile.md](docs/dockerfile.md) — bring-your-own-Dockerfile contract.
- [docs/endpoint-envs.md](docs/endpoint-envs.md) — tenant envs/secrets.
- [docs/compile-cache.md](docs/compile-cache.md) — compiled cells: the graph
  digest vs the declared envelope, kernel lanes, JIT intake.

## Examples

- `examples/marco-polo/` — minimal inference endpoint (sync, async, streaming)
