fix(job-orchestration)!: Add defensive timeouts to prevent scheduler hang (fixes y-scope#2252). (y-scope#2271)

Co-authored-by: hoophalab <200652805+hoophalab@users.noreply.github.com>
Co-authored-by: Junhao Liao <junhao.liao@yscope.com>

Author: 3 people

components/clp-py-utils/clp_py_utils/clp_config.py

@@ -445,12 +445,18 @@ def transform_for_container(self):
         self.port = self.DEFAULT_PORT
 
 
-class CompressionWorker(BaseModel):
+class WorkerConfigBase(BaseModel):
     logging_level: LoggingLevel = "INFO"
+    task_soft_time_limit: NonNegativeInt = 600
+    task_time_limit: NonNegativeInt = 1200
 
 
-class QueryWorker(BaseModel):
-    logging_level: LoggingLevel = "INFO"
+class CompressionWorker(WorkerConfigBase):
+    pass
+
+
+class QueryWorker(WorkerConfigBase):
+    pass
 
 
 class Redis(BaseModel):

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

@@ -1,5 +1,6 @@
 from celery import signals
 from celery.app.task import Task
+from celery.exceptions import SoftTimeLimitExceeded
 from celery.utils.log import get_task_logger
 
 from job_orchestration.executor.compress.celery import app
@@ -23,11 +24,17 @@ def compress(
     paths_to_compress_json: str,
     clp_metadata_db_connection_config,
 ):
-    return compression_entry_point(
-        job_id,
-        task_id,
-        clp_io_config_json,
-        paths_to_compress_json,
-        clp_metadata_db_connection_config,
-        logger,
-    )
+    try:
+        return compression_entry_point(
+            job_id,
+            task_id,
+            clp_io_config_json,
+            paths_to_compress_json,
+            clp_metadata_db_connection_config,
+            logger,
+        )
+    except SoftTimeLimitExceeded:
+        logger.exception(
+            f"Compression task job_id={job_id} task_id={task_id} exceeded soft time limit."
+        )
+        raise

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

@@ -1,5 +1,6 @@
 import os
 
+from job_orchestration.executor.utils import load_clp_config_from_config_path_env_var
 from job_orchestration.scheduler.constants import SchedulerType, TASK_QUEUE_HIGHEST_PRIORITY
 
 # Worker settings
@@ -18,6 +19,11 @@
 
 # Results backend settings
 result_persistent = True
+result_expires = 7200
 
 broker_url = os.getenv("BROKER_URL")
 result_backend = os.getenv("RESULT_BACKEND")
+
+_clp_config = load_clp_config_from_config_path_env_var()
+task_soft_time_limit = _clp_config.compression_worker.task_soft_time_limit
+task_time_limit = _clp_config.compression_worker.task_time_limit

components/job-orchestration/job_orchestration/executor/query/celeryconfig.py

@@ -1,5 +1,6 @@
 import os
 
+from job_orchestration.executor.utils import load_clp_config_from_config_path_env_var
 from job_orchestration.scheduler.constants import SchedulerType
 
 imports = (
@@ -17,6 +18,11 @@
 result_backend = os.getenv("RESULT_BACKEND")
 
 result_persistent = True
+result_expires = 7200
+
+_clp_config = load_clp_config_from_config_path_env_var()
+task_soft_time_limit = _clp_config.query_worker.task_soft_time_limit
+task_time_limit = _clp_config.query_worker.task_time_limit
 
 # Differentiate between tasks that have started v.s. tasks still in queue
 task_track_started = True

components/job-orchestration/job_orchestration/executor/query/extract_stream_task.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 from celery.app.task import Task
+from celery.exceptions import SoftTimeLimitExceeded
 from celery.utils.log import get_task_logger
 from clp_py_utils.clp_config import (
     Database,
@@ -181,9 +182,7 @@ def _make_command_and_env_vars(
     return command, env_vars
 
 
-@app.task(bind=True)
-def extract_stream(
-    self: Task,
+def extract_stream_entry_point(
     job_id: str,
     task_id: int,
     job_config: dict,
@@ -298,3 +297,31 @@ def extract_stream(
             logger.info("Finished uploading streams.")
 
     return task_results.model_dump()
+
+
+@app.task(bind=True)
+def extract_stream(
+    self: Task,
+    job_id: str,
+    task_id: int,
+    job_config: dict,
+    archive_id: str,
+    clp_metadata_db_conn_params: dict,
+    results_cache_uri: str,
+    dataset: str | None = None,
+) -> dict[str, Any]:
+    try:
+        return extract_stream_entry_point(
+            job_id,
+            task_id,
+            job_config,
+            archive_id,
+            clp_metadata_db_conn_params,
+            results_cache_uri,
+            dataset,
+        )
+    except SoftTimeLimitExceeded:
+        logger.exception(
+            f"Stream extraction task job_id={job_id} task_id={task_id} exceeded soft time limit."
+        )
+        raise

components/job-orchestration/job_orchestration/executor/query/fs_search_task.py

@@ -5,6 +5,7 @@
 
 import msgpack
 from celery.app.task import Task
+from celery.exceptions import SoftTimeLimitExceeded
 from celery.utils.log import get_task_logger
 from clp_py_utils.clp_config import (
     Database,
@@ -208,9 +209,7 @@ def upload_results_to_s3(
     return
 
 
-@app.task(bind=True)
-def search(
-    self: Task,
+def search_entry_point(
     job_id: str,
     task_id: int,
     job_config_blob: bytes,
@@ -286,3 +285,29 @@ def search(
         upload_results_to_s3(task_results, s3_config, src_file, dest_path)
 
     return task_results.model_dump()
+
+
+@app.task(bind=True)
+def search(
+    self: Task,
+    job_id: str,
+    task_id: int,
+    job_config_blob: bytes,
+    archive_id: str,
+    clp_metadata_db_conn_params: dict,
+    results_cache_uri: str,
+    dataset: str | None = None,
+) -> dict[str, Any]:
+    try:
+        return search_entry_point(
+            job_id,
+            task_id,
+            job_config_blob,
+            archive_id,
+            clp_metadata_db_conn_params,
+            results_cache_uri,
+            dataset,
+        )
+    except SoftTimeLimitExceeded:
+        logger.exception(f"Search task job_id={job_id} task_id={task_id} exceeded soft time limit.")
+        raise

components/job-orchestration/job_orchestration/executor/utils.py

@@ -1,10 +1,28 @@
+"""Utilities shared by job-orchestration executors."""
+
+import os
 from logging import Logger
 from pathlib import Path
 
-from clp_py_utils.clp_config import WorkerConfig
+from clp_py_utils.clp_config import ClpConfig, WorkerConfig
 from clp_py_utils.core import read_yaml_config_file
 
 
+def load_clp_config_from_config_path_env_var() -> ClpConfig:
+    """
+    Loads the CLP config from `CLP_CONFIG_PATH`.
+
+    :return: The loaded `ClpConfig` object.
+    """
+    config_path = os.getenv("CLP_CONFIG_PATH")
+    if config_path is None or "" == config_path:
+        return ClpConfig()
+    config = read_yaml_config_file(Path(config_path))
+    if config is None:
+        return ClpConfig()
+    return ClpConfig.model_validate(config)
+
+
 def load_worker_config(
     config_path: Path,
     logger: Logger,

components/job-orchestration/job_orchestration/scheduler/compress/task_manager/celery_task_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from typing import Any
 
 import celery
@@ -8,6 +9,8 @@
 from job_orchestration.scheduler.compress.task_manager.task_manager import TaskManager
 from job_orchestration.scheduler.task_result import CompressionTaskResult
 
+logger = logging.getLogger(__name__)
+
 
 class CeleryTaskManager(TaskManager):
     class ResultHandle(TaskManager.ResultHandle):
@@ -20,6 +23,9 @@ def get_result(self, timeout: float = 0.1) -> list[CompressionTaskResult] | None
                 return [CompressionTaskResult.model_validate(res) for res in results]
             except celery.exceptions.TimeoutError:
                 return None
+            except celery.exceptions.SoftTimeLimitExceeded:
+                logger.exception("Compression task exceeded soft time limit.")
+                raise
 
     def submit(self, task_params: list[dict[str, Any]]) -> TaskManager.ResultHandle:
         task_instances = [compress.s(**params) for params in task_params]

components/job-orchestration/job_orchestration/scheduler/query/query_scheduler.py

@@ -77,6 +77,9 @@
 )
 from job_orchestration.scheduler.utils import kill_hanging_jobs
 
+TASK_RESULT_GET_TIMEOUT = 10
+TASK_RESULT_GET_INTERVAL = 0.005
+
 # Setup logging
 logger = get_logger("search-job-handler")
 
@@ -837,7 +840,13 @@ def handle_pending_query_jobs(
 def try_getting_task_result(async_task_result):
     if not async_task_result.ready():
         return None
-    return async_task_result.get(interval=0.005)
+    try:
+        return async_task_result.get(
+            timeout=TASK_RESULT_GET_TIMEOUT, interval=TASK_RESULT_GET_INTERVAL
+        )
+    except celery.exceptions.TimeoutError:
+        logger.exception("Timed out waiting for task result.")
+        raise
 
 
 def found_max_num_latest_results(

components/package-template/src/etc/clp-config.template.json.yaml

@@ -104,9 +104,13 @@ telemetry:
 #
 #compression_worker:
 #  logging_level: "INFO"
+#  task_soft_time_limit: 600  # A value of 0 disables the limit
+#  task_time_limit: 1200  # A value of 0 disables the limit
 #
 #query_worker:
 #  logging_level: "INFO"
+#  task_soft_time_limit: 600  # A value of 0 disables the limit
+#  task_time_limit: 1200  # A value of 0 disables the limit
 #
 #webui:
 #  host: "localhost"