Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix custom task run name rendering #15773

Merged
merged 7 commits into from
Oct 23, 2024
Merged

Fix custom task run name rendering #15773

merged 7 commits into from
Oct 23, 2024

Conversation

zzstoatzz
Copy link
Collaborator

@zzstoatzz zzstoatzz commented Oct 22, 2024

MRE i was using
from random import uniform
from time import sleep
from typing import Any

from prefect import flow, task
from prefect.runtime import task_run


def generate_task_run_name(parameters: dict[str, Any]) -> str:
    """previously we got parameters here from prefect.runtime.task_run.parameters, but as discussed below the values are not guaranteed to be resolved"""
    return f'{task_run.task_name} - input["val"]:  {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[str, Any]) -> dict[str, Any]:
    input_val = input["val"]
    input["val"] = input_val + 1
    print(f'task_1 - input["val"]:  {input_val}, output:  {input}')

    sleep(uniform(1, 5))
    return input


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

    sleep(uniform(1, 5))
    return input


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

    sleep(uniform(1, 5))
    return


@flow
def my_flow() -> None:
    pre_task()
    inputs: list[dict[str, Any]] = [
        {"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()

closes #15747

the issue (which #15770 does not address) is that the TaskRunContext may receive a non-resolved dict of parameters when it is created, and the TaskRunContext is treated as immutable. this itself feels like an incongruence to address - but this PR is mostly concerned with the custom task run name

with CSTRO I'm not sure that we can move the future resolution before the creation / entrance of the TaskRunContext (happy to talk this through sync) while also correctly handling UpstreamTaskError exceptions, so in addition to fixing the templating syntax like task_run_name="using {input[val]}" by deferring the attempt to render the custom name, this PR proposes a new DX when providing a callable task_run_name that needs access to the rendered parameters

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

if the user provides a function that accepts a parameters dict, then we will populate it with the dict of parameters, which at that point is guaranteed to be resolved

if we are okay with this slight pivot, then I think we should deprecate / remove the get_parameters from runtime.task_run, since it does not currently work for cases where futures need to be resolved

Copy link

codspeed-hq bot commented Oct 22, 2024

CodSpeed Performance Report

Merging #15773 will not alter performance

Comparing fix-custom-name-rendering (32348bb) with main (d09ef03)

Summary

✅ 3 untouched benchmarks

@github-actions github-actions bot added the bug Something isn't working label Oct 22, 2024
@zzstoatzz zzstoatzz added the fix A fix for a bug in an existing feature label Oct 22, 2024
@zzstoatzz zzstoatzz self-assigned this Oct 22, 2024
@zzstoatzz zzstoatzz marked this pull request as ready for review October 22, 2024 18:30
@desertaxle desertaxle changed the title Fix custom name rendering Fix custom task run name rendering Oct 23, 2024
Copy link
Member

@desertaxle desertaxle left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense to me!

src/prefect/tasks.py Outdated Show resolved Hide resolved
@zzstoatzz zzstoatzz force-pushed the fix-custom-name-rendering branch 2 times, most recently from fc43d42 to 8dfcc06 Compare October 23, 2024 02:35
src/prefect/tasks.py Outdated Show resolved Hide resolved
@zzstoatzz zzstoatzz force-pushed the fix-custom-name-rendering branch 2 times, most recently from 4339df9 to 08bf7d2 Compare October 23, 2024 16:11
@zzstoatzz zzstoatzz merged commit 3e5d0bd into main Oct 23, 2024
35 checks passed
@zzstoatzz zzstoatzz deleted the fix-custom-name-rendering branch October 23, 2024 17:03
Comment on lines +690 to +695
# If the callable accepts a 'parameters' kwarg, pass the entire parameters dict
if "parameters" in sig.parameters:
task_run_name = task.task_run_name(parameters=parameters)
else:
# If it doesn't expect parameters, call it without arguments
task_run_name = task.task_run_name()
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we pass the parameters if the author of the callable requests parameters

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working fix A fix for a bug in an existing feature
Projects
None yet
Development

Successfully merging this pull request may close these issues.

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