Temporal - Fan-out, Rate-Limited Async Activities example

async_api_concurrency.py

"""Concurrency + rate limiting for an activity that calls an ASYNCHRONOUS API.

The API starts a job and returns immediately; the job finishes later and calls
you back — sometimes reporting success, sometimes failure. We model that with
Async Activity Completion: the activity submits the job, saves its task token,
and returns via raise_complete_async(). A callback later completes OR fails the
activity via that token. A failure is retried per the activity's retry policy
(exponential backoff); each retry re-submits the job, so it usually succeeds
within a couple of attempts.

Because the worker slot is freed on return, max_concurrent_activities can't limit
jobs in flight. So concurrency is capped in the WORKFLOW with a plain counter of
how many activities are currently running. The rate is capped by
max_task_queue_activities_per_second on the worker.

Run:  temporal server start-dev      (in another terminal)
      pip install temporalio
      python async_api_concurrency.py
"""

import asyncio
import random
import uuid
from datetime import timedelta

from temporalio import activity, workflow
from temporalio.client import Client
from temporalio.common import RetryPolicy
from temporalio.exceptions import ApplicationError
from temporalio.worker import Worker

TASK_QUEUE = "async-api"

ITEMS = 20              # number of jobs to fan out
MAX_CONCURRENT = 5      # concurrency cap: max jobs in flight at once
STARTS_PER_SECOND = 100 # rate limit: job starts per second
JOB_DURATION = 2.0      # seconds each remote job takes (simulator only)
FAILURE_RATE = 0.3      # fraction of job attempts that report failure (simulator only)

# Each remote job must be completed (via its task token) within this window, or
# the activity times out and is retried. Async completion does NOT stop this
# clock — it runs until the callback completes the activity — so size it above
# your worst-case job duration.
START_TO_CLOSE_TIMEOUT = timedelta(minutes=10)

# On a failure/timeout the activity retries with exponential backoff. Each retry
# re-runs the submit code with a new task token, so keep submission idempotent.
RETRY_POLICY = RetryPolicy(
    initial_interval=timedelta(seconds=1),
    backoff_coefficient=2.0,               # 1s, 2s, 4s, 8s, ...
    maximum_interval=timedelta(seconds=30),
    maximum_attempts=5,
)


# --- Stand-in for your real async API; also records peak in-flight + breaches --
class FakeAPI:
    def __init__(self, max_concurrent, job_duration, failure_rate):
        self.max_concurrent = max_concurrent
        self.job_duration = job_duration
        self.failure_rate = failure_rate
        self.in_flight = self.peak = self.started = self.violations = self.failures = 0
        self.status = {}   # job_id -> "pending" | "succeeded" | "failed"

    def start_job(self):
        self.in_flight += 1
        self.started += 1
        self.peak = max(self.peak, self.in_flight)
        if self.in_flight > self.max_concurrent:
            self.violations += 1
        job_id = f"job-{self.started}"
        self.status[job_id] = "pending"
        # Decide this attempt's fate up front; report it when the job "finishes".
        outcome = "failed" if random.random() < self.failure_rate else "succeeded"
        asyncio.get_running_loop().call_later(self.job_duration, self._finish, job_id, outcome)
        return job_id

    def _finish(self, job_id, outcome):
        self.status[job_id] = outcome
        self.in_flight -= 1
        if outcome == "failed":
            self.failures += 1

    def job_done(self, job_id):
        return self.status[job_id] != "pending"

    def job_succeeded(self, job_id):
        return self.status[job_id] == "succeeded"


api = FakeAPI(MAX_CONCURRENT, JOB_DURATION, FAILURE_RATE)
temporal_client: Client = None  # set in main(), used to complete activities


# --- Activity: submit the async job, complete/fail it later via its task token --
@activity.defn
async def call_async_api(item: int) -> None:
    task_token = activity.info().task_token   # save this to find the activity later
    job_id = api.start_job()

    # Stand-in for your provider's completion webhook: when the remote job
    # finishes, report success or failure out of band. The worker slot is already
    # free by then, which is the whole point of async completion. A reported
    # failure fails the activity, so Temporal retries it per RETRY_POLICY.
    async def report_when_done():
        while not api.job_done(job_id):
            await asyncio.sleep(0.25)
        handle = temporal_client.get_async_activity_handle(task_token=task_token)
        if api.job_succeeded(job_id):
            await handle.complete()
        else:
            await handle.fail(ApplicationError(f"{job_id} failed"))

    asyncio.create_task(report_when_done())
    activity.raise_complete_async()           # don't complete when this returns


# --- Workflow: cap concurrency with a running-count gate -----------------------
@workflow.defn
class FanOutWorkflow:
    def __init__(self):
        self._running = 0            # activities currently in flight

    @workflow.run
    async def run(self, items: int, max_concurrent: int) -> None:
        self._max_concurrent = max_concurrent

        async def process(item: int) -> None:
            await self._acquire_slot()
            try:
                await workflow.execute_activity(
                    call_async_api,
                    item,
                    start_to_close_timeout=START_TO_CLOSE_TIMEOUT,
                    retry_policy=RETRY_POLICY,
                )
            finally:
                self._running -= 1   # release the slot (held across retries too)

        await asyncio.gather(*(process(i) for i in range(items)))

    async def _acquire_slot(self) -> None:
        # Wait until a slot is free, then claim it by incrementing the counter.
        # The re-check after waking matters: when one slot frees up, EVERY waiter's
        # condition becomes true at once, so we re-test and only one coroutine
        # actually takes the slot; the rest loop back and keep waiting. The check
        # and increment happen with no await between them, so they're atomic.
        while True:
            await workflow.wait_condition(lambda: self._running < self._max_concurrent)
            if self._running < self._max_concurrent:
                self._running += 1
                return


async def main() -> None:
    global temporal_client
    temporal_client = await Client.connect("localhost:7233")

    worker = Worker(
        temporal_client,
        task_queue=TASK_QUEUE,
        workflows=[FanOutWorkflow],
        activities=[call_async_api],
        # Rate limit: cap job starts per second (enforced across all workers).
        max_task_queue_activities_per_second=STARTS_PER_SECOND,
    )

    async with worker:
        await temporal_client.execute_workflow(
            FanOutWorkflow.run,
            args=[ITEMS, MAX_CONCURRENT],
            id=f"async-api-{uuid.uuid4()}",
            task_queue=TASK_QUEUE,
        )

    ok = "OK" if api.violations == 0 else "VIOLATED"
    print(f"DONE — items={ITEMS} | attempts={api.started} | failed & retried={api.failures} | "
          f"peak in-flight={api.peak} | cap={MAX_CONCURRENT} | violations={api.violations}  =>  {ok}")


if __name__ == "__main__":
    asyncio.run(main())