feat: Support extra PyPI index (#229)

Signed-off-by: terrytangyuan <terrytangyuan@gmail.com>
tensorchord · Jun 2, 2022 · 8c0c98d · 8c0c98d
1 parent 31514df
commit 8c0c98d
Show file tree

Hide file tree

Showing 7 changed files with 47 additions and 34 deletions.
diff --git 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
@@ -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,27 +186,31 @@ 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
 	}
 
 	modeStr := ""
 	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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -28,8 +28,9 @@ type Graph struct {
 	CUDA     *string
 	CUDNN    *string
 
-	UbuntuAPTSource *string
-	PyPIMirror      *string
+	UbuntuAPTSource   *string
+	PyPIIndexURL      *string
+	PyPIExtraIndexURL *string
 
 	PublicKeyPath string