Fix: Handle "Annotations" and "Labels" trial's metadata references

pkg/webhook/v1beta1/experiment/validator/validator.go

@@ -311,10 +311,7 @@ func (g *DefaultValidator) validateTrialTemplate(instance *experimentsv1beta1.Ex
 
 		// Check if parameter reference exist in experiment parameters
 		if len(experimentParameterNames) > 0 {
-			// Check if parameter is trial metadata
-			regex := regexp.MustCompile(consts.TrialTemplateMetaReplaceFormatRegex)
-			match := regex.FindStringSubmatch(parameter.Reference)
-			if !(len(match) > 0 && contains(consts.TrialTemplateMetaKeys, match[1])) {
+			if !isMetaKey(parameter.Reference) {
 				if _, ok := experimentParameterNames[parameter.Reference]; !ok {
 					return fmt.Errorf("parameter reference %v does not exist in spec.parameters: %v", parameter.Reference, instance.Spec.Parameters)
 				}
@@ -487,6 +484,25 @@ func (g *DefaultValidator) validateMetricsCollector(inst *experimentsv1beta1.Exp
 	return nil
 }
 
+func isMetaKey(parameter string) bool {
+	// Check if parameter is trial metadata reference as ${trailSpec.Name}, ${trialSpec.Labels[label]}, etc. used for substitution
+	match := regexp.MustCompile(consts.TrialTemplateMetaReplaceFormatRegex).FindStringSubmatch(parameter)
+	isMeta := false
+	if len(match) > 0 {
+		matchedKey := match[1]
+		if contains(consts.TrialTemplateMetaKeys, matchedKey) {
+			isMeta = true
+		} else {
+			// Check if it's Labels[label] or Annotations[annotation]
+			subMatch := regexp.MustCompile(consts.TrialTemplateMetaParseFormatRegex).FindStringSubmatch(matchedKey)
+			if len(subMatch) == 3 && contains(consts.TrialTemplateMetaKeys, subMatch[1]) {
+				isMeta = true
+			}
+		}
+	}
+	return isMeta
+}
+
 func contains(slice []string, item string) bool {
 	for _, s := range slice {
 		if s == item {

pkg/webhook/v1beta1/experiment/validator/validator_test.go

@@ -416,6 +416,8 @@ spec:
 	validTemplate5 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
 	validTemplate6 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
 	validTemplate7 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
+	validTemplate8 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
+	validTemplate9 := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(validJobStr, nil)
 
 	missedParameterTemplate := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(missedParameterJobStr, nil)
 	oddParameterTemplate := p.EXPECT().GetTrialTemplate(gomock.Any()).Return(oddParameterJobStr, nil)
@@ -433,6 +435,8 @@ spec:
 		validTemplate5,
 		validTemplate6,
 		validTemplate7,
+		validTemplate8,
+		validTemplate9,
 		missedParameterTemplate,
 		oddParameterTemplate,
 		invalidParameterTemplate,
@@ -582,6 +586,24 @@ spec:
 			Err:             false,
 			testDescription: "Trial template contains Trial metadata reference as parameter",
 		},
+		{
+			Instance: func() *experimentsv1beta1.Experiment {
+				i := newFakeInstance()
+				i.Spec.TrialTemplate.TrialParameters[1].Reference = "${trialSpec.Annotations[test-annotation]}"
+				return i
+			}(),
+			Err:             false,
+			testDescription: "Trial template contains Trial annotation reference as parameter",
+		},
+		{
+			Instance: func() *experimentsv1beta1.Experiment {
+				i := newFakeInstance()
+				i.Spec.TrialTemplate.TrialParameters[1].Reference = "${trialSpec.Labels[test-label]}"
+				return i
+			}(),
+			Err:             false,
+			testDescription: "Trial template contains Trial's label reference as parameter",
+		},
 		// Trial Template doesn't contain parameter from trialParameters
 		// missedParameterTemplate case
 		{