Set save_main_session default to true for cloudpickle and introduce overwrite flag

Author: Abacn

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -254,8 +254,8 @@ def _lock_reducer(obj):
 def dump_session(file_path):
   # Since References are saved (https://s.apache.org/beam-picklers), we only
   # dump supported Beam Registries (currently only logical type registry)
-  from apache_beam.typehints import schemas
   from apache_beam.coders import typecoders
+  from apache_beam.typehints import schemas
 
   with _pickle_lock, open(file_path, 'wb') as file:
     coder_reg = typecoders.registry.get_custom_type_coder_tuples()
@@ -268,8 +268,8 @@ def dump_session(file_path):
 
 
 def load_session(file_path):
-  from apache_beam.typehints import schemas
   from apache_beam.coders import typecoders
+  from apache_beam.typehints import schemas
 
   with _pickle_lock, open(file_path, 'rb') as file:
     registries = cloudpickle.load(file)

sdks/python/apache_beam/options/pipeline_options.py

@@ -1672,14 +1672,22 @@ def _add_argparse_args(cls, parser):
         choices=['cloudpickle', 'default', 'dill', 'dill_unsafe'])
     parser.add_argument(
         '--save_main_session',
-        default=False,
+        default=None,
         action='store_true',
         help=(
             'Save the main session state so that pickled functions and classes '
             'defined in __main__ (e.g. interactive session) can be unpickled. '
             'Some workflows do not need the session state if for instance all '
             'their functions/classes are defined in proper modules '
             '(not __main__) and the modules are importable in the worker. '))
+    parser.add_argument(
+        '--disable_save_main_session',
+        action='store_false',
+        dest='save_main_session',
+        help=(
+            'Disable saving the main session state. It is enabled/disabled by'
+            'default for cloudpickle/dill pickler. See "save_main_session".'))
+
     parser.add_argument(
         '--sdk_location',
         default='default',
@@ -1780,10 +1788,23 @@ def _add_argparse_args(cls, parser):
             'If not specified, the default Maven Central repository will be '
             'used.'))
 
+  def _handle_load_main_session(self, validator):
+    save_main_session = getattr(self, 'save_main_session')
+    if save_main_session is None:
+      # save_main_session default to False for dill, while default to true
+      # for cloudpickle
+      pickle_library = getattr(self, 'pickle_library')
+      if pickle_library in ['default', 'cloudpickle']:
+        setattr(self, 'save_main_session', True)
+      else:
+        setattr(self, 'save_main_session', False)
+    return []
+
   def validate(self, validator):
     errors = []
     errors.extend(validator.validate_container_prebuilding_options(self))
     errors.extend(validator.validate_pickle_library(self))
+    errors.extend(self._handle_load_main_session(validator))
     return errors