feat: Add prefix for cache id (#204)

* feat: Add prefix for cache id Signed-off-by: Ce Gao <cegao@tensorchord.ai> * fix: Add license Signed-off-by: Ce Gao <cegao@tensorchord.ai>
tensorchord · May 27, 2022 · e0984bb · e0984bb
1 parent ee31696
commit e0984bb
Show file tree

Hide file tree

Showing 6 changed files with 36 additions and 8 deletions.
diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -32,6 +32,7 @@ import (
 	"github.com/tensorchord/envd/pkg/lang/frontend/starlark"
 	"github.com/tensorchord/envd/pkg/lang/ir"
 	"github.com/tensorchord/envd/pkg/progress/progresswriter"
+	"github.com/tensorchord/envd/pkg/util/fileutil"
 )
 
 type Builder interface {
@@ -113,7 +114,7 @@ func (b generalBuilder) compile(ctx context.Context) (*llb.Definition, error) {
 	if err := b.interpret(); err != nil {
 		return nil, errors.Wrap(err, "failed to interpret")
 	}
-	def, err := ir.Compile(ctx)
+	def, err := ir.Compile(ctx, fileutil.Base(b.buildContextDir))
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to compile build.envd")
 	}

diff --git a/pkg/lang/ir/cache.go b/pkg/lang/ir/cache.go
@@ -0,0 +1,25 @@
+// Copyright 2022 The envd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ir
+
+import "fmt"
+
+func (g Graph) CacheID(filename string) string {
+	gpu := g.CUDA == nil || g.CUDNN == nil
+	if gpu {
+		return fmt.Sprintf("/%s-gpu-%s", g.CachePrefix, filename)
+	}
+	return fmt.Sprintf("/%s-cpu-%s", g.CachePrefix, filename)
+}
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -54,12 +54,13 @@ func GPUEnabled() bool {
 	return DefaultGraph.CUDA != nil
 }
 
-func Compile(ctx context.Context) (*llb.Definition, error) {
+func Compile(ctx context.Context, cachePrefix string) (*llb.Definition, error) {
 	w, err := compileui.New(ctx, os.Stdout, "auto")
 	if err != nil {
 		return nil, errors.Wrap(err, "failed to create compileui")
 	}
 	DefaultGraph.Writer = w
+	DefaultGraph.CachePrefix = cachePrefix
 	state, err := DefaultGraph.Compile()
 	if err != nil {
 		return nil, err

diff --git a/pkg/lang/ir/python.go b/pkg/lang/ir/python.go
@@ -48,7 +48,7 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {
 		Run(llb.Shlex(fmt.Sprintf(`sh -c "%s"`, cmd)), llb.WithCustomNamef("pip install %s",
 			strings.Join(g.PyPIPackages, " ")))
 	run.AddMount(cacheDir, llb.Scratch(),
-		llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared))
 	return run.Root()
 }
 

diff --git a/pkg/lang/ir/system.go b/pkg/lang/ir/system.go
@@ -77,9 +77,9 @@ func (g Graph) compileBuiltinSystemPackages(root llb.State) llb.State {
 		llb.WithCustomNamef("(built-in packages) apt-get install %s",
 			strings.Join(g.BuiltinSystemPackages, " ")))
 	run.AddMount(cacheDir, llb.Scratch(),
-		llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared))
 	run.AddMount(cacheLibDir, llb.Scratch(),
-		llb.AsPersistentCacheDir("/"+cacheLibDir, llb.CacheMountShared))
+		llb.AsPersistentCacheDir(g.CacheID(cacheLibDir), llb.CacheMountShared))
 
 	// TODO(gaocegege): Refactor user to a seperate stage.
 	res := run.
@@ -119,9 +119,9 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
 		llb.WithCustomNamef("(user-defined packages) apt-get install %s",
 			strings.Join(g.SystemPackages, " ")))
 	run.AddMount(cacheDir, llb.Scratch(),
-		llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
+		llb.AsPersistentCacheDir(g.CacheID(cacheDir), llb.CacheMountShared))
 	run.AddMount(cacheLibDir, llb.Scratch(),
-		llb.AsPersistentCacheDir("/"+cacheLibDir, llb.CacheMountShared))
+		llb.AsPersistentCacheDir(g.CacheID(cacheLibDir), llb.CacheMountShared))
 	return run.Root()
 }
 

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -39,7 +39,8 @@ type Graph struct {
 	Exec []string
 	*JupyterConfig
 
-	Writer compileui.Writer
+	Writer      compileui.Writer
+	CachePrefix string
 }
 
 type JupyterConfig struct {