canvas.py 35 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.canvas
  4. ~~~~~~~~~~~~~
  5. Composing task workflows.
  6. Documentation for some of these types are in :mod:`celery`.
  7. You should import these from :mod:`celery` and not this module.
  8. """
  9. from __future__ import absolute_import, unicode_literals
  10. import sys
  11. from collections import MutableSequence, deque
  12. from copy import deepcopy
  13. from functools import partial as _partial, reduce
  14. from operator import itemgetter
  15. from itertools import chain as _chain
  16. from kombu.utils import cached_property, fxrange, reprcall, uuid
  17. from vine import barrier
  18. from celery._state import current_app
  19. from celery.local import try_import
  20. from celery.result import GroupResult
  21. from celery.utils import abstract
  22. from celery.utils.functional import (
  23. maybe_list, is_list, _regen, regen, chunks as _chunks,
  24. )
  25. from celery.utils.text import truncate
  26. __all__ = ['Signature', 'chain', 'xmap', 'xstarmap', 'chunks',
  27. 'group', 'chord', 'signature', 'maybe_signature']
  28. PY3 = sys.version_info[0] == 3
  29. # json in Python 2.7 borks if dict contains byte keys.
  30. JSON_NEEDS_UNICODE_KEYS = PY3 and not try_import('simplejson')
  31. class _getitem_property(object):
  32. """Attribute -> dict key descriptor.
  33. The target object must support ``__getitem__``,
  34. and optionally ``__setitem__``.
  35. Example:
  36. >>> from collections import defaultdict
  37. >>> class Me(dict):
  38. ... deep = defaultdict(dict)
  39. ...
  40. ... foo = _getitem_property('foo')
  41. ... deep_thing = _getitem_property('deep.thing')
  42. >>> me = Me()
  43. >>> me.foo
  44. None
  45. >>> me.foo = 10
  46. >>> me.foo
  47. 10
  48. >>> me['foo']
  49. 10
  50. >>> me.deep_thing = 42
  51. >>> me.deep_thing
  52. 42
  53. >>> me.deep
  54. defaultdict(<type 'dict'>, {'thing': 42})
  55. """
  56. def __init__(self, keypath):
  57. path, _, self.key = keypath.rpartition('.')
  58. self.path = path.split('.') if path else None
  59. def _path(self, obj):
  60. return (reduce(lambda d, k: d[k], [obj] + self.path) if self.path
  61. else obj)
  62. def __get__(self, obj, type=None):
  63. if obj is None:
  64. return type
  65. return self._path(obj).get(self.key)
  66. def __set__(self, obj, value):
  67. self._path(obj)[self.key] = value
  68. def maybe_unroll_group(g):
  69. """Unroll group with only one member."""
  70. # Issue #1656
  71. try:
  72. size = len(g.tasks)
  73. except TypeError:
  74. try:
  75. size = g.tasks.__length_hint__()
  76. except (AttributeError, TypeError):
  77. return g
  78. else:
  79. return list(g.tasks)[0] if size == 1 else g
  80. else:
  81. return g.tasks[0] if size == 1 else g
  82. def task_name_from(task):
  83. return getattr(task, 'name', task)
  84. def _upgrade(fields, sig):
  85. """Used by custom signatures in .from_dict, to keep common fields."""
  86. sig.update(chord_size=fields.get('chord_size'))
  87. return sig
  88. class Signature(dict):
  89. """Class that wraps the arguments and execution options
  90. for a single task invocation.
  91. Used as the parts in a :class:`group` and other constructs,
  92. or to pass tasks around as callbacks while being compatible
  93. with serializers with a strict type subset.
  94. :param task: Either a task class/instance, or the name of a task.
  95. :keyword args: Positional arguments to apply.
  96. :keyword kwargs: Keyword arguments to apply.
  97. :keyword options: Additional options to :meth:`Task.apply_async`.
  98. Note that if the first argument is a :class:`dict`, the other
  99. arguments will be ignored and the values in the dict will be used
  100. instead.
  101. >>> s = signature('tasks.add', args=(2, 2))
  102. >>> signature(s)
  103. {'task': 'tasks.add', args=(2, 2), kwargs={}, options={}}
  104. """
  105. TYPES = {}
  106. _app = _type = None
  107. @classmethod
  108. def register_type(cls, subclass, name=None):
  109. cls.TYPES[name or subclass.__name__] = subclass
  110. return subclass
  111. @classmethod
  112. def from_dict(self, d, app=None):
  113. typ = d.get('subtask_type')
  114. if typ:
  115. return self.TYPES[typ].from_dict(d, app=app)
  116. return Signature(d, app=app)
  117. def __init__(self, task=None, args=None, kwargs=None, options=None,
  118. type=None, subtask_type=None, immutable=False,
  119. app=None, **ex):
  120. self._app = app
  121. init = dict.__init__
  122. if isinstance(task, dict):
  123. return init(self, task) # works like dict(d)
  124. # Also supports using task class/instance instead of string name.
  125. try:
  126. task_name = task.name
  127. except AttributeError:
  128. task_name = task
  129. else:
  130. self._type = task
  131. init(self,
  132. task=task_name, args=tuple(args or ()),
  133. kwargs=kwargs or {},
  134. options=dict(options or {}, **ex),
  135. subtask_type=subtask_type,
  136. immutable=immutable,
  137. chord_size=None)
  138. def __call__(self, *partial_args, **partial_kwargs):
  139. args, kwargs, _ = self._merge(partial_args, partial_kwargs, None)
  140. return self.type(*args, **kwargs)
  141. def delay(self, *partial_args, **partial_kwargs):
  142. return self.apply_async(partial_args, partial_kwargs)
  143. def apply(self, args=(), kwargs={}, **options):
  144. """Apply this task locally."""
  145. # For callbacks: extra args are prepended to the stored args.
  146. args, kwargs, options = self._merge(args, kwargs, options)
  147. return self.type.apply(args, kwargs, **options)
  148. def _merge(self, args=(), kwargs={}, options={}, force=False):
  149. if self.immutable and not force:
  150. return (self.args, self.kwargs,
  151. dict(self.options, **options) if options else self.options)
  152. return (tuple(args) + tuple(self.args) if args else self.args,
  153. dict(self.kwargs, **kwargs) if kwargs else self.kwargs,
  154. dict(self.options, **options) if options else self.options)
  155. def clone(self, args=(), kwargs={}, **opts):
  156. # need to deepcopy options so origins links etc. is not modified.
  157. if args or kwargs or opts:
  158. args, kwargs, opts = self._merge(args, kwargs, opts)
  159. else:
  160. args, kwargs, opts = self.args, self.kwargs, self.options
  161. s = Signature.from_dict({'task': self.task, 'args': tuple(args),
  162. 'kwargs': kwargs, 'options': deepcopy(opts),
  163. 'subtask_type': self.subtask_type,
  164. 'chord_size': self.chord_size,
  165. 'immutable': self.immutable}, app=self._app)
  166. s._type = self._type
  167. return s
  168. partial = clone
  169. def freeze(self, _id=None, group_id=None, chord=None,
  170. root_id=None, parent_id=None):
  171. opts = self.options
  172. try:
  173. tid = opts['task_id']
  174. except KeyError:
  175. tid = opts['task_id'] = _id or uuid()
  176. if root_id:
  177. opts['root_id'] = root_id
  178. if parent_id:
  179. opts['parent_id'] = parent_id
  180. if 'reply_to' not in opts:
  181. opts['reply_to'] = self.app.oid
  182. if group_id:
  183. opts['group_id'] = group_id
  184. if chord:
  185. opts['chord'] = chord
  186. return self.AsyncResult(tid)
  187. _freeze = freeze
  188. def replace(self, args=None, kwargs=None, options=None):
  189. s = self.clone()
  190. if args is not None:
  191. s.args = args
  192. if kwargs is not None:
  193. s.kwargs = kwargs
  194. if options is not None:
  195. s.options = options
  196. return s
  197. def set(self, immutable=None, **options):
  198. if immutable is not None:
  199. self.set_immutable(immutable)
  200. self.options.update(options)
  201. return self
  202. def set_immutable(self, immutable):
  203. self.immutable = immutable
  204. def set_parent_id(self, parent_id):
  205. self.parent_id = parent_id
  206. def apply_async(self, args=(), kwargs={}, route_name=None, **options):
  207. try:
  208. _apply = self._apply_async
  209. except IndexError: # pragma: no cover
  210. # no tasks for chain, etc to find type
  211. return
  212. # For callbacks: extra args are prepended to the stored args.
  213. if args or kwargs or options:
  214. args, kwargs, options = self._merge(args, kwargs, options)
  215. else:
  216. args, kwargs, options = self.args, self.kwargs, self.options
  217. return _apply(args, kwargs, **options)
  218. def _with_list_option(self, key):
  219. items = self.options.setdefault(key, [])
  220. if not isinstance(items, MutableSequence):
  221. items = self.options[key] = [items]
  222. return items
  223. def append_to_list_option(self, key, value):
  224. items = self._with_list_option(key)
  225. if value not in items:
  226. items.append(value)
  227. return value
  228. def extend_list_option(self, key, value):
  229. items = self._with_list_option(key)
  230. items.extend(maybe_list(value))
  231. def link(self, callback):
  232. return self.append_to_list_option('link', callback)
  233. def link_error(self, errback):
  234. return self.append_to_list_option('link_error', errback)
  235. def flatten_links(self):
  236. return list(_chain.from_iterable(_chain(
  237. [[self]],
  238. (link.flatten_links()
  239. for link in maybe_list(self.options.get('link')) or [])
  240. )))
  241. def __or__(self, other):
  242. if isinstance(self, group):
  243. if isinstance(other, group):
  244. return group(_chain(self.tasks, other.tasks), app=self.app)
  245. return chord(self, body=other, app=self._app)
  246. elif isinstance(other, group):
  247. other = maybe_unroll_group(other)
  248. if not isinstance(self, chain) and isinstance(other, chain):
  249. return chain((self,) + other.tasks, app=self._app)
  250. elif isinstance(other, chain):
  251. return chain(*self.tasks + other.tasks, app=self._app)
  252. elif isinstance(other, Signature):
  253. if isinstance(self, chain):
  254. return chain(*self.tasks + (other,), app=self._app)
  255. return chain(self, other, app=self._app)
  256. return NotImplemented
  257. def __deepcopy__(self, memo):
  258. memo[id(self)] = self
  259. return dict(self)
  260. def __invert__(self):
  261. return self.apply_async().get()
  262. def __reduce__(self):
  263. # for serialization, the task type is lazily loaded,
  264. # and not stored in the dict itself.
  265. return signature, (dict(self),)
  266. def __json__(self):
  267. return dict(self)
  268. def reprcall(self, *args, **kwargs):
  269. args, kwargs, _ = self._merge(args, kwargs, {}, force=True)
  270. return reprcall(self['task'], args, kwargs)
  271. def election(self):
  272. type = self.type
  273. app = type.app
  274. tid = self.options.get('task_id') or uuid()
  275. with app.producer_or_acquire(None) as P:
  276. props = type.backend.on_task_call(P, tid)
  277. app.control.election(tid, 'task', self.clone(task_id=tid, **props),
  278. connection=P.connection)
  279. return type.AsyncResult(tid)
  280. def __repr__(self):
  281. return self.reprcall()
  282. if JSON_NEEDS_UNICODE_KEYS: # pragma: no cover
  283. def items(self):
  284. for k, v in dict.items(self):
  285. yield k.decode() if isinstance(k, bytes) else k, v
  286. @property
  287. def name(self):
  288. # for duck typing compatibility with Task.name
  289. return self.task
  290. @cached_property
  291. def type(self):
  292. return self._type or self.app.tasks[self['task']]
  293. @cached_property
  294. def app(self):
  295. return self._app or current_app
  296. @cached_property
  297. def AsyncResult(self):
  298. try:
  299. return self.type.AsyncResult
  300. except KeyError: # task not registered
  301. return self.app.AsyncResult
  302. @cached_property
  303. def _apply_async(self):
  304. try:
  305. return self.type.apply_async
  306. except KeyError:
  307. return _partial(self.app.send_task, self['task'])
  308. id = _getitem_property('options.task_id')
  309. parent_id = _getitem_property('options.parent_id')
  310. root_id = _getitem_property('options.root_id')
  311. task = _getitem_property('task')
  312. args = _getitem_property('args')
  313. kwargs = _getitem_property('kwargs')
  314. options = _getitem_property('options')
  315. subtask_type = _getitem_property('subtask_type')
  316. chord_size = _getitem_property('chord_size')
  317. immutable = _getitem_property('immutable')
  318. abstract.CallableSignature.register(Signature)
  319. @Signature.register_type
  320. class chain(Signature):
  321. tasks = _getitem_property('kwargs.tasks')
  322. def __init__(self, *tasks, **options):
  323. tasks = (regen(tasks[0]) if len(tasks) == 1 and is_list(tasks[0])
  324. else tasks)
  325. Signature.__init__(
  326. self, 'celery.chain', (), {'tasks': tasks}, **options
  327. )
  328. self._use_link = options.pop('use_link', None)
  329. self.subtask_type = 'chain'
  330. self._frozen = None
  331. def __call__(self, *args, **kwargs):
  332. if self.tasks:
  333. return self.apply_async(args, kwargs)
  334. def clone(self, *args, **kwargs):
  335. s = Signature.clone(self, *args, **kwargs)
  336. s.kwargs['tasks'] = [sig.clone() for sig in s.kwargs['tasks']]
  337. return s
  338. def apply_async(self, args=(), kwargs={}, **options):
  339. # python is best at unpacking kwargs, so .run is here to do that.
  340. app = self.app
  341. if app.conf.task_always_eager:
  342. return self.apply(args, kwargs, **options)
  343. return self.run(args, kwargs, app=app, **(
  344. dict(self.options, **options) if options else self.options))
  345. def run(self, args=(), kwargs={}, group_id=None, chord=None,
  346. task_id=None, link=None, link_error=None, publisher=None,
  347. producer=None, root_id=None, parent_id=None, app=None, **options):
  348. app = app or self.app
  349. use_link = self._use_link
  350. if use_link is None and app.conf.task_protocol == 1:
  351. use_link = True
  352. args = (tuple(args) + tuple(self.args)
  353. if args and not self.immutable else self.args)
  354. if self._frozen:
  355. tasks, results = self._frozen
  356. else:
  357. tasks, results = self.prepare_steps(
  358. args, self.tasks, root_id, parent_id, link_error, app,
  359. task_id, group_id, chord,
  360. )
  361. if results:
  362. if link:
  363. tasks[0].extend_list_option('link', link)
  364. first_task = tasks.pop()
  365. first_task.apply_async(
  366. chain=tasks if not use_link else None, **options)
  367. return results[0]
  368. def freeze(self, _id=None, group_id=None, chord=None,
  369. root_id=None, parent_id=None):
  370. _, results = self._frozen = self.prepare_steps(
  371. self.args, self.tasks, root_id, parent_id, None,
  372. self.app, _id, group_id, chord, clone=False,
  373. )
  374. return results[0]
  375. def prepare_steps(self, args, tasks,
  376. root_id=None, parent_id=None, link_error=None, app=None,
  377. last_task_id=None, group_id=None, chord_body=None,
  378. clone=True, from_dict=Signature.from_dict):
  379. app = app or self.app
  380. # use chain message field for protocol 2 and later.
  381. # this avoids pickle blowing the stack on the recursion
  382. # required by linking task together in a tree structure.
  383. # (why is pickle using recursion? or better yet why cannot python
  384. # do tail call optimization making recursion actually useful?)
  385. use_link = self._use_link
  386. if use_link is None and app.conf.task_protocol == 1:
  387. use_link = True
  388. steps = deque(tasks)
  389. steps_pop = steps.pop
  390. steps_extend = steps.extend
  391. prev_task = None
  392. prev_res = prev_prev_res = None
  393. tasks, results = [], []
  394. i = 0
  395. while steps:
  396. task = steps_pop()
  397. is_first_task, is_last_task = not steps, not i
  398. if not isinstance(task, abstract.CallableSignature):
  399. task = from_dict(task, app=app)
  400. if isinstance(task, group):
  401. task = maybe_unroll_group(task)
  402. # first task gets partial args from chain
  403. if clone:
  404. task = task.clone(args) if is_first_task else task.clone()
  405. elif is_first_task:
  406. task.args = tuple(args) + tuple(task.args)
  407. if isinstance(task, chain):
  408. # splice the chain
  409. steps_extend(task.tasks)
  410. continue
  411. if isinstance(task, group) and prev_task:
  412. # automatically upgrade group(...) | s to chord(group, s)
  413. # for chords we freeze by pretending it's a normal
  414. # signature instead of a group.
  415. tasks.pop()
  416. results.pop()
  417. task = chord(
  418. task, body=prev_task,
  419. task_id=prev_res.task_id, root_id=root_id, app=app,
  420. )
  421. prev_res = prev_prev_res
  422. if is_last_task:
  423. # chain(task_id=id) means task id is set for the last task
  424. # in the chain. If the chord is part of a chord/group
  425. # then that chord/group must synchronize based on the
  426. # last task in the chain, so we only set the group_id and
  427. # chord callback for the last task.
  428. res = task.freeze(
  429. last_task_id,
  430. root_id=root_id, group_id=group_id, chord=chord_body,
  431. )
  432. else:
  433. res = task.freeze(root_id=root_id)
  434. i += 1
  435. if prev_task:
  436. prev_task.set_parent_id(task.id)
  437. if use_link:
  438. # link previous task to this task.
  439. task.link(prev_task)
  440. if prev_res:
  441. prev_res.parent = res
  442. if is_first_task and parent_id is not None:
  443. task.set_parent_id(parent_id)
  444. if link_error:
  445. for errback in maybe_list(link_error):
  446. task.link_error(errback)
  447. tasks.append(task)
  448. results.append(res)
  449. prev_task, prev_prev_res, prev_res = (
  450. task, prev_res, res,
  451. )
  452. if root_id is None and tasks:
  453. root_id = tasks[-1].id
  454. for task in reversed(tasks):
  455. task.options['root_id'] = root_id
  456. return tasks, results
  457. def apply(self, args=(), kwargs={}, **options):
  458. last, fargs = None, args
  459. for task in self.tasks:
  460. res = task.clone(fargs).apply(
  461. last and (last.get(),), **dict(self.options, **options))
  462. res.parent, last, fargs = last, res, None
  463. return last
  464. @classmethod
  465. def from_dict(self, d, app=None):
  466. tasks = d['kwargs']['tasks']
  467. if tasks:
  468. if isinstance(tasks, tuple): # aaaargh
  469. tasks = d['kwargs']['tasks'] = list(tasks)
  470. # First task must be signature object to get app
  471. tasks[0] = maybe_signature(tasks[0], app=app)
  472. return _upgrade(d, chain(*tasks, app=app, **d['options']))
  473. @property
  474. def app(self):
  475. app = self._app
  476. if app is None:
  477. try:
  478. app = self.tasks[0]._app
  479. except (KeyError, IndexError):
  480. pass
  481. return app or current_app
  482. def __repr__(self):
  483. return ' | '.join(repr(t) for t in self.tasks)
  484. class _basemap(Signature):
  485. _task_name = None
  486. _unpack_args = itemgetter('task', 'it')
  487. def __init__(self, task, it, **options):
  488. Signature.__init__(
  489. self, self._task_name, (),
  490. {'task': task, 'it': regen(it)}, immutable=True, **options
  491. )
  492. def apply_async(self, args=(), kwargs={}, **opts):
  493. # need to evaluate generators
  494. task, it = self._unpack_args(self.kwargs)
  495. return self.type.apply_async(
  496. (), {'task': task, 'it': list(it)},
  497. route_name=task_name_from(self.kwargs.get('task')), **opts
  498. )
  499. @classmethod
  500. def from_dict(cls, d, app=None):
  501. return _upgrade(
  502. d, cls(*cls._unpack_args(d['kwargs']), app=app, **d['options']),
  503. )
  504. @Signature.register_type
  505. class xmap(_basemap):
  506. _task_name = 'celery.map'
  507. def __repr__(self):
  508. task, it = self._unpack_args(self.kwargs)
  509. return '[{0}(x) for x in {1}]'.format(task.task,
  510. truncate(repr(it), 100))
  511. @Signature.register_type
  512. class xstarmap(_basemap):
  513. _task_name = 'celery.starmap'
  514. def __repr__(self):
  515. task, it = self._unpack_args(self.kwargs)
  516. return '[{0}(*x) for x in {1}]'.format(task.task,
  517. truncate(repr(it), 100))
  518. @Signature.register_type
  519. class chunks(Signature):
  520. _unpack_args = itemgetter('task', 'it', 'n')
  521. def __init__(self, task, it, n, **options):
  522. Signature.__init__(
  523. self, 'celery.chunks', (),
  524. {'task': task, 'it': regen(it), 'n': n},
  525. immutable=True, **options
  526. )
  527. @classmethod
  528. def from_dict(self, d, app=None):
  529. return _upgrade(
  530. d, chunks(*self._unpack_args(
  531. d['kwargs']), app=app, **d['options']),
  532. )
  533. def apply_async(self, args=(), kwargs={}, **opts):
  534. return self.group().apply_async(
  535. args, kwargs,
  536. route_name=task_name_from(self.kwargs.get('task')), **opts
  537. )
  538. def __call__(self, **options):
  539. return self.apply_async(**options)
  540. def group(self):
  541. # need to evaluate generators
  542. task, it, n = self._unpack_args(self.kwargs)
  543. return group((xstarmap(task, part, app=self._app)
  544. for part in _chunks(iter(it), n)),
  545. app=self._app)
  546. @classmethod
  547. def apply_chunks(cls, task, it, n, app=None):
  548. return cls(task, it, n, app=app)()
  549. def _maybe_group(tasks, app):
  550. if isinstance(tasks, dict):
  551. tasks = signature(tasks, app=app)
  552. if isinstance(tasks, group):
  553. tasks = tasks.tasks
  554. elif isinstance(tasks, abstract.CallableSignature):
  555. tasks = [tasks]
  556. else:
  557. tasks = [signature(t, app=app) for t in tasks]
  558. return tasks
  559. @Signature.register_type
  560. class group(Signature):
  561. tasks = _getitem_property('kwargs.tasks')
  562. def __init__(self, *tasks, **options):
  563. if len(tasks) == 1:
  564. tasks = tasks[0]
  565. if isinstance(tasks, group):
  566. tasks = tasks.tasks
  567. if not isinstance(tasks, _regen):
  568. tasks = regen(tasks)
  569. Signature.__init__(
  570. self, 'celery.group', (), {'tasks': tasks}, **options
  571. )
  572. self.subtask_type = 'group'
  573. @classmethod
  574. def from_dict(self, d, app=None):
  575. return _upgrade(
  576. d, group(d['kwargs']['tasks'], app=app, **d['options']),
  577. )
  578. def __len__(self):
  579. return len(self.tasks)
  580. def _prepared(self, tasks, partial_args, group_id, root_id, app,
  581. CallableSignature=abstract.CallableSignature,
  582. from_dict=Signature.from_dict,
  583. isinstance=isinstance, tuple=tuple):
  584. for task in tasks:
  585. if isinstance(task, CallableSignature):
  586. # local sigs are always of type Signature, and we
  587. # clone them to make sure we do not modify the originals.
  588. task = task.clone()
  589. else:
  590. # serialized sigs must be converted to Signature.
  591. task = from_dict(task, app=app)
  592. if isinstance(task, group):
  593. # needs yield_from :(
  594. unroll = task._prepared(
  595. task.tasks, partial_args, group_id, root_id, app,
  596. )
  597. for taskN, resN in unroll:
  598. yield taskN, resN
  599. else:
  600. if partial_args and not task.immutable:
  601. task.args = tuple(partial_args) + tuple(task.args)
  602. yield task, task.freeze(group_id=group_id, root_id=root_id)
  603. def _apply_tasks(self, tasks, producer=None, app=None, p=None,
  604. add_to_parent=None, chord=None, **options):
  605. app = app or self.app
  606. with app.producer_or_acquire(producer) as producer:
  607. for sig, res in tasks:
  608. sig.apply_async(producer=producer, add_to_parent=False,
  609. chord=sig.options.get('chord') or chord,
  610. **options)
  611. if p:
  612. p.add(res)
  613. res.backend.add_pending_result(res)
  614. yield res # <-- r.parent, etc set in the frozen result.
  615. def _freeze_gid(self, options):
  616. # remove task_id and use that as the group_id,
  617. # if we don't remove it then every task will have the same id...
  618. options = dict(self.options, **options)
  619. options['group_id'] = group_id = (
  620. options.pop('task_id', uuid()))
  621. return options, group_id, options.get('root_id')
  622. def set_parent_id(self, parent_id):
  623. for task in self.tasks:
  624. task.set_parent_id(parent_id)
  625. def apply_async(self, args=(), kwargs=None, add_to_parent=True,
  626. producer=None, **options):
  627. app = self.app
  628. if app.conf.task_always_eager:
  629. return self.apply(args, kwargs, **options)
  630. if not self.tasks:
  631. return self.freeze()
  632. options, group_id, root_id = self._freeze_gid(options)
  633. tasks = self._prepared(self.tasks, args, group_id, root_id, app)
  634. p = barrier()
  635. results = list(self._apply_tasks(tasks, producer, app, p, **options))
  636. result = self.app.GroupResult(group_id, results, ready_barrier=p)
  637. p.finalize()
  638. # - Special case of group(A.s() | group(B.s(), C.s()))
  639. # That is, group with single item that is a chain but the
  640. # last task in that chain is a group.
  641. #
  642. # We cannot actually support arbitrary GroupResults in chains,
  643. # but this special case we can.
  644. if len(result) == 1 and isinstance(result[0], GroupResult):
  645. result = result[0]
  646. parent_task = app.current_worker_task
  647. if add_to_parent and parent_task:
  648. parent_task.add_trail(result)
  649. return result
  650. def apply(self, args=(), kwargs={}, **options):
  651. app = self.app
  652. if not self.tasks:
  653. return self.freeze() # empty group returns GroupResult
  654. options, group_id, root_id = self._freeze_gid(options)
  655. tasks = self._prepared(self.tasks, args, group_id, root_id, app)
  656. return app.GroupResult(group_id, [
  657. sig.apply(**options) for sig, _ in tasks
  658. ])
  659. def set_immutable(self, immutable):
  660. for task in self.tasks:
  661. task.set_immutable(immutable)
  662. def link(self, sig):
  663. # Simply link to first task
  664. sig = sig.clone().set(immutable=True)
  665. return self.tasks[0].link(sig)
  666. def link_error(self, sig):
  667. sig = sig.clone().set(immutable=True)
  668. return self.tasks[0].link_error(sig)
  669. def __call__(self, *partial_args, **options):
  670. return self.apply_async(partial_args, **options)
  671. def _freeze_unroll(self, new_tasks, group_id, chord, root_id, parent_id):
  672. stack = deque(self.tasks)
  673. while stack:
  674. task = maybe_signature(stack.popleft(), app=self._app).clone()
  675. if isinstance(task, group):
  676. stack.extendleft(task.tasks)
  677. else:
  678. new_tasks.append(task)
  679. yield task.freeze(group_id=group_id,
  680. chord=chord, root_id=root_id,
  681. parent_id=parent_id)
  682. def freeze(self, _id=None, group_id=None, chord=None,
  683. root_id=None, parent_id=None):
  684. opts = self.options
  685. try:
  686. gid = opts['task_id']
  687. except KeyError:
  688. gid = opts['task_id'] = uuid()
  689. if group_id:
  690. opts['group_id'] = group_id
  691. if chord:
  692. opts['chord'] = chord
  693. root_id = opts.setdefault('root_id', root_id)
  694. parent_id = opts.setdefault('parent_id', parent_id)
  695. new_tasks = []
  696. # Need to unroll subgroups early so that chord gets the
  697. # right result instance for chord_unlock etc.
  698. results = list(self._freeze_unroll(
  699. new_tasks, group_id, chord, root_id, parent_id,
  700. ))
  701. if isinstance(self.tasks, MutableSequence):
  702. self.tasks[:] = new_tasks
  703. else:
  704. self.tasks = new_tasks
  705. return self.app.GroupResult(gid, results)
  706. _freeze = freeze
  707. def skew(self, start=1.0, stop=None, step=1.0):
  708. it = fxrange(start, stop, step, repeatlast=True)
  709. for task in self.tasks:
  710. task.set(countdown=next(it))
  711. return self
  712. def __iter__(self):
  713. return iter(self.tasks)
  714. def __repr__(self):
  715. return 'group({0.tasks!r})'.format(self)
  716. @property
  717. def app(self):
  718. app = self._app
  719. if app is None:
  720. try:
  721. app = self.tasks[0].app
  722. except (KeyError, IndexError):
  723. pass
  724. return app if app is not None else current_app
  725. @Signature.register_type
  726. class chord(Signature):
  727. def __init__(self, header, body=None, task='celery.chord',
  728. args=(), kwargs={}, app=None, **options):
  729. Signature.__init__(
  730. self, task, args,
  731. dict(kwargs, header=_maybe_group(header, app),
  732. body=maybe_signature(body, app=app)), app=app, **options
  733. )
  734. self.subtask_type = 'chord'
  735. def freeze(self, _id=None, group_id=None, chord=None,
  736. root_id=None, parent_id=None):
  737. if not isinstance(self.tasks, group):
  738. self.tasks = group(self.tasks, app=self.app)
  739. bodyres = self.body.freeze(_id, parent_id=self.id, root_id=root_id)
  740. self.tasks.freeze(
  741. parent_id=parent_id, root_id=root_id, chord=self.body)
  742. self.id = self.tasks.id
  743. self.body.set_parent_id(self.id)
  744. return bodyres
  745. def set_parent_id(self, parent_id):
  746. tasks = self.tasks
  747. if isinstance(tasks, group):
  748. tasks = tasks.tasks
  749. for task in tasks:
  750. task.set_parent_id(parent_id)
  751. self.parent_id = parent_id
  752. @classmethod
  753. def from_dict(self, d, app=None):
  754. args, d['kwargs'] = self._unpack_args(**d['kwargs'])
  755. return _upgrade(d, self(*args, app=app, **d))
  756. @staticmethod
  757. def _unpack_args(header=None, body=None, **kwargs):
  758. # Python signatures are better at extracting keys from dicts
  759. # than manually popping things off.
  760. return (header, body), kwargs
  761. @cached_property
  762. def app(self):
  763. return self._get_app(self.body)
  764. def _get_app(self, body=None):
  765. app = self._app
  766. if app is None:
  767. try:
  768. tasks = self.tasks.tasks # is a group
  769. except AttributeError:
  770. tasks = self.tasks
  771. app = tasks[0]._app
  772. if app is None and body is not None:
  773. app = body._app
  774. return app if app is not None else current_app
  775. def apply_async(self, args=(), kwargs={}, task_id=None,
  776. producer=None, publisher=None, connection=None,
  777. router=None, result_cls=None, **options):
  778. args = (tuple(args) + tuple(self.args)
  779. if args and not self.immutable else self.args)
  780. body = kwargs.get('body') or self.kwargs['body']
  781. kwargs = dict(self.kwargs, **kwargs)
  782. body = body.clone(**options)
  783. app = self._get_app(body)
  784. tasks = (self.tasks.clone() if isinstance(self.tasks, group)
  785. else group(self.tasks, app=app))
  786. if app.conf.task_always_eager:
  787. return self.apply((), kwargs,
  788. body=body, task_id=task_id, **options)
  789. return self.run(tasks, body, args, task_id=task_id, **options)
  790. def apply(self, args=(), kwargs={}, propagate=True, body=None, **options):
  791. body = self.body if body is None else body
  792. tasks = (self.tasks.clone() if isinstance(self.tasks, group)
  793. else group(self.tasks, app=self.app))
  794. return body.apply(
  795. args=(tasks.apply().get(propagate=propagate),),
  796. )
  797. def _traverse_tasks(self, tasks, value=None):
  798. stack = deque(list(tasks))
  799. while stack:
  800. task = stack.popleft()
  801. if isinstance(task, group):
  802. stack.extend(task.tasks)
  803. else:
  804. yield task if value is None else value
  805. def __length_hint__(self):
  806. return sum(self._traverse_tasks(self.tasks, 1))
  807. def run(self, header, body, partial_args, app=None, interval=None,
  808. countdown=1, max_retries=None, eager=False,
  809. task_id=None, **options):
  810. app = app or self._get_app(body)
  811. group_id = uuid()
  812. root_id = body.options.get('root_id')
  813. body.chord_size = self.__length_hint__()
  814. options = dict(self.options, **options) if options else self.options
  815. if options:
  816. options.pop('task_id', None)
  817. body.options.update(options)
  818. results = header.freeze(
  819. group_id=group_id, chord=body, root_id=root_id).results
  820. bodyres = body.freeze(task_id, root_id=root_id)
  821. parent = app.backend.apply_chord(
  822. header, partial_args, group_id, body,
  823. interval=interval, countdown=countdown,
  824. options=options, max_retries=max_retries,
  825. result=results)
  826. bodyres.parent = parent
  827. return bodyres
  828. def __call__(self, body=None, **options):
  829. return self.apply_async((), {'body': body} if body else {}, **options)
  830. def clone(self, *args, **kwargs):
  831. s = Signature.clone(self, *args, **kwargs)
  832. # need to make copy of body
  833. try:
  834. s.kwargs['body'] = s.kwargs['body'].clone()
  835. except (AttributeError, KeyError):
  836. pass
  837. return s
  838. def link(self, callback):
  839. self.body.link(callback)
  840. return callback
  841. def link_error(self, errback):
  842. self.body.link_error(errback)
  843. return errback
  844. def set_immutable(self, immutable):
  845. # changes mutability of header only, not callback.
  846. for task in self.tasks:
  847. task.set_immutable(immutable)
  848. def __repr__(self):
  849. if self.body:
  850. return self.body.reprcall(self.tasks)
  851. return '<chord without body: {0.tasks!r}>'.format(self)
  852. tasks = _getitem_property('kwargs.header')
  853. body = _getitem_property('kwargs.body')
  854. def signature(varies, *args, **kwargs):
  855. app = kwargs.get('app')
  856. if isinstance(varies, dict):
  857. if isinstance(varies, abstract.CallableSignature):
  858. return varies.clone()
  859. return Signature.from_dict(varies, app=app)
  860. return Signature(varies, *args, **kwargs)
  861. subtask = signature # XXX compat
  862. def maybe_signature(d, app=None):
  863. if d is not None:
  864. if (isinstance(d, dict) and
  865. not isinstance(d, abstract.CallableSignature)):
  866. d = signature(d)
  867. if app is not None:
  868. d._app = app
  869. return d
  870. maybe_subtask = maybe_signature # XXX compat