Save Logical Type Registry and Coder Registry on cloudpickle save main session

Author: Abacn

CHANGES.md

@@ -94,6 +94,10 @@
 * PulsarIO has now changed support status from incomplete to experimental. Both read and writes should now minimally
   function (un-partitioned topics, without schema support, timestamp ordered messages for read) (Java)
   ([#36141](https://github.com/apache/beam/issues/36141)).
+* (Python) Logical type and coder registry are saved for pipelines with `save_main_session` pipeline option enabled in
+  the case of default pickler. This fixes a side effect of switching to cloudpickle as default pickler in Beam 2.65.0.
+  Previously for cloudpickle `--save_main_session` pipeline option was not honored
+  ([#35738](https://github.com/apache/beam/issues/35738)).
 
 ## Known Issues
 

sdks/python/apache_beam/coders/typecoders.py

@@ -114,6 +114,14 @@ def _register_coder_internal(
       typehint_coder_class: Type[coders.Coder]) -> None:
     self._coders[typehint_type] = typehint_coder_class
 
+  @staticmethod
+  def _normalize_typehint_type(typehint_type):
+    if typehint_type.__module__ == '__main__':
+      # See https://github.com/apache/beam/issues/21541
+      # TODO(robertwb): Remove once all runners are portable.
+      return getattr(typehint_type, '__name__', str(typehint_type))
+    return typehint_type
+
   def register_coder(
       self, typehint_type: Any,
       typehint_coder_class: Type[coders.Coder]) -> None:
@@ -123,11 +131,8 @@ def register_coder(
           'Received %r instead.' % typehint_coder_class)
     if typehint_type not in self.custom_types:
       self.custom_types.append(typehint_type)
-    if typehint_type.__module__ == '__main__':
-      # See https://github.com/apache/beam/issues/21541
-      # TODO(robertwb): Remove once all runners are portable.
-      typehint_type = getattr(typehint_type, '__name__', str(typehint_type))
-    self._register_coder_internal(typehint_type, typehint_coder_class)
+    self._register_coder_internal(
+        self._normalize_typehint_type(typehint_type), typehint_coder_class)
 
   def get_coder(self, typehint: Any) -> coders.Coder:
     if typehint and typehint.__module__ == '__main__':
@@ -170,9 +175,15 @@ def get_coder(self, typehint: Any) -> coders.Coder:
       coder = self._fallback_coder
     return coder.from_type_hint(typehint, self)
 
-  def get_custom_type_coder_tuples(self, types):
+  def get_custom_type_coder_tuples(self, types=None):
     """Returns type/coder tuples for all custom types passed in."""
-    return [(t, self._coders[t]) for t in types if t in self.custom_types]
+    return [(t, self._coders[self._normalize_typehint_type(t)])
+            for t in self.custom_types if (types is None or t in types)]
+
+  def load_custom_type_coder_tuples(self, type_coder):
+    """Load type/coder tuples into coder registry."""
+    for t, c in type_coder:
+      self.register_coder(t, c)
 
   def verify_deterministic(self, key_coder, op_name, silent=True):
     if not key_coder.is_deterministic():

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -196,12 +196,35 @@ def _lock_reducer(obj):
 
 
 def dump_session(file_path):
-  # It is possible to dump session with cloudpickle. However, since references
-  # are saved it should not be necessary. See https://s.apache.org/beam-picklers
-  pass
+  # 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
+
+  with _pickle_lock, open(file_path, 'wb') as file:
+    coder_reg = typecoders.registry.get_custom_type_coder_tuples()
+    logicaltype_reg = schemas.LogicalType._known_logical_types.copy()
+
+    pickler = cloudpickle.CloudPickler(file)
+    # TODO(https://github.com/apache/beam/issues/18500) add file system registry
+    # once implemented
+    pickler.dump({"coder": coder_reg, "logicaltype": logicaltype_reg})
 
 
 def load_session(file_path):
-  # It is possible to load_session with cloudpickle. However, since references
-  # are saved it should not be necessary. See https://s.apache.org/beam-picklers
-  pass
+  from apache_beam.typehints import schemas
+  from apache_beam.coders import typecoders
+
+  with _pickle_lock, open(file_path, 'rb') as file:
+    registries = cloudpickle.load(file)
+    if type(registries) != dict:
+      raise ValueError(
+          "Faled loading session: expected dict, get {}", type(registries))
+    if "coder" in registries:
+      typecoders.registry.load_custom_type_coder_tuples(registries["coder"])
+    else:
+      _LOGGER.warning('No coder registry found in saved session')
+    if "logicaltype" in registries:
+      schemas.LogicalType._known_logical_types.load(registries["logicaltype"])
+    else:
+      _LOGGER.warning('No logical type registry found in saved session')

sdks/python/apache_beam/internal/pickler.py

@@ -76,6 +76,14 @@ def load_session(file_path):
   return desired_pickle_lib.load_session(file_path)
 
 
+def is_currently_dill():
+  return desired_pickle_lib == dill_pickler
+
+
+def is_currently_cloudpickle():
+  return desired_pickle_lib == cloudpickle_pickler
+
+
 def set_library(selected_library=DEFAULT_PICKLE_LIB):
   """ Sets pickle library that will be used. """
   global desired_pickle_lib
@@ -93,12 +101,11 @@ def set_library(selected_library=DEFAULT_PICKLE_LIB):
         "Pipeline option pickle_library=dill_unsafe is set, but dill is not "
         "installed. Install dill in job submission and runtime environments.")
 
-  is_currently_dill = (desired_pickle_lib == dill_pickler)
   dill_is_requested = (
       selected_library == USE_DILL or selected_library == USE_DILL_UNSAFE)
 
   # If switching to or from dill, update the pickler hook overrides.
-  if is_currently_dill != dill_is_requested:
+  if is_currently_dill() != dill_is_requested:
     dill_pickler.override_pickler_hooks(selected_library == USE_DILL)
 
   if dill_is_requested:

sdks/python/apache_beam/runners/portability/stager.py

@@ -376,7 +376,6 @@ def create_job_resources(
       pickled_session_file = os.path.join(
           temp_dir, names.PICKLED_MAIN_SESSION_FILE)
       pickler.dump_session(pickled_session_file)
-      # for pickle_library: cloudpickle, dump_session is no op
       if os.path.exists(pickled_session_file):
         resources.append(
             Stager._create_file_stage_to_artifact(

sdks/python/apache_beam/runners/portability/stager_test.py

@@ -200,7 +200,7 @@ def test_with_main_session(self):
   # (https://github.com/apache/beam/issues/21457): Remove the decorator once
   # cloudpickle is default pickle library
   @pytest.mark.no_xdist
-  def test_main_session_not_staged_when_using_cloudpickle(self):
+  def test_main_session_staged_when_using_cloudpickle(self):
     staging_dir = self.make_temp_dir()
     options = PipelineOptions()
 
@@ -209,7 +209,10 @@ def test_main_session_not_staged_when_using_cloudpickle(self):
     # session is saved when pickle_library==cloudpickle.
     options.view_as(SetupOptions).pickle_library = pickler.USE_CLOUDPICKLE
     self.update_options(options)
-    self.assertEqual([stager.SUBMISSION_ENV_DEPENDENCIES_FILE],
+    self.assertEqual([
+        names.PICKLED_MAIN_SESSION_FILE,
+        stager.SUBMISSION_ENV_DEPENDENCIES_FILE
+    ],
                      self.stager.create_and_stage_job_resources(
                          options, staging_location=staging_dir)[1])
 

sdks/python/apache_beam/runners/worker/sdk_worker_main.py

@@ -113,26 +113,25 @@ def create_harness(environment, dry_run=False):
   _LOGGER.info('semi_persistent_directory: %s', semi_persistent_directory)
   _worker_id = environment.get('WORKER_ID', None)
 
-  if pickle_library != pickler.USE_CLOUDPICKLE:
-    try:
-      _load_main_session(semi_persistent_directory)
-    except LoadMainSessionException:
-      exception_details = traceback.format_exc()
-      _LOGGER.error(
-          'Could not load main session: %s', exception_details, exc_info=True)
-      raise
-    except Exception:  # pylint: disable=broad-except
-      summary = (
-          "Could not load main session. Inspect which external dependencies "
-          "are used in the main module of your pipeline. Verify that "
-          "corresponding packages are installed in the pipeline runtime "
-          "environment and their installed versions match the versions used in "
-          "pipeline submission environment. For more information, see: https://"
-          "beam.apache.org/documentation/sdks/python-pipeline-dependencies/")
-      _LOGGER.error(summary, exc_info=True)
-      exception_details = traceback.format_exc()
-      deferred_exception = LoadMainSessionException(
-          f"{summary} {exception_details}")
+  try:
+    _load_main_session(semi_persistent_directory)
+  except LoadMainSessionException:
+    exception_details = traceback.format_exc()
+    _LOGGER.error(
+        'Could not load main session: %s', exception_details, exc_info=True)
+    raise
+  except Exception:  # pylint: disable=broad-except
+    summary = (
+        "Could not load main session. Inspect which external dependencies "
+        "are used in the main module of your pipeline. Verify that "
+        "corresponding packages are installed in the pipeline runtime "
+        "environment and their installed versions match the versions used in "
+        "pipeline submission environment. For more information, see: https://"
+        "beam.apache.org/documentation/sdks/python-pipeline-dependencies/")
+    _LOGGER.error(summary, exc_info=True)
+    exception_details = traceback.format_exc()
+    deferred_exception = LoadMainSessionException(
+        f"{summary} {exception_details}")
 
   _LOGGER.info(
       'Pipeline_options: %s',
@@ -352,6 +351,14 @@ class LoadMainSessionException(Exception):
 
 def _load_main_session(semi_persistent_directory):
   """Loads a pickled main session from the path specified."""
+  if pickler.is_currently_dill():
+    warn_msg = ' Functions defined in __main__ (interactive session) may fail.'
+    err_msg = ' Functions defined in __main__ (interactive session) will ' \
+      'almost certainly fail.'
+  elif pickler.is_currently_cloudpickle():
+    warn_msg = ' User registered objects (e.g. schema, logical type) through' \
+        'registeries may not be effective'
+    err_msg = ''
   if semi_persistent_directory:
     session_file = os.path.join(
         semi_persistent_directory, 'staged', names.PICKLED_MAIN_SESSION_FILE)
@@ -361,21 +368,18 @@ def _load_main_session(semi_persistent_directory):
       # This can happen if the worker fails to download the main session.
       # Raise a fatal error and crash this worker, forcing a restart.
       if os.path.getsize(session_file) == 0:
-        # Potenitally transient error, unclear if still happening.
-        raise LoadMainSessionException(
-            'Session file found, but empty: %s. Functions defined in __main__ '
-            '(interactive session) will almost certainly fail.' %
-            (session_file, ))
-      pickler.load_session(session_file)
+        if pickler.is_currently_dill():
+          # Potenitally transient error, unclear if still happening.
+          raise LoadMainSessionException(
+              'Session file found, but empty: %s.%s' % (session_file, err_msg))
+        else:
+          _LOGGER.warning('Empty session file: %s.%s', warn_msg, session_file)
+      else:
+        pickler.load_session(session_file)
     else:
-      _LOGGER.warning(
-          'No session file found: %s. Functions defined in __main__ '
-          '(interactive session) may fail.',
-          session_file)
+      _LOGGER.warning('No session file found: %s.%s', warn_msg, session_file)
   else:
-    _LOGGER.warning(
-        'No semi_persistent_directory found: Functions defined in __main__ '
-        '(interactive session) may fail.')
+    _LOGGER.warning('No semi_persistent_directory found: %s', warn_msg)
 
 
 if __name__ == '__main__':

sdks/python/apache_beam/typehints/schemas.py

@@ -706,6 +706,11 @@ def copy(self):
     copy.by_language_type.update(self.by_language_type)
     return copy
 
+  def load(self, another):
+    self.by_urn.update(another.by_urn)
+    self.by_logical_type.update(another.by_logical_type)
+    self.by_language_type.update(another.by_language_type)
+
 
 LanguageT = TypeVar('LanguageT')
 RepresentationT = TypeVar('RepresentationT')