This commit is contained in:
Nuno Campos
2025-05-08 16:49:01 -07:00
parent 5cff35d1c3
commit 1d977f1c09
7 changed files with 32 additions and 18 deletions
+5 -3
View File
@@ -40,7 +40,7 @@ def task(
*,
name: Optional[str] = None,
retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
cache: Optional[CachePolicy[P]] = None,
cache_policy: Optional[CachePolicy[P]] = None,
) -> Callable[
[Union[Callable[P, Awaitable[T]], Callable[P, T]]],
Callable[P, SyncAsyncFuture[T]],
@@ -58,7 +58,7 @@ def task(
*,
name: Optional[str] = None,
retry: Optional[Union[RetryPolicy, Sequence[RetryPolicy]]] = None,
cache: Optional[CachePolicy[P]] = None,
cache_policy: Optional[CachePolicy[P]] = None,
) -> Union[
Callable[
[Union[Callable[P, Awaitable[T]], Callable[P, T]]],
@@ -144,7 +144,9 @@ def task(
# handle regular functions / partials / callable classes, etc.
func.__name__ = name
call_func = functools.partial(call, func, retry=retry_policies, cache=cache)
call_func = functools.partial(
call, func, retry=retry_policies, cache=cache_policy
)
object.__setattr__(call_func, "_is_pregel_task", True)
return functools.update_wrapper(call_func, func)
+10 -7
View File
@@ -113,12 +113,12 @@ class PregelTaskWrites(NamedTuple):
class Call:
__slots__ = ("func", "input", "retry", "cache", "callbacks")
__slots__ = ("func", "input", "retry", "cache_policy", "callbacks")
func: Callable
input: tuple[tuple[Any, ...], dict[str, Any]]
retry: Optional[Sequence[RetryPolicy]]
cache: Optional[CachePolicy]
cache_policy: Optional[CachePolicy]
callbacks: Callbacks
def __init__(
@@ -127,13 +127,13 @@ class Call:
input: tuple[tuple[Any, ...], dict[str, Any]],
*,
retry: Optional[Sequence[RetryPolicy]],
cache: Optional[CachePolicy],
cache_policy: Optional[CachePolicy],
callbacks: Callbacks,
) -> None:
self.func = func
self.input = input
self.retry = retry
self.cache = cache
self.cache_policy = cache_policy
self.callbacks = callbacks
@@ -639,13 +639,14 @@ def prepare_single_task(
if parent_ns
else b"",
(identifier(call.func) or "__dynamic__").encode(),
call.cache.key(*call.input[0], **call.input[1]),
call.cache_policy.key(*call.input[0], **call.input[1]),
)
)
),
call.cache.ttl,
call.cache_policy.ttl,
call.cache_policy.refresh,
)
if call.cache
if call.cache_policy
else None,
task_id,
task_path,
@@ -771,6 +772,7 @@ def prepare_single_task(
)
),
proc.cache_policy.ttl,
proc.cache_policy.refresh,
)
if proc.cache_policy
else None,
@@ -915,6 +917,7 @@ def prepare_single_task(
)
),
proc.cache_policy.ttl,
proc.cache_policy.refresh,
)
if proc.cache_policy
else None,
+6 -2
View File
@@ -242,12 +242,16 @@ def call(
func: Callable[P, T],
*args: Any,
retry: Optional[Sequence[RetryPolicy]] = None,
cache: Optional[CachePolicy] = None,
cache_policy: Optional[CachePolicy] = None,
**kwargs: Any,
) -> SyncAsyncFuture[T]:
config = get_config()
impl = config[CONF][CONFIG_KEY_CALL]
fut = impl(
func, (args, kwargs), retry=retry, cache=cache, callbacks=config["callbacks"]
func,
(args, kwargs),
retry=retry,
cache=cache_policy,
callbacks=config["callbacks"],
)
return fut
+2 -2
View File
@@ -1055,7 +1055,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
if cached := {
t.cache_key.key: t
for t in self.tasks.values()
if t.cache_key and not t.writes
if t.cache_key and not t.cache_key.refresh and not t.writes
}:
for key, values in self.cache.get(cached.keys()).items():
cached[key].writes.extend(values)
@@ -1231,7 +1231,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
if cached := {
t.cache_key.key: t
for t in self.tasks.values()
if t.cache_key and not t.writes
if t.cache_key and not t.cache_key.refresh and not t.writes
}:
for key, values in (await self.cache.aget(cached.keys())).items():
cached[key].writes.extend(values)
+3 -3
View File
@@ -523,7 +523,7 @@ def _call(
input: Any,
*,
retry: Optional[Sequence[RetryPolicy]] = None,
cache: Optional[CachePolicy] = None,
cache_policy: Optional[CachePolicy] = None,
callbacks: Callbacks = None,
futures: weakref.ref[FuturesDict],
schedule_task: weakref.ref[
@@ -545,7 +545,7 @@ def _call(
if next_task := schedule_task()( # type: ignore[misc]
task(), # type: ignore[arg-type]
scratchpad.call_counter(),
Call(func, input, retry=retry, cache=cache, callbacks=callbacks),
Call(func, input, retry=retry, cache_policy=cache_policy, callbacks=callbacks),
):
if match_cached_writes:
match_cached_writes()
@@ -630,7 +630,7 @@ def _acall(
if next_task := schedule_task()( # type: ignore[misc]
task(), # type: ignore[arg-type]
scratchpad.call_counter(),
Call(func, input, retry=retry, cache=cache, callbacks=callbacks),
Call(func, input, retry=retry, cache_policy=cache, callbacks=callbacks),
):
if fut := next(
(
+5
View File
@@ -140,6 +140,9 @@ class CachePolicy(NamedTuple, Generic[P]):
ttl: Optional[int] = None
"""Time to live for the cache entry in seconds. If None, the entry never expires."""
refresh: bool = False
"""Whether to force a refresh of the cache entry when it is accessed."""
@dataclasses.dataclass(**_DC_KWARGS)
class Interrupt:
@@ -191,6 +194,8 @@ class CacheKey(NamedTuple):
"""Key for the cache entry."""
ttl: Optional[int]
"""Time to live for the cache entry in seconds."""
refresh: bool
"""Whether to force a refresh of the cache entry when it is accessed."""
@dataclasses.dataclass(**_T_DC_KWARGS)
+1 -1
View File
@@ -6607,7 +6607,7 @@ def test_multiple_interrupts_functional_cache(
counter = 0
@task(cache=CachePolicy())
@task(cache_policy=CachePolicy())
def double(x: int) -> int:
"""Increment the counter."""
nonlocal counter