def create_server_and_stub(
self, grads_to_wait, lr_staleness_modulation, use_async, **kwargs
):
args = PserverArgs(
grads_to_wait=grads_to_wait,
lr_staleness_modulation=lr_staleness_modulation,
use_async=use_async,
port=self._port,
model_zoo=_test_model_zoo_path,
model_def="test_module.custom_model",
**kwargs
)
pserver = ParameterServer(args)
pserver.prepare()
self._parameters = pserver.parameters
self._server = pserver.server
self._stub = elasticdl_pb2_grpc.PserverStub(self._channel)
self._lr = 0.1
def create_pserver(model_zoo_path, model_def, grads_to_wait, use_async,
num_ps_pods):
ports = [i + 12345 for i in range(num_ps_pods)]
channels = []
for port in ports:
addr = "localhost:%d" % port
channel = build_channel(addr)
channels.append(channel)
pservers = []
for port in ports:
args = PserverArgs(
grads_to_wait=grads_to_wait,
use_async=True,
port=port,
model_zoo=model_zoo_path,
model_def=model_def,
)
pserver = ParameterServer(args)
pserver.prepare()
pservers.append(pserver)
return ports, channels, pservers
def main():
args = parse_ps_args()
pserver = ParameterServer(args)
pserver.prepare()
pserver.run()