feat: inject brick variable (#659)

Author: lucarin91

internal/orchestrator/bricksindex/bricks_index.go

@@ -2,6 +2,7 @@ package bricksindex
 
 import (
 	"io"
+	"iter"
 	"slices"
 
 	"github.com/arduino/go-paths-helper"
@@ -56,6 +57,16 @@ func (b Brick) GetVariable(name string) (BrickVariable, bool) {
 	return b.Variables[idx], true
 }
 
+func (b Brick) GetDefaultVariables() iter.Seq2[string, string] {
+	return func(yield func(string, string) bool) {
+		for _, v := range b.Variables {
+			if !yield(v.Name, v.DefaultValue) {
+				return
+			}
+		}
+	}
+}
+
 func unmarshalBricksIndex(content io.Reader) (*BricksIndex, error) {
 	var index BricksIndex
 	if err := yaml.NewDecoder(content).Decode(&index); err != nil {

internal/orchestrator/orchestrator.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"iter"
 	"log/slog"
+	"maps"
 	"net"
 	"os"
 	"os/user"
@@ -34,6 +35,7 @@ import (
 	"github.com/bcmi-labs/orchestrator/internal/store"
 	"github.com/bcmi-labs/orchestrator/pkg/helpers"
 	"github.com/bcmi-labs/orchestrator/pkg/micro"
+	"github.com/bcmi-labs/orchestrator/pkg/x"
 	"github.com/bcmi-labs/orchestrator/pkg/x/fatomic"
 )
 
@@ -137,15 +139,14 @@ func StartApp(
 			}
 		}
 		if app.MainPythonFile != nil {
-			// Override the compose Variables with the app's variables and model configuration.
-			envs, mapped_env := handleEnvironmentVariables(app, modelsIndex)
+			envs := getAppEnvironmentVariables(app, bricksIndex, modelsIndex)
 
 			if !yield(StreamMessage{data: "Provisioning app..."}) {
 				cancel()
 				return
 			}
 
-			if err := provisioner.App(ctx, bricksIndex, &app, cfg, mapped_env, staticStore); err != nil {
+			if err := provisioner.App(ctx, bricksIndex, &app, cfg, envs, staticStore); err != nil {
 				yield(StreamMessage{error: err})
 				return
 			}
@@ -177,7 +178,7 @@ func StartApp(
 				}
 			})
 
-			process, err := paths.NewProcess(envs, commands...)
+			process, err := paths.NewProcess(envs.AsList(), commands...)
 			if err != nil {
 				yield(StreamMessage{error: err})
 				return
@@ -197,35 +198,44 @@ func StartApp(
 	}
 }
 
-func handleEnvironmentVariables(app app.ArduinoApp, modelsIndex *modelsindex.ModelsIndex) ([]string, map[string]string) {
-	envs := []string{}
-	mapped_env := map[string]string{}
-	addMapToEnv := func(m map[string]string) {
-		for k, v := range m {
-			envs = append(envs, fmt.Sprintf("%s=%s", k, v))
-			mapped_env[k] = v
-		}
-	}
+// getAppEnvironmentVariables returns the environment variables for the app by merging variables and config in the following order:
+// - brick default variables (variables defined in the brick definition)
+// - brick instance variables (variables defined in the app.yaml for the brick instance)
+// - model configuration variables (variables defined in the model configuration)
+// In addition, it adds some useful environment variables like APP_HOME and HOST_IP.
+func getAppEnvironmentVariables(app app.ArduinoApp, brickIndex *bricksindex.BricksIndex, modelsIndex *modelsindex.ModelsIndex) x.EnvVars {
+	envs := make(x.EnvVars)
+
 	for _, brick := range app.Descriptor.Bricks {
-		addMapToEnv(brick.Variables)
+		if brickDef, found := brickIndex.FindBrickByID(brick.ID); found {
+			maps.Insert(envs, brickDef.GetDefaultVariables())
+		}
+		maps.Insert(envs, maps.All(brick.Variables))
+
 		if m, found := modelsIndex.GetModelByID(brick.Model); found {
-			addMapToEnv(m.ModelConfiguration)
+			maps.Insert(envs, maps.All(m.ModelConfiguration))
 		}
 	}
 
 	// Add the APP_HOME directory to the environment variables
-	addMapToEnv(map[string]string{"APP_HOME": app.FullPath.String()})
-	slog.Debug("Configuring app environment", slog.String("APP_HOME", app.FullPath.String()), slog.Any("envs", envs))
+	envs["APP_HOME"] = app.FullPath.String()
 
 	// Pre-select default camera device if available. This can be overridden by the app environment variables (or in future by applab)
 	// This is required because there are some video devices for HW acceleration that are auto registered in /dev but are not real cameras.
 	if videoDevices := getVideoDevices(); len(videoDevices) > 0 {
 		// VIDEO_DEVICE will be the first device in /dev/v4l/by-id
-		addMapToEnv(map[string]string{"VIDEO_DEVICE": videoDevices[0]})
-		slog.Info("Configuring default video device", slog.String("VIDEO_DEVICE", videoDevices[0]))
+		envs["VIDEO_DEVICE"] = videoDevices[0]
 	}
 
-	return envs, mapped_env
+	if hostIP, err := helpers.GetHostIP(); err == nil {
+		envs["HOST_IP"] = hostIP
+	} else {
+		slog.Warn("unable to get host IP", slog.String("error", err.Error()))
+	}
+
+	slog.Debug("Current environment variables", slog.Any("envs", envs))
+
+	return envs
 }
 
 func extractIndexFromVideoDeviceName(device string) (int, error) {

internal/orchestrator/orchestrator_test.go

@@ -14,7 +14,9 @@ import (
 	"go.bug.st/f"
 
 	"github.com/bcmi-labs/orchestrator/internal/orchestrator/app"
+	"github.com/bcmi-labs/orchestrator/internal/orchestrator/bricksindex"
 	"github.com/bcmi-labs/orchestrator/internal/orchestrator/config"
+	"github.com/bcmi-labs/orchestrator/internal/orchestrator/modelsindex"
 )
 
 func TestCloneApp(t *testing.T) {
@@ -355,7 +357,7 @@ func createApp(
 	isExample bool,
 	idProvider *app.IDProvider,
 	cfg config.Configuration,
-) {
+) app.ID {
 	t.Helper()
 
 	res, err := CreateApp(t.Context(), CreateAppRequest{
@@ -371,7 +373,10 @@ func createApp(
 		newID, err := idProvider.IDFromPath(newPath)
 		require.NoError(t, err)
 		assert.Empty(t, gCmp.Diff(f.Must(idProvider.ParseID("examples:"+name)), newID))
+		res.ID = newID
 	}
+
+	return res.ID
 }
 
 func TestSortV4LVideoDevices(t *testing.T) {
@@ -387,3 +392,88 @@ func TestSortV4LVideoDevices(t *testing.T) {
 	assert.Equal(t, "usb-Generic_GENERAL_-_UVC-video-index1", devices[1])
 	assert.Equal(t, "usb-046d_0825-video-index2", devices[2])
 }
+
+func TestGetAppEnvironmentVariablesWithDefaults(t *testing.T) {
+	cfg := setTestOrchestratorConfig(t)
+	idProvider := app.NewAppIDProvider(cfg)
+
+	docker, err := dockerClient.NewClientWithOpts(
+		dockerClient.FromEnv,
+		dockerClient.WithAPIVersionNegotiation(),
+	)
+	require.NoError(t, err)
+	dockerCli, err := command.NewDockerCli(
+		command.WithAPIClient(docker),
+		command.WithBaseContext(t.Context()),
+	)
+	require.NoError(t, err)
+
+	err = dockerCli.Initialize(&flags.ClientOptions{})
+	require.NoError(t, err)
+
+	appId := createApp(t, "app1", false, idProvider, cfg)
+	appDesc, err := app.Load(appId.ToPath().String())
+	require.NoError(t, err)
+	appDesc.Descriptor.Bricks = []app.Brick{
+		{
+			ID:        "arduino:object_detection",
+			Model:     "",                  // use the default model
+			Variables: map[string]string{}, // use the default variables
+		},
+	}
+
+	bricksIndexContent := []byte(`
+bricks:
+- id: arduino:object_detection
+  name: Object Detection
+  description: "Brick for object detection using a pre-trained model. It processes\
+    \ images and returns the predicted class label, bounding-boxes and confidence\
+    \ score.\nBrick is designed to work with pre-trained models provided by framework\
+    \ or with custom object detection models trained on Edge Impulse platform. \n"
+  require_container: true
+  require_model: true
+  require_devices: false
+  ports: []
+  category: video
+  model_name: yolox-object-detection
+  variables:
+  - name: CUSTOM_MODEL_PATH
+    default_value: /home/arduino/.arduino-bricks/ei-models
+    description: path to the custom model directory
+  - name: EI_OBJ_DETECTION_MODEL
+    default_value: /models/ootb/ei/yolo-x-nano.eim
+    description: path to the model file
+`)
+	err = cfg.AssetsDir().Join("bricks-list.yaml").WriteFile(bricksIndexContent)
+	require.NoError(t, err)
+	bricksIndex, err := bricksindex.GenerateBricksIndexFromFile(cfg.AssetsDir())
+	assert.NoError(t, err)
+
+	modelsIndexContent := []byte(`
+models:
+- yolox-object-detection:
+    runner: brick
+    name : "General purpose object detection - YoloX"
+    description: "General purpose object detection model based on YoloX Nano. This model is trained on the COCO dataset and can detect 80 different object classes."
+    model_configuration:
+      "EI_OBJ_DETECTION_MODEL": "/models/ootb/ei/yolo-x-nano.eim"
+    metadata:
+    source: "edgeimpulse"
+    ei-project-id: 717280
+    source-model-id: "YOLOX-Nano"
+    source-model-url: "https://github.com/Megvii-BaseDetection/YOLOX"
+    bricks:
+    - arduino:object_detection
+    - arduino:video_object_detection
+`)
+	err = cfg.AssetsDir().Join("models-list.yaml").WriteFile(modelsIndexContent)
+	require.NoError(t, err)
+	modelIndex, err := modelsindex.GenerateModelsIndexFromFile(cfg.AssetsDir())
+	require.NoError(t, err)
+
+	env := getAppEnvironmentVariables(appDesc, bricksIndex, modelIndex)
+	require.Equal(t, cfg.AppsDir().Join("app1").String(), env["APP_HOME"])
+	require.Equal(t, "/models/ootb/ei/yolo-x-nano.eim", env["EI_OBJ_DETECTION_MODEL"])
+	require.Equal(t, "/home/arduino/.arduino-bricks/ei-models", env["CUSTOM_MODEL_PATH"])
+	// we ignore HOST_IP since it's dynamic
+}

internal/orchestrator/provision.go

@@ -21,7 +21,7 @@ import (
 	"github.com/bcmi-labs/orchestrator/internal/orchestrator/bricksindex"
 	"github.com/bcmi-labs/orchestrator/internal/orchestrator/config"
 	"github.com/bcmi-labs/orchestrator/internal/store"
-	"github.com/bcmi-labs/orchestrator/pkg/helpers"
+	"github.com/bcmi-labs/orchestrator/pkg/x"
 )
 
 type volume struct {
@@ -189,7 +189,7 @@ func generateMainComposeFile(
 	bricksIndex *bricksindex.BricksIndex,
 	pythonImage string,
 	cfg config.Configuration,
-	mapped_env map[string]string,
+	envs x.EnvVars,
 	staticStore *store.StaticStore,
 ) error {
 	slog.Debug("Generating main compose file for the App")
@@ -329,18 +329,7 @@ func generateMainComposeFile(
 				DockerAppMainLabel: "true",
 				DockerAppPathLabel: app.FullPath.String(),
 			},
-			// Add the host ip to the environment variables if we can retrieve it.
-			// This is used to show the network ip address of the board in the apps.
-			Environment: map[string]string{
-				"HOST_IP": func() string {
-					if hostIP, err := helpers.GetHostIP(); err != nil {
-						slog.Warn("Failed to get host IP", slog.Any("error", err))
-						return ""
-					} else {
-						return hostIP
-					}
-				}(),
-			},
+			Environment: envs,
 		},
 	}
 
@@ -355,7 +344,7 @@ func generateMainComposeFile(
 
 	// If there are services that require devices, we need to generate an override compose file
 	// Write additional file to override devices section in included compose files
-	if e := generateServicesOverrideFile(app, slices.Collect(maps.Keys(services)), servicesThatRequireDevices, devices.devicePaths, getCurrentUser(), groups, overrideComposeFile); e != nil {
+	if e := generateServicesOverrideFile(app, slices.Collect(maps.Keys(services)), servicesThatRequireDevices, devices.devicePaths, getCurrentUser(), groups, overrideComposeFile, envs); e != nil {
 		return e
 	}
 
@@ -371,7 +360,7 @@ func generateMainComposeFile(
 			slog.Warn("Failed to extract volumes from compose file", slog.String("compose_file", composeFilePath), slog.Any("error", err))
 			continue
 		}
-		provisionComposeVolumes(composeFilePath, volumes, app, mapped_env)
+		provisionComposeVolumes(composeFilePath, volumes, app, envs)
 	}
 
 	// Done!
@@ -409,7 +398,7 @@ func extractServicesFromComposeFile(composeFile *paths.Path) (map[string]service
 	return services, nil
 }
 
-func generateServicesOverrideFile(arduinoApp *app.ArduinoApp, services []string, servicesThatRequireDevices []string, devices []string, user string, groups []string, overrideComposeFile *paths.Path) error {
+func generateServicesOverrideFile(arduinoApp *app.ArduinoApp, services []string, servicesThatRequireDevices []string, devices []string, user string, groups []string, overrideComposeFile *paths.Path, envs x.EnvVars) error {
 	if overrideComposeFile.Exist() {
 		if err := overrideComposeFile.Remove(); err != nil {
 			return fmt.Errorf("failed to remove existing override compose file: %w", err)
@@ -422,10 +411,11 @@ func generateServicesOverrideFile(arduinoApp *app.ArduinoApp, services []string,
 	}
 
 	type serviceOverride struct {
-		User     string            `yaml:"user,omitempty"`
-		Devices  *[]string         `yaml:"devices,omitempty"`
-		GroupAdd *[]string         `yaml:"group_add,omitempty"`
-		Labels   map[string]string `yaml:"labels,omitempty"`
+		User        string            `yaml:"user,omitempty"`
+		Devices     *[]string         `yaml:"devices,omitempty"`
+		GroupAdd    *[]string         `yaml:"group_add,omitempty"`
+		Labels      map[string]string `yaml:"labels,omitempty"`
+		Environment map[string]string `yaml:"environment,omitempty"`
 	}
 	var overrideCompose struct {
 		Services map[string]serviceOverride `yaml:"services,omitempty"`
@@ -443,6 +433,7 @@ func generateServicesOverrideFile(arduinoApp *app.ArduinoApp, services []string,
 			override.Devices = &devices
 			override.GroupAdd = &groups
 		}
+		override.Environment = envs
 		overrideCompose.Services[svc] = override
 	}
 	writeOverrideCompose := func() error {

internal/orchestrator/provision_test.go

@@ -20,9 +20,6 @@ func TestProvisionAppWithOverrides(t *testing.T) {
 	cfg := setTestOrchestratorConfig(t)
 	tempDirectory := t.TempDir()
 
-	// TODO: hack to skip the preEmbargo check
-	cfg.UsedPythonImageTag = "latest"
-
 	staticStore := store.NewStaticStore(cfg.AssetsDir().String())
 
 	// Define a mock app with bricks that require overrides
@@ -33,6 +30,9 @@ func TestProvisionAppWithOverrides(t *testing.T) {
 				{
 					ID:    "arduino:video_object_detection",
 					Model: "yolox-object-detection",
+					Variables: map[string]string{
+						"CUSTOM_MODEL_PATH": "/models/custom/ei/",
+					},
 				},
 				{
 					ID: "arduino:web_ui",
@@ -98,7 +98,9 @@ bricks:
 	require.Equal(t, "Object Detection", br.Name, "Brick name should match")
 
 	// Run the provision function to generate the main compose file
-	env := map[string]string{}
+	env := map[string]string{
+		"FOO": "bar",
+	}
 	err = generateMainComposeFile(&app, bricksIndex, "app-bricks:python-apps-base:dev-latest", cfg, env, staticStore)
 
 	// Validate that the main compose file and overrides are created
@@ -110,7 +112,8 @@ bricks:
 
 	// Open override file and check for the expected override
 	overridesContent, err := overridesFilePath.ReadFile()
-	require.Nil(t, err, "Failed to read overrides file")
+	require.NoError(t, err)
+
 	type services struct {
 		Services map[string]map[string]interface{} `yaml:"services"`
 	}
@@ -119,6 +122,7 @@ bricks:
 	require.Nil(t, err, "Failed to unmarshal overrides content")
 	require.NotNil(t, content.Services["ei-video-obj-detection-runner"], "Override for ei-video-obj-detection-runner should exist")
 	require.NotNil(t, content.Services["ei-video-obj-detection-runner"]["devices"], "Override for ei-video-obj-detection-runner devices should exist")
+	require.Equal(t, "bar", content.Services["ei-video-obj-detection-runner"]["environment"].(map[string]interface{})["FOO"])
 }
 
 func TestVolumeParser(t *testing.T) {

pkg/x/x.go

@@ -6,3 +6,13 @@ import "iter"
 func EmptyIter[V any]() iter.Seq[V] {
 	return func(yield func(V) bool) {}
 }
+
+type EnvVars map[string]string
+
+func (e EnvVars) AsList() []string {
+	list := make([]string, 0, len(e))
+	for k, v := range e {
+		list = append(list, k+"="+v)
+	}
+	return list
+}