feat(ir): all in llb (#941)

* change python dockerfile to llb Signed-off-by: Keming <[email protected]> * change to docker hub Signed-off-by: Keming <[email protected]> * copy envd-sshd Signed-off-by: Keming <[email protected]> * add llb log Signed-off-by: Keming <[email protected]> * delete outdated comments Signed-off-by: Keming <[email protected]> * cuda version Signed-off-by: Keming <[email protected]> * del python dockerfile Signed-off-by: Keming <[email protected]> * fix test Signed-off-by: Keming <[email protected]> Signed-off-by: Keming <[email protected]>
tensorchord · Sep 28, 2022 · b704029 · b704029
1 parent f5f70e0
commit b704029
Show file tree

Hide file tree

Showing 13 changed files with 112 additions and 111 deletions.
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
@@ -26,9 +26,8 @@ jobs:
       - name: Docker Login
         uses: docker/login-action@v2
         with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
+          username: ${{ secrets.DOCKERIO_USERNAME }}
+          password: ${{ secrets.DOCKERIO_TOKEN }}
       - name: Run GoReleaser
         uses: goreleaser/goreleaser-action@v3
         with:
@@ -109,12 +108,13 @@ jobs:
         key: ${{ runner.os }}-buildx-${{ github.sha }}
         restore-keys: |
           ${{ runner.os }}-buildx-
+    - name: Docker Login
+      uses: docker/login-action@v2
+      with:
+        username: ${{ secrets.DOCKERIO_USERNAME }}
+        password: ${{ secrets.DOCKERIO_TOKEN }}
     - name: Docker Buildx
-      env:
-        DOCKERIO_USERNAME: ${{ secrets.DOCKERIO_USERNAME }}
-        DOCKERIO_PASSWORD: ${{ secrets.DOCKERIO_PASSWORD }}
       run: |
-          docker login --username "${DOCKERIO_USERNAME}" --password "${DOCKERIO_PASSWORD}"
           ./base-images/build.sh
           DOCKER_IMAGE_TAG=latest ./base-images/build.sh
   cache_publish:
@@ -145,10 +145,11 @@ jobs:
       run: |
         mv dist/envd_linux_amd64_v1/envd /usr/local/bin/envd
         chmod +x /usr/local/bin/envd
+    - name: Docker Login
+      uses: docker/login-action@v2
+      with:
+        username: ${{ secrets.DOCKERIO_USERNAME }}
+        password: ${{ secrets.DOCKERIO_TOKEN }}
     - name: Build and push
-      env:
-        DOCKERIO_USERNAME: ${{ secrets.DOCKERIO_USERNAME }}
-        DOCKERIO_PASSWORD: ${{ secrets.DOCKERIO_PASSWORD }}
       run: |
-          docker login --username "${DOCKERIO_USERNAME}" --password "${DOCKERIO_PASSWORD}"
           ./base-images/remote-cache/build-and-push-remote-cache.sh
diff --git a/.goreleaser.yaml b/.goreleaser.yaml
@@ -69,15 +69,15 @@ changelog:
       - '^chore:'
 dockers:
 - image_templates:
-  - "ghcr.io/tensorchord/envd-sshd-from-scratch:{{ .Version }}-amd64"
+  - "tensorchord/envd-sshd-from-scratch:v{{ .Version }}-amd64"
   use: buildx
   dockerfile: .goreleaser/envd-sshd.Dockerfile
   ids:
   - envd-sshd
   build_flag_templates:
   - "--platform=linux/amd64"
 - image_templates:
-  - "ghcr.io/tensorchord/envd-sshd-from-scratch:{{ .Version }}-arm64v8"
+  - "tensorchord/envd-sshd-from-scratch:v{{ .Version }}-arm64v8"
   use: buildx
   goarch: arm64
   ids:
@@ -86,10 +86,10 @@ dockers:
   build_flag_templates:
   - "--platform=linux/arm64/v8"
 docker_manifests:
-- name_template: ghcr.io/tensorchord/envd-sshd-from-scratch:{{ .Version }}
+- name_template: tensorchord/envd-sshd-from-scratch:v{{ .Version }}
   image_templates:
-  - ghcr.io/tensorchord/envd-sshd-from-scratch:{{ .Version }}-amd64
-  - ghcr.io/tensorchord/envd-sshd-from-scratch:{{ .Version }}-arm64v8
+  - tensorchord/envd-sshd-from-scratch:v{{ .Version }}-amd64
+  - tensorchord/envd-sshd-from-scratch:v{{ .Version }}-arm64v8
 # See https://github.com/tensorchord/envd/issues/908
 # brews:
 #   - name: envd

diff --git a/base-images/build.sh b/base-images/build.sh
@@ -16,25 +16,12 @@ RLANG_VERSION="${RLANG_VERSION:-4.2}"
 cd ${ROOT_DIR}
 # ubuntu 22.04 build require moby/buildkit version greater than 0.8.1
 if ! docker buildx inspect cuda; then
-    docker buildx create --use --platform linux/x86_64,linux/arm64,linux/ppc64le --driver-opt image=moby/buildkit:v0.10.3 --name cuda --node cuda
+    docker buildx create --use --platform linux/x86_64,linux/arm64,linux/ppc64le --driver-opt image=moby/buildkit:v0.10.3
 fi
 
 # https://github.com/docker/buildx/issues/495#issuecomment-754688157
 docker run --rm --privileged multiarch/qemu-user-static --reset -p yes
 
-docker buildx build \
-    --build-arg ENVD_VERSION=${ENVD_VERSION} \
-    --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-sshd-from-scratch \
-    --pull --push --platform linux/x86_64,linux/arm64 \
-    -t ${DOCKER_HUB_ORG}/python:${PYTHON_VERSION}-${ENVD_OS}-envd-${DOCKER_IMAGE_TAG} \
-    -f python${PYTHON_VERSION}-${ENVD_OS}.Dockerfile .
-docker buildx build --build-arg IMAGE_NAME=docker.io/nvidia/cuda \
-    --build-arg ENVD_VERSION=${ENVD_VERSION} \
-    --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-sshd-from-scratch \
-    --pull --push --platform linux/x86_64,linux/arm64 \
-    -t ${DOCKER_HUB_ORG}/python:${PYTHON_VERSION}-${ENVD_OS}-cuda11.2-cudnn8-envd-${DOCKER_IMAGE_TAG} \
-    -f python${PYTHON_VERSION}-${ENVD_OS}-cuda11.2.Dockerfile .
-
 # TODO(gaocegege): Support linux/arm64
 docker buildx build \
     --build-arg ENVD_VERSION=${ENVD_VERSION} \

diff --git a/base-images/python3.9-ubuntu20.04-cuda11.2.Dockerfile b/base-images/python3.9-ubuntu20.04-cuda11.2.Dockerfile
diff --git a/base-images/python3.9-ubuntu20.04.Dockerfile b/base-images/python3.9-ubuntu20.04.Dockerfile
diff --git a/base-images/remote-cache/build-and-push-remote-cache.sh b/base-images/remote-cache/build-and-push-remote-cache.sh
@@ -11,5 +11,6 @@ DOCKER_HUB_ORG="${DOCKER_HUB_ORG:-tensorchord}"
 cd ${ROOT_DIR}
 
 envd build --export-cache type=registry,ref=docker.io/${DOCKER_HUB_ORG}/python-cache:envd-v${ENVD_VERSION} --force
+envd build -f build.envd:build_gpu --export-cache type=registry,ref=docker.io/${DOCKER_HUB_ORG}/python-cache:envd-v${ENVD_VERSION} --force
 
 cd - > /dev/null
diff --git a/base-images/remote-cache/build.envd b/base-images/remote-cache/build.envd
@@ -1,2 +1,6 @@
 def build():
     base(os="ubuntu20.04", language="python3")
+
+def build_gpu():
+    base(os="ubuntu20.04", language="python3")
+    install.cuda(version="11.2", cudnn="8")
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -399,7 +399,7 @@ func (b generalBuilder) checkDepsFileUpdate(ctx context.Context, tag string, man
 			return true, err
 		}
 		modifiedtime := file.ModTime().Unix()
-		// Only needt o use the latest modified time
+		// Only need to use the latest modified time
 		if modifiedtime > latestTimestamp {
 			latestTimestamp = modifiedtime
 		}

diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -86,7 +86,7 @@ func Compile(ctx context.Context, envName string, pub string) (*llb.Definition,
 	}
 	state, err := DefaultGraph.Compile(uid, gid)
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to compile")
+		return nil, errors.Wrap(err, "failed to compile the graph")
 	}
 	// TODO(gaocegege): Support multi platform.
 	def, err := state.Marshal(ctx, llb.LinuxAmd64)

diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -112,6 +112,7 @@ func JuliaPackageServer(url string) error {
 
 func Shell(shell string) error {
 	DefaultGraph.Shell = shell
+	DefaultGraph.SystemPackages = append(DefaultGraph.SystemPackages, shell)
 	return nil
 }
 

diff --git a/pkg/lang/ir/shell.go b/pkg/lang/ir/shell.go
@@ -91,7 +91,7 @@ func (g Graph) compileZSH(root llb.State) (llb.State, error) {
 		File(llb.Mkfile(installPath,
 			0644, []byte(m.InstallScript()), llb.WithUIDGID(g.uid, g.gid)))
 	zshrc := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)),
-		llb.WithCustomName("install oh-my-zsh")).
+		llb.WithCustomName("[internal] install oh-my-zsh")).
 		File(llb.Mkfile(zshrcPath,
 			0644, []byte(m.ZSHRC()), llb.WithUIDGID(g.uid, g.gid)))
 	return zshrc, nil

diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -91,14 +91,15 @@ func (g Graph) compileCopy(root llb.State) llb.State {
 	return result
 }
 
-func (g *Graph) compileCUDAPackages(org, version string) llb.State {
-	return llb.Image(fmt.Sprintf(
-		"docker.io/%s/python:3.9-%s-cuda%s-cudnn%s-envd-%s",
-		org, g.OS, *g.CUDA, *g.CUDNN, version))
+func (g *Graph) compileCUDAPackages(org string) llb.State {
+	return g.preparePythonBase(llb.Image(fmt.Sprintf(
+		"docker.io/%s/%s-cudnn%s-devel-%s",
+		org, *g.CUDA, *g.CUDNN, g.OS)))
 }
 
 func (g Graph) compileSystemPackages(root llb.State) llb.State {
 	if len(g.SystemPackages) == 0 {
+		logrus.Debug("skip the apt since system package is not specified")
 		return root
 	}
 
@@ -143,6 +144,39 @@ func (g *Graph) compileExtraSource(root llb.State) (llb.State, error) {
 	return llb.Merge(inputs, llb.WithCustomName("[internal] build source layers")), nil
 }
 
+func (g *Graph) preparePythonBase(root llb.State) llb.State {
+	for _, env := range types.BaseEnvironment {
+		root = root.AddEnv(env.Name, env.Value)
+	}
+
+	// envd-sshd
+	sshd := root.File(llb.Copy(
+		llb.Image(types.EnvdSshdImage), "/usr/bin/envd-sshd", "/var/envd/bin/envd-sshd",
+		&llb.CopyInfo{CreateDestPath: true}),
+		llb.WithCustomName(fmt.Sprintf("[internal] add envd-sshd from %s", types.EnvdSshdImage)))
+
+	// apt packages
+	var sb strings.Builder
+	sb.WriteString("apt-get update && apt-get install -y apt-utils && ")
+	sb.WriteString("apt-get install -y --no-install-recommends --no-install-suggests --fix-missing ")
+	sb.WriteString(strings.Join(types.BaseAptPackage, " "))
+	sb.WriteString("&& rm -rf /var/lib/apt/lists/* ")
+	// shell prompt
+	sb.WriteString("&& curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y")
+
+	cacheDir := "/var/cache/apt"
+	cacheLibDir := "/var/lib/apt"
+
+	run := sshd.Run(llb.Shlex(fmt.Sprintf("bash -c \"%s\"", sb.String())),
+		llb.WithCustomName("[internal] install system packages"))
+	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()
+}
+
 func (g *Graph) compileBase() (llb.State, error) {
 	logger := logrus.WithFields(logrus.Fields{
 		"os":       g.OS,
@@ -173,16 +207,15 @@ func (g *Graph) compileBase() (llb.State, error) {
 				g.uid = 1001
 			}
 		case "python":
-			base = llb.Image(fmt.Sprintf(
-				"docker.io/%s/python:3.9-ubuntu20.04-envd-%s", org, v))
+			// TODO(keming) use user input `base(os="")`
+			base = g.preparePythonBase(llb.Image(types.PythonBaseImage))
 		case "julia":
 			base = llb.Image(fmt.Sprintf(
 				"docker.io/%s/julia:1.8rc1-ubuntu20.04-envd-%s", org, v))
 		}
 	} else {
-		base = g.compileCUDAPackages(org, v)
+		base = g.compileCUDAPackages("nvidia/cuda")
 	}
-	var res llb.ExecState
 
 	// Install conda first.
 	condaStage, err := g.installConda(base)
@@ -191,6 +224,7 @@ func (g *Graph) compileBase() (llb.State, error) {
 	}
 
 	// TODO(gaocegege): Refactor user to a separate stage.
+	var res llb.ExecState
 	if g.uid == 0 {
 		res = condaStage.
 			Run(llb.Shlex(fmt.Sprintf("groupadd -g %d envd", 1001)),

diff --git a/pkg/types/envd.go b/pkg/types/envd.go
@@ -16,9 +16,12 @@ package types
 
 import (
 	"encoding/json"
+	"fmt"
 
 	"github.com/docker/docker/api/types"
 	"github.com/moby/buildkit/util/system"
+
+	"github.com/tensorchord/envd/pkg/version"
 )
 
 // DefaultPathEnvUnix is unix style list of directories to search for
@@ -31,6 +34,46 @@ const DefaultPathEnvUnix = "/opt/conda/envs/envd/bin:/opt/conda/bin:/usr/local/j
 // ';' character .
 const DefaultPathEnvWindows = system.DefaultPathEnvWindows
 
+const PythonBaseImage = "ubuntu:20.04"
+
+var EnvdSshdImage = fmt.Sprintf(
+	"tensorchord/envd-sshd-from-scratch:%s",
+	version.GetVersionForImageTag())
+
+var BaseEnvironment = []struct {
+	Name  string
+	Value string
+}{
+	{"DEBIAN_FRONTEND", "noninteractive"},
+	{"PATH", DefaultPathEnvUnix},
+	{"LANG", "C.UTF-8"},
+	{"LC_ALL", "C.UTF-8"},
+}
+var BaseAptPackage = []string{
+	"bash-static",
+	"libtinfo5",
+	"libncursesw5",
+	// conda dependencies
+	"bzip2",
+	"ca-certificates",
+	"libglib2.0-0",
+	"libsm6",
+	"libxext6",
+	"libxrender1",
+	"mercurial",
+	"procps",
+	"subversion",
+	"wget",
+	// envd dependencies
+	"curl",
+	"openssh-client",
+	"git",
+	"tini",
+	"sudo",
+	"vim",
+	"zsh",
+}
+
 type EnvdImage struct {
 	types.ImageSummary