From 07bed530f7eaf2419ed0e755e0f24d7afd814a46 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Wed, 20 Nov 2024 20:56:11 +0100
Subject: [PATCH] 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 | 34 ++++++++++++++++++++++++++++++++
 nix/impure-shell.nix             |  2 ++
 nix/server.nix                   |  3 +++
 3 files changed, 39 insertions(+)
 create mode 100644 .github/workflows/nix_cache.yaml

diff --git a/.github/workflows/nix_cache.yaml b/.github/workflows/nix_cache.yaml
new file mode 100644
index 00000000000..967a5982e05
--- /dev/null
+++ b/.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
diff --git a/nix/impure-shell.nix b/nix/impure-shell.nix
index 92e14bc3cb0..a13fd711146 100644
--- a/nix/impure-shell.nix
+++ b/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"
diff --git a/nix/server.nix b/nix/server.nix
index 5903a65a1b8..237102a8c56 100644
--- a/nix/server.nix
+++ b/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