Add fast build flag to cog (#2086)

* Add an experimental —x-flag to the commands to
denote whether fast build/run/push is used.

pkg/cli/build.go

@@ -24,6 +24,7 @@ var buildDockerfileFile string
 var buildUseCogBaseImage bool
 var buildStrip bool
 var buildPrecompile bool
+var buildFast bool
 
 const useCogBaseImageFlagKey = "use-cog-base-image"
 
@@ -46,6 +47,7 @@ func newBuildCommand() *cobra.Command {
 	addBuildTimestampFlag(cmd)
 	addStripFlag(cmd)
 	addPrecompileFlag(cmd)
+	addFastFlag(cmd)
 	cmd.Flags().StringVarP(&buildTag, "tag", "t", "", "A name for the built image in the form 'repository:tag'")
 	return cmd
 }
@@ -69,7 +71,7 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast); err != nil {
 		return err
 	}
 
@@ -136,6 +138,12 @@ func addPrecompileFlag(cmd *cobra.Command) {
 	_ = cmd.Flags().MarkHidden(precompileFlag)
 }
 
+func addFastFlag(cmd *cobra.Command) {
+	const fastFlag = "x-fast"
+	cmd.Flags().BoolVar(&buildFast, fastFlag, false, "Whether to use the experimental fast features")
+	_ = cmd.Flags().MarkHidden(fastFlag)
+}
+
 func checkMutuallyExclusiveFlags(cmd *cobra.Command, args []string) error {
 	flags := []string{useCogBaseImageFlagKey, "use-cuda-base-image", "dockerfile"}
 	var flagsSet []string

pkg/cli/predict.go

@@ -55,6 +55,7 @@ the prediction on that.`,
 	addDockerfileFlag(cmd)
 	addGpusFlag(cmd)
 	addSetupTimeoutFlag(cmd)
+	addFastFlag(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 [email protected]")
 	cmd.Flags().StringVarP(&outPath, "output", "o", "", "Output path")
@@ -126,7 +127,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 		Image:   imageName,
 		Volumes: volumes,
 		Env:     envFlags,
-	}, false)
+	}, false, buildFast)
 
 	go func() {
 		captureSignal := make(chan os.Signal, 1)
@@ -152,7 +153,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 				Image:   imageName,
 				Volumes: volumes,
 				Env:     envFlags,
-			}, false)
+			}, false, buildFast)
 
 			if err := predictor.Start(os.Stderr, timeout); err != nil {
 				return err

pkg/cli/push.go

@@ -32,6 +32,7 @@ func newPushCommand() *cobra.Command {
 	addUseCogBaseImageFlag(cmd)
 	addStripFlag(cmd)
 	addPrecompileFlag(cmd)
+	addFastFlag(cmd)
 
 	return cmd
 }
@@ -58,12 +59,15 @@ func push(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast); err != nil {
 
 		return err
 	}
 
 	console.Infof("\nPushing image '%s'...", imageName)
+	if buildFast {
+		console.Info("Fast push enabled.")
+	}
 
 	err = docker.Push(imageName)
 	if err != nil {

pkg/cli/run.go

@@ -36,6 +36,7 @@ func newRunCommand() *cobra.Command {
 	addUseCudaBaseImageFlag(cmd)
 	addUseCogBaseImageFlag(cmd)
 	addGpusFlag(cmd)
+	addFastFlag(cmd)
 
 	flags := cmd.Flags()
 	// Flags after first argment are considered args and passed to command
@@ -91,6 +92,10 @@ func run(cmd *cobra.Command, args []string) error {
 	console.Info("")
 	console.Infof("Running '%s' in Docker with the current directory mounted as a volume...", strings.Join(args, " "))
 
+	if buildFast {
+		console.Info("Fast run enabled.")
+	}
+
 	err = docker.Run(runOptions)
 	// Only retry if we're using a GPU but but the user didn't explicitly select a GPU with --gpus
 	// If the user specified the wrong GPU, they are explicitly selecting a GPU and they'll want to hear about it

pkg/cli/serve.go

@@ -4,12 +4,13 @@ import (
 	"runtime"
 	"strings"
 
+	"github.com/spf13/cobra"
+
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/image"
 	"github.com/replicate/cog/pkg/util"
 	"github.com/replicate/cog/pkg/util/console"
-	"github.com/spf13/cobra"
 )
 
 var (
@@ -32,6 +33,7 @@ Generate and run an HTTP server based on the declared model inputs and outputs.`
 	addUseCudaBaseImageFlag(cmd)
 	addUseCogBaseImageFlag(cmd)
 	addGpusFlag(cmd)
+	addFastFlag(cmd)
 
 	cmd.Flags().IntVarP(&port, "port", "p", port, "Port on which to listen")
 
@@ -49,6 +51,10 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
 		return err
 	}
 
+	if buildFast {
+		console.Info("Fast serve enabled.")
+	}
+
 	gpus := ""
 	if gpusFlag != "" {
 		gpus = gpusFlag

pkg/cli/train.go

@@ -42,6 +42,7 @@ Otherwise, it will build the model in the current directory and train it.`,
 	addUseCudaBaseImageFlag(cmd)
 	addGpusFlag(cmd)
 	addUseCogBaseImageFlag(cmd)
+	addFastFlag(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 [email protected]")
 	cmd.Flags().StringArrayVarP(&trainEnvFlags, "env", "e", []string{}, "Environment variables, in the form name=value")
@@ -108,7 +109,7 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		Volumes: volumes,
 		Env:     trainEnvFlags,
 		Args:    []string{"python", "-m", "cog.server.http", "--x-mode", "train"},
-	}, true)
+	}, true, buildFast)
 
 	go func() {
 		captureSignal := make(chan os.Signal, 1)

pkg/image/build.go

@@ -33,8 +33,11 @@ var errGit = errors.New("git error")
 // 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, useCudaBaseImage string, progressOutput string, schemaFile string, dockerfileFile string, useCogBaseImage *bool, strip bool, precompile bool) error {
+func Build(cfg *config.Config, dir, imageName string, secrets []string, noCache, separateWeights bool, useCudaBaseImage string, progressOutput string, schemaFile string, dockerfileFile string, useCogBaseImage *bool, strip bool, precompile bool, fastFlag bool) error {
 	console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName)
+	if fastFlag {
+		console.Info("Fast build enabled.")
+	}
 
 	// remove bundled schema files that may be left from previous builds
 	_ = os.Remove(bundledSchemaFile)

pkg/predict/predictor.go

@@ -50,7 +50,11 @@ type Predictor struct {
 	port        int
 }
 
-func NewPredictor(runOptions docker.RunOptions, isTrain bool) Predictor {
+func NewPredictor(runOptions docker.RunOptions, isTrain bool, fastFlag bool) Predictor {
+	if fastFlag {
+		console.Info("Fast predictor enabled.")
+	}
+
 	if global.Debug {
 		runOptions.Env = append(runOptions.Env, "COG_LOG_LEVEL=debug")
 	} else {