experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add an option in multi_worker_test_base.create_in_process_cluster to choose different rpc_layer.

Browse Source 
PiperOrigin-RevId: 254239865
This commit is contained in:
Yuefeng Zhou 2019-06-20 11:41:03 -07:00 committed by TensorFlower Gardener
parent c0b21f701c
commit 470b7c46fe
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/python/distribute/multi_worker_test_base.py
View File

@ -140,7 +140,8 @@ def _create_cluster(num_workers,
def create_in_process_cluster(num_workers, def create_in_process_cluster(num_workers,
num_ps, num_ps,
has_chief=False, has_chief=False,
has_eval=False): has_eval=False,
rpc_layer='grpc'):
"""Create an in-process cluster that consists of only standard server.""" """Create an in-process cluster that consists of only standard server."""
# Leave some memory for cuda runtime. # Leave some memory for cuda runtime.
gpu_mem_frac = 0.7 / (num_workers + int(has_chief) + int(has_eval)) gpu_mem_frac = 0.7 / (num_workers + int(has_chief) + int(has_eval))
@ -180,7 +181,7 @@ def create_in_process_cluster(num_workers,
worker_config=worker_config, worker_config=worker_config,
ps_config=ps_config, ps_config=ps_config,
eval_config=eval_config, eval_config=eval_config,
protocol='grpc') protocol=rpc_layer)
def create_cluster_spec(test_obj, def create_cluster_spec(test_obj,