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

* change python dockerfile to llb

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* change to docker hub

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* copy envd-sshd

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* add llb log

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* delete outdated comments

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* cuda version

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* del python dockerfile

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* fix test

Signed-off-by: Keming <kemingyang@tensorchord.ai>

Signed-off-by: Keming <kemingyang@tensorchord.ai>

.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

.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

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} \

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

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")

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
  }

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)

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
 }
 

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

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)),

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