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()