Repository: tensorchord/envd (language: python)

Issue:
bug(lang): custom base image does not work as expected.
## Description

We always use conda in the compilePython function, and some base images does not have it. We should maintain a separate logic for custom images.

/assign

## Reproduction

## Additional Info

---
<!-- Issue Author: Don't delete this message to encourage other users to support your issue! -->
**Message from the maintainers**:

Impacted by this bug? Give it a 👍. We prioritise the issues with the most 👍.

chore: Separate e2e test with unit test cases
## Description

Now our e2e test cases are placed in pkg/app. It took too long time. We should separate the e2e test to a separate Github Action.

---
<!-- Issue Author: Don't delete this message to encourage other users to support your issue! -->
**Message from the maintainers**:

Love this enhancement proposal? Give it a 👍. We prioritise the proposals with the most 👍.



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index fe0506fa4..9c56a2484 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -93,13 +93,47 @@ jobs:
           git diff --cached --exit-code || (echo 'Please run "make generate" to verify generate' && exit 1);
       - name: Test
         run: make test
-      - name: send coverage
-        uses: shogo82148/actions-goveralls@v1
+      - name: Upload coverage report
+        uses: actions/upload-artifact@v2
         with:
-          path-to-profile: coverage.out
-          working-directory: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
-          flag-name: ${{ matrix.os }}
-          parallel: true
+          name: coverage-out
+          path: ${{ env.GOPATH }}/src/github.com/tensorchord/envd/coverage.out
+  e2e:
+    name: e2e
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+    runs-on: ${{ matrix.os }}
+    env:
+      GOPATH: ${{ github.workspace }}/go
+    defaults:
+      run:
+        working-directory: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v2
+        with:
+          path: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
+      - name: Setup Go
+        uses: actions/setup-go@v2
+        with:
+          go-version: 1.18
+      - name: Cache Go modules
+        uses: actions/cache@preview
+        with:
+          path: ~/go/pkg/mod
+          key: ${{ runner.os }}-build-${{ hashFiles('**/go.sum') }}
+          restore-keys: |
+            ${{ runner.OS }}-build-${{ env.cache-name }}-
+            ${{ runner.OS }}-build-
+            ${{ runner.OS }}-
+      - name: e2e test
+        run: make e2e-test
+      - name: Upload coverage report
+        uses: actions/upload-artifact@v2
+        with:
+          name: e2e-coverage-out
+          path: ${{ env.GOPATH }}/src/github.com/tensorchord/envd/e2e-coverage.out
   build:
     name: build
     strategy:
@@ -133,10 +167,42 @@ jobs:
         run: make
 
   # notifies that all test jobs are finished.
-  finish:
-    needs: test
+  report:
+    needs:
+      - test
+      - e2e
+    env:
+      GOPATH: ${{ github.workspace }}/go
     runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
     steps:
-      - uses: shogo82148/actions-goveralls@v1
+      - name: Check out code
+        uses: actions/checkout@v2
+        with:
+          path: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
+      - name: Setup Go
+        uses: actions/setup-go@v2
         with:
-          parallel-finished: true
+          go-version: 1.18
+      - name: Install bins
+        run: |
+          go install github.com/mattn/goveralls@latest
+          go install github.com/wadey/gocovmerge@latest
+      - name: Get coverage report
+        uses: actions/download-artifact@v3
+        with:
+          name: coverage-out
+          path: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
+      - name: Get coverage report
+        uses: actions/download-artifact@v3
+        with:
+          name: e2e-coverage-out
+          path: ${{ env.GOPATH }}/src/github.com/tensorchord/envd
+      # - name: Send coverage
+      #   env:
+      #     COVERALLS_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+      #   run: |
+      #     gocovmerge e2e-coverage.out coverage.out > final.out
+      #     goveralls -coverprofile=final.out -service=github
diff --git a/Makefile b/Makefile
index b726381b3..e1f487009 100644
--- a/Makefile
+++ b/Makefile
@@ -150,9 +150,12 @@ test-local:
 	@go test -v -race -coverprofile=coverage.out ./...
 
 test: generate  ## Run the tests
-	@go test -race -coverpkg=./pkg/... -coverprofile=coverage.out ./...
+	@go test -race -coverpkg=./pkg/... -coverprofile=coverage.out $(shell go list ./... | grep -v app)
 	@go tool cover -func coverage.out | tail -n 1 | awk '{ print "Total coverage: " $$3 }'
 
+e2e-test: generate
+	@go test -race -coverpkg=./pkg/app -coverprofile=e2e-coverage.out ./pkg/app
+
 clean:  ## Clean the outputs and artifacts
 	@-rm -vrf ${OUTPUT_DIR}
 	@-rm -vrf ${DEBUG_DIR}
diff --git a/examples/custom-image/build.envd b/examples/custom-image/build.envd
index 1bd514361..418dc759b 100644
--- a/examples/custom-image/build.envd
+++ b/examples/custom-image/build.envd
@@ -1,5 +1,5 @@
 def build():
     base(language="python", image="python:3.8")
     install.python_packages(name=[
-      "ormb",
+      "via",
     ])
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
index 69e3344c0..cf7fdf466 100644
--- a/pkg/lang/ir/compile.go
+++ b/pkg/lang/ir/compile.go
@@ -194,21 +194,29 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 	aptStage := g.compileUbuntuAPT(base)
 	var merged llb.State
 	var err error
-	switch g.Language.Name {
-	case "r":
-		merged, err = g.compileRLang(aptStage)
-		if err != nil {
-			return llb.State{}, errors.Wrap(err, "failed to compile r language")
-		}
-	case "python":
-		merged, err = g.compilePython(aptStage)
+	// Use custom logic when image is specified.
+	if g.Image != nil {
+		merged, err = g.compileCustomPython(aptStage)
 		if err != nil {
-			return llb.State{}, errors.Wrap(err, "failed to compile python")
+			return llb.State{}, errors.Wrap(err, "failed to compile custom python image")
 		}
-	case "julia":
-		merged, err = g.compileJulia(aptStage)
-		if err != nil {
-			return llb.State{}, errors.Wrap(err, "failed to compile julia")
+	} else {
+		switch g.Language.Name {
+		case "r":
+			merged, err = g.compileRLang(aptStage)
+			if err != nil {
+				return llb.State{}, errors.Wrap(err, "failed to compile r language")
+			}
+		case "python":
+			merged, err = g.compilePython(aptStage)
+			if err != nil {
+				return llb.State{}, errors.Wrap(err, "failed to compile python")
+			}
+		case "julia":
+			merged, err = g.compileJulia(aptStage)
+			if err != nil {
+				return llb.State{}, errors.Wrap(err, "failed to compile julia")
+			}
 		}
 	}
 
diff --git a/pkg/lang/ir/custom.go b/pkg/lang/ir/custom.go
new file mode 100644
index 000000000..fa3f4fb58
--- /dev/null
+++ b/pkg/lang/ir/custom.go
@@ -0,0 +1,87 @@
+// 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 ir
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/moby/buildkit/client/llb"
+)
+
+func (g Graph) compileCustomPython(aptStage llb.State) (llb.State, error) {
+	pypiMirrorStage := g.compilePyPIIndex(aptStage)
+
+	builtinSystemStage := pypiMirrorStage
+
+	systemStage := g.compileCustomSystemPackages(builtinSystemStage)
+	pypiStage := g.compileCustomPyPIPackages(systemStage)
+
+	return pypiStage, nil
+}
+
+func (g Graph) compileCustomPyPIPackages(root llb.State) llb.State {
+	if len(g.PyPIPackages) == 0 {
+		return root
+	}
+
+	cacheDir := "/home/root/.cache"
+
+	// Compose the package install command.
+	var sb strings.Builder
+	// Always use the conda's pip.
+	sb.WriteString("pip install")
+	for _, pkg := range g.PyPIPackages {
+		sb.WriteString(fmt.Sprintf(" %s", pkg))
+	}
+
+	cmd := sb.String()
+	// 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.WithCustomName("[internal] settings pip cache mount permissions"))
+	run := root.
+		Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
+			strings.Join(g.PyPIPackages, " ")))
+	run.AddMount(cacheDir, cache,
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared),
+		llb.SourcePath("/cache"))
+	return run.Root()
+}
+
+func (g Graph) compileCustomSystemPackages(root llb.State) llb.State {
+	if len(g.SystemPackages) == 0 {
+		return root
+	}
+
+	// Compose the package install command.
+	var sb strings.Builder
+	sb.WriteString("apt-get update && apt-get install -y --no-install-recommends")
+
+	for _, pkg := range g.SystemPackages {
+		sb.WriteString(fmt.Sprintf(" %s", pkg))
+	}
+
+	cacheDir := "/var/cache/apt"
+	cacheLibDir := "/var/lib/apt"
+
+	run := root.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())),
+		llb.WithCustomNamef("apt-get install %s",
+			strings.Join(g.SystemPackages, " ")))
+	run.AddMount(cacheDir, llb.Scratch(),
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared))
+	run.AddMount(cacheLibDir, llb.Scratch(),
+		llb.AsPersistentCacheDir(g.CacheID(cacheLibDir), llb.CacheMountShared))
+	return run.Root()
+}
diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
index 7524b951c..a472fc0a3 100644
--- a/pkg/lang/ir/system.go
+++ b/pkg/lang/ir/system.go
@@ -98,8 +98,9 @@ func (g *Graph) compileBase() llb.State {
 	logger.Debug("compile base image")
 
 	var base llb.State
+	// Do not update user permission in the base image.
 	if g.Image != nil {
-		base = llb.Image(*g.Image)
+		return llb.Image(*g.Image)
 	} else if g.CUDA == nil && g.CUDNN == nil {
 		switch g.Language.Name {
 		case "r":
@@ -135,11 +136,9 @@ func (g *Graph) compileBase() llb.State {
 			Run(llb.Shlex("sed -i \"s./root./home/envd.g\" /etc/passwd"),
 				llb.WithCustomName("[internal] set root home dir to /home/envd")).
 			Run(llb.Shlex("sed -i \"s/envd:x:1001/envd:x:0/g\" /etc/group"),
-				llb.WithCustomName("[internal] set envd group to 0 as root group"))
-		if g.Image == nil && g.Language.Name == "python" {
-			res = res.Run(llb.Shlex("chown -R root:root /opt/conda"),
+				llb.WithCustomName("[internal] set envd group to 0 as root group")).
+			Run(llb.Shlex("chown -R root:root /opt/conda"),
 				llb.WithCustomName("[internal] configure user permissions"))
-		}
 	} else {
 		res = base.
 			Run(llb.Shlex(fmt.Sprintf("groupadd -g %d envd", g.gid)),
@@ -149,11 +148,9 @@ func (g *Graph) compileBase() llb.State {
 			Run(llb.Shlex("adduser envd sudo"),
 				llb.WithCustomName("[internal] add user envd to sudoers")).
 			Run(llb.Shlex("chown -R envd:envd /usr/local/lib"),
+				llb.WithCustomName("[internal] configure user permissions")).
+			Run(llb.Shlex("chown -R envd:envd /opt/conda"),
 				llb.WithCustomName("[internal] configure user permissions"))
-		if g.Image == nil && g.Language.Name == "python" {
-			res = res.Run(llb.Shlex("chown -R envd:envd /opt/conda"),
-				llb.WithCustomName("[internal] configure user permissions"))
-		}
 	}
 	return llb.User("envd")(res.Root())
 }


Test patch:
diff --git a/pkg/app/build_test.go b/pkg/app/build_test.go
index 4aa38129e..1cdecf177 100644
--- a/pkg/app/build_test.go
+++ b/pkg/app/build_test.go
@@ -24,26 +24,55 @@ import (
 	"github.com/tensorchord/envd/pkg/home"
 )
 
-var _ = Describe("build command", func() {
-	buildContext := "testdata/build-test"
-	args := []string{
-		"envd.test", "--debug", "build", "--path", buildContext,
-	}
-	BeforeEach(func() {
+var _ = Describe("build command", Ordered, func() {
+	buildTestName := "testdata/build-test"
+	customImageTestName := "testdata/custom-image-test"
+	When("given the right arguments", func() {
+		It("should build successfully", func() {
+			app := New()
+			err := app.Run([]string{"envd.test", "--debug", "bootstrap"})
+			Expect(err).NotTo(HaveOccurred())
+			cli, err := docker.NewClient(context.TODO())
+			Expect(err).NotTo(HaveOccurred())
+			_, err = cli.Destroy(context.TODO(), buildTestName)
+			Expect(err).NotTo(HaveOccurred())
+
+			app = New()
+			args := []string{
+				"envd.test", "--debug", "build", "--path", buildTestName,
+			}
+			err = app.Run(args)
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+	When("given the custom image", func() {
+		It("should build successfully", func() {
+			app := New()
+			err := app.Run([]string{"envd.test", "--debug", "bootstrap"})
+			Expect(err).NotTo(HaveOccurred())
+			cli, err := docker.NewClient(context.TODO())
+			Expect(err).NotTo(HaveOccurred())
+			_, err = cli.Destroy(context.TODO(), customImageTestName)
+			Expect(err).NotTo(HaveOccurred())
+
+			app = New()
+			args := []string{
+				"envd.test", "--debug", "build", "--path", customImageTestName,
+			}
+			err = app.Run(args)
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+	AfterAll(func() {
 		Expect(home.Initialize()).NotTo(HaveOccurred())
 		app := New()
 		err := app.Run([]string{"envd.test", "--debug", "bootstrap"})
 		Expect(err).NotTo(HaveOccurred())
 		cli, err := docker.NewClient(context.TODO())
 		Expect(err).NotTo(HaveOccurred())
-		_, err = cli.Destroy(context.TODO(), buildContext)
+		_, err = cli.Destroy(context.TODO(), buildTestName)
+		Expect(err).NotTo(HaveOccurred())
+		_, err = cli.Destroy(context.TODO(), customImageTestName)
 		Expect(err).NotTo(HaveOccurred())
-	})
-	When("given the right arguments", func() {
-		It("should build successfully", func() {
-			app := New()
-			err := app.Run(args)
-			Expect(err).NotTo(HaveOccurred())
-		})
 	})
 })
diff --git a/pkg/app/testdata/custom-image-test/build.envd b/pkg/app/testdata/custom-image-test/build.envd
new file mode 100644
index 000000000..418dc759b
--- /dev/null
+++ b/pkg/app/testdata/custom-image-test/build.envd
@@ -0,0 +1,5 @@
+def build():
+    base(language="python", image="python:3.8")
+    install.python_packages(name=[
+      "via",
+    ])
diff --git a/pkg/app/testdata/up-test/build.envd b/pkg/app/testdata/up-test/build.envd
index 9228bc397..5ca54710d 100644
--- a/pkg/app/testdata/up-test/build.envd
+++ b/pkg/app/testdata/up-test/build.envd
@@ -6,7 +6,7 @@ def build():
     install.system_packages(name = ["screenfetch"])
     shell("zsh")
     config.pip_index(url = "https://pypi.tuna.tsinghua.edu.cn/simple")
-    git_config(name="envd", email="envd@envd", editor="vim")
+    # git_config(name="envd", email="envd@envd", editor="vim")
     install.vscode_extensions([
         "ms-python.python"
     ])
diff --git a/pkg/app/up_test.go b/pkg/app/up_test.go
index 813455507..31a13b3cc 100644
--- a/pkg/app/up_test.go
+++ b/pkg/app/up_test.go
@@ -22,15 +22,16 @@ import (
 
 	"github.com/tensorchord/envd/pkg/docker"
 	"github.com/tensorchord/envd/pkg/home"
+	"github.com/tensorchord/envd/pkg/lang/ir"
 )
 
-var _ = Describe("up command", func() {
+var _ = Describe("up command", Ordered, func() {
 	buildContext := "testdata/up-test"
 	env := "up-test"
 	baseArgs := []string{
 		"envd.test", "--debug",
 	}
-	BeforeEach(func() {
+	BeforeAll(func() {
 		Expect(home.Initialize()).NotTo(HaveOccurred())
 		app := New()
 		err := app.Run(append(baseArgs, "bootstrap"))
@@ -63,4 +64,8 @@ var _ = Describe("up command", func() {
 			Expect(err).NotTo(HaveOccurred())
 		})
 	})
+	AfterAll(func() {
+		// Cleanup graph.
+		ir.DefaultGraph = ir.NewGraph()
+	})
 })
diff --git a/pkg/home/manager_test.go b/pkg/home/manager_test.go
index 8c7a4b6b6..f469f1a2c 100644
--- a/pkg/home/manager_test.go
+++ b/pkg/home/manager_test.go
@@ -28,7 +28,20 @@ import (
 var _ = Describe("home manager", func() {
 	When("initialized", func() {
 		It("should initialized successfully", func() {
-			Expect(Initialize()).NotTo(HaveOccurred())
+			defaultManager = &generalManager{
+				cacheMap: make(map[string]bool),
+				context: types.EnvdContext{
+					Current: "default",
+					Contexts: []types.Context{
+						{
+							Name:          "default",
+							Builder:       types.BuilderTypeDocker,
+							BuilderSocket: "envd_buildkitd",
+						},
+					},
+				},
+			}
+			Expect(defaultManager.init()).NotTo(HaveOccurred())
 			m := GetManager()
 			Expect(m.CacheDir()).To(Equal(filepath.Join(xdg.CacheHome, "envd")))
 			Expect(m.ConfigFile()).To(Equal(filepath.Join(xdg.ConfigHome, "envd/config.envd")))