From 290463b3e2227e23fabee7711b4c342b18a6d7aa Mon Sep 17 00:00:00 2001
From: akirchhoff-modular <akirchhoff@modular.com>
Date: Wed, 5 Feb 2025 16:46:16 -0800
Subject: [PATCH] [Pipelines] Work around BuildBuddy CLI signal bug (#55279)

Due to a BuildBuddy CLI bug
(https://github.com/buildbuddy-io/buildbuddy/issues/8326), when using
`./bazelw run`, we get invoked with the `SIGINT` handler being
`SIG_IGN`. This is annoying because it becomes much harder to terminate
the process when `SIGINT` is disabled. While less-commonly used,
BuildBuddy also disables `SIGTERM` and `SIGQUIT`. Detect this situation
and return the state to what it would have been if we had been invoked
without the BuildBuddy CLI.

Part of AITLIB-132.

MODULAR_ORIG_COMMIT_REV_ID: bf1a3105755564fa03eaad2cf7c6e2636418e7a0
---
 pipelines/python/max/entrypoints/pipelines.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/pipelines/python/max/entrypoints/pipelines.py b/pipelines/python/max/entrypoints/pipelines.py
index b8b2543c..09956422 100644
--- a/pipelines/python/max/entrypoints/pipelines.py
+++ b/pipelines/python/max/entrypoints/pipelines.py
@@ -7,6 +7,7 @@
 import functools
 import logging
 import os
+import signal
 
 import click
 from max.pipelines import PIPELINE_REGISTRY, PipelineConfig
@@ -213,4 +214,18 @@ def cli_list():
     if directory := os.getenv("BUILD_WORKSPACE_DIRECTORY"):
         os.chdir(directory)
 
+    # Workaround for https://github.com/buildbuddy-io/buildbuddy/issues/8326
+    if (
+        signal.getsignal(signal.SIGINT) == signal.SIG_IGN
+        and signal.getsignal(signal.SIGTERM) == signal.SIG_IGN
+        and signal.getsignal(signal.SIGQUIT) == signal.SIG_IGN
+    ):
+        # For SIGINT, Python remaps SIG_DFL to default_int_handler on startup.
+        # We do the same here to retain the same behavior we would get if we
+        # started normally.  (SIG_DFL terminates the process immediately;
+        # default_int_handler raises KeyboardInterrupt.)
+        signal.signal(signal.SIGINT, signal.default_int_handler)
+        signal.signal(signal.SIGTERM, signal.SIG_DFL)
+        signal.signal(signal.SIGQUIT, signal.SIG_DFL)
+
     main()