diff options
| author |  Wei Ho <weiho@google.com> | 2016-11-16 13:03:00 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2016-11-16 13:22:51 -0800 |
| commit | 92ab01f8708b72baf3148abc7252f89c71792bb5 (patch) | |
| tree | 6614ee3e060248e7c5676bcbbdd85eae1301fcf9 /tensorflow/tools/dist_test | |
| parent | 8c46bff60dc32d7708653f023556b6d30bd85277 (diff) | |
Split ClusterConfig out of RunConfig
Change: 139368020
Diffstat (limited to 'tensorflow/tools/dist_test')
| -rw-r--r-- | tensorflow/tools/dist_test/python/census_widendeep.py | 14 |
1 files changed, 11 insertions, 3 deletions
diff --git a/tensorflow/tools/dist_test/python/census_widendeep.py b/tensorflow/tools/dist_test/python/census_widendeep.py index b9e0aa1433..f5510c374a 100644 --- a/tensorflow/tools/dist_test/python/census_widendeep.py +++ b/tensorflow/tools/dist_test/python/census_widendeep.py @@ -20,6 +20,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import json import os from six.moves import urllib @@ -240,9 +241,16 @@ def _create_experiment_fn(output_dir): # pylint: disable=unused-argument categorical_columns, continuous_columns) - config = run_config.RunConfig(master=FLAGS.master_grpc_url, - num_ps_replicas=FLAGS.num_parameter_servers, - task=FLAGS.worker_index) + os.environ["TF_CONFIG"] = json.dumps({ + "cluster": { + tf.contrib.learn.TaskType.PS: ["fake_ps"] * + FLAGS.num_parameter_servers + }, + "task": { + "index": FLAGS.worker_index + } + }) + config = run_config.RunConfig(master=FLAGS.master_grpc_url) estimator = tf.contrib.learn.DNNLinearCombinedClassifier( model_dir=FLAGS.model_dir, |