feat(lang): support include other git repo for envd functions/variabl…

…es (#808)

* feat(lang): support include other git repo for envd functions/variables

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* format

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* fix thread

Signed-off-by: Keming <kemingyang@tensorchord.ai>

* add git and example

Signed-off-by: Keming <kemingyang@tensorchord.ai>

Signed-off-by: Keming <kemingyang@tensorchord.ai>

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)

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,28 +33,33 @@ 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 {
  Eval(script string) (interface{}, error)
  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
 }

pkg/lang/frontend/starlark/universe/const.go

@@ -19,4 +19,7 @@ const (
  ruleShell = "shell"
  ruleRun = "run"
  ruleGitConfig = "git_config"
+ ruleInclude = "include"
+
+ GitPrefix = "git@"
 )

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
+}

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
+}