chore: show the compile stage progress similar to build stage (#79)

* feat: Print progress when compile

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

* chore: Update license

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

* fix: Update license

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

.gitignore

@@ -33,5 +33,6 @@ bin/
 debug-bin/
 /build.MIDI
 .ipynb_checkpoints/
+cover.html
 
 dist/

Makefile

@@ -117,6 +117,7 @@ build-local:
 generate: mockgen-install
  @mockgen -source pkg/buildkitd/buildkitd.go -destination pkg/buildkitd/mock/mock.go -package mock
  @mockgen -source pkg/lang/frontend/starlark/interpreter.go -destination pkg/lang/frontend/starlark/mock/mock.go -package mock
+ @mockgen -source pkg/progress/compileui/display.go -destination pkg/progress/compileui/mock/mock.go -package mock
 
 # It is used by vscode to attach into the process.
 debug-local:
@@ -128,7 +129,7 @@ debug-local:
  done
 
 addlicense:
- addlicense -c "The MIDI Authors" **/*.go **/**/*.go
+ addlicense -c "The MIDI Authors" **/*.go **/**/*.go **/**/**/*.go
 
 test: generate
  @go test -race -coverprofile=coverage.out ./...

build.MIDI

@@ -5,3 +5,4 @@ install_package(name = [
 cuda(version="11.6", cudnn="8")
 shell("zsh")
 jupyter(password="", port=8888)
+vscode(plugins = ["dbaeumer.vscode-eslint-2.2.3"])

go.mod

@@ -5,26 +5,33 @@ go 1.17
 require (
  github.com/alessio/shellescape v1.4.1
  github.com/cockroachdb/errors v1.9.0
+ github.com/containerd/console v1.0.3
  github.com/containerd/containerd v1.6.3-0.20220401172941-5ff8fce1fcc6
  github.com/creack/pty v1.1.18
  github.com/docker/docker v20.10.7+incompatible
- github.com/fatih/color v1.13.0
+ github.com/docker/go-connections v0.4.0
  github.com/gliderlabs/ssh v0.3.3
  github.com/go-git/go-git/v5 v5.4.2
  github.com/golang/mock v1.6.0
  github.com/google/uuid v1.3.0
  github.com/moby/buildkit v0.10.1
+ github.com/morikuni/aec v1.0.0
  github.com/onsi/ginkgo/v2 v2.1.4
  github.com/onsi/gomega v1.19.0
  github.com/opencontainers/go-digest v1.0.0
+ github.com/pkg/errors v0.9.1
  github.com/pkg/sftp v1.13.4
  github.com/sirupsen/logrus v1.8.1
  github.com/spf13/viper v1.4.0
+ github.com/stretchr/testify v1.7.0
+ github.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea
+ github.com/tonistiigi/vt100 v0.0.0-20210615222946-8066bb97264f
  github.com/urfave/cli/v2 v2.4.0
  go.starlark.net v0.0.0-20220328144851-d1966c6b9fcd
  golang.org/x/crypto v0.0.0-20211202192323-5770296d904e
  golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
  golang.org/x/term v0.0.0-20210927222741-03fcf44c2211
+ golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac
 )
 
 require (
@@ -38,13 +45,12 @@ require (
  github.com/chzyer/readline v0.0.0-20180603132655-2972be24d48e // indirect
  github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f // indirect
  github.com/cockroachdb/redact v1.1.3 // indirect
- github.com/containerd/console v1.0.3 // indirect
  github.com/containerd/continuity v0.2.3-0.20220330195504-d132b287edc8 // indirect
  github.com/containerd/typeurl v1.0.2 // indirect
  github.com/cpuguy83/go-md2man/v2 v2.0.1 // indirect
+ github.com/davecgh/go-spew v1.1.1 // indirect
  github.com/docker/cli v20.10.13+incompatible // indirect
  github.com/docker/distribution v2.8.0+incompatible // indirect
- github.com/docker/go-connections v0.4.0 // indirect
  github.com/docker/go-units v0.4.0 // indirect
  github.com/emirpasic/gods v1.12.0 // indirect
  github.com/fsnotify/fsnotify v1.4.9 // indirect
@@ -69,15 +75,12 @@ require (
  github.com/kr/pretty v0.3.0 // indirect
  github.com/kr/text v0.2.0 // indirect
  github.com/magiconair/properties v1.8.1 // indirect
- github.com/mattn/go-colorable v0.1.11 // indirect
- github.com/mattn/go-isatty v0.0.14 // indirect
  github.com/mitchellh/go-homedir v1.1.0 // indirect
  github.com/mitchellh/mapstructure v1.1.2 // indirect
  github.com/moby/sys/signal v0.6.0 // indirect
- github.com/morikuni/aec v1.0.0 // indirect
  github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799 // indirect
  github.com/pelletier/go-toml v1.9.4 // indirect
- github.com/pkg/errors v0.9.1 // indirect
+ github.com/pmezard/go-difflib v1.0.0 // indirect
  github.com/rogpeppe/go-internal v1.8.1 // indirect
  github.com/russross/blackfriday/v2 v2.1.0 // indirect
  github.com/sergi/go-diff v1.1.0 // indirect
@@ -86,8 +89,6 @@ require (
  github.com/spf13/jwalterweatherman v1.0.0 // indirect
  github.com/spf13/pflag v1.0.5 // indirect
  github.com/tonistiigi/fsutil v0.0.0-20220115021204-b19f7f9cb274 // indirect
- github.com/tonistiigi/units v0.0.0-20180711220420-6950e57a87ea // indirect
- github.com/tonistiigi/vt100 v0.0.0-20210615222946-8066bb97264f // indirect
  github.com/xanzy/ssh-agent v0.3.0 // indirect
  go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.29.0 // indirect
  go.opentelemetry.io/otel v1.4.1 // indirect
@@ -98,13 +99,13 @@ require (
  golang.org/x/net v0.0.0-20220225172249-27dd8689420f // indirect
  golang.org/x/sys v0.0.0-20220422013727-9388b58f7150 // indirect
  golang.org/x/text v0.3.7 // indirect
- golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac // indirect
  golang.org/x/xerrors v0.0.0-20220411194840-2f41105eb62f // indirect
  google.golang.org/genproto v0.0.0-20211208223120-3a66f561d7aa // indirect
  google.golang.org/grpc v1.45.0 // indirect
  google.golang.org/protobuf v1.27.1 // indirect
  gopkg.in/warnings.v0 v0.1.2 // indirect
  gopkg.in/yaml.v2 v2.4.0 // indirect
+ gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )
 
 // Copied from buildkit to make github.com/tonistiigi/fsutil happy.

go.sum

@@ -168,8 +168,6 @@ github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021/go.
 github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c=
 github.com/etcd-io/bbolt v1.3.3/go.mod h1:ZF2nL25h33cCyBtcyWeZ2/I3HQOfTP+0PIEvHjkjCrw=
 github.com/fasthttp-contrib/websocket v0.0.0-20160511215533-1f3b11f56072/go.mod h1:duJ4Jxv5lDcvg4QuQr0oowTf7dz4/CR8NtyCooz9HL8=
-github.com/fatih/color v1.13.0 h1:8LOYc1KYPPmyKMuN8QV2DNRWNbLo6LZ0iLs8+mlH53w=
-github.com/fatih/color v1.13.0/go.mod h1:kLAiJbzzSOZDVNGyDpeOxJ47H46qBXwg5ILebYFFOfk=
 github.com/fatih/structs v1.1.0/go.mod h1:9NiDSp5zOcgEDl+j00MP/WkGVPOlPRLejGD8Ga6PJ7M=
 github.com/flosch/pongo2 v0.0.0-20190707114632-bbf5a6c351f4/go.mod h1:T9YF2M40nIgbVgp3rreNmTged+9HrbNTIQf1PsaIiTA=
 github.com/flynn/go-shlex v0.0.0-20150515145356-3f9db97f8568/go.mod h1:xEzjJPgXI435gkrCt3MPfRiAkVrwSbHsst4LCFVfpJc=
@@ -389,14 +387,11 @@ github.com/matryer/is v1.2.0 h1:92UTHpy8CDwaJ08GqLDzhhuixiBUUD1p3AU6PHddz4A=
 github.com/matryer/is v1.2.0/go.mod h1:2fLPjFQM9rhQ15aVEtbuwhJinnOqrmgXPNdZsdwlWXA=
 github.com/mattn/go-colorable v0.1.2/go.mod h1:U0ppj6V5qS13XJ6of8GYAs25YV2eR4EVcfRqFIhoBtE=
 github.com/mattn/go-colorable v0.1.8/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
-github.com/mattn/go-colorable v0.1.9/go.mod h1:u6P/XSegPjTcexA+o6vUJrdnUu04hMope9wVRipJSqc=
-github.com/mattn/go-colorable v0.1.11 h1:nQ+aFkoE2TMGc0b68U2OKSexC+eq46+XwZzWXHRmPYs=
 github.com/mattn/go-colorable v0.1.11/go.mod h1:u5H1YNBxpqRaxsYJYSkiCWKzEfiAb1Gb520KVy5xxl4=
 github.com/mattn/go-isatty v0.0.7/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
 github.com/mattn/go-isatty v0.0.8/go.mod h1:Iq45c/XA43vh69/j3iqttzPXn0bhXyGjM0Hdxcsrc5s=
 github.com/mattn/go-isatty v0.0.9/go.mod h1:YNRxwqDuOph6SZLI9vUUz6OYw3QyUt7WiY2yME+cCiQ=
 github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Kysco4FUpU=
-github.com/mattn/go-isatty v0.0.14 h1:yVuAays6BHfxijgZPzw+3Zlu5yQgKGP2/hcQbHb7S9Y=
 github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27kJ6hsGG94=
 github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=

pkg/builder/builder.go

@@ -21,7 +21,7 @@ import (
 
  "github.com/cockroachdb/errors"
  "github.com/moby/buildkit/client"
- "github.com/moby/buildkit/util/progress/progresswriter"
+ "github.com/moby/buildkit/client/llb"
  "github.com/sirupsen/logrus"
  "golang.org/x/sync/errgroup"
 
@@ -31,6 +31,7 @@ import (
  "github.com/tensorchord/MIDI/pkg/home"
  "github.com/tensorchord/MIDI/pkg/lang/frontend/starlark"
  "github.com/tensorchord/MIDI/pkg/lang/ir"
+ "github.com/tensorchord/MIDI/pkg/progress/progresswriter"
 )
 
 type Builder interface {
@@ -47,7 +48,6 @@ type generalBuilder struct {
  logger *logrus.Entry
  starlark.Interpreter
  buildkitd.Client
- progresswriter.Writer
 }
 
 func New(ctx context.Context, configFilePath, manifestFilePath, tag string) (Builder, error) {
@@ -68,12 +68,6 @@ func New(ctx context.Context, configFilePath, manifestFilePath, tag string) (Bui
  }
  b.Client = cli
 
- pw, err := progresswriter.NewPrinter(ctx, os.Stdout, b.progressMode)
- if err != nil {
- return nil, errors.Wrap(err, "failed to create progress writer")
- }
- b.Writer = pw
-
  b.Interpreter = starlark.NewInterpreter()
  return b, nil
 }
@@ -84,20 +78,47 @@ func (b generalBuilder) GPUEnabled() bool {
 }
 
 func (b generalBuilder) Build(ctx context.Context) error {
+ def, err := b.compile(ctx)
+ if err != nil {
+ return errors.Wrap(err, "failed to compile")
+ }
+
+ 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 {
+ return errors.Wrap(err, "failed to build")
+ }
+ return nil
+}
+
+func (b generalBuilder) interpret() error {
  // Evaluate config first.
  if _, err := b.ExecFile(b.configFilePath); err != nil {
- return err
+ return errors.Wrap(err, "failed to exec starlark file")
  }
 
  if _, err := b.ExecFile(b.manifestFilePath); err != nil {
- return err
+ return errors.Wrap(err, "failed to exec starlark file")
  }
+ return nil
+}
 
+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)
  if err != nil {
- return errors.Wrap(err, "failed to compile build.MIDI")
+ return nil, errors.Wrap(err, "failed to compile build.MIDI")
  }
+ b.logger.Debug("compiled build.MIDI")
+ return def, nil
+}
 
+func (b generalBuilder) build(ctx context.Context, def *llb.Definition, pw progresswriter.Writer) error {
  ctx, cancel := context.WithCancel(ctx)
  defer cancel()
  eg, ctx := errgroup.WithContext(ctx)
@@ -130,7 +151,7 @@ func (b generalBuilder) Build(ctx context.Context) error {
  FrontendAttrs: map[string]string{
  "build-arg:HTTPS_PROXY": os.Getenv("HTTPS_PROXY"),
  },
- }, b.Status())
+ }, pw.Status())
  if err != nil {
  err = errors.Wrap(err, "failed to solve LLB")
  b.logger.Error(err)
@@ -143,8 +164,8 @@ func (b generalBuilder) Build(ctx context.Context) error {
  // Watch the progress.
  eg.Go(func() error {
  // not using shared context to not disrupt display but let is finish reporting errors
- <-b.Done()
- return b.Err()
+ <-pw.Done()
+ return pw.Err()
  })
 
  // Load the image to docker host.
@@ -164,7 +185,7 @@ func (b generalBuilder) Build(ctx context.Context) error {
  return nil
  })
 
- err = eg.Wait()
+ err := eg.Wait()
  if err != nil {
  if errors.Is(err, context.Canceled) {
  b.logger.Debug("cancelling the error group")
@@ -175,6 +196,5 @@ func (b generalBuilder) Build(ctx context.Context) error {
  return errors.Wrap(err, "failed to wait error group")
  }
  }
-
  return nil
 }

pkg/builder/builder_test.go

@@ -20,7 +20,7 @@ import (
  "os"
 
  "github.com/golang/mock/gomock"
- "github.com/moby/buildkit/util/progress/progresswriter"
+ "github.com/moby/buildkit/client/llb"
  . "github.com/onsi/ginkgo/v2"
  . "github.com/onsi/gomega"
  "github.com/sirupsen/logrus"
@@ -30,6 +30,10 @@ import (
  "github.com/tensorchord/MIDI/pkg/flag"
  "github.com/tensorchord/MIDI/pkg/home"
  mockstarlark "github.com/tensorchord/MIDI/pkg/lang/frontend/starlark/mock"
+ "github.com/tensorchord/MIDI/pkg/lang/ir"
+ "github.com/tensorchord/MIDI/pkg/progress/compileui"
+ compileuimock "github.com/tensorchord/MIDI/pkg/progress/compileui/mock"
+ "github.com/tensorchord/MIDI/pkg/progress/progresswriter"
 )
 
 var _ = Describe("Builder", func() {
@@ -58,6 +62,7 @@ var _ = Describe("Builder", func() {
  })
  When("building the manifest", func() {
  var b *generalBuilder
+ var w compileui.Writer
  BeforeEach(func() {
  ctrl := gomock.NewController(GinkgoT())
  ctrlStarlark := gomock.NewController(GinkgoT())
@@ -72,11 +77,10 @@ var _ = Describe("Builder", func() {
  }
  b.Client = mockbuildkitd.NewMockClient(ctrl)
  b.Interpreter = mockstarlark.NewMockInterpreter(ctrlStarlark)
- pw, err := progresswriter.NewPrinter(context.TODO(), os.Stdout, b.progressMode)
- if err != nil {
- Fail(err.Error())
- }
- b.Writer = pw
+
+ ctrlWriter := gomock.NewController(GinkgoT())
+ w = compileuimock.NewMockWriter(ctrlWriter)
+ ir.DefaultGraph.Writer = w
  })
 
  When("failed to interpret config", func() {
@@ -104,23 +108,23 @@ var _ = Describe("Builder", func() {
  })
  })
  It("should build successfully", func() {
- b.Interpreter.(*mockstarlark.MockInterpreter).EXPECT().ExecFile(
- gomock.Eq(configFilePath),
- ).Return(nil, nil).Times(1)
- b.Interpreter.(*mockstarlark.MockInterpreter).EXPECT().ExecFile(
- gomock.Eq(b.manifestFilePath),
- ).Return(nil, nil).Times(1)
  err := home.Initialize()
  Expect(err).ToNot(HaveOccurred())
- close(b.Writer.Status())
 
  b.Client.(*mockbuildkitd.MockClient).EXPECT().Solve(
  gomock.Any(),
  gomock.Any(),
  gomock.Any(),
  gomock.Any(),
  ).Return(nil, nil).AnyTimes()
- err = b.Build(context.TODO())
+
+ var def *llb.Definition
+ pw, err := progresswriter.NewPrinter(context.TODO(), os.Stdout, b.progressMode)
+ if err != nil {
+ Fail(err.Error())
+ }
+ close(pw.Status())
+ err = b.build(context.TODO(), def, pw)
  Expect(err).ToNot(HaveOccurred())
  })
  })

pkg/lang/ir/consts.go

@@ -1,3 +1,17 @@
+// Copyright 2022 The MIDI 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 ir
 
 const (

pkg/lang/ir/graph.go

@@ -17,6 +17,7 @@ package ir
 import (
  "context"
  "fmt"
+ "os"
  "path/filepath"
  "strings"
 
@@ -27,6 +28,7 @@ import (
 
  "github.com/tensorchord/MIDI/pkg/editor/vscode"
  "github.com/tensorchord/MIDI/pkg/flag"
+ "github.com/tensorchord/MIDI/pkg/progress/compileui"
  "github.com/tensorchord/MIDI/pkg/shell"
 )
 
@@ -56,6 +58,11 @@ func GPUEnabled() bool {
 }
 
 func Compile(ctx context.Context) (*llb.Definition, error) {
+ w, err := compileui.New(ctx, os.Stdout, "auto")
+ if err != nil {
+ return nil, errors.Wrap(err, "failed to create compileui")
+ }
+ DefaultGraph.Writer = w
  state, err := DefaultGraph.Compile()
  if err != nil {
  return nil, err
@@ -105,6 +112,7 @@ func (g Graph) Compile() (llb.State, error) {
 
  // TODO(gaocegege): Support order-based exec.
  run := g.compileRun(merged)
+ g.Writer.Finish()
  return run, nil
 }