nix: build and cache impure devshells (#2765)

* nix: build and cache all devshells

* nix: add poetry to the impure shell

This shouldn't be used to manage dependencies in a Nix devshell, but can
be handy to update `poetry.lock`.

* Fix Nix build, disable pure shell (covered by Nix tests)

.github/workflows/nix_cache.yaml

@@ -0,0 +1,34 @@
+name: "Cache devshells"
+on:
+  pull_request:
+    paths:
+      - "flake.nix"
+      - "flake.lock"
+      - "nix/**"
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  tests:
+    runs-on:
+      group: aws-highmemory-32-plus-priv
+    steps:
+      - uses: actions/checkout@v4
+      - uses: cachix/install-nix-action@v27
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+      - uses: cachix/cachix-action@v14
+        with:
+          name: text-generation-inference
+          # If you chose signing key for write access
+          authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+        env:
+          USER: github_runner
+      - name: Build impure devshell
+        run: nix build .\#devShells.x86_64-linux.impure
+      - name: Build impure devshell (CUDA dev)
+        run: nix build .\#devShells.x86_64-linux.impureWithCuda
+      # Pure shell dependencies are covered by Nix tests.
+      # - name: Build pure devshell
+      #   run: nix build .\#devShells.x86_64-linux.pure

nix/impure-shell.nix

@@ -9,6 +9,7 @@
   cudaPackages,
   openssl,
   pkg-config,
+  poetry,
   protobuf,
   python3,
   pyright,
@@ -28,6 +29,7 @@ mkShell {
       black
       isort
       pkg-config
+      poetry
       (rust-bin.stable.latest.default.override {
         extensions = [
           "rust-analyzer"

nix/server.nix

@@ -30,6 +30,7 @@
   opentelemetry-semantic-conventions,
   outlines,
   peft,
+  pillow,
   prometheus-client,
   punica-kernels,
   py-cpuinfo,
@@ -69,6 +70,7 @@ buildPythonPackage {
     "huggingface-hub"
     "loguru"
     "opentelemetry-instrumentation-grpc"
+    "pillow"
     "sentencepiece"
     "typer"
   ];
@@ -102,6 +104,7 @@ buildPythonPackage {
     opentelemetry-semantic-conventions
     outlines
     peft
+    pillow
     prometheus-client
     punica-kernels
     py-cpuinfo