Skip to content

Commit

Permalink
feat(CLI): fine-grained cache management (#155)
Browse files Browse the repository at this point in the history 
* feat(CLI): fine-grained cache management

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

* fix: Fix lint

Signed-off-by: Ce Gao <cegao@tensorchord.ai>
  • Loading branch information
@gaocegege
gaocegege committed May 17, 2022
1 parent 33cdee2 commit a7735f1
Show file tree
Hide file tree
Showing 4 changed files with 126 additions and 38 deletions.
82 changes: 67 additions & 15 deletions pkg/home/manager.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,23 +15,31 @@
package home

import (
"encoding/gob"
"os"
"path/filepath"
"sync"

"github.com/adrg/xdg"
"github.com/cockroachdb/errors"
"github.com/sirupsen/logrus"
"github.com/tensorchord/envd/pkg/util/fileutil"
)

type Manager interface {
CacheDir() string
MarkCache(string, bool) error
Cached(string) bool
ConfigFile() string
}

type generalManager struct {
cacheDir string
configFile string
cacheDir string
cacheStatusFile string
configFile string

// TODO(gaocegege): Abstract CacheManager.
cacheMap map[string]bool

logger *logrus.Entry
}
Expand All @@ -43,7 +51,9 @@ var (

func Initialize() error {
once.Do(func() {
defaultManager = &generalManager{}
defaultManager = &generalManager{
cacheMap: make(map[string]bool),
}
})
if err := defaultManager.init(); err != nil {
return err
Expand All @@ -63,23 +73,38 @@ func (m generalManager) ConfigFile() string {
return m.configFile
}

func (m generalManager) MarkCache(key string, cached bool) error {
m.cacheMap[key] = cached
return m.dumpCacheStatus()
}

func (m generalManager) Cached(key string) bool {
return m.cacheMap[key]
}

func (m *generalManager) dumpCacheStatus() error {
file, err := os.Create(m.cacheStatusFile)
if err != nil {
return errors.Wrap(err, "failed to open cache status file")
}
defer file.Close()

e := gob.NewEncoder(file)
if err := e.Encode(m.cacheMap); err != nil {
return errors.Wrap(err, "failed to encode cache map")
}
return nil
}

func (m *generalManager) init() error {
// Create $XDG_CONFIG_HOME/envd/config.envd
config, err := xdg.ConfigFile("envd/config.envd")
if err != nil {
return errors.Wrap(err, "failed to get config file")
}

_, err = os.Stat(config)
if err != nil {
if os.IsNotExist(err) {
logrus.WithField("config", config).Info("Creating config file")
if _, err := os.Create(config); err != nil {
return errors.Wrap(err, "failed to create config file")
}
} else {
return errors.Wrap(err, "failed to stat config file")
}
if err := fileutil.CreateIfNotExist(config); err != nil {
return errors.Wrap(err, "failed to create config file")
}
m.configFile = config

Expand All @@ -90,9 +115,36 @@ func (m *generalManager) init() error {
}
m.cacheDir = filepath.Join(xdg.CacheHome, "envd")

m.cacheStatusFile = filepath.Join(m.cacheDir, "cache.status")
_, err = os.Stat(m.cacheStatusFile)
if err != nil {
if os.IsNotExist(err) {
logrus.WithField("filename", m.cacheStatusFile).Debug("Creating file")
if _, err := os.Create(m.cacheStatusFile); err != nil {
return errors.Wrap(err, "failed to create file")
}
if err := m.dumpCacheStatus(); err != nil {
return errors.Wrap(err, "failed to dump cache status")
}
} else {
return errors.Wrap(err, "failed to stat file")
}
}

file, err := os.Open(m.cacheStatusFile)
if err != nil {
return errors.Wrap(err, "failed to open cache status file")
}
defer file.Close()
e := gob.NewDecoder(file)
if err := e.Decode(&m.cacheMap); err != nil {
return errors.Wrap(err, "failed to decode cache map")
}

m.logger = logrus.WithFields(logrus.Fields{
"cacheDir": m.cacheDir,
"config": m.configFile,
"cache-dir": m.cacheDir,
"config": m.configFile,
"cache-status": m.cacheStatusFile,
})

m.logger.Debug("home manager initialized")
Expand Down
39 changes: 23 additions & 16 deletions pkg/lang/ir/graph.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,9 @@ func (g Graph) Compile() (llb.State, error) {
if err != nil {
return llb.State{}, errors.Wrap(err, "failed to compile shell")
}
diffShellStage := llb.Diff(builtinSystemStage, shellStage)
pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage))
systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage))
diffShellStage := llb.Diff(builtinSystemStage, shellStage, llb.WithCustomName("install shell"))
pypiStage := llb.Diff(builtinSystemStage, g.compilePyPIPackages(builtinSystemStage), llb.WithCustomName("install PyPI packages"))
systemStage := llb.Diff(builtinSystemStage, g.compileSystemPackages(builtinSystemStage), llb.WithCustomName("install system packages"))
sshStage := g.copyenvdSSHServer()

vscodeStage, err := g.compileVSCode()
Expand All @@ -105,11 +105,11 @@ func (g Graph) Compile() (llb.State, error) {
if vscodeStage != nil {
merged = llb.Merge([]llb.State{
builtinSystemStage, systemStage, pypiStage, sshStage, *vscodeStage, diffShellStage,
})
}, llb.WithCustomName("merging all components into one"))
} else {
merged = llb.Merge([]llb.State{
builtinSystemStage, systemStage, pypiStage, sshStage, diffShellStage,
})
}, llb.WithCustomName("merging all components into one"))
}

// TODO(gaocegege): Support order-based exec.
Expand Down Expand Up @@ -164,7 +164,8 @@ func (g Graph) compilePyPIPackages(root llb.State) llb.State {

cacheDir := "/home/envd/.cache/pip"

run := root.Run(llb.Shlex(sb.String()))
run := root.Run(llb.Shlex(sb.String()), llb.WithCustomNamef("pip install %s",
strings.Join(g.PyPIPackages, " ")))
run.AddMount(cacheDir, llb.Scratch(),
llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
return run.Root()
Expand Down Expand Up @@ -193,7 +194,9 @@ func (g Graph) compileBuiltinSystemPackages(root llb.State) llb.State {
cacheDir := "/var/cache/apt"
cacheLibDir := "/var/lib/apt"

run := root.Run(llb.Shlex(sb.String()))
run := root.Run(llb.Shlex(sb.String()),
llb.WithCustomNamef("(built-in packages) apt-get install %s",
strings.Join(g.BuiltinSystemPackages, " ")))
run.AddMount(cacheDir, llb.Scratch(),
llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
run.AddMount(cacheLibDir, llb.Scratch(),
Expand All @@ -218,7 +221,9 @@ func (g Graph) compileSystemPackages(root llb.State) llb.State {
cacheDir := "/var/cache/apt"
cacheLibDir := "/var/lib/apt"

run := root.Run(llb.Shlex(sb.String()))
run := root.Run(llb.Shlex(sb.String()),
llb.WithCustomNamef("(user-defined packages) apt-get install %s",
strings.Join(g.SystemPackages, " ")))
run.AddMount(cacheDir, llb.Scratch(),
llb.AsPersistentCacheDir("/"+cacheDir, llb.CacheMountShared))
run.AddMount(cacheLibDir, llb.Scratch(),
Expand All @@ -231,7 +236,7 @@ func (g Graph) copyenvdSSHServer() llb.State {
run := llb.Image(viper.GetString(flag.FlagSSHImage)).
File(llb.Copy(llb.Image(viper.GetString(flag.FlagSSHImage)),
"usr/bin/envd-ssh", "/var/envd/bin/envd-ssh",
&llb.CopyInfo{CreateDestPath: true}))
&llb.CopyInfo{CreateDestPath: true}), llb.WithCustomName("install envd-ssh"))
return run
}

Expand All @@ -251,10 +256,11 @@ func (g Graph) compileVSCode() (*llb.State, error) {
ext := llb.Scratch().File(llb.Copy(llb.Local(flag.FlagCacheDir),
vscodeClient.PluginPath(p),
"/home/envd/.vscode-server/extensions/"+p.String(),
&llb.CopyInfo{CreateDestPath: true}))
&llb.CopyInfo{CreateDestPath: true}),
llb.WithCustomNamef("install vscode plugin %s", p.String()))
inputs = append(inputs, ext)
}
layer := llb.Merge(inputs)
layer := llb.Merge(inputs, llb.WithCustomName("merging plugins for vscode"))
return &layer, nil
}

Expand All @@ -268,9 +274,9 @@ func (g Graph) compileUbuntuAPT(root llb.State) llb.State {
if g.UbuntuAPTSource != nil {
logrus.WithField("source", *g.UbuntuAPTSource).Debug("using custom APT source")
aptSource := llb.Scratch().
File(llb.Mkdir(filepath.Dir(aptSourceFilePath), 0755, llb.WithParents(true))).
File(llb.Mkfile(aptSourceFilePath, 0644, []byte(*g.UbuntuAPTSource)))
return llb.Merge([]llb.State{root, aptSource})
File(llb.Mkdir(filepath.Dir(aptSourceFilePath), 0755, llb.WithParents(true)), llb.WithCustomName("create apt source dir")).
File(llb.Mkfile(aptSourceFilePath, 0644, []byte(*g.UbuntuAPTSource)), llb.WithCustomName("create apt source file"))
return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add apt source"))
}
return root
}
Expand All @@ -282,7 +288,7 @@ func (g Graph) compilePyPIMirror(root llb.State) llb.State {
aptSource := llb.Scratch().
File(llb.Mkdir(filepath.Dir(pypiMirrorFilePath), 0755, llb.WithParents(true))).
File(llb.Mkfile(pypiMirrorFilePath, 0644, []byte(content)))
return llb.Merge([]llb.State{root, aptSource})
return llb.Merge([]llb.State{root, aptSource}, llb.WithCustomName("add PyPI mirror"))
}
return root
}
Expand All @@ -300,7 +306,8 @@ func (g Graph) compileZSH(root llb.State) (llb.State, error) {
File(llb.Copy(llb.Local(flag.FlagCacheDir), "oh-my-zsh", "/home/envd/.oh-my-zsh",
&llb.CopyInfo{CreateDestPath: true})).
File(llb.Mkfile(installPath, 0644, []byte(m.InstallScript())))
run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)))
run := zshStage.Run(llb.Shlex(fmt.Sprintf("bash %s", installPath)),
llb.WithCustomName("install oh-my-zsh"))
return run.Root(), nil
}

Expand Down
21 changes: 15 additions & 6 deletions pkg/shell/zsh.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import (
"path/filepath"

"github.com/go-git/go-git/v5"
"github.com/pkg/errors"
"github.com/sirupsen/logrus"
"github.com/tensorchord/envd/pkg/home"
"github.com/tensorchord/envd/pkg/util/fileutil"
Expand All @@ -45,26 +46,34 @@ func (m generalManager) InstallScript() string {
}

func (m generalManager) DownloadOrCache() (bool, error) {
if ok, err := fileutil.DirExists(m.OHMyZSHDir()); err != nil {
return false, err
} else if ok {
if home.GetManager().Cached("oh-my-zsh") {
logrus.WithFields(logrus.Fields{
"cache-dir": m.OHMyZSHDir(),
}).Debug("found cached oh-my-zsh")
}).Debug("oh-my-zsh already exists in cache")
return true, nil
}
url := "https://github.com/ohmyzsh/ohmyzsh"
l := logrus.WithFields(logrus.Fields{
"cache-dir": m.OHMyZSHDir(),
"URL": url,
})
l.Debug("downloading oh-my-zsh")

// Cleanup the cache dir.
if fileutil.RemoveAll(m.OHMyZSHDir()) != nil {
return false, errors.New("failed to remove oh-my-zsh dir")
}
l.Debug("cache miss, downloading oh-my-zsh")
_, err := git.PlainClone(m.OHMyZSHDir(), false, &git.CloneOptions{
URL: url,
URL: url,
Depth: 1,
})
if err != nil {
return false, err
}

if err := home.GetManager().MarkCache("oh-my-zsh", true); err != nil {
return false, errors.Wrap(err, "failed to update cache status")
}
l.Debug("oh-my-zsh is downloaded")
return false, nil
}
Expand Down
22 changes: 21 additions & 1 deletion pkg/util/fileutil/file.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@ import (
"path/filepath"

"github.com/cockroachdb/errors"
"github.com/sirupsen/logrus"
)

// FileExists returns true if the file exists
Expand All @@ -35,7 +36,11 @@ func FileExists(filename string) (bool, error) {
return !info.IsDir(), nil
}

// DirExists returns true if the directory exists
func RemoveAll(dirname string) error {
return os.RemoveAll(dirname)
}

// DirExists returns true if the directory exists.
func DirExists(filename string) (bool, error) {
info, err := os.Stat(filename)
if err != nil {
Expand All @@ -47,6 +52,21 @@ func DirExists(filename string) (bool, error) {
return info.IsDir(), nil
}

func CreateIfNotExist(f string) error {
_, err := os.Stat(f)
if err != nil {
if os.IsNotExist(err) {
logrus.WithField("filename", f).Debug("Creating file")
if _, err := os.Create(f); err != nil {
return errors.Wrap(err, "failed to create file")
}
} else {
return errors.Wrap(err, "failed to stat file")
}
}
return nil
}

func CWD() (string, error) {
return os.Getwd()
}
Expand Down

0 comments on commit a7735f1

Please sign in to comment.