From 9aaa12e7ac54d859ebfee4b023f66010cef1bca6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Thu, 15 Aug 2024 10:21:51 +0200
Subject: [PATCH] nix: build router incrementally (#2422)

---
 flake.lock | 22 ----------------------
 flake.nix  | 53 +++++++++++++++++++++++++----------------------------
 2 files changed, 25 insertions(+), 50 deletions(-)

diff --git a/flake.lock b/flake.lock
index dab582a5..ef05c258 100644
--- a/flake.lock
+++ b/flake.lock
@@ -579,27 +579,6 @@
         "type": "github"
       }
     },
-    "naersk": {
-      "inputs": {
-        "nixpkgs": [
-          "tgi-nix",
-          "nixpkgs"
-        ]
-      },
-      "locked": {
-        "lastModified": 1721727458,
-        "narHash": "sha256-r/xppY958gmZ4oTfLiHN0ZGuQ+RSTijDblVgVLFi1mw=",
-        "owner": "nix-community",
-        "repo": "naersk",
-        "rev": "3fb418eaf352498f6b6c30592e3beb63df42ef11",
-        "type": "github"
-      },
-      "original": {
-        "owner": "nix-community",
-        "repo": "naersk",
-        "type": "github"
-      }
-    },
     "nix-github-actions": {
       "inputs": {
         "nixpkgs": [
@@ -901,7 +880,6 @@
       "inputs": {
         "crate2nix": "crate2nix",
         "flake-utils": "flake-utils_6",
-        "naersk": "naersk",
         "nixpkgs": [
           "tgi-nix",
           "nixpkgs"
diff --git a/flake.nix b/flake.nix
index 229184d2..168c8649 100644
--- a/flake.nix
+++ b/flake.nix
@@ -7,10 +7,6 @@
     tgi-nix.url = "github:danieldk/tgi-nix";
     nixpkgs.follows = "tgi-nix/nixpkgs";
     flake-utils.url = "github:numtide/flake-utils";
-    naersk = {
-      url = "github:nix-community/naersk";
-      inputs.nixpkgs.follows = "tgi-nix/nixpkgs";
-    };
     poetry2nix.url = "github:nix-community/poetry2nix";
     rust-overlay = {
       url = "github:oxalica/rust-overlay";
@@ -21,7 +17,6 @@
     {
       self,
       crate2nix,
-      naersk,
       nixpkgs,
       flake-utils,
       rust-overlay,
@@ -34,6 +29,7 @@
         cargoNix = crate2nix.tools.${system}.appliedCargoNix {
           name = "tgi";
           src = ./.;
+          additionalCargoNixArgs = [ "--all-features" ];
         };
         config = {
           allowUnfree = true;
@@ -46,32 +42,10 @@
             tgi-nix.overlay
           ];
         };
-        naersk' = pkgs.callPackage naersk { };
-        router =
-          with pkgs;
-          naersk'.buildPackage {
-            name = "router";
-            src = ./.;
-            cargoBuildOptions =
-              x:
-              x
-              ++ [
-                "-p"
-                "text-generation-router-v3"
-              ];
-            nativeBuildInputs = [ pkg-config ];
-            buildInputs = [
-              openssl.dev
-              protobuf
-            ];
-            doCheck = false;
-          };
-
         inherit (poetry2nix.lib.mkPoetry2Nix { inherit pkgs; }) mkPoetryEditablePackage;
         text-generation-server = mkPoetryEditablePackage { editablePackageSources = ./server; };
       in
       {
-        defaultPackage = router;
         devShells.default =
           with pkgs;
           mkShell {
@@ -118,7 +92,30 @@
                 vllm
 
                 cargoNix.workspaceMembers.text-generation-launcher.build
-                router
+
+                (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";