default to not including TRT

neggles · Sep 26, 2023 · 9e1971b · 9e1971b
1 parent 904d2d0
commit 9e1971b
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 7 deletions.
diff --git a/.github/workflows/build-push-base.yaml b/.github/workflows/build-push-base.yaml
@@ -61,7 +61,8 @@ jobs:
       fail-fast: false
       max-parallel: 1
       matrix:
-        target: ["base"]
+        target:
+          - base
         cuda:
           - cu118
           - cu121
@@ -73,6 +74,12 @@ jobs:
             cuda: cu121
             torch: torch201
         include:
+          - target: trt
+            cuda: cu121
+            torch: torch201
+          - target: trt
+            cuda: cu121
+            torch: torch210
           - target: base
             cuda: cu121
             torch: nightly

diff --git a/docker-bake.hcl b/docker-bake.hcl
@@ -94,7 +94,11 @@ target "common" {
 }
 
 target "base" {
-  name     = stripName("base-${cuda.name}-torch${torch.version}")
+  name = stripName(
+    cuda.with-trt
+    ? "trt-${cuda.name}-torch${torch.version}"
+    : "base-${cuda.name}-torch${torch.version}"
+  )
   inherits = ["common", "docker-metadata-action"]
   context  = "docker/base"
   target   = equal(torch.xformers, "") ? "base" : "xformers-binary"
@@ -124,12 +128,19 @@ target "base" {
     ],
     cuda = [
       {
-        name    = "cu118"
-        version = "11.8.0"
+        name     = "cu118"
+        version  = "11.8.0"
+        with-trt = false
       },
       {
-        name    = "cu121"
-        version = "12.1.1"
+        name     = "cu121"
+        version  = "12.1.1"
+        with-trt = false
+      },
+      {
+        name     = "cu121"
+        version  = "12.1.1"
+        with-trt = true
       }
     ]
   }
@@ -142,6 +153,7 @@ target "base" {
     TORCH_PACKAGE    = "torch==${torch.version}+${cudaName(cuda.version)}"
     TRITON_PACKAGE   = torch.triton
     XFORMERS_PACKAGE = torch.xformers
+    INCLUDE_TRT      = cuda.with-trt
   }
 }
 

diff --git a/docker/base/Dockerfile b/docker/base/Dockerfile
@@ -105,7 +105,7 @@ ARG INCLUDE_TRT=1
 RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
     --mount=type=cache,target=/var/lib/apt,sharing=locked \
     apt-get update \
-  && if [ "${INCLUDE_TRT}" == "1" ]; then  \
+  && if [ "${INCLUDE_TRT}" == "true" ]; then  \
     apt-get -y install --no-install-recommends \
         libnvinfer-dev \
         python3-libnvinfer-dev \