Build with Buildx on M1 Macs

pkg/docker/local_builder.go

@@ -7,9 +7,11 @@ import (
 	"os/exec"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"strings"
 
 	"github.com/replicate/cog/pkg/console"
+	"github.com/replicate/cog/pkg/global"
 	"github.com/replicate/cog/pkg/logger"
 	"github.com/replicate/cog/pkg/shell"
 )
@@ -28,6 +30,8 @@ func NewLocalImageBuilder(registry string) *LocalImageBuilder {
 }
 
 func (b *LocalImageBuilder) Build(dir string, dockerfileContents string, name string, logWriter logger.Logger) (tag string, err error) {
+	buildRequiresGPU := false // TODO(andreas)
+
 	console.Debugf("Building in %s", dir)
 
 	// TODO(andreas): pipe dockerfile contents to builder
@@ -37,19 +41,27 @@ func (b *LocalImageBuilder) Build(dir string, dockerfileContents string, name st
 		return "", fmt.Errorf("Failed to write Dockerfile")
 	}
 
-	// shelling out to docker build because it's easier to get logs this way
-	// than when using the sdk
-	cmd := exec.Command(
-		"docker", "build", ".",
-		"--progress", "plain",
-		"-f", dockerfilePath,
-		// "--build-arg", "BUILDKIT_INLINE_CACHE=1",
-	)
-	cmd.Dir = dir
-	// TODO(andreas): follow https://github.com/moby/buildkit/issues/1436, hopefully buildkit will be able to use GPUs soon
-	cmd.Env = append(os.Environ(), "DOCKER_BUILDKIT=0")
+	var cmd *exec.Cmd
+	var outputPipe shell.PipeFunc
+	if global.IsM1Mac(runtime.GOOS, runtime.GOARCH) {
+		cmd, outputPipe = b.buildxCommand(dir, dockerfilePath, logWriter)
+		if err != nil {
+			return "", err
+		}
+	} else if buildRequiresGPU {
+		// TODO(andreas): follow https://github.com/moby/buildkit/issues/1436, hopefully buildkit will be able to use GPUs soon
+		cmd, outputPipe = b.legacyCommand(dir, dockerfilePath, logWriter)
+		if err != nil {
+			return "", err
+		}
+	} else {
+		cmd, outputPipe = b.buildKitCommand(dir, dockerfilePath, logWriter)
+		if err != nil {
+			return "", err
+		}
+	}
 
-	lastLogsChan, tagChan, err := buildPipe(cmd.StdoutPipe, logWriter)
+	lastLogsChan, tagChan, err := buildPipe(outputPipe, logWriter)
 	if err != nil {
 		return "", err
 	}
@@ -65,15 +77,14 @@ func (b *LocalImageBuilder) Build(dir string, dockerfileContents string, name st
 		}
 		return "", err
 	}
-
 	dockerTag := <-tagChan
 
+	logWriter.Infof("Successfully built %s", dockerTag)
+
 	if err != nil {
 		return "", err
 	}
 
-	logWriter.Infof("Successfully built %s", dockerTag)
-
 	tag = dockerTag
 	if name != "" {
 		tag = fmt.Sprintf("%s/%s:%s", b.registry, strings.ToLower(name), dockerTag)
@@ -85,6 +96,57 @@ func (b *LocalImageBuilder) Build(dir string, dockerfileContents string, name st
 	return tag, nil
 }
 
+func (b *LocalImageBuilder) legacyCommand(dir string, dockerfilePath string, logWriter logger.Logger) (*exec.Cmd, shell.PipeFunc) {
+	// shelling out to docker build because it's easier to get logs this way
+	// than when using the sdk
+	cmd := exec.Command(
+		"docker", "build", ".",
+		"--progress", "plain",
+		"-f", dockerfilePath,
+	)
+	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), "DOCKER_BUILDKIT=0")
+
+	console.Debug("Using legacy builder")
+
+	return cmd, cmd.StdoutPipe
+}
+
+func (b *LocalImageBuilder) buildKitCommand(dir string, dockerfilePath string, logWriter logger.Logger) (*exec.Cmd, shell.PipeFunc) {
+	// shelling out to docker build because it's easier to get logs this way
+	// than when using the sdk
+	cmd := exec.Command(
+		"docker", "build", ".",
+		"--progress", "plain",
+		"-f", dockerfilePath,
+		"--build-arg", "BUILDKIT_INLINE_CACHE=1",
+	)
+	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), "DOCKER_BUILDKIT=1")
+
+	console.Debug("Using BuildKit")
+
+	return cmd, cmd.StderrPipe
+}
+
+func (b *LocalImageBuilder) buildxCommand(dir string, dockerfilePath string, logWriter logger.Logger) (*exec.Cmd, shell.PipeFunc) {
+	// shelling out to docker build because it's easier to get logs this way
+	// than when using the sdk
+	cmd := exec.Command(
+		"docker", "buildx", "build", ".",
+		"--progress", "plain",
+		"-f", dockerfilePath,
+		"--build-arg", "BUILDKIT_INLINE_CACHE=1",
+		"--platform", "linux/amd64",
+	)
+	cmd.Dir = dir
+	cmd.Env = append(os.Environ(), "DOCKER_BUILDKIT=1")
+
+	console.Debug("Using Buildx")
+
+	return cmd, cmd.StderrPipe
+}
+
 func (b *LocalImageBuilder) tag(dockerTag string, tag string, logWriter logger.Logger) error {
 	console.Debugf("Tagging %s as %s", dockerTag, tag)
 

pkg/global/global.go

@@ -13,3 +13,7 @@ var (
 	ConfigFilename   = "cog.yaml"
 	CogServerAddress = "https://cog.replicate.ai"
 )
+
+func IsM1Mac(goos string, goarch string) bool {
+	return goos == "darwin" && goarch == "arm64"
+}

pkg/model/compatibility.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/replicate/cog/pkg/console"
+	"github.com/replicate/cog/pkg/global"
 
 	"github.com/replicate/cog/pkg/version"
 )
@@ -345,7 +346,7 @@ func torchvisionGPUPackage(ver string, cuda string) (name string, cpuVersion str
 // TODO(andreas): clean up this hack by actually parsing the torch_stable.html list in the generator
 func torchStripCPUSuffixForM1(version string, goos string, goarch string) string {
 	// TODO(andreas): clean up this hack
-	if goos == "darwin" && goarch == "arm64" {
+	if global.IsM1Mac(goos, goarch) {
 		return strings.ReplaceAll(version, "+cpu", "")
 	}
 	return version