feat(package): Add support for gracefully shutting down the compression scheduler and compression workers (resolves #1037). (#1169)

Author: AVMatthews

components/clp-package-utils/clp_package_utils/scripts/stop_clp.py

@@ -34,10 +34,12 @@
 logger = logging.getLogger(__file__)
 
 
-def stop_running_container(container_name: str, already_exited_containers: List[str], force: bool):
+def stop_running_container(
+    container_name: str, already_exited_containers: List[str], force: bool, timeout: int = 10
+):
     if is_container_running(container_name):
         logger.info(f"Stopping {container_name}...")
-        cmd = ["docker", "stop", container_name]
+        cmd = ["docker", "stop", "--timeout", str(timeout), container_name]
         subprocess.run(cmd, stdout=subprocess.DEVNULL, check=True)
 
         logger.info(f"Removing {container_name}...")
@@ -148,9 +150,6 @@ def main(argv):
         if target in (ALL_TARGET_NAME, QUERY_WORKER_COMPONENT_NAME):
             container_name = f"clp-{QUERY_WORKER_COMPONENT_NAME}-{instance_id}"
             stop_running_container(container_name, already_exited_containers, force)
-        if target in (ALL_TARGET_NAME, COMPRESSION_WORKER_COMPONENT_NAME):
-            container_name = f"clp-{COMPRESSION_WORKER_COMPONENT_NAME}-{instance_id}"
-            stop_running_container(container_name, already_exited_containers, force)
         if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, QUERY_SCHEDULER_COMPONENT_NAME):
             container_name = f"clp-{QUERY_SCHEDULER_COMPONENT_NAME}-{instance_id}"
             stop_running_container(container_name, already_exited_containers, force)
@@ -164,11 +163,14 @@ def main(argv):
             COMPRESSION_SCHEDULER_COMPONENT_NAME,
         ):
             container_name = f"clp-{COMPRESSION_SCHEDULER_COMPONENT_NAME}-{instance_id}"
-            stop_running_container(container_name, already_exited_containers, force)
+            stop_running_container(container_name, already_exited_containers, force, timeout=300)
 
             container_config_file_path = logs_dir / f"{container_name}.yml"
             if container_config_file_path.exists():
                 container_config_file_path.unlink()
+        if target in (ALL_TARGET_NAME, COMPRESSION_WORKER_COMPONENT_NAME):
+            container_name = f"clp-{COMPRESSION_WORKER_COMPONENT_NAME}-{instance_id}"
+            stop_running_container(container_name, already_exited_containers, force, timeout=60)
         if target in (ALL_TARGET_NAME, CONTROLLER_TARGET_NAME, REDIS_COMPONENT_NAME):
             container_name = f"clp-{REDIS_COMPONENT_NAME}-{instance_id}"
             stop_running_container(container_name, already_exited_containers, force)

components/job-orchestration/job_orchestration/executor/compress/compression_task.py

@@ -6,6 +6,7 @@
 from contextlib import closing
 from typing import Any, Dict, List, Optional, Tuple
 
+from celery import signals
 from celery.app.task import Task
 from celery.utils.log import get_task_logger
 from clp_py_utils.clp_config import (
@@ -523,6 +524,11 @@ def run_clp(
         return CompressionTaskStatus.FAILED, worker_output
 
 
+@signals.worker_shutdown.connect
+def worker_shutdown_handler(signal=None, sender=None, **kwargs):
+    logger.info("Shutdown signal received.")
+
+
 @app.task(bind=True)
 def compress(
     self: Task,

components/job-orchestration/job_orchestration/scheduler/compress/compression_scheduler.py

@@ -2,6 +2,7 @@
 import datetime
 import logging
 import os
+import signal
 import sys
 import time
 from contextlib import closing
@@ -55,6 +56,14 @@
 
 scheduled_jobs = {}
 
+received_sigterm = False
+
+
+def sigterm_handler(signal_number, frame):
+    global received_sigterm
+    received_sigterm = True
+    logger.info("Received SIGTERM.")
+
 
 def fetch_new_jobs(db_cursor):
     db_cursor.execute(
@@ -451,6 +460,10 @@ def poll_running_jobs(db_conn, db_cursor):
     for job_id in jobs_to_delete:
         del scheduled_jobs[job_id]
 
+    if received_sigterm and 0 == len(scheduled_jobs):
+        logger.info("Recieved SIGTERM and there're no more running jobs. Exiting.")
+        sys.exit(0)
+
 
 def main(argv):
     args_parser = argparse.ArgumentParser()
@@ -466,6 +479,9 @@ def main(argv):
     # Update logging level based on config
     set_logging_level(logger, os.getenv("CLP_LOGGING_LEVEL"))
 
+    # Register the SIGTERM handler
+    signal.signal(signal.SIGTERM, sigterm_handler)
+
     # Load configuration
     config_path = Path(args.config)
     try:
@@ -500,16 +516,17 @@ def main(argv):
         # Start Job Processing Loop
         while True:
             try:
-                search_and_schedule_new_tasks(
-                    clp_config,
-                    db_conn,
-                    db_cursor,
-                    clp_metadata_db_connection_config,
-                )
+                if not received_sigterm:
+                    search_and_schedule_new_tasks(
+                        clp_config,
+                        db_conn,
+                        db_cursor,
+                        clp_metadata_db_connection_config,
+                    )
                 poll_running_jobs(db_conn, db_cursor)
                 time.sleep(clp_config.compression_scheduler.jobs_poll_delay)
             except KeyboardInterrupt:
-                logger.info("Gracefully shutting down")
+                logger.info("Forcefully shutting down")
                 return -1
             except Exception:
                 logger.exception(f"Error in scheduling.")