From 019467500eee9542833aedd6e2c17bd62ffa4b40 Mon Sep 17 00:00:00 2001 From: Domenic Barbuzzi Date: Fri, 31 May 2024 10:44:40 -0400 Subject: [PATCH] Handle server startup failure in __enter__ (#274) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With a context manager class, the `__exit__` method is not called when an exception is raised during the context manager’s `__enter__` method. This PR addresses that by manually calling that method if an exception is raised. --- tests/utils/server.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/utils/server.py b/tests/utils/server.py index 4f1a2accaaa1b..4c0111b21cf9e 100644 --- a/tests/utils/server.py +++ b/tests/utils/server.py @@ -97,10 +97,15 @@ def __enter__(self): ray.init(ignore_reinit_error=True) log_banner(self._logger, "server startup command args", shlex.join(self._args)) - self.server_runner = ServerRunner.remote(self._args, - logger=self._logger) - ray.get(self.server_runner.ready.remote()) - return self.server_runner + + try: + self.server_runner = ServerRunner.remote(self._args, + logger=self._logger) + ray.get(self.server_runner.ready.remote()) + return self.server_runner + except Exception as e: + self.__exit__(*sys.exc_info()) + raise e def __exit__(self, exc_type, exc_value, exc_traceback): """