Fix: add no wait in endpoint to skip graceful timeout

beam-cloud · Dec 19, 2024 · 2cce7bd · 2cce7bd
1 parent ef1d97f
commit 2cce7bd
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/sdk/src/beta9/runner/endpoint.py b/sdk/src/beta9/runner/endpoint.py
@@ -277,4 +277,7 @@ def shutdown(self, signum=None, frame=None):
         "timeout": cfg.timeout,
     }
 
+    if os.environ.get("no_wait") == "true":
+        options["graceful_timeout"] = 0
+
     GunicornApplication(Starlette(), options).run()
diff --git a/sdk/src/beta9/runner/serve.py b/sdk/src/beta9/runner/serve.py
@@ -86,7 +86,10 @@ def run(self, *, command: List[str]) -> None:
                 " ".join(command),
                 shell=True,
                 preexec_fn=os.setsid,
-                env=os.environ,
+                env={
+                    **os.environ,
+                    "no_wait": "true",
+                },
                 stdout=sys.stdout,
                 stderr=sys.stdout,
             )