Frontend - Show customized task display names (#1463)

* Frontend - Show customized task display names

* Added customized name test

* Added ContainerOp.set_display_name(name) method

* Stopped writing human_name to display_name annotation for now
Reason: It's a change to existing pipelines.

* Added test for op.set_display_name

* Fix for tests that have workflows with status nodes, but without any spec or templates

* Fixed the test workflow

* Fix linter error
Error: "The key 'metadata' is not sorted alphabetically"

frontend/src/lib/WorkflowParser.test.ts

@@ -303,6 +303,37 @@ describe('WorkflowParser', () => {
       expect(g.node('node1').label).toEqual('node1');
       expect(g.node('exitNode').label).toEqual('onExit - clean');
     });
+
+    it('gives nodes customized labels based on template annotation', () => {
+      const workflow = {
+        metadata: { name: 'testWorkflow' },
+        spec: {
+          templates: [
+            {
+              metadata: {
+                annotations: {
+                  'kubeflow.org/pipelines/task_display_name': 'Customized name',
+                }
+              },
+              name: 'some-template',
+            }
+          ],
+        },
+        status: {
+          nodes: {
+            node1: {
+              id: 'node1',
+              name: 'node1',
+              phase: 'Succeeded',
+              templateName: 'some-template',
+              type: 'Pod',
+            },
+          },
+        }
+      };
+      const g = WorkflowParser.createRuntimeGraph(workflow as any);
+      expect(g.node('node1').label).toEqual('Customized name');
+    });
   });
 
   describe('getNodeInputOutputParams', () => {

frontend/src/lib/WorkflowParser.ts

@@ -79,6 +79,17 @@ export default class WorkflowParser {
         if (node.name === `${workflowName}.onExit`) {
           nodeLabel = `onExit - ${node.templateName}`;
         }
+
+        if (workflow.spec && workflow.spec.templates) {
+          const tmpl = workflow.spec.templates.find(t => !!t && !!t.name && t.name === node.templateName);
+          if (tmpl && tmpl.metadata && tmpl.metadata.annotations) {
+            const displayName = tmpl.metadata.annotations['kubeflow.org/pipelines/task_display_name'];
+            if (displayName) {
+              nodeLabel = displayName;
+            }
+          }
+        }
+
         g.setNode(node.id, {
           height: Constants.NODE_HEIGHT,
           icon: statusToIcon(node.phase as NodePhase, node.startedAt, node.finishedAt, node.message),

frontend/third_party/argo-ui/argo_template.ts

@@ -175,6 +175,13 @@ export interface Inputs {
    */
   parameters?: Parameter[];
 }
+/**
+ * Pod metdata
+ */
+export interface Metadata {
+  annotations?: { [key: string]: string; };
+  labels?: { [key: string]: string; };
+}
 /**
  * Outputs hold parameters, artifacts, and results from a step
  */
@@ -543,6 +550,11 @@ export interface Template {
    * Inputs describe what inputs parameters and artifacts are supplied to this template
    */
   inputs?: Inputs;
+
+  /**
+   * Metdata sets the pods's metadata, i.e. annotations and labels
+   */
+  metadata?: Metadata;
   /**
    * Name is the name of the template
    */

sdk/python/kfp/compiler/_op_to_template.py

@@ -249,4 +249,8 @@ def _op_to_template(op: BaseOp):
     if processed_op.sidecars:
         template['sidecars'] = processed_op.sidecars
 
+    # Display name
+    if processed_op.display_name:
+        template.setdefault('metadata', {}).setdefault('annotations', {})['kubeflow.org/pipelines/task_display_name'] = processed_op.display_name
+
     return template

sdk/python/kfp/dsl/_container_op.py

@@ -674,6 +674,7 @@ class BaseOp(object):
 
         # human_name must exist to construct operator's name
         self.human_name = name
+        self.display_name = None #TODO Set display_name to human_name
         # ID of the current Op. Ideally, it should be generated by the compiler that sees the bigger context.
         # However, the ID is used in the task output references (PipelineParams) which can be serialized to strings.
         # Because of this we must obtain a unique ID right now.
@@ -816,6 +817,9 @@ class BaseOp(object):
         self.sidecars.append(sidecar)
         return self
 
+    def set_display_name(self, name: str):
+        self.display_name = name
+
     def __repr__(self):
         return str({self.__class__.__name__: self.__dict__})
 

sdk/python/tests/compiler/compiler_tests.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import kfp
 import kfp.compiler as compiler
 import kfp.dsl as dsl
 import os
@@ -497,6 +498,27 @@ class TestCompiler(unittest.TestCase):
 
     self._test_op_to_template_yaml(op1, file_base_name='tolerations')
 
+  def test_set_display_name(self):
+    """Test a pipeline with a customized task names."""
+
+    import kfp
+    op1 = kfp.components.load_component_from_text(
+      '''
+name: Component name
+implementation:
+  container:
+    image: busybox
+'''
+    )
+
+    @dsl.pipeline()
+    def some_pipeline():
+      op1().set_display_name('Custom name')
+
+    workflow_dict = kfp.compiler.Compiler()._compile(some_pipeline)
+    template = workflow_dict['spec']['templates'][0]
+    self.assertEqual(template['metadata']['annotations']['kubeflow.org/pipelines/task_display_name'], 'Custom name')
+
   def test_op_transformers(self):
     def some_op():
         return dsl.ContainerOp(