From b3c1678ed78f519376e5e9850d08c769f2c5ab24 Mon Sep 17 00:00:00 2001 From: andreasjansson Date: Mon, 3 May 2021 18:15:43 -0700 Subject: [PATCH] Build with Buildx on M1 Macs Signed-off-by: andreasjansson --- pkg/docker/local_builder.go | 92 +++++++++++++++++++++++++++++++------ pkg/global/global.go | 4 ++ pkg/model/compatibility.go | 3 +- 3 files changed, 83 insertions(+), 16 deletions(-) diff --git a/pkg/docker/local_builder.go b/pkg/docker/local_builder.go index 02c9bd1f31..ddc75b5237 100644 --- a/pkg/docker/local_builder.go +++ b/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) diff --git a/pkg/global/global.go b/pkg/global/global.go index 69a3045087..60dd363c4d 100644 --- a/pkg/global/global.go +++ b/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" +} diff --git a/pkg/model/compatibility.go b/pkg/model/compatibility.go index f82f101272..07026803ce 100644 --- a/pkg/model/compatibility.go +++ b/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