message: Create a multi-process unit test for MultiWorkerMirroredStrategy + Estimator
diff: @@ -23,6 +23,8 @@ import contextlib
  import copy
  import json
  import os
  +import subprocess
  +import sys
  import threading
  import numpy as np
  @@ -45,6 +47,7 @@ from tensorflow.python.platform import test
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.training import coordinator
  from tensorflow.python.training import server_lib
  +from tensorflow.python.util import nest
  original_run_std_server = dc._run_std_server # pylint: disable=protected-access
  @@ -446,6 +449,62 @@ class IndependentWorkerTestBase(test.TestCase):
  self._coord.join(worker_threads)
  +class MultiWorkerMultiProcessTest(test.TestCase):
  + """Testing infra for independent workers using multiple processes."""
  +
  + def _run_task_in_process(self, cmd_args, cluster_spec, task_type, task_id):
  + env = os.environ.copy()
  + env['TF_CONFIG'] = json.dumps({
  + 'cluster': cluster_spec,
  + 'task': {
  + 'type': task_type,
  + 'index': task_id
  + }
  + })
  + return subprocess.Popen(
  + cmd_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
  +
  + def run_multiple_tasks_in_processes(self, cmd_args, cluster_spec):
  + """Run `cmd_args` in a process for each task in `cluster_spec`."""
  + processes = {}
  + for task_type in cluster_spec.keys():
  + processes[task_type] = []
  + for task_id in range(len(cluster_spec[task_type])):
  + p = self._run_task_in_process(cmd_args, cluster_spec, task_type,
  + task_id)
  + processes[task_type].append(p)
  + return processes
  +
  + def join_independent_workers(self, worker_processes):
  + return_codes = []
  + for p in nest.flatten(worker_processes):
  + try:
  + # Calling p.wait() will hang if we don't consume its output.
  + p.communicate()
  + except ValueError:
  + # The output of the process may have been consumed, in which case
  + # calling `p.communicate()` will raise a ValueError.
  + pass
  + finally:
  + return_codes.append(p.returncode)
  + for return_code in return_codes:
  + self.assertEqual(return_code, 0)
  +
  + def stream_stderr(self, process):
  + # TODO(yuefengz): calling stream_stderr on a single process will probably
  + # make all processes hang if they have too much output e.g. adding
  + # --vmodule=execute=2 to cmd_args. But this method is useful for debugging
  + # purposes. We should figure out the hanging problem, probably by consuming
  + # outputs of all processes at the same time.
  + while True:
  + output = process.stderr.readline()
  + if not output and process.poll() is not None:
  + break
  + if output:
  + print(output.strip())
  + sys.stdout.flush()
  +
  +
  def get_tf_config_task():
  return json.loads(os.environ['TF_CONFIG'])['task']
  