feat(v1): copy the requirements content if possible (#1738)

Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Aug 14, 2023 · be78326 · be78326
1 parent 7441e38
commit be78326
Show file tree

Hide file tree

Showing 8 changed files with 79 additions and 15 deletions.
diff --git a/pkg/builder/build.go b/pkg/builder/build.go
@@ -18,7 +18,6 @@ import (
 	"context"
 	"io"
 	"os"
-	"path/filepath"
 	"strings"
 
 	"github.com/cockroachdb/errors"
@@ -160,12 +159,11 @@ func (b generalBuilder) Interpret() error {
 }
 
 func (b generalBuilder) Compile(ctx context.Context) (*llb.Definition, error) {
-	envName := filepath.Base(b.BuildContextDir)
 	platform, err := parsePlatform(b.Platform)
 	if err != nil {
 		return nil, err
 	}
-	def, err := b.graph.Compile(ctx, envName, b.PubKeyPath, platform, b.Options.ProgressMode)
+	def, err := b.graph.Compile(ctx, b.BuildContextDir, b.PubKeyPath, platform, b.Options.ProgressMode)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to compile build.envd")
 	}

diff --git a/pkg/lang/ir/graph.go b/pkg/lang/ir/graph.go
@@ -24,7 +24,7 @@ import (
 )
 
 type Graph interface {
-	Compile(ctx context.Context, envName string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error)
+	Compile(ctx context.Context, envPath string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error)
 
 	graphDebugger
 	graphVisitor

diff --git a/pkg/lang/ir/v0/compile.go b/pkg/lang/ir/v0/compile.go
@@ -112,13 +112,14 @@ func (g generalGraph) GetPlatform() *ocispecs.Platform {
 	return g.Platform
 }
 
-func (g *generalGraph) Compile(ctx context.Context, envName string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error) {
+func (g *generalGraph) Compile(ctx context.Context, envPath string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error) {
 	w, err := compileui.New(ctx, os.Stdout, progressMode)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create compileui")
 	}
 	g.Writer = w
-	g.EnvironmentName = envName
+	g.EnvironmentPath = envPath
+	g.EnvironmentName = filepath.Base(envPath)
 	g.PublicKeyPath = pub
 	g.Platform = platform
 

diff --git a/pkg/lang/ir/v0/types.go b/pkg/lang/ir/v0/types.go
@@ -77,6 +77,8 @@ type generalGraph struct {
 	// It is the BaseDir(BuildContextDir)
 	// e.g. mnist, streamlit-mnist
 	EnvironmentName string
+	// EnvironmentPath is the full path of this environment.
+	EnvironmentPath string
 
 	ir.RuntimeGraph
 

diff --git a/pkg/lang/ir/v1/compile.go b/pkg/lang/ir/v1/compile.go
@@ -109,7 +109,7 @@ func (g generalGraph) GetPlatform() *ocispecs.Platform {
 	return g.Platform
 }
 
-func (g *generalGraph) Compile(ctx context.Context, envName string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error) {
+func (g *generalGraph) Compile(ctx context.Context, envPath string, pub string, platform *ocispecs.Platform, progressMode string) (*llb.Definition, error) {
 	w, err := compileui.New(ctx, os.Stdout, progressMode)
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create compileui")
@@ -120,7 +120,8 @@ func (g *generalGraph) Compile(ctx context.Context, envName string, pub string,
 	}
 
 	g.Writer = w
-	g.EnvironmentName = envName
+	g.EnvironmentPath = envPath
+	g.EnvironmentName = filepath.Base(envPath)
 	g.PublicKeyPath = pub
 	g.Platform = platform
 	if c.Builder == types.BuilderTypeMoby {

diff --git a/pkg/lang/ir/v1/python.go b/pkg/lang/ir/v1/python.go
@@ -132,13 +132,25 @@ func (g generalGraph) compilePyPIPackages(root llb.State) llb.State {
 		logrus.WithField("file", *g.RequirementsFile).
 			Debug("Configure pip install requirements statements")
 		root = root.Dir(g.getWorkingDir())
-		run := root.
-			Run(llb.Shlexf("python -m pip install -r %s", *g.RequirementsFile),
-				llb.WithCustomNamef("pip install -r %s", *g.RequirementsFile))
-		run.AddMount(cacheDir, cache,
-			llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
-		run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext))
-		root = run.Root()
+		dependencies, safeToCopy := g.IsRequirementsFileSafeToCopyContent()
+		logrus.WithField("safeToCopy", safeToCopy).WithField("dependencies", dependencies).
+			Debug("Is requirements file safe to copy")
+		if safeToCopy {
+			// avoid mounting host directory to make it cache friendly
+			root = root.Run(llb.Shlexf(
+				"python -m pip install %s",
+				strings.Join(dependencies, " ")),
+				llb.WithCustomNamef("[internal] pip install from %s with %s", *g.RequirementsFile, strings.Join(dependencies, " ")),
+			).Root()
+		} else {
+			run := root.
+				Run(llb.Shlexf("python -m pip install -r %s", *g.RequirementsFile),
+					llb.WithCustomNamef("[internal] pip install -r %s", *g.RequirementsFile))
+			run.AddMount(cacheDir, cache,
+				llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared), llb.SourcePath("/cache/pip"))
+			run.AddMount(g.getWorkingDir(), llb.Local(flag.FlagBuildContext))
+			root = run.Root()
+		}
 	}
 
 	if len(g.PythonWheels) > 0 {

diff --git a/pkg/lang/ir/v1/types.go b/pkg/lang/ir/v1/types.go
@@ -78,6 +78,8 @@ type generalGraph struct {
 	// It is the BaseDir(BuildContextDir)
 	// e.g. mnist, streamlit-mnist
 	EnvironmentName string
+	// EnvironmentPath is the full path of this environment.
+	EnvironmentPath string
 
 	// (v1) disable `merge` op for `moby` builder
 	// check https://github.com/tensorchord/envd/issues/1693

diff --git a/pkg/lang/ir/v1/util.go b/pkg/lang/ir/v1/util.go
@@ -15,13 +15,16 @@
 package v1
 
 import (
+	"bufio"
 	"bytes"
 	"crypto/md5"
 	"encoding/gob"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
+	"os"
 	"os/user"
+	"path/filepath"
 	"regexp"
 	"runtime"
 	"strconv"
@@ -175,3 +178,48 @@ func (g generalGraph) GeneralGraphFromLabel(label []byte) (ir.Graph, error) {
 	}
 	return &newg, nil
 }
+
+// IsRequirementsFileSafeToCopyContent returns true and the dependencies in this file if the
+// requirements file is safe to copy. Otherwise, it returns false and empty dependencies.
+// check https://github.com/tensorchord/envd/issues/1629
+func (g generalGraph) IsRequirementsFileSafeToCopyContent() (dependencies []string, safe bool) {
+	if g.RequirementsFile == nil {
+		return
+	}
+	filePath := filepath.Join(g.EnvironmentPath, *g.RequirementsFile)
+	file, err := os.Open(filePath)
+	if err != nil {
+		logrus.WithError(err).Debugf("failed to open requirements file: %s", filePath)
+		return
+	}
+	defer file.Close()
+
+	forbidden := []string{
+		"-e",
+		"-r",
+		"-c",
+		"--find-links",
+	}
+
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if strings.HasPrefix(line, "#") {
+			continue
+		}
+		for _, f := range forbidden {
+			if strings.Contains(line, f) {
+				logrus.Debugf("requirements file contains %s, which is not safe to copy", f)
+				return []string{}, false
+			}
+		}
+		exist, err := fileutil.FileExists(line)
+		// has a local file to install
+		if err == nil && exist {
+			logrus.Debugf("requirements file contains local file %s, which is not safe to copy", line)
+			return []string{}, false
+		}
+		dependencies = append(dependencies, line)
+	}
+	return dependencies, true
+}