Add local image flag for fast context loading

pkg/cli/baseimage.go

@@ -122,7 +122,7 @@ func newBaseImageBuildCommand() *cobra.Command {
 			}
 			baseImageName := dockerfile.BaseImageName(baseImageCUDAVersion, baseImagePythonVersion, baseImageTorchVersion)
 
-			err = docker.Build(cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp)
+			err = docker.Build(cwd, dockerfileContents, baseImageName, []string{}, buildNoCache, buildProgressOutput, config.BuildSourceEpochTimestamp, dockerfile.StandardBuildDirectory, nil)
 			if err != nil {
 				return err
 			}

pkg/cli/build.go

@@ -25,6 +25,7 @@ var buildUseCogBaseImage bool
 var buildStrip bool
 var buildPrecompile bool
 var buildFast bool
+var buildLocalImage bool
 
 const useCogBaseImageFlagKey = "use-cog-base-image"
 
@@ -48,6 +49,7 @@ func newBuildCommand() *cobra.Command {
 	addStripFlag(cmd)
 	addPrecompileFlag(cmd)
 	addFastFlag(cmd)
+	addLocalImage(cmd)
 	cmd.Flags().StringVarP(&buildTag, "tag", "t", "", "A name for the built image in the form 'repository:tag'")
 	return cmd
 }
@@ -74,7 +76,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, buildFast, nil); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, nil, buildLocalImage); err != nil {
 		return err
 	}
 
@@ -147,6 +149,12 @@ func addFastFlag(cmd *cobra.Command) {
 	_ = cmd.Flags().MarkHidden(fastFlag)
 }
 
+func addLocalImage(cmd *cobra.Command) {
+	const localImage = "x-localimage"
+	cmd.Flags().BoolVar(&buildLocalImage, localImage, false, "Whether to use the experimental local image features")
+	_ = cmd.Flags().MarkHidden(localImage)
+}
+
 func checkMutuallyExclusiveFlags(cmd *cobra.Command, args []string) error {
 	flags := []string{useCogBaseImageFlagKey, "use-cuda-base-image", "dockerfile"}
 	var flagsSet []string

pkg/cli/debug.go

@@ -27,6 +27,8 @@ func newDebugCommand() *cobra.Command {
 	addDockerfileFlag(cmd)
 	addUseCogBaseImageFlag(cmd)
 	addBuildTimestampFlag(cmd)
+	addFastFlag(cmd)
+	addLocalImage(cmd)
 	cmd.Flags().StringVarP(&imageName, "image-name", "", "", "The image name to use for the generated Dockerfile")
 
 	return cmd
@@ -39,7 +41,7 @@ func cmdDockerfile(cmd *cobra.Command, args []string) error {
 	}
 
 	command := docker.NewDockerCommand()
-	generator, err := dockerfile.NewGenerator(cfg, projectDir, false, command)
+	generator, err := dockerfile.NewGenerator(cfg, projectDir, buildFast, command, buildLocalImage)
 	if err != nil {
 		return fmt.Errorf("Error creating Dockerfile generator: %w", err)
 	}

pkg/cli/predict.go

@@ -56,6 +56,7 @@ the prediction on that.`,
 	addGpusFlag(cmd)
 	addSetupTimeoutFlag(cmd)
 	addFastFlag(cmd)
+	addLocalImage(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")
@@ -76,22 +77,27 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 		if err != nil {
 			return err
 		}
+
 		if cfg.Build.Fast {
 			buildFast = cfg.Build.Fast
 		}
 
-		if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil {
-			return err
-		}
+		if buildFast {
+			imageName = config.DockerImageName(projectDir)
+		} else {
+			if imageName, err = image.BuildBase(cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput); err != nil {
+				return err
+			}
 
-		// Base image doesn't have /src in it, so mount as volume
-		volumes = append(volumes, docker.Volume{
-			Source:      projectDir,
-			Destination: "/src",
-		})
+			// Base image doesn't have /src in it, so mount as volume
+			volumes = append(volumes, docker.Volume{
+				Source:      projectDir,
+				Destination: "/src",
+			})
 
-		if gpus == "" && cfg.Build.GPU {
-			gpus = "all"
+			if gpus == "" && cfg.Build.GPU {
+				gpus = "all"
+			}
 		}
 
 	} else {
@@ -127,13 +133,17 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 
 	console.Info("")
 	console.Infof("Starting Docker image %s and running setup()...", imageName)
+	dockerCommand := docker.NewDockerCommand()
 
-	predictor := predict.NewPredictor(docker.RunOptions{
+	predictor, err := predict.NewPredictor(docker.RunOptions{
 		GPUs:    gpus,
 		Image:   imageName,
 		Volumes: volumes,
 		Env:     envFlags,
-	}, false, buildFast)
+	}, false, buildFast, dockerCommand)
+	if err != nil {
+		return err
+	}
 
 	go func() {
 		captureSignal := make(chan os.Signal, 1)
@@ -155,11 +165,14 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 			console.Info("Missing device driver, re-trying without GPU")
 
 			_ = predictor.Stop()
-			predictor = predict.NewPredictor(docker.RunOptions{
+			predictor, err = predict.NewPredictor(docker.RunOptions{
 				Image:   imageName,
 				Volumes: volumes,
 				Env:     envFlags,
-			}, false, buildFast)
+			}, false, buildFast, dockerCommand)
+			if err != nil {
+				return err
+			}
 
 			if err := predictor.Start(os.Stderr, timeout); err != nil {
 				return err
@@ -177,7 +190,7 @@ func cmdPredict(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
-	return predictIndividualInputs(predictor, inputFlags, outPath, false)
+	return predictIndividualInputs(*predictor, inputFlags, outPath, false)
 }
 
 func isURI(ref *openapi3.Schema) bool {

pkg/cli/push.go

@@ -36,6 +36,7 @@ func newPushCommand() *cobra.Command {
 	addStripFlag(cmd)
 	addPrecompileFlag(cmd)
 	addFastFlag(cmd)
+	addLocalImage(cmd)
 
 	return cmd
 }
@@ -63,6 +64,10 @@ func push(cmd *cobra.Command, args []string) error {
 		if err := docker.ManifestInspect(imageName); err != nil && strings.Contains(err.Error(), `"code":"NAME_UNKNOWN"`) {
 			return fmt.Errorf("Unable to find Replicate existing model for %s. Go to replicate.com and create a new model before pushing.", imageName)
 		}
+	} else {
+		if buildLocalImage {
+			return fmt.Errorf("Unable to push a local image model to a non replicate host, please disable the local image flag before pushing to this host.")
+		}
 	}
 
 	annotations := map[string]string{}
@@ -76,7 +81,7 @@ func push(cmd *cobra.Command, args []string) error {
 
 	startBuildTime := time.Now()
 
-	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, annotations); err != nil {
+	if err := image.Build(cfg, projectDir, imageName, buildSecrets, buildNoCache, buildSeparateWeights, buildUseCudaBaseImage, buildProgressOutput, buildSchemaFile, buildDockerfileFile, DetermineUseCogBaseImage(cmd), buildStrip, buildPrecompile, buildFast, annotations, buildLocalImage); err != nil {
 		return err
 	}
 

pkg/cli/run.go

@@ -37,6 +37,7 @@ func newRunCommand() *cobra.Command {
 	addUseCogBaseImageFlag(cmd)
 	addGpusFlag(cmd)
 	addFastFlag(cmd)
+	addLocalImage(cmd)
 
 	flags := cmd.Flags()
 	// Flags after first argument are considered args and passed to command
@@ -67,6 +68,8 @@ func run(cmd *cobra.Command, args []string) error {
 		gpus = "all"
 	}
 
+	dockerCommand := docker.NewDockerCommand()
+
 	runOptions := docker.RunOptions{
 		Args:    args,
 		Env:     envFlags,
@@ -75,6 +78,10 @@ func run(cmd *cobra.Command, args []string) error {
 		Volumes: []docker.Volume{{Source: projectDir, Destination: "/src"}},
 		Workdir: "/src",
 	}
+	runOptions, err = docker.FillInWeightsManifestVolumes(dockerCommand, runOptions)
+	if err != nil {
+		return err
+	}
 
 	if util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) {
 		runOptions.Platform = "linux/amd64"

pkg/cli/serve.go

@@ -69,6 +69,7 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
 		"--await-explicit-shutdown", "true",
 	}
 
+	dockerCommand := docker.NewDockerCommand()
 	runOptions := docker.RunOptions{
 		Args:    args,
 		Env:     envFlags,
@@ -77,6 +78,10 @@ func cmdServe(cmd *cobra.Command, arg []string) error {
 		Volumes: []docker.Volume{{Source: projectDir, Destination: "/src"}},
 		Workdir: "/src",
 	}
+	runOptions, err = docker.FillInWeightsManifestVolumes(dockerCommand, runOptions)
+	if err != nil {
+		return err
+	}
 
 	if util.IsAppleSiliconMac(runtime.GOOS, runtime.GOARCH) {
 		runOptions.Platform = "linux/amd64"

pkg/cli/train.go

@@ -56,13 +56,14 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 	volumes := []docker.Volume{}
 	gpus := gpusFlag
 
+	cfg, projectDir, err := config.GetConfig(projectDirFlag)
+	if err != nil {
+		return err
+	}
+
 	if len(args) == 0 {
 		// Build image
 
-		cfg, projectDir, err := config.GetConfig(projectDirFlag)
-		if err != nil {
-			return err
-		}
 		if cfg.Build.Fast {
 			buildFast = cfg.Build.Fast
 		}
@@ -108,14 +109,18 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 
 	console.Info("")
 	console.Infof("Starting Docker image %s...", imageName)
+	dockerCommand := docker.NewDockerCommand()
 
-	predictor := predict.NewPredictor(docker.RunOptions{
+	predictor, err := predict.NewPredictor(docker.RunOptions{
 		GPUs:    gpus,
 		Image:   imageName,
 		Volumes: volumes,
 		Env:     trainEnvFlags,
 		Args:    []string{"python", "-m", "cog.server.http", "--x-mode", "train"},
-	}, true, buildFast)
+	}, true, buildFast, dockerCommand)
+	if err != nil {
+		return err
+	}
 
 	go func() {
 		captureSignal := make(chan os.Signal, 1)
@@ -141,5 +146,5 @@ func cmdTrain(cmd *cobra.Command, args []string) error {
 		}
 	}()
 
-	return predictIndividualInputs(predictor, trainInputFlags, trainOutPath, true)
+	return predictIndividualInputs(*predictor, trainInputFlags, trainOutPath, true)
 }

pkg/docker/build.go

@@ -13,7 +13,7 @@ import (
 	"github.com/replicate/cog/pkg/util/console"
 )
 
-func Build(dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, epoch int64) error {
+func Build(dir, dockerfileContents, imageName string, secrets []string, noCache bool, progressOutput string, epoch int64, contextDir string, buildContexts map[string]string) error {
 	var args []string
 
 	args = append(args, "buildx", "build")
@@ -52,11 +52,15 @@ func Build(dir, dockerfileContents, imageName string, secrets []string, noCache
 		args = append(args, "--cache-to", "type=inline")
 	}
 
+	for name, dir := range buildContexts {
+		args = append(args, "--build-context", name+"="+dir)
+	}
+
 	args = append(args,
 		"--file", "-",
 		"--tag", imageName,
 		"--progress", progressOutput,
-		".",
+		contextDir,
 	)
 
 	cmd := exec.Command("docker", args...)

pkg/docker/command/manifest.go

@@ -20,3 +20,4 @@ const R8PythonVersionEnvVarName = "R8_PYTHON_VERSION"
 var CogConfigLabelKey = global.LabelNamespace + "config"
 var CogVersionLabelKey = global.LabelNamespace + "version"
 var CogOpenAPISchemaLabelKey = global.LabelNamespace + "openapi_schema"
+var CogWeightsManifestLabelKey = global.LabelNamespace + "r8_weights_manifest"

pkg/docker/run.go

@@ -3,6 +3,7 @@ package docker
 import (
 	"bufio"
 	"bytes"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -15,8 +16,10 @@ import (
 
 	"github.com/mattn/go-isatty"
 
+	"github.com/replicate/cog/pkg/docker/command"
 	"github.com/replicate/cog/pkg/util"
 	"github.com/replicate/cog/pkg/util/console"
+	"github.com/replicate/cog/pkg/weights"
 )
 
 type Port struct {
@@ -207,3 +210,27 @@ func GetPort(containerID string, containerPort int) (int, error) {
 	return 0, fmt.Errorf("did not find port bound to 0.0.0.0 in `docker port` output")
 
 }
+
+func FillInWeightsManifestVolumes(dockerCommand command.Command, runOptions RunOptions) (RunOptions, error) {
+	// Check if the image has a weights manifest
+	manifest, err := dockerCommand.Inspect(runOptions.Image)
+	if err != nil {
+		return runOptions, err
+	}
+	weightsManifest, ok := manifest.Config.Labels[command.CogWeightsManifestLabelKey]
+	if ok {
+		var weightsPaths []weights.WeightManifest
+		err = json.Unmarshal([]byte(weightsManifest), &weightsPaths)
+		if err != nil {
+			return runOptions, err
+		}
+		for _, weightPath := range weightsPaths {
+			runOptions.Volumes = append(runOptions.Volumes, Volume{
+				Source:      weightPath.Source,
+				Destination: "/src/" + weightPath.Destination,
+			})
+		}
+	}
+
+	return runOptions, nil
+}

pkg/dockerfile/base.go

@@ -180,7 +180,7 @@ func (g *BaseImageGenerator) GenerateDockerfile() (string, error) {
 		return "", err
 	}
 
-	generator, err := NewGenerator(conf, "", false, g.command)
+	generator, err := NewGenerator(conf, "", false, g.command, true)
 	if err != nil {
 		return "", err
 	}

pkg/dockerfile/build_tempdir.go

@@ -6,9 +6,11 @@ import (
 	"time"
 )
 
+const CogBuildArtifactsFolder = ".cog"
+
 func BuildCogTempDir(dir string, subDir string) (string, error) {
-	rootTmp := path.Join(dir, ".cog", "tmp", subDir)
-	if err := os.MkdirAll(rootTmp, 0o755); err != nil {
+	rootTmp := path.Join(dir, CogBuildArtifactsFolder, "tmp", subDir)
+	if err := os.MkdirAll(rootTmp, 0o777); err != nil {
 		return "", err
 	}
 	return rootTmp, nil