diff --git a/examples/include_pkg/build.envd b/examples/include_pkg/build.envd new file mode 100644 index 000000000..6adc899bd --- /dev/null +++ b/examples/include_pkg/build.envd @@ -0,0 +1,6 @@ +envdlib = include("https://github.com/kemingy/envdlib") + + +def build(): + base(os="ubuntu20.04", language="python") + envdlib.tensorboard(8888) diff --git a/pkg/lang/frontend/starlark/interpreter.go b/pkg/lang/frontend/starlark/interpreter.go index 9dec9e7b0..02cacb6b5 100644 --- a/pkg/lang/frontend/starlark/interpreter.go +++ b/pkg/lang/frontend/starlark/interpreter.go @@ -17,12 +17,14 @@ package starlark import ( "bytes" "hash/fnv" + "io/fs" "io/ioutil" + "path/filepath" "strconv" + "strings" "github.com/cockroachdb/errors" "github.com/sirupsen/logrus" - "go.starlark.net/repl" "go.starlark.net/starlark" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/config" @@ -31,6 +33,7 @@ import ( "github.com/tensorchord/envd/pkg/lang/frontend/starlark/io" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/runtime" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/universe" + "github.com/tensorchord/envd/pkg/util/fileutil" ) type Interpreter interface { @@ -38,21 +41,25 @@ type Interpreter interface { ExecFile(filename string, funcname string) (interface{}, error) } +type entry struct { + globals starlark.StringDict + err error +} + // generalInterpreter is the interpreter implementation for Starlark. // Please refer to https://github.com/google/starlark-go type generalInterpreter struct { - *starlark.Thread predeclared starlark.StringDict buildContextDir string + cache map[string]*entry } func NewInterpreter(buildContextDir string) Interpreter { // Register envd rules and built-in variables to Starlark. - universe.RegisterenvdRules() + universe.RegisterEnvdRules() universe.RegisterBuildContext(buildContextDir) return &generalInterpreter{ - Thread: &starlark.Thread{Load: repl.MakeLoad()}, predeclared: starlark.StringDict{ "install": install.Module, "config": config.Module, @@ -61,37 +68,84 @@ func NewInterpreter(buildContextDir string) Interpreter { "data": data.Module, }, buildContextDir: buildContextDir, + cache: make(map[string]*entry), } } -func GetEnvdProgramHash(filename string) (string, error) { - envdSrc, err := ioutil.ReadFile(filename) - if err != nil { - return "", err +func (s *generalInterpreter) NewThread(module string) *starlark.Thread { + thread := &starlark.Thread{ + Name: module, + Load: s.load, } - // No Check builtin or predeclared for now - funcAlwaysHas := func(x string) bool { - return true + return thread +} + +func (s *generalInterpreter) load(thread *starlark.Thread, module string) (starlark.StringDict, error) { + return s.exec(thread, module) +} + +func (s *generalInterpreter) exec(thread *starlark.Thread, module string) (starlark.StringDict, error) { + e, ok := s.cache[module] + if e != nil { + return e.globals, e.err } - _, prog, err := starlark.SourceProgram(filename, envdSrc, funcAlwaysHas) - if err != nil { - return "", err + if ok { + return nil, errors.Newf("Detect cycling import during parsing %s", module) } - buf := new(bytes.Buffer) - err = prog.Write(buf) - if err != nil { - return "", err + + s.cache[module] = nil + + if !strings.HasPrefix(module, universe.GitPrefix) { + var data interface{} + globals, err := starlark.ExecFile(thread, module, data, s.predeclared) + e = &entry{globals, err} + } else { + // exec remote git repo + url := module[len(universe.GitPrefix):] + path, err := fileutil.DownloadOrUpdateGitRepo(url) + if err != nil { + return nil, err + } + globals, err := s.loadGitModule(thread, path) + e = &entry{globals, err} } - h := fnv.New64a() - h.Write(buf.Bytes()) - hashsum := h.Sum64() - return strconv.FormatUint(hashsum, 16), nil + + return e.globals, e.err +} + +func (s *generalInterpreter) loadGitModule(thread *starlark.Thread, path string) (globals starlark.StringDict, err error) { + var src interface{} + globals = starlark.StringDict{} + logger := logrus.WithField("file", thread.Name) + logger.Debugf("load git module from: %s", path) + err = filepath.WalkDir(path, func(path string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if d.IsDir() || !strings.HasSuffix(d.Name(), ".envd") { + return nil + } + dict, err := starlark.ExecFile(thread, path, src, s.predeclared) + if err != nil { + return err + } + for key, val := range dict { + if _, exist := globals[key]; exist { + return errors.Newf("found duplicated object name: %s in %s", key, path) + } + if !strings.HasPrefix(key, "_") { + globals[key] = val + } + } + return nil + }) + return } func (s generalInterpreter) ExecFile(filename string, funcname string) (interface{}, error) { logrus.WithField("filename", filename).Debug("interprete the file") - var src interface{} - globals, err := starlark.ExecFile(s.Thread, filename, src, s.predeclared) + thread := s.NewThread(filename) + globals, err := s.exec(thread, filename) if err != nil { return nil, err } @@ -100,7 +154,7 @@ func (s generalInterpreter) ExecFile(filename string, funcname string) (interfac if globals.Has(funcname) { buildVar := globals[funcname] if fn, ok := buildVar.(*starlark.Function); ok { - _, err := starlark.Call(s.Thread, fn, nil, nil) + _, err := starlark.Call(thread, fn, nil, nil) if err != nil { return nil, errors.Wrapf(err, "Exception when exec %s func", funcname) } @@ -116,5 +170,30 @@ func (s generalInterpreter) ExecFile(filename string, funcname string) (interfac } func (s generalInterpreter) Eval(script string) (interface{}, error) { - return starlark.ExecFile(s.Thread, "", script, s.predeclared) + thread := s.NewThread(script) + return starlark.ExecFile(thread, "", script, s.predeclared) +} + +func GetEnvdProgramHash(filename string) (string, error) { + envdSrc, err := ioutil.ReadFile(filename) + if err != nil { + return "", err + } + // No Check builtin or predeclared for now + funcAlwaysHas := func(x string) bool { + return true + } + _, prog, err := starlark.SourceProgram(filename, envdSrc, funcAlwaysHas) + if err != nil { + return "", err + } + buf := new(bytes.Buffer) + err = prog.Write(buf) + if err != nil { + return "", err + } + h := fnv.New64a() + h.Write(buf.Bytes()) + hashsum := h.Sum64() + return strconv.FormatUint(hashsum, 16), nil } diff --git a/pkg/lang/frontend/starlark/universe/const.go b/pkg/lang/frontend/starlark/universe/const.go index bd3bb6840..7e2b74e68 100644 --- a/pkg/lang/frontend/starlark/universe/const.go +++ b/pkg/lang/frontend/starlark/universe/const.go @@ -19,4 +19,7 @@ const ( ruleShell = "shell" ruleRun = "run" ruleGitConfig = "git_config" + ruleInclude = "include" + + GitPrefix = "git@" ) diff --git a/pkg/lang/frontend/starlark/universe/universe.go b/pkg/lang/frontend/starlark/universe/universe.go index 3c33fd7d2..71b75dede 100644 --- a/pkg/lang/frontend/starlark/universe/universe.go +++ b/pkg/lang/frontend/starlark/universe/universe.go @@ -15,8 +15,11 @@ package universe import ( + "fmt" + "github.com/sirupsen/logrus" "go.starlark.net/starlark" + "go.starlark.net/starlarkstruct" "github.com/tensorchord/envd/pkg/lang/frontend/starlark/builtin" "github.com/tensorchord/envd/pkg/lang/ir" @@ -26,12 +29,13 @@ var ( logger = logrus.WithField("frontend", "starlark") ) -// RegisterenvdRules registers built-in envd rules into the global namespace. -func RegisterenvdRules() { +// RegisterEnvdRules registers built-in envd rules into the global namespace. +func RegisterEnvdRules() { starlark.Universe[ruleBase] = starlark.NewBuiltin(ruleBase, ruleFuncBase) starlark.Universe[ruleShell] = starlark.NewBuiltin(ruleShell, ruleFuncShell) starlark.Universe[ruleRun] = starlark.NewBuiltin(ruleRun, ruleFuncRun) starlark.Universe[ruleGitConfig] = starlark.NewBuiltin(ruleGitConfig, ruleFuncGitConfig) + starlark.Universe[ruleInclude] = starlark.NewBuiltin(ruleInclude, ruleFuncInclude) } func RegisterBuildContext(buildContextDir string) { @@ -117,3 +121,25 @@ func ruleFuncGitConfig(thread *starlark.Thread, _ *starlark.Builtin, err := ir.Git(nameStr, emailStr, editorStr) return starlark.None, err } + +func ruleFuncInclude(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var gitRepo string + + if err := starlark.UnpackArgs(ruleInclude, + args, kwargs, "git?", &gitRepo); err != nil { + return nil, err + } + + logger.Debugf("rule `%s` is invoked, git=%s", ruleInclude, gitRepo) + + globals, err := thread.Load(thread, fmt.Sprintf("%s%s", GitPrefix, gitRepo)) + if err != nil { + return nil, err + } + module := &starlarkstruct.Module{ + Name: gitRepo, + Members: globals, + } + return module, nil +} diff --git a/pkg/util/fileutil/file.go b/pkg/util/fileutil/file.go index 1ece58a2d..7c5336702 100644 --- a/pkg/util/fileutil/file.go +++ b/pkg/util/fileutil/file.go @@ -22,12 +22,14 @@ import ( "strings" "github.com/cockroachdb/errors" + "github.com/go-git/go-git/v5" "github.com/sirupsen/logrus" ) var ( - DefaultConfigDir string - DefaultCacheDir string + DefaultConfigDir string + DefaultCacheDir string + DefaultEnvdLibDir string ) func init() { @@ -37,6 +39,7 @@ func init() { } DefaultConfigDir = filepath.Join(home, ".config", "envd") DefaultCacheDir = filepath.Join(home, ".cache", "envd") + DefaultEnvdLibDir = filepath.Join(DefaultCacheDir, "envdlib") } // FileExists returns true if the file exists @@ -139,3 +142,42 @@ func validateAndJoin(dir, file string) (string, error) { } return filepath.Join(dir, file), nil } + +// DownloadOrUpdateGitRepo downloads (if not exist) or update (if exist) +func DownloadOrUpdateGitRepo(url string) (path string, err error) { + logger := logrus.WithField("git", url) + path = filepath.Join(DefaultEnvdLibDir, strings.ReplaceAll(url, "/", "_")) + var repo *git.Repository + exist, err := DirExists(path) + if err != nil { + return + } + if !exist { + logger.Debugf("clone repo to %s", path) + // check https://github.com/go-git/go-git/issues/305 + _, err = git.PlainClone(path, false, &git.CloneOptions{ + URL: url, + }) + if err != nil { + return + } + } else { + logger.Debugf("repo already exists in %s", path) + repo, err = git.PlainOpen(path) + if err != nil { + return + } + var wt *git.Worktree + wt, err = repo.Worktree() + if err != nil { + return + } + logger.Debug("try to pull latest") + err = wt.Pull(&git.PullOptions{}) + if err != nil && errors.Is(err, git.NoErrAlreadyUpToDate) { + return + } + } + + return path, nil +}