From 4d48767ce033a3d60fac74d810ea546f741bc174 Mon Sep 17 00:00:00 2001 From: Keming Date: Wed, 17 Aug 2022 17:43:23 +0800 Subject: [PATCH] feat: add runtime environments (#787) Signed-off-by: Keming Signed-off-by: Keming --- envd/api/runtime/__init__.py | 13 +++++++++++ examples/python-basic/build.envd | 1 + pkg/builder/builder.go | 5 +++- pkg/builder/util.go | 6 ++--- pkg/lang/frontend/starlark/runtime/const.go | 1 + pkg/lang/frontend/starlark/runtime/runtime.go | 23 +++++++++++++++++++ pkg/lang/ir/compile.go | 13 +++++++++++ pkg/lang/ir/interface.go | 6 +++++ pkg/lang/ir/types.go | 1 + 9 files changed, 65 insertions(+), 4 deletions(-) diff --git a/envd/api/runtime/__init__.py b/envd/api/runtime/__init__.py index 6901c995a..df594344a 100644 --- a/envd/api/runtime/__init__.py +++ b/envd/api/runtime/__init__.py @@ -61,3 +61,16 @@ def daemon(commands: List[List[str]]): ]) ``` """ + + +def environ(env: Dict[str, str]): + """Add runtime environments + + Args: + env (Dict[str, str]): environment name to value + + Example usage: + ``` + runtime.environ(env={"ENVD_MODE": "DEV"}) + ``` + """ diff --git a/examples/python-basic/build.envd b/examples/python-basic/build.envd index ec29788e9..30d1a81d6 100644 --- a/examples/python-basic/build.envd +++ b/examples/python-basic/build.envd @@ -15,4 +15,5 @@ def build(): runtime.command(commands={ "test": "ls /", }) + runtime.environ(env={"ENVD_MODE": "DEV"}) jupyter_lab() diff --git a/pkg/builder/builder.go b/pkg/builder/builder.go index 65d68e520..7012436f9 100644 --- a/pkg/builder/builder.go +++ b/pkg/builder/builder.go @@ -206,7 +206,10 @@ func (b generalBuilder) imageConfig(ctx context.Context) (string, error) { return "", errors.Wrap(err, "failed to get entrypoint") } b.logger.Debugf("final entrypoint: {%s}\n", ep) - data, err := ImageConfigStr(labels, ports, ep) + + env := ir.CompileEnviron() + + data, err := ImageConfigStr(labels, ports, ep, env) if err != nil { return "", errors.Wrap(err, "failed to get image config") } diff --git a/pkg/builder/util.go b/pkg/builder/util.go index 0bb246cd8..253b3217a 100644 --- a/pkg/builder/util.go +++ b/pkg/builder/util.go @@ -35,14 +35,14 @@ const ( defaultFunc = "build" ) -func ImageConfigStr(labels map[string]string, - ports map[string]struct{}, entrypoint []string) (string, error) { +func ImageConfigStr(labels map[string]string, ports map[string]struct{}, + entrypoint []string, env []string) (string, error) { pl := platforms.Normalize(platforms.DefaultSpec()) img := v1.Image{ Config: v1.ImageConfig{ Labels: labels, WorkingDir: "/", - Env: []string{"PATH=" + DefaultPathEnv(pl.OS)}, + Env: append(env, "PATH="+DefaultPathEnv(pl.OS)), ExposedPorts: ports, Entrypoint: entrypoint, }, diff --git a/pkg/lang/frontend/starlark/runtime/const.go b/pkg/lang/frontend/starlark/runtime/const.go index 70b0d3bb8..89e285a70 100644 --- a/pkg/lang/frontend/starlark/runtime/const.go +++ b/pkg/lang/frontend/starlark/runtime/const.go @@ -18,4 +18,5 @@ const ( ruleCommand = "runtime.command" ruleExpose = "runtime.expose" ruleDaemon = "runtime.daemon" + ruleEnviron = "runtime.environ" ) diff --git a/pkg/lang/frontend/starlark/runtime/runtime.go b/pkg/lang/frontend/starlark/runtime/runtime.go index e849af02d..523aba43b 100644 --- a/pkg/lang/frontend/starlark/runtime/runtime.go +++ b/pkg/lang/frontend/starlark/runtime/runtime.go @@ -35,6 +35,7 @@ var Module = &starlarkstruct.Module{ "command": starlark.NewBuiltin(ruleCommand, ruleFuncCommand), "daemon": starlark.NewBuiltin(ruleDaemon, ruleFuncDaemon), "expose": starlark.NewBuiltin(ruleExpose, ruleFuncExpose), + "environ": starlark.NewBuiltin(ruleEnviron, ruleFuncEnviron), }, } @@ -119,3 +120,25 @@ func ruleFuncExpose(thread *starlark.Thread, _ *starlark.Builtin, err := ir.RuntimeExpose(int(envdPortInt), int(hostPortInt), serviceNameStr) return starlark.None, err } + +func ruleFuncEnviron(thread *starlark.Thread, _ *starlark.Builtin, + args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) { + var env starlark.IterableMapping + + if err := starlark.UnpackArgs(ruleCommand, args, kwargs, "env", &env); err != nil { + return nil, err + } + + envMap := make(map[string]string) + for _, tuple := range env.Items() { + if len(tuple) != 2 { + return nil, fmt.Errorf("invalid env (%s)", tuple.String()) + } + + envMap[tuple[0].(starlark.String).GoString()] = tuple[1].(starlark.String).GoString() + } + + logger.Debugf("rule `%s` is invoked, env: %v", ruleEnviron, envMap) + ir.RuntimeEnviron(envMap) + return starlark.None, nil +} diff --git a/pkg/lang/ir/compile.go b/pkg/lang/ir/compile.go index 8dab31a56..077cde0e9 100644 --- a/pkg/lang/ir/compile.go +++ b/pkg/lang/ir/compile.go @@ -50,6 +50,7 @@ func NewGraph() *Graph { SystemPackages: []string{}, Exec: []string{}, RuntimeCommands: make(map[string]string), + RuntimeEnviron: make(map[string]string), Shell: shellBASH, } } @@ -97,6 +98,10 @@ func CompileEntrypoint(buildContextDir string) ([]string, error) { return DefaultGraph.GetEntrypoint(buildContextDir) } +func CompileEnviron() []string { + return DefaultGraph.EnvString() +} + func (g Graph) GPUEnabled() bool { return g.CUDA != nil } @@ -149,6 +154,14 @@ func (g Graph) ExposedPorts() (map[string]struct{}, error) { return ports, nil } +func (g Graph) EnvString() []string { + var envs []string + for k, v := range g.RuntimeEnviron { + envs = append(envs, fmt.Sprintf("%s=%s", k, v)) + } + return envs +} + func (g Graph) DefaultCacheImporter() (*string, error) { // The base remote cache should work for all languages. res := fmt.Sprintf( diff --git a/pkg/lang/ir/interface.go b/pkg/lang/ir/interface.go index 35f6f19ae..703d4f2f7 100644 --- a/pkg/lang/ir/interface.go +++ b/pkg/lang/ir/interface.go @@ -231,3 +231,9 @@ func RuntimeExpose(envdPort, hostPort int, serviceName string) error { }) return nil } + +func RuntimeEnviron(env map[string]string) { + for k, v := range env { + DefaultGraph.RuntimeEnviron[k] = v + } +} diff --git a/pkg/lang/ir/types.go b/pkg/lang/ir/types.go index 3cf6f1f37..1736c8604 100644 --- a/pkg/lang/ir/types.go +++ b/pkg/lang/ir/types.go @@ -56,6 +56,7 @@ type Graph struct { Entrypoint []string RuntimeCommands map[string]string RuntimeDaemon [][]string + RuntimeEnviron map[string]string *JupyterConfig *GitConfig