From 584b4d7a68a00991a61134b602364f6db33587a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Tue, 1 Oct 2024 18:02:06 +0200
Subject: [PATCH]  nix: experimental support for building a Docker container
 (#2470)

* nix: experimental support for building a Docker image

Run using something like:

```
docker run \
  --device nvidia.com/gpu=all \
  -it --rm -p 8080:80 \
  -v $PWD/data:/data \
  -v $PWD/tmp:/tmp \
  tgi-docker:latest \
  --model-id <model_id>
```

* Example of building the Docker image using Nix inside Docker

* Stream to make the builder image smaller

This avoids storing a Docker image tarball in the image. Instead,
stream the layers while doing `docker run`.

* Don't spam journalctl on Linux

* Other dockerfile.

---------

Co-authored-by: Nicolas Patry <patry.nicolas@protonmail.com>
---
 .dockerignore  |  1 +
 Dockerfile.nix | 24 ++++++++++++++++++++++++
 flake.nix      | 29 ++++++++++++++++++++---------
 nix/docker.nix | 23 +++++++++++++++++++++++
 4 files changed, 68 insertions(+), 9 deletions(-)
 create mode 100644 Dockerfile.nix
 create mode 100644 nix/docker.nix

diff --git a/.dockerignore b/.dockerignore
index 1c641e7a..0b495924 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -4,3 +4,4 @@ server/transformers
 server/flash-attention
 cmake-build-debug/
 cmake-build-release/
+Dockerfile*
diff --git a/Dockerfile.nix b/Dockerfile.nix
new file mode 100644
index 00000000..f1e7e0f5
--- /dev/null
+++ b/Dockerfile.nix
@@ -0,0 +1,24 @@
+# Build the image and get out the docker file:
+#
+# docker build -t tgi-nix-builder -f Dockerfile.nix
+# docker run --log-driver=none tgi-nix-builder | docker load
+
+FROM nixos/nix:2.18.8 AS builder
+RUN echo "experimental-features = nix-command flakes" >> /etc/nix/nix.conf
+RUN nix profile install nixpkgs#cachix
+RUN cachix use text-generation-inference
+WORKDIR /root
+ADD . .
+RUN nix build .
+RUN mkdir /tmp/nix-store-closure
+RUN cp -R $(nix-store -qR result/) /tmp/nix-store-closure
+
+FROM ubuntu:24.04
+
+WORKDIR /app
+
+# Copy /nix/store
+COPY --from=builder /tmp/nix-store-closure /nix/store
+COPY --from=builder /root/result /app
+RUN ldconfig
+CMD ["ldconfig", "/app/bin/text-generation-launcher"]
diff --git a/flake.nix b/flake.nix
index be19e908..b18a5765 100644
--- a/flake.nix
+++ b/flake.nix
@@ -141,15 +141,26 @@
           };
         };
 
-        packages.default = pkgs.writeShellApplication {
-          name = "text-generation-inference";
-          runtimeInputs = [
-            server
-            router
-          ];
-          text = ''
-            ${launcher}/bin/text-generation-launcher "$@"
-          '';
+        packages = rec {
+          default = pkgs.writeShellApplication {
+            name = "text-generation-inference";
+            runtimeInputs = [
+              server
+              router
+            ];
+            text = ''
+              ${launcher}/bin/text-generation-launcher "$@"
+            '';
+          };
+
+          dockerImage = pkgs.callPackage nix/docker.nix {
+            text-generation-inference = default;
+          };
+
+          dockerImageStreamed = pkgs.callPackage nix/docker.nix {
+            text-generation-inference = default;
+            stream = true;
+          };
         };
       }
     );
diff --git a/nix/docker.nix b/nix/docker.nix
new file mode 100644
index 00000000..c4b1d899
--- /dev/null
+++ b/nix/docker.nix
@@ -0,0 +1,23 @@
+{
+  dockerTools,
+  cacert,
+  text-generation-inference,
+  stream ? false,
+}:
+
+let
+  build = if stream then dockerTools.streamLayeredImage else dockerTools.buildLayeredImage;
+in
+build {
+  name = "tgi-docker";
+  tag = "latest";
+  config = {
+    EntryPoint = [ "${text-generation-inference}/bin/text-generation-inference" ];
+    Env = [
+      "HF_HOME=/data"
+      "PORT=80"
+    ];
+
+  };
+  contents = [ cacert ];
+}