Use standard pickle, and make pickler an option

We used cloudpickle to make it easy to author tests where the
test worker is defined in the test function. The problem is that
cloudpickle will also try to serialize by value classes defined in the
__main__ package, which leads to serialization problems when e.g. the worker class
references in its members unpickleable things.

Also, revert #452 as it is addressed by this change.

Pull Request: #455

Author: mtrofin

compiler_opt/distributed/buffered_scheduler_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 """Test for buffered_scheduler."""
 
+import cloudpickle
 import concurrent.futures
 import threading
 import time
@@ -35,7 +36,7 @@ def wait_seconds(self, n: int):
         return n + 1
 
     with local_worker_manager.LocalWorkerPoolManager(
-        WaitingWorker, count=2) as pool:
+        WaitingWorker, count=2, pickle_func=cloudpickle.dumps) as pool:
       _, futures = buffered_scheduler.schedule_on_worker_pool(
           lambda w, v: w.wait_seconds(v), range(4), pool)
       not_done = futures
@@ -54,7 +55,7 @@ def square(self, the_value, extra_factor=1):
         return the_value * the_value * extra_factor
 
     with local_worker_manager.LocalWorkerPoolManager(
-        TheWorker, count=2) as pool:
+        TheWorker, count=2, pickle_func=cloudpickle.dumps) as pool:
       workers, futures = buffered_scheduler.schedule_on_worker_pool(
           lambda w, v: w.square(v), range(10), pool)
       self.assertLen(workers, 2)

compiler_opt/distributed/local/local_worker_manager.py

@@ -28,10 +28,10 @@
 """
 import concurrent.futures
 import sys
-import cloudpickle
 import dataclasses
 import functools
 import multiprocessing
+import pickle
 import psutil
 import threading
 
@@ -80,7 +80,7 @@ def _run_impl(pipe: connection.Connection, worker_class: SerializedClass,
   pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
   # unpickle the type, which triggers all the necessary imports. This registers
   # abseil flags, which can then be parsed.
-  cls = cloudpickle.loads(worker_class)
+  cls = pickle.loads(worker_class)
   if parse_argv:
     flags.FLAGS(sys.argv, known_only=True)
   obj = cls(*args, **kwargs)
@@ -126,7 +126,8 @@ def _run(pipe: connection.Connection, worker_class: SerializedClass,
     raise e
 
 
-def _make_stub(cls: 'type[worker.Worker]', parse_argv: bool, *args, **kwargs):
+def _make_stub(cls: 'type[worker.Worker]', parse_argv: bool, pickle_func, *args,
+               **kwargs):
 
   class _Stub:
     """Client stub to a worker hosted by a process."""
@@ -141,7 +142,7 @@ def __init__(self):
       # to handle high priority requests. The expectation is that the user
       # achieves concurrency through multiprocessing, not multithreading.
       self._process = _get_context().Process(
-          target=functools.partial(_run, child_pipe, cloudpickle.dumps(cls),
+          target=functools.partial(_run, child_pipe, pickle_func(cls),
                                    parse_argv, *args, **kwargs))
       # lock for the msgid -> reply future map. The map will be set to None
       # when we stop.
@@ -254,14 +255,14 @@ def __dir__(self):
 
 
 def _create_local_worker_pool(worker_cls: 'type[worker.Worker]',
-                              count: int | None, parse_argv: bool, *args,
-                              **kwargs) -> worker.FixedWorkerPool:
+                              count: int | None, parse_argv: bool, pickle_func,
+                              *args, **kwargs) -> worker.FixedWorkerPool:
   """Create a local worker pool for worker_cls."""
   if not count:
     count = _get_context().cpu_count()
   final_kwargs = worker.get_full_worker_args(worker_cls, **kwargs)
   stubs = [
-      _make_stub(worker_cls, parse_argv, *args, **final_kwargs)
+      _make_stub(worker_cls, parse_argv, pickle_func, *args, **final_kwargs)
       for _ in range(count)
   ]
   return worker.FixedWorkerPool(workers=stubs, worker_concurrency=16)
@@ -283,13 +284,15 @@ class LocalWorkerPoolManager(AbstractContextManager):
 
   def __init__(self,
                worker_class: 'type[worker.Worker]',
+               pickle_func=pickle.dumps,
                *,
                count: int | None,
                worker_args: tuple = (),
                worker_kwargs: dict | None = None):
     worker_kwargs = {} if worker_kwargs is None else worker_kwargs
     self._pool = _create_local_worker_pool(worker_class, count, True,
-                                           *worker_args, **worker_kwargs)
+                                           pickle_func, *worker_args,
+                                           **worker_kwargs)
 
   def __enter__(self) -> worker.FixedWorkerPool:
     return self._pool

compiler_opt/distributed/local/local_worker_manager_test.py

@@ -24,7 +24,7 @@
 from compiler_opt.distributed.local import local_worker_manager
 from tf_agents.system import system_multiprocessing as multiprocessing
 
-flags.DEFINE_integer(
+_TEST_FLAG = flags.DEFINE_integer(
     'test_only_flag', default=1, help='A flag used by some tests.')
 
 
@@ -74,7 +74,7 @@ def method(self):
 class JobGetFlags(Worker):
 
   def method(self):
-    return {'argv': sys.argv, 'the_flag': flags.FLAGS.test_only_flag}
+    return {'argv': sys.argv, 'the_flag': _TEST_FLAG.value}
 
 
 class LocalWorkerManagerTest(absltest.TestCase):

compiler_opt/es/blackbox_learner_test.py

@@ -15,6 +15,7 @@
 
 import os
 from absl.testing import absltest
+import cloudpickle
 import gin
 import tempfile
 import numpy as np
@@ -148,6 +149,7 @@ def test_run_step(self):
     with local_worker_manager.LocalWorkerPoolManager(
         blackbox_test_utils.ESWorker,
         count=3,
+        pickle_func=cloudpickle.dumps,
         worker_args=('',),
         worker_kwargs=dict(kwarg='')) as pool:
       self._learner.run_step(pool)  # pylint: disable=protected-access