import inspect
from typing import Any
from celery._state import get_current_task
from celery.canvas import Signature
from celery.utils.log import get_task_logger
from firexkit.bag_of_goodies import AutoInjectRegistry, BagOfGoodies
from firexkit.result import FireXResults, FxAsyncResult
logger = get_task_logger(__name__)
# this is where most code import from!!
returns = FireXResults.returns
# Route to whatever queue the currently running task would route to.
AUTO_QUEUE = "auto"
def _resolve_queue_for_current_task(queue: str | None) -> str | None:
"""Resolve ``queue`` the way :meth:`FireXTask.enqueue_child` does.
Signatures are enqueued from plain functions that have no task to call
``enqueue_child`` on, so the task celery is currently running stands in for
the explicit ``self``.
"""
# Deliberately duck-typed: firexkit.task imports this module, so FireXTask
# cannot be imported here to isinstance() against.
resolve_queue = getattr(get_current_task(), "_resolve_queue", None)
if resolve_queue is None:
# No FireX task is running on this thread (the submit process, a background
# thread, a plain celery task), so there is no worker queue to route back
# to. Leaving the queue unset lets the task's own queue -- else
# task_default_queue -- decide, which beats guessing at a queue name.
return None if queue == AUTO_QUEUE else queue
return resolve_queue(queue)
[docs]
class InvalidChainArgsException(Exception):
def __init__(self, msg, wrong_args: dict | None = None):
super().__init__(msg)
self.wrong_args = wrong_args or {}
def _simulate_chain_args_kwargs(
prev_task_return_args: dict[str, Any] | None,
task_pos_args: tuple,
task_kwargs: dict[str, Any],
chain_depth: int,
) -> tuple[
tuple[Any, ...],
dict[str, Any],
]:
simulated_pos_args = task_pos_args
if prev_task_return_args is not None and chain_depth:
simulated_pos_args = tuple(
# don't know why, but in a chain the previous task's results are supplied
# as the fist positional arg. Crazy.
[prev_task_return_args] + list(task_pos_args)
)
simulated_kwargs = task_kwargs | {"chain_depth": chain_depth}
else:
simulated_pos_args = task_pos_args
simulated_kwargs = task_kwargs
return simulated_pos_args, simulated_kwargs
def _task_python_signature(task_obj) -> inspect.Signature:
py_sig = getattr(task_obj, "sig", None)
if py_sig is None:
# the UT app sometimes doesn't set the FireXTask base class, so 'sig' is missing.
py_sig = inspect.signature(task_obj.run)
return py_sig
def _with_adopted_defaults(
params: list[inspect.Parameter],
defaults_by_arg: dict[str, Any],
) -> list[inspect.Parameter]:
"""
Fill in the defaults an overridden service gives to args that this service
declares without one, skipping the ones python's Signature won't accept.
A positional arg can only take a default when every positional arg of the
same kind after it already has one, so the params are walked right to left.
"""
adopted: list[inspect.Parameter] = []
trailing_kind = None
trailing_all_have_defaults = True
for param in reversed(params):
if param.kind != trailing_kind:
# each kind is a run of params that's ordered independently of the others.
trailing_kind = param.kind
trailing_all_have_defaults = True
if (
param.default is param.empty
and param.name in defaults_by_arg
# *args/**kwargs can never have a default.
and param.kind not in (param.VAR_POSITIONAL, param.VAR_KEYWORD)
and (
trailing_all_have_defaults
or param.kind
not in (param.POSITIONAL_ONLY, param.POSITIONAL_OR_KEYWORD)
)
):
param = param.replace(default=defaults_by_arg[param.name])
trailing_all_have_defaults &= param.default is not param.empty
adopted.append(param)
adopted.reverse()
return adopted
def _merged_python_signature(task_obj) -> inspect.Signature:
"""
Get a service's signature extended with the args of the services it overrides.
An override is expected to be given the args of the service it replaces, so
chain arg validation needs to consider both. Args only the overridden service
declares are added as optional keyword-only params, which is always a valid
position and never claims a positional arg the overriding service wants.
"""
# FIXME: bad relationship between core abstractions
from firexkit.task import FireXTask
task_obj: FireXTask
py_sig = _task_python_signature(task_obj)
own_params: dict[str, inspect.Parameter] = dict(py_sig.parameters)
added_params: dict[str, inspect.Parameter] = {}
# the nearest override wins, matching how the args themselves resolve.
orig_defaults: dict[str, Any] = {}
orig_task: FireXTask | None = task_obj.orig
while orig_task:
for orig_arg, orig_param in _task_python_signature(
orig_task
).parameters.items():
# auto-inject args must not have defaults since it makes
# priority resolution too complicated for the user.
if BagOfGoodies.get_auto_inject_type(orig_param.annotation):
pass
elif orig_arg in own_params:
if orig_param.default is not orig_param.empty:
orig_defaults.setdefault(orig_arg, orig_param.default)
elif orig_arg not in added_params and orig_param.kind not in (
orig_param.VAR_POSITIONAL,
orig_param.VAR_KEYWORD,
):
# *args/**kwargs are skipped: they can't carry the default that keeps
# the arg optional, and this service doesn't accept what they'd accept.
added_params[orig_arg] = orig_param.replace(
kind=orig_param.KEYWORD_ONLY,
# an arg this service doesn't declare can't be mandatory for it.
default=(
orig_param.default
if orig_param.default is not orig_param.empty
else None
),
)
orig_task = orig_task.orig
return py_sig.replace(
parameters=sorted(
_with_adopted_defaults(list(own_params.values()), orig_defaults)
+ list(added_params.values()),
# a stable sort keeps each kind's params in their declared order while
# putting the added keyword-only params ahead of any **kwargs.
key=lambda p: p.kind,
),
)
def _fake_validation_bog(
task_obj,
simulated_pos_args: tuple[Any, ...],
simulated_kwargs: dict[str, Any],
) -> BagOfGoodies:
# FIXME: bad relationship between core abstractions
from firexkit.task import FireXTask
task_obj: FireXTask
merged_py_sig = _merged_python_signature(task_obj)
bog = BagOfGoodies(
merged_py_sig,
simulated_pos_args,
simulated_kwargs,
root_orig_task_sig=task_obj.root_orig.sig
if task_obj.is_overriding_task()
else None,
pydantic_validate=task_obj.get_pydantic_validate(),
)
bog.update(
# add fake stuff that can be generated.
{
k: True # we ignore falsy values when resolving????
for k in bog.get_post_pydantic_convert_supplied_arg_names()
}
)
return bog
[docs]
class SignatureX(Signature):
"""
Fake class to make intended support of monkey-patched celery.Signature
more clear. Ideally this class would actually be used, but this is step1.
Ideally FireX would use Celery's builtin extension mechanisms instead of
monkey patching.
"""
[docs]
def clone(self, args=None, kwargs=None, **opts) -> "SignatureX":
cloned = super().clone(args=args, kwargs=kwargs, **opts)
return type(self)(
task=self.type,
args=cloned.args,
kwargs=cloned.kwargs,
options=cloned.options,
subtask_type=cloned.subtask_type,
immutable=cloned.immutable,
app=self.app,
)
[docs]
def verify_args(self) -> None:
from firexkit.task import (
FireXTask, # FIXME: bad relationship between core abstractions
)
prev_task_return_args: dict[str, Any] | None = None
task_names_to_missing_required_arg_names: dict[str, set[str]] = {}
for chain_depth, task_sig in enumerate(self._get_sigs()):
simulated_pos_args, simulated_kwargs = _simulate_chain_args_kwargs(
prev_task_return_args,
task_sig.args,
task_sig.kwargs,
chain_depth,
)
# I think task_sig/task_obj is bound/unbound distinction, or maybe something
# to do with plugins, but it's not clear.
task_obj: FireXTask = self.app.tasks[task_sig.task]
task_bog = _fake_validation_bog(
task_obj, simulated_pos_args, simulated_kwargs
)
unbound_required_arg_names = task_bog.get_unbound_required_arg_names()
if unbound_required_arg_names:
task_names_to_missing_required_arg_names[task_obj.name] = (
unbound_required_arg_names
)
# If any of the previous keys has a dynamic return, then we can't do any validation
if task_obj.has_dynamic_returns():
break
undefined_indirect = task_bog.get_args_to_indirect_value_keys()
if undefined_indirect:
txt = "\n".join([f"{k}: {v}" for k, v in undefined_indirect.items()])
raise InvalidChainArgsException(
msg=f"Service {task_obj.name} indirectly references the following unavailable parameters: \n{txt}",
wrong_args=undefined_indirect,
)
prev_task_return_args = dict(
task_bog.all_supplied_args()
| {
rk: True # we ignore falsy values when resolving???
for rk in task_obj.return_keys
}
)
if task_names_to_missing_required_arg_names:
service_msgs = []
for (
task_name,
arg_names,
) in task_names_to_missing_required_arg_names.items():
service_msgs.append(
f" {', '.join(arg_names)} \t required by {task_name}"
)
msg = "\n".join(service_msgs)
raise InvalidChainArgsException(
msg=f"Missing mandatory arguments: \n{msg}",
wrong_args=task_names_to_missing_required_arg_names,
)
[docs]
def injectArgs(self, **kwargs):
task_sig = self.get_first_sig()
# Don't update kwargs in-place, because that may change the kwargs in the original object
# even where the object was cloned in __or__ above (kwargs isn't copied and is still the same object)
task_sig.kwargs = kwargs | task_sig.kwargs
[docs]
def get_first_sig(self) -> "SignatureX":
try:
return self.tasks[0] # This might be a chain
except AttributeError:
return self # This might be a signature
def _get_sigs(self) -> list["SignatureX"]:
try:
tasks: list[SignatureX] = self.tasks
except AttributeError:
return [self]
else:
if len(tasks) > 1:
return [
t
for t in tasks
# one of many hacks needed due to InjectArgs hack.
if t.name is not None
]
return tasks
[docs]
def remove_inject_args(self):
if self._is_chain() and len(tasks := self.tasks) > 1:
for s in tasks:
if s.name is None and not isinstance(s, InjectArgs):
logger.error(f"removing {s} with name {s.name}")
self.tasks = [
s
for s in tasks
# one of many hacks need due to InjectArgs hack.
if s.name is not None
]
[docs]
def is_multi_chain(self) -> bool:
return len(self._get_sigs()) > 1
[docs]
def declares_forget(self) -> bool:
"""Whether any task of this chain was declared with ``@app.task(forget=True)``.
A chain's results are forgotten as a whole, so one task asking to be forgotten
makes the whole chain forgettable.
"""
return any(
# a task of this chain isn't necessarily registered here, e.g. a signature
# built for a service that only lives on another worker.
getattr(s.app.tasks.get(s.task), "forget", False)
for s in self._get_sigs()
)
[docs]
def delay(self, *partial_args, **partial_kwargs) -> FxAsyncResult:
return super(*partial_args, **partial_kwargs).delay()
[docs]
def enqueue(
self,
block: bool = False,
raise_exception_on_failure: bool = True,
queue: str | None = None,
priority: int | None = None,
soft_time_limit: int | None = None,
) -> FxAsyncResult:
self.remove_inject_args()
self.verify_args()
if resolved_queue := _resolve_queue_for_current_task(queue):
self.set_queue(resolved_queue)
if priority:
self.set_priority(priority)
if soft_time_limit:
self.set_soft_time_limit(soft_time_limit)
result_promise: FxAsyncResult = self.delay()
if block:
result_promise.fx_wait(
raise_on_failure=raise_exception_on_failure,
)
return result_promise
def _set(self, **kwargs):
# one of many hacks needed for InjectArgs
if self.options is not None:
Signature.set(self, **kwargs)
[docs]
def set_use_cache(self, use_cache: bool):
"""Set the :attr:`use_cache` execution option in every task in :attr:`sig`"""
for s in self._get_sigs():
s._set(use_cache=use_cache)
[docs]
def set_priority(self, priority: int):
"""Set the :attr:`priority` execution option in every task in :attr:`sig`"""
for s in self._get_sigs():
s._set(priority=priority)
[docs]
def set_queue(self, queue: str):
"""Set the :attr:`queue` execution option in every task in :attr:`sig`"""
for s in self._get_sigs():
s._set(queue=queue)
[docs]
def set_soft_time_limit(self, soft_time_limit):
"""Set the :attr:`soft_time_limit` execution option in every task in :attr:`sig`"""
for s in self._get_sigs():
s._set(soft_time_limit=soft_time_limit)
[docs]
def set_label(self, label):
self.set(label=label)
[docs]
def get_label(self):
try:
if self.options:
return self.options["label"]
except KeyError:
pass
return "|".join([s.name for s in self._get_sigs()])
def _is_chain(self) -> bool:
return hasattr(self, "tasks")
[docs]
def apply_async_x(
self,
auto_inject_reg: AutoInjectRegistry | None,
) -> FxAsyncResult:
self.remove_inject_args()
first_sig = self.get_first_sig()
if auto_inject_reg and not first_sig.kwargs.get(
AutoInjectRegistry.AUTO_IN_REG_ABOG_KEY
):
first_sig.kwargs[AutoInjectRegistry.AUTO_IN_REG_ABOG_KEY] = auto_inject_reg
if self._is_chain():
for chain_depth, chain_sig in enumerate(self._get_sigs()):
chain_sig.kwargs["chain_depth"] = chain_depth
self.verify_args()
# args & kwargs are expected to be set by prior kludges
fx_r: FxAsyncResult = self.apply_async()
return fx_r
[docs]
def enqueue_and_extract(
self,
queue: str | None = None,
return_keys: str | tuple = (),
raise_exception_on_failure: bool = True,
priority: int | None = None,
soft_time_limit: int | None = None,
block=True,
**_kwargs,
) -> tuple | dict:
if _kwargs:
logger.warning(f"Unexpected kwargs: {_kwargs}")
if not block:
logger.warning(
f"enqueue_and_extract ignored block={block}, "
"since it needs to block in order to extract results"
)
result_promise = self.enqueue(
queue=queue or AUTO_QUEUE,
block=True,
raise_exception_on_failure=raise_exception_on_failure,
priority=priority,
soft_time_limit=soft_time_limit,
)
return result_promise.legacy_extract_results(
return_keys=return_keys,
)
[docs]
class InjectArgs(SignatureX):
def __init__(self, *args, **kwargs):
assert not args, "Inject args accepts no positional args."
self.args = ()
self.kwargs = kwargs
[docs]
def clone(self, args=None, kwargs=None, **opts) -> "InjectArgs":
assert not args, "InjectArgs accepts no positional args."
return type(self)(**(self.kwargs | (kwargs or {})))
def __str__(self):
return f"InjectArgs({', '.join(self.kwargs.keys())})"
@property
def options(self):
return {}
def __or__(self, other) -> SignatureX:
if isinstance(other, InjectArgs):
r = InjectArgs(**(self.kwargs | other.kwargs))
else:
r = other.clone()
# chains and signatures are both handled by this
SignatureX.injectArgs(r, **self.kwargs)
return r
Signature.injectArgs = SignatureX.injectArgs
Signature.set_priority = SignatureX.set_priority
Signature.set_use_cache = SignatureX.set_use_cache
Signature.set_queue = SignatureX.set_queue
Signature.set_soft_time_limit = SignatureX.set_soft_time_limit
Signature.set_label = SignatureX.set_label
Signature.get_label = SignatureX.get_label
Signature.enqueue = SignatureX.enqueue
Signature.verify_args = SignatureX.verify_args
Signature.get_first_sig = SignatureX.get_first_sig
Signature._get_sigs = SignatureX._get_sigs
Signature.apply_async_x = SignatureX.apply_async_x
Signature._is_chain = SignatureX._is_chain
Signature.remove_inject_args = SignatureX.remove_inject_args
Signature._set = SignatureX._set
Signature.is_multi_chain = SignatureX.is_multi_chain
Signature.declares_forget = SignatureX.declares_forget