提交 20259f6f 编辑于 作者: Keith Rush's avatar Keith Rush 提交者: tensorflow-copybara
浏览文件

Moves server construction outside of try-block; if construction fails, cleanup should not happen.

PiperOrigin-RevId: 344555142
上级 8b03066c
......@@ -63,14 +63,14 @@ def server_context(ex_factory: executor_factory.ExecutorFactory,
raise ValueError('The number of threads must be a positive integer.')
if port < 1:
raise ValueError('The server port must be a positive integer.')
service = executor_service.ExecutorService(ex_factory)
server_kwargs = {}
if options is not None:
server_kwargs['options'] = options
thread_pool_executor = concurrent.futures.ThreadPoolExecutor(
max_workers=num_threads)
server = grpc.server(thread_pool_executor, **server_kwargs)
try:
service = executor_service.ExecutorService(ex_factory)
server_kwargs = {}
if options is not None:
server_kwargs['options'] = options
thread_pool_executor = concurrent.futures.ThreadPoolExecutor(
max_workers=num_threads)
server = grpc.server(thread_pool_executor, **server_kwargs)
full_port_string = '[::]:{}'.format(port)
if credentials is not None:
server.add_secure_port(full_port_string, credentials)
......
......@@ -57,6 +57,19 @@ class ServerUtilsTest(test_case.TestCase):
mock_logging_info.assert_called_once_with('Shutting down server.')
@mock.patch(
'tensorflow_federated.python.core.impl.executors.executor_service.ExecutorService',
side_effect=ValueError)
def test_failure_on_construction_fails_as_expected(self, mock_service):
ex = eager_tf_executor.EagerTFExecutor()
ex_factory = executor_stacks.ResourceManagingExecutorFactory(lambda _: ex)
with self.assertRaises(ValueError):
with server_utils.server_context(ex_factory, 1,
portpicker.pick_unused_port()):
time.sleep(1)
if __name__ == '__main__':
test_case.main()
Supports Markdown
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册