Add --use-cuda-base-image flag

pkg/cli/build.go

@@ -15,6 +15,7 @@ var buildSeparateWeights bool
 var buildSecrets []string
 var buildNoCache bool
 var buildProgressOutput string
+var buildUseCudaBaseImage string
 
 func newBuildCommand() *cobra.Command {
 	cmd := &cobra.Command{
@@ -27,6 +28,7 @@ func newBuildCommand() *cobra.Command {
 	addSecretsFlag(cmd)
 	addNoCacheFlag(cmd)
 	addSeparateWeightsFlag(cmd)
+	addUseCudaBaseImageFlag(cmd)
 	cmd.Flags().StringVarP(&buildTag, "tag", "t", "", "A name for the built image in the form 'repository:tag'")
 	return cmd
 }
@@ -45,7 +47,7 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 		imageName = config.DockerImageName(projectDir)
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildProgressOutput); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput); err != nil {
 		return err
 	}
 
@@ -73,3 +75,7 @@ func addNoCacheFlag(cmd *cobra.Command) {
 func addSeparateWeightsFlag(cmd *cobra.Command) {
 	cmd.Flags().BoolVar(&buildSeparateWeights, "separate-weights", false, "Separate model weights from code in image layers")
 }
+
+func addUseCudaBaseImageFlag(cmd *cobra.Command) {
+	cmd.Flags().StringVar(&buildUseCudaBaseImage, "use-cuda-base-image", "auto", "Use Nvidia CUDA base image, 'true' (default) or 'false' (use python base image). False results in a smaller image but may cause problems for non-torch projects")
+}

pkg/cli/debug.go

@@ -22,6 +22,7 @@ func newDebugCommand() *cobra.Command {
 	}
 
 	addSeparateWeightsFlag(cmd)
+	addUseCudaBaseImageFlag(cmd)
 	cmd.Flags().StringVarP(&imageName, "image-name", "", "", "The image name to use for the generated Dockerfile")
 
 	return cmd
@@ -43,6 +44,8 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
+	generator.SetUseCudaBaseImage(buildUseCudaBaseImage)
+
 	if buildSeparateWeights {
 		if imageName == "" {
 			imageName = config.DockerImageName(projectDir)

pkg/cli/predict.go

@@ -43,7 +43,10 @@ the prediction on that.`,
 		Args:       cobra.MaximumNArgs(1),
 		SuggestFor: []string{"infer"},
 	}
+
+	addUseCudaBaseImageFlag(cmd)
 	addBuildProgressOutputFlag(cmd)
+
 	cmd.Flags().StringArrayVarP(&inputFlags, "input", "i", []string{}, "Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg")
 	cmd.Flags().StringVarP(&outPath, "output", "o", "", "Output path")
 
@@ -63,7 +66,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 			return err
 		}
 
-		if imageName, err = image.BuildBase(cfg, projectDir, buildProgressOutput); err != nil {
+		if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, buildProgressOutput); err != nil {
 			return err
 		}
 

pkg/cli/push.go

@@ -22,10 +22,11 @@ func newPushCommand() *cobra.Command {
 		RunE:    push,
 		Args:    cobra.MaximumNArgs(1),
 	}
-	addBuildProgressOutputFlag(cmd)
 	addSecretsFlag(cmd)
 	addNoCacheFlag(cmd)
 	addSeparateWeightsFlag(cmd)
+	addUseCudaBaseImageFlag(cmd)
+	addBuildProgressOutputFlag(cmd)
 
 	return cmd
 }
@@ -45,7 +46,7 @@ func push(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push registry.hooli.corp/hotdog-detector'")
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildProgressOutput); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput); err != nil {
 		return err
 	}
 

pkg/cli/run.go

@@ -24,6 +24,7 @@ func newRunCommand() *cobra.Command {
 		Args:  cobra.MinimumNArgs(1),
 	}
 	addBuildProgressOutputFlag(cmd)
+	addUseCudaBaseImageFlag(cmd)
 
 	flags := cmd.Flags()
 	// Flags after first argment are considered args and passed to command
@@ -42,7 +43,7 @@ func run(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	imageName, err := image.BuildBase(cfg, projectDir, buildProgressOutput)
+	imageName, err := image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, buildProgressOutput)
 	if err != nil {
 		return err
 	}

pkg/cli/train.go

@@ -29,7 +29,10 @@ It will build the model in the current directory and train it.`,
 		Args:   cobra.MaximumNArgs(1),
 		Hidden: true,
 	}
+
 	addBuildProgressOutputFlag(cmd)
+	addUseCudaBaseImageFlag(cmd)
+
 	cmd.Flags().StringArrayVarP(&trainInputFlags, "input", "i", []string{}, "Inputs, in the form name=value. if value is prefixed with @, then it is read from a file on disk. E.g. -i path=@image.jpg")
 
 	return cmd
@@ -48,7 +51,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if imageName, err = image.BuildBase(cfg, projectDir, buildProgressOutput); err != nil {
+	if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, buildProgressOutput); err != nil {
 		return err
 	}
 

pkg/dockerfile/generator.go

@@ -48,6 +48,8 @@ type Generator struct {
 	GOOS   string
 	GOARCH string
 
+	useCudaBaseImage bool
+
 	// absolute path to tmpDir, a directory that will be cleaned up
 	tmpDir string
 	// tmpDir relative to Dir
@@ -73,16 +75,22 @@ func NewGenerator(config *config.Config, dir string) (*Generator, error) {
 	}
 
 	return &Generator{
-		Config:         config,
-		Dir:            dir,
-		GOOS:           runtime.GOOS,
-		GOARCH:         runtime.GOOS,
-		tmpDir:         tmpDir,
-		relativeTmpDir: relativeTmpDir,
-		fileWalker:     filepath.Walk,
+		Config:           config,
+		Dir:              dir,
+		GOOS:             runtime.GOOS,
+		GOARCH:           runtime.GOOS,
+		tmpDir:           tmpDir,
+		relativeTmpDir:   relativeTmpDir,
+		fileWalker:       filepath.Walk,
+		useCudaBaseImage: true,
 	}, nil
 }
 
+func (g *Generator) SetUseCudaBaseImage(argumentValue string) {
+	// "false" -> false, "true" -> true, "auto" -> true, "asdf" -> true
+	g.useCudaBaseImage = argumentValue != "false"
+}
+
 func (g *Generator) GenerateBase() (string, error) {
 	baseImage, err := g.baseImage()
 	if err != nil {
@@ -158,7 +166,7 @@ func (g *Generator) Generate(imageName string) (weightsBase string, dockerfile s
 		return "", "", "", err
 	}
 	installPython := ""
-	if g.Config.Build.GPU {
+	if g.Config.Build.GPU && g.useCudaBaseImage {
 		installPython, err = g.installPythonCUDA()
 		if err != nil {
 			return "", "", "", err
@@ -245,10 +253,10 @@ func (g *Generator) Cleanup() error {
 }
 
 func (g *Generator) baseImage() (string, error) {
-	if g.Config.Build.GPU {
+	if g.Config.Build.GPU && g.useCudaBaseImage {
 		return g.Config.CUDABaseImageTag()
 	}
-	return "python:" + g.Config.Build.PythonVersion, nil
+	return "python:" + g.Config.Build.PythonVersion + "-slim", nil
 }
 
 func (g *Generator) preamble() string {

pkg/dockerfile/generator_test.go

@@ -81,7 +81,7 @@ predict: predict.py:Predictor
 	expected := `#syntax=docker/dockerfile:1.4
 FROM r8.im/replicate/cog-test-weights AS weights
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
@@ -152,7 +152,7 @@ predict: predict.py:Predictor
 	expected := `#syntax=docker/dockerfile:1.4
 FROM r8.im/replicate/cog-test-weights AS weights
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
@@ -248,7 +248,7 @@ build:
 	expected := `#syntax=docker/dockerfile:1.4
 FROM r8.im/replicate/cog-test-weights AS weights
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin
@@ -428,7 +428,7 @@ predict: predict.py:Predictor
 
 	expected := `#syntax=docker/dockerfile:1.4
 ` + testTiniStage() +
-		`FROM python:3.8
+		`FROM python:3.8-slim
 ENV DEBIAN_FRONTEND=noninteractive
 ENV PYTHONUNBUFFERED=1
 ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64:/usr/local/nvidia/bin

pkg/image/build.go

@@ -20,7 +20,7 @@ const dockerignoreBackupPath = ".dockerignore.cog.bak"
 // Build a Cog model from a config
 //
 // This is separated out from docker.Build(), so that can be as close as possible to the behavior of 'docker build'.
-func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, progressOutput string) error {
+func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string) error {
 	console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName)
 
 	generator, err := dockerfile.NewGenerator(cfg, dir)
@@ -32,6 +32,7 @@ func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache,
 			console.Warnf("Error cleaning up Dockerfile generator: %s", err)
 		}
 	}()
+	generator.SetUseCudaBaseImage(useCudaBaseImage)
 
 	if separateWeights {
 		weightsDockerfile, runnerDockerfile, dockerignore, err := generator.Generate(imageName)
@@ -110,7 +111,7 @@ func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache,
 	return nil
 }
 
-func BuildBase(cfg *config.Config, dir string, progressOutput string) (string, error) {
+func BuildBase(cfg *config.Config, dir string, useCudaBaseImage string, progressOutput string) (string, error) {
 	// TODO: better image management so we don't eat up disk space
 	// https://github.com/replicate/cog/issues/80
 	imageName := config.BaseDockerImageName(dir)
@@ -125,6 +126,9 @@ func BuildBase(cfg *config.Config, dir string, progressOutput string) (string, e
 			console.Warnf("Error cleaning up Dockerfile generator: %s", err)
 		}
 	}()
+
+	generator.SetUseCudaBaseImage(useCudaBaseImage)
+
 	dockerfileContents, err := generator.GenerateBase()
 	if err != nil {
 		return "", fmt.Errorf("Failed to generate Dockerfile: %w", err)

pkg/util/mime/mime.go

@@ -59,6 +59,7 @@ var typeToExtension = map[string]string{
 	"font/woff2": ".woff2",
 
 	"image/bmp":                ".bmp",
+	"image/x-ms-bmp":           ".bmp",
 	"image/gif":                ".gif",
 	"image/jpeg":               ".jpg",
 	"image/png":                ".png",