From 1411bfb9897bafc4d5fbc61d94b7d295db0ba70a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 16 Aug 2024 10:01:01 +0200
Subject: [PATCH] nix: try to reduce the number of Rust rebuilds (#2424)

Try to reduce the number of router/launcher rebuilds by filtering
sources. In this way, recompiles should only be triggered by changes
in Cargo or Rust files.
---
 flake.lock              | 16 ++++++++++
 flake.nix               | 30 ++++---------------
 nix/crate-overrides.nix | 66 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 25 deletions(-)
 create mode 100644 nix/crate-overrides.nix

diff --git a/flake.lock b/flake.lock
index ef05c2588..7c7723770 100644
--- a/flake.lock
+++ b/flake.lock
@@ -579,6 +579,21 @@
         "type": "github"
       }
     },
+    "nix-filter": {
+      "locked": {
+        "lastModified": 1710156097,
+        "narHash": "sha256-1Wvk8UP7PXdf8bCCaEoMnOT1qe5/Duqgj+rL8sRQsSM=",
+        "owner": "numtide",
+        "repo": "nix-filter",
+        "rev": "3342559a24e85fc164b295c3444e8a139924675b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "nix-filter",
+        "type": "github"
+      }
+    },
     "nix-github-actions": {
       "inputs": {
         "nixpkgs": [
@@ -880,6 +895,7 @@
       "inputs": {
         "crate2nix": "crate2nix",
         "flake-utils": "flake-utils_6",
+        "nix-filter": "nix-filter",
         "nixpkgs": [
           "tgi-nix",
           "nixpkgs"
diff --git a/flake.nix b/flake.nix
index 168c86494..cf05746a8 100644
--- a/flake.nix
+++ b/flake.nix
@@ -4,6 +4,7 @@
       url = "github:nix-community/crate2nix";
       inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
     };
+    nix-filter.url = "github:numtide/nix-filter";
     tgi-nix.url = "github:danieldk/tgi-nix";
     nixpkgs.follows = "tgi-nix/nixpkgs";
     flake-utils.url = "github:numtide/flake-utils";
@@ -17,6 +18,7 @@
     {
       self,
       crate2nix,
+      nix-filter,
       nixpkgs,
       flake-utils,
       rust-overlay,
@@ -44,6 +46,7 @@
         };
         inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryEditablePackage;
         text-generation-server = mkPoetryEditablePackage { editablePackageSources = ./server; };
+        crateOverrides = import ./nix/crate-overrides.nix { inherit pkgs nix-filter; };
       in
       {
         devShells.default =
@@ -91,31 +94,8 @@
                 transformers
                 vllm
 
-                cargoNix.workspaceMembers.text-generation-launcher.build
-
-                (cargoNix.workspaceMembers.text-generation-router-v3.build.override {
-                  crateOverrides = defaultCrateOverrides // {
-                    aws-lc-rs = attrs: {
-                      # aws-lc-rs does its own custom parsing of Cargo environment
-                      # variables like DEP_.*_INCLUDE. However buildRustCrate does
-                      # not use the version number, so the parsing fails.
-                      postPatch = ''
-                        substituteInPlace build.rs \
-                          --replace-fail \
-                          "assert!(!selected.is_empty()" \
-                          "// assert!(!selected.is_empty()"
-                      '';
-                    };
-                    rav1e = attrs: { env.CARGO_ENCODED_RUSTFLAGS = "-C target-feature=-crt-static"; };
-                    text-generation-router-v3 = attrs: {
-                      # We need to do the src/source root dance so that the build
-                      # has access to the protobuf file.
-                      src = ./.;
-                      postPatch = "cd backends/v3";
-                      buildInputs = [ protobuf ];
-                    };
-                  };
-                })
+                (cargoNix.workspaceMembers.text-generation-launcher.build.override { inherit crateOverrides; })
+                (cargoNix.workspaceMembers.text-generation-router-v3.build.override { inherit crateOverrides; })
               ]);
 
             venvDir = "./.venv";
diff --git a/nix/crate-overrides.nix b/nix/crate-overrides.nix
new file mode 100644
index 000000000..343b3b25b
--- /dev/null
+++ b/nix/crate-overrides.nix
@@ -0,0 +1,66 @@
+{ pkgs, nix-filter }:
+
+let
+  filter = nix-filter.lib;
+in
+with pkgs;
+defaultCrateOverrides
+// {
+  aws-lc-rs = attrs: {
+    # aws-lc-rs does its own custom parsing of Cargo environment
+    # variables like DEP_.*_INCLUDE. However buildRustCrate does
+    # not use the version number, so the parsing fails.
+    postPatch = ''
+      substituteInPlace build.rs \
+        --replace-fail \
+        "assert!(!selected.is_empty()" \
+        "// assert!(!selected.is_empty()"
+    '';
+  };
+  rav1e = attrs: { env.CARGO_ENCODED_RUSTFLAGS = "-C target-feature=-crt-static"; };
+
+  grpc-metadata = attrs: {
+    src =
+      filter {
+        root = ../backends/grpc-metadata;
+        include = with filter; [
+          isDirectory
+          (matchExt "rs")
+        ];
+      };
+  };
+  text-generation-launcer = attrs: {
+    src =
+      filter {
+        root = ../launcher;
+        include = with filter; [
+          isDirectory
+          (matchExt "rs")
+        ];
+      };
+  };
+  text-generation-router = attrs: {
+    src =
+      filter {
+        root = ../router;
+        include = with filter; [
+          isDirectory
+          (matchExt "rs")
+        ];
+      };
+  };
+  text-generation-router-v3 = attrs: {
+    # We need to do the src/source root dance so that the build
+    # has access to the protobuf file.
+    src = filter {
+      root = ../.;
+      include = with filter; [
+        isDirectory
+        (and (inDirectory "backends/v3") (matchExt "rs"))
+        (and (inDirectory "proto") (matchExt "proto"))
+      ];
+    };
+    postPatch = "cd backends/v3";
+    buildInputs = [ protobuf ];
+  };
+}