def __init__(self, config, start=True):
"""Creates a new worker server.
Args:
config: A `tf.data.experimental.service.WorkerConfig` configration.
start: (Optional.) Boolean, indicating whether to start the server after
creating it. Defaults to True.
"""
if config.dispatcher_address is None:
raise ValueError(
"Must specify a `dispatcher_address` in the `config` passed "
"to `WorkerServer`.")
if isinstance(config, service_config_pb2.WorkerConfig):
config_proto = config
else:
config_proto = service_config_pb2.WorkerConfig(
dispatcher_address=config.dispatcher_address,
worker_address=config.worker_address,
port=config.port,
protocol=config.protocol,
heartbeat_interval_ms=config.heartbeat_interval_ms,
dispatcher_timeout_ms=config.dispatcher_timeout_ms,
data_transfer_protocol=None)
self._server = _pywrap_server_lib.TF_DATA_NewWorkerServer(
config_proto.SerializeToString())
if start:
self._server.start()
def _make_worker(dispatcher_address, shutdown_quiet_period_ms=0, port=0):
"""Creates a worker server."""
defaults = server_lib.WorkerConfig(dispatcher_address=dispatcher_address)
config_proto = service_config_pb2.WorkerConfig(
dispatcher_address=dispatcher_address,
worker_address=defaults.worker_address,
port=port,
protocol=PROTOCOL,
heartbeat_interval_ms=TEST_HEARTBEAT_INTERVAL_MS,
dispatcher_timeout_ms=TEST_DISPATCHER_TIMEOUT_MS,
data_transfer_protocol=None,
shutdown_quiet_period_ms=shutdown_quiet_period_ms)
return server_lib.WorkerServer(config_proto, start=False)
