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

@@ -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"

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-launcher.build.override { inherit crateOverrides; })
-
+                (cargoNix.workspaceMembers.text-generation-router-v3.build.override { inherit crateOverrides; })
-                (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 ];
-                    };
-                  };
-                })
               ]);
 
             venvDir = "./.venv";

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 ];
+  };
+}