From 728a421840ba10c6ba78c836a22dc7af140d370d Mon Sep 17 00:00:00 2001 From: Keming Date: Fri, 29 Jul 2022 11:20:07 +0800 Subject: [PATCH] feat: add starship as the prompt manager (#681) * feat: add starship as the prompt manager Signed-off-by: Keming * rm llb.merge Signed-off-by: Keming * fix py version Signed-off-by: Keming * fix conflict Signed-off-by: Keming * chore: rm zsh theme func, parameterize version Signed-off-by: Keming * fix lint Signed-off-by: Keming --- Makefile | 5 ++- base-images/build.sh | 28 +++++++------ .../julia1.8rc1-ubuntu20.04.Dockerfile | 6 ++- ...python3.9-ubuntu20.04-cuda11.6.Dockerfile} | 6 ++- ...rfile => python3.9-ubuntu20.04.Dockerfile} | 8 +++- base-images/r4.2.Dockerfile | 6 ++- pkg/lang/ir/compile.go | 3 +- pkg/lang/ir/consts.go | 4 ++ pkg/lang/ir/shell.go | 41 ++++++++++++++++++- pkg/lang/ir/system.go | 4 +- pkg/shell/zsh.go | 27 ------------ pkg/shell/zshrc | 2 +- 12 files changed, 89 insertions(+), 51 deletions(-) rename base-images/{python3.8-ubuntu20.04-cuda11.6.Dockerfile => python3.9-ubuntu20.04-cuda11.6.Dockerfile} (96%) rename base-images/{python3.8-ubuntu20.04.Dockerfile => python3.9-ubuntu20.04.Dockerfile} (92%) diff --git a/Makefile b/Makefile index a25af72f8..ea1a14f69 100644 --- 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 index ba85c73e9..a2dfcad28 100755 --- 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,22 +19,22 @@ 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 \ @@ -38,15 +42,15 @@ docker buildx build \ --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 index ec0b7dbeb..d44f466b5 100644 --- a/base-images/julia1.8rc1-ubuntu20.04.Dockerfile +++ b/base-images/julia1.8rc1-ubuntu20.04.Dockerfile @@ -14,6 +14,8 @@ LABEL maintainer "envd-maintainers " 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/base-images/python3.8-ubuntu20.04-cuda11.6.Dockerfile b/base-images/python3.9-ubuntu20.04-cuda11.6.Dockerfile similarity index 96% rename from base-images/python3.8-ubuntu20.04-cuda11.6.Dockerfile rename to base-images/python3.9-ubuntu20.04-cuda11.6.Dockerfile index a3cf5a3aa..03e7b7800 100644 --- a/base-images/python3.8-ubuntu20.04-cuda11.6.Dockerfile +++ b/base-images/python3.9-ubuntu20.04-cuda11.6.Dockerfile @@ -53,6 +53,8 @@ ARG TARGETARCH LABEL maintainer "envd-maintainers " 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 b/base-images/python3.9-ubuntu20.04.Dockerfile similarity index 92% rename from base-images/python3.8-ubuntu20.04.Dockerfile rename to base-images/python3.9-ubuntu20.04.Dockerfile index c28105612..b9e4a283e 100644 --- a/base-images/python3.8-ubuntu20.04.Dockerfile +++ b/base-images/python3.9-ubuntu20.04.Dockerfile @@ -16,9 +16,11 @@ LABEL maintainer "envd-maintainers " 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 @@ -26,7 +28,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 # 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 index 87b8c9c00..16d004cee 100644 --- a/base-images/r4.2.Dockerfile +++ b/base-images/r4.2.Dockerfile @@ -16,6 +16,8 @@ LABEL maintainer "envd-maintainers " 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 index 0167537e9..6f03b79f8 100644 --- 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 index b971a1ad5..989a4c29c 100644 --- 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 index b32d9f011..fe18bd5a6 100644 --- a/pkg/lang/ir/shell.go +++ b/pkg/lang/ir/shell.go @@ -25,6 +25,24 @@ 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) @@ -32,6 +50,25 @@ func (g *Graph) compileShell(root llb.State) (llb.State, error) { 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 index ab34a8585..ad56ab9f3 100644 --- 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 index 82ca5d96f..87a3a10e9 100644 --- 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,10 +136,6 @@ 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") } @@ -148,28 +143,6 @@ func (m generalManager) DownloadOrCache() (bool, error) { 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 index 476dc491c..6082f43ff 100644 --- 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