Backend - Annotate Workflow with the run name (kubeflow#2902)

Jeffwan · Dec 9, 2020 · 616788e · 616788e
1 parent 2a02e8f
commit 616788e
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 2 deletions.
diff --git a/backend/src/apiserver/resource/resource_manager.go b/backend/src/apiserver/resource/resource_manager.go
@@ -285,6 +285,8 @@ func (r *ResourceManager) CreateRun(apiRun *api.Run) (*model.RunDetail, error) {
 	workflow.OverrideParameters(parameters)
 	// Add label to the workflow so it can be persisted by persistent agent later.
 	workflow.SetLabels(util.LabelKeyWorkflowRunId, runId)
+	// Add run name annotation to the workflow so that it can be logged by the Metadata Writer.
+	workflow.SetAnnotations(util.AnnotationKeyRunName, apiRun.Name)
 	// Replace {{workflow.uid}} with runId
 	err = workflow.ReplaceUID(runId)
 	if err != nil {

diff --git a/backend/src/apiserver/resource/resource_manager_test.go b/backend/src/apiserver/resource/resource_manager_test.go
@@ -297,6 +297,7 @@ func TestCreateRun_ThroughPipelineID(t *testing.T) {
 	expectedRuntimeWorkflow.Spec.Arguments.Parameters = []v1alpha1.Parameter{
 		{Name: "param1", Value: util.StringPointer("world")}}
 	expectedRuntimeWorkflow.Labels = map[string]string{util.LabelKeyWorkflowRunId: "123e4567-e89b-12d3-a456-426655440000"}
+	expectedRuntimeWorkflow.Annotations = map[string]string{util.AnnotationKeyRunName: "run1"}
 	expectedRuntimeWorkflow.Spec.ServiceAccountName = defaultPipelineRunnerServiceAccount
 
 	expectedRunDetail := &model.RunDetail{
@@ -341,6 +342,7 @@ func TestCreateRun_ThroughWorkflowSpec(t *testing.T) {
 	expectedRuntimeWorkflow.Spec.Arguments.Parameters = []v1alpha1.Parameter{
 		{Name: "param1", Value: util.StringPointer("world")}}
 	expectedRuntimeWorkflow.Labels = map[string]string{util.LabelKeyWorkflowRunId: "123e4567-e89b-12d3-a456-426655440000"}
+	expectedRuntimeWorkflow.Annotations = map[string]string{util.AnnotationKeyRunName: "run1"}
 	expectedRuntimeWorkflow.Spec.ServiceAccountName = defaultPipelineRunnerServiceAccount
 	expectedRunDetail := &model.RunDetail{
 		Run: model.Run{
@@ -422,6 +424,7 @@ func TestCreateRun_ThroughPipelineVersion(t *testing.T) {
 	expectedRuntimeWorkflow.Spec.Arguments.Parameters = []v1alpha1.Parameter{
 		{Name: "param1", Value: util.StringPointer("world")}}
 	expectedRuntimeWorkflow.Labels = map[string]string{util.LabelKeyWorkflowRunId: "123e4567-e89b-12d3-a456-426655440000"}
+	expectedRuntimeWorkflow.Annotations = map[string]string{util.AnnotationKeyRunName: "run1"}
 	expectedRuntimeWorkflow.Spec.ServiceAccountName = defaultPipelineRunnerServiceAccount
 
 	expectedRunDetail := &model.RunDetail{

diff --git a/backend/src/apiserver/server/run_server_test.go b/backend/src/apiserver/server/run_server_test.go
@@ -34,6 +34,7 @@ func TestCreateRun(t *testing.T) {
 	expectedRuntimeWorkflow.Spec.Arguments.Parameters = []v1alpha1.Parameter{
 		{Name: "param1", Value: util.StringPointer("world")}}
 	expectedRuntimeWorkflow.Labels = map[string]string{util.LabelKeyWorkflowRunId: "123e4567-e89b-12d3-a456-426655440000"}
+	expectedRuntimeWorkflow.Annotations = map[string]string{util.AnnotationKeyRunName: "123"}
 	expectedRuntimeWorkflow.Spec.ServiceAccountName = "pipeline-runner"
 	expectedRunDetail := api.RunDetail{
 		Run: &api.Run{

diff --git a/backend/src/common/util/consts.go b/backend/src/common/util/consts.go
@@ -44,6 +44,10 @@ const (
 	LabelKeyWorkflowRunId = "pipeline/runid"
 	LabelKeyWorkflowPersistedFinalState = "pipeline/persistedFinalState"
 
+	// LabelKeyWorkflowEpoch is a Workflow annotation key.
+	// It captures the the name of the Run.
+	AnnotationKeyRunName = "pipelines.kubeflow.org/run_name"
+
 	AnnotationKeyIstioSidecarInject = "sidecar.istio.io/inject"
 	AnnotationValueIstioSidecarInjectEnabled = "true"
 	AnnotationValueIstioSidecarInjectDisabled = "false"

diff --git a/backend/src/common/util/workflow.go b/backend/src/common/util/workflow.go
@@ -115,8 +115,8 @@ func isScheduledWorkflow(reference metav1.OwnerReference) bool {
 	}
 
 	if reference.APIVersion == gvk.GroupVersion().String() &&
-			reference.Kind == gvk.Kind &&
-			reference.UID != "" {
+		reference.Kind == gvk.Kind &&
+		reference.UID != "" {
 		return true
 	}
 	return false
@@ -217,6 +217,13 @@ func (w *Workflow) SetLabels(key string, value string) {
 	w.Labels[key] = value
 }
 
+func (w *Workflow) SetAnnotations(key string, value string) {
+	if w.Annotations == nil {
+		w.Annotations = make(map[string]string)
+	}
+	w.Annotations[key] = value
+}
+
 func (w *Workflow) ReplaceUID(id string) error {
 	newWorkflowString := strings.Replace(w.ToStringForStore(), "{{workflow.uid}}", id, -1)
 	var workflow *workflowapi.Workflow