feat(debug): Add llb export (#1124)

* feat(debug): Add llb dup

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

* fix: Update

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

* fix: Add mount in dot plot

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

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

pkg/app/app.go

@@ -87,6 +87,7 @@ func New() EnvdApp {
 		CommandRun,
 		CommandResume,
 		CommandUp,
+		CommandDebug,
 		CommandVersion,
 		CommandTop,
 	}

pkg/app/debug.go

@@ -0,0 +1,31 @@
+// 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 app
+
+import (
+	"github.com/urfave/cli/v2"
+)
+
+var CommandDebug = &cli.Command{
+	Name:        "debug",
+	Category:    CategoryOther,
+	Aliases:     []string{"b"},
+	Usage:       "Debug commands",
+	Description: ``,
+
+	Subcommands: []*cli.Command{
+		CommandDebugLLB,
+	},
+}

pkg/app/debug_llb.go

@@ -0,0 +1,221 @@
+// 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 app
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"strings"
+
+	"github.com/cockroachdb/errors"
+	"github.com/moby/buildkit/client/llb"
+	"github.com/moby/buildkit/solver/pb"
+	"github.com/opencontainers/go-digest"
+	"github.com/sirupsen/logrus"
+	"github.com/urfave/cli/v2"
+
+	"github.com/tensorchord/envd/pkg/app/telemetry"
+	sshconfig "github.com/tensorchord/envd/pkg/ssh/config"
+)
+
+var CommandDebugLLB = &cli.Command{
+	Name:        "llb",
+	Category:    CategoryOther,
+	Aliases:     []string{"b"},
+	Usage:       "dump buildkit LLB in human-readable format.",
+	Description: ``,
+
+	Flags: []cli.Flag{
+		&cli.BoolFlag{
+			Name:  "dot",
+			Usage: "Output dot format",
+		},
+		&cli.PathFlag{
+			Name:    "path",
+			Usage:   "Path to the directory containing the build.envd",
+			Aliases: []string{"p"},
+			Value:   ".",
+		},
+		&cli.PathFlag{
+			Name:    "from",
+			Usage:   "Function to execute, format `file:func`",
+			Aliases: []string{"f"},
+			Value:   "build.envd:build",
+		},
+		&cli.PathFlag{
+			Name:    "public-key",
+			Usage:   "Path to the public key",
+			Aliases: []string{"pubk"},
+			Value:   sshconfig.GetPublicKeyOrPanic(),
+			Hidden:  true,
+		},
+	},
+
+	Action: debugLLB,
+}
+
+func debugLLB(clicontext *cli.Context) error {
+	telemetry.GetReporter().Telemetry("debug-llb", nil)
+	opt, err := ParseBuildOpt(clicontext)
+	if err != nil {
+		return err
+	}
+
+	logger := logrus.WithFields(logrus.Fields{
+		"build-context": opt.BuildContextDir,
+		"build-file":    opt.ManifestFilePath,
+		"config":        opt.ConfigFilePath,
+		"tag":           opt.Tag,
+	})
+	logger.WithFields(logrus.Fields{
+		"builder-options": opt,
+	}).Debug("starting debug llb command")
+
+	builder, err := GetBuilder(clicontext, opt)
+	if err != nil {
+		return err
+	}
+	if err = InterpretEnvdDef(builder); err != nil {
+		return err
+	}
+
+	def, err := builder.Compile(clicontext.Context)
+	if err != nil {
+		return errors.Wrap(err, "failed to compile envd IR to LLB")
+	}
+
+	ops, err := loadLLB(def)
+	if err != nil {
+		return errors.Wrap(err, "failed to load LLB")
+	}
+
+	if clicontext.Bool("dot") {
+		writeDot(ops, os.Stdout)
+	} else {
+		enc := json.NewEncoder(os.Stdout)
+		for _, op := range ops {
+			if err := enc.Encode(op); err != nil {
+				return errors.Wrap(err, "failed to encode LLB op")
+			}
+		}
+	}
+	return nil
+}
+
+type llbOp struct {
+	Op         pb.Op
+	Digest     digest.Digest
+	OpMetadata pb.OpMetadata
+}
+
+// Refer to https://github.com/moby/buildkit/blob/master/cmd/buildctl/debug/dumpllb.go#L17:5
+func loadLLB(def *llb.Definition) ([]llbOp, error) {
+	var ops []llbOp
+	for _, dt := range def.Def {
+		var op pb.Op
+		if err := (&op).Unmarshal(dt); err != nil {
+			return nil, errors.Wrap(err, "failed to parse op")
+		}
+		dgst := digest.FromBytes(dt)
+		ent := llbOp{Op: op, Digest: dgst, OpMetadata: def.Metadata[dgst]}
+		ops = append(ops, ent)
+	}
+	return ops, nil
+}
+
+func writeDot(ops []llbOp, w io.Writer) {
+	// TODO: print OpMetadata
+	fmt.Fprintln(w, "digraph {")
+	defer fmt.Fprintln(w, "}")
+	for _, op := range ops {
+		name, shape := attr(op.Digest, op.Op)
+		fmt.Fprintf(w, "  %q [label=%q shape=%q];\n", op.Digest, name, shape)
+	}
+	for _, op := range ops {
+		for i, inp := range op.Op.Inputs {
+			label := ""
+			if eo, ok := op.Op.Op.(*pb.Op_Exec); ok {
+				for _, m := range eo.Exec.Mounts {
+					if int(m.Input) == i && m.Dest != "/" {
+						label = m.Dest
+					}
+				}
+			}
+			fmt.Fprintf(w, "  %q -> %q [label=%q];\n", inp.Digest, op.Digest, label)
+		}
+	}
+}
+
+func attr(dgst digest.Digest, op pb.Op) (string, string) {
+	switch op := op.Op.(type) {
+	case *pb.Op_Source:
+		return op.Source.Identifier, "ellipse"
+	case *pb.Op_Exec:
+		return generateExecNode(op.Exec)
+	case *pb.Op_Build:
+		return "build", "box3d"
+	case *pb.Op_Merge:
+		return "merge", "invtriangle"
+	case *pb.Op_Diff:
+		return "diff", "doublecircle"
+	case *pb.Op_File:
+		names := []string{}
+
+		for _, action := range op.File.Actions {
+			var name string
+
+			switch act := action.Action.(type) {
+			case *pb.FileAction_Copy:
+				name = fmt.Sprintf("copy{src=%s, dest=%s}", act.Copy.Src, act.Copy.Dest)
+			case *pb.FileAction_Mkfile:
+				name = fmt.Sprintf("mkfile{path=%s}", act.Mkfile.Path)
+			case *pb.FileAction_Mkdir:
+				name = fmt.Sprintf("mkdir{path=%s}", act.Mkdir.Path)
+			case *pb.FileAction_Rm:
+				name = fmt.Sprintf("rm{path=%s}", act.Rm.Path)
+			}
+
+			names = append(names, name)
+		}
+		return strings.Join(names, ","), "note"
+	default:
+		return dgst.String(), "plaintext"
+	}
+}
+
+func generateExecNode(op *pb.ExecOp) (string, string) {
+	mounts := []string{}
+	for _, m := range op.Mounts {
+		mstr := fmt.Sprintf("selector=%s, target=%s, mount-type=%s", m.Selector,
+			m.Dest, m.MountType)
+		if m.CacheOpt != nil {
+			mstr = mstr + fmt.Sprintf(" cache-id=%s, cache-share-mode = %s",
+				m.CacheOpt.ID, m.CacheOpt.Sharing)
+		}
+		mounts = append(mounts, mstr)
+	}
+
+	name := fmt.Sprintf("user=%s, cwd=%s, args={%s}, mounts={%s}, env={%s}",
+		op.Meta.User,
+		op.Meta.Cwd,
+		strings.Join(op.Meta.Args, " "),
+		strings.Join(mounts, " "),
+		strings.Join(op.Meta.Env, " "),
+	)
+
+	return name, "box"
+}

pkg/app/telemetry/reporter.go

@@ -49,9 +49,16 @@ var (
 
 func Initialize(enabled bool, token string) error {
 	once.Do(func() {
+		// Ref https://segment.com/docs/connections/sources/catalog/libraries/server/go/#development-settings
+		c, err := segmentio.NewWithConfig(token, segmentio.Config{
+			BatchSize: 1,
+		})
+		if err != nil {
+			panic(err)
+		}
 		reporter = &defaultReporter{
 			enabled: enabled,
-			client:  segmentio.New(token),
+			client:  c,
 		}
 	})
 	return reporter.init()
@@ -104,6 +111,23 @@ func (r *defaultReporter) init() error {
 	file.Close()
 	r.UID = string(uid)
 
+	r.Identify()
+	return nil
+}
+
+func (r *defaultReporter) dumpTelemetry() error {
+	file, err := os.Create(r.telemetryFile)
+	if err != nil {
+		return errors.Wrap(err, "failed to create cache telemetry file")
+	}
+	defer file.Close()
+
+	// Write uid to file.
+	_, err = file.Write([]byte(r.UID))
+	return err
+}
+
+func (r *defaultReporter) Identify() {
 	logrus.WithField("UID", r.UID).Debug("telemetry initialization")
 	if r.enabled {
 		logrus.Debug("sending telemetry")
@@ -124,22 +148,9 @@ func (r *defaultReporter) init() error {
 			Traits:    segmentio.NewTraits(),
 		}); err != nil {
 			logrus.Warn("telemetry failed")
-			return nil
+			return
 		}
 	}
-	return nil
-}
-
-func (r *defaultReporter) dumpTelemetry() error {
-	file, err := os.Create(r.telemetryFile)
-	if err != nil {
-		return errors.Wrap(err, "failed to create cache telemetry file")
-	}
-	defer file.Close()
-
-	// Write uid to file.
-	_, err = file.Write([]byte(r.UID))
-	return err
 }
 
 func (r *defaultReporter) Telemetry(command string, runner *string) {

pkg/builder/builder.go

@@ -42,6 +42,8 @@ import (
 type Builder interface {
 	Build(ctx context.Context, force bool) error
 	Interpret() error
+	// Compile compiles envd IR to LLB.
+	Compile(ctx context.Context) (*llb.Definition, error)
 	GPUEnabled() bool
 	NumGPUs() int
 }
@@ -155,7 +157,7 @@ func (b generalBuilder) Build(ctx context.Context, force bool) error {
 		return nil
 	}
 
-	def, err := b.compile(ctx)
+	def, err := b.Compile(ctx)
 	if err != nil {
 		return errors.Wrap(err, "failed to compile")
 	}
@@ -187,7 +189,7 @@ func (b generalBuilder) Interpret() error {
 	return nil
 }
 
-func (b generalBuilder) compile(ctx context.Context) (*llb.Definition, error) {
+func (b generalBuilder) Compile(ctx context.Context) (*llb.Definition, error) {
 	envName := filepath.Base(b.BuildContextDir)
 	def, err := ir.Compile(ctx, envName, b.PubKeyPath)
 	if err != nil {