build(web_api): slim CPU/GPU extras + image size reduction

[magic-vladyslav]

Summary

Carves the web_api deploy images away from the monolithic modules extra and onto purpose-built, slim web_api (CPU) / web_api-gpu extras that contain only what web_api/app and internal.alignment actually import. Also makes the CPU image ship CPU-only torch (no CUDA wheels) and the GPU image drop the unused TensorRT/CUDA-13 stack, and restores the newest torch for everyone else.

Net effect: dramatically smaller, faster web_api images, and a web_api dev on Apple Silicon no longer has to resolve the other team's tensorrt/training deps.

What web_api actually imports → covering extra

Verified by tracing every import in web_api/app/*.py and recursively through internal.alignment (sift → misalignment_detector → manual_correspondence → field → online_finetuner). Nothing reachable from web_api imports mazepa, mazepa_addons, training, lightning, wandb, torchmetrics, meshing/skeletonization/chunkedgraph/montaging/calcada, or TensorRT.

web_api import	external pkg	covering extra
task_management.* (tasks.py)	—	task_management (→ databackends, sql, tenacity, pcg_skel→caveclient)
db_annotations.* (annotations/collections/layers/...)	—	databackends+tenacity (via task_management)
layer.volumetric + .cloudvol + .annotation (painting/precomputed)	einops	cloudvol + tensorstore (volumetric __init__ loads .tensorstore) + tensor_ops
internal.alignment.sift	scipy, cv2, numpy, torch	tensor_ops; scipy declared explicitly
internal.alignment.misalignment_detector	einops, torch	convnet
internal.alignment.manual_correspondence/field/online_finetuner	torch, einops, torchfields	tensor_ops
segmentation.py	cutie, hydra-core, omegaconf, torch, google.cloud.storage	new cutie sub-extra; hydra/omegaconf explicit; gcs via base cloud-files
main.py	google-auth	google-cloud-iap (web stack)

Hidden / transitive gaps closed

• scipy (internal/alignment/sift.py, alignment.py) — only transitive via scikit-image → declared explicitly.
• hydra-core / omegaconf (segmentation.py) — only transitive via the git cutie → declared explicitly.
• caveclient / tenacity / google-cloud-storage — verified not gaps (caveclient via pcg_skel, tenacity via task_management, gcs via cloud-files).
• cchardet — pruned (doesn't build on 3.12+); cutie needs it → CPU image keeps the faust-cchardet/stub shim.

New / changed extras (pyproject.toml)

• New cutie sub-extra; segmentation now references zetta_utils[cutie] (modules/segmentation semantics unchanged for the other team).
• New web-api-base (shared deps) + web-api / web-api-gpu leaf extras.
• CPU-only torch for web-api: a pytorch-cpu index + [tool.uv.sources] bind torch to it for the web-api extra, with [tool.uv] conflicts between web-api/web-api-gpu. requirements.web_api.txt resolves torch==…+cpu with 0 nvidia-* CUDA wheels. uv-only — plain pip install '.[web-api-gpu]' ignores it, so the GPU image keeps its base cu121 torch.
• web-api-gpu omits the gpu/tensorrt extra — web_api only calls convnet.load_model(tensorrt_enabled=False), so TensorRT is never imported; layering a CUDA-13 runtime on the CUDA-12.1 base was pure bloat + a version mismatch.
• torch floors: restored >= 2.11 on training/alignment/montaging; kept >= 2.5 on the web_api-path extras (convnet, tensor-typing, web-api) so the cu121 base's torch 2.5.1 is honored. --resolution highest still pins 2.12 everywhere else, so non-GPU consumers get the newest torch.

Removed from the web_api images

training (lightning/wandb/torchmetrics), mazepa_addons (kubernetes/awscli/mitmproxy/gcloud SDKs), meshing, skeletonization, montaging, chunkedgraph, calcada, and (CPU) the native abiss/waterz/lsds builds + their numpy==1.26.4/cython/nanobind machinery — plus the full nvidia-* CUDA stack (CPU) and tensorrt-cu13 + CUDA-13 libs (GPU).

Image size wins

• CPU image: CPU-only torch, no CUDA libs / no triton → ~4–5 GB lighter.
• GPU image: no CUDA-13 TensorRT stack → ~3–5 GB lighter, and the CUDA-12.1/13 mismatch is gone.

Dockerfiles

• web_api/Dockerfile (CPU): install requirements.web_api.txt --no-deps with PIP_EXTRA_INDEX_URL=…/whl/cpu; keep the cchardet shim; drop libboost/unixodbc + standalone cutie + abiss/waterz/lsd machinery; replace RUN zetta --help (pulls kubernetes) with a python -c "import app.main" smoke test.
• web_api/gpu.Dockerfile: pip install '.[web-api-gpu]' on the pytorch/pytorch:2.5.1-cuda12.1-cudnn9-runtime base, keeping PIP_EXTRA_INDEX_URL=…/cu121; drop numpy/cython/lsd/nanobind.
• Deleted web_api/requirements.txt — single source of truth is the extra.

Scripts & CI

• update_pinned_requirements.sh: exports requirements.web_api.txt / requirements.web_api_gpu.txt (lock/prune/fork-strategy untouched).
• install_zutils.py: --mode gains web_api / web_api_gpu.
• build_web_api.py: builds the CPU & GPU variants in parallel by default (--no-parallel to opt out) with per-variant prefixed output.
• .github/workflows/testing.yaml: new web-api-extras-build job (py 3.11/3.12/3.14: clean CPU install with the cpu torch index → smoke-import app.main → assert lightning/wandb/torchmetrics/mitmproxy/awscli/kubernetes/tensorrt absent) and web-api-gpu-build job (full GPU docker build). Both added to all-checks-test.

Verification

• web_api.app.main imports on CPU (darwin) including internal.alignment + cutie + hydra/omegaconf/scipy; no hardcoded .cuda().
• Pinned files confirmed: web_api → torch==…+cpu, 0 CUDA libs, no heavy pkgs; web_api_gpu → CUDA torch, no tensorrt; modules/all → CUDA torch + tensorrt unchanged (other team unaffected).
• web_api-gpu resolves cleanly with torch==2.5.1; modules now correctly rejects torch==2.5.1 (requires >=2.11).
• update_pinned_requirements.sh runs on Apple Silicon without choking on tensorrt (static-metadata workaround).

Note for the other team

pyproject.toml is shared, but modules/training/segmentation/all semantics are preserved: modules still resolves with tensorrt + CUDA torch, segmentation still includes cutie (via the cutie sub-extra). No undeclared internal dependency was added — scipy/hydra-core/omegaconf were already resolved transitively and are now declared explicitly in the web_api extra only.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 100.00%. Comparing base (72cd1e9) to head (06ea78c).
⚠️ Report is 4 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff            @@
##              main     #1263   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files          211       211           
  Lines        11292     11314   +22     
=========================================
+ Hits         11292     11314   +22     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dodamih]

This would break tensorrt which requires CUDA 13. What needs CUDA 12.1?

[nkemnitz]

I had to modify the dependency installation partially due to cutie (but not only). pip respects the requires_python settings that maintainers specify in their packages, but their upper bounds are often too limiting (often they just haven't been tested, but work fine on newer Python, and/or the package simply has been abandonded - like cutie).
For update_pinned_requirements.sh and install_zutils.sh I am now using uv, which apparently ignores the requires_python upper bounds by design. That should also resolve the need to rebuild the super-outdated cchardet: https://github.com/ZettaAI/zetta_utils/blob/main/install_zutils.py#L779-L780

Also please make sure to rerun update_pinned_requirements.sh whenever you really need to update pyproject.toml dependencies.

[nkemnitz]

This would break tensorrt which requires CUDA 13. What needs CUDA 12.1?

I think it's because Cloud Run driver's for L4 GPU can't be updated: https://github.com/ZettaAI/zetta_utils/blob/main/web_api/gpu.Dockerfile#L3-L6

[nkemnitz]

See comments, but a major one regarding the web_api/gpu build:

The R535 limitation on Cloud Run does not prevent you from using newer Pytorch / CUDA driver.

• Updating to Pytorch 2.12 + CUDA 12.6 in the requirements should be straightforward, thanks to minor-version compatibility. That already resolves the issue with Pytorch 2.5 <-> 2.12
• CUDA 13.0 is also possible with apt install cuda-compat-13-0, and prepending LD_LIBRARY_PATH with /usr/local/cuda-13.0/compat, which should contain the libcuda.so.
• I would drop the pytorch/pytorch image in either case. and rely on the pinned requirements files to install the version zutils actually prefers. Otherwise we need to keep paying close attention to the requirements and update the base image version to stay in sync with the resolved dependencies

[nkemnitz]

The install scripts explicitly use --no-deps, so this change drops tensorrt libs and bindings from the pinned requirements and break the main image. Try preserving the dependencies for the metapackage, that might still bypass the macOS issue.

Suggested change

	# backend; pip still pulls sub-package deps from the index at install time.
	[[tool.uv.dependency-metadata]]
	name = "tensorrt-cu13"
	requires-dist = []
	# backend
	[[tool.uv.dependency-metadata]]
	name = "tensorrt-cu13"
	requires-dist = ["tensorrt-cu13-libs", "tensorrt-cu13-bindings"]

[nkemnitz]

cchardet is declared by cutie, but never used. You neither need cchardet nor faust-cchardet.

[nkemnitz]

cchardet is declared by cutie, but never used. You neither need cchardet nor faust-cchardet.

[nkemnitz]

Nothing consumes this file. In gpu.Dockerfile you are relying on pip+pyproject.toml to resolve dependencies. Also note that this requirements file here pins torch==2.12 and CUDA 13 libraries, which you wanted to avoid in the web_api/gpu.Dockerfile.

[nkemnitz]

Switching to uv pip here (like the base image does via install_zutils.sh) may help with consistency/reproducibility. It's also faster than pip

[supersergiy]

Could you please address Nicos input?