From 3b7195079892e5e3f33e20642de7995fa9c3276a Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Thu, 19 May 2022 14:09:33 +0800 Subject: [PATCH] feat: Use UID 1000 to build (#173) * feat: Use UID 1000 to build Signed-off-by: Ce Gao * fix: Fix test cases Signed-off-by: Ce Gao --- pkg/docker/docker.go | 1 + pkg/editor/jupyter/util.go | 2 +- pkg/editor/jupyter/util_test.go | 6 +- pkg/lang/ir/compile.go | 112 +++++++++++ pkg/lang/ir/consts.go | 2 + pkg/lang/ir/editor.go | 40 ++++ pkg/lang/ir/graph.go | 328 -------------------------------- pkg/lang/ir/python.go | 44 +++++ pkg/lang/ir/shell.go | 36 ++++ pkg/lang/ir/system.go | 131 +++++++++++++ 10 files changed, 369 insertions(+), 333 deletions(-) create mode 100644 pkg/lang/ir/compile.go create mode 100644 pkg/lang/ir/editor.go delete mode 100644 pkg/lang/ir/graph.go create mode 100644 pkg/lang/ir/python.go create mode 100644 pkg/lang/ir/shell.go create mode 100644 pkg/lang/ir/system.go diff --git a/pkg/docker/docker.go b/pkg/docker/docker.go index c9758c97f..6d1d57fb5 100644 --- a/pkg/docker/docker.go +++ b/pkg/docker/docker.go @@ -173,6 +173,7 @@ func (c generalClient) StartEnvd(ctx context.Context, tag, name, buildContext st }) config := &container.Config{ Image: tag, + User: "envd", Entrypoint: []string{ "tini", "--", diff --git a/pkg/editor/jupyter/util.go b/pkg/editor/jupyter/util.go index 8330442f1..78fa57ecd 100644 --- a/pkg/editor/jupyter/util.go +++ b/pkg/editor/jupyter/util.go @@ -26,7 +26,7 @@ func GenerateCommand(g ir.Graph, notebookDir string) []string { } cmd := []string{ - "jupyter", "notebook", "--allow-root", + "python3", "-m", "notebook", "--ip", "0.0.0.0", "--notebook-dir", notebookDir, } if g.JupyterConfig.Password != "" { diff --git a/pkg/editor/jupyter/util_test.go b/pkg/editor/jupyter/util_test.go index 6bc523857..9f8821c11 100644 --- a/pkg/editor/jupyter/util_test.go +++ b/pkg/editor/jupyter/util_test.go @@ -21,8 +21,7 @@ func TestGenerateCommand(t *testing.T) { }, dir: "test", expected: []string{ - "jupyter", "notebook", "--allow-root", - "--ip", "0.0.0.0", "--notebook-dir", "test", + "python3", "-m", "notebook", "--ip", "0.0.0.0", "--notebook-dir", "test", "--NotebookApp.password", "''", "--NotebookApp.token", "''", "--port", "8888", }, @@ -36,8 +35,7 @@ func TestGenerateCommand(t *testing.T) { }, dir: "test", expected: []string{ - "jupyter", "notebook", "--allow-root", - "--ip", "0.0.0.0", "--notebook-dir", "test", + "python3", "-m", "notebook", "--ip", "0.0.0.0", "--notebook-dir", "test", "--NotebookApp.password", "test", "--NotebookApp.token", "''", "--port", "8888", }, diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go new file mode 100644 index 000000000..33651398b --- /dev/null +++ b/pkg/lang/ir/compile.go @@ -0,0 +1,112 @@ +// 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 ( + "context" + "os" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + + "github.com/tensorchord/envd/pkg/progress/compileui" +) + +func NewGraph() *Graph { + return &Graph{ + OS: osDefault, + Language: languageDefault, + CUDA: nil, + CUDNN: nil, + BuiltinSystemPackages: []string{ + // TODO(gaocegege): Move them into the base image. + "curl", + "openssh-client", + "git", + "sudo", + "tini", + }, + + PyPIPackages: []string{}, + SystemPackages: []string{}, + Exec: []string{}, + Shell: shellBASH, + } +} + +var DefaultGraph = NewGraph() + +func GPUEnabled() bool { + return DefaultGraph.CUDA != nil +} + +func Compile(ctx context.Context) (*llb.Definition, error) { + w, err := compileui.New(ctx, os.Stdout, "auto") + if err != nil { + return nil, errors.Wrap(err, "failed to create compileui") + } + DefaultGraph.Writer = w + state, err := DefaultGraph.Compile() + if err != nil { + return nil, err + } + // TODO(gaocegege): Support multi platform. + def, err := state.Marshal(ctx, llb.LinuxAmd64) + if err != nil { + return nil, err + } + return def, nil +} + +func (g Graph) Compile() (llb.State, error) { + // TODO(gaocegege): Support more OS and langs. + base := g.compileBase() + aptStage := g.compileUbuntuAPT(base) + pypiMirrorStage := g.compilePyPIMirror(aptStage) + + g.compileJupyter() + // TODO(gaocegege): Make apt update a seperate stage to + // parallel system and user-defined package installation. + builtinSystemStage := g.compileBuiltinSystemPackages(pypiMirrorStage) + shellStage, err := g.compileShell(builtinSystemStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to compile shell") + } + diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell")) + pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage), llb.WithCustomName("install PyPI packages")) + systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages")) + sshStage := g.copyEnvdSSHServer() + + vscodeStage, err := g.compileVSCode() + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get vscode plugins") + } + + var merged llb.State + if vscodeStage != nil { + merged = llb.Merge([]llb.State{ + builtinSystemStage, systemStage, pypiStage, sshStage, *vscodeStage, diffShellStage, + }, llb.WithCustomName("merging all components into one")) + } else { + merged = llb.Merge([]llb.State{ + builtinSystemStage, systemStage, pypiStage, sshStage, diffShellStage, + }, llb.WithCustomName("merging all components into one")) + } + + // TODO(gaocegege): Support order-based exec. + run := g.compileRun(merged) + g.Writer.Finish() + return run, nil +} diff --git a/pkg/lang/ir/consts.go b/pkg/lang/ir/consts.go index b83404b44..15a541958 100644 --- a/pkg/lang/ir/consts.go +++ b/pkg/lang/ir/consts.go @@ -25,4 +25,6 @@ const ( pypiConfigTemplate = ` [global] index-url=%s` + + defaultUID = "1000" ) diff --git a/pkg/lang/ir/editor.go b/pkg/lang/ir/editor.go new file mode 100644 index 000000000..6179536f8 --- /dev/null +++ b/pkg/lang/ir/editor.go @@ -0,0 +1,40 @@ +package ir + +import ( + "github.com/moby/buildkit/client/llb" + "github.com/tensorchord/envd/pkg/editor/vscode" + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/progress/compileui" +) + +func (g Graph) compileVSCode() (*llb.State, error) { + if len(g.VSCodePlugins) == 0 { + return nil, nil + } + inputs := []llb.State{} + for _, p := range g.VSCodePlugins { + vscodeClient := vscode.NewClient() + g.Writer.LogVSCodePlugin(p, compileui.ActionStart, false) + if cached, err := vscodeClient.DownloadOrCache(p); err != nil { + return nil, err + } else { + g.Writer.LogVSCodePlugin(p, compileui.ActionEnd, cached) + } + ext := llb.Scratch().File(llb.Copy(llb.Local(flag.FlagCacheDir), + vscodeClient.PluginPath(p), + "/home/envd/.vscode-server/extensions/"+p.String(), + &llb.CopyInfo{ + CreateDestPath: true, + }, llb.WithUser(defaultUID)), + llb.WithCustomNamef("install vscode plugin %s", p.String())) + inputs = append(inputs, ext) + } + layer := llb.Merge(inputs, llb.WithCustomName("merging plugins for vscode")) + return &layer, nil +} + +func (g *Graph) compileJupyter() { + if g.JupyterConfig != nil { + g.PyPIPackages = append(g.PyPIPackages, "jupyter") + } +} diff --git a/pkg/lang/ir/graph.go b/pkg/lang/ir/graph.go deleted file mode 100644 index 7aa0f8b8b..000000000 --- a/pkg/lang/ir/graph.go +++ /dev/null @@ -1,328 +0,0 @@ -// 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 ( - "context" - "fmt" - "os" - "path/filepath" - "strings" - - "github.com/cockroachdb/errors" - "github.com/moby/buildkit/client/llb" - "github.com/sirupsen/logrus" - "github.com/spf13/viper" - - "github.com/tensorchord/envd/pkg/editor/vscode" - "github.com/tensorchord/envd/pkg/flag" - "github.com/tensorchord/envd/pkg/progress/compileui" - "github.com/tensorchord/envd/pkg/shell" -) - -func NewGraph() *Graph { - return &Graph{ - OS: osDefault, - Language: languageDefault, - CUDA: nil, - CUDNN: nil, - BuiltinSystemPackages: []string{ - // TODO(gaocegege): Move them into the base image. - "curl", - "openssh-client", - "git", - "sudo", - "tini", - }, - - PyPIPackages: []string{}, - SystemPackages: []string{}, - Exec: []string{}, - Shell: shellBASH, - } -} - -var DefaultGraph = NewGraph() - -func GPUEnabled() bool { - return DefaultGraph.CUDA != nil -} - -func Compile(ctx context.Context) (*llb.Definition, error) { - w, err := compileui.New(ctx, os.Stdout, "auto") - if err != nil { - return nil, errors.Wrap(err, "failed to create compileui") - } - DefaultGraph.Writer = w - state, err := DefaultGraph.Compile() - if err != nil { - return nil, err - } - // TODO(gaocegege): Support multi platform. - def, err := state.Marshal(ctx, llb.LinuxAmd64) - if err != nil { - return nil, err - } - return def, nil -} - -func (g Graph) Compile() (llb.State, error) { - // TODO(gaocegege): Support more OS and langs. - base := g.compileBase() - aptStage := g.compileUbuntuAPT(base) - pypiMirrorStage := g.compilePyPIMirror(aptStage) - - g.compileJupyter() - // TODO(gaocegege): Make apt update a seperate stage to - // parallel system and user-defined package installation. - builtinSystemStage := g.compileBuiltinSystemPackages(pypiMirrorStage) - shellStage, err := g.compileShell(builtinSystemStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to compile shell") - } - diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell")) - pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage), llb.WithCustomName("install PyPI packages")) - systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages")) - sshStage := g.copyEnvdSSHServer() - - vscodeStage, err := g.compileVSCode() - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to get vscode plugins") - } - - var merged llb.State - if vscodeStage != nil { - merged = llb.Merge([]llb.State{ - builtinSystemStage, systemStage, pypiStage, sshStage, *vscodeStage, diffShellStage, - }, llb.WithCustomName("merging all components into one")) - } else { - merged = llb.Merge([]llb.State{ - builtinSystemStage, systemStage, pypiStage, sshStage, diffShellStage, - }, llb.WithCustomName("merging all components into one")) - } - - // TODO(gaocegege): Support order-based exec. - run := g.compileRun(merged) - g.Writer.Finish() - return run, nil -} - -func (g *Graph) compileBase() llb.State { - var base llb.State - if g.CUDA == nil && g.CUDNN == nil { - base = llb.Image("docker.io/library/python:3.8") - } else { - base = g.compileCUDAPackages() - } - return base -} - -func (g *Graph) compileCUDAPackages() llb.State { - root := llb.Image( - fmt.Sprintf("nvidia/cuda:%s.0-cudnn%s-devel-%s", *g.CUDA, *g.CUDNN, g.OS)) - g.BuiltinSystemPackages = append(g.BuiltinSystemPackages, []string{ - g.Language, - fmt.Sprintf("%s-pip", g.Language), - }...) - return root -} - -func (g *Graph) compileShell(root llb.State) (llb.State, error) { - if g.Shell == shellZSH { - return g.compileZSH(root) - } - return root, nil -} - -func (g Graph) compilePyPIPackages(root llb.State) llb.State { - if len(g.PyPIPackages) == 0 { - return root - } - - // Compose the package install command. - var sb strings.Builder - // TODO(gaocegege): Support per-user config to keep the mirror. - sb.WriteString("pip install") - for _, pkg := range g.PyPIPackages { - sb.WriteString(fmt.Sprintf(" %s", pkg)) - } - - cacheDir := "/home/envd/.cache/pip" - cmd := sb.String() - run := root.Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s", - strings.Join(g.PyPIPackages, " "))) - run.AddMount(cacheDir, llb.Scratch(), - llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared)) - return run.Root() -} - -func (g Graph) compileBuiltinSystemPackages(root llb.State) llb.State { - // TODO(gaocegege): Refactor it to avoid shell configuration in built-in system packages. - // Do not need to install bash or sh since it is built-in - if g.Shell == shellZSH { - g.BuiltinSystemPackages = append(g.BuiltinSystemPackages, shellZSH) - } - - if len(g.BuiltinSystemPackages) == 0 { - return root - } - - // Compose the package install command. - var sb strings.Builder - sb.WriteString( - "sh -c \"apt-get update && apt-get install -y --no-install-recommends") - for _, pkg := range g.BuiltinSystemPackages { - sb.WriteString(fmt.Sprintf(" %s", pkg)) - } - sb.WriteString("\"") - - cacheDir := "/var/cache/apt" - cacheLibDir := "/var/lib/apt" - - run := root.Run(llb.Shlex(sb.String()), - llb.WithCustomNamef("(built-in packages) apt-get install %s", - strings.Join(g.BuiltinSystemPackages, " "))) - run.AddMount(cacheDir, llb.Scratch(), - llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared)) - run.AddMount(cacheLibDir, llb.Scratch(), - llb.AsPersistentCacheDir("/"+cacheLibDir, llb.CacheMountShared)) - - // TODO(gaocegege): Refactor user to a seperate stage. - run = run. - Run(llb.Shlex("groupadd -g 1000 envd"), llb.WithCustomName("create user group envd")). - Run(llb.Shlex("useradd -p \"\" -u 1000 -g envd -s /bin/sh -m envd"), llb.WithCustomName("create user envd")). - Run(llb.Shlex("adduser envd sudo"), llb.WithCustomName("add user envd to sudoers")) - return llb.User("envd")(run.Root()) -} - -func (g Graph) compileSystemPackages(root llb.State) llb.State { - if len(g.SystemPackages) == 0 { - return root - } - - // Compose the package install command. - var sb strings.Builder - sb.WriteString("sudo 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(sb.String()), - llb.WithCustomNamef("(user-defined packages) apt-get install %s", - strings.Join(g.SystemPackages, " "))) - run.AddMount(cacheDir, llb.Scratch(), - llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared)) - run.AddMount(cacheLibDir, llb.Scratch(), - llb.AsPersistentCacheDir("/"+cacheLibDir, llb.CacheMountShared)) - return run.Root() -} - -func (g Graph) copyEnvdSSHServer() llb.State { - // TODO(gaocegege): Remove global var ssh image. - run := llb.Image(viper.GetString(flag.FlagSSHImage)). - File(llb.Copy(llb.Image(viper.GetString(flag.FlagSSHImage)), - "usr/bin/envd-ssh", "/var/envd/bin/envd-ssh", - &llb.CopyInfo{CreateDestPath: true}), llb.WithCustomName("install envd-ssh")) - return run -} - -func (g Graph) compileVSCode() (*llb.State, error) { - if len(g.VSCodePlugins) == 0 { - return nil, nil - } - inputs := []llb.State{} - for _, p := range g.VSCodePlugins { - vscodeClient := vscode.NewClient() - g.Writer.LogVSCodePlugin(p, compileui.ActionStart, false) - if cached, err := vscodeClient.DownloadOrCache(p); err != nil { - return nil, err - } else { - g.Writer.LogVSCodePlugin(p, compileui.ActionEnd, cached) - } - ext := llb.Scratch().File(llb.Copy(llb.Local(flag.FlagCacheDir), - vscodeClient.PluginPath(p), - "/home/envd/.vscode-server/extensions/"+p.String(), - &llb.CopyInfo{CreateDestPath: true}), - llb.WithCustomNamef("install vscode plugin %s", p.String())) - inputs = append(inputs, ext) - } - layer := llb.Merge(inputs, llb.WithCustomName("merging plugins for vscode")) - return &layer, nil -} - -func (g *Graph) compileJupyter() { - if g.JupyterConfig != nil { - g.PyPIPackages = append(g.PyPIPackages, "jupyter") - } -} - -func (g Graph) compileUbuntuAPT(root llb.State) llb.State { - if g.UbuntuAPTSource != nil { - logrus.WithField("source", *g.UbuntuAPTSource).Debug("using custom APT source") - aptSource := llb.Scratch(). - File(llb.Mkdir(filepath.Dir(aptSourceFilePath), 0755, llb.WithParents(true)), llb.WithCustomName("create apt source dir")). - File(llb.Mkfile(aptSourceFilePath, 0644, []byte(*g.UbuntuAPTSource)), llb.WithCustomName("create apt source file")) - return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add apt source")) - } - return root -} - -func (g Graph) compilePyPIMirror(root llb.State) llb.State { - if g.PyPIMirror != nil { - logrus.WithField("mirror", *g.PyPIMirror).Debug("using custom PyPI mirror") - content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIMirror) - aptSource := llb.Scratch(). - File(llb.Mkdir(filepath.Dir(pypiMirrorFilePath), 0755, llb.WithParents(true))). - File(llb.Mkfile(pypiMirrorFilePath, 0644, []byte(content))) - return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add PyPI mirror")) - } - return root -} - -func (g Graph) compileZSH(root llb.State) (llb.State, error) { - installPath := "/home/envd/install.sh" - m := shell.NewManager() - g.Writer.LogZSH(compileui.ActionStart, false) - if cached, err := m.DownloadOrCache(); err != nil { - return llb.State{}, errors.Wrap(err, "failed to download oh-my-zsh") - } else { - g.Writer.LogZSH(compileui.ActionEnd, cached) - } - zshStage := root. - File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", "/home/envd/.oh-my-zsh", - &llb.CopyInfo{CreateDestPath: true})). - File(llb.Mkfile(installPath, 0644, []byte(m.InstallScript()))) - run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)), - llb.WithCustomName("install oh-my-zsh")) - return run.Root(), nil -} - -func (g Graph) compileRun(root llb.State) llb.State { - if len(g.Exec) == 0 { - return root - } else if len(g.Exec) == 1 { - return root.Run(llb.Shlex(g.Exec[0])).Root() - } - - run := root.Run(llb.Shlex(g.Exec[0])) - for _, c := range g.Exec[1:] { - run = run.Run(llb.Shlex(c)) - } - return run.Root() -} diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go new file mode 100644 index 000000000..6371c4bfe --- /dev/null +++ b/pkg/lang/ir/python.go @@ -0,0 +1,44 @@ +package ir + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/moby/buildkit/client/llb" + "github.com/sirupsen/logrus" +) + +func (g Graph) compilePyPIPackages(root llb.State) llb.State { + if len(g.PyPIPackages) == 0 { + return root + } + + // Compose the package install command. + var sb strings.Builder + // TODO(gaocegege): Support per-user config to keep the mirror. + sb.WriteString("pip install") + for _, pkg := range g.PyPIPackages { + sb.WriteString(fmt.Sprintf(" %s", pkg)) + } + + cacheDir := "/home/envd/.cache/pip" + cmd := sb.String() + run := root.Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s", + strings.Join(g.PyPIPackages, " "))) + run.AddMount(cacheDir, llb.Scratch(), + llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared)) + return run.Root() +} + +func (g Graph) compilePyPIMirror(root llb.State) llb.State { + if g.PyPIMirror != nil { + logrus.WithField("mirror", *g.PyPIMirror).Debug("using custom PyPI mirror") + content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIMirror) + aptSource := llb.Scratch(). + File(llb.Mkdir(filepath.Dir(pypiMirrorFilePath), 0755, llb.WithParents(true))). + File(llb.Mkfile(pypiMirrorFilePath, 0644, []byte(content))) + return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add PyPI mirror")) + } + return root +} diff --git a/pkg/lang/ir/shell.go b/pkg/lang/ir/shell.go new file mode 100644 index 000000000..810639513 --- /dev/null +++ b/pkg/lang/ir/shell.go @@ -0,0 +1,36 @@ +package ir + +import ( + "fmt" + + "github.com/cockroachdb/errors" + "github.com/moby/buildkit/client/llb" + "github.com/tensorchord/envd/pkg/flag" + "github.com/tensorchord/envd/pkg/progress/compileui" + "github.com/tensorchord/envd/pkg/shell" +) + +func (g *Graph) compileShell(root llb.State) (llb.State, error) { + if g.Shell == shellZSH { + return g.compileZSH(root) + } + return root, nil +} + +func (g Graph) compileZSH(root llb.State) (llb.State, error) { + installPath := "/home/envd/install.sh" + m := shell.NewManager() + g.Writer.LogZSH(compileui.ActionStart, false) + if cached, err := m.DownloadOrCache(); err != nil { + return llb.State{}, errors.Wrap(err, "failed to download oh-my-zsh") + } else { + g.Writer.LogZSH(compileui.ActionEnd, cached) + } + zshStage := root. + File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", "/home/envd/.oh-my-zsh", + &llb.CopyInfo{CreateDestPath: true}, llb.WithUser(defaultUID))). + File(llb.Mkfile(installPath, 0644, []byte(m.InstallScript()), llb.WithUser(defaultUID))) + run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)), + llb.WithCustomName("install oh-my-zsh")) + return run.Root(), nil +} diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go new file mode 100644 index 000000000..cffdb60c6 --- /dev/null +++ b/pkg/lang/ir/system.go @@ -0,0 +1,131 @@ +package ir + +import ( + "fmt" + "path/filepath" + "strings" + + "github.com/moby/buildkit/client/llb" + "github.com/sirupsen/logrus" + "github.com/spf13/viper" + "github.com/tensorchord/envd/pkg/flag" +) + +func (g Graph) compileUbuntuAPT(root llb.State) llb.State { + if g.UbuntuAPTSource != nil { + logrus.WithField("source", *g.UbuntuAPTSource).Debug("using custom APT source") + aptSource := llb.Scratch(). + File(llb.Mkdir(filepath.Dir(aptSourceFilePath), 0755, llb.WithParents(true)), llb.WithCustomName("create apt source dir")). + File(llb.Mkfile(aptSourceFilePath, 0644, []byte(*g.UbuntuAPTSource)), llb.WithCustomName("create apt source file")) + return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add apt source")) + } + return root +} + +func (g Graph) compileRun(root llb.State) llb.State { + if len(g.Exec) == 0 { + return root + } else if len(g.Exec) == 1 { + return root.Run(llb.Shlex(g.Exec[0])).Root() + } + + run := root.Run(llb.Shlex(g.Exec[0])) + for _, c := range g.Exec[1:] { + run = run.Run(llb.Shlex(c)) + } + return run.Root() +} + +func (g Graph) compileBuiltinSystemPackages(root llb.State) llb.State { + // TODO(gaocegege): Refactor it to avoid shell configuration in built-in system packages. + // Do not need to install bash or sh since it is built-in + if g.Shell == shellZSH { + g.BuiltinSystemPackages = append(g.BuiltinSystemPackages, shellZSH) + } + + if len(g.BuiltinSystemPackages) == 0 { + return root + } + + // Compose the package install command. + var sb strings.Builder + sb.WriteString( + "sh -c \"apt-get update && apt-get install -y --no-install-recommends") + for _, pkg := range g.BuiltinSystemPackages { + sb.WriteString(fmt.Sprintf(" %s", pkg)) + } + sb.WriteString("\"") + + cacheDir := "/var/cache/apt" + cacheLibDir := "/var/lib/apt" + + run := root.Run(llb.Shlex(sb.String()), + llb.WithCustomNamef("(built-in packages) apt-get install %s", + strings.Join(g.BuiltinSystemPackages, " "))) + run.AddMount(cacheDir, llb.Scratch(), + llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared)) + run.AddMount(cacheLibDir, llb.Scratch(), + llb.AsPersistentCacheDir("/"+cacheLibDir, llb.CacheMountShared)) + + // TODO(gaocegege): Refactor user to a seperate stage. + res := run. + Run(llb.Shlex("groupadd -g 1000 envd"), llb.WithCustomName("create user group envd")). + Run(llb.Shlex("useradd -p \"\" -u 1000 -g envd -s /bin/sh -m envd"), llb.WithCustomName("create user envd")). + Run(llb.Shlex("adduser envd sudo"), llb.WithCustomName("add user envd to sudoers")) + return llb.User("envd")(res.Root()) +} + +func (g *Graph) compileCUDAPackages() llb.State { + root := llb.Image( + fmt.Sprintf("nvidia/cuda:%s.0-cudnn%s-devel-%s", *g.CUDA, *g.CUDNN, g.OS)) + g.BuiltinSystemPackages = append(g.BuiltinSystemPackages, []string{ + g.Language, + fmt.Sprintf("%s-pip", g.Language), + }...) + return root +} + +func (g Graph) compileSystemPackages(root llb.State) llb.State { + if len(g.SystemPackages) == 0 { + return root + } + + // Compose the package install command. + var sb strings.Builder + sb.WriteString("sudo 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(sb.String()), + llb.WithCustomNamef("(user-defined packages) apt-get install %s", + strings.Join(g.SystemPackages, " "))) + run.AddMount(cacheDir, llb.Scratch(), + llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared)) + run.AddMount(cacheLibDir, llb.Scratch(), + llb.AsPersistentCacheDir("/"+cacheLibDir, llb.CacheMountShared)) + return run.Root() +} + +func (g *Graph) compileBase() llb.State { + var base llb.State + if g.CUDA == nil && g.CUDNN == nil { + base = llb.Image("docker.io/library/python:3.8") + } else { + base = g.compileCUDAPackages() + } + return base +} + +func (g Graph) copyEnvdSSHServer() llb.State { + // TODO(gaocegege): Remove global var ssh image. + run := llb.Image(viper.GetString(flag.FlagSSHImage)). + File(llb.Copy(llb.Image(viper.GetString(flag.FlagSSHImage)), + "usr/bin/envd-ssh", "/var/envd/bin/envd-ssh", + &llb.CopyInfo{CreateDestPath: true}), llb.WithCustomName("install envd-ssh")) + return run +}