add short_description to Task and Workflow model (#3224)

Signed-off-by: Niels Bantilan <[email protected]>

Author: cosmicBboy

flytekit/models/admin/workflow.py

@@ -72,13 +72,14 @@ def from_flyte_idl(cls, pb2_object):
 
 
 class Workflow(_common.FlyteIdlEntity):
-    def __init__(self, id, closure):
+    def __init__(self, id, closure, short_description=None):
         """
         :param flytekit.models.core.identifier.Identifier id:
         :param WorkflowClosure closure:
         """
         self._id = id
         self._closure = closure
+        self._short_description = short_description
 
     @property
     def id(self):
@@ -94,11 +95,22 @@ def closure(self):
         """
         return self._closure
 
+    @property
+    def short_description(self):
+        """
+        :rtype: str
+        """
+        return self._short_description
+
     def to_flyte_idl(self):
         """
         :rtype: flyteidl.admin.workflow_pb2.Workflow
         """
-        return _admin_workflow.Workflow(id=self.id.to_flyte_idl(), closure=self.closure.to_flyte_idl())
+        return _admin_workflow.Workflow(
+            id=self.id.to_flyte_idl(),
+            closure=self.closure.to_flyte_idl(),
+            short_description=self.short_description,
+        )
 
     @classmethod
     def from_flyte_idl(cls, pb2_object):
@@ -109,6 +121,7 @@ def from_flyte_idl(cls, pb2_object):
         return cls(
             id=_identifier.Identifier.from_flyte_idl(pb2_object.id),
             closure=WorkflowClosure.from_flyte_idl(pb2_object.closure),
+            short_description=pb2_object.short_description,
         )
 
 

flytekit/models/task.py

@@ -701,13 +701,14 @@ def from_flyte_idl(cls, pb2_object):
 
 
 class Task(_common.FlyteIdlEntity):
-    def __init__(self, id, closure):
+    def __init__(self, id, closure, short_description=None):
         """
         :param flytekit.models.core.identifier.Identifier id: The (project, domain, name) identifier for this task.
         :param TaskClosure closure: The closure for the underlying workload.
         """
         self._id = id
         self._closure = closure
+        self._short_description = short_description
 
     @property
     def id(self):
@@ -725,13 +726,22 @@ def closure(self):
         """
         return self._closure
 
+    @property
+    def short_description(self):
+        """
+        The short description of the task.
+        :rtype: str
+        """
+        return self._short_description
+
     def to_flyte_idl(self):
         """
         :rtype: flyteidl.admin.task_pb2.Task
         """
         return _admin_task.Task(
             closure=self.closure.to_flyte_idl(),
             id=self.id.to_flyte_idl(),
+            short_description=self.short_description,
         )
 
     @classmethod
@@ -743,6 +753,7 @@ def from_flyte_idl(cls, pb2_object):
         return cls(
             closure=TaskClosure.from_flyte_idl(pb2_object.closure),
             id=_identifier.Identifier.from_flyte_idl(pb2_object.id),
+            short_description=pb2_object.short_description,
         )
 
 

tests/flytekit/unit/models/test_tasks.py

@@ -299,6 +299,7 @@ def test_task(task_closure):
     obj = task.Task(
         identifier.Identifier(identifier.ResourceType.TASK, "project", "domain", "name", "version"),
         task_closure,
+        "my short description",
     )
     assert obj.id.project == "project"
     assert obj.id.domain == "domain"

tests/flytekit/unit/remote/test_remote.py

@@ -390,12 +390,15 @@ def get_compiled_workflow_closure():
 def test_fetch_lazy(remote):
     mock_client = remote._client
     mock_client.get_task.return_value = Task(
-        id=Identifier(ResourceType.TASK, "p", "d", "n", "v"), closure=LIST_OF_TASK_CLOSURES[0]
+        id=Identifier(ResourceType.TASK, "p", "d", "n", "v"),
+        closure=LIST_OF_TASK_CLOSURES[0],
+        short_description="task description",
     )
 
     mock_client.get_workflow.return_value = Workflow(
         id=Identifier(ResourceType.TASK, "p", "d", "n", "v"),
         closure=WorkflowClosure(compiled_workflow=get_compiled_workflow_closure()),
+        short_description="workflow description",
     )
 
     lw = remote.fetch_workflow_lazy(name="wn", version="v")
@@ -453,8 +456,9 @@ def test_launch_backfill(remote):
     mock_client.get_workflow.return_value = Workflow(
         id=Identifier(ResourceType.WORKFLOW, "p", "d", "daily2", "v"),
         closure=WorkflowClosure(
-            compiled_workflow=CompiledWorkflowClosure(primary=ser_wf, sub_workflows=[], tasks=tasks)
+            compiled_workflow=CompiledWorkflowClosure(primary=ser_wf, sub_workflows=[], tasks=tasks),
         ),
+        short_description="workflow description",
     )
 
     wf = remote.launch_backfill(
@@ -478,6 +482,7 @@ def test_fetch_workflow_with_branch(mock_promote, mock_workflow, remote):
     mock_client.get_workflow.return_value = Workflow(
         id=Identifier(ResourceType.TASK, "p", "d", "n", "v"),
         closure=WorkflowClosure(compiled_workflow=MagicMock()),
+        short_description="workflow description",
     )
 
     admin_launch_plan = MagicMock()
@@ -496,6 +501,7 @@ def test_fetch_workflow_with_nested_branch(mock_promote, mock_workflow, remote):
     mock_client.get_workflow.return_value = Workflow(
         id=Identifier(ResourceType.TASK, "p", "d", "n", "v"),
         closure=WorkflowClosure(compiled_workflow=MagicMock()),
+        short_description="workflow description",
     )
     admin_launch_plan = MagicMock()
     admin_launch_plan.spec = {"workflow_id": 123}
@@ -866,7 +872,6 @@ def workflow1():
     assert isinstance(registered_workflow, FlyteWorkflow)
     assert registered_workflow.id == Identifier(ResourceType.WORKFLOW, "flytesnacks", "development", "tests.flytekit.unit.remote.test_remote.workflow1", "dummy_version")
 
-
 @mock.patch("flytekit.remote.remote.get_serializable")
 @mock.patch("flytekit.remote.remote.FlyteRemote.fetch_launch_plan")
 @mock.patch("flytekit.remote.remote.FlyteRemote.raw_register")