feat(builder): Abstract BuildFunc to use gateway client (#606)

* feat(builder): Abstract BuildFunc to use gateway client

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update license

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Update

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Remove unused logic

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Fix test cases

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

pkg/app/build.go

@@ -120,9 +120,9 @@ func build(clicontext *cli.Context) error {
  "output": output,
  }).Debug("starting build command")
  builder, err := builder.New(clicontext.Context, cfg,
- manifest, funcname, buildContext, tag, output, debug)
+ manifest, funcname, buildContext, tag, output, debug, clicontext.Path("public-key"))
  if err != nil {
  return errors.Wrap(err, "failed to create the builder")
  }
- return builder.Build(clicontext.Context, clicontext.Path("public-key"))
+ return builder.Build(clicontext.Context)
 }

pkg/app/up.go

@@ -153,12 +153,12 @@ func up(clicontext *cli.Context) error {
  debug := clicontext.Bool("debug")
  output := ""
  builder, err := builder.New(clicontext.Context, config, manifest, funcname,
- buildContext, tag, output, debug)
+ buildContext, tag, output, debug, clicontext.Path("public-key"))
  if err != nil {
  return errors.Wrap(err, "failed to create the builder")
  }
 
- if err := builder.Build(clicontext.Context, clicontext.Path("public-key")); err != nil {
+ if err := builder.Build(clicontext.Context); err != nil {
  return errors.Wrap(err, "failed to build the image")
  }
  // Do not attach GPU if the flag is set.

pkg/builder/build.go

@@ -0,0 +1,54 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package builder
+
+import (
+ "context"
+
+ "github.com/cockroachdb/errors"
+ "github.com/moby/buildkit/frontend/gateway/client"
+)
+
+func (b generalBuilder) BuildFunc() func(ctx context.Context, c client.Client) (*client.Result, error) {
+ return func(ctx context.Context, c client.Client) (*client.Result, error) {
+ depsFiles := []string{
+ b.pubKeyPath,
+ b.configFilePath,
+ b.manifestFilePath,
+ }
+ isUpdated, err := b.CheckDepsFileUpdate(ctx, b.tag, depsFiles)
+ if err != nil {
+ b.logger.Debugf("failed to check manifest update: %s", err)
+ }
+ if !isUpdated {
+ b.logger.Infof("manifest is not updated, skip building")
+ return nil, nil
+ }
+
+ def, err := b.compile(ctx)
+ if err != nil {
+ return nil, errors.Wrap(err, "failed to compile")
+ }
+
+ res, err := c.Solve(ctx, client.SolveRequest{
+ Definition: def.ToPB(),
+ })
+ if err != nil {
+ return nil, err
+ }
+
+ return res, nil
+ }
+}

pkg/builder/builder.go

@@ -39,7 +39,7 @@ import (
 )
 
 type Builder interface {
- Build(ctx context.Context, pub string) error
+ Build(ctx context.Context) error
  GPUEnabled() bool
  NumGPUs() int
 }
@@ -51,6 +51,7 @@ type generalBuilder struct {
  tag string
  buildContextDir string
  buildfuncname string
+ pubKeyPath string
 
  entries []client.ExportEntry
 
@@ -60,7 +61,8 @@ type generalBuilder struct {
 }
 
 func New(ctx context.Context, configFilePath, manifestFilePath, funcname,
- buildContextDir, tag string, output string, debug bool) (Builder, error) {
+ buildContextDir, tag string, output string, debug bool,
+ pubKeyPath string) (Builder, error) {
  entries, err := parseOutput(output)
  if err != nil {
  return nil, errors.Wrap(err, "failed to parse output")
@@ -89,6 +91,7 @@ func New(ctx context.Context, configFilePath, manifestFilePath, funcname,
  configFilePath: configFilePath,
  buildContextDir: buildContextDir,
  entries: entries,
+ pubKeyPath: pubKeyPath,
  progressMode: mode,
  tag: tag,
  logger: logrus.WithFields(logrus.Fields{
@@ -150,31 +153,13 @@ func (b generalBuilder) CheckDepsFileUpdate(ctx context.Context, tag string, dep
  return false, nil
 }
 
-func (b generalBuilder) Build(ctx context.Context, pub string) error {
- depsFiles := []string{
- pub,
- b.configFilePath,
- b.manifestFilePath,
- }
- isUpdated, err := b.CheckDepsFileUpdate(ctx, b.tag, depsFiles)
- if err != nil {
- b.logger.Debugf("failed to check manifest update: %s", err)
- }
- if !isUpdated {
- b.logger.Infof("manifest is not updated, skip building")
- return nil
- }
- def, err := b.compile(ctx, pub)
- if err != nil {
- return errors.Wrap(err, "failed to compile")
- }
-
+func (b generalBuilder) Build(ctx context.Context) error {
  pw, err := progresswriter.NewPrinter(ctx, os.Stdout, b.progressMode)
  if err != nil {
  return errors.Wrap(err, "failed to create progress writer")
  }
 
- if err = b.build(ctx, def, pw); err != nil {
+ if err = b.build(ctx, pw); err != nil {
  return errors.Wrap(err, "failed to build")
  }
  return nil
@@ -192,11 +177,11 @@ func (b generalBuilder) interpret() error {
  return nil
 }
 
-func (b generalBuilder) compile(ctx context.Context, pub string) (*llb.Definition, error) {
+func (b generalBuilder) compile(ctx context.Context) (*llb.Definition, error) {
  if err := b.interpret(); err != nil {
  return nil, errors.Wrap(err, "failed to interpret")
  }
- def, err := ir.Compile(ctx, fileutil.Base(b.buildContextDir), pub)
+ def, err := ir.Compile(ctx, fileutil.Base(b.buildContextDir), b.pubKeyPath)
  if err != nil {
  return nil, errors.Wrap(err, "failed to compile build.envd")
  }
@@ -226,7 +211,7 @@ func (b generalBuilder) imageConfig(ctx context.Context) (string, error) {
  return data, nil
 }
 
-func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progresswriter.Writer) error {
+func (b generalBuilder) build(ctx context.Context, pw progresswriter.Writer) error {
  imageConfig, err := b.imageConfig(ctx)
  if err != nil {
  return errors.Wrap(err, "failed to get labels")
@@ -260,22 +245,22 @@ func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progr
  }
  }
  defer pipeW.Close()
- _, err := b.Solve(ctx, def, client.SolveOpt{
+ _, err := b.Client.Build(ctx, client.SolveOpt{
  Exports: []client.ExportEntry{entry},
  LocalDirs: map[string]string{
- flag.FlagContextDir: b.buildContextDir,
- flag.FlagCacheDir: home.GetManager().CacheDir(),
+ // TODO(gaocegege): Move it to BuildFunc with the help
+ // of llb.Local
+ flag.FlagCacheDir: home.GetManager().CacheDir(),
  },
  Session: attachable,
  // TODO(gaocegege): Use llb.WithProxy to implement it.
  FrontendAttrs: map[string]string{
  "build-arg:HTTPS_PROXY": os.Getenv("HTTPS_PROXY"),
  },
- }, pw.Status())
+ }, "envd", b.BuildFunc(), pw.Status())
 
  if err != nil {
  err = errors.Wrap(err, "failed to solve LLB")
- b.logger.Error(err)
  return err
  }
  b.logger.Debug("llb def is solved successfully")
@@ -299,22 +284,20 @@ func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progr
  })
  default:
  eg.Go(func() error {
- _, err := b.Solve(ctx, def, client.SolveOpt{
+ _, err := b.Client.Build(ctx, client.SolveOpt{
  Exports: []client.ExportEntry{entry},
  LocalDirs: map[string]string{
- flag.FlagContextDir: b.buildContextDir,
- flag.FlagCacheDir: home.GetManager().CacheDir(),
+ flag.FlagCacheDir: home.GetManager().CacheDir(),
  },
  Session: attachable,
  // TODO(gaocegege): Use llb.WithProxy to implement it.
  FrontendAttrs: map[string]string{
  "build-arg:HTTPS_PROXY": os.Getenv("HTTPS_PROXY"),
  },
- }, pw.Status())
+ }, "envd", b.BuildFunc(), pw.Status())
 
  if err != nil {
  err = errors.Wrap(err, "failed to solve LLB")
- b.logger.Error(err)
  return err
  }
  b.logger.Debug("llb def is solved successfully")

pkg/builder/builder_test.go

@@ -20,7 +20,7 @@ import (
  "os"
 
  "github.com/golang/mock/gomock"
- "github.com/moby/buildkit/client/llb"
+ "github.com/moby/buildkit/client"
  . "github.com/onsi/ginkgo/v2"
  . "github.com/onsi/gomega"
  "github.com/sirupsen/logrus"
@@ -50,12 +50,14 @@ var _ = Describe("Builder", func() {
  BeforeEach(func() {
  ctrl := gomock.NewController(GinkgoT())
  ctrlStarlark := gomock.NewController(GinkgoT())
+ pub := sshconfig.GetPublicKey()
  b = &generalBuilder{
  manifestFilePath: manifestFilePath,
  configFilePath: configFilePath,
- progressMode: "auto",
+ progressMode: "plain",
  tag: tag,
  buildfuncname: "build",
+ pubKeyPath: pub,
  logger: logrus.WithFields(logrus.Fields{
  "tag": tag,
  }),
@@ -68,50 +70,46 @@ var _ = Describe("Builder", func() {
  ir.DefaultGraph.Writer = w
  })
 
- When("failed to interpret config", func() {
+ When("build error", func() {
  It("should get an error", func() {
- expected := errors.New("failed to interpret config")
- b.Interpreter.(*mockstarlark.MockInterpreter).EXPECT().ExecFile(
- gomock.Eq(configFilePath), "",
- ).Return(nil, expected)
- pub := sshconfig.GetPublicKey()
- err := b.Build(context.TODO(), pub)
- Expect(err).To(HaveOccurred())
- })
- })
+ b.entries = []client.ExportEntry{
+ {
+ Type: client.ExporterDocker,
+ },
+ }
 
- When("failed to interpret manifest", func() {
- It("should get an error", func() {
- expected := errors.New("failed to interpret manifest")
- pub := sshconfig.GetPublicKey()
- b.Interpreter.(*mockstarlark.MockInterpreter).EXPECT().ExecFile(
- gomock.Eq(configFilePath), gomock.Eq(""),
- ).Return(nil, nil)
- b.Interpreter.(*mockstarlark.MockInterpreter).EXPECT().ExecFile(
- gomock.Eq(b.manifestFilePath), gomock.Eq("build"),
- ).Return(nil, expected)
- err := b.Build(context.TODO(), pub)
+ b.Client.(*mockbuildkitd.MockClient).EXPECT().Build(gomock.Any(),
+ gomock.Any(), gomock.Eq("envd"), gomock.Any(), gomock.Any()).
+ Return(nil, errors.New("build error"))
+
+ pw, err := progresswriter.NewPrinter(context.TODO(), os.Stdout, b.progressMode)
+ Expect(err).NotTo(HaveOccurred())
+
+ close(pw.Status())
+ err = b.build(context.TODO(), pw)
  Expect(err).To(HaveOccurred())
  })
  })
+
  It("should build successfully", func() {
  err := home.Initialize()
  Expect(err).ToNot(HaveOccurred())
 
- b.Client.(*mockbuildkitd.MockClient).EXPECT().Solve(
- gomock.Any(),
- gomock.Any(),
- gomock.Any(),
- gomock.Any(),
- ).Return(nil, nil).AnyTimes()
+ b.entries = []client.ExportEntry{
+ {
+ Type: client.ExporterDocker,
+ },
+ }
+
+ b.Client.(*mockbuildkitd.MockClient).EXPECT().Build(gomock.Any(),
+ gomock.Any(), gomock.Eq("envd"), gomock.Any(), gomock.Any()).
+ Return(nil, nil)
 
- var def *llb.Definition
  pw, err := progresswriter.NewPrinter(context.TODO(), os.Stdout, b.progressMode)
- if err != nil {
- Fail(err.Error())
- }
+ Expect(err).NotTo(HaveOccurred())
+
  close(pw.Status())
- err = b.build(context.TODO(), def, pw)
+ err = b.build(context.TODO(), pw)
  Expect(err).ToNot(HaveOccurred())
  })
  })

pkg/buildkitd/buildkitd.go

@@ -24,6 +24,7 @@ import (
  "github.com/cockroachdb/errors"
  "github.com/moby/buildkit/client"
  "github.com/moby/buildkit/client/llb"
+ gateway "github.com/moby/buildkit/frontend/gateway/client"
  "github.com/sirupsen/logrus"
  "github.com/spf13/viper"
  "github.com/tonistiigi/units"
@@ -46,6 +47,9 @@ type Client interface {
  // Solve calls Solve on the controller.
  Solve(ctx context.Context, def *llb.Definition,
  opt client.SolveOpt, statusChan chan *client.SolveStatus) (*client.SolveResponse, error)
+ Build(ctx context.Context, opt client.SolveOpt, product string,
+ buildFunc gateway.BuildFunc, statusChan chan *client.SolveStatus,
+ ) (*client.SolveResponse, error)
  Prune(ctx context.Context, keepDuration time.Duration,
  keepStorage float64, filter []string, verbose, all bool) error
  Close() error
@@ -144,7 +148,8 @@ func (c *generalClient) maybeStart(ctx context.Context,
  }
 
  if err := c.waitUntilConnected(ctx, connectingTimeout); err != nil {
- return "", errors.Wrap(err, "failed to connect to buildkitd")
+ return "", errors.Wrapf(err,
+ "failed to connect to buildkitd %s", c.BuildkitdAddr())
  }
 
  return c.BuildkitdAddr(), nil

pkg/flag/consts.go

@@ -16,7 +16,6 @@ package flag
 
 const (
  FlagCacheDir = "cache-dir"
- FlagContextDir = "context-dir"
  FlagBuildkitdImage = "buildkitd-image"
  FlagDebug = "debug"
 )