support trial meta injection in trial template rendering

kubeflow · Jul 9, 2020 · f3d58b9 · f3d58b9
1 parent 4145c4f
commit f3d58b9
Show file tree

Hide file tree

Showing 5 changed files with 52 additions and 5 deletions.
diff --git a/pkg/controller.v1beta1/consts/const.go b/pkg/controller.v1beta1/consts/const.go
@@ -144,6 +144,11 @@ const (
 	// TrialTemplateReplaceFormatRegex is the regex for Trial template format
 	TrialTemplateReplaceFormatRegex = "\\{trialParameters\\..+?\\}"
 
+	// TrialTemplateTrialName is the placeholder of trial name in trial template
+	TrialTemplateTrialName = "Name"
+	// TrialTemplateTrialNamespace is the placeholder of trial's namespace in trial template
+	TrialTemplateTrialNamespace = "Namespace"
+
 	// UnavailableMetricValue is the value when metric was not reported or metric value can't be converted to float64
 	UnavailableMetricValue = "unavailable"
 )

diff --git a/pkg/controller.v1beta1/experiment/manifest/generator.go b/pkg/controller.v1beta1/experiment/manifest/generator.go
@@ -68,8 +68,14 @@ func (g *DefaultGenerator) GetRunSpecWithHyperParameters(experiment *experiments
 		return nil, err
 	}
 
+	// Trial meta data which may be required by TrialTemplate
+	trialMeta := map[string]string{
+		consts.TrialTemplateTrialName:      trialName,
+		consts.TrialTemplateTrialNamespace: trialNamespace,
+	}
+
 	// Apply parameters to Trial Template from assignment
-	replacedTemplate, err := g.applyParameters(trialTemplate, experiment.Spec.TrialTemplate.TrialParameters, assignments)
+	replacedTemplate, err := g.applyParameters(trialTemplate, experiment.Spec.TrialTemplate.TrialParameters, assignments, trialMeta)
 	if err != nil {
 		return nil, err
 	}
@@ -86,7 +92,7 @@ func (g *DefaultGenerator) GetRunSpecWithHyperParameters(experiment *experiments
 	return runSpec, nil
 }
 
-func (g *DefaultGenerator) applyParameters(trialTemplate string, trialParams []experimentsv1beta1.TrialParameterSpec, assignments []commonapiv1beta1.ParameterAssignment) (string, error) {
+func (g *DefaultGenerator) applyParameters(trialTemplate string, trialParams []experimentsv1beta1.TrialParameterSpec, assignments []commonapiv1beta1.ParameterAssignment, trialMeta map[string]string) (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))
@@ -106,6 +112,13 @@ func (g *DefaultGenerator) applyParameters(trialTemplate string, trialParams []e
 			return "", fmt.Errorf("Unable to find parameter: %v in parameter assignment %v", parameter.Reference, assignmentsMap)
 		}
 	}
+	// Inject TrialMeta if needed
+	for key, value := range trialMeta {
+		if strings.Contains(trialTemplate, fmt.Sprintf(consts.TrialTemplateReplaceFormat, key)) {
+			trialTemplate = strings.Replace(trialTemplate, fmt.Sprintf(consts.TrialTemplateReplaceFormat, key), value, -1)
+		}
+	}
+
 	return trialTemplate, nil
 }
 

diff --git a/pkg/controller.v1beta1/experiment/manifest/generator_test.go b/pkg/controller.v1beta1/experiment/manifest/generator_test.go
@@ -50,6 +50,10 @@ func TestGetRunSpecWithHP(t *testing.T) {
 								"--lr=0.05",
 								"--num-layers=5",
 							},
+							Env: []v1.EnvVar{
+								{Name: "trial_name", Value: "trial-name"},
+								{Name: "trial_ns", Value: "trial-namespace"},
+							},
 						},
 					},
 				},
@@ -330,6 +334,10 @@ func newFakeInstance() *experimentsv1beta1.Experiment {
 								"--lr=${trialParameters.learningRate}",
 								"--num-layers=${trialParameters.numberLayers}",
 							},
+							Env: []v1.EnvVar{
+								{Name: "trial_name", Value: "${trialParameters.Name}"},
+								{Name: "trial_ns", Value: "${trialParameters.Namespace}"},
+							},
 						},
 					},
 				},

diff --git a/pkg/webhook/v1beta1/experiment/validator/validator.go b/pkg/webhook/v1beta1/experiment/validator/validator.go
@@ -184,6 +184,13 @@ func (g *DefaultValidator) validateTrialTemplate(instance *experimentsv1beta1.Ex
 		return fmt.Errorf("spec.trialTemplate.trialParameters must be specified")
 	}
 
+	// Check if parameter names conflict with preserved names
+	for _, parameter := range trialTemplate.TrialParameters {
+		if parameter.Name == consts.TrialTemplateTrialName || parameter.Name == consts.TrialTemplateTrialNamespace {
+			return fmt.Errorf("Name of trialParameters should not be %s or %s", consts.TrialTemplateTrialName, consts.TrialTemplateTrialNamespace)
+		}
+	}
+
 	// Check if trialSpec or configMap exists
 	if trialTemplate.TrialSource.TrialSpec == nil && trialTemplate.TrialSource.ConfigMap == nil {
 		return fmt.Errorf("spec.trialTemplate.trialSpec or spec.trialTemplate.configMap must be specified")

diff --git a/pkg/webhook/v1beta1/experiment/validator/validator_test.go b/pkg/webhook/v1beta1/experiment/validator/validator_test.go
@@ -6,7 +6,7 @@ import (
 
 	"github.com/golang/mock/gomock"
 	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"
 	"k8s.io/apimachinery/pkg/util/intstr"
@@ -15,7 +15,7 @@ import (
 	commonv1beta1 "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"
 	manifestmock "github.com/kubeflow/katib/pkg/mock/v1beta1/experiment/manifest"
 )
 
@@ -361,7 +361,7 @@ spec:
 		Err             bool
 		testDescription string
 	}{
-		// TrialParamters is nil
+		// TrialParameters is nil
 		{
 			Instance: func() *experimentsv1beta1.Experiment {
 				i := newFakeInstance()
@@ -371,6 +371,20 @@ spec:
 			Err:             true,
 			testDescription: "Trial parameters is nil",
 		},
+		// TrialParameters should not be equal to preserved trial information
+		{
+			Instance: func() *experimentsv1beta1.Experiment {
+				i := newFakeInstance()
+				trialParameters := append(
+					i.Spec.TrialTemplate.TrialParameters, experimentsv1beta1.TrialParameterSpec{Name: consts.TrialTemplateTrialName})
+				trialParameters = append(
+					trialParameters, experimentsv1beta1.TrialParameterSpec{Name: consts.TrialTemplateTrialNamespace})
+				i.Spec.TrialTemplate.TrialParameters = trialParameters
+				return i
+			}(),
+			Err:             true,
+			testDescription: "TrialParameters should not be equal to preserved trial information",
+		},
 		// TrialSpec and ConfigMap is nil
 		{
 			Instance: func() *experimentsv1beta1.Experiment {