support trial meta injection in trial template rendering

pkg/controller.v1beta1/consts/const.go

@@ -137,12 +137,23 @@ const (
 	// LabelTrialTemplateConfigMapValue is the label value for the Trial templates configMap
 	LabelTrialTemplateConfigMapValue = "katib-trial-templates"
 
-	// TrialTemplateReplaceFormat is the format to make substitution in Trial template from Names in TrialParameters
+	// TrialTemplateParamReplaceFormat is the format to make substitution in Trial template from Names in TrialParameters
 	// E.g if Name = learningRate, according value in Trial template must be ${trialParameters.learningRate}
-	TrialTemplateReplaceFormat = "${trialParameters.%v}"
+	TrialTemplateParamReplaceFormat = "${trialParameters.%v}"
 
-	// TrialTemplateReplaceFormatRegex is the regex for Trial template format
-	TrialTemplateReplaceFormatRegex = "\\{trialParameters\\..+?\\}"
+	// TrialTemplateParamReplaceFormatRegex is the regex for TrialParameters format in Trial template
+	TrialTemplateParamReplaceFormatRegex = "\\$\\{trialParameters\\..+?\\}"
+
+	// TrialTemplateMetaReplaceFormatRegex is the regex for TrialMetadata format in Trial template
+	TrialTemplateMetaReplaceFormatRegex = "\\$\\{trialSpec\\.(.+?)\\}"
+
+	// valid keys of trial metadata which are used to make substitution in Trial template
+	TrialTemplateMetaKeyOfName        = "Name"
+	TrialTemplateMetaKeyOfNamespace   = "Namespace"
+	TrialTemplateMetaKeyOfKind        = "Kind"
+	TrialTemplateMetaKeyOfAPIVersion  = "APIVersion"
+	TrialTemplateMetaKeyOfAnnotations = "Annotations"
+	TrialTemplateMetaKeyOfLabels      = "Labels"
 
 	// UnavailableMetricValue is the value when metric was not reported or metric value can't be converted to float64
 	UnavailableMetricValue = "unavailable"

pkg/controller.v1beta1/experiment/experiment_controller_test.go

@@ -201,8 +201,7 @@ func TestReconcileExperiment(t *testing.T) {
 	if err != nil {
 		t.Errorf("ConvertObjectToUnstructured failed: %v", err)
 	}
-	generator.EXPECT().GetRunSpecWithHyperParameters(gomock.Any(), gomock.Any(),
-		gomock.Any(), gomock.Any()).Return(
+	generator.EXPECT().GetRunSpecWithHyperParameters(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(
 		returnedUnstructured,
 		nil).AnyTimes()
 

pkg/controller.v1beta1/experiment/experiment_util.go

@@ -2,6 +2,7 @@ package experiment
 
 import (
 	"context"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 
 	"k8s.io/apimachinery/pkg/api/errors"
 
@@ -61,6 +62,15 @@ func (r *ReconcileExperiment) createTrialInstance(expInstance *experimentsv1beta
 
 }
 
+func buildTrialMetaForRunSpec(trial *trialsv1beta1.Trial) map[string]string {
+	return map[string]string{
+		consts.TrialTemplateMetaKeyOfName:       trial.Name,
+		consts.TrialTemplateMetaKeyOfNamespace:  trial.Namespace,
+		consts.TrialTemplateMetaKeyOfKind:       trial.Kind,
+		consts.TrialTemplateMetaKeyOfAPIVersion: trial.APIVersion,
+	}
+}
+
 func needUpdateFinalizers(exp *experimentsv1beta1.Experiment) (bool, []string) {
 	deleted := !exp.ObjectMeta.DeletionTimestamp.IsZero()
 	pendingFinalizers := exp.GetFinalizers()

pkg/controller.v1beta1/experiment/manifest/generator.go

@@ -3,6 +3,7 @@ package manifest
 import (
 	"errors"
 	"fmt"
+	"regexp"
 	"strings"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -12,7 +13,7 @@ import (
 	commonapiv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
 	experimentsv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/experiments/v1beta1"
 	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
-	util "github.com/kubeflow/katib/pkg/controller.v1beta1/util"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/util"
 	"github.com/kubeflow/katib/pkg/util/v1beta1/katibclient"
 	"github.com/kubeflow/katib/pkg/util/v1beta1/katibconfig"
 )
@@ -62,14 +63,8 @@ func (g *DefaultGenerator) GetSuggestionConfigData(algorithmName string) (map[st
 // GetRunSpecWithHyperParameters returns the specification for trial with hyperparameters.
 func (g *DefaultGenerator) GetRunSpecWithHyperParameters(experiment *experimentsv1beta1.Experiment, trialName, trialNamespace string, assignments []commonapiv1beta1.ParameterAssignment) (*unstructured.Unstructured, error) {
 
-	// Get string Trial template from Experiment spec
-	trialTemplate, err := g.GetTrialTemplate(experiment)
-	if err != nil {
-		return nil, err
-	}
-
 	// Apply parameters to Trial Template from assignment
-	replacedTemplate, err := g.applyParameters(trialTemplate, experiment.Spec.TrialTemplate.TrialParameters, assignments)
+	replacedTemplate, err := g.applyParameters(experiment, trialName, trialNamespace, assignments)
 	if err != nil {
 		return nil, err
 	}
@@ -86,26 +81,95 @@ func (g *DefaultGenerator) GetRunSpecWithHyperParameters(experiment *experiments
 	return runSpec, nil
 }
 
-func (g *DefaultGenerator) applyParameters(trialTemplate string, trialParams []experimentsv1beta1.TrialParameterSpec, assignments []commonapiv1beta1.ParameterAssignment) (string, error) {
-	// Number of parameters must be equal
-	if len(assignments) != len(trialParams) {
-		return "", fmt.Errorf("Number of Trial assignment from Suggestion: %v not equal to number Trial parameters from Experiment: %v", len(assignments), len(trialParams))
+func (g *DefaultGenerator) applyParameters(experiment *experimentsv1beta1.Experiment, trialName, trialNamespace string, assignments []commonapiv1beta1.ParameterAssignment) (string, error) {
+	// Get string Trial template from Experiment spec
+	trialTemplate, err := g.GetTrialTemplate(experiment)
+	if err != nil {
+		return "", err
 	}
+
+	trialSpec := experiment.Spec.TrialTemplate.TrialSpec
+	// If trialSpec is not defined in TrialTemplate, deserialize templateString to fetch it
+	if trialSpec == nil {
+		trialSpec, err = util.ConvertStringToUnstructured(trialTemplate)
+		if err != nil {
+			return "", fmt.Errorf("ConvertStringToUnstructured failed: %v", err)
+		}
+	}
+
 	// Convert parameter assignment to map key = parameter name, value = parameter value
 	assignmentsMap := make(map[string]string)
 	for _, assignment := range assignments {
 		assignmentsMap[assignment.Name] = assignment.Value
 	}
 
-	// Replacing parameters from Trial parameters
-	for _, parameter := range trialParams {
+	placeHolderToValueMap := make(map[string]string)
+	var metaRefKey []string
+	var metaKey, metaIndex string
+	nonMetaParamCount := 0
+	for _, param := range experiment.Spec.TrialTemplate.TrialParameters {
+		metaMatchRegex := regexp.MustCompile(consts.TrialTemplateMetaReplaceFormatRegex)
+		metaRefKey = metaMatchRegex.FindStringSubmatch(param.Reference)
+		// handle trial parameters which consume trial assignments
+		if len(metaRefKey) == 0 {
+			if value, ok := assignmentsMap[param.Reference]; ok {
+				placeHolderToValueMap[param.Name] = value
+				nonMetaParamCount += 1
+				continue
+			} else {
+				return "", fmt.Errorf("illegal reference of trial metadata: %v", param.Reference)
+			}
+		}
 
-		if parameterValue, ok := assignmentsMap[parameter.Reference]; ok {
-			trialTemplate = strings.Replace(trialTemplate, fmt.Sprintf(consts.TrialTemplateReplaceFormat, parameter.Name), parameterValue, -1)
+		// handle trial parameters which consume trial meta data
+		// extract index (key) of Labels and Annotations if exists
+		if sub := regexp.MustCompile("(.+)\\[(.+)]").FindStringSubmatch(metaRefKey[1]); len(sub) > 0 {
+			if len(sub) != 3 {
+				return "", fmt.Errorf("illegal reference of trial metadata: %v", param.Reference)
+			}
+			metaKey = sub[1]
+			metaIndex = sub[2]
 		} else {
-			return "", fmt.Errorf("Unable to find parameter: %v in parameter assignment %v", parameter.Reference, assignmentsMap)
+			metaKey = metaRefKey[1]
+			metaIndex = ""
+		}
+		// fetch metadata value
+		switch metaKey {
+		case consts.TrialTemplateMetaKeyOfName:
+			placeHolderToValueMap[param.Name] = trialName
+		case consts.TrialTemplateMetaKeyOfNamespace:
+			placeHolderToValueMap[param.Name] = trialNamespace
+		case consts.TrialTemplateMetaKeyOfKind:
+			placeHolderToValueMap[param.Name] = trialSpec.GetKind()
+		case consts.TrialTemplateMetaKeyOfAPIVersion:
+			placeHolderToValueMap[param.Name] = trialSpec.GetAPIVersion()
+		case consts.TrialTemplateMetaKeyOfAnnotations:
+			if value, ok := trialSpec.GetAnnotations()[metaIndex]; !ok {
+				return "", fmt.Errorf("illegal reference of trial metadata: %v; failed to fetch Annotation: %v", param.Reference, metaIndex)
+			} else {
+				placeHolderToValueMap[param.Name] = value
+			}
+		case consts.TrialTemplateMetaKeyOfLabels:
+			if value, ok := trialSpec.GetLabels()[metaIndex]; !ok {
+				return "", fmt.Errorf("illegal reference of trial metadata: %v; failed to fetch Label: %v", param.Reference, metaIndex)
+			} else {
+				placeHolderToValueMap[param.Name] = value
+			}
+		default:
+			return "", fmt.Errorf("illegal reference of trial metadata: %v", param.Reference)
 		}
 	}
+
+	// Number of parameters must be equal
+	if len(assignments) != nonMetaParamCount {
+		return "", fmt.Errorf("Number of TrialAssignment: %v != number of nonMetaTrialParameters in TrialSpec: %v", len(assignments), nonMetaParamCount)
+	}
+
+	// Replacing placeholders with parameter values
+	for placeHolder, paramValue := range placeHolderToValueMap {
+		trialTemplate = strings.Replace(trialTemplate, fmt.Sprintf(consts.TrialTemplateParamReplaceFormat, placeHolder), paramValue, -1)
+	}
+
 	return trialTemplate, nil
 }
 

pkg/controller.v1beta1/experiment/manifest/generator_test.go

@@ -2,17 +2,18 @@ package manifest
 
 import (
 	"errors"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/consts"
 	"math"
 	"reflect"
 	"testing"
 
 	"github.com/golang/mock/gomock"
 	commonapiv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/common/v1beta1"
 	experimentsv1beta1 "github.com/kubeflow/katib/pkg/apis/controller/experiments/v1beta1"
-	util "github.com/kubeflow/katib/pkg/controller.v1beta1/util"
+	"github.com/kubeflow/katib/pkg/controller.v1beta1/util"
 	katibclientmock "github.com/kubeflow/katib/pkg/mock/v1beta1/util/katibclient"
 	batchv1 "k8s.io/api/batch/v1"
-	v1 "k8s.io/api/core/v1"
+	"k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 )
@@ -50,6 +51,12 @@ func TestGetRunSpecWithHP(t *testing.T) {
 								"--lr=0.05",
 								"--num-layers=5",
 							},
+							Env: []v1.EnvVar{
+								{Name: consts.TrialTemplateMetaKeyOfName, Value: "trial-name"},
+								{Name: consts.TrialTemplateMetaKeyOfNamespace, Value: "trial-namespace"},
+								{Name: consts.TrialTemplateMetaKeyOfKind, Value: "Job"},
+								{Name: consts.TrialTemplateMetaKeyOfAPIVersion, Value: "batch/v1"},
+							},
 						},
 					},
 				},
@@ -330,6 +337,12 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 								"--lr=${trialParameters.learningRate}",
 								"--num-layers=${trialParameters.numberLayers}",
 							},
+							Env: []v1.EnvVar{
+								{Name: consts.TrialTemplateMetaKeyOfName, Value: "${trialParameters.trialName}"},
+								{Name: consts.TrialTemplateMetaKeyOfNamespace, Value: "${trialParameters.trialNamespace}"},
+								{Name: consts.TrialTemplateMetaKeyOfKind, Value: "${trialParameters.jobKind}"},
+								{Name: consts.TrialTemplateMetaKeyOfAPIVersion, Value: "${trialParameters.jobAPIVersion}"},
+							},
 						},
 					},
 				},
@@ -355,6 +368,26 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 						Description: "Number of layers",
 						Reference:   "num-layers",
 					},
+					{
+						Name:        "trialName",
+						Description: "name of current trial",
+						Reference:   "${trialSpec.Name}",
+					},
+					{
+						Name:        "trialNamespace",
+						Description: "namespace of current trial",
+						Reference:   "${trialSpec.Namespace}",
+					},
+					{
+						Name:        "jobKind",
+						Description: "job kind of current trial",
+						Reference:   "${trialSpec.Kind}",
+					},
+					{
+						Name:        "jobAPIVersion",
+						Description: "job API Version of current trial",
+						Reference:   "${trialSpec.APIVersion}",
+					},
 				},
 			},
 		},