From 8e0cb7a47b850bfdc1fe7cdc5823650b2d5d6136 Mon Sep 17 00:00:00 2001
From: Yorick van Pelt <yorick@yorickvanpelt.nl>
Date: Tue, 14 May 2024 11:51:09 +0200
Subject: [PATCH] Switch from fetched cog-trt-llm to subdirectory

---
 default.nix | 6 ------
 flake.nix   | 6 ++----
 2 files changed, 2 insertions(+), 10 deletions(-)

diff --git a/default.nix b/default.nix
index 176e5bc..90a8b97 100644
--- a/default.nix
+++ b/default.nix
@@ -128,12 +128,6 @@ in
       patchelf --add-rpath ${pythonDrvs.nvidia-pytriton.public}/${site}/nvidia_pytriton.libs $f
     done
   '';
-  # TODO: open-source, switch to fetchFromGitHub
-  deps.cog-trt-llm = builtins.fetchGit {
-    url = "git@github.com:replicate/cog-trt-llm.git";
-    rev = "1f092d891b3cefeea5e0b4d39eb4406ebc60d99a";
-    ref = "main";
-  };
   deps.tensorrt-src = pkgs.fetchFromGitHub {
     owner = "NVIDIA";
     repo = "TensorRT";
diff --git a/flake.nix b/flake.nix
index 5a7145c..1919c0d 100644
--- a/flake.nix
+++ b/flake.nix
@@ -25,6 +25,7 @@
           rootDependencies = [ "nvidia-pytriton" "transformers" "tokenizers" ];
         };
         cognix.environment.TRITONSERVER_BACKEND_DIR = "${config.deps.backend_dir}/backends";
+        cognix.sourceIgnores = "cog-trt-llm/";
         # don't need this file in a runner
         python-env.pip.drvs.tensorrt-libs.mkDerivation.postInstall = lib.mkAfter ''
           rm $out/lib/python*/site-packages/tensorrt_libs/libnvinfer_builder_resource*
@@ -38,10 +39,7 @@
         };
         # override cog.yaml:
         cog.concurrency = lib.mkForce 1;
-        # copy cog-trt-llm source into /src
-        cognix.postCopyCommands = ''
-          cp ${config.deps.cog-trt-llm}/{*.py,cog-trt-llm-config.yaml} $out/src/
-        '';
+        cognix.rootPath = lib.mkForce "${./cog-trt-llm}";
         # this just needs the examples/ dir
         cognix.environment.TRTLLM_DIR = config.deps.tensorrt-llm.examples;
       });