feat: portal sigterm handler

zetta_utils/mazepa/progress_tracker.py

@@ -13,7 +13,6 @@
 from rich.console import Console
 
 from zetta_utils.common import custom_signal_handler_ctx, get_user_confirmation
-from zetta_utils.common.ctx_managers import noop_ctx_mngr
 
 from .execution_state import ProgressReport
 
@@ -42,6 +41,10 @@ def handler(_, __):
     return handler
 
 
+def graceful_shutdown(_, __):  # pragma: no cover
+    sys.exit(0)
+
+
 class ProgressUpdateFN(Protocol):
     def __call__(self, progress_reports: dict[str, ProgressReport]) -> None:
         ...
@@ -98,7 +101,7 @@ def custom_debugger_hook():
                 get_confirm_sigint_fn(progress_bar), signal.SIGINT
             )
         else:
-            handler_ctx = noop_ctx_mngr()
+            handler_ctx = custom_signal_handler_ctx(graceful_shutdown, signal.SIGTERM)
         with handler_ctx:
             submission_tracker_ids = {
                 k: progress_bar.add_task(

zetta_utils/mazepa_addons/configurations/execute_locally.py

@@ -37,6 +37,7 @@ def execute_locally(
     semaphores_spec: dict[SemaphoreType, int] | None = None,
     debug: bool = False,
     write_progress_summary: bool = False,
+    require_interrupt_confirm: bool = True,
 ):
 
     queues_dir_ = queues_dir if queues_dir else ""
@@ -75,4 +76,5 @@ def execute_locally(
             checkpoint_interval_sec=checkpoint_interval_sec,
             raise_on_failed_checkpoint=raise_on_failed_checkpoint,
             write_progress_summary=write_progress_summary,
+            require_interrupt_confirm=require_interrupt_confirm,
         )

zetta_utils/mazepa_addons/configurations/execute_on_gcp_with_sqs.py

@@ -225,6 +225,7 @@ def execute_on_gcp_with_sqs(  # pylint: disable=too-many-locals
     checkpoint_interval_sec: float = 300.0,
     raise_on_failed_checkpoint: bool = True,
     write_progress_summary: bool = False,
+    require_interrupt_confirm: bool = True,
 ):
     if debug and not local_test:
         raise ValueError("`debug` can only be set to `True` when `local_test` is also `True`.")
@@ -245,6 +246,7 @@ def execute_on_gcp_with_sqs(  # pylint: disable=too-many-locals
             semaphores_spec=semaphores_spec,
             debug=debug,
             write_progress_summary=write_progress_summary,
+            require_interrupt_confirm=require_interrupt_confirm,
         )
     else:
         assert gcloud.check_image_exists(worker_image), worker_image
@@ -306,5 +308,5 @@ def execute_on_gcp_with_sqs(  # pylint: disable=too-many-locals
                 checkpoint_interval_sec=checkpoint_interval_sec,
                 raise_on_failed_checkpoint=raise_on_failed_checkpoint,
                 write_progress_summary=write_progress_summary,
-                require_interrupt_confirm=False,
+                require_interrupt_confirm=require_interrupt_confirm,
             )