add entrypoint in custom image (#739)

* WIP: add entrypoint Signed-off-by: Keming <kemingyang@tensorchord.ai> * fix entrypoint Signed-off-by: Keming <kemingyang@tensorchord.ai> * add test Signed-off-by: Keming <kemingyang@tensorchord.ai> * make entrypoint public Signed-off-by: Keming <kemingyang@tensorchord.ai>
tensorchord · Aug 9, 2022 · b492e83 · b492e83
1 parent 4df956b
commit b492e83
Show file tree

Hide file tree

Showing 9 changed files with 38 additions and 10 deletions.
diff --git a/e2e/testdata/custom-image-test/build.envd b/e2e/testdata/custom-image-test/build.envd
@@ -1,5 +1,6 @@
 def build():
-    base(language="python", image="python:3.8")
+    base(language="python", image="python:3.9-slim")
     install.python_packages(name=[
         "via",
     ])
+    config.entrypoint(["date", "-u"])
diff --git a/examples/custom-image/build.envd b/examples/custom-image/build.envd
@@ -1,5 +1,6 @@
 def build():
-    base(language="python", image="python:3.8")
+    base(language="python", image="python:3.9-slim")
     install.python_packages(name=[
       "via",
     ])
+    config.entrypoint(["date", "-u"])
diff --git a/pkg/app/up.go b/pkg/app/up.go
@@ -216,7 +216,7 @@ func up(clicontext *cli.Context) error {
 
 	err = dockerClient.CleanEnvdIfExists(clicontext.Context, ctr, force)
 	if err != nil {
-		return errors.Wrap(err, "failed to start the envd environment")
+		return errors.Wrap(err, "failed to clean the envd environment")
 	}
 	containerID, containerIP, err := dockerClient.StartEnvd(clicontext.Context,
 		tag, ctr, buildContext, gpu, numGPUs, sshPortInHost, *ir.DefaultGraph, clicontext.Duration("timeout"),

diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go
@@ -209,10 +209,11 @@ func (b generalBuilder) imageConfig(ctx context.Context) (string, error) {
 	}
 	labels[types.ImageLabelContext] = b.BuildContextDir
 
-	ep, err := ir.Entrypoint(b.BuildContextDir)
+	ep, err := ir.CompileEntrypoint(b.BuildContextDir)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to get entrypoint")
 	}
+	b.logger.Debugf("final entrypoint: {%s}\n", ep)
 	data, err := ImageConfigStr(labels, ports, ep)
 	if err != nil {
 		return "", errors.Wrap(err, "failed to get image config")

diff --git a/pkg/lang/frontend/starlark/config/config.go b/pkg/lang/frontend/starlark/config/config.go
@@ -42,6 +42,7 @@ var Module = &starlarkstruct.Module{
 		"julia_pkg_server": starlark.NewBuiltin(
 			ruleJuliaPackageServer, ruleFuncJuliaPackageServer),
 		"rstudio_server": starlark.NewBuiltin(ruleRStudioServer, ruleFuncRStudioServer),
+		"entrypoint":     starlark.NewBuiltin(ruleEntrypoint, ruleFuncEntrypoint),
 	},
 }
 
@@ -196,3 +197,23 @@ func ruleFuncCondaChannel(thread *starlark.Thread, _ *starlark.Builtin,
 
 	return starlark.None, nil
 }
+
+func ruleFuncEntrypoint(thread *starlark.Thread, _ *starlark.Builtin,
+	args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var argv *starlark.List
+
+	if err := starlark.UnpackArgs(ruleEntrypoint, args, kwargs, "name", &argv); err != nil {
+		return nil, err
+	}
+
+	argList := []string{}
+	if argv != nil {
+		for i := 0; i < argv.Len(); i++ {
+			argList = append(argList, argv.Index(i).(starlark.String).GoString())
+		}
+	}
+
+	logger.Debugf("user defined entrypoints: {%s}\n", argList)
+	ir.Entrypoint(argList)
+	return starlark.None, nil
+}
diff --git a/pkg/lang/frontend/starlark/config/const.go b/pkg/lang/frontend/starlark/config/const.go
@@ -23,4 +23,5 @@ const (
 	ruleGPU                = "config.gpu"
 	ruleJuliaPackageServer = "config.julia_pkg_server"
 	ruleRStudioServer      = "config.rstudio_server"
+	ruleEntrypoint         = "config.entrypoint"
 )
diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go
@@ -91,8 +91,8 @@ func ExposedPorts() (map[string]struct{}, error) {
 	return DefaultGraph.ExposedPorts()
 }
 
-func Entrypoint(buildContextDir string) ([]string, error) {
-	return DefaultGraph.Entrypoint(buildContextDir)
+func CompileEntrypoint(buildContextDir string) ([]string, error) {
+	return DefaultGraph.GetEntrypoint(buildContextDir)
 }
 
 func (g Graph) GPUEnabled() bool {
@@ -149,11 +149,9 @@ func (g Graph) DefaultCacheImporter() (*string, error) {
 	return &res, nil
 }
 
-func (g Graph) Entrypoint(buildContextDir string) ([]string, error) {
-	// Do not set entrypoint if the image is customized.
+func (g Graph) GetEntrypoint(buildContextDir string) ([]string, error) {
 	if g.Image != nil {
-		logrus.Debug("skip entrypoint because the image is customized")
-		return []string{}, nil
+		return g.Entrypoint, nil
 	}
 
 	ep := []string{

diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go
@@ -204,6 +204,10 @@ func Mount(src, dest string) {
 	})
 }
 
+func Entrypoint(args []string) {
+	DefaultGraph.Entrypoint = append(DefaultGraph.Entrypoint, args...)
+}
+
 func RuntimeCommands(commands map[string]string) {
 	for k, v := range commands {
 		DefaultGraph.RuntimeCommands[k] = v

diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go
@@ -53,6 +53,7 @@ type Graph struct {
 	Exec            []string
 	Copy            []CopyInfo
 	Mount           []MountInfo
+	Entrypoint      []string
 	RuntimeCommands map[string]string
 
 	*JupyterConfig