From 8c0c98d6741b70a0afb17bd93b500df826d1a2a9 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Wed, 1 Jun 2022 20:57:39 -0400 Subject: [PATCH] feat: Support extra PyPI index (#229) Signed-off-by: terrytangyuan --- pkg/lang/frontend/starlark/const.go | 2 +- pkg/lang/frontend/starlark/rules.go | 26 +++++++++++++++----------- pkg/lang/ir/compile.go | 2 +- pkg/lang/ir/consts.go | 14 ++++++++------ pkg/lang/ir/interface.go | 15 ++++++++------- pkg/lang/ir/python.go | 17 +++++++++++------ pkg/lang/ir/types.go | 5 +++-- 7 files changed, 47 insertions(+), 34 deletions(-) diff --git a/pkg/lang/frontend/starlark/const.go b/pkg/lang/frontend/starlark/const.go index f4de60ece..533e10c8b 100644 --- a/pkg/lang/frontend/starlark/const.go +++ b/pkg/lang/frontend/starlark/const.go @@ -21,7 +21,7 @@ const ( ruleCUDA = "cuda" ruleVSCode = "vscode" ruleUbuntuAPT = "ubuntu_apt" - rulePyPIMirror = "pip_mirror" + rulePyPIIndex = "pip_index" ruleShell = "shell" ruleJupyter = "jupyter" ruleRun = "run" diff --git a/pkg/lang/frontend/starlark/rules.go b/pkg/lang/frontend/starlark/rules.go index 035f13c7b..ca0710266 100644 --- a/pkg/lang/frontend/starlark/rules.go +++ b/pkg/lang/frontend/starlark/rules.go @@ -36,7 +36,7 @@ func registerenvdRules() { starlark.Universe[ruleCUDA] = starlark.NewBuiltin(ruleCUDA, ruleFuncCUDA) starlark.Universe[ruleVSCode] = starlark.NewBuiltin(ruleVSCode, ruleFuncVSCode) starlark.Universe[ruleUbuntuAPT] = starlark.NewBuiltin(ruleUbuntuAPT, ruleFuncUbuntuAPT) - starlark.Universe[rulePyPIMirror] = starlark.NewBuiltin(rulePyPIMirror, ruleFuncPyPIMirror) + starlark.Universe[rulePyPIIndex] = starlark.NewBuiltin(rulePyPIIndex, ruleFuncPyPIIndex) starlark.Universe[ruleShell] = starlark.NewBuiltin(ruleShell, ruleFuncShell) starlark.Universe[ruleJupyter] = starlark.NewBuiltin(ruleJupyter, ruleFuncJupyter) starlark.Universe[ruleRun] = starlark.NewBuiltin(ruleRun, ruleFuncRun) @@ -186,12 +186,12 @@ func ruleFuncUbuntuAPT(thread *starlark.Thread, _ *starlark.Builtin, return starlark.None, nil } -func ruleFuncPyPIMirror(thread *starlark.Thread, _ *starlark.Builtin, +func ruleFuncPyPIIndex(thread *starlark.Thread, _ *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { - var mode, mirror starlark.String + var mode, url, extraURL starlark.String - if err := starlark.UnpackArgs(rulePyPIMirror, args, kwargs, - "mode?", &mode, "mirror?", &mirror); err != nil { + if err := starlark.UnpackArgs(rulePyPIIndex, args, kwargs, + "mode?", &mode, "url?", &url, "extra_url?", &extraURL); err != nil { return nil, err } @@ -199,14 +199,18 @@ func ruleFuncPyPIMirror(thread *starlark.Thread, _ *starlark.Builtin, if mode != starlark.String("") { modeStr = mode.GoString() } - mirrorStr := "" - if mirror != starlark.String("") { - mirrorStr = mirror.GoString() + indexStr := "" + if url != starlark.String("") { + indexStr = url.GoString() + } + extraIndexStr := "" + if extraURL != starlark.String("") { + extraIndexStr = extraURL.GoString() } - logger.Debugf("rule `%s` is invoked, mode=%s, mirror=%s", rulePyPIMirror, - modeStr, mirrorStr) - if err := ir.PyPIMirror(modeStr, mirrorStr); err != nil { + logger.Debugf("rule `%s` is invoked, mode=%s, index=%s, extraIndex=%s", rulePyPIIndex, + modeStr, indexStr, extraIndexStr) + if err := ir.PyPIIndex(modeStr, indexStr, extraIndexStr); err != nil { return nil, err } diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go index 1743682bb..d5e841ab2 100644 --- a/pkg/lang/ir/compile.go +++ b/pkg/lang/ir/compile.go @@ -102,7 +102,7 @@ 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) + pypiMirrorStage := g.compilePyPIIndex(aptStage) g.compileJupyter() builtinSystemStage := pypiMirrorStage diff --git a/pkg/lang/ir/consts.go b/pkg/lang/ir/consts.go index 99b8941e6..c6c2db93e 100644 --- a/pkg/lang/ir/consts.go +++ b/pkg/lang/ir/consts.go @@ -15,16 +15,18 @@ package ir const ( - osDefault = "ubuntu20.04" - languageDefault = "python3.8" - mirrorModeAuto = "auto" + osDefault = "ubuntu20.04" + languageDefault = "python3.8" + pypiIndexModeAuto = "auto" - aptSourceFilePath = "/etc/apt/sources.list" - pypiMirrorFilePath = "/etc/pip.conf" + aptSourceFilePath = "/etc/apt/sources.list" + pypiIndexFilePath = "/etc/pip.conf" pypiConfigTemplate = ` [global] -index-url=%s` +index-url=%s +%s +` defaultUID = 1000 defaultGID = 1000 diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go index 35cccac3e..0ced54072 100644 --- a/pkg/lang/ir/interface.go +++ b/pkg/lang/ir/interface.go @@ -52,7 +52,7 @@ func VSCodePlugins(plugins []string) error { // UbuntuAPT updates the Ubuntu apt source.list in the image. func UbuntuAPT(mode, source string) error { if source == "" { - if mode == mirrorModeAuto { + if mode == pypiIndexModeAuto { // If the mode is set to `auto`, envd detects the location of the run // then set to the nearest mirror return errors.New("auto-mode not implemented") @@ -64,17 +64,18 @@ func UbuntuAPT(mode, source string) error { return nil } -func PyPIMirror(mode, mirror string) error { - if mirror == "" { - if mode == mirrorModeAuto { +func PyPIIndex(mode, url, extraURL string) error { + if url == "" { + if mode == pypiIndexModeAuto { // If the mode is set to `auto`, envd detects the location of the run - // then set to the nearest mirror. + // then set to the nearest index URL. return errors.New("auto-mode not implemented") } - return errors.New("mirror is required") + return errors.New("url is required") } - DefaultGraph.PyPIMirror = &mirror + DefaultGraph.PyPIIndexURL = &url + DefaultGraph.PyPIExtraIndexURL = &extraURL return nil } diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go index b9cf649e6..547938cf1 100644 --- a/pkg/lang/ir/python.go +++ b/pkg/lang/ir/python.go @@ -52,14 +52,19 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State { 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) +func (g Graph) compilePyPIIndex(root llb.State) llb.State { + if g.PyPIIndexURL != nil { + logrus.WithField("index", *g.PyPIIndexURL).Debug("using custom PyPI index") + var extraIndex string + if g.PyPIExtraIndexURL != nil { + logrus.WithField("index", *g.PyPIIndexURL).Debug("using extra PyPI index") + extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL + } + content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex) pypiMirror := llb.Scratch(). - File(llb.Mkdir(filepath.Dir(pypiMirrorFilePath), + File(llb.Mkdir(filepath.Dir(pypiIndexFilePath), 0755, llb.WithParents(true), llb.WithUIDGID(defaultUID, defaultGID))). - File(llb.Mkfile(pypiMirrorFilePath, + File(llb.Mkfile(pypiIndexFilePath, 0644, []byte(content), llb.WithUIDGID(defaultUID, defaultGID))) return llb.Merge([]llb.State{root, pypiMirror}, llb.WithCustomName("add PyPI mirror")) } diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go index 2118dea37..4fa24ed0e 100644 --- a/pkg/lang/ir/types.go +++ b/pkg/lang/ir/types.go @@ -28,8 +28,9 @@ type Graph struct { CUDA *string CUDNN *string - UbuntuAPTSource *string - PyPIMirror *string + UbuntuAPTSource *string + PyPIIndexURL *string + PyPIExtraIndexURL *string PublicKeyPath string