feat: Support specify build target (#497)

* support specify function Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * change example Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix test Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix naming Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix lint Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai> * fix Signed-off-by: Jinjing.Zhou <allenzhou@tensorchord.ai>
tensorchord · Jul 5, 2022 · 6e9e44d · 6e9e44d
1 parent e443784
commit 6e9e44d
Show file tree

Hide file tree

Showing 7 changed files with 48 additions and 26 deletions.
diff --git a/examples/mnist/README.md b/examples/mnist/README.md
@@ -13,7 +13,7 @@ Then you can open jupyter notebook at [`http://localhost:8888`](http://localhost
 ## Quick start (GPU)
 
 ```bash
-$ envd up -f build_gpu.envd
+$ envd up -f :build_gpu # Call build_gpu function in build.envd
 ```
 
 Also you can use `-f` option to specify the file to build

diff --git a/examples/mnist/build_gpu.envd b/examples/mnist/build_gpu.envd
diff --git a/pkg/app/build.go b/pkg/app/build.go
@@ -48,10 +48,10 @@ To build and push the image to a registry:
  DefaultText: "PROJECT:dev",
  },
  &cli.PathFlag{
- Name: "file",
- Usage: "Name of the build.envd",
+ Name: "from",
+ Usage: "Function to execute, format `file:func`",
  Aliases: []string{"f"},
- Value: "build.envd",
+ Value: "build.envd:build",
  },
  &cli.PathFlag{
  Name: "path",
@@ -82,7 +82,11 @@ func build(clicontext *cli.Context) error {
  return errors.Wrap(err, "failed to get absolute path of the build context")
  }
 
- manifest, err := filepath.Abs(filepath.Join(buildContext, clicontext.Path("file")))
+ filename, funcname, err := builder.ParseFromStr(clicontext.String("from"))
+ if err != nil {
+ return err
+ }
+ manifest, err := filepath.Abs(filepath.Join(buildContext, filename))
  if err != nil {
  return errors.Wrap(err, "failed to get absolute path of the build file")
  }
@@ -111,7 +115,7 @@ func build(clicontext *cli.Context) error {
  "output": output,
  }).Debug("starting build command")
  builder, err := builder.New(clicontext.Context, cfg,
- manifest, buildContext, tag, output, debug)
+ manifest, funcname, buildContext, tag, output, debug)
  if err != nil {
  return errors.Wrap(err, "failed to create the builder")
  }

diff --git a/pkg/app/up.go b/pkg/app/up.go
@@ -61,10 +61,10 @@ var CommandUp = &cli.Command{
  Aliases: []string{"v"},
  },
  &cli.PathFlag{
- Name: "file",
- Usage: "Name of the build.envd",
+ Name: "from",
+ Usage: "Function to execute, format `file:func`",
  Aliases: []string{"f"},
- Value: "build.envd",
+ Value: "build.envd:build",
  },
  &cli.PathFlag{
  Name: "private-key",
@@ -105,8 +105,12 @@ func up(clicontext *cli.Context) error {
  if err != nil {
  return errors.Wrap(err, "failed to get absolute path of the build context")
  }
+ filename, funcname, err := builder.ParseFromStr(clicontext.String("from"))
+ if err != nil {
+ return err
+ }
 
- manifest, err := filepath.Abs(filepath.Join(buildContext, clicontext.Path("file")))
+ manifest, err := filepath.Abs(filepath.Join(buildContext, filename))
  if err != nil {
  return errors.Wrap(err, "failed to get absolute path of the build file")
  }
@@ -137,7 +141,7 @@ func up(clicontext *cli.Context) error {
  logger.Debug("starting up command")
  debug := clicontext.Bool("debug")
  output := ""
- builder, err := builder.New(clicontext.Context, config, manifest,
+ builder, err := builder.New(clicontext.Context, config, manifest, funcname,
  buildContext, tag, output, debug)
  if err != nil {
  return errors.Wrap(err, "failed to create the builder")

diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -50,6 +50,7 @@ type generalBuilder struct {
  progressMode string
  tag string
  buildContextDir string
+ buildfuncname string
 
  entries []client.ExportEntry
 
@@ -58,7 +59,7 @@ type generalBuilder struct {
  buildkitd.Client
 }
 
-func New(ctx context.Context, configFilePath, manifestFilePath,
+func New(ctx context.Context, configFilePath, manifestFilePath, funcname,
  buildContextDir, tag string, output string, debug bool) (Builder, error) {
  entries, err := parseOutput(output)
  if err != nil {
@@ -84,6 +85,7 @@ func New(ctx context.Context, configFilePath, manifestFilePath,
 
  b := &generalBuilder{
  manifestFilePath: manifestFilePath,
+ buildfuncname: funcname,
  configFilePath: configFilePath,
  buildContextDir: buildContextDir,
  entries: entries,
@@ -141,7 +143,7 @@ func (b generalBuilder) interpret() error {
  return errors.Wrap(err, "failed to exec starlark file")
  }
 
- if _, err := b.ExecFile(b.manifestFilePath, "build"); err != nil {
+ if _, err := b.ExecFile(b.manifestFilePath, b.buildfuncname); err != nil {
  return errors.Wrap(err, "failed to exec starlark file")
  }
  return nil

diff --git a/pkg/builder/builder_test.go b/pkg/builder/builder_test.go
@@ -55,6 +55,7 @@ var _ = Describe("Builder", func() {
  configFilePath: configFilePath,
  progressMode: "auto",
  tag: tag,
+ buildfuncname: "build",
  logger: logrus.WithFields(logrus.Fields{
  "tag": tag,
  }),

diff --git a/pkg/builder/util.go b/pkg/builder/util.go
@@ -29,6 +29,11 @@ import (
  "github.com/sirupsen/logrus"
 )
 
+const (
+ defaultFile = "build.envd"
+ defaultFunc = "build"
+)
+
 func ImageConfigStr(labels map[string]string,
  ports map[string]struct{}, entrypoint []string) (string, error) {
  pl := platforms.Normalize(platforms.DefaultSpec())
@@ -172,3 +177,22 @@ func resolveExporterDest(exporter, dest string) (func(map[string]string) (io.Wri
  return nil, "", nil
  }
 }
+
+func ParseFromStr(fromStr string) (string, string, error) {
+ filename := defaultFile
+ funcname := defaultFunc
+ if strings.Contains(fromStr, ":") {
+ fromArr := strings.Split(fromStr, ":")
+
+ if len(fromArr) != 2 {
+ return "", "", errors.New("invalid from format, expected `file:func`")
+ }
+ if fromArr[0] != "" {
+ filename = fromArr[0]
+ }
+ if fromArr[1] != "" {
+ funcname = fromArr[1]
+ }
+ }
+ return filename, funcname, nil
+}