Skip to content

Commit

Permalink
Fix task run name
Browse files Browse the repository at this point in the history
  • Loading branch information
jlowin committed Jun 16, 2024
1 parent 837ab75 commit e678bf8
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 4 deletions.
6 changes: 5 additions & 1 deletion src/controlflow/core/controller/controller.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,11 @@ def _setup_run(self):
# start tracking tasks
for task in ready_tasks:
if not task._prefect_task.is_started:
task._prefect_task.start()
task._prefect_task.start(
depends_on=[
t.result for t in task.depends_on if t.result is not None
]
)

# if there are no ready tasks, return. This will usually happen because
# all the tasks are complete.
Expand Down
10 changes: 7 additions & 3 deletions src/controlflow/core/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
)

import prefect
from prefect.context import TaskRunContext
from pydantic import (
Field,
PydanticSchemaGenerationError,
Expand Down Expand Up @@ -54,6 +55,11 @@
logger = get_logger(__name__)


def get_task_run_name() -> str:
context = TaskRunContext.get()
return f'Run {context.parameters['self'].friendly_name()}'


class TaskStatus(Enum):
INCOMPLETE = "INCOMPLETE"
SUCCESSFUL = "SUCCESSFUL"
Expand Down Expand Up @@ -307,7 +313,7 @@ async def run_once_async(self, agent: "Agent" = None, flow: "Flow" = None):
controller = controlflow.Controller(tasks=[self], agents=agent, flow=flow)
await controller.run_once_async()

@prefect.task(task_run_name=lambda _, args: f"Run {args['self'].friendly_name()}")
@prefect.task(task_run_name=get_task_run_name)
def _run(
self,
raise_on_error: bool = True,
Expand All @@ -320,8 +326,6 @@ def _run(
"""
from controlflow.core.flow import Flow, get_flow

self._prefect_task.is_started = True

if max_iterations == NOTSET:
max_iterations = controlflow.settings.max_task_iterations
if max_iterations is None:
Expand Down

0 comments on commit e678bf8

Please sign in to comment.