Add local image flag for fast context loading (#2236)

* Add local image flag for fast context loading

* Add a new experimental flag —x-localimage for use
with —x-fastpush
* When activated, the flag uses the .cog tmp
directory as an rsync for the source code without
weights
* It separates the different build steps into
different build contexts
* On run/predict/train/serve it loads the weights
in as volume mounts
* This allows it to not reload the context each
time a small file is changed for weights

* Fix no tmpMount usage

* Potential fix for code scanning alert no. 24: Potentially unsafe quoting

Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Signed-off-by: Will Sackfield <[email protected]>

* Fix lint

* Use . in non local image for copy

* Add back weights copy exclusion from non local images

* Send standard build directory if not local image

* Fix calculating rel dir on relative dir

* Build then predict in int test

* Add fast and local image to predict test

* Use the image name from the config in predict

* Only if we use build fast

* Fix cog predict build on fast build

* Feed docker image in directly

* Use projectDir to fill in weights path

* Put debug in the right position

* Remove -t in cog predict

* Capture output of test and check return code

* Use absolute path for weights manifest

* Use hard links rather than file copies

* Mirror the local directories permissions

* Follow symlink to target

* Make sure directory permissions are aligned

* Fix fast push

* Synchronise directories between different golang
modules

* Use MkdirAll

* Use a set for weight path check

---------

Signed-off-by: Will Sackfield <[email protected]>
Signed-off-by: Will Sackfield <[email protected]>
Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

Author: 8W9aG

pkg/cli/baseimage.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/docker"
+	"github.com/replicate/cog/pkg/dockercontext"
 	"github.com/replicate/cog/pkg/dockerfile"
 	"github.com/replicate/cog/pkg/global"
 	"github.com/replicate/cog/pkg/update"
@@ -122,7 +123,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, dockercontext.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/fast_push.go

@@ -13,13 +13,16 @@ import (
 	"github.com/vbauerster/mpb/v8"
 
 	"github.com/replicate/cog/pkg/docker/command"
+	"github.com/replicate/cog/pkg/dockercontext"
 	"github.com/replicate/cog/pkg/monobeam"
 	"github.com/replicate/cog/pkg/requirements"
 	"github.com/replicate/cog/pkg/util"
 	"github.com/replicate/cog/pkg/web"
 	"github.com/replicate/cog/pkg/weights"
 )
 
+var TarballsDir = filepath.Join(dockercontext.CogBuildArtifactsFolder, "tarballs")
+
 const weightsObjectType = "weights"
 const filesObjectType = "files"
 const requirementsTarFile = "requirements.tar.zst"
@@ -46,8 +49,7 @@ func FastPush(ctx context.Context, image string, projectDir string, command comm
 	// Create a list of files to upload.
 	files := []web.File{}
 
-	tmpAptDir := filepath.Join(projectDir, ".cog", "tmp", "apt")
-	aptTarFile, err := CurrentAptTarball(tmpAptDir)
+	aptTarFile, err := CurrentAptTarball(dockercontext.CogTempDir(projectDir, dockercontext.AptBuildDir))
 	if err != nil {
 		return fmt.Errorf("current apt tarball error: %w", err)
 	}
@@ -67,15 +69,15 @@ func FastPush(ctx context.Context, image string, projectDir string, command comm
 		})
 	}
 
-	tmpRequirementsDir := filepath.Join(projectDir, ".cog", "tmp", "requirements")
+	tmpRequirementsDir := dockercontext.CogTempDir(projectDir, dockercontext.RequirementsBuildDir)
 	requirementsFile, err := requirements.CurrentRequirements(tmpRequirementsDir)
 	if err != nil {
 		return err
 	}
 
 	// Temp directory for tarballs extracted from the image
 	// Separate from other temp directories so that they don't cause cache invalidation
-	tmpTarballsDir := filepath.Join(projectDir, ".cog", "tmp", "tarballs")
+	tmpTarballsDir := filepath.Join(projectDir, TarballsDir)
 	// Upload python packages.
 	if requirementsFile != "" {
 		pythonTar, err := createPythonPackagesTarFile(image, tmpTarballsDir, command)