feat: add runtime environments (#787)

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

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

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"})
+ ```
+ """

examples/python-basic/build.envd

@@ -15,4 +15,5 @@ def build():
  runtime.command(commands={
  "test": "ls /",
  })
+ runtime.environ(env={"ENVD_MODE": "DEV"})
  jupyter_lab()

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")
  }

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

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

@@ -18,4 +18,5 @@ const (
  ruleCommand = "runtime.command"
  ruleExpose = "runtime.expose"
  ruleDaemon = "runtime.daemon"
+ ruleEnviron = "runtime.environ"
 )

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

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(

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

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