Use env variable for graceful shutdown time period

Signed-off-by: Sreekanth <prsreekanth920@gmail.com>

Author: BulkBeing

packages/pynumaflow/pynumaflow/_constants.py

@@ -61,6 +61,11 @@
 # If NUM_THREADS_DEFAULT env is not set default to 4
 NUM_THREADS_DEFAULT = int(os.getenv("MAX_THREADS", "4"))
 
+# Grace period in seconds for the GRPC server to shutdown
+NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS = int(
+    os.getenv("NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS", "60")
+)
+
 _LOGGER = setup_logging(__name__)
 if os.getenv("PYTHONDEBUG"):
     _LOGGER.setLevel(logging.DEBUG)

packages/pynumaflow/pynumaflow/sinker/async_server.py

@@ -26,6 +26,7 @@
     ON_SUCCESS_SINK_SOCK_PATH,
     ON_SUCCESS_SINK_SERVER_INFO_FILE_PATH,
     MAX_NUM_THREADS,
+    NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS,
 )
 
 from pynumaflow.shared.server import NumaflowServer
@@ -168,9 +169,9 @@ async def _watch_for_shutdown():
             """Wait for the shutdown event and stop the server with a grace period."""
             await shutdown_event.wait()
             _LOGGER.info("Shutdown signal received, stopping server gracefully...")
-            # Stop accepting new requests and wait for a maximum of 5 seconds
-            # for in-flight requests to complete
-            await server.stop(5)
+            # Stop accepting new requests and wait for a maximum of
+            # NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS seconds for in-flight requests to complete
+            await server.stop(NUMAFLOW_GRPC_SHUTDOWN_GRACE_PERIOD_SECONDS)
 
         shutdown_task = asyncio.create_task(_watch_for_shutdown())
         await server.wait_for_termination()