Issues pickling/unpickling Flags with dill

[rasmi]

I came across this issue when using the save_main_session in Apache Beam 2.4.0 with absl.flags (absl-py version 0.2.1). save_main_session pickles then unpickles the variables in the main session, flags.FLAGS included.

The following error is raised:

Traceback (most recent call last):
  File "test.py", line 46, in <module>
    app.run(main)
  File "venv/local/lib/python2.7/site-packages/absl/app.py", line 274, in run
    _run_main(main, args)
  File "venv/local/lib/python2.7/site-packages/absl/app.py", line 238, in _run_main
    sys.exit(main(argv))
  File "test.py", line 43, in main
    _ = output | beam.io.WriteToText(FLAGS.output_file)
  File "venv/local/lib/python2.7/site-packages/apache_beam/pipeline.py", line 389, in __exit__
    self.run().wait_until_finish()
  File "venv/local/lib/python2.7/site-packages/apache_beam/pipeline.py", line 369, in run
    self.to_runner_api(), self.runner, self._options).run(False)
  File "venv/local/lib/python2.7/site-packages/apache_beam/pipeline.py", line 379, in run
    pickler.dump_session(os.path.join(tmpdir, 'main_session.pickle'))
  File "venv/local/lib/python2.7/site-packages/apache_beam/internal/pickler.py", line 242, in dump_session
    dill.load_session(file_path)
  File "venv/local/lib/python2.7/site-packages/dill/dill.py", line 363, in load_session
    module = unpickler.load()
  File "/usr/lib/python2.7/pickle.py", line 864, in load
    dispatch[key](self)
  File "/usr/lib/python2.7/pickle.py", line 1221, in load_build
    setstate = getattr(inst, "__setstate__", None)
  File "venv/local/lib/python2.7/site-packages/absl/flags/_flagvalues.py", line 468, in __getattr__
    fl = self._flags()
  File "venv/local/lib/python2.7/site-packages/absl/flags/_flagvalues.py", line 141, in _flags
    return self.__dict__['__flags']
KeyError: '__flags'

Example code:

"""
Demonstrating incompatibility between absl.flags and
save_main_session in Apache Beam's Python SDK.
"""
import re
import six
import apache_beam as beam

from absl import app
from absl import flags


FLAGS = flags.FLAGS

flags.DEFINE_string(
    'output_file',
    'output.txt',
    help='Output filename.')

def main(argv):
  del argv # Unused.

  pipeline_options = beam.pipeline.PipelineOptions()
  # Uncomment the next line to break things.
  #pipeline_options.view_as(beam.pipeline.SetupOptions).save_main_session = True

  with beam.Pipeline(options=pipeline_options) as p:
    lines = p | beam.Create(
      ['This is a test of compatibility issues between absl and beam.',
       'Since absl.flags cannot be pickled correctly, it cannot be used',
       'with the save_main_session pipeline option.'])

    counts = (
        lines
        | 'Split' >> (beam.FlatMap(lambda x: re.findall(r'[A-Za-z\']+', x))
                      .with_output_types(six.text_type))
        | 'PairWithOne' >> beam.Map(lambda x: (x, 1))
        | 'GroupAndSum' >> beam.CombinePerKey(sum))

    output = counts | 'Format' >> beam.Map(
      lambda (word, count): '%s: %s' % (word, count))

    _ = output | beam.io.WriteToText(FLAGS.output_file)

if __name__ == '__main__':
  app.run(main)

[shoyer]

I'm not sure it makes sense to pickle FLAGS at all.

If you pickle/unpickle absl.flags.FLAGS, should it be a separate independent copy or should it be the same as the global FLAGS?

[rasmi]

I can see how pickling/unpickling FLAGS in those two ways would lead to issues with state management in different cases.

In this case, Beam is pickling absl.flags.FLAGS after the program is invoked and flags are parsed. After that, FLAGS is only being accessed/referenced and not otherwise modified, so an independent copy of FLAGS at the time it is pickled would work fine. But I can imagine cases where additional calls to absl.flags methods after unpickling may cause issues with FLAGS, in which case you'd want the global FLAGS.

I am biased towards my use case, but I think the first case would be sufficient in most situations.

[yuhc]

So any ideas about how to pickle an independent copy of a global FLAGS?

[shoyer]

My tentative plan here is to make pickling FLAGS an error. You'll be to able to use copy.deepcopy() on FLAGS, though.

[yuhc]

My tentative plan here is to make pickling FLAGS an error. You'll be to able to use copy.deepcopy() on FLAGS, though.

copy.deepcopy() seems not working now (Python 2.7) :( .

DUP=copy.deepcopy(FLAGS)
print pickle.loads(pickle.dumps(DUP))

throws the same error. Please let me know if I'm wrong.
But maybe one alternative way is to implement the pickler/unpickler methods based on the output of FLAGS.flag_value_dict().

[shoyer]

I have an internal change (still under review inside Google for now) that will fix deepcopy() and will hopefully make it into public version of abseil-py soon.

[yuhc]

That's awesome! I'm looking forward to that release. I also saw the commit #d08b66 which makes much sense to me.
Is there an ETA for the fix?

[shoyer]

d08b66f is the original version of this fix but it broke some internal users so it was rolled back.

I don't have an ETA, but it will be coming soon.

[yuhc]

Good to know. Thanks!
And sorry that I have one more question: Can I assume that the __deepcopy__ is enabled by fixing result.__dict__ = copy.deepcopy(self.__dict__) with correct assignment? If so, can I implement the __getstate__ to return the __dict__ and __setstate__ to update the destination object's __dict__ in the same way as __deepcopy__? I'm curious what memo is and whether it can help fix the deepcopy.
I ask this because I'm in a hurry to pickle FLAGS before it's officially supported.

[yuhc]

Please ignore my last question. I applied a quick hack which meets my need link .

[shoyer]

This has been merged in 8526d52

Note that you still can't pickle flags, but you do get a useful error message.

[rasmi]

Fantastic, thanks @shoyer. I don't especially mind the inability to pickle since save_main_session in Beam is an antipattern. Deepcopy functionality is sufficient, you can close this issue if you don't have separate plans/needs to address pickling.

[sehargul-123]

File "DeepSpeech.py", line 11, in
import absl.app
File "/home/sehar/venv/lib/python3.6/site-packages/absl/app.py", line 40, in
from absl import flags
File "/home/sehar/venv/lib/python3.6/site-packages/absl/flags/init.py", line 41, in
from absl.flags import _defines
File "/home/sehar/venv/lib/python3.6/site-packages/absl/flags/_defines.py", line 31, in
from absl.flags import _flagvalues
File "/home/sehar/venv/lib/python3.6/site-packages/absl/flags/_flagvalues.py", line 27, in
import logging
File "/home/sehar/DeepSpeech/logging.py", line 6, in
from util.flags import FLAGS
File "/home/sehar/DeepSpeech/util/flags.py", line 6, in
FLAGS = absl.flags.FLAGS
AttributeError: module 'absl' has no attribute 'flags'

[sehargul-123]

i am using tensorflow 1.14 gpu based kindly help me to resolve this issue

[yilei]

Could you file a separate issue? This isn't related to pickling/unpickling flags.

[yilei]

Closing this issue according to rasmi's comment

[sehargul-123]

but will u still help me with this issue