Repository: tensorchord/envd (language: python)

Issue:
feat(buildkit): Export cache to registry
## Description

The build cache can be exported to the registry. We need to investigate the benefits of AI/ML scenarios.


Discussion / hints:
If we implement this, we can remove conda in base images. The dependencies can be installed during build process, and we can cache them in the registry.
/cc @kemingy 
Ref 

- https://github.com/moby/buildkit/issues/2383
- https://github.com/moby/buildkit/issues/1251

Tried to push the cache to the registry and import it in a clean build, it does not work. Not sure what happens:

```go
eg.Go(func() error {
				_, err := b.Client.Build(ctx, client.SolveOpt{
					CacheExports: []client.CacheOptionsEntry{
						{
							Type: "inline",
							Attrs: map[string]string{
								"mode": "max",
							},
						},
					},
					CacheImports: []client.CacheOptionsEntry{
						{
							Type: "registry",
							Attrs: map[string]string{
								"ref": b.tag,
							},
						},
					},
					Exports: []client.ExportEntry{entry},
					LocalDirs: map[string]string{
						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"),
					},
				}, "envd", b.BuildFunc(), pw.Status())

				if err != nil {
					err = errors.Wrap(err, "failed to solve LLB")
					return err
				}
				b.logger.Debug("llb def is solved successfully")
				return nil
			})
```
https://github.com/moby/buildkit/issues/2383#issuecomment-930627760

Maybe it is related to the umask. We need to re-compile buildkit to know what happens here.
I simplified the LLB to:

```
func (g Graph) compilePython(aptStage llb.State) (llb.State, error) {
	systemStage := g.compileSystemPackages(aptStage)

	return systemStage, nil
}
```

There are only two layers: llb.Image and llb.Run. The remote still does not work. Then I added some debug info into buildit: 

- Print cachemap of run state.

The cachemap is always the same.

I need to have a deep dive into the implementation about the remote cache.

Fix patch:
diff --git a/pkg/app/build.go b/pkg/app/build.go
index 842a09906..e9cac1d39 100644
--- a/pkg/app/build.go
+++ b/pkg/app/build.go
@@ -77,8 +77,18 @@ To build and push the image to a registry:
 			Usage: "Force rebuild the image",
 			Value: false,
 		},
+		// https://github.com/urfave/cli/issues/1134#issuecomment-1191407527
+		&cli.StringFlag{
+			Name:    "export-cache",
+			Usage:   "Export the cache (e.g. type=registry,ref=<image>)",
+			Aliases: []string{"ec"},
+		},
+		&cli.StringFlag{
+			Name:    "import-cache",
+			Usage:   "Import the cache (e.g. type=registry,ref=<image>)",
+			Aliases: []string{"ic"},
+		},
 	},
-
 	Action: build,
 }
 
@@ -122,6 +132,8 @@ func build(clicontext *cli.Context) error {
 	debug := clicontext.Bool("debug")
 	output := clicontext.String("output")
 	force := clicontext.Bool("force")
+	exportCache := clicontext.String("export-cache")
+	importCache := clicontext.String("import-cache")
 
 	opt := builder.Options{
 		ManifestFilePath: manifest,
@@ -132,6 +144,8 @@ func build(clicontext *cli.Context) error {
 		OutputOpts:       output,
 		PubKeyPath:       clicontext.Path("public-key"),
 		ProgressMode:     "auto",
+		ExportCache:      exportCache,
+		ImportCache:      importCache,
 	}
 	if debug {
 		opt.ProgressMode = "plain"
diff --git a/pkg/app/up.go b/pkg/app/up.go
index 1df6eb79b..a1df97817 100644
--- a/pkg/app/up.go
+++ b/pkg/app/up.go
@@ -99,6 +99,17 @@ var CommandUp = &cli.Command{
 			Usage: "Force rebuild and run the container although the previous container is running",
 			Value: false,
 		},
+		// https://github.com/urfave/cli/issues/1134#issuecomment-1191407527
+		&cli.StringFlag{
+			Name:    "export-cache",
+			Usage:   "Export the cache (e.g. type=registry,ref=<image>)",
+			Aliases: []string{"ec"},
+		},
+		&cli.StringFlag{
+			Name:    "import-cache",
+			Usage:   "Import the cache (e.g. type=registry,ref=<image>)",
+			Aliases: []string{"ic"},
+		},
 	},
 
 	Action: up,
@@ -140,6 +151,8 @@ func up(clicontext *cli.Context) error {
 	debug := clicontext.Bool("debug")
 	force := clicontext.Bool("force")
 	output := ""
+	exportCache := clicontext.String("export-cache")
+	importCache := clicontext.String("import-cache")
 
 	opt := builder.Options{
 		ManifestFilePath: manifest,
@@ -150,6 +163,8 @@ func up(clicontext *cli.Context) error {
 		OutputOpts:       output,
 		PubKeyPath:       clicontext.Path("public-key"),
 		ProgressMode:     "auto",
+		ExportCache:      exportCache,
+		ImportCache:      importCache,
 	}
 	if debug {
 		opt.ProgressMode = "plain"
diff --git a/pkg/builder/build.go b/pkg/builder/build.go
index fe32de509..6e87f6efd 100644
--- a/pkg/builder/build.go
+++ b/pkg/builder/build.go
@@ -28,9 +28,17 @@ func (b generalBuilder) BuildFunc() func(ctx context.Context, c client.Client) (
 			return nil, errors.Wrap(err, "failed to compile")
 		}
 
-		res, err := c.Solve(ctx, client.SolveRequest{
+		sreq := client.SolveRequest{
 			Definition: def.ToPB(),
-		})
+		}
+		if b.Options.ImportCache != "" {
+			ci, err := ParseImportCache([]string{b.Options.ImportCache})
+			if err != nil {
+				return nil, errors.Wrap(err, "failed to get the import cache")
+			}
+			sreq.CacheImports = ci
+		}
+		res, err := c.Solve(ctx, sreq)
 		if err != nil {
 			return nil, err
 		}
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
index 6633b949b..239c19e30 100644
--- a/pkg/builder/builder.go
+++ b/pkg/builder/builder.go
@@ -61,6 +61,12 @@ type Options struct {
 	PubKeyPath string
 	// OutputOpts is the output options.
 	OutputOpts string
+	// ExportCache is the option to export cache.
+	// e.g. type=registry,ref=docker.io/username/image
+	ExportCache string
+	// ImportCache is the option to import cache.
+	// e.g. type=registry,ref=docker.io/username/image
+	ImportCache string
 }
 
 type generalBuilder struct {
@@ -188,6 +194,11 @@ func (b generalBuilder) build(ctx context.Context, pw progresswriter.Writer) err
 	if err != nil {
 		return errors.Wrap(err, "failed to get labels")
 	}
+
+	ce, err := ParseExportCache([]string{b.ExportCache}, nil)
+	if err != nil {
+		return errors.Wrap(err, "failed to parse export cache")
+	}
 	// k := platforms.Format(platforms.DefaultSpec())
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
@@ -199,6 +210,9 @@ func (b generalBuilder) build(ctx context.Context, pw progresswriter.Writer) err
 	for _, entry := range b.entries {
 		// Set up docker config auth.
 		attachable := []session.Attachable{authprovider.NewDockerAuthProvider(os.Stderr)}
+		b.logger.WithFields(logrus.Fields{
+			"type": entry.Type,
+		}).Debug("build image with buildkit")
 		switch entry.Type {
 		// Create default build.
 		case client.ExporterDocker:
@@ -217,20 +231,15 @@ func (b generalBuilder) build(ctx context.Context, pw progresswriter.Writer) err
 					}
 				}
 				defer pipeW.Close()
-				_, err := b.Client.Build(ctx, client.SolveOpt{
-					Exports: []client.ExportEntry{entry},
+				solveOpt := client.SolveOpt{
+					CacheExports: ce,
+					Exports:      []client.ExportEntry{entry},
 					LocalDirs: map[string]string{
-						// 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"),
-					},
-				}, "envd", b.BuildFunc(), pw.Status())
-
+				}
+				_, err := b.Client.Build(ctx, solveOpt, "envd", b.BuildFunc(), pw.Status())
 				if err != nil {
 					err = errors.Wrap(err, "failed to solve LLB")
 					return err
@@ -256,18 +265,15 @@ func (b generalBuilder) build(ctx context.Context, pw progresswriter.Writer) err
 			})
 		default:
 			eg.Go(func() error {
-				_, err := b.Client.Build(ctx, client.SolveOpt{
-					Exports: []client.ExportEntry{entry},
+				solveOpt := client.SolveOpt{
+					CacheExports: ce,
+					Exports:      []client.ExportEntry{entry},
 					LocalDirs: map[string]string{
 						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"),
-					},
-				}, "envd", b.BuildFunc(), pw.Status())
-
+				}
+				_, err := b.Client.Build(ctx, solveOpt, "envd", b.BuildFunc(), pw.Status())
 				if err != nil {
 					err = errors.Wrap(err, "failed to solve LLB")
 					return err
diff --git a/pkg/builder/util.go b/pkg/builder/util.go
index 8e14320f4..3d13ab513 100644
--- a/pkg/builder/util.go
+++ b/pkg/builder/util.go
@@ -25,6 +25,7 @@ import (
 	"github.com/containerd/console"
 	"github.com/containerd/containerd/platforms"
 	"github.com/moby/buildkit/client"
+	gatewayclient "github.com/moby/buildkit/frontend/gateway/client"
 	v1 "github.com/opencontainers/image-spec/specs-go/v1"
 	"github.com/sirupsen/logrus"
 )
@@ -76,6 +77,142 @@ func DefaultPathEnv(os string) string {
 	return DefaultPathEnvUnix
 }
 
+func parseImportCacheCSV(s string) (gatewayclient.CacheOptionsEntry, error) {
+	im := gatewayclient.CacheOptionsEntry{
+		Type:  "",
+		Attrs: map[string]string{},
+	}
+	csvReader := csv.NewReader(strings.NewReader(s))
+	fields, err := csvReader.Read()
+	if err != nil {
+		return im, err
+	}
+	for _, field := range fields {
+		parts := strings.SplitN(field, "=", 2)
+		if len(parts) != 2 {
+			return im, errors.Errorf("invalid value %s", field)
+		}
+		key := strings.ToLower(parts[0])
+		value := parts[1]
+		switch key {
+		case "type":
+			im.Type = value
+		default:
+			im.Attrs[key] = value
+		}
+	}
+	if im.Type == "" {
+		return im, errors.New("--import-cache requires type=<type>")
+	}
+	return im, nil
+}
+
+// ParseImportCache parses --import-cache
+func ParseImportCache(importCaches []string) ([]gatewayclient.CacheOptionsEntry, error) {
+	var imports []gatewayclient.CacheOptionsEntry
+	for _, importCache := range importCaches {
+		legacy := !strings.Contains(importCache, "type=")
+		if legacy {
+			logrus.Warn("--import-cache <ref> is deprecated. Please use --import-cache type=registry,ref=<ref>,<opt>=<optval>[,<opt>=<optval>] instead.")
+			imports = append(imports, gatewayclient.CacheOptionsEntry{
+				Type:  "registry",
+				Attrs: map[string]string{"ref": importCache},
+			})
+		} else {
+			im, err := parseImportCacheCSV(importCache)
+			if err != nil {
+				return nil, err
+			}
+			imports = append(imports, im)
+		}
+	}
+	return imports, nil
+}
+
+// ParseExportCache parses --export-cache (and legacy --export-cache-opt)
+// Refer to github.com/moby/buildkit/cmd/buildctl/build/exportcache.go
+func ParseExportCache(exportCaches, legacyExportCacheOpts []string) ([]client.CacheOptionsEntry, error) {
+	var exports []client.CacheOptionsEntry
+	if len(legacyExportCacheOpts) > 0 {
+		if len(exportCaches) != 1 {
+			return nil, errors.New("--export-cache-opt requires exactly single --export-cache")
+		}
+	}
+	for _, exportCache := range exportCaches {
+		legacy := !strings.Contains(exportCache, "type=")
+		if legacy {
+			logrus.Warnf("--export-cache <ref> --export-cache-opt <opt>=<optval> is deprecated. Please use --export-cache type=registry,ref=<ref>,<opt>=<optval>[,<opt>=<optval>] instead")
+			attrs, err := attrMap(legacyExportCacheOpts)
+			if err != nil {
+				return nil, err
+			}
+			if _, ok := attrs["mode"]; !ok {
+				attrs["mode"] = "min"
+			}
+			attrs["ref"] = exportCache
+			exports = append(exports, client.CacheOptionsEntry{
+				Type:  "registry",
+				Attrs: attrs,
+			})
+		} else {
+			if len(legacyExportCacheOpts) > 0 {
+				return nil, errors.New("--export-cache-opt is not supported for the specified --export-cache. Please use --export-cache type=<type>,<opt>=<optval>[,<opt>=<optval>] instead")
+			}
+			ex, err := parseExportCacheCSV(exportCache)
+			if err != nil {
+				return nil, err
+			}
+			exports = append(exports, ex)
+		}
+	}
+	return exports, nil
+}
+
+func attrMap(sl []string) (map[string]string, error) {
+	m := map[string]string{}
+	for _, v := range sl {
+		parts := strings.SplitN(v, "=", 2)
+		if len(parts) != 2 {
+			return nil, errors.Errorf("invalid value %s", v)
+		}
+		m[parts[0]] = parts[1]
+	}
+	return m, nil
+}
+
+func parseExportCacheCSV(s string) (client.CacheOptionsEntry, error) {
+	ex := client.CacheOptionsEntry{
+		Type:  "",
+		Attrs: map[string]string{},
+	}
+	csvReader := csv.NewReader(strings.NewReader(s))
+	fields, err := csvReader.Read()
+	if err != nil {
+		return ex, err
+	}
+	for _, field := range fields {
+		parts := strings.SplitN(field, "=", 2)
+		if len(parts) != 2 {
+			return ex, errors.Errorf("invalid value %s", field)
+		}
+		key := strings.ToLower(parts[0])
+		value := parts[1]
+		switch key {
+		case "type":
+			ex.Type = value
+		default:
+			ex.Attrs[key] = value
+		}
+	}
+	if ex.Type == "" {
+		return ex, errors.New("--export-cache requires type=<type>")
+	}
+	if _, ok := ex.Attrs["mode"]; !ok {
+		ex.Attrs["mode"] = "min"
+	}
+	return ex, nil
+}
+
 // parseOutput parses --output
 // Refer to https://github.com/moby/buildkit/blob/master/cmd/buildctl/build/output.go#L56
 func parseOutput(exports string) ([]client.ExportEntry, error) {
diff --git a/pkg/lang/ir/conda.go b/pkg/lang/ir/conda.go
index 51516925a..47dadc535 100644
--- a/pkg/lang/ir/conda.go
+++ b/pkg/lang/ir/conda.go
@@ -36,7 +36,7 @@ func (g Graph) compileCondaChannel(root llb.State) llb.State {
 		logrus.WithField("conda-channel", *g.CondaChannel).Debug("using custom connda channel")
 		stage := root.
 			File(llb.Mkfile(condarc,
-				0644, []byte(*g.CondaChannel), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] settings conda channel"))
+				0644, []byte(*g.CondaChannel), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] setting conda channel"))
 		return stage
 	}
 	return root
@@ -72,7 +72,7 @@ func (g Graph) compileCondaPackages(root llb.State) llb.State {
 	// Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
 	cache := root.File(llb.Mkdir("/cache",
 		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),
-		llb.WithCustomName("[internal] settings conda cache mount permissions"))
+		llb.WithCustomName("[internal] setting conda cache mount permissions"))
 	run := root.
 		Run(llb.Shlex(cmd), llb.WithCustomNamef("conda install %s",
 			strings.Join(g.CondaPackages, " ")))
diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
index d96398709..b6b458bc6 100644
--- a/pkg/lang/ir/python.go
+++ b/pkg/lang/ir/python.go
@@ -122,7 +122,7 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 	root = llb.User("envd")(root)
 	// Refer to https://github.com/moby/buildkit/blob/31054718bf775bf32d1376fe1f3611985f837584/frontend/dockerfile/dockerfile2llb/convert_runmount.go#L46
 	cache := root.File(llb.Mkdir("/cache",
-		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] settings pip cache mount permissions"))
+		0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomName("[internal] setting pip cache mount permissions"))
 	run := root.
 		Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
 			strings.Join(g.PyPIPackages, " ")))
@@ -143,10 +143,10 @@ func (g Graph) compilePyPIIndex(root llb.State) llb.State {
 		pypiMirror := root.
 			File(llb.Mkdir(filepath.Dir(pypiIndexFilePath),
 				0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),
-				llb.WithCustomName("[internal] settings PyPI index")).
+				llb.WithCustomName("[internal] setting PyPI index")).
 			File(llb.Mkfile(pypiIndexFilePath,
 				0644, []byte(content), llb.WithUIDGID(g.uid, g.gid)),
-				llb.WithCustomName("[internal] settings PyPI index"))
+				llb.WithCustomName("[internal] setting PyPI index"))
 		return pypiMirror
 	}
 	return root
diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
index a472fc0a3..40f62822e 100644
--- a/pkg/lang/ir/system.go
+++ b/pkg/lang/ir/system.go
@@ -32,11 +32,11 @@ func (g Graph) compileUbuntuAPT(root llb.State) llb.State {
 		logrus.WithField("source", *g.UbuntuAPTSource).Debug("using custom APT source")
 		aptSource := llb.Scratch().
 			File(llb.Mkdir(filepath.Dir(aptSourceFilePath), 0755, llb.WithParents(true)),
-				llb.WithCustomName("[internal] settings apt source")).
+				llb.WithCustomName("[internal] setting apt source")).
 			File(llb.Mkfile(aptSourceFilePath, 0644, []byte(*g.UbuntuAPTSource)),
-				llb.WithCustomName("[internal] settings apt source"))
+				llb.WithCustomName("[internal] setting apt source"))
 		return llb.Merge([]llb.State{root, aptSource},
-			llb.WithCustomName("[internal] settings apt source"))
+			llb.WithCustomName("[internal] setting apt source"))
 	}
 	return root
 }


Test patch:
diff --git a/pkg/builder/util_test.go b/pkg/builder/util_test.go
index 049699c7d..740bb8d4e 100644
--- a/pkg/builder/util_test.go
+++ b/pkg/builder/util_test.go
@@ -16,10 +16,118 @@ package builder
 import (
 	"testing"
 
+	"github.com/moby/buildkit/client"
+	gatewayclient "github.com/moby/buildkit/frontend/gateway/client"
 	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/require"
 )
 
-func Test_parseOutput(t *testing.T) {
+func TestParseImportCache(t *testing.T) {
+	type testCase struct {
+		importCaches []string // --import-cache
+		expected     []gatewayclient.CacheOptionsEntry
+		expectedErr  string
+	}
+	testCases := []testCase{
+		{
+			importCaches: []string{"type=registry,ref=example.com/foo/bar", "type=local,src=/path/to/store"},
+			expected: []gatewayclient.CacheOptionsEntry{
+				{
+					Type: "registry",
+					Attrs: map[string]string{
+						"ref": "example.com/foo/bar",
+					},
+				},
+				{
+					Type: "local",
+					Attrs: map[string]string{
+						"src": "/path/to/store",
+					},
+				},
+			},
+		},
+		{
+			importCaches: []string{"example.com/foo/bar", "example.com/baz/qux"},
+			expected: []gatewayclient.CacheOptionsEntry{
+				{
+					Type: "registry",
+					Attrs: map[string]string{
+						"ref": "example.com/foo/bar",
+					},
+				},
+				{
+					Type: "registry",
+					Attrs: map[string]string{
+						"ref": "example.com/baz/qux",
+					},
+				},
+			},
+		},
+	}
+	for _, tc := range testCases {
+		im, err := ParseImportCache(tc.importCaches)
+		if tc.expectedErr == "" {
+			require.EqualValues(t, tc.expected, im)
+		} else {
+			require.Error(t, err)
+			require.Contains(t, err.Error(), tc.expectedErr)
+		}
+	}
+}
+
+func TestParseExportCache(t *testing.T) {
+	type testCase struct {
+		exportCaches          []string // --export-cache
+		legacyExportCacheOpts []string // --export-cache-opt (legacy)
+		expected              []client.CacheOptionsEntry
+		expectedErr           string
+	}
+	testCases := []testCase{
+		{
+			exportCaches: []string{"type=registry,ref=example.com/foo/bar"},
+			expected: []client.CacheOptionsEntry{
+				{
+					Type: "registry",
+					Attrs: map[string]string{
+						"ref":  "example.com/foo/bar",
+						"mode": "min",
+					},
+				},
+			},
+		},
+		{
+			exportCaches:          []string{"example.com/foo/bar"},
+			legacyExportCacheOpts: []string{"mode=max"},
+			expected: []client.CacheOptionsEntry{
+				{
+					Type: "registry",
+					Attrs: map[string]string{
+						"ref":  "example.com/foo/bar",
+						"mode": "max",
+					},
+				},
+			},
+		},
+		{
+			exportCaches:          []string{"type=registry,ref=example.com/foo/bar"},
+			legacyExportCacheOpts: []string{"mode=max"},
+			expectedErr:           "--export-cache-opt is not supported for the specified --export-cache",
+		},
+		// TODO: test multiple exportCaches (valid for CLI but not supported by solver)
+
+	}
+	for _, tc := range testCases {
+		ex, err := ParseExportCache(tc.exportCaches, tc.legacyExportCacheOpts)
+		if tc.expectedErr == "" {
+			require.EqualValues(t, tc.expected, ex)
+		} else {
+			require.Error(t, err)
+			require.Contains(t, err.Error(), tc.expectedErr)
+		}
+	}
+}
+
+func TestParseOutput(t *testing.T) {
 	type args struct {
 		output string
 	}