diff --git a/flytekit/models/admin/workflow.py b/flytekit/models/admin/workflow.py index e40307b6ba..50eb8fafd5 100644 --- a/flytekit/models/admin/workflow.py +++ b/flytekit/models/admin/workflow.py @@ -70,13 +70,14 @@ def from_flyte_idl(cls, pb2_object): class Workflow(_common.FlyteIdlEntity): - def __init__(self, id, closure): + def __init__(self, id, closure, short_description): """ :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): @@ -92,11 +93,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): @@ -107,6 +119,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, ) diff --git a/flytekit/models/task.py b/flytekit/models/task.py index 1bd5eda2d6..e443116243 100644 --- a/flytekit/models/task.py +++ b/flytekit/models/task.py @@ -686,13 +686,14 @@ def from_flyte_idl(cls, pb2_object): class Task(_common.FlyteIdlEntity): - def __init__(self, id, closure): + def __init__(self, id, closure, short_description): """ :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): @@ -710,6 +711,13 @@ 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.task_pb2.Task @@ -717,6 +725,7 @@ def to_flyte_idl(self): return _admin_task.Task( closure=self.closure.to_flyte_idl(), id=self.id.to_flyte_idl(), + short_description=self.short_description, ) @classmethod @@ -728,6 +737,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, ) diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index 2e2dcdc22b..e25ee07174 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/tests/flytekit/unit/remote/test_remote.py @@ -380,12 +380,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") @@ -443,8 +446,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( @@ -468,6 +472,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() @@ -486,6 +491,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} @@ -856,7 +862,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")