From e0984bb112d21c4aab36a48cd3a04b9cb27f59a6 Mon Sep 17 00:00:00 2001 From: Ce Gao Date: Fri, 27 May 2022 14:19:17 +0800 Subject: [PATCH] feat: Add prefix for cache id (#204) * feat: Add prefix for cache id Signed-off-by: Ce Gao * fix: Add license Signed-off-by: Ce Gao --- pkg/builder/builder.go | 3 ++- pkg/lang/ir/cache.go | 25 +++++++++++++++++++++++++ pkg/lang/ir/compile.go | 3 ++- pkg/lang/ir/python.go | 2 +- pkg/lang/ir/system.go | 8 ++++---- pkg/lang/ir/types.go | 3 ++- 6 files changed, 36 insertions(+), 8 deletions(-) create mode 100644 pkg/lang/ir/cache.go diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go index ba37a846e..0bfc90198 100644 --- 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 new file mode 100644 index 000000000..708f1cc12 --- /dev/null +++ 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 index 34ec2c25d..7549a846f 100644 --- 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 index 034bea17d..f5723469c 100644 --- 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 index 5e1975479..b8aa07e7d 100644 --- 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 index 4a9b606ce..9de932dce 100644 --- 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 {