fix: Eliminate the exception during shutdown (#3826)

bentoml · May 5, 2023 · 5248eee · 5248eee
1 parent f3ef39a
commit 5248eee
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/src/bentoml/_internal/marshal/dispatcher.py b/src/bentoml/_internal/marshal/dispatcher.py
@@ -235,7 +235,7 @@ async def train_optimizer(
                     info.enqueue_time = now
                 self._loop.create_task(self.outbound_call(inputs_info))
         except asyncio.CancelledError:
-            return
+            raise
         except Exception as e:  # pylint: disable=broad-except
             logger.error(traceback.format_exc(), exc_info=e)
 
@@ -313,7 +313,7 @@ async def controller(self):
                 inputs_info = tuple(self._queue.pop() for _ in range(n_call_out))
                 self._loop.create_task(self.outbound_call(inputs_info))
             except asyncio.CancelledError:
-                return
+                raise
             except Exception as e:  # pylint: disable=broad-except
                 logger.error(traceback.format_exc(), exc_info=e)
 

diff --git a/tests/e2e/bento_server_http/tests/test_serve.py b/tests/e2e/bento_server_http/tests/test_serve.py
@@ -20,7 +20,7 @@ def test_http_server(bentoml_home: str):
 
     server.stop()
 
-    timeout = 60
+    timeout = 10
     start_time = time.time()
     while time.time() - start_time < timeout:
         retcode = server.process.poll()
@@ -48,7 +48,7 @@ def test_http_server_ctx(bentoml_home: str):
 
         assert res == {"more_test": "and more json"}
 
-    timeout = 60
+    timeout = 10
     start_time = time.time()
     while time.time() - start_time < timeout:
         retcode = server.process.poll()