diff --git a/flytekit/models/admin/workflow.py b/flytekit/models/admin/workflow.py index 0bb179a799..8cc0583995 100644 --- a/flytekit/models/admin/workflow.py +++ b/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, ) diff --git a/flytekit/models/task.py b/flytekit/models/task.py index c9328f9704..6d19f4a260 100644 --- a/flytekit/models/task.py +++ b/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,6 +726,14 @@ 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 @@ -732,6 +741,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 @@ -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, ) diff --git a/tests/flytekit/unit/models/test_tasks.py b/tests/flytekit/unit/models/test_tasks.py index 23e85ed60a..8e4bc3df02 100644 --- a/tests/flytekit/unit/models/test_tasks.py +++ b/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" diff --git a/tests/flytekit/unit/remote/test_remote.py b/tests/flytekit/unit/remote/test_remote.py index e3ae715c9a..8c9cafd8be 100644 --- a/tests/flytekit/unit/remote/test_remote.py +++ b/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")