feat(lang): install.python_packages(local_wheels=[]) (#838)

* feat(lang): install.python_packages(local_wheels=[]) Signed-off-by: Keming <kemingyang@tensorchord.ai> * use mount Signed-off-by: Keming <kemingyang@tensorchord.ai> * add doc Signed-off-by: Keming <kemingyang@tensorchord.ai> Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Sep 1, 2022 · 190ee76 · 190ee76
1 parent 977dd47
commit 190ee76
Show file tree

Hide file tree

Showing 5 changed files with 33 additions and 10 deletions.
diff --git a/envd/api/install/__init__.py b/envd/api/install/__init__.py
@@ -32,12 +32,14 @@ def system_packages(name: List[str]):
     """
 
 
-def python_packages(name: List[str], requirements: str):
+def python_packages(name: List[str], requirements: str, local_wheels: List[str]):
     """Install python package by pip
 
     Args:
         name (List[str]): package name list
-        requirements: (str): requirements file path
+        requirements (str): requirements file path
+        local_wheels (List[str]): local wheels
+            (wheel files should be placed under the current directory)
     """
 
 

diff --git a/pkg/lang/frontend/starlark/install/install.go b/pkg/lang/frontend/starlark/install/install.go
@@ -46,9 +46,10 @@ func ruleFuncPyPIPackage(thread *starlark.Thread, _ *starlark.Builtin,
 	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
 	var name *starlark.List
 	var requirementsFile starlark.String
+	var wheels *starlark.List
 
 	if err := starlark.UnpackArgs(rulePyPIPackage, args, kwargs,
-		"name?", &name, "requirements?", &requirementsFile); err != nil {
+		"name?", &name, "requirements?", &requirementsFile, "local_wheels?", &wheels); err != nil {
 		return nil, err
 	}
 
@@ -61,10 +62,17 @@ func ruleFuncPyPIPackage(thread *starlark.Thread, _ *starlark.Builtin,
 
 	requirementsFileStr := requirementsFile.GoString()
 
-	logger.Debugf("rule `%s` is invoked, name=%v, requirements=%s",
-		rulePyPIPackage, nameList, requirementsFileStr)
+	localWheels := []string{}
+	if wheels != nil {
+		for i := 0; i < wheels.Len(); i++ {
+			localWheels = append(localWheels, wheels.Index(i).(starlark.String).GoString())
+		}
+	}
+
+	logger.Debugf("rule `%s` is invoked, name=%v, requirements=%s, local_wheels=%s",
+		rulePyPIPackage, nameList, requirementsFileStr, localWheels)
 
-	err := ir.PyPIPackage(nameList, requirementsFileStr)
+	err := ir.PyPIPackage(nameList, requirementsFileStr, localWheels)
 	return starlark.None, err
 }
 

diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -37,8 +37,9 @@ func Base(os, language, image string) error {
 	return nil
 }
 
-func PyPIPackage(deps []string, requirementsFile string) error {
+func PyPIPackage(deps []string, requirementsFile string, wheels []string) error {
 	DefaultGraph.PyPIPackages = append(DefaultGraph.PyPIPackages, deps...)
+	DefaultGraph.PythonWheels = append(DefaultGraph.PythonWheels, wheels...)
 
 	if requirementsFile != "" {
 		DefaultGraph.RequirementsFile = &requirementsFile

diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -115,7 +115,7 @@ func (g Graph) compileAlternative(root llb.State) llb.State {
 }
 
 func (g Graph) compilePyPIPackages(root llb.State) llb.State {
-	if len(g.PyPIPackages) == 0 && g.RequirementsFile == nil {
+	if len(g.PyPIPackages) == 0 && g.RequirementsFile == nil && len(g.PythonWheels) == 0 {
 		return root
 	}
 
@@ -158,14 +158,25 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 			Debug("Configure pip install requirements statements")
 		root = root.Dir(g.getWorkingDir())
 		run := root.
-			Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s",
-				strings.Join(g.PyPIPackages, " ")))
+			Run(llb.Shlex(cmd), llb.WithCustomNamef("pip install %s", *g.RequirementsFile))
 		run.AddMount(cacheDir, cache,
 			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache"))
 		run.AddMount(g.getWorkingDir(),
 			llb.Local(flag.FlagBuildContext), llb.Readonly)
 		root = run.Root()
 	}
+
+	if len(g.PythonWheels) > 0 {
+		root = root.Dir(g.getWorkingDir())
+		cmdTemplate := "/opt/conda/envs/envd/bin/python -m pip install %s"
+		for _, wheel := range g.PythonWheels {
+			run := root.Run(llb.Shlex(fmt.Sprintf(cmdTemplate, wheel)), llb.WithCustomNamef("pip install %s", wheel))
+			run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext), llb.Readonly)
+			run.AddMount(cacheDir, cache,
+				llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache"))
+			root = run.Root()
+		}
+	}
 	return root
 }
 

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -45,6 +45,7 @@ type Graph struct {
 
 	PyPIPackages     []string
 	RequirementsFile *string
+	PythonWheels     []string
 	RPackages        []string
 	JuliaPackages    []string
 	SystemPackages   []string