feat: support trust the pypi index in v1, add api doc (#1412)

* feat: support trust the pypi index in v1, add api doc Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix lint Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix py lint Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Jan 12, 2023 · e1bbd56 · e1bbd56
1 parent d1179a6
commit e1bbd56
Show file tree

Hide file tree

Showing 8 changed files with 36 additions and 19 deletions.
diff --git a/envd/api/v0/config.py b/envd/api/v0/config.py
@@ -56,13 +56,14 @@ def jupyter(token: str, port: int):
     """
 
 
-def pip_index(url: str, extra_url: str):
+def pip_index(url: str, extra_url: str = "", trust: bool = False):
     """Configure pypi index mirror
 
     Args:
         url (str): PyPI index URL (i.e. https://mirror.sjtu.edu.cn/pypi/web/simple)
         extra_url (str): PyPI extra index URL. `url` and `extra_url` will be
             treated equally, see https://github.com/pypa/pip/issues/8606
+        trust (bool): trust the provided index
     """
 
 

diff --git a/envd/api/v1/config.py b/envd/api/v1/config.py
@@ -62,13 +62,14 @@ def jupyter(token: str, port: int):
     """
 
 
-def pip_index(url: str, extra_url: str):
+def pip_index(url: str, extra_url: str = "", trust: bool = False):
     """Configure pypi index mirror
 
     Args:
         url (str): PyPI index URL (i.e. https://mirror.sjtu.edu.cn/pypi/web/simple)
         extra_url (str): PyPI extra index URL. `url` and `extra_url` will be
             treated equally, see https://github.com/pypa/pip/issues/8606
+        trust (bool): trust the provided index
     """
 
 

diff --git a/envd/api/v1/install.py b/envd/api/v1/install.py
@@ -50,17 +50,11 @@ def conda(use_mamba: bool = False):
 
 
 def r_lang():
-    """Install R Lang.
-
-    Not implemented yet. Please use v0 if you need R.
-    """
+    """Install R Lang."""
 
 
 def julia():
-    """Install Julia.
-
-    Not implemented yet. Please use v0 if you need Julia.
-    """
+    """Install Julia."""
 
 
 def apt_packages(name: List[str] = []):

diff --git a/pkg/lang/frontend/starlark/v1/config/config.go b/pkg/lang/frontend/starlark/v1/config/config.go
@@ -96,18 +96,19 @@ func ruleFuncJupyter(thread *starlark.Thread, _ *starlark.Builtin,
 func ruleFuncPyPIIndex(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var url, extraURL starlark.String
+	var trust bool
 
 	if err := starlark.UnpackArgs(rulePyPIIndex, args, kwargs,
-		"url?", &url, "extra_url?", &extraURL); err != nil {
+		"url", &url, "extra_url?", &extraURL, "trust?", &trust); err != nil {
 		return nil, err
 	}
 
 	indexStr := url.GoString()
 	extraIndexStr := extraURL.GoString()
 
-	logger.Debugf("rule `%s` is invoked, index=%s, extraIndex=%s",
-		rulePyPIIndex, indexStr, extraIndexStr)
-	if err := ir.PyPIIndex(indexStr, extraIndexStr); err != nil {
+	logger.Debugf("rule `%s` is invoked, index=%s, extraIndex=%s, trust=%t",
+		rulePyPIIndex, indexStr, extraIndexStr, trust)
+	if err := ir.PyPIIndex(indexStr, extraIndexStr, trust); err != nil {
 		return nil, err
 	}
 

diff --git a/pkg/lang/ir/v1/consts.go b/pkg/lang/ir/v1/consts.go
@@ -27,6 +27,7 @@ const (
 [global]
 index-url=%s
 %s
+%s
 
 [install]
 src = /tmp

diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go
@@ -160,14 +160,17 @@ func UbuntuAPT(source string) error {
 	return nil
 }
 
-func PyPIIndex(url, extraURL string) error {
+func PyPIIndex(url, extraURL string, trust bool) error {
 	if url == "" {
 		return errors.New("url is required")
 	}
 	g := DefaultGraph.(*generalGraph)
 
 	g.PyPIIndexURL = &url
-	g.PyPIExtraIndexURL = &extraURL
+	if len(extraURL) > 0 {
+		g.PyPIExtraIndexURL = &extraURL
+	}
+	g.PyPITrust = trust
 	return nil
 }
 

diff --git a/pkg/lang/ir/v1/python.go b/pkg/lang/ir/v1/python.go
@@ -16,6 +16,7 @@ package v1
 
 import (
 	"fmt"
+	"net/url"
 	"path/filepath"
 	"strings"
 
@@ -168,12 +169,26 @@ func (g generalGraph) compilePyPIIndex(root llb.State) llb.State {
 		return root
 	}
 	logrus.WithField("index", *g.PyPIIndexURL).Debug("using custom PyPI index")
-	var extraIndex string
+	var extra, trusted string
 	if g.PyPIExtraIndexURL != nil {
 		logrus.WithField("index", *g.PyPIIndexURL).Debug("using extra PyPI index")
-		extraIndex = "extra-index-url=" + *g.PyPIExtraIndexURL
+		extra = "extra-index-url=" + *g.PyPIExtraIndexURL
 	}
-	content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extraIndex)
+	if g.PyPITrust {
+		var hosts []string
+		for _, p := range []*string{g.PyPIIndexURL, g.PyPIExtraIndexURL} {
+			if p != nil {
+				u, err := url.Parse(*p)
+				if err == nil && u != nil && u.Hostname() != "" {
+					hosts = append(hosts, u.Hostname())
+				}
+			}
+		}
+		if len(hosts) > 0 {
+			trusted = fmt.Sprintf("trusted-host=%s", strings.Join(hosts, " "))
+		}
+	}
+	content := fmt.Sprintf(pypiConfigTemplate, *g.PyPIIndexURL, extra, trusted)
 	dir := filepath.Dir(pypiIndexFilePath)
 	pypiMirror := root.
 		File(llb.Mkdir(dir, 0755, llb.WithParents(true), llb.WithUIDGID(g.uid, g.gid)),

diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
@@ -43,6 +43,7 @@ type generalGraph struct {
 	JuliaPackageServer *string
 	PyPIIndexURL       *string
 	PyPIExtraIndexURL  *string
+	PyPITrust          bool
 
 	PublicKeyPath string