From 2cce7bd2ff486cda768bdaef0d22f7a605a0f6fd Mon Sep 17 00:00:00 2001
From: John <john@slai.io>
Date: Thu, 19 Dec 2024 12:53:26 -0800
Subject: [PATCH] Fix: add no wait in endpoint to skip graceful timeout

---
 sdk/src/beta9/runner/endpoint.py | 3 +++
 sdk/src/beta9/runner/serve.py    | 5 ++++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/sdk/src/beta9/runner/endpoint.py b/sdk/src/beta9/runner/endpoint.py
index 6a70a312a..f9fa4e539 100644
--- 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
index 587c3c1bd..deff93b96 100644
--- 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,
             )