Added device_filters to session_config
device_filters should be set to avoid communication between master and all the workers. Rather the communication must be made between master <-> ps and worker <-> ps.
import argparse
+import json
import os
import trainer.model as model
@@ -8,6 +9,25 @@ from tensorflow.contrib.learn.python.learn.utils import (
saved_model_export_utils)
from tensorflow.contrib.training.python.training import hparam
+def _get_session_config_from_env_var():
+ """Returns a tf.ConfigProto instance that has appropriate device_filters set.
+ """
+
+ tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
+
+ if (tf_config and 'task' in tf_config and 'type' in tf_config['task'] and
+ 'index' in tf_config['task']):
+ # Master should only communicate with itself and ps
+ if tf_config['task']['type'] == 'master':
+ return tf.ConfigProto(device_filters=['/job:ps', '/job:master'])
+ # Worker should only communicate with itself and ps
+ elif tf_config['task']['type'] == 'worker':
+ return tf.ConfigProto(device_filters=[
+ '/job:ps',
+ '/job:worker/task:%d' % tf_config['task']['index']
+ ])
+ return None
+
def run_experiment(hparams):
"""Run the training and evaluate using the high level API"""
@@ -37,7 +57,7 @@ def run_experiment(hparams):
name='census-eval'
)
- run_config = tf.estimator.RunConfig()
+ run_config = tf.estimator.RunConfig(session_config=_get_session_config_from_env_var())
run_config = run_config.replace(model_dir=hparams.job_dir)
print('model dir {}'.format(run_config.model_dir))
estimator = model.build_estimator(
