feat: Support configuring CRAN mirror for R environment (#405)

* feat: Support configuring CRAN mirror for R environment Signed-off-by: terrytangyuan <terrytangyuan@gmail.com> * fix panic Signed-off-by: terrytangyuan <terrytangyuan@gmail.com>
tensorchord · Jun 17, 2022 · 5f3b16b · 5f3b16b
1 parent 8e27e99
commit 5f3b16b
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 1 deletion.
diff --git a/pkg/lang/frontend/starlark/config/config.go b/pkg/lang/frontend/starlark/config/config.go
@@ -32,6 +32,8 @@ var Module = &starlarkstruct.Module{
 		"apt_source": starlark.NewBuiltin(ruleUbuntuAptSource, ruleFuncUbuntuAptSource),
 		"gpu":        starlark.NewBuiltin(ruleGPU, ruleFuncGPU),
 		"jupyter":    starlark.NewBuiltin(ruleJupyter, ruleFuncJupyter),
+		"cran_mirror": starlark.NewBuiltin(
+			ruleCRANMirror, ruleFuncCRANMirror),
 		"pip_index": starlark.NewBuiltin(
 			rulePyPIIndex, ruleFuncPyPIIndex),
 		"conda_channel": starlark.NewBuiltin(
@@ -117,6 +119,27 @@ func ruleFuncPyPIIndex(thread *starlark.Thread, _ *starlark.Builtin,
 	return starlark.None, nil
 }
 
+func ruleFuncCRANMirror(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var url starlark.String
+
+	if err := starlark.UnpackArgs(ruleCRANMirror, args, kwargs,
+		"url?", &url); err != nil {
+		return nil, err
+	}
+
+	urlStr := ""
+	if url != starlark.String("") {
+		urlStr = url.GoString()
+	}
+
+	logger.Debugf("rule `%s` is invoked, url=%s", ruleCRANMirror, urlStr)
+	if err := ir.CRANMirror(urlStr); err != nil {
+		return nil, err
+	}
+	return starlark.None, nil
+}
+
 func ruleFuncUbuntuAptSource(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var mode, source starlark.String

diff --git a/pkg/lang/frontend/starlark/config/const.go b/pkg/lang/frontend/starlark/config/const.go
@@ -17,6 +17,7 @@ package config
 const (
 	ruleUbuntuAptSource = "config.apt_source"
 	rulePyPIIndex       = "config.pip_index"
+	ruleCRANMirror      = "config.cran_mirror"
 	ruleJupyter         = "config.jupyter"
 	ruleCondaChannel    = "config.conda_channel"
 	ruleGPU             = "config.gpu"

diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -95,6 +95,11 @@ func PyPIIndex(mode, url, extraURL string) error {
 	return nil
 }
 
+func CRANMirror(url string) error {
+	DefaultGraph.CRANMirrorURL = &url
+	return nil
+}
+
 func Shell(shell string) error {
 	DefaultGraph.Shell = shell
 	return nil

diff --git a/pkg/lang/ir/r.go b/pkg/lang/ir/r.go
@@ -24,7 +24,11 @@ import (
 func (g Graph) installRPackages(root llb.State) llb.State {
 	// TODO(terrytangyuan): Support different CRAN mirrors
 	var sb strings.Builder
-	sb.WriteString(`R -e 'install.packages(c(`)
+	mirrorURL := "https://cran.rstudio.com"
+	if g.CRANMirrorURL != nil {
+		mirrorURL = *g.CRANMirrorURL
+	}
+	sb.WriteString(fmt.Sprintf(`R -e 'options(repos = c(CRAN = "%s")); install.packages(c(`, mirrorURL))
 	for i, pkg := range g.RPackages {
 		sb.WriteString(fmt.Sprintf(`"%s"`, pkg))
 		if i != len(g.RPackages)-1 {

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -34,6 +34,7 @@ type Graph struct {
 	NumGPUs int
 
 	UbuntuAPTSource   *string
+	CRANMirrorURL     *string
 	PyPIIndexURL      *string
 	PyPIExtraIndexURL *string