PSv2: Replace an assert with a more informative ValueError.

PiperOrigin-RevId: 329046260
Change-Id: Ifdad03e35b12dfbbd70feedcd9f96f9e20b202d2
This commit is contained in:
Rick Chao 2020-08-28 18:46:11 -07:00 committed by TensorFlower Gardener
parent d6850c4a4f
commit c8df8f1471

View File

@ -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,