diff --git a/pkg/cli/build.go b/pkg/cli/build.go index 92dda89c6b..b3050396b7 100644 --- a/pkg/cli/build.go +++ b/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 diff --git a/pkg/cli/predict.go b/pkg/cli/predict.go index ea8d4bde62..c3c325c54b 100644 --- a/pkg/cli/predict.go +++ b/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 path=@image.jpg") 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 diff --git a/pkg/cli/push.go b/pkg/cli/push.go index 3a6463a071..d56b8e3b57 100644 --- a/pkg/cli/push.go +++ b/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 { diff --git a/pkg/cli/run.go b/pkg/cli/run.go index 32f4b2cc29..bbb28b8a2f 100644 --- a/pkg/cli/run.go +++ b/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 diff --git a/pkg/cli/serve.go b/pkg/cli/serve.go index cb05124164..4801276d33 100644 --- a/pkg/cli/serve.go +++ b/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 diff --git a/pkg/cli/train.go b/pkg/cli/train.go index 312d6c76e5..725792c5f7 100644 --- a/pkg/cli/train.go +++ b/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 path=@image.jpg") 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) diff --git a/pkg/image/build.go b/pkg/image/build.go index 6173a5309f..7dac6b8c38 100644 --- a/pkg/image/build.go +++ b/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) diff --git a/pkg/predict/predictor.go b/pkg/predict/predictor.go index 3fd3049e95..72525d1b97 100644 --- a/pkg/predict/predictor.go +++ b/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 {