forked from temporalio/sdk-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathworker.py
More file actions
331 lines (291 loc) · 13.1 KB
/
Copy pathworker.py
File metadata and controls
331 lines (291 loc) · 13.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
"""Worker using SDK Core. (unstable)
Nothing in this module should be considered stable. The API may change.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import TYPE_CHECKING, Awaitable, Callable, List, Optional, Sequence, Tuple
import google.protobuf.internal.containers
from typing_extensions import TypeAlias
import temporalio.api.common.v1
import temporalio.api.history.v1
import temporalio.bridge.client
import temporalio.bridge.proto
import temporalio.bridge.proto.activity_task
import temporalio.bridge.proto.workflow_activation
import temporalio.bridge.proto.workflow_completion
import temporalio.bridge.runtime
import temporalio.bridge.temporal_sdk_bridge
import temporalio.converter
import temporalio.exceptions
from temporalio.bridge.temporal_sdk_bridge import PollShutdownError
@dataclass
class WorkerConfig:
"""Python representation of the Rust struct for configuring a worker."""
namespace: str
task_queue: str
build_id: str
identity_override: Optional[str]
max_cached_workflows: int
max_outstanding_workflow_tasks: int
max_outstanding_activities: int
max_outstanding_local_activities: int
max_concurrent_workflow_task_polls: int
nonsticky_to_sticky_poll_ratio: float
max_concurrent_activity_task_polls: int
no_remote_activities: bool
sticky_queue_schedule_to_start_timeout_millis: int
max_heartbeat_throttle_interval_millis: int
default_heartbeat_throttle_interval_millis: int
max_activities_per_second: Optional[float]
max_task_queue_activities_per_second: Optional[float]
graceful_shutdown_period_millis: int
use_worker_versioning: bool
class Worker:
"""SDK Core worker."""
@staticmethod
def create(client: temporalio.bridge.client.Client, config: WorkerConfig) -> Worker:
"""Create a bridge worker from a bridge client."""
return Worker(
temporalio.bridge.temporal_sdk_bridge.new_worker(
client._runtime._ref, client._ref, config
)
)
@staticmethod
def for_replay(
runtime: temporalio.bridge.runtime.Runtime,
config: WorkerConfig,
) -> Tuple[Worker, temporalio.bridge.temporal_sdk_bridge.HistoryPusher]:
"""Create a bridge replay worker."""
[
replay_worker,
pusher,
] = temporalio.bridge.temporal_sdk_bridge.new_replay_worker(
runtime._ref, config
)
return Worker(replay_worker), pusher
def __init__(self, ref: temporalio.bridge.temporal_sdk_bridge.WorkerRef) -> None:
"""Create SDK core worker from a bridge worker."""
self._ref = ref
async def poll_workflow_activation(
self,
) -> temporalio.bridge.proto.workflow_activation.WorkflowActivation:
"""Poll for a workflow activation."""
return (
temporalio.bridge.proto.workflow_activation.WorkflowActivation.FromString(
await self._ref.poll_workflow_activation()
)
)
async def poll_activity_task(
self,
) -> temporalio.bridge.proto.activity_task.ActivityTask:
"""Poll for an activity task."""
return temporalio.bridge.proto.activity_task.ActivityTask.FromString(
await self._ref.poll_activity_task()
)
async def complete_workflow_activation(
self,
comp: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
) -> None:
"""Complete a workflow activation."""
await self._ref.complete_workflow_activation(comp.SerializeToString())
async def complete_activity_task(
self, comp: temporalio.bridge.proto.ActivityTaskCompletion
) -> None:
"""Complete an activity task."""
await self._ref.complete_activity_task(comp.SerializeToString())
def record_activity_heartbeat(
self, comp: temporalio.bridge.proto.ActivityHeartbeat
) -> None:
"""Record an activity heartbeat."""
self._ref.record_activity_heartbeat(comp.SerializeToString())
def request_workflow_eviction(self, run_id: str) -> None:
"""Request a workflow be evicted."""
self._ref.request_workflow_eviction(run_id)
def initiate_shutdown(self) -> None:
"""Start shutdown of the worker."""
self._ref.initiate_shutdown()
async def finalize_shutdown(self) -> None:
"""Finalize the worker.
This will fail if shutdown hasn't completed fully due to internal
reference count checks.
"""
ref = self._ref
self._ref = None
await ref.finalize_shutdown()
# See https://mypy.readthedocs.io/en/stable/runtime_troubles.html#using-classes-that-are-generic-in-stubs-but-not-at-runtime
if TYPE_CHECKING:
PayloadContainer: TypeAlias = (
google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
temporalio.api.common.v1.Payload
]
)
else:
PayloadContainer: TypeAlias = (
google.protobuf.internal.containers.RepeatedCompositeFieldContainer
)
async def _apply_to_payloads(
payloads: PayloadContainer,
cb: Callable[
[Sequence[temporalio.api.common.v1.Payload]],
Awaitable[List[temporalio.api.common.v1.Payload]],
],
) -> None:
"""Apply API payload callback to payloads."""
if len(payloads) == 0:
return
new_payloads = await cb(payloads)
del payloads[:]
# TODO(cretz): Copy too expensive?
payloads.extend(new_payloads)
async def _apply_to_payload(
payload: temporalio.api.common.v1.Payload,
cb: Callable[
[Sequence[temporalio.api.common.v1.Payload]],
Awaitable[List[temporalio.api.common.v1.Payload]],
],
) -> None:
"""Apply API payload callback to payload."""
new_payload = (await cb([payload]))[0]
payload.metadata.clear()
payload.metadata.update(new_payload.metadata)
payload.data = new_payload.data
async def _decode_payloads(
payloads: PayloadContainer,
codec: temporalio.converter.PayloadCodec,
) -> None:
"""Decode payloads with the given codec."""
return await _apply_to_payloads(payloads, codec.decode)
async def _decode_payload(
payload: temporalio.api.common.v1.Payload,
codec: temporalio.converter.PayloadCodec,
) -> None:
"""Decode a payload with the given codec."""
return await _apply_to_payload(payload, codec.decode)
async def _encode_payloads(
payloads: PayloadContainer,
codec: temporalio.converter.PayloadCodec,
) -> None:
"""Encode payloads with the given codec."""
return await _apply_to_payloads(payloads, codec.encode)
async def _encode_payload(
payload: temporalio.api.common.v1.Payload,
codec: temporalio.converter.PayloadCodec,
) -> None:
"""Decode a payload with the given codec."""
return await _apply_to_payload(payload, codec.encode)
async def decode_activation(
act: temporalio.bridge.proto.workflow_activation.WorkflowActivation,
codec: temporalio.converter.PayloadCodec,
) -> None:
"""Decode the given activation with the codec."""
for job in act.jobs:
if job.HasField("cancel_workflow"):
await _decode_payloads(job.cancel_workflow.details, codec)
elif job.HasField("query_workflow"):
await _decode_payloads(job.query_workflow.arguments, codec)
elif job.HasField("resolve_activity"):
if job.resolve_activity.result.HasField("cancelled"):
await codec.decode_failure(
job.resolve_activity.result.cancelled.failure
)
elif job.resolve_activity.result.HasField("completed"):
if job.resolve_activity.result.completed.HasField("result"):
await _decode_payload(
job.resolve_activity.result.completed.result, codec
)
elif job.resolve_activity.result.HasField("failed"):
await codec.decode_failure(job.resolve_activity.result.failed.failure)
elif job.HasField("resolve_child_workflow_execution"):
if job.resolve_child_workflow_execution.result.HasField("cancelled"):
await codec.decode_failure(
job.resolve_child_workflow_execution.result.cancelled.failure
)
elif job.resolve_child_workflow_execution.result.HasField(
"completed"
) and job.resolve_child_workflow_execution.result.completed.HasField(
"result"
):
await _decode_payload(
job.resolve_child_workflow_execution.result.completed.result, codec
)
elif job.resolve_child_workflow_execution.result.HasField("failed"):
await codec.decode_failure(
job.resolve_child_workflow_execution.result.failed.failure
)
elif job.HasField("resolve_child_workflow_execution_start"):
if job.resolve_child_workflow_execution_start.HasField("cancelled"):
await codec.decode_failure(
job.resolve_child_workflow_execution_start.cancelled.failure
)
elif job.HasField("resolve_request_cancel_external_workflow"):
if job.resolve_request_cancel_external_workflow.HasField("failure"):
await codec.decode_failure(
job.resolve_request_cancel_external_workflow.failure
)
elif job.HasField("resolve_signal_external_workflow"):
if job.resolve_signal_external_workflow.HasField("failure"):
await codec.decode_failure(job.resolve_signal_external_workflow.failure)
elif job.HasField("signal_workflow"):
await _decode_payloads(job.signal_workflow.input, codec)
elif job.HasField("start_workflow"):
await _decode_payloads(job.start_workflow.arguments, codec)
if job.start_workflow.HasField("continued_failure"):
await codec.decode_failure(job.start_workflow.continued_failure)
for val in job.start_workflow.memo.fields.values():
# This uses API payload not bridge payload
new_payload = (await codec.decode([val]))[0]
val.metadata.clear()
val.metadata.update(new_payload.metadata)
val.data = new_payload.data
elif job.HasField("do_update"):
await _decode_payloads(job.do_update.input, codec)
async def encode_completion(
comp: temporalio.bridge.proto.workflow_completion.WorkflowActivationCompletion,
codec: temporalio.converter.PayloadCodec,
) -> None:
"""Recursively encode the given completion with the codec."""
if comp.HasField("failed"):
await codec.encode_failure(comp.failed.failure)
elif comp.HasField("successful"):
for command in comp.successful.commands:
if command.HasField("complete_workflow_execution"):
if command.complete_workflow_execution.HasField("result"):
await _encode_payload(
command.complete_workflow_execution.result, codec
)
elif command.HasField("continue_as_new_workflow_execution"):
await _encode_payloads(
command.continue_as_new_workflow_execution.arguments, codec
)
for val in command.continue_as_new_workflow_execution.memo.values():
await _encode_payload(val, codec)
elif command.HasField("fail_workflow_execution"):
await codec.encode_failure(command.fail_workflow_execution.failure)
elif command.HasField("respond_to_query"):
if command.respond_to_query.HasField("failed"):
await codec.encode_failure(command.respond_to_query.failed)
elif command.respond_to_query.HasField(
"succeeded"
) and command.respond_to_query.succeeded.HasField("response"):
await _encode_payload(
command.respond_to_query.succeeded.response, codec
)
elif command.HasField("schedule_activity"):
await _encode_payloads(command.schedule_activity.arguments, codec)
elif command.HasField("schedule_local_activity"):
await _encode_payloads(command.schedule_local_activity.arguments, codec)
elif command.HasField("signal_external_workflow_execution"):
await _encode_payloads(
command.signal_external_workflow_execution.args, codec
)
elif command.HasField("start_child_workflow_execution"):
await _encode_payloads(
command.start_child_workflow_execution.input, codec
)
for val in command.start_child_workflow_execution.memo.values():
await _encode_payload(val, codec)
elif command.HasField("update_response"):
if command.update_response.HasField("completed"):
await _encode_payload(command.update_response.completed, codec)
elif command.update_response.HasField("rejected"):
await codec.encode_failure(command.update_response.rejected)