Task in CombineFn is not properly finalized Apache beam

Viewed 380

I'm using Python 3.7 SDK of apache beam 2.17.0 for dataflow. Code is running locally, but I gather data from pubsub. I try to combine per key and everything goes fine until the pipeline calls the "merge_accumulators" function. From this point on, all the underlying code is executed twice. After debugging and going deep in the source code, I found the task is not properly finalized and that is why it is executed twice.

This is the pipeline code:

        options = {
        "runner": "DirectRunner",
        "streaming": True,
        "save_main_session": True
    }

    p = beam.Pipeline(options = PipelineOptions(flags=[], **options))

    processRows = (p
                  |'Read from topic' >> beam.io.ReadFromPubSub(subscription=get_subscription_address())
                  |'Filter do not track' >> beam.ParDo(TakeOutNoTrack)
                  |'Map Data' >> beam.ParDo(mapData)
                   |'Filter metatags' >> beam.ParDo(filterMetatags)
                  |'Label admin' >> beam.ParDo(labelAdmin)
                  |'Process row' >> beam.ParDo(processRow)
                 )

    sessionRow = (processRows
                  |'Add timestamp' >> beam.Map(lambda x: window.TimestampedValue(x, x['timestamp']))
                  |'Key on uuid' >> beam.Map(lambda x: (x['capture_uuid'], x))
                  |'User session window' >> beam.WindowInto(window.Sessions(config_triggers['session_gap']),
                                                              trigger=trigger.AfterWatermark(
                                                                  early=trigger.AfterCount(config_triggers['after_count'])),
                                                              accumulation_mode=trigger.AccumulationMode.ACCUMULATING)
                  |'CombineValues' >> beam.CombinePerKey(JoinSessions())
                 )
    printing = (sessionRow
                |'Printing' >> beam.Map(lambda x: print(x))
                )
    print('running pipeline')
    p.run().wait_until_finish()
    print('done running the pipeline')
    return

This is the config_triggers:

    config_triggers = {
    "session_gap": 1320,
    "after_count": 1,
    "session_length": 20
}

This is the combine class:

class JoinSessions(beam.CombineFn):
def define_format(self):
    try:
        data = {
            "session_uuid": [],
            "capture_uuid": "",
            "metatags": [],
            "timestamps": [],
            "admin": []
        }

        return data

    except Exception:
        logging.error("error at define data: \n%s" % traceback.format_exc())

def create_accumulator(self):
    try:
        return self.define_format()
    except Exception:
        logging.error("error at create accumulator: \n%s " % traceback.format_exc())

def add_input(self, metatags, input):
    try:
        metatags["session_uuid"].append(input.get('session_uuid'))
        metatags["capture_uuid"] = input.get('capture_uuid')
        metatags["metatags"].append(input.get('metatags'))
        metatags["timestamps"].append(input.get('timestamp'))
        metatags["admin"].append(input.get('admin'))
        print('test add_input')

        return metatags
    except Exception:
        logging.error("error at add input: \n%s" % traceback.format_exc())

def merge_accumulators(self, accumulators):
    # print(accumulators)
    try:
        global test_counter
        tags_accumulated = self.define_format()

        for tags in accumulators:
            tags_accumulated["session_uuid"] += tags['session_uuid']
            tags_accumulated["capture_uuid"] += tags['capture_uuid']
            tags_accumulated["metatags"] += tags['metatags']
            tags_accumulated["timestamps"] += tags['timestamps']
            tags_accumulated["admin"] += tags['admin']
        test_counter += 1
        print('counter = ', test_counter)
        return tags_accumulated

    except Exception:
        logging.error("Error at merge Accumulators: \n%s" % traceback.format_exc())

def extract_output(self, metatags):
    try:
        # print('New input in the pipeline:')
        # print('Extract_output: ')
        # print(metatags, '\n')

        return metatags
    except Exception:
        logging.error("error at return input: \n%s" % traceback.format_exc())

No errors are thrown nor exceptions or some kind of information. Just the output of the 'printing' label is printed twice. Also the global counter goes up two times, but there is just one data entry in the pipeline. The print on the add_input function is executed just once.

I'm new to dataflow, so, sorry if I made a dumb mistake.

1 Answers

I think this is due to the trigger you have set.

trigger=trigger.AfterWatermark(early=trigger.AfterCount(config_triggers['after_count'])),accumulation_mode=trigger.AccumulationMode.ACCUMULATING)

config_triggers['after_count] is 1.

So you've set up a trigger that fires after every 1 element and you also accumulate elements produced by trigger firings. So a second trigger within the same window will include elements from the first trigger and so on. See following for details regarding setting the trigger of your pipeline correctly according to your use-case.

https://beam.apache.org/documentation/programming-guide/#triggers

Related