diff --git a/pkg/cli/build.go b/pkg/cli/build.go index f08d456738..d9cd93be9a 100644 --- a/pkg/cli/build.go +++ b/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") +} diff --git a/pkg/cli/debug.go b/pkg/cli/debug.go index d7674def05..ac3a8e9555 100644 --- a/pkg/cli/debug.go +++ b/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) diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index 01a0ba1793..d4e8dec674 100644 --- a/pkg/cli/predict.go +++ b/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 } diff --git a/pkg/cli/push.go b/pkg/cli/push.go index cccbeccd9b..13a018acac 100644 --- a/pkg/cli/push.go +++ b/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 } diff --git a/pkg/cli/run.go b/pkg/cli/run.go index f7543b31f7..06c29f490d 100644 --- a/pkg/cli/run.go +++ b/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 } diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 0aca2041bb..e1f92ccc5c 100644 --- a/pkg/cli/train.go +++ b/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 } diff --git a/pkg/dockerfile/generator.go b/pkg/dockerfile/generator.go index 99b1f93eb5..7c5d3adedc 100644 --- a/pkg/dockerfile/generator.go +++ b/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 { diff --git a/pkg/dockerfile/generator_test.go b/pkg/dockerfile/generator_test.go index 350b132028..693b6a1211 100644 --- a/pkg/dockerfile/generator_test.go +++ b/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 diff --git a/pkg/image/build.go b/pkg/image/build.go index 1a4856e431..b6d6408691 100644 --- a/pkg/image/build.go +++ b/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) diff --git a/pkg/util/mime/mime.go b/pkg/util/mime/mime.go index bc91e6fb7c..924f402d72 100644 --- a/pkg/util/mime/mime.go +++ b/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",