How to log a progress bar on the console in Apache Beam?

user_34935 0
Viewed 76

I am trying out Apache Beam with the DirectRunner, i.e. just on a UNIX command line. I'm using the python SDK.

How do I get a particular transform to report how much it has processed at regular intervals? Over time, this gives me an idea of how fast it is running.

My model is tqdm, which shows a lightweight counter that can be a progress bar if you give a total.

P.S. I'm using bounded datasets (read from a file), and apache-beam 2.40.0.

apache-beam python
1 Answers

If you only want to track your own DoFn then you could simply update an existing tqdm instance from your process function.

But if you want to track all of them then you would need to get a bit hacky. In Beam everything is a DoFn, including stuff like beam.Map and built-in IOs and the stuff you create yourself. The very important part of a DoFn is a process function that gets called on each element – this is a good place to count stuff, except that you cannot hook into it since lots of classes override it. But they have a make_fn function which seems to be public and we can use it to wrap anything else with our tracker:

import apache_beam as beam
from tqdm import tqdm
from time import sleep
from unittest import mock

tqdms = {}


def make_fn(self, fn, has_side_inputs):
    dofn = None
    if isinstance(fn, beam.DoFn):
        dofn = fn
    else:
        dofn = beam.core.CallableWrapperDoFn(fn)

    def wrapper(fn):
        def wrapped(*args, **kwargs):
            if self.label not in tqdms:
                tqdms[self.label] = tqdm(desc=self.label, delay=0)
            # Show running status in case something is really slow. Won't work
            # when fn is a generator.
            tqdms[self.label].set_postfix(running=True)
            result = fn(*args, **kwargs)
            tqdms[self.label].update(1)
            tqdms[self.label].set_postfix(running=False)
            return result

        return wrapped

    dofn.process = wrapper(dofn.process)
    return dofn


def main(pipeline):
    (
        pipeline
        | beam.Create(range(100))
        | beam.Reshuffle()
        | beam.FlatMap(lambda x: sleep(0.01))
        | beam.combiners.Count.Globally()
        | beam.Map(lambda x: sleep(1))
    )


if __name__ == '__main__':
    with beam.Pipeline() as pipeline:
        with mock.patch.object(beam.transforms.core.ParDo, 'make_fn', make_fn):
            main(pipeline)

This works in Beam 2.40.0 and it outputs about 9 progress bars (some are internals) that get dynamically updated as the job runs.

user_12790050
Related