1from contextlib import contextmanager
2
3from celery import chain, chord, group
4
5
[docs]
6class CanvasTask:
7 """
8 support for canvas tasks
9
10 See Also:
11 * om.runtime.sequence
12 * om.runtime.parallel
13 * om.runtime.mapreduce
14 """
15
16 def __init__(self, canvasfn):
17 self.sigs = []
18 self.canvasfn = canvasfn
19 self.runtime = None
20
21 def add(self, task):
22 self.sigs.append(task)
23
24 def delay(self, *args, **kwargs):
25 return self.apply_async(args=args, kwargs=kwargs)
26
27 def apply_async(self, args=None, kwargs=None, **celery_kwargs):
28 task = self.sigs[-1]
29 if self.canvasfn is chord:
30 sig = task.signature(args=args, kwargs=kwargs, **celery_kwargs, immutable=False)
31 else:
32 # immutable means results are not passed on from task to task
33 sig = task.signature(args=args, kwargs=kwargs, **celery_kwargs, immutable=True)
34 self.sigs[-1] = sig
35 return sig
36
37 def run(self):
38 if self.canvasfn is chord:
39 result = self.canvasfn(self.sigs[:-1])(self.sigs[-1])
40 else:
41 result = self.canvasfn(*self.sigs).apply_async()
42 # add easy result collection
43 self._easy_collect(result)
44 return result
45
46 def _easy_collect(self, result):
47 # traverse graph of results and return a single list
48 # see https://docs.celeryproject.org/en/stable/userguide/canvas.html
49 collect = lambda r: set(
50 {r} | collect(r.parent) if r.parent else {r},
51 )
52 flatten = lambda l: l[0] if isinstance(l[0], list) else l
53 result.collect = lambda: collect(result)
54 result.getall = lambda: flatten([r.get() for r in collect(result)])
55
56
57def make_canvased(canvasfn):
58 @contextmanager
59 def canvased(self):
60 """
61 context manager to support sequenced, parallel and mapreduce tasks
62
63 Usage::
64
65 # run tasks async, in sequence
66 with om.runtime.sequence() as crt:
67 crt.model('mymodel').fit(...)
68 crt.model('mymodel').predict(...)
69 result = crt.run()
70
71 # run tasks async, in parallel
72 with om.runtime.parallel() as crt:
73 crt.model('mymodel').predict(...)
74 crt.model('mymodel').predict(...)
75 result = crt.run()
76
77 # run tasks async, in parallel with a final step
78 with om.runtime.mapreduce() as crt:
79 # map tasks
80 crt.model('mymodel').predict(...)
81 crt.model('mymodel').predict(...)
82 # reduce results - combined is a virtualobj function
83 crt.model('combined').reduce(...)
84 result = crt.run()
85
86 # combined is a virtual obj function, e.g.
87 @virtualobj
88 def combined(data=None, **kwargs):
89 # data is the list of results from each map step
90 return data
91
92 Note that the statements inside the context are
93 executed in sequence, as any normal python code. However,
94 the actual tasks are only executed on calling crt.run()
95
96 Args:
97 self: the runtime
98
99 Returns:
100 OmegaRuntime, for use within context
101 """
102 canvas = CanvasTask(canvasfn)
103 _orig_task = self.task
104
105 def canvas_task(*args, **kwargs):
106 task = _orig_task(*args, **kwargs)
107 canvas.add(task)
108 return canvas
109
110 self.task = canvas_task
111 canvas.runtime = self
112 canvas.runtime.run = canvas.run
113 try:
114 yield canvas.runtime
115 finally:
116 canvas.runtime.task = _orig_task
117 canvas.runtime.run = None
118
119 return canvased
120
121
122canvas_chain = make_canvased(chain)
123canvas_group = make_canvased(group)
124canvas_chord = make_canvased(chord)