Repository: tensorchord/envd (language: python)

Issue:
feat: disable buildkit merge op when using moby builder
### Describe the feature

Currently, we use `docker` instead of `moby` builder because https://github.com/tensorchord/envd/issues/1547.

- track the upstream issue https://github.com/moby/moby/issues/45111

The downside is that we need to convert the image type which may take a long time for large CUDA-based images.

A possible solution is to disable the merge op when using moby builder.

Currently, merge op is used by:
- system package & language installation
- extra HTTP sources (rare)
- language packages (when the build has multiple languages, i.e. `python` + `r`) (rare)


### Why do you need this feature?

_No response_

### Additional context

_No response_


Discussion / hints:
cc @pingsutw 
thanks, Let me know once you open a pr. I can test it

Fix patch:
diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go
index a188d3992..fcadd94b8 100644
--- a/pkg/lang/ir/v1/compile.go
+++ b/pkg/lang/ir/v1/compile.go
@@ -29,6 +29,7 @@ import (
 	servertypes "github.com/tensorchord/envd-server/api/types"
 
 	"github.com/tensorchord/envd/pkg/config"
+	"github.com/tensorchord/envd/pkg/home"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/progress/compileui"
 	"github.com/tensorchord/envd/pkg/types"
@@ -113,10 +114,18 @@ func (g *generalGraph) Compile(ctx context.Context, envName string, pub string,
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create compileui")
 	}
+	c, err := home.GetManager().ContextGetCurrent()
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to get the current envd context")
+	}
+
 	g.Writer = w
 	g.EnvironmentName = envName
 	g.PublicKeyPath = pub
 	g.Platform = platform
+	if c.Builder == types.BuilderTypeMoby {
+		g.DisableMergeOp = true
+	}
 
 	uid, gid, err := g.getUIDGID()
 	if err != nil {
@@ -308,24 +317,29 @@ func (g *generalGraph) CompileLLB(uid, gid int) (llb.State, error) {
 	}
 
 	systemPackages := g.compileSystemPackages(aptMirror)
-	lang, err := g.compileLanguage(aptMirror)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to compile language")
+	var language llb.State
+	if g.DisableMergeOp {
+		language, err = g.compileLanguage(systemPackages)
+		if err != nil {
+			return llb.State{}, errors.Wrap(err, "failed to compile language from system packages")
+		}
+	} else {
+		lang, err := g.compileLanguage(aptMirror)
+		if err != nil {
+			return llb.State{}, errors.Wrap(err, "failed to compile language from apt mirror")
+		}
+		language = llb.Merge([]llb.State{
+			base,
+			llb.Diff(base, systemPackages, llb.WithCustomName("[internal] install system packages")),
+			llb.Diff(base, lang, llb.WithCustomName("[internal] prepare language")),
+		}, llb.WithCustomName("[internal] language environment and system packages"))
 	}
-	merge := llb.Merge([]llb.State{
-		base,
-		llb.Diff(base, systemPackages, llb.WithCustomName("[internal] install system packages")),
-		llb.Diff(base, lang, llb.WithCustomName("[internal] prepare language")),
-	}, llb.WithCustomName("[internal] language environment and system packages"))
-	packages := g.compileLanguagePackages(merge)
+	packages := g.compileLanguagePackages(language)
 	if err != nil {
 		return llb.State{}, errors.Wrap(err, "failed to compile language")
 	}
 
-	source, err := g.compileExtraSource(packages)
-	if err != nil {
-		return llb.State{}, errors.Wrap(err, "failed to compile extra source")
-	}
+	source := g.compileExtraSource(packages)
 	copy := g.compileCopy(source)
 
 	// dev postprocessing: related to UID, which may not be cached
diff --git a/pkg/lang/ir/v1/editor.go b/pkg/lang/ir/v1/editor.go
index c081fd90d..613d2be2b 100644
--- a/pkg/lang/ir/v1/editor.go
+++ b/pkg/lang/ir/v1/editor.go
@@ -58,18 +58,13 @@ func (g generalGraph) compileVSCode(root llb.State) (llb.State, error) {
 }
 
 // nolint:unused
-func (g *generalGraph) compileJupyter() error {
+func (g *generalGraph) compileJupyter() {
 	if g.JupyterConfig == nil {
-		return nil
+		return
 	}
 
+	// no need to check if `python` is installed since v1 should support user costumed image
 	g.PyPIPackages = append(g.PyPIPackages, []string{"jupyter"})
-	for _, language := range g.Languages {
-		if language.Name == "python" {
-			return nil
-		}
-	}
-	return errors.Newf("Jupyter is not supported in other languages yet")
 }
 
 func (g generalGraph) generateJupyterCommand(workingDir string) []string {
diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go
index f6ec6a7e8..b555a5da2 100644
--- a/pkg/lang/ir/v1/system.go
+++ b/pkg/lang/ir/v1/system.go
@@ -208,9 +208,23 @@ func (g generalGraph) compileSystemPackages(root llb.State) llb.State {
 }
 
 // nolint:unparam
-func (g *generalGraph) compileExtraSource(root llb.State) (llb.State, error) {
+func (g *generalGraph) compileExtraSource(root llb.State) llb.State {
 	if len(g.HTTP) == 0 {
-		return root, nil
+		return root
+	}
+	if g.DisableMergeOp {
+		for _, httpInfo := range g.HTTP {
+			src := llb.HTTP(
+				httpInfo.URL,
+				llb.Checksum(httpInfo.Checksum),
+				llb.Filename(httpInfo.Filename),
+				llb.Chown(g.uid, g.gid),
+			)
+			root = root.File(llb.Copy(
+				src, "/", g.getExtraSourceDir(), &llb.CopyInfo{CreateDestPath: true},
+			))
+		}
+		return root
 	}
 	inputs := []llb.State{}
 	for _, httpInfo := range g.HTTP {
@@ -225,7 +239,7 @@ func (g *generalGraph) compileExtraSource(root llb.State) (llb.State, error) {
 		))
 	}
 	inputs = append(inputs, root)
-	return llb.Merge(inputs, llb.WithCustomName("[internal] build source layers")), nil
+	return llb.Merge(inputs, llb.WithCustomName("[internal] build source layers"))
 }
 
 func (g *generalGraph) compileLanguage(root llb.State) (llb.State, error) {
@@ -255,9 +269,8 @@ func (g *generalGraph) compileLanguage(root llb.State) (llb.State, error) {
 }
 
 func (g *generalGraph) compileLanguagePackages(root llb.State) llb.State {
-	packs := []llb.State{}
-
 	// Use default python in the base image if install.python() is not specified.
+	g.compileJupyter()
 	index := g.compilePyPIIndex(root)
 	pack := g.compilePyPIPackages(index)
 	if g.CondaConfig != nil {
@@ -268,21 +281,12 @@ func (g *generalGraph) compileLanguagePackages(root llb.State) llb.State {
 	for _, language := range g.Languages {
 		switch language.Name {
 		case "r":
-			pack = g.installRPackages(root)
+			pack = g.installRPackages(pack)
 		case "julia":
-			pack = g.installJuliaPackages(root)
+			pack = g.installJuliaPackages(pack)
 		}
-		packs = append(packs, pack)
-	}
-	if len(packs) <= 1 {
-		// there is only one language needs to be installed, thus no need to merge
-		return pack
-	}
-	for i, lang := range g.Languages {
-		packs[i] = llb.Diff(root, packs[i], llb.WithCustomNamef("[internal] get diff of %s's packages", lang.Name))
 	}
-	return llb.Merge(append([]llb.State{root}, packs...),
-		llb.WithCustomName("[internal] merge packages for all language environments"))
+	return pack
 }
 
 func (g *generalGraph) compileDevPackages(root llb.State) llb.State {
diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
index cc869d6a3..7abb1b2ba 100644
--- a/pkg/lang/ir/v1/types.go
+++ b/pkg/lang/ir/v1/types.go
@@ -79,6 +79,10 @@ type generalGraph struct {
 	// e.g. mnist, streamlit-mnist
 	EnvironmentName string
 
+	// (v1) disable `merge` op for `moby` builder
+	// check https://github.com/tensorchord/envd/issues/1693
+	DisableMergeOp bool
+
 	ir.RuntimeGraph
 
 	Platform *ocispecs.Platform


Test patch:
diff --git a/e2e/v1/cli/moby_test.go b/e2e/v1/cli/moby_test.go
new file mode 100644
index 000000000..692fcbff6
--- /dev/null
+++ b/e2e/v1/cli/moby_test.go
@@ -0,0 +1,63 @@
+// Copyright 2023 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 cli
+
+import (
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+
+	"github.com/tensorchord/envd/e2e/v1"
+	"github.com/tensorchord/envd/pkg/home"
+	"github.com/tensorchord/envd/pkg/types"
+)
+
+var _ = Describe("e2e moby builder test", Ordered, func() {
+	exampleName := "build-test"
+	defaultContext := "default"
+	mobyContext := "envd-test-moby"
+	ctx := types.Context{
+		Name:    mobyContext,
+		Builder: types.BuilderTypeMoby,
+		Runner:  types.RunnerTypeDocker,
+	}
+
+	BeforeAll(func() {
+		Expect(home.Initialize()).To(Succeed())
+		err := home.GetManager().ContextCreate(ctx, true)
+		Expect(err).NotTo(HaveOccurred())
+	})
+
+	It("should find a new context", func() {
+		ctx, err := home.GetManager().ContextGetCurrent()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(ctx.Name).To(Equal(mobyContext))
+		Expect(ctx.Builder).To(Equal(types.BuilderTypeMoby))
+	})
+
+	It("build images with moby", func() {
+		e := e2e.NewExample(e2e.BuildContextDirWithName(exampleName), "e2e")
+		e.BuildImage(true)()
+	})
+
+	AfterAll(func() {
+		err := home.GetManager().ContextUse(defaultContext)
+		Expect(err).NotTo(HaveOccurred())
+		err = home.GetManager().ContextRemove(mobyContext)
+		Expect(err).NotTo(HaveOccurred())
+		context, err := home.GetManager().ContextGetCurrent()
+		Expect(err).NotTo(HaveOccurred())
+		Expect(context.Name).To(Equal(defaultContext))
+	})
+})
diff --git a/e2e/v1/cli/testdata/build-test/build.envd b/e2e/v1/cli/testdata/build-test/build.envd
index d1f3175fa..27814172b 100644
--- a/e2e/v1/cli/testdata/build-test/build.envd
+++ b/e2e/v1/cli/testdata/build-test/build.envd
@@ -2,6 +2,7 @@
 
 
 def build():
-    base(dev=True)
+    base()
+    install.apt_packages(name=["htop"])
     install.conda()
     install.python()