Set useCogBaseImage to false in base-image (#1892)

* We never want to useCogBaseImage when generating base images.
replicate · Aug 21, 2024 · 2394064 · 2394064
1 parent 8c50fea
commit 2394064
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 0 deletions.
diff --git a/pkg/dockerfile/base.go b/pkg/dockerfile/base.go
@@ -180,6 +180,8 @@ func (g *BaseImageGenerator) GenerateDockerfile() (string, error) {
 	if err != nil {
 		return "", err
 	}
+	useCogBaseImage := false
+	generator.useCogBaseImage = &useCogBaseImage
 
 	dockerfile, err := generator.generateInitialSteps()
 	if err != nil {

diff --git a/pkg/dockerfile/base_test.go b/pkg/dockerfile/base_test.go
@@ -1,6 +1,7 @@
 package dockerfile
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -28,3 +29,15 @@ func TestBaseImageName(t *testing.T) {
 		require.Equal(t, tt.expected, actual)
 	}
 }
+
+func TestGenerateDockerfile(t *testing.T) {
+	generator, err := NewBaseImageGenerator(
+		"12.1",
+		"3.8",
+		"2.1.0",
+	)
+	require.NoError(t, err)
+	dockerfile, err := generator.GenerateDockerfile()
+	require.NoError(t, err)
+	require.True(t, strings.Contains(dockerfile, "FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04"))
+}