Adds support to specify a map of environment variables to the AI endp…

…oint deployment
Azure · Apr 25, 2024 · 3432cb8 · 3432cb8
1 parent 6a2aa89
commit 3432cb8
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 6 deletions.
diff --git a/cli/azd/pkg/ai/config.go b/cli/azd/pkg/ai/config.go
@@ -14,13 +14,19 @@ type ComponentConfig struct {
 	Overrides map[string]osutil.ExpandableString `yaml:"overrides,omitempty"`
 }
 
+type DeploymentConfig struct {
+	ComponentConfig `yaml:",inline"`
+	// A map of environment variables to set for the deployment
+	Environment map[string]osutil.ExpandableString `yaml:"environment,omitempty"`
+}
+
 // EndpointDeploymentConfig is a configuration structure for an ML online endpoint deployment
 type EndpointDeploymentConfig struct {
 	Workspace   osutil.ExpandableString `yaml:"workspace,omitempty"`
 	Environment *ComponentConfig        `yaml:"environment,omitempty"`
 	Model       *ComponentConfig        `yaml:"model,omitempty"`
 	Flow        *ComponentConfig        `yaml:"flow,omitempty"`
-	Deployment  *ComponentConfig        `yaml:"deployment,omitempty"`
+	Deployment  *DeploymentConfig       `yaml:"deployment,omitempty"`
 }
 
 // Flow is a configuration to defined a Prompt flow component

diff --git a/cli/azd/pkg/project/ai_helper.go b/cli/azd/pkg/project/ai_helper.go
@@ -425,6 +425,16 @@ func (a *aiHelper) DeployToEndpoint(
 		deploymentArgs = append(deploymentArgs, "--set", fmt.Sprintf("model=%s", modelVersionName))
 	}
 
+	if config.Deployment.Overrides == nil {
+		config.Deployment.Overrides = map[string]osutil.ExpandableString{}
+	}
+
+	// Transform any referenced environment variables into override expressions
+	for key, value := range config.Deployment.Environment {
+		envVarOverrideKey := fmt.Sprintf("environment_variables.%s", key)
+		config.Deployment.Overrides[envVarOverrideKey] = value
+	}
+
 	deploymentArgs, err = a.applyOverrides(deploymentArgs, config.Deployment.Overrides)
 	if err != nil {
 		return nil, err
@@ -738,7 +748,7 @@ func (a *aiHelper) applyOverrides(args []string, overrides map[string]osutil.Exp
 	for key, value := range overrides {
 		expandedValue, err := value.Envsubst(a.env.Getenv)
 		if err != nil {
-			return nil, fmt.Errorf("failed parsing environment override %s: %w", key, err)
+			return nil, fmt.Errorf("failed parsing override %s: %w", key, err)
 		}
 
 		args = append(args, "--set", fmt.Sprintf("%s=%s", key, expandedValue))

diff --git a/cli/azd/pkg/project/ai_helper_test.go b/cli/azd/pkg/project/ai_helper_test.go
@@ -279,9 +279,11 @@ func Test_AiHelper_DeployToEndpoint(t *testing.T) {
 			Name: osutil.NewExpandableString(modelName),
 			Path: "./deployment/model.yaml",
 		},
-		Deployment: &ai.ComponentConfig{
-			Name: osutil.NewExpandableString(deploymentName),
-			Path: "./deployment/deployment.yaml",
+		Deployment: &ai.DeploymentConfig{
+			ComponentConfig: ai.ComponentConfig{
+				Name: osutil.NewExpandableString(deploymentName),
+				Path: "./deployment/deployment.yaml",
+			},
 		},
 	}
 

diff --git a/schemas/alpha/azure.yaml.json b/schemas/alpha/azure.yaml.json
@@ -960,6 +960,24 @@
                 "path"
             ]
         },
+        "aiDeploymentConfig": {
+            "allOf": [
+                { "$ref": "#/definitions/aiComponentConfig" },
+                {
+                    "type": "object",
+                    "properties": {
+                        "environment": {
+                            "type": "object",
+                            "title": "A map of key/value pairs to set as environment variables for the deployment.",
+                            "description": "Optional. Values support OS & AZD environment variable substitution.",
+                            "additionalProperties":{
+                                "type": "string"
+                            }
+                        }
+                    }
+                }
+            ]
+        },
         "aiEndpointConfig": {
             "type": "object",
             "additionalProperties": false,
@@ -985,7 +1003,7 @@
                     "description": "Optional. When omitted a model will not be created."
                 },
                 "deployment": {
-                    "$ref": "#/definitions/aiComponentConfig",
+                    "$ref": "#/definitions/aiDeploymentConfig",
                     "title": "The Azure AI Studio online endpoint deployment configuration.",
                     "description": "Required. A new online endpoint deployment will be created and traffic will automatically to shifted to the new deployment upon successful completion."
                 }