Simplify error handling (#585)

beam-cloud · Oct 3, 2024 · fb3b519 · fb3b519
1 parent 4ef8eda
commit fb3b519
Showing 1 changed file with 5 additions and 10 deletions.
diff --git a/sdk/src/beta9/runner/function.py b/sdk/src/beta9/runner/function.py
@@ -158,7 +158,7 @@ def main(channel: Channel):
 
         # Start monitoring the task
         thread_pool = ThreadPoolExecutor()
-        monitor_future = thread_pool.submit(
+        thread_pool.submit(
             _monitor_task,
             context=context,
             stub_id=config.stub_id,
@@ -169,16 +169,11 @@ def main(channel: Channel):
         )
 
         # Invoke the function and handle its result
-        result = InvokeResult()
-        try:
-            result = invoke_function(function_stub, context, task_id)
-            if result.exception:
-                raise result.exception
-        except BaseException:
+        result = invoke_function(function_stub, context, task_id)
+        if result.exception:
             handle_task_failure(result, gateway_stub, task_id, container_id, container_hostname)
-            raise
-        finally:
-            monitor_future.cancel()
+            thread_pool.shutdown(wait=False)
+            raise result.exception
 
         # End the task and send callback
         complete_task(