Preserve task dependencies to futures passed as parameters in .map

[robalar]

Preserves relationships between a mapped task and its parameters if that parameter is a future, closes #6686.

Example

from prefect import task, flow, get_run_logger


@task
def log_task(x: int):
    logger = get_run_logger()
    logger.info(x)


@task
def flat_numbers() -> list[int]:
    return [0, 1, 2, 3, 4]


@flow
def test_flow():
    x = flat_numbers.submit()
    log_task.map(x)

Will now correctly produce this graph:

Checklist

• This pull request references any related issue by including "closes <link to issue>"
  • If no issue exists and your change is not a small fix, please create an issue first.
• This pull request includes tests or only affects documentation.
• This pull request includes a fix, feature, enhancement, docs, or maintenance label categorizing the change.

[netlify]

❌ Deploy Preview for prefect-orion failed.

Built without sensitive environment variables

Name	Link
🔨 Latest commit	835fd04
🔍 Latest deploy log	https://app.netlify.com/sites/prefect-orion/deploys/631a3bf88806ba0008148a3a

[robalar]

@madkinsz Have pushed a commit that passes all the test cases. It obviously uses wait_for but feel like it should extend when we start persisting the keys too.

I don't love that I had to add a special case to visit_collection but could not find any other way to recuse into resolved values and avoid infinite recursion.

I also found some tests in the same file that test relationships but only via values returned, I believe my tests would sit better next to them. As a side note, have you considered breaking up this file? It's really quite large and hard to orient/find things in.

[zanieb]

Hi again! visit_collection actually has a max_depth field. We just need to pass it through.

Here's a commit that:

• Adds max depth to task input collection and resolution
• Separates collection from resolution to avoid introducing new resolve future type
• Uses max depth to avoid resolving futures below the top level
• Adds extra_task_inputs passing through the engine so we do not have to use wait_for.

See 00d8647

I did not remove any of the functions you introduced and there may be some additional polish here. Are you interested in cherry-picking that commit into here and cleaning things up?

I also found some tests in the same file that test relationships but only via values returned, I believe my tests would sit better next to them.

I'd prefer to keep all the mapping tests separate from those other inputs tests for now since they're a special case. If you want to create a TestTaskMapInputs section that's fine with me.

As a side note, have you considered breaking up this file? It's really quite large and hard to orient/find things in.

We haven't yet! It seems reasonable in the long run but I mostly just navigate by class. We don't have strong norms about using classes vs files. I think I'd prefer files as we grow the test suite though.

[robalar]

Ah, that makes so much sense! The nesting was giving me a massive headache, but it didn't need to be! Feels obvious looking at it..... 😄

I did not remove any of the functions you introduced and there may be some additional polish here. Are you interested in cherry-picking that commit into here and cleaning things up?

Yes very happy to tidy up!

I'd prefer to keep all the mapping tests separate from those other inputs tests for now since they're a special case. If you want to create a TestTaskMapInputs section that's fine with me.

Noted, will split into its own section.

[zanieb]

Excellent work! Since I wrote some of this I'm going to request @bunchesofdonald as a reviewer. This looks good to me though.

[robalar]

@madkinsz thank you for your guidance on this! Not sure I can claim to have done too much of it, but glad to have helped it along!

[bunchesofdonald]

I think this looks great! That test suite is fantastic.

[bunchesofdonald]

I'm not excited about extra_task_inputs, but I also can't think of a better approach at the moment. It may point at a possible need to refactor the way we validate the parameters to a map call or the way we gather at task inputs at some point.

[zanieb]

Yeah me neither. I think ideally we have a way to track lengths on futures and create indexed children of futures. When the future has no known length, we resolve it to determine the length for mapping then access its indices (which creates child futures) that we pass downstream. That way, our existing tracking will work.

That said, users want to create dependencies between tasks in their graph without necessarily waiting for them to finish. extra_task_inputs may be useful for user-defined relationships that do not have runtime implications.

[zanieb]

cc @WillRaphaelson for reference

[robalar]

@madkinsz do you require any more action from me to get this merged? all tests are passing, netlify seems to be failing due to authentication issues to a private repo.

[zanieb]

See #6728 for some notes on that failure, it's expected right now. Thanks a bunch for contributing!