From c8df8f147158d012245de8c0fc062401507df4ad Mon Sep 17 00:00:00 2001 From: Rick Chao Date: Fri, 28 Aug 2020 18:46:11 -0700 Subject: [PATCH] PSv2: Replace an assert with a more informative ValueError. PiperOrigin-RevId: 329046260 Change-Id: Ifdad03e35b12dfbbd70feedcd9f96f9e20b202d2 --- tensorflow/python/distribute/client/utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tensorflow/python/distribute/client/utils.py b/tensorflow/python/distribute/client/utils.py index 6c595579863..51d82630d6e 100644 --- a/tensorflow/python/distribute/client/utils.py +++ b/tensorflow/python/distribute/client/utils.py @@ -28,8 +28,11 @@ def start_server(cluster_resolver, protocol): """Start a server and block the process from exiting.""" # This function is for multi-processing test or users who would like to have # every job run the same binary for simplicity. - assert (cluster_resolver.task_type == 'worker' or - cluster_resolver.task_type == 'ps') + if not (cluster_resolver.task_type == 'worker' or + cluster_resolver.task_type == 'ps'): + raise ValueError('Unexpected task_type to start a server: {}'.format( + cluster_resolver.task_type)) + server = server_lib.Server( cluster_resolver.cluster_spec().as_cluster_def(), job_name=cluster_resolver.task_type,