feat(job-orchestration): Add Spider compression task; Move Celery compression task into a dedicated submodule. (#1340)

Co-authored-by: LinZhihao-723 <[email protected]>

Author: sitaowang1998

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

@@ -0,0 +1,34 @@
+from celery import signals
+from celery.app.task import Task
+from celery.utils.log import get_task_logger
+from job_orchestration.executor.compress.celery import app
+from job_orchestration.executor.compress.compression_task import compression_entry_point
+
+# Setup logging
+logger = get_task_logger(__name__)
+
+
+@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,
+    job_id: int,
+    task_id: int,
+    tag_ids,
+    clp_io_config_json: str,
+    paths_to_compress_json: str,
+    clp_metadata_db_connection_config,
+):
+    return compression_entry_point(
+        job_id,
+        task_id,
+        tag_ids,
+        clp_io_config_json,
+        paths_to_compress_json,
+        clp_metadata_db_connection_config,
+        logger,
+    )

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

@@ -6,13 +6,13 @@
 # Force workers to consume only one task at a time
 worker_prefetch_multiplier = 1
 imports = [
-    "job_orchestration.executor.compress.compression_task",
+    "job_orchestration.executor.compress.celery_compress",
 ]
 
 # Queue settings
 task_queue_max_priority = TASK_QUEUE_HIGHEST_PRIORITY
 task_routes = {
-    "job_orchestration.executor.compress.compression_task.compress": SchedulerType.COMPRESSION,
+    "job_orchestration.executor.compress.celery_compress.compress": SchedulerType.COMPRESSION,
 }
 task_create_missing_queues = True
 

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

@@ -6,9 +6,6 @@
 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 (
     CLP_DB_PASS_ENV_VAR_NAME,
     CLP_DB_USER_ENV_VAR_NAME,
@@ -32,7 +29,6 @@
     s3_put,
 )
 from clp_py_utils.sql_adapter import SQL_Adapter
-from job_orchestration.executor.compress.celery import app
 from job_orchestration.scheduler.constants import CompressionTaskStatus
 from job_orchestration.scheduler.job_config import (
     ClpIoConfig,
@@ -42,9 +38,6 @@
 )
 from job_orchestration.scheduler.scheduler_data import CompressionTaskResult
 
-# Setup logging
-logger = get_task_logger(__name__)
-
 
 def update_compression_task_metadata(db_cursor, task_id, kv):
     if not len(kv):
@@ -322,6 +315,7 @@ def run_clp(
     paths_to_compress: PathsToCompress,
     sql_adapter: SQL_Adapter,
     clp_metadata_db_connection_config,
+    logger,
 ):
     """
     Compresses logs into archives.
@@ -336,6 +330,7 @@ def run_clp(
     :param paths_to_compress: PathToCompress
     :param sql_adapter: SQL_Adapter
     :param clp_metadata_db_connection_config
+    :param logger
     :return: tuple -- (whether compression was successful, output messages)
     """
     instance_id_str = f"compression-job-{job_id}-task-{task_id}"
@@ -523,20 +518,14 @@ 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,
+def compression_entry_point(
     job_id: int,
     task_id: int,
     tag_ids,
     clp_io_config_json: str,
     paths_to_compress_json: str,
-    clp_metadata_db_connection_config,
+    clp_metadata_db_connection_config: Dict[str, Any],
+    logger,
 ):
     clp_home = pathlib.Path(os.getenv("CLP_HOME"))
 
@@ -578,6 +567,7 @@ def compress(
         paths_to_compress,
         sql_adapter,
         clp_metadata_db_connection_config,
+        logger,
     )
     duration = (datetime.datetime.now() - start_time).total_seconds()
     logger.info(f"[job_id={job_id} task_id={task_id}] COMPRESSION COMPLETED.")

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

@@ -0,0 +1,49 @@
+import json
+
+from clp_py_utils.clp_logging import get_logger
+from job_orchestration.executor.compress.compression_task import compression_entry_point
+from job_orchestration.utils.spider_utils import int8_list_to_utf8_str, utf8_str_to_int8_list
+from spider_py import Int8, Int64, TaskContext
+
+# Setup logging
+logger = get_logger("spider_compression")
+
+
+def compress(
+    _: TaskContext,
+    job_id: Int64,
+    task_id: Int64,
+    tag_ids: list[Int64],
+    clp_io_config_json: list[Int8],
+    paths_to_compress_json: list[Int8],
+    clp_metadata_db_connection_config_json: list[Int8],
+) -> list[Int8]:
+    """
+    Compresses files using the general compression entry point.
+    :param _: Spider's task context. Not used in the function.
+    :param job_id:
+    :param task_id:
+    :param tag_ids:
+    :param clp_io_config_json: A JSON string representation of
+        `job_orchestration.scheduler.constants.ClpIoConfig`.
+    :param paths_to_compress_json: A JSON string representation of
+        `job_orchestration.scheduler.constants.PathToCompress`.
+    :param clp_metadata_db_connection_config_json: A JSON string representation of
+        `clp_py_utils.clp_config.Database`.
+    :return: A JSON string representation of
+        `job_orchestration.scheduler.constants.CompressionTaskResult`, encoded as a list of
+        `spider_py.Int8`.
+    """
+    result_as_json_str = json.dumps(
+        compression_entry_point(
+            int(job_id),
+            int(task_id),
+            [int(tag_id) for tag_id in tag_ids],
+            int8_list_to_utf8_str(clp_io_config_json),
+            int8_list_to_utf8_str(paths_to_compress_json),
+            json.loads(int8_list_to_utf8_str(clp_metadata_db_connection_config_json)),
+            logger,
+        )
+    )
+
+    return utf8_str_to_int8_list(result_as_json_str)

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

@@ -30,7 +30,7 @@
 from clp_py_utils.core import read_yaml_config_file
 from clp_py_utils.s3_utils import s3_get_object_metadata
 from clp_py_utils.sql_adapter import SQL_Adapter
-from job_orchestration.executor.compress.compression_task import compress
+from job_orchestration.executor.compress.celery_compress import compress
 from job_orchestration.scheduler.compress.partition import PathsToCompressBuffer
 from job_orchestration.scheduler.constants import (
     CompressionJobStatus,

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

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import spider_py
+
+
+def int8_list_to_utf8_str(byte_list: list[spider_py.Int8]) -> str:
+    """
+    Converts a list of `spider_py.Int8` values to a UTF-8 encoded string.
+
+    :param byte_list:
+    :return: Decoded UTF-8 string constructed from the input byte list.
+    """
+    return bytes(int(byte) for byte in byte_list).decode("utf-8")
+
+
+def utf8_str_to_int8_list(utf8_str: str) -> list[spider_py.Int8]:
+    """
+    Converts a UTF-8 encoded string to a list of `spider_py.Int8` values.
+
+    :param utf8_str:
+    :return: A list of `spider_py.Int8` values representing the input string.
+    """
+    return [spider_py.Int8(byte) for byte in utf8_str.encode("utf-8")]