feat: build image by platform (#1582)

* feat: build image by platform

Signed-off-by: hang lv <xlv20@fudan.edu.cn>

* feat: support sequential multi-platform image building

Signed-off-by: hang lv <xlv20@fudan.edu.cn>

* feat: set goos/goarch as default target platform

Signed-off-by: hang lv <xlv20@fudan.edu.cn>

* fix: lint

Signed-off-by: hang lv <xlv20@fudan.edu.cn>

* fix: pull correct base image

Signed-off-by: hang lv <xlv20@fudan.edu.cn>

---------

Signed-off-by: hang lv <xlv20@fudan.edu.cn>

Author: n063h

go.mod

@@ -180,7 +180,7 @@ require (
 	go.opentelemetry.io/otel/sdk v1.14.0 // indirect
 	go.opentelemetry.io/otel/trace v1.14.0 // indirect
 	go.opentelemetry.io/proto/otlp v0.19.0 // indirect
-	golang.org/x/exp v0.0.0-20230321023759-10a507213a29 // indirect
+  golang.org/x/exp v0.0.0-20230321023759-10a507213a29 // indirect
 	golang.org/x/mod v0.10.0 // indirect
 	golang.org/x/net v0.9.0 // indirect
 	golang.org/x/oauth2 v0.6.0 // indirect

pkg/app/build.go

@@ -15,6 +15,7 @@
 package app
 
 import (
+	"strings"
 	"time"
 
 	"github.com/sirupsen/logrus"
@@ -23,6 +24,7 @@ import (
 	buildutil "github.com/tensorchord/envd/pkg/app/build"
 	"github.com/tensorchord/envd/pkg/app/telemetry"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
+	"github.com/tensorchord/envd/pkg/util/runtimeutil"
 )
 
 var CommandBuild = &cli.Command{
@@ -89,6 +91,12 @@ To build and push the image to a registry:
 			Usage:   "Import the cache (e.g. type=registry,ref=<image>)",
 			Aliases: []string{"ic"},
 		},
+		&cli.StringFlag{
+			Name: "platform",
+			Usage: `Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64).
+	Build images with same tags could cause image overwriting, platform suffixes will be added to differentiate the images.`,
+			DefaultText: runtimeutil.GetRuntimePlatform(),
+		},
 	},
 	Action: build,
 }
@@ -106,12 +114,24 @@ func build(clicontext *cli.Context) error {
 	logger := logrus.WithField("builder-options", opt)
 	logger.Debug("starting build command")
 
-	builder, err := buildutil.GetBuilder(clicontext, opt)
-	if err != nil {
-		return err
-	}
-	if err = buildutil.InterpretEnvdDef(builder); err != nil {
-		return err
+	platforms := strings.Split(opt.Platform, ",")
+	for _, platform := range platforms {
+		o := opt
+		o.Platform = platform
+		if len(platforms) > 1 {
+			// Transform the platform suffix to comply with the tag naming rule.
+			o.Tag += "-" + strings.Replace(platform, "/", "-", 1)
+		}
+		builder, err := buildutil.GetBuilder(clicontext, o)
+		if err != nil {
+			return err
+		}
+		if err = buildutil.InterpretEnvdDef(builder); err != nil {
+			return err
+		}
+		if err := buildutil.BuildImage(clicontext, builder); err != nil {
+			return err
+		}
 	}
-	return buildutil.BuildImage(clicontext, builder)
+	return nil
 }

pkg/app/build/build.go

@@ -135,6 +135,7 @@ func ParseBuildOpt(clicontext *cli.Context) (builder.Options, error) {
 	exportCache := clicontext.String("export-cache")
 	importCache := clicontext.String("import-cache")
 	useProxy := clicontext.Bool("use-proxy")
+	platform := clicontext.String("platform")
 
 	opt := builder.Options{
 		ManifestFilePath: manifest,
@@ -148,6 +149,7 @@ func ParseBuildOpt(clicontext *cli.Context) (builder.Options, error) {
 		ExportCache:      exportCache,
 		ImportCache:      importCache,
 		UseHTTPProxy:     useProxy,
+		Platform:         platform,
 	}
 
 	debug := clicontext.Bool("debug")

pkg/app/up.go

@@ -29,6 +29,7 @@ import (
 	"github.com/tensorchord/envd/pkg/home"
 	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
 	"github.com/tensorchord/envd/pkg/types"
+	"github.com/tensorchord/envd/pkg/util/runtimeutil"
 )
 
 var CommandUp = &cli.Command{
@@ -141,6 +142,11 @@ var CommandUp = &cli.Command{
 			Usage:   "Import the cache (e.g. type=registry,ref=<image>)",
 			Aliases: []string{"ic"},
 		},
+		&cli.StringFlag{
+			Name:        "platform",
+			Usage:       "Specify the target platform for the build output, (for example, windows/amd64, linux/amd64, or darwin/arm64)",
+			DefaultText: runtimeutil.GetRuntimePlatform(),
+		},
 	},
 
 	Action: up,

pkg/builder/build.go

@@ -19,13 +19,15 @@ import (
 	"io"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/cockroachdb/errors"
 	"github.com/docker/cli/cli/config"
 	"github.com/moby/buildkit/client"
 	"github.com/moby/buildkit/client/llb"
 	"github.com/moby/buildkit/session"
 	"github.com/moby/buildkit/session/auth/authprovider"
+	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sync/errgroup"
 
@@ -157,7 +159,11 @@ func (b generalBuilder) Interpret() error {
 
 func (b generalBuilder) Compile(ctx context.Context) (*llb.Definition, error) {
 	envName := filepath.Base(b.BuildContextDir)
-	def, err := b.graph.Compile(ctx, envName, b.PubKeyPath, b.Options.ProgressMode)
+	platform, err := parsePlatform(b.Platform)
+	if err != nil {
+		return nil, err
+	}
+	def, err := b.graph.Compile(ctx, envName, b.PubKeyPath, platform, b.Options.ProgressMode)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to compile build.envd")
 	}
@@ -191,8 +197,9 @@ func (b generalBuilder) imageConfig(ctx context.Context) (string, error) {
 
 	env := b.graph.GetEnviron()
 	user := b.graph.GetUser()
+	platform := b.graph.GetPlatform()
 
-	data, err := ImageConfigStr(labels, ports, ep, env, user)
+	data, err := ImageConfigStr(labels, ports, ep, env, user, platform)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to get image config")
 	}
@@ -335,3 +342,19 @@ func constructSolveOpt(ce []client.CacheOptionsEntry, entry client.ExportEntry,
 	}
 	return opt
 }
+
+func parsePlatform(platform string) (*ocispecs.Platform, error) {
+	os, arch, variant := "linux", "amd64", ""
+	if platform == "" {
+		return &ocispecs.Platform{Architecture: arch, OS: os, Variant: variant}, nil
+	}
+	arr := strings.Split(platform, "/")
+	if len(arr) < 2 {
+		return nil, errors.New("invalid platform format, expected `os/arch[/variant]`")
+	}
+	os, arch = arr[0], arr[1]
+	if len(arr) >= 3 {
+		variant = arr[2]
+	}
+	return &ocispecs.Platform{Architecture: arch, OS: os, Variant: variant}, nil
+}

pkg/builder/types.go

@@ -49,6 +49,9 @@ type Options struct {
 	ImportCache string
 	// UseHTTPProxy uses HTTPS_PROXY/HTTP_PROXY/NO_PROXY in the build process.
 	UseHTTPProxy bool
+	// Specify the target platform for the build output.
+	// e.g. platform=linux/arm64,linux/amd64
+	Platform string
 }
 
 type generalBuilder struct {

pkg/builder/util.go

@@ -25,7 +25,7 @@ import (
 	"github.com/containerd/console"
 	"github.com/moby/buildkit/client"
 	gatewayclient "github.com/moby/buildkit/frontend/gateway/client"
-	v1 "github.com/opencontainers/image-spec/specs-go/v1"
+	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sirupsen/logrus"
 )
 
@@ -35,20 +35,20 @@ const (
 )
 
 func ImageConfigStr(labels map[string]string, ports map[string]struct{},
-	entrypoint []string, env []string, user string) (string, error) {
-	img := v1.Image{
-		Config: v1.ImageConfig{
+	entrypoint []string, env []string, user string, platform *ocispecs.Platform) (string, error) {
+	img := ocispecs.Image{
+		Config: ocispecs.ImageConfig{
 			Labels:       labels,
 			User:         user,
 			WorkingDir:   "/",
 			Env:          env,
 			ExposedPorts: ports,
 			Entrypoint:   entrypoint,
 		},
-		Architecture: "amd64",
+		Architecture: platform.Architecture,
 		// Refer to https://github.com/tensorchord/envd/issues/269#issuecomment-1152944914
-		OS: "linux",
-		RootFS: v1.RootFS{
+		OS: platform.OS,
+		RootFS: ocispecs.RootFS{
 			Type: "layers",
 		},
 	}

pkg/editor/vscode/util.go

@@ -22,7 +22,7 @@ import (
 	"strings"
 
 	"github.com/cockroachdb/errors"
-	v1 "github.com/opencontainers/image-spec/specs-go/v1"
+	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sirupsen/logrus"
 )
 
@@ -38,7 +38,7 @@ const (
 	PLATFORM_ALPINE_X64   = "alpine-x64"
 )
 
-func ConvertLLBPlatform(platform *v1.Platform) (string, error) {
+func ConvertLLBPlatform(platform *ocispecs.Platform) (string, error) {
 	// Convert opencontainers style platform to VSCode extension style platform.
 	switch platform.OS {
 	case "windows":

pkg/lang/ir/graph.go

@@ -18,12 +18,13 @@ import (
 	"context"
 
 	"github.com/moby/buildkit/client/llb"
+	ocispecs "github.com/opencontainers/image-spec/specs-go/v1"
 
 	"github.com/tensorchord/envd/pkg/progress/compileui"
 )
 
 type Graph interface {
-	Compile(ctx context.Context, envName string, pub string, progressMode string) (*llb.Definition, error)
+	Compile(ctx context.Context, envName string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error)
 
 	graphDebugger
 	graphVisitor
@@ -56,4 +57,5 @@ type graphVisitor interface {
 	GetHTTP() []HTTPInfo
 	GetRuntimeCommands() map[string]string
 	GetUser() string
+	GetPlatform() *ocispecs.Platform
 }

pkg/lang/ir/util.go

@@ -26,12 +26,16 @@ import (
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
-func FetchImageConfig(ctx context.Context, imageName string) (config v1.ImageConfig, err error) {
+func FetchImageConfig(ctx context.Context, imageName string, platform *v1.Platform) (config v1.ImageConfig, err error) {
 	ref, err := docker.ParseReference(fmt.Sprintf("//%s", imageName))
 	if err != nil {
 		return config, errors.Wrap(err, "failed to parse image reference")
 	}
 	sys := types.SystemContext{}
+	if platform != nil {
+		sys.ArchitectureChoice = platform.Architecture
+		sys.OSChoice = platform.OS
+	}
 	src, err := ref.NewImageSource(ctx, &sys)
 	if err != nil {
 		return config, errors.Wrap(err, "failed to get image source from ref")