feat: add starship as the prompt manager (#681)

* feat: add starship as the prompt manager Signed-off-by: Keming <[email protected]> * rm llb.merge Signed-off-by: Keming <[email protected]> * fix py version Signed-off-by: Keming <[email protected]> * fix conflict Signed-off-by: Keming <[email protected]> * chore: rm zsh theme func, parameterize version Signed-off-by: Keming <[email protected]> * fix lint Signed-off-by: Keming <[email protected]>
tensorchord · Jul 29, 2022 · 728a421 · 728a421
1 parent f82c54d
commit 728a421
Show file tree

Hide file tree

Showing 12 changed files with 89 additions and 51 deletions.
diff --git a/Makefile b/Makefile
@@ -95,7 +95,7 @@ export GOFLAGS ?= -count=1
 #
 
 # All targets.
-.PHONY: help lint test build container push addlicense debug debug-local build-local generate clean test-local addlicense-install mockgen-install pypi-build
+.PHONY: help lint test build dev container push addlicense debug debug-local build-local generate clean test-local addlicense-install mockgen-install pypi-build
 
 .DEFAULT_GOAL:=build
 
@@ -129,6 +129,9 @@ build-local:
 pypi-build: clean
 	@python setup.py sdist bdist_wheel
 
+dev: pypi-build  ## install envd command for local debug
+	@pip install --force-reinstall dist/*.whl
+
 generate: mockgen-install  ## Generate mocks
 	@mockgen -source pkg/buildkitd/buildkitd.go -destination pkg/buildkitd/mock/mock.go -package mock
 	@mockgen -source pkg/lang/frontend/starlark/interpreter.go -destination pkg/lang/frontend/starlark/mock/mock.go -package mock

diff --git a/base-images/build.sh b/base-images/build.sh
@@ -4,6 +4,10 @@ ROOT_DIR=`dirname $0`
 
 ENVD_VERSION="${ENVD_VERSION:-0.0.1}"
 DOCKER_HUB_ORG="${DOCKER_HUB_ORG:-tensorchord}"
+PYTHON_VERSION="${PYTHON_VERSION:-3.9}"
+ENVD_OS="${ENVD_OS:-ubuntu20.04}"
+JULIA_VERSION="${JULIA_VERSION:-1.8rc1}"
+RLANG_VERSION="${RLANG_VERSION:-4.2}"
 
 
 cd ${ROOT_DIR}
@@ -15,38 +19,38 @@ 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 IMAGE_NAME=docker.io/nvidia/cuda \
+docker buildx build \
     --build-arg ENVD_VERSION=${ENVD_VERSION} \
     --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-ssh-from-scratch \
     --build-arg HTTP_PROXY=${HTTP_PROXY} \
-     --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
+    --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
     --pull --push --platform linux/x86_64,linux/arm64 \
-    -t ${DOCKER_HUB_ORG}/python:3.8-ubuntu20.04-cuda11.6-cudnn8 \
-    -f python3.8-ubuntu20.04-cuda11.6.Dockerfile .
-docker buildx build \
+    -t ${DOCKER_HUB_ORG}/python:${PYTHON_VERSION}-${ENVD_OS} \
+    -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-ssh-from-scratch \
     --build-arg HTTP_PROXY=${HTTP_PROXY} \
-     --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
+    --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
     --pull --push --platform linux/x86_64,linux/arm64 \
-    -t ${DOCKER_HUB_ORG}/python:3.8-ubuntu20.04 \
-    -f python3.8-ubuntu20.04.Dockerfile .
+    -t ${DOCKER_HUB_ORG}/python:${PYTHON_VERSION}-${ENVD_OS}-cuda11.6-cudnn8 \
+    -f python${PYTHON_VERSION}-${ENVD_OS}-cuda11.6.Dockerfile .
 
 # TODO(gaocegege): Support linux/arm64
 docker buildx build \
     --build-arg ENVD_VERSION=${ENVD_VERSION} \
     --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-ssh-from-scratch \
     --build-arg HTTP_PROXY=${HTTP_PROXY} \
     --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
-    -t ${DOCKER_HUB_ORG}/r-base:4.2 \
+    -t ${DOCKER_HUB_ORG}/r-base:${RLANG_VERSION} \
     --pull --push --platform linux/x86_64 \
-    -f r4.2.Dockerfile .
+    -f r${RLANG_VERSION}.Dockerfile .
 docker buildx build \
     --build-arg ENVD_VERSION=${ENVD_VERSION} \
     --build-arg ENVD_SSH_IMAGE=ghcr.io/tensorchord/envd-ssh-from-scratch \
     --build-arg HTTP_PROXY=${HTTP_PROXY} \
     --build-arg HTTPS_PROXY=${HTTPS_PROXY} \
-    -t ${DOCKER_HUB_ORG}/julia:1.8rc1-ubuntu20.04 \
+    -t ${DOCKER_HUB_ORG}/julia:${JULIA_VERSION}-${ENVD_OS} \
     --pull --push --platform linux/x86_64,linux/arm64 \
-    -f julia1.8rc1-ubuntu20.04.Dockerfile .
+    -f julia${JULIA_VERSION}-${ENVD_OS}.Dockerfile .
 cd - > /dev/null
diff --git a/base-images/julia1.8rc1-ubuntu20.04.Dockerfile b/base-images/julia1.8rc1-ubuntu20.04.Dockerfile
@@ -14,6 +14,8 @@ LABEL maintainer "envd-maintainers <[email protected]>"
 
 ENV DEBIAN_FRONTEND noninteractive
 ENV PATH="/usr/bin:${PATH}"
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends --no-install-suggests --fix-missing \
@@ -23,7 +25,9 @@ RUN apt-get update && \
     procps subversion wget \
     # envd dependencies
     python3 curl openssh-client git tini sudo python3-pip zsh vim \
-    && rm -rf /var/lib/apt/lists/*
+    && rm -rf /var/lib/apt/lists/* \
+    # prompt
+    && curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y
 
 COPY --from=envd /usr/bin/envd-ssh /var/envd/bin/envd-ssh
 

diff --git a/...python3.8-ubuntu20.04-cuda11.6.Dockerfile → ...python3.9-ubuntu20.04-cuda11.6.Dockerfile b/...python3.8-ubuntu20.04-cuda11.6.Dockerfile → ...python3.9-ubuntu20.04-cuda11.6.Dockerfile
@@ -53,6 +53,8 @@ ARG TARGETARCH
 LABEL maintainer "envd-maintainers <[email protected]>"
 
 ENV DEBIAN_FRONTEND noninteractive
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends --no-install-suggests --fix-missing bash-static \
@@ -77,7 +79,9 @@ RUN apt-get update && \
     procps subversion wget \
     # envd dependencies
     curl openssh-client git tini sudo zsh vim \
-    && rm -rf /var/lib/apt/lists/*
+    && rm -rf /var/lib/apt/lists/* \
+    # prompt
+    && curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y
 
 # Keep apt from auto upgrading the cublas and nccl packages. See https://gitlab.com/nvidia/container-images/cuda/-/issues/88
 RUN apt-mark hold ${NV_LIBCUBLAS_DEV_PACKAGE_NAME} ${NV_LIBNCCL_DEV_PACKAGE_NAME}

diff --git a/base-images/python3.8-ubuntu20.04.Dockerfile → base-images/python3.9-ubuntu20.04.Dockerfile b/base-images/python3.8-ubuntu20.04.Dockerfile → base-images/python3.9-ubuntu20.04.Dockerfile
@@ -16,17 +16,21 @@ LABEL maintainer "envd-maintainers <[email protected]>"
 
 ENV DEBIAN_FRONTEND noninteractive
 ENV PATH="/usr/bin:${PATH}"
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
 
 RUN apt-get update && \
-    apt-get install apt-utils && \
+    apt-get install -y apt-utils && \
     apt-get install -y --no-install-recommends --no-install-suggests --fix-missing \
     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 zsh vim \
-    && rm -rf /var/lib/apt/lists/*
+    && rm -rf /var/lib/apt/lists/* \
+    # prompt
+    && curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y
 
 # Leave these args here to better use the Docker build cache
 ARG CONDA_VERSION=py39_4.11.0

diff --git a/base-images/r4.2.Dockerfile b/base-images/r4.2.Dockerfile
@@ -16,6 +16,8 @@ LABEL maintainer "envd-maintainers <[email protected]>"
 
 ENV DEBIAN_FRONTEND noninteractive
 ENV PATH="/usr/bin:${PATH}"
+ENV LANG C.UTF-8
+ENV LC_ALL C.UTF-8
 
 RUN apt-get update && \
     apt-get install -y --no-install-recommends --no-install-suggests --fix-missing \
@@ -25,7 +27,9 @@ RUN apt-get update && \
     libpq5 psmisc procps python-setuptools pwgen lsb-release \
     # envd dependencies
     python3 curl openssh-client git tini sudo zsh vim \
-    && rm -rf /var/lib/apt/lists/*
+    && rm -rf /var/lib/apt/lists/* \
+    # prompt
+    && curl --proto '=https' --tlsv1.2 -sSf https://starship.rs/install.sh | sh -s -- -y
 
 RUN set -x && \
     UNAME_M="$(uname -m)" && \

diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -223,7 +223,8 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) {
 		}
 	}
 
-	copy := g.compileCopy(merged)
+	prompt := g.compilePrompt(merged)
+	copy := g.compileCopy(prompt)
 	// TODO(gaocegege): Support order-based exec.
 	run := g.compileRun(copy)
 	finalStage, err := g.compileGit(run)

diff --git a/pkg/lang/ir/consts.go b/pkg/lang/ir/consts.go
@@ -19,6 +19,10 @@ const (
 	languageDefault   = "python"
 	pypiIndexModeAuto = "auto"
 
+	// used inside the container
+	defaultConfigDir   = "/home/envd/.config"
+	starshipConfigPath = "/home/envd/.config/starship.toml"
+
 	aptSourceFilePath = "/etc/apt/sources.list"
 	pypiIndexFilePath = "/etc/pip.conf"
 

diff --git a/pkg/lang/ir/shell.go b/pkg/lang/ir/shell.go
@@ -25,13 +25,50 @@ import (
 	"github.com/tensorchord/envd/pkg/shell"
 )
 
+const (
+	starshipConfig = `
+[container]
+format = "[$symbol \\[envd\\]]($style)"
+
+[sudo]
+disabled = false
+symbol = "sudo "
+
+[python]
+symbol = "Py "
+
+[status]
+map_symbol = true
+disabled = false
+`
+)
+
 func (g *Graph) compileShell(root llb.State) (llb.State, error) {
 	if g.Shell == shellZSH {
 		return g.compileZSH(root)
 	}
 	return root, nil
 }
 
+func (g *Graph) compilePrompt(root llb.State) llb.State {
+	// starship config
+	config := root.
+		File(llb.Mkdir(defaultConfigDir, 0755, llb.WithParents(true)),
+			llb.WithCustomName("[internal] creating config dir")).
+		File(llb.Mkfile(starshipConfigPath, 0644, []byte(starshipConfig), llb.WithUIDGID(g.uid, g.gid)),
+			llb.WithCustomName("[internal] setting prompt config"))
+
+	run := config.Run(llb.Shlex(`bash -c 'echo "eval \"\$(starship init bash)\"" >> /home/envd/.bashrc'`),
+		llb.WithCustomName("[internal] setting prompt config")).Root()
+
+	if g.Shell == shellZSH {
+		run = run.Run(
+			llb.Shlex(`bash -c 'echo "eval \"\$(starship init zsh)\"" >> /home/envd/.zshrc'`),
+			llb.WithCustomName("[internal] setting prompt config")).Root()
+	}
+	return run
+}
+
 func (g Graph) compileZSH(root llb.State) (llb.State, error) {
 	installPath := "/home/envd/install.sh"
 	zshrcPath := "/home/envd/.zshrc"
@@ -48,9 +85,9 @@ func (g Graph) compileZSH(root llb.State) (llb.State, error) {
 			&llb.CopyInfo{CreateDestPath: true}, llb.WithUIDGID(g.uid, g.gid))).
 		File(llb.Mkfile(installPath,
 			0644, []byte(m.InstallScript()), llb.WithUIDGID(g.uid, g.gid)))
-	run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)),
+	zshrc := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)),
 		llb.WithCustomName("install oh-my-zsh")).
 		File(llb.Mkfile(zshrcPath,
 			0644, []byte(m.ZSHRC()), llb.WithUIDGID(g.uid, g.gid)))
-	return run, nil
+	return zshrc, nil
 }
diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -74,7 +74,7 @@ func (g Graph) compileCopy(root llb.State) llb.State {
 }
 
 func (g *Graph) compileCUDAPackages() llb.State {
-	root := llb.Image(fmt.Sprintf("docker.io/tensorchord/python:3.8-%s-cuda%s-cudnn%s", g.OS, *g.CUDA, *g.CUDNN))
+	root := llb.Image(fmt.Sprintf("docker.io/tensorchord/python:3.9-%s-cuda%s-cudnn%s", g.OS, *g.CUDA, *g.CUDNN))
 	return root
 }
 
@@ -132,7 +132,7 @@ func (g *Graph) compileBase() llb.State {
 				g.uid = 1001
 			}
 		case "python":
-			base = llb.Image("docker.io/tensorchord/python:3.8-ubuntu20.04")
+			base = llb.Image("docker.io/tensorchord/python:3.9-ubuntu20.04")
 		case "julia":
 			base = llb.Image("docker.io/tensorchord/julia:1.8rc1-ubuntu20.04")
 		}

diff --git a/pkg/shell/zsh.go b/pkg/shell/zsh.go
@@ -16,7 +16,6 @@ package shell
 
 import (
 	_ "embed"
-	"os"
 	"path/filepath"
 
 	"github.com/cockroachdb/errors"
@@ -137,39 +136,13 @@ func (m generalManager) DownloadOrCache() (bool, error) {
 		return false, errors.Wrap(err, "failed to checkout master")
 	}
 
-	if err := m.createEnvdPromptTheme(); err != nil {
-		return false, errors.Wrap(err, "failed to create envd.zsh-theme")
-	}
-
 	if err := home.GetManager().MarkCache(cacheKey, true); err != nil {
 		return false, errors.Wrap(err, "failed to update cache status")
 	}
 	l.Debug("oh-my-zsh is downloaded")
 	return false, nil
 }
 
-// Refer to https://github.com/tensorchord/envd/issues/183#issuecomment-1148172564
-func (m generalManager) createEnvdPromptTheme() error {
-	path := "themes/envd.zsh-theme"
-	content := `
-PROMPT="(envd) %(?:%{$fg_bold[green]%}%{%G➜%} :%{$fg_bold[red]%}%{%G➜%} )"
-PROMPT+=' %{$fg[cyan]%}%c%{$reset_color%} $(git_prompt_info)'
-
-ZSH_THEME_GIT_PROMPT_PREFIX="%{$fg_bold[blue]%}git:(%{$fg[red]%}"
-ZSH_THEME_GIT_PROMPT_SUFFIX="%{$reset_color%} "
-ZSH_THEME_GIT_PROMPT_DIRTY="%{$fg[blue]%}) %{$fg[yellow]%}%{%G✗%}"
-ZSH_THEME_GIT_PROMPT_CLEAN="%{$fg[blue]%})"
-`
-	absolutePath := filepath.Join(m.OHMyZSHDir(), path)
-	logger := logrus.WithField("path", absolutePath)
-	logger.Debug("creating envd.zsh-theme")
-	if err := os.WriteFile(
-		absolutePath, []byte(content), 0644); err != nil {
-		return err
-	}
-	return nil
-}
-
 func (m generalManager) OHMyZSHDir() string {
 	return filepath.Join(home.GetManager().CacheDir(), "oh-my-zsh")
 }
diff --git a/pkg/shell/zshrc b/pkg/shell/zshrc
@@ -8,7 +8,7 @@ export ZSH="$HOME/.oh-my-zsh"
 # load a random theme each time oh-my-zsh is loaded, in which case,
 # to know which specific one was loaded, run: echo $RANDOM_THEME
 # See https://github.com/ohmyzsh/ohmyzsh/wiki/Themes
-ZSH_THEME="envd"
+# ZSH_THEME="envd"
 
 # Set list of themes to pick from when loading at random
 # Setting this variable when ZSH_THEME=random will cause zsh to load