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

Author: Abacn

CHANGES.md

@@ -94,6 +94,9 @@
 * 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 schema registry are saved for pipeline cloudpickle with `save_main_session` pipeline option.
+  enabled. This fixes a side effect of switching to cloudpickle as default pickler in Beam 2.65.0. Previously
+  `--save_main_session` was not honored for cloudpickle ([#35738](https://github.com/apache/beam/issues/35738)).
 
 ## Known Issues
 

sdks/python/apache_beam/internal/cloudpickle_pickler.py

@@ -196,12 +196,28 @@ 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 known Beam Registries
+  from apache_beam.typehints.schemas import LogicalType
+  from apache_beam.typehints import schema_registry
+
+  with _pickle_lock, open(file_path, 'wb') as file:
+    pickler = cloudpickle.CloudPickler(file)
+    # Logical Type registry
+    pickler.dump(LogicalType._known_logical_types)
+    # Schema registry
+    pickler.dump(schema_registry.SCHEMA_REGISTRY)
 
 
 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
+  with _pickle_lock, open(file_path, 'rb') as file:
+    from apache_beam.typehints.schemas import LogicalType
+    from apache_beam.typehints import schema_registry
+
+    known_logical_type = cloudpickle.load(file)
+    schema_registry = cloudpickle.load(file)
+
+    LogicalType._known_logical_types = known_logical_type
+    schema_registry.SCHEMA_REGISTRY = schema_registry

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__':