Callabe task name generators (3.0 vs. 2.x)

[dmichaelcarter]

Bug summary

I am doing some spike work on migrating our flows from v2.19 to v3.0. In doing this, I have found what appears to be a bug with callabe task name generators. I often use task vars (more specifically, values in a task dictionary arg) for naming task threads. This is essential in making sense of complex graphs including dozens, if not hundreds of mapped tasks. I often name the task run based on a value such as a table name during ETL processes; thus it would be more of a hassle to gauge impact of failed tasks when they are named by thread ID instead.

• It seems that prefect.runtime.task_run may not be properly awaiting task args

The error I am seeing with my callable name generator (which works in Prefect v2.19):
TypeError: 'PrefectConcurrentFuture' object is not subscriptable

A much more generic error is also printed during these flow failures:

Please wait for all submitted tasks to complete before exiting your flow by calling `.wait()` on the `PrefectFuture` returned from your `.submit()` calls.

Example:

from prefect import flow, task

@task
def say_hello(name):
    print(f"Hello, {name}!")

@flow
def example_flow():
    future = say_hello.submit(name="Marvin")
    future.wait()

example_flow()

Here is a simple flow which replicates my problem:

from time import sleep
from random import randint
from prefect import task, flow
from prefect.runtime import task_run


def generate_task_run_name() -> str:
    return f'{task_run.task_name} - input:  {task_run.parameters["input"]["val"]}'


@task(log_prints=True)
def pre_task() -> bool:
    return True


@task(log_prints=True, task_run_name=generate_task_run_name)
def task_1(input: dict) -> int:
    input_val = input['val']
    output = input_val + 1
    print(f'task_1 - input:  {input_val}, output:  {output}')

    sleep(randint(5,20))
    return output


@task(log_prints=True, task_run_name=generate_task_run_name)
def task_2(input: dict) -> int:
    input_val = input['val']
    output = input_val * 10
    print(f'task_2 - input:  {input_val}, output:  {output}')

    sleep(randint(5,20))
    return output


@task(log_prints=True, task_run_name=generate_task_run_name)
def task_3(input: dict) -> None:
    input_val = input['val']
    print(f'task_3 - input: {input_val}')

    sleep(randint(5,20))
    return


@flow
def my_flow() -> None:
    pre_task()
    inputs: list = [
        {'val': 1, 'something_else': True}
        ,{'val': 2, 'something_else': True}
        ,{'val': 3, 'something_else': True}
        ,{'val': 4, 'something_else': True}
        ,{'val': 5, 'something_else': True}
        ,{'val': 6, 'something_else': True}
        ,{'val': 7, 'something_else': True}
        ,{'val': 8, 'something_else': True}
    ]

    result_1 = task_1.map(input=inputs)
    result_2 = task_2.map(input=result_1)
    final_result = task_3.map(input=result_2)

    final_result.wait()


if __name__ == '__main__':
    my_flow()

Version info (prefect version output)

Git commit:          0894bad4
Built:               Thu, Oct 10, 2024 10:17 AM
OS/Arch:             darwin/arm64
Profile:             default
Server type:         cloud
Pydantic version:    2.9.2
Integrations:
  prefect-dask:      0.2.9

Additional context

No response