Metadata-Version: 2.4
Name: gen-worker
Version: 0.1.4
Summary: Add your description here
Author-email: Paul Fidika <paul@fidika.com>
Requires-Python: >=3.10
Requires-Dist: aiohttp>=3.11.14
Requires-Dist: backoff>=2.2.1
Requires-Dist: grpcio>=1.71.0
Requires-Dist: msgspec>=0.18.6
Requires-Dist: protobuf>=5.26.1
Requires-Dist: psutil>=7.0.0
Requires-Dist: pyjwt[crypto]>=2.8.0
Requires-Dist: tqdm>=4.66.0
Provides-Extra: dev
Requires-Dist: devpi-client>=7.2.0; extra == 'dev'
Requires-Dist: grpcio-tools>=1.71.0; extra == 'dev'
Provides-Extra: torch
Requires-Dist: flashpack>=0.1.2; extra == 'torch'
Requires-Dist: numpy>=1.26.0; extra == 'torch'
Requires-Dist: safetensors>=0.4.3; extra == 'torch'
Requires-Dist: torch>=2.6.0; extra == 'torch'
Requires-Dist: torchaudio>=2.6.0; extra == 'torch'
Requires-Dist: torchvision>=0.21.0; extra == 'torch'
Description-Content-Type: text/markdown

This is a python package, called gen_worker, which provides the worker runtime SDK:

- Orchestrator gRPC client + job loop
- Function discovery via @worker_function
- ActionContext + errors + progress events
- Model downloading from the Cozy hub (async + retries + progress)
- Output uploads (presigned PUT or S3 creds)

Torch-based model memory management is optional and installed via extras.

---

Files in src/gen_worker/pb must be auto-generated in the gen-orchestrator repo, using the proto files. Go in there and run `task proto`

Install modes:

- Core only: `gen-worker`
- Torch runtime add-on: `gen-worker[torch]` (torch + torchvision + torchaudio + safetensors + flashpack + numpy)

Example tenant projects live in `../worker-example-functions`. They use:

- `pyproject.toml` + `uv.lock` for dependencies (no requirements.txt)
- `cozy.toml` TOML manifest for deployment config (functions.modules, runtime.base_image, etc.)

Dependency policy:

- Require `pyproject.toml` and/or `uv.lock`
- Do not use `requirements.txt`
- Put Cozy deployment config in `cozy.toml`

Example:

```toml
[functions]
modules = ["functions"]

[runtime]
base_image = "ghcr.io/cozy/python-worker:cuda12.1-torch2.6"
```

Function signature:

```python
from gen_worker import worker_function, ResourceRequirements, ActionContext

@worker_function(ResourceRequirements(model_family="sdxl", requires_gpu=True))
def generate(ctx: ActionContext, payload: dict) -> dict:
    return {"ok": True}
```

Dynamic checkpoints:

- Use `ResourceRequirements(model_family=...)` to declare a family (e.g., "sdxl")
- Pass the exact checkpoint at runtime via request payload (e.g., `model_ref`)

Build contract (gen-builder):

- Tenant code + `pyproject.toml`/`uv.lock` + `cozy.toml` are packaged together
- gen-builder layers tenant code + deps on top of a python-worker base image
- gen-orchestrator deploys the resulting worker image

---

Env hints:

- `SCHEDULER_ADDR` sets the primary scheduler address.
- `SCHEDULER_ADDRS` (comma-separated) provides seed addresses for leader discovery.
- `WORKER_JWT` is accepted as the auth token if `AUTH_TOKEN` is not set.
- `SCHEDULER_JWKS_URL` enables verification of `WORKER_JWT` before connecting.
- JWT verification uses RSA and requires PyJWT crypto support (installed by default via `PyJWT[crypto]`).
- `WORKER_MAX_INPUT_BYTES`, `WORKER_MAX_OUTPUT_BYTES`, `WORKER_MAX_UPLOAD_BYTES` cap payload sizes.
- `WORKER_MAX_CONCURRENCY` limits concurrent runs; `ResourceRequirements(max_concurrency=...)` limits per-function.
- `COZY_HUB_URL` base URL for Cozy hub downloads (used by core downloader).
- `COZY_HUB_TOKEN` optional bearer token for Cozy hub downloads.
- `MODEL_MANAGER_CLASS` optional ModelManager plugin (module:Class) loaded at startup.

Error hints:

- Use `gen_worker.errors.RetryableError` in worker functions to flag retryable failures.

Output upload hints:

- To upload raw bytes, include an `output_upload` object in the job input:
  - Presigned PUT:
    - `{"output_upload":{"put_url":"...","headers":{"Content-Type":"image/png"},"public_url":"https://..."}}`
  - S3 creds:
    - `{"output_upload":{"s3":{"bucket":"...","key":"...","region":"...","access_key_id":"...","secret_access_key":"..."}}}`
