Use correct CUDA base image

Signed-off-by: andreasjansson <[email protected]>

Author: andreasjansson

pkg/docker/generate.go

@@ -79,7 +79,7 @@ func (g *DockerfileGenerator) baseImage() (string, error) {
 	case "cpu":
 		return "ubuntu:20.04", nil
 	case "gpu":
-		return g.gpuBaseImage()
+		return g.Config.CUDABaseImageTag()
 	}
 	return "", fmt.Errorf("Invalid architecture: %s", g.Arch)
 }
@@ -89,11 +89,6 @@ func (g *DockerfileGenerator) preamble() string {
 	return "ENV DEBIAN_FRONTEND=noninteractive"
 }
 
-func (g *DockerfileGenerator) gpuBaseImage() (string, error) {
-	// TODO: return correct ubuntu version for tf / torch
-	return "nvidia/cuda:11.0-devel-ubuntu20.04", nil
-}
-
 func (g *DockerfileGenerator) aptInstalls() (string, error) {
 	packages := g.Config.Environment.SystemPackages
 	if len(packages) == 0 {

pkg/docker/generate_test.go

@@ -55,6 +55,7 @@ func TestGenerateEmpty(t *testing.T) {
 model: infer.py:Model
 `))
 	require.NoError(t, err)
+	require.NoError(t, conf.ValidateAndCompleteConfig())
 
 	expectedCPU := `FROM ubuntu:20.04
 ENV DEBIAN_FRONTEND=noninteractive
@@ -64,7 +65,7 @@ COPY . /code
 WORKDIR /code
 CMD ["python", "-c", "from infer import Model; Model().start_server()"]`
 
-	expectedGPU := `FROM nvidia/cuda:11.0-devel-ubuntu20.04
+	expectedGPU := `FROM nvidia/cuda:11.0-cudnn8-devel-ubuntu16.04
 ENV DEBIAN_FRONTEND=noninteractive
 ` + installPython("3.8") + installCog() + `
 RUN ### --> Copying code
@@ -96,9 +97,7 @@ environment:
 model: infer.py:Model
 `))
 	require.NoError(t, err)
-
-	err = conf.ValidateAndCompleteConfig()
-	require.NoError(t, err)
+	require.NoError(t, conf.ValidateAndCompleteConfig())
 
 	expectedCPU := `FROM ubuntu:20.04
 ENV DEBIAN_FRONTEND=noninteractive
@@ -115,7 +114,7 @@ COPY . /code
 WORKDIR /code
 CMD ["python", "-c", "from infer import Model; Model().start_server()"]`
 
-	expectedGPU := `FROM nvidia/cuda:11.0-devel-ubuntu20.04
+	expectedGPU := `FROM nvidia/cuda:10.2-cudnn8-devel-ubuntu18.04
 ENV DEBIAN_FRONTEND=noninteractive
 ` + installPython("3.8") + `RUN ### --> Installing system packages
 RUN apt-get update -qq && apt-get install -qy ffmpeg cowsay && rm -rf /var/lib/apt/lists/*

pkg/model/compatibility.go

@@ -12,6 +12,11 @@ import (
 	"github.com/replicate/cog/pkg/version"
 )
 
+// TODO(andreas): check tf/py versions. tf 1.5.0 didn't install on py 3.8
+// TODO(andreas): support more tf versions. No matching tensorflow CPU package for version 1.15.4, etc.
+// TODO(andreas): allow user to install versions that aren't compatible
+// TODO(andreas): allow user to install tf cpu package on gpu
+
 type TFCompatibility struct {
 	TF           string
 	TFCPUPackage string
@@ -105,13 +110,13 @@ var CUDABaseImages []CUDABaseImage
 
 func init() {
 	if err := json.Unmarshal(tfCompatibilityMatrixData, &TFCompatibilityMatrix); err != nil {
-		console.Fatal("Failed to load embedded Tensorflow compatibility matrix: %s", err)
+		console.Fatalf("Failed to load embedded Tensorflow compatibility matrix: %s", err)
 	}
 	if err := json.Unmarshal(torchCompatibilityMatrixData, &TorchCompatibilityMatrix); err != nil {
-		console.Fatal("Failed to load embedded PyTorch compatibility matrix: %s", err)
+		console.Fatalf("Failed to load embedded PyTorch compatibility matrix: %s", err)
 	}
 	if err := json.Unmarshal(cudaBaseImageTagsData, &CUDABaseImages); err != nil {
-		console.Fatal("Failed to load embedded CUDA base images: %s", err)
+		console.Fatalf("Failed to load embedded CUDA base images: %s", err)
 	}
 }
 
@@ -186,10 +191,11 @@ func latestCuDNNForCUDA(cuda string) string {
 func latestTF() TFCompatibility {
 	var latest *TFCompatibility
 	for _, compat := range TFCompatibilityMatrix {
+		compat := compat
 		if latest == nil {
 			latest = &compat
 		} else {
-			greater, err := versionGreater(latest.TF, compat.TF)
+			greater, err := versionGreater(compat.TF, latest.TF)
 			if err != nil {
 				// should never happen
 				panic(fmt.Sprintf("Invalid tensorflow version: %s", err))