Manually request a retry in Apache Beam

user_555515 0
Viewed 349

I'm using Apache Beam to do ETLs, and have a problem working with a certain set of report API endpoints.

The endpoints are designed to not return data immediately when requested, but just report IDs. Then the client is supposed to check the reports' state using the IDs. If a report is COMPLETE then it is downloaded, otherwise if it is UNDER_CONSTRUCT the client is supposed to retry later.

My attempt is to decorate the process of a DoFn with apache_beam.utils.retry.with_exponential_backoff, and manually raise an exception to make it retry when report is not complete.

However it doesn't seem to work in that way, and after the exception is raised the program terminates immediately without any wait or retry.

I made a minimal example below, where the actual report endpoint replaced with some mocks.

from apache_beam import Pipeline, PTransform, DoFn, ParDo
from apache_beam.utils.retry import with_exponential_backoff


class CreateReports(PTransform):
    def expand(self, pvalue):
        return [
            '100000',
            '100001',
            '100002',
        ]


class _MockEndpoint:
    def __init__(self):
        self.count = 0

    def check_report(self, report_id):
        self.count += 1
        if self.count % 3 != 2:
            return {'status': 'UNDER_CONSTRUCT'}
        return {'status': 'COMPLETE', 'data': (report_id, self.count)}


class _ReportNotReady(Exception):
    @classmethod
    def _retry_filter(cls, ex):
        print('Filtering', ex, cls, isinstance(ex, cls))
        return isinstance(ex, cls)


class CheckReport(DoFn):
    endpoint = _MockEndpoint()

    @with_exponential_backoff(retry_filter=_ReportNotReady._retry_filter)
    def process(self, element):
        response = self.endpoint.check_report(report_id=element)
        if response['status'] == 'COMPLETE':
            yield response['data']
        else:
            raise _ReportNotReady('Retry later')



with Pipeline() as p:
    m = p | CreateReports() | ParDo(CheckReport())
    print('done', m)

My questions are

  1. Is there better way to implement the manual retry under certain conditions?
  2. If retry.with_exponential_backoff is the right solution, any critical arguments missing there so it doesn't work?
apache-beam python
2 Answers

Here is my takes on your issue. For it I look at the github apache beam code and try to update the code to work of what I believe it should be handled base on the current retries documentation page.

  1. I look at your code and perform the following updates which I think is what you are looking for. Under this update I was able to perform the operation with the retries.
from apache_beam import Pipeline, PTransform, DoFn, ParDo
from apache_beam.utils.retry import with_exponential_backoff


class CreateReports(PTransform):
    def expand(self, pvalue):
        return [
            '100000',
            '100001',
            '100002',
        ]

def retry_on_value_error_filter(exception):
    return isinstance(exception, Exception)

class _MockEndpoint:
    def __init__(self):
        self.count = 0

    @with_exponential_backoff(num_retries=3,retry_filter=retry_on_value_error_filter)
    def check_report(self, report_id):
        self.count += 1
        if self.count % 3 != 2:
            raise Exception('Retry later, at the moment is under construct')
        return {'status': 'COMPLETE', 'data': (report_id, self.count)}

class CheckReport(DoFn):
    endpoint = _MockEndpoint()

    def process(self, element):
      try:
        response = self.endpoint.check_report(report_id=element)
        if response['status'] == 'COMPLETE':
            yield response['data']
      except Exception as ex:
        print(ex)
        return

with Pipeline() as p:
    m = p | CreateReports() | ParDo(CheckReport())
    print('done', m)
  1. This work as intended, but due the lack of proper documentation and samples its kind of hard to decide if its the recommended option. Regardless, I strongly recommend to check the available documentation about retries and retry.with_exponential_backoff ( on the same page ) to better understand its behaviour.
user_145515670

I would do this by handling the exception in the process method, possibly via invoking a decorated function from your process method. E.g., spreading everything out explicitly.

class MyDoFn(beam.DoFn):
  def call_flaky_function(self, element):
    [stuff that may raise an exception]

  @retry.with_exponential_backoff(...)
  def process_with_retries(self, element):
    self.call_flaky_function(element)

  def process(self, element):
    yield self.process_with_retries(element)

Your call_flaky_function could itself raise an exception if the status is not complete. Alternatively, you could just have a do or while loop and implement the exponential backoff yourself with ever-increasing sleeps.

user_5823330
Related