feat: Support installing multi-language environment at the same time (#…

…1501) * support installing multi-language environment Signed-off-by: Botong Ou <[email protected]> * support installing multi-language environment in parallel Signed-off-by: Botong Ou <[email protected]> * support installing multi-language environment in parallel Signed-off-by: Botong Ou <[email protected]> * support installing multi-language environment in parallel Signed-off-by: Botong Ou <[email protected]> * support installing multi-language environment in parallel Signed-off-by: Botong Ou <[email protected]> --------- Signed-off-by: Botong Ou <[email protected]>
tensorchord · Feb 24, 2023 · b641b0c · b641b0c
1 parent 6cbe558
commit b641b0c
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 48 deletions.
diff --git a/pkg/lang/ir/v1/editor.go b/pkg/lang/ir/v1/editor.go
@@ -65,12 +65,12 @@ func (g *generalGraph) compileJupyter() error {
 	}
 
 	g.PyPIPackages = append(g.PyPIPackages, []string{"jupyter"})
-	switch g.Language.Name {
-	case "python":
-		return nil
-	default:
-		return errors.Newf("Jupyter is not supported in %s yet", g.Language.Name)
+	for _, language := range g.Languages {
+		if language.Name == "python" {
+			return nil
+		}
 	}
+	return errors.Newf("Jupyter is not supported in other languages yet")
 }
 
 func (g generalGraph) generateJupyterCommand(workingDir string) []string {

diff --git a/pkg/lang/ir/v1/interface.go b/pkg/lang/ir/v1/interface.go
@@ -43,10 +43,10 @@ func Python(version string) error {
 	}
 	g := DefaultGraph.(*generalGraph)
 
-	g.Language = ir.Language{
+	g.Languages = append(g.Languages, ir.Language{
 		Name:    "python",
 		Version: &version,
-	}
+	})
 	return nil
 }
 
@@ -61,17 +61,17 @@ func Conda(mamba bool) {
 func RLang() {
 	g := DefaultGraph.(*generalGraph)
 
-	g.Language = ir.Language{
+	g.Languages = append(g.Languages, ir.Language{
 		Name: "r",
-	}
+	})
 }
 
 func Julia() {
 	g := DefaultGraph.(*generalGraph)
 
-	g.Language = ir.Language{
+	g.Languages = append(g.Languages, ir.Language{
 		Name: "julia",
-	}
+	})
 }
 
 func PyPIPackage(deps []string, requirementsFile string, wheels []string) error {

diff --git a/pkg/lang/ir/v1/python.go b/pkg/lang/ir/v1/python.go
@@ -68,11 +68,16 @@ func (g *generalGraph) installPython(root llb.State) (llb.State, error) {
 }
 
 func (g generalGraph) getAppropriatePythonVersion() (string, error) {
-	if g.Language.Version == nil {
-		return PythonVersionDefault, nil
+	var version string
+	for _, language := range g.Languages {
+		if language.Name == "python" {
+			if language.Version == nil {
+				return PythonVersionDefault, nil
+			}
+			version = *language.Version
+		}
 	}
 
-	version := *g.Language.Version
 	if version == "3" || version == "" {
 		return PythonVersionDefault, nil
 	}

diff --git a/pkg/lang/ir/v1/system.go b/pkg/lang/ir/v1/system.go
@@ -229,44 +229,51 @@ func (g *generalGraph) compileExtraSource(root llb.State) (llb.State, error) {
 }
 
 func (g *generalGraph) compileLanguage(root llb.State) (llb.State, error) {
+	langs := []llb.State{root}
 	lang := root
 	var err error
-	switch g.Language.Name {
-	case "python":
-		lang, err = g.installPython(root)
-	case "r":
-		rSrc := g.compileRLang(root)
-		lang = g.installRLang(rSrc)
-	case "julia":
-		lang = g.installJulia(root)
+	for _, language := range g.Languages {
+		switch language.Name {
+		case "python":
+			lang, err = g.installPython(root)
+		case "r":
+			rSrc := g.compileRLang(root)
+			lang = g.installRLang(rSrc)
+		case "julia":
+			lang = g.installJulia(root)
+		}
+		langs = append(langs, llb.Diff(root, lang, llb.WithCustomNamef("[internal] build %s environments", language.Name)))
 	}
-
-	return lang, err
+	return llb.Merge(langs, llb.WithCustomName("[internal] build all language environments")), err
 }
 
 func (g *generalGraph) compileLanguagePackages(root llb.State) llb.State {
+	packs := []llb.State{root}
 	pack := root
-	switch g.Language.Name {
-	case "python":
-		index := g.compilePyPIIndex(root)
-		pypi := g.compilePyPIPackages(index)
-		if g.CondaConfig == nil {
-			pack = pypi
-		} else {
-			channel := g.compileCondaChannel(root)
-			conda := g.compileCondaPackages(channel)
-			pack = llb.Merge([]llb.State{
-				root,
-				llb.Diff(root, pypi, llb.WithCustomName("[internal] PyPI packages")),
-				llb.Diff(root, conda, llb.WithCustomName("[internal] conda packages")),
-			}, llb.WithCustomName("[internal] Python packages"))
+	for _, language := range g.Languages {
+		switch language.Name {
+		case "python":
+			index := g.compilePyPIIndex(root)
+			pypi := g.compilePyPIPackages(index)
+			if g.CondaConfig == nil {
+				pack = pypi
+			} else {
+				channel := g.compileCondaChannel(root)
+				conda := g.compileCondaPackages(channel)
+				pack = llb.Merge([]llb.State{
+					root,
+					llb.Diff(root, pypi, llb.WithCustomName("[internal] PyPI packages")),
+					llb.Diff(root, conda, llb.WithCustomName("[internal] conda packages")),
+				}, llb.WithCustomName("[internal] Python packages"))
+			}
+		case "r":
+			pack = g.installRPackages(root)
+		case "julia":
+			pack = g.installJuliaPackages(root)
 		}
-	case "r":
-		pack = g.installRPackages(root)
-	case "julia":
-		pack = g.installJuliaPackages(root)
+		packs = append(packs, llb.Diff(root, pack, llb.WithCustomNamef("[internal] install %s's packages", language.Name)))
 	}
-	return pack
+	return llb.Merge(packs, llb.WithCustomName("[internal] install packages for all language environments"))
 }
 
 func (g *generalGraph) compileDevPackages(root llb.State) llb.State {
@@ -313,11 +320,8 @@ func (g *generalGraph) compileBaseImage() (llb.State, error) {
 
 	logger := logrus.WithFields(logrus.Fields{
 		"image":    g.Image,
-		"language": g.Language.Name,
+		"language": g.Languages,
 	})
-	if g.Language.Version != nil {
-		logger = logger.WithField("version", *g.Language.Version)
-	}
 	logger.Debug("compile base image")
 
 	// Fix https://github.com/tensorchord/envd/issues/1147.

diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
@@ -28,7 +28,7 @@ type generalGraph struct {
 	uid int `default:"-1"`
 	gid int `default:"-1"`
 
-	ir.Language
+	Languages         []ir.Language
 	EnvdSyntaxVersion string
 	Image             string
 	User              string