Dataflow breaks using TaggedOutputs, "can't pickle WeakDictionary"

Question

we are trying to deploy an Streaming pipeline to Dataflow where we separate in few different "routes" that we manipulate differently the data. We did the complete development with the DirectRunner, and works smoothly as we tested but now, that we did deployed it to Dataflow, it does not work.

The code fails when yielding on the following doFn

class SplitByRoute(beam.DoFn):
   OUTPUT_TAG_ROUTE_ONE= "route_one"
   OUTPUT_TAG_ROUTE_TWO = "route_two"
   OUTPUT_NOT_SUPPORTED = "not_supported"

   def __init__(self):
       beam.DoFn.__init__(self)

   def process(self, elem):
       try:
           route = self.define_route(elem["param"]) # Just tag it depending on param
       except Exception:
           route = None
       logging.info(f"Routed to {route}")
       if route == self.OUTPUT_TAG_ROUTE_ONE:
           yield TaggedOutput(self.OUTPUT_TAG_ROUTE_ONE, elem)
       elif route == self.OUTPUT_TAG_ROUTE_TWO:
           logging.info(f"Element: {elem}")
           yield TaggedOutput(self.OUTPUT_TAG_ROUTE_TWO, elem)
       else:
           yield TaggedOutput(self.OUTPUT_NOT_SUPPORTED, elem)

It does log the element, yield the output and fails with the following error AttributeError: Can't pickle local object 'WeakValueDictionary.__init__..remove' [while running 'generatedPtransform-3196']

Other considerations are that we use taggedOutputs on the pipeline before this DoFn, and it works on Dataflow but this one in particularly fails with the error mentioned. Could it be the memory cache? or something related to it? Where Weakrefs are used? Far as I know, this error happens when you have a class inside another one. Maybe not(?)

Any suggestions so how we could manage this? It's been very frustrating error.

Thank you!!! :)

user_12850364 · Answer

We found the error

As you might know, apache-beam uses dill package to serialize the data between the modules. This let us pickle an instance of a object and send it through the pipeline. The problem was that in self.define_route(elem["param"]), we used that instance of the class and we modified one of it's attributes. As the answer from Samuel Romero says, you can pickle a class, but I didn't really know (and probably someone has to) that if you modify the class instance it can not be pickle again. that's an strage behaviour, I know, so I opened an issue on BEAM https://issues.apache.org/jira/browse/BEAM-10384 if you want to check it out.

I will probably get into it (to understand better the problem) soon or later, but if someone had the same error, the workaround, as I mentioned is to do not modify the instance of a class beeing serialized.

Thanks to anyone who tried to help!

user_11821225 · Answer

As you can read here, Python uses the pickle library for data serialization and it is subject to its limitations. Data serialization is the way processes transfer data between them since they do not share memory space.

Here I found a suggestion about using a fork of multiprocessing module that uses the dill package instead of pickle. This fork is part of the pathos framework (as is the dill package too) and is now called pathos.multiprocess and not pathos.multiprocessing as seen in the reference I mentioned previously.

Dataflow breaks using TaggedOutputs, "can't pickle WeakDictionary"

2 Answers