dex/examples/python/dex_examples/patterns/polling/polling_flow.py at main · superdurable/dex · GitHub

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
# Copyright (c) 2022-2026 Super Durable, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
from datetime import timedelta

from dex import (
    AsyncContext,
    Context,
    Flow,
    PersistenceSchema,
    RetryPolicy,
    Step,
    StepDecision,
    StepList,
    StepOptions,
    Stream,
    force_fail,
    go_to,
    graceful_complete,
)

from dex_examples.patterns.polling.job_service import FakeJobService, JobState, JobStatus

# Keep POLL_INTERVAL + JOB_CALL_TIMEOUT <= heartbeat_timeout - 10s.
# MAX_JOB_WAIT bounds one attempt and all retries; Dex enforces it.
POLL_INTERVAL = timedelta(seconds=2)
JOB_CALL_TIMEOUT = timedelta(seconds=10)
MAX_JOB_WAIT = timedelta(minutes=10)

JOB_PROGRESS = Stream("JobProgress", JobStatus, 1 << 20)


# dex:group group-id:failure group-label:"Failure"
# dex:explanation text:"Fail the Flow with the reason polling stopped."
class RecordJobWaitFailure(Step[str]):
    def execute(self, context: Context, job_id: str) -> StepDecision:
        failure = context.recovery_error
        assert failure is not None
        return force_fail(
            f"waiting for job {job_id} failed: {failure.error_type}: {failure.detail}"
        )


# dex:group group-id:polling group-label:"Polling"
# dex:explanation text:"Poll the job with heartbeats until it succeeds or fails."
class AwaitJob(Step[str]):
    def __init__(self, jobs: FakeJobService) -> None:
        self.jobs = jobs

    def get_step_options(self) -> StepOptions:
        return StepOptions(
            execute_method_timeout=MAX_JOB_WAIT,
            heartbeat_timeout=timedelta(minutes=1),
            execute_retry=RetryPolicy(
                initial_interval=timedelta(seconds=1),
                backoff_coefficient=2.0,
                maximum_attempts=3,
                total_duration=MAX_JOB_WAIT,
            ),
        ).on_execute_failure_proceed_to(RecordJobWaitFailure)

    async def execute(self, context: AsyncContext, job_id: str) -> StepDecision:
        reported_state = context.get_last_heartbeat_value(JobState)
        while True:
            status = await self.get_job_status(job_id)
            if status.state is JobState.SUCCEEDED:
                return graceful_complete(status.result)
            if status.state is JobState.FAILED:
                return force_fail(f"job {job_id} failed")
            if status.state != reported_state:
                JOB_PROGRESS.write(context, status)
                reported_state = status.state
            else:
                await context.heartbeat(reported_state)
            await asyncio.sleep(POLL_INTERVAL.total_seconds())

    async def get_job_status(self, job_id: str) -> JobStatus:
        return await asyncio.wait_for(
            self.jobs.get_job_status(job_id), JOB_CALL_TIMEOUT.total_seconds()
        )


# dex:group group-id:start-job group-label:"Start job"
# dex:explanation text:"Start the external job."
class StartJob(Step[None]):
    def __init__(self, jobs: FakeJobService) -> None:
        self.jobs = jobs

    async def execute(self, context: AsyncContext, input: None) -> StepDecision:
        job_id = context.flow_id
        await asyncio.wait_for(
            self.jobs.start_job(job_id), JOB_CALL_TIMEOUT.total_seconds()
        )
        return go_to(AwaitJob, job_id)


class PollingFlow(Flow[None]):
    def __init__(self, jobs: FakeJobService) -> None:
        self.start_job = StartJob(jobs)
        self.await_job = AwaitJob(jobs)
        self.record_job_wait_failure = RecordJobWaitFailure()

    def get_steps(self) -> StepList[None]:
        return StepList.start_step(self.start_job).other_steps(
            self.await_job, self.record_job_wait_failure
        )

    def get_persistence_schema(self) -> PersistenceSchema:
        return PersistenceSchema.of(JOB_PROGRESS)