Skip to content

datakami/cognix

Repository files navigation

Cognix [WIP]

Build cog images deterministically using Nix.

Example (./examples/torch-demo/default.nix):

# build the image and load it into docker
$ nix build github:yorickvp/cognix#torch-demo && ./result load

$ cog predict torch-demo:......

Starting Docker image torch-demo:.... and running setup()...
Running prediction...
cuda works!

$ docker image ls torch-demo:....
REPOSITORY   TAG                                IMAGE ID       CREATED        SIZE
torch-demo   c209d6h86w7j7ksnjks4rkynjfw3ahwb   8fe343a42975   53 years ago   4.7GB

# explore the contribution of various packages to the image size
# format: store-path                                                           own-size  size-with-deps
$ nix path-info ./result -rSsh | sort -hk3
[..]
/nix/store/rbrw4jb0bz54kznbwwf47lj7k77jg00j-python3.10-nvidia-cudnn-cu11-8.5.0.96      	 868.8M	   1.4G
/nix/store/qqrzrbsgzk3bbg1pfficq5l2qnyz2b2k-python3.10-torch-2.0.1                     	   1.3G	   4.3G

Example (./examples/ebsynth-cpu/default.nix):

$ git clone ...
$ nix build .#ebsynth-cpu && ./result load
$ cog predict ebsynth-cpu:...... -i style=@...

Why?

  • Smaller images, they only contain the packages needed at runtime.
  • Determinism! Building an image twice on different machines should give you the exact same image, even after years.
  • Completeness: The dependencies for an image are fully specified, so there are no build steps you don't know about.
  • Better, granular caching: Nix generates a docker image with a layer for every dependency (up to ~limit, then it starts merging cleverly), images sharing dependencies will share layers as well.

Checklist

  • GPU support
  • Automatically parsing cog.yaml
  • Auto-resolving torch/cuda versions
  • Better ergonomics: module system
  • Fully sort out cog/r8 compat
  • Other python versions than 3.10
  • System dependencies mapping
  • cog.yaml run commands
  • Support upgrading cog at runtime
  • Downloading weights at build-time (example)
  • Uploading weights to gcs bucket (nix run .#gte-large.push-weights)
  • Downloading weights at run-time (example)
  • Generating openapi.yaml during the build
  • Automatically download weights at run-time
  • Ergonomic interface