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 task key comp for 2.x #14707

Merged
merged 3 commits into from
Jul 23, 2024
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
37 changes: 25 additions & 12 deletions src/prefect/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

import datetime
import inspect
import os
from copy import copy
from functools import partial, update_wrapper
from typing import (
Expand Down Expand Up @@ -76,6 +75,8 @@

logger = get_logger("tasks")

NUM_CHARS_DYNAMIC_KEY: int = 8


def task_input_hash(
context: "TaskRunContext", arguments: Dict[str, Any]
Expand Down Expand Up @@ -125,6 +126,28 @@ def retry_backoff_callable(retries: int) -> List[float]:
return retry_backoff_callable


def _generate_task_key(fn: Callable[..., Any]) -> str:
"""Generate a task key based on the function name and source code.
We may eventually want some sort of top-level namespace here to
disambiguate tasks with the same function name in different modules,
in a more human-readable way, while avoiding relative import problems (see #12337).
As long as the task implementations are unique (even if named the same), we should
not have any collisions.
Args:
fn: The function to generate a task key for.
"""
if not hasattr(fn, "__qualname__"):
return to_qualified_name(type(fn))

qualname = fn.__qualname__.split(".")[-1]

code_hash = (
h[:NUM_CHARS_DYNAMIC_KEY] if (h := hash_objects(fn.__code__)) else "unknown"
)

return f"{qualname}-{code_hash}"


@PrefectObjectRegistry.register_instances
class Task(Generic[P, R]):
"""
Expand Down Expand Up @@ -292,17 +315,7 @@ def __init__(

self.tags = set(tags if tags else [])

if not hasattr(self.fn, "__qualname__"):
self.task_key = to_qualified_name(type(self.fn))
else:
try:
task_origin_hash = hash_objects(
self.name, os.path.abspath(inspect.getsourcefile(self.fn))
)
except TypeError:
task_origin_hash = "unknown-source-file"

self.task_key = f"{self.fn.__qualname__}-{task_origin_hash}"
self.task_key = _generate_task_key(self.fn)

self.cache_key_fn = cache_key_fn
self.cache_expiration = cache_expiration
Expand Down
22 changes: 0 additions & 22 deletions tests/test_background_tasks.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,4 @@
import asyncio
import inspect
import os
from pathlib import Path
from typing import AsyncGenerator, Iterable, Tuple
from unittest import mock
Expand All @@ -26,7 +24,6 @@
)
from prefect.task_server import TaskServer
from prefect.utilities.asyncutils import sync_compatible
from prefect.utilities.hashing import hash_objects


@sync_compatible
Expand Down Expand Up @@ -395,22 +392,3 @@ async def bar(x: int, mappable: Iterable) -> Tuple[int, Iterable]:
assert await result_factory.read_parameters(
task_run.state.state_details.task_parameters_id
) == {"x": i + 1, "mappable": ["some", "iterable"]}


class TestTaskKey:
def test_task_key_includes_qualname_and_source_file_hash(self):
def some_fn():
pass

t = Task(fn=some_fn)
source_file = os.path.abspath(inspect.getsourcefile(some_fn))
task_origin_hash = hash_objects(t.name, source_file)
assert t.task_key == f"{some_fn.__qualname__}-{task_origin_hash}"

def test_task_key_handles_unknown_source_file(self, monkeypatch):
def some_fn():
pass

monkeypatch.setattr(inspect, "getsourcefile", lambda x: None)
t = Task(fn=some_fn)
assert t.task_key == f"{some_fn.__qualname__}-unknown-source-file"
14 changes: 14 additions & 0 deletions tests/test_tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,20 @@ def my_task():
assert my_task.name == "another_name"


class TestTaskKey:
def test_task_key_typical_case(self):
@task
def my_task():
pass

assert my_task.task_key.startswith("my_task-")

def test_task_key_after_import(self):
from tests.generic_tasks import noop

assert noop.task_key.startswith("noop-")


class TestTaskRunName:
def test_run_name_default(self):
@task
Expand Down