Output image ID at end of build/push

go.mod

@@ -90,6 +90,7 @@ require (
 	github.com/golangci/unconvert v0.0.0-20180507085042-28b1c447d1f4 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
+	github.com/google/uuid v1.3.0 // indirect
 	github.com/gordonklaus/ineffassign v0.0.0-20210914165742-4cc7213b9bc8 // indirect
 	github.com/gostaticanalysis/analysisutil v0.7.1 // indirect
 	github.com/gostaticanalysis/comment v1.4.2 // indirect

go.sum

@@ -279,6 +279,8 @@ github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm4
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510/go.mod h1:pupxD2MaaD3pAXIBCelhxNneeOaAeabZDe5s4K6zSpQ=
 github.com/google/uuid v1.1.2/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
+github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+vpHVxEJEs9eg=
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/googleapis/google-cloud-go-testing v0.0.0-20200911160855-bcd43fbb19e8/go.mod h1:dvDLG8qkwmyD9a/MJJN3XJcT3xFxOKAvTZGvuZmac9g=

pkg/cli/build.go

@@ -38,12 +38,13 @@ func buildCommand(cmd *cobra.Command, args []string) error {
 		imageName = config.DockerImageName(projectDir)
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildProgressOutput); err != nil {
+	imageID, err := image.Build(cfg, projectDir, imageName, buildProgressOutput)
+	if err != nil {
 		return err
 	}
 
 	console.Infof("\nImage built as %s", imageName)
-
+	console.Output(imageID)
 	return nil
 }
 

pkg/cli/push.go

@@ -42,7 +42,8 @@ func push(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("To push images, you must either set the 'image' option in cog.yaml or pass an image name as an argument. For example, 'cog push registry.hooli.corp/hotdog-detector'")
 	}
 
-	if err := image.Build(cfg, projectDir, imageName, buildProgressOutput); err != nil {
+	imageID, err := image.Build(cfg, projectDir, imageName, buildProgressOutput)
+	if err != nil {
 		return err
 	}
 
@@ -57,5 +58,6 @@ func push(cmd *cobra.Command, args []string) error {
 			console.Infof("\nRun your model on Replicate:\n    %s", replicatePage)
 		}
 	}
+	console.Output(imageID)
 	return exitStatus
 }

pkg/docker/image_rm.go

@@ -0,0 +1,15 @@
+package docker
+
+import (
+	"os"
+	"os/exec"
+)
+
+func ImageRm(name string) error {
+	cmd := exec.Command("docker", "image", "rm", name)
+	cmd.Env = os.Environ()
+	cmd.Stderr = os.Stderr
+
+	_, err := cmd.Output()
+	return err
+}

pkg/docker/tag.go

@@ -0,0 +1,15 @@
+package docker
+
+import (
+	"os"
+	"os/exec"
+)
+
+func Tag(source, target string) error {
+	cmd := exec.Command("docker", "image", "tag", source, target)
+	cmd.Env = os.Environ()
+	cmd.Stderr = os.Stderr
+
+	_, err := cmd.Output()
+	return err
+}

pkg/image/build.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/google/uuid"
 	"github.com/replicate/cog/pkg/config"
 	"github.com/replicate/cog/pkg/docker"
 	"github.com/replicate/cog/pkg/dockerfile"
@@ -15,12 +16,16 @@ import (
 // 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, progressOutput string) error {
+func Build(cfg *config.Config, dir, imageName string, progressOutput string) (string, error) {
 	console.Infof("Building Docker image from environment in cog.yaml as %s...", imageName)
 
+	// Use a unique temporary ID for builds so there isn't a race condition where two builds could happen in parallel with the same image name,
+	// which may cause generating OpenAPI schema or adding labels to collide.
+	tempImageName := fmt.Sprintf("cog-build-%s", uuid.New().String())
+
 	generator, err := dockerfile.NewGenerator(cfg, dir)
 	if err != nil {
-		return fmt.Errorf("Error creating Dockerfile generator: %w", err)
+		return "", fmt.Errorf("Error creating Dockerfile generator: %w", err)
 	}
 	defer func() {
 		if err := generator.Cleanup(); err != nil {
@@ -30,21 +35,21 @@ func Build(cfg *config.Config, dir, imageName string, progressOutput string) err
 
 	dockerfileContents, err := generator.Generate()
 	if err != nil {
-		return fmt.Errorf("Failed to generate Dockerfile: %w", err)
+		return "", fmt.Errorf("Failed to generate Dockerfile: %w", err)
 	}
 
-	if err := docker.Build(dir, dockerfileContents, imageName, progressOutput); err != nil {
-		return fmt.Errorf("Failed to build Docker image: %w", err)
+	if err := docker.Build(dir, dockerfileContents, tempImageName, progressOutput); err != nil {
+		return "", fmt.Errorf("Failed to build Docker image: %w", err)
 	}
 
 	console.Info("Adding labels to image...")
-	schema, err := GenerateOpenAPISchema(imageName, cfg.Build.GPU)
+	schema, err := GenerateOpenAPISchema(tempImageName, cfg.Build.GPU)
 	if err != nil {
-		return fmt.Errorf("Failed to get type signature: %w", err)
+		return "", fmt.Errorf("Failed to get type signature: %w", err)
 	}
 	configJSON, err := json.Marshal(cfg)
 	if err != nil {
-		return fmt.Errorf("Failed to convert config to JSON: %w", err)
+		return "", fmt.Errorf("Failed to convert config to JSON: %w", err)
 	}
 	// We used to set the cog_version and config labels in Dockerfile, because we didn't require running the
 	// built image to get those. But, the escaping of JSON inside a label inside a Dockerfile was gnarly, and
@@ -62,16 +67,31 @@ func Build(cfg *config.Config, dir, imageName string, progressOutput string) err
 	if len((*schema).(map[string]interface{})) != 0 {
 		schemaJSON, err := json.Marshal(schema)
 		if err != nil {
-			return fmt.Errorf("Failed to convert type signature to JSON: %w", err)
+			return "", fmt.Errorf("Failed to convert type signature to JSON: %w", err)
 		}
 		labels[global.LabelNamespace+"openapi_schema"] = string(schemaJSON)
 		labels["org.cogmodel.openapi_schema"] = string(schemaJSON)
 	}
 
-	if err := docker.BuildAddLabelsToImage(imageName, labels); err != nil {
-		return fmt.Errorf("Failed to add labels to image: %w", err)
+	if err := docker.BuildAddLabelsToImage(tempImageName, labels); err != nil {
+		return "", fmt.Errorf("Failed to add labels to image: %w", err)
 	}
-	return nil
+
+	image, err := docker.ImageInspect(tempImageName)
+	if err != nil {
+		return "", fmt.Errorf("Failed to inspect image: %w", err)
+	}
+
+	if err := docker.Tag(tempImageName, imageName); err != nil {
+		return "", fmt.Errorf("Failed to tag image: %w", err)
+	}
+
+	// We've just tagged the image with a different name, so this removes the temporary tag, not the actual image
+	if err := docker.ImageRm(tempImageName); err != nil {
+		return "", fmt.Errorf("Failed to remove temporary image: %w", err)
+	}
+
+	return image.ID, nil
 }
 
 func BuildBase(cfg *config.Config, dir string, progressOutput string) (string, error) {