From 6b6760705b68f3a063159e878a52878a36b18b1b Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Tue, 25 Oct 2022 21:07:55 +0800 Subject: [PATCH] chore: Refactor the IR (#1088) Signed-off-by: Ce Gao Signed-off-by: Ce Gao --- pkg/lang/ir/compile.go | 7 +---- pkg/lang/ir/julia.go | 4 +-- pkg/lang/ir/python.go | 60 +++++++++++++++++++++------------------ pkg/lang/ir/r.go | 4 +-- pkg/lang/ir/supervisor.go | 2 +- pkg/lang/ir/system.go | 17 +++++++---- 6 files changed, 50 insertions(+), 44 deletions(-) diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go index 7b446b00b..2b0c4be75 100644 --- a/pkg/lang/ir/compile.go +++ b/pkg/lang/ir/compile.go @@ -253,15 +253,10 @@ func (g Graph) Compile(uid, gid int) (llb.State, error) { }).Debug("compile LLB") // TODO(gaocegege): Support more OS and langs. - base, err := g.compileBase() + aptStage, err := g.compileBase() if err != nil { return llb.State{}, errors.Wrap(err, "failed to get the base image") } - source, err := g.compileExtraSource(base) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to get extra sources") - } - aptStage := g.compileUbuntuAPT(source) var merged llb.State // Use custom logic when image is specified. if g.Image != nil { diff --git a/pkg/lang/ir/julia.go b/pkg/lang/ir/julia.go index de85af966..1b89e990f 100644 --- a/pkg/lang/ir/julia.go +++ b/pkg/lang/ir/julia.go @@ -58,12 +58,12 @@ func (g Graph) compileJulia(aptStage llb.State) (llb.State, error) { merged = llb.Merge([]llb.State{ builtinSystemStage, systemStage, diffShellStage, diffSSHStage, juliaStage, *vscodeStage, - }, llb.WithCustomName("merging all components into one")) + }, llb.WithCustomName("[internal] generating the image")) } else { merged = llb.Merge([]llb.State{ builtinSystemStage, systemStage, diffShellStage, diffSSHStage, juliaStage, - }, llb.WithCustomName("merging all components into one")) + }, llb.WithCustomName("[internal] generating the image")) } return merged, nil } diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go index 4cae46b2b..d655129e9 100644 --- a/pkg/lang/ir/python.go +++ b/pkg/lang/ir/python.go @@ -52,16 +52,8 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) { if err := g.compileJupyter(); err != nil { return llb.State{}, errors.Wrap(err, "failed to compile jupyter") } - builtinSystemStage := pypiMirrorStage - - sshStage, err := g.copySSHKey(builtinSystemStage) - if err != nil { - return llb.State{}, errors.Wrap(err, "failed to copy ssh keys") - } - diffSSHStage := llb.Diff(builtinSystemStage, sshStage, llb.WithCustomName("install ssh keys")) - // Conda affects shell and python, thus we cannot do it in parallel. - shellStage, err := g.compileShell(builtinSystemStage) + shellStage, err := g.compileShell(pypiMirrorStage) if err != nil { return llb.State{}, errors.Wrap(err, "failed to compile shell") } @@ -73,34 +65,42 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) { return llb.State{}, errors.Wrap(err, "failed to compile conda environment") } - diffCondaEnvStage := llb.Diff(shellStage, condaEnvStage) - diffSystemStage := llb.Diff(shellStage, systemStage) + diffCondaEnvStage := llb.Diff(shellStage, condaEnvStage, + llb.WithCustomName("[internal] conda python environment")) + diffSystemStage := llb.Diff(shellStage, systemStage, + llb.WithCustomName("[internal] install system packages")) prePythonStage := llb.Merge([]llb.State{diffSystemStage, diffCondaEnvStage, shellStage}, llb.WithCustomName("pre-python stage")) - condaStage := llb.Diff(builtinSystemStage, + condaStage := llb.Diff(prePythonStage, g.compileCondaPackages(prePythonStage), - llb.WithCustomName("install conda packages")) + llb.WithCustomName("[internal] install conda packages")) - pypiStage := llb.Diff(condaEnvStage, + pypiStage := llb.Diff(prePythonStage, g.compilePyPIPackages(prePythonStage), - llb.WithCustomName("install PyPI packages")) + llb.WithCustomName("[internal] install PyPI packages")) vscodeStage, err := g.compileVSCode() if err != nil { return llb.State{}, errors.Wrap(err, "failed to get vscode plugins") } + sshStage, err := g.copySSHKey(prePythonStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to copy ssh keys") + } + diffSSHStage := llb.Diff(prePythonStage, sshStage, + llb.WithCustomName("[internal] install ssh key")) var merged llb.State if vscodeStage != nil { merged = llb.Merge([]llb.State{ - builtinSystemStage, condaStage, - diffSSHStage, pypiStage, *vscodeStage, - }, llb.WithCustomName("merging all components into one")) + prePythonStage, condaStage, pypiStage, + diffSSHStage, *vscodeStage, + }, llb.WithCustomName("[internal] generating the image")) } else { merged = llb.Merge([]llb.State{ - builtinSystemStage, condaStage, + prePythonStage, condaStage, diffSSHStage, pypiStage, - }, llb.WithCustomName("merging all components into one")) + }, llb.WithCustomName("[internal] generating the image")) } merged = g.compileAlternative(merged) return merged, nil @@ -110,10 +110,14 @@ func (g Graph) compilePython(aptStage llb.State) (llb.State, error) { func (g Graph) compileAlternative(root llb.State) llb.State { envdPrefix := "/opt/conda/envs/envd/bin" run := root. - Run(llb.Shlexf("update-alternatives --install /usr/bin/python python %s/python 1", envdPrefix), llb.WithCustomName("update alternative python to envd")). - Run(llb.Shlexf("update-alternatives --install /usr/bin/python3 python3 %s/python3 1", envdPrefix), llb.WithCustomName("update alternative python3 to envd")). - Run(llb.Shlexf("update-alternatives --install /usr/bin/pip pip %s/pip 1", envdPrefix), llb.WithCustomName("update alternative pip to envd")). - Run(llb.Shlexf("update-alternatives --install /usr/bin/pip3 pip3 %s/pip3 1", envdPrefix), llb.WithCustomName("update alternative pip3 to envd")) + Run(llb.Shlexf("update-alternatives --install /usr/bin/python python %s/python 1", envdPrefix), + llb.WithCustomName("[internal] update alternative python to envd")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/python3 python3 %s/python3 1", envdPrefix), + llb.WithCustomName("[internal] update alternative python3 to envd")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/pip pip %s/pip 1", envdPrefix), + llb.WithCustomName("[internal] update alternative pip to envd")). + Run(llb.Shlexf("update-alternatives --install /usr/bin/pip3 pip3 %s/pip3 1", envdPrefix), + llb.WithCustomName("[internal] update alternative pip3 to envd")) return run.Root() } @@ -200,13 +204,13 @@ func (g Graph) compilePyPIIndex(root llb.State) llb.State { extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL } content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex) + dir := filepath.Dir(pypiIndexFilePath) pypiMirror := root. - File(llb.Mkdir(filepath.Dir(pypiIndexFilePath), - 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), - llb.WithCustomName("[internal] setting PyPI index")). + File(llb.Mkdir(dir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomNamef("[internal] setting PyPI index dir %s", dir)). File(llb.Mkfile(pypiIndexFilePath, 0644, []byte(content), llb.WithUIDGID(g.uid, g.gid)), - llb.WithCustomName("[internal] setting PyPI index")) + llb.WithCustomNamef("[internal] setting PyPI index file %s", pypiIndexFilePath)) return pypiMirror } return root diff --git a/pkg/lang/ir/r.go b/pkg/lang/ir/r.go index b379d225f..2f268358b 100644 --- a/pkg/lang/ir/r.go +++ b/pkg/lang/ir/r.go @@ -58,12 +58,12 @@ func (g Graph) compileRLang(aptStage llb.State) (llb.State, error) { merged = llb.Merge([]llb.State{ builtinSystemStage, systemStage, diffShellStage, diffSSHStage, rPackageInstallStage, *vscodeStage, - }, llb.WithCustomName("merging all components into one")) + }, llb.WithCustomName("[internal] generating the image")) } else { merged = llb.Merge([]llb.State{ builtinSystemStage, systemStage, diffShellStage, diffSSHStage, rPackageInstallStage, - }, llb.WithCustomName("merging all components into one")) + }, llb.WithCustomName("[internal] generating the image")) } return merged, nil } diff --git a/pkg/lang/ir/supervisor.go b/pkg/lang/ir/supervisor.go index 3bf3ed332..7956bd96a 100644 --- a/pkg/lang/ir/supervisor.go +++ b/pkg/lang/ir/supervisor.go @@ -51,7 +51,7 @@ wait = "5s" func (g Graph) addNewProcess(root llb.State, name, command string) llb.State { template := fmt.Sprintf(horustTemplate, name, command, types.EnvdWorkDir) filename := filepath.Join(types.HorustServiceDir, fmt.Sprintf("%s.toml", name)) - supervisor := root.File(llb.Mkfile(filename, 0644, []byte(template), llb.WithUIDGID(g.uid, g.gid))) + supervisor := root.File(llb.Mkfile(filename, 0644, []byte(template), llb.WithUIDGID(g.uid, g.gid)), llb.WithCustomNamef("[internal] create file %s", filename)) return supervisor } diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go index 781787d02..6097c712e 100644 --- a/pkg/lang/ir/system.go +++ b/pkg/lang/ir/system.go @@ -164,7 +164,7 @@ func (g *Graph) preparePythonBase(root llb.State) llb.State { return run.Root() } -func (g Graph) compileSshd(root llb.State) llb.State { +func (g Graph) compileSSHD(root llb.State) llb.State { sshd := root.File(llb.Copy( llb.Image(types.EnvdSshdImage), "/usr/bin/envd-sshd", "/var/envd/bin/envd-sshd", &llb.CopyInfo{CreateDestPath: true}), @@ -212,14 +212,20 @@ func (g *Graph) compileBase() (llb.State, error) { base = g.compileCUDAPackages("nvidia/cuda") } - base = g.compileUserGroup(base) // Install conda first. condaStage, err := g.installConda(base) if err != nil { return llb.State{}, errors.Wrap(err, "failed to install conda") } supervisor := g.installHorust(condaStage) - return g.compileSshd(supervisor), nil + sshdStage := g.compileSSHD(supervisor) + source, err := g.compileExtraSource(sshdStage) + if err != nil { + return llb.State{}, errors.Wrap(err, "failed to get extra sources") + } + aptStage := g.compileUbuntuAPT(source) + final := g.compileUserGroup(aptStage) + return final, nil } func (g Graph) installHorust(root llb.State) llb.State { @@ -240,9 +246,10 @@ func (g Graph) copySSHKey(root llb.State) (llb.State, error) { } run := root. File(llb.Mkdir("/var/envd", 0755, llb.WithParents(true), - llb.WithUIDGID(g.uid, g.gid))). + llb.WithUIDGID(g.uid, g.gid)), + llb.WithCustomName("[internal] create dir for ssh key")). File(llb.Mkfile(config.ContainerAuthorizedKeysPath, 0644, []byte(dat+" envd"), llb.WithUIDGID(g.uid, g.gid)), - llb.WithCustomName("install ssh keys")) + llb.WithCustomName("[internal] install ssh keys")) return run, nil }