Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions flytekit/models/admin/workflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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):
Expand All @@ -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,
)


Expand Down
13 changes: 12 additions & 1 deletion flytekit/models/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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
Expand All @@ -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,
)


Expand Down
1 change: 1 addition & 0 deletions tests/flytekit/unit/models/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
11 changes: 8 additions & 3 deletions tests/flytekit/unit/remote/test_remote.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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(
Expand All @@ -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()
Expand All @@ -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}
Expand Down Expand Up @@ -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")
Expand Down