commit 5ee6de057595d71d9e1ef33bff8e5c5bb035a432
author Ran Chen <crccw@google.com> Fri Sep 25 17:58:59 2020 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Fri Sep 25 18:03:12 2020 -0700
tree de87d8320efd71350948b211cf2073d8a57eed95
parent ba6d71b874c954bdd4219da118318fe2710515cf

Support global and nested MultiProcessRunnerPool

We used to launch MultiProcessRunnerPool as daemon processes to blocking the
program from exiting. We also requires whoever creates global
MultiProcessRunnerPool to register atexit callbacks to shut them down to avoid
thread sanitizer complaints.

This change avoids such pitfalls and allows nested MultiProcessRunnerPool. One
use case of nested pools is that we need to spawn a new process with
accelerators hidden to test saving/loading logics. We already use a pool for
multi worker testing, we need a nested pool of CPU only processes to test the
loading part.

As you can see this is tricky for two reasons:
1) processes spawned by multiprocessing don't execute atexit callbacks [1]. Even
if they do, the multiprocessing exit hook that waits for all processes executes
before atexit callbacks. [2]
2) multiprocessing doesn't register the hook as an atexit callback when being
imported. We need to make sure our callback is registered after the exit hook of
multiprocessing library.

[1] https://bugs.python.org/issue39675
[2] https://github.com/python/cpython/blob/069b8d20be8018fbd49ed5aaf64c4caba311e48f/Lib/multiprocessing/process.py#L261

PiperOrigin-RevId: 333843641
Change-Id: I8305ef8b8e76aa0e9a9623c70ea1c026fb501f15