Add support for fast build images in cog run (#2408)

* Cog run fast builds the image then runs the
command on that image if in fast mode

Author: 8W9aG

pkg/cli/run.go

@@ -59,15 +59,44 @@ func run(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	client := registry.NewRegistryClient()
 
 	cfg, projectDir, err := config.GetConfig(configFilename)
 	if err != nil {
 		return err
 	}
-	client := registry.NewRegistryClient()
-	imageName, err := image.BuildBase(ctx, dockerClient, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput, client)
-	if err != nil {
-		return err
+
+	var imageName string
+	if cfg.Build.Fast || buildFast {
+		imageName = config.DockerImageName(projectDir)
+		err = image.Build(
+			ctx,
+			cfg,
+			projectDir,
+			imageName,
+			buildSecrets,
+			buildNoCache,
+			buildSeparateWeights,
+			buildUseCudaBaseImage,
+			buildProgressOutput,
+			buildSchemaFile,
+			buildDockerfileFile,
+			DetermineUseCogBaseImage(cmd),
+			buildStrip,
+			buildPrecompile,
+			cfg.Build.Fast || buildFast,
+			nil,
+			buildLocalImage,
+			dockerClient,
+			client)
+		if err != nil {
+			return err
+		}
+	} else {
+		imageName, err = image.BuildBase(ctx, dockerClient, cfg, projectDir, buildUseCudaBaseImage, DetermineUseCogBaseImage(cmd), buildProgressOutput, client)
+		if err != nil {
+			return err
+		}
 	}
 
 	gpus := ""
@@ -102,10 +131,6 @@ 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(ctx, dockerClient, 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

test-integration/test_integration/test_run.py

@@ -1,4 +1,5 @@
 import subprocess
+from pathlib import Path
 
 
 def test_run(tmpdir_factory, cog_binary):
@@ -46,3 +47,16 @@ def test_run_with_secret(tmpdir_factory, cog_binary):
     )
     assert b"RUN echo hello world" in result.stdout
     assert b"RUN --mount=type=secret,id=foo,target=secret.txt echo shh" in result.stdout
+
+
+def test_run_fast_build(cog_binary):
+    project_dir = Path(__file__).parent / "fixtures/fast-build"
+    result = subprocess.run(
+        [cog_binary, "run", "echo", "hello world"],
+        cwd=project_dir,
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    assert result.returncode == 0
+    assert result.stdout == "hello world\n"