test

Author: Nathan903

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

@@ -1,9 +1,12 @@
 from celery import Celery
 
 from job_orchestration.executor.compress import celeryconfig
+from job_orchestration.executor.telemetry import init_telemetry
 
 app = Celery("compress")
 app.config_from_object(celeryconfig)
 
+init_telemetry()
+
 if "__main__" == __name__:
     app.start()

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

@@ -4,6 +4,7 @@
 
 from job_orchestration.executor.compress.celery import app
 from job_orchestration.executor.compress.compression_task import compression_entry_point
+from job_orchestration.executor.telemetry import shutdown_telemetry
 
 # Setup logging
 logger = get_task_logger(__name__)
@@ -12,6 +13,7 @@
 @signals.worker_shutdown.connect
 def worker_shutdown_handler(signal=None, sender=None, **kwargs):
     logger.info("Shutdown signal received.")
+    shutdown_telemetry()
 
 
 @app.task(bind=True)

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

@@ -31,6 +31,7 @@
     s3_put,
 )
 from clp_py_utils.sql_adapter import SqlAdapter
+from opentelemetry.api.metrics import get_meter
 
 from job_orchestration.scheduler.constants import CompressionTaskStatus
 from job_orchestration.scheduler.job_config import (
@@ -43,6 +44,21 @@
 from job_orchestration.scheduler.task_result import CompressionTaskResult
 from job_orchestration.scheduler.utils import is_s3_based_input
 
+# OpenTelemetry counters for compression metrics.
+# Created at module-import time; when telemetry is disabled the meter is a
+# no-op so these counters silently accept ``add()`` calls.
+_compression_meter = get_meter("compression-worker")
+_bytes_input_counter = _compression_meter.create_counter(
+    "clp.compression.bytes_input_total",
+    description="Total bytes of uncompressed log data input for compression",
+    unit="By",
+)
+_bytes_output_counter = _compression_meter.create_counter(
+    "clp.compression.bytes_output_total",
+    description="Total bytes of compressed archive data output from compression",
+    unit="By",
+)
+
 
 def update_compression_task_metadata(db_cursor, task_id, kv):
     if not len(kv):
@@ -628,6 +644,11 @@ def compression_entry_point(
     duration = (datetime.datetime.now() - start_time).total_seconds()
     logger.info(f"[job_id={job_id} task_id={task_id}] COMPRESSION COMPLETED.")
 
+    # Emit telemetry counters for bytes input/output.
+    # Even on partial failure these represent actual work done.
+    _bytes_input_counter.add(worker_output["total_uncompressed_size"])
+    _bytes_output_counter.add(worker_output["total_compressed_size"])
+
     with (
         closing(sql_adapter.create_connection(True)) as db_conn,
         closing(db_conn.cursor(dictionary=True)) as db_cursor,

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

@@ -1,9 +1,18 @@
-from celery import Celery
+from celery import Celery, signals
 
 from job_orchestration.executor.query import celeryconfig
+from job_orchestration.executor.telemetry import init_telemetry, shutdown_telemetry
 
 app = Celery("query")
 app.config_from_object(celeryconfig)
 
+init_telemetry()
+
+
+@signals.worker_shutdown.connect
+def worker_shutdown_handler(**kwargs):
+    shutdown_telemetry()
+
+
 if "__main__" == __name__:
     app.start()

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

@@ -23,6 +23,8 @@
 
 from job_orchestration.executor.query.celery import app
 from job_orchestration.executor.query.utils import (
+    _bytes_output_counter,
+    emit_bytes_scanned,
     report_task_failure,
     run_query_task,
 )
@@ -243,7 +245,7 @@ def extract_stream(
             start_time=start_time,
         )
 
-    task_results, task_stdout_str = run_query_task(
+    task_results, task_stdout_str, stdout_byte_count = run_query_task(
         sql_adapter=sql_adapter,
         logger=logger,
         clp_logs_dir=clp_logs_dir,
@@ -255,6 +257,17 @@ def extract_stream(
         start_time=start_time,
     )
 
+    # Emit telemetry counters.
+    # bytes_scanned: the uncompressed size of the archive that was extracted.
+    emit_bytes_scanned(
+        sql_adapter=sql_adapter,
+        clp_metadata_db_conn_params=clp_metadata_db_conn_params,
+        archive_id=archive_id,
+        dataset=dataset,
+        logger=logger,
+    )
+    _bytes_output_counter.add(stdout_byte_count)
+
     if enable_s3_upload and QueryTaskStatus.SUCCEEDED == task_results.status:
         logger.info("Uploading streams to S3...")
 

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

@@ -22,6 +22,8 @@
 
 from job_orchestration.executor.query.celery import app
 from job_orchestration.executor.query.utils import (
+    _bytes_output_counter,
+    emit_bytes_scanned,
     report_task_failure,
     run_query_task,
 )
@@ -262,7 +264,7 @@ def search(
             start_time=start_time,
         )
 
-    task_results, _ = run_query_task(
+    task_results, _, stdout_byte_count = run_query_task(
         sql_adapter=sql_adapter,
         logger=logger,
         clp_logs_dir=clp_logs_dir,
@@ -274,6 +276,17 @@ def search(
         start_time=start_time,
     )
 
+    # Emit telemetry counters.
+    # bytes_scanned: the uncompressed size of the archive that was queried.
+    emit_bytes_scanned(
+        sql_adapter=sql_adapter,
+        clp_metadata_db_conn_params=clp_metadata_db_conn_params,
+        archive_id=archive_id,
+        dataset=dataset,
+        logger=logger,
+    )
+    _bytes_output_counter.add(stdout_byte_count)
+
     storage_config = worker_config.stream_output.storage
     if (
         StorageType.S3 == storage_config.type

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

@@ -9,10 +9,27 @@
 from typing import Any
 
 from clp_py_utils.clp_config import QUERY_TASKS_TABLE_NAME
+from clp_py_utils.clp_metadata_db_utils import get_archives_table_name
 from clp_py_utils.sql_adapter import SqlAdapter
+from opentelemetry.api.metrics import get_meter
 
 from job_orchestration.scheduler.scheduler_data import QueryTaskResult, QueryTaskStatus
 
+# OpenTelemetry counters for query metrics.
+# Created at module-import time; when telemetry is disabled the meter is a
+# no-op so these counters silently accept ``add()`` calls.
+_query_meter = get_meter("query-worker")
+_bytes_scanned_counter = _query_meter.create_counter(
+    "clp.query.bytes_scanned_total",
+    description="Total bytes of uncompressed log data scanned during queries",
+    unit="By",
+)
+_bytes_output_counter = _query_meter.create_counter(
+    "clp.query.bytes_output_total",
+    description="Total bytes of query results output",
+    unit="By",
+)
+
 
 def get_task_log_file_path(clp_logs_dir: Path, job_id: str, task_id: int) -> Path:
     worker_logs_dir = clp_logs_dir / job_id
@@ -49,7 +66,11 @@ def run_query_task(
     job_id: str,
     task_id: int,
     start_time: datetime.datetime,
-) -> tuple[QueryTaskResult, str]:
+) -> tuple[QueryTaskResult, str, int]:
+    """Run a query subprocess and return the result, stdout string, and stdout byte count.
+
+    :return: Tuple of (task_result, stdout_str, stdout_byte_count).
+    """
     clo_log_path = get_task_log_file_path(clp_logs_dir, job_id, task_id)
     clo_log_file = open(clo_log_path, "w")
 
@@ -113,7 +134,7 @@ def sigterm_handler(_signo, _stack_frame):
     if QueryTaskStatus.FAILED == task_status:
         task_result.error_log_path = str(clo_log_path)
 
-    return task_result, stdout_data.decode("utf-8")
+    return task_result, stdout_data.decode("utf-8"), len(stdout_data)
 
 
 def update_query_task_metadata(
@@ -134,3 +155,32 @@ def update_query_task_metadata(
             WHERE id = {task_id}
         """
         db_cursor.execute(query)
+
+
+def emit_bytes_scanned(
+    sql_adapter: SqlAdapter,
+    clp_metadata_db_conn_params: dict,
+    archive_id: str,
+    dataset: str | None,
+    logger: Logger,
+) -> None:
+    """Emit the ``clp.query.bytes_scanned_total`` counter by looking up the
+    archive's ``uncompressed_size`` from the metadata database.
+    """
+    try:
+        with (
+            closing(sql_adapter.create_connection(True)) as db_conn,
+            closing(db_conn.cursor(dictionary=True)) as db_cursor,
+        ):
+            table_name = get_archives_table_name(
+                clp_metadata_db_conn_params["table_prefix"], dataset
+            )
+            db_cursor.execute(
+                f"SELECT uncompressed_size FROM {table_name} WHERE id = %s",
+                (archive_id,),
+            )
+            row = db_cursor.fetchone()
+            if row is not None:
+                _bytes_scanned_counter.add(row["uncompressed_size"])
+    except Exception:
+        logger.exception("Failed to emit bytes_scanned_total telemetry")

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

@@ -0,0 +1,88 @@
+"""OpenTelemetry metrics initialization for CLP workers.
+
+This module mirrors the Rust implementation in clp-rust-utils/src/telemetry.rs,
+providing equivalent telemetry initialization, shutdown, and consent-checking
+behaviour for Python Celery workers.
+
+NOTE: TELEMETRY_DISABLE_VALUES must be kept consistent with the Rust
+implementation in ``clp-rust-utils/src/telemetry.rs`` and the Python
+implementation in ``clp_package_utils/scripts/start_clp.py``.
+"""
+
+import os
+
+from opentelemetry.api.metrics import set_meter_provider
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+from opentelemetry.sdk.resources import Resource
+
+# Values accepted by CLP_DISABLE_TELEMETRY and DO_NOT_TRACK to disable telemetry.
+TELEMETRY_DISABLE_VALUES = ("1", "true", "yes", "y")
+
+# Default OTLP HTTP endpoint for the OpenTelemetry Collector.
+_DEFAULT_OTEL_ENDPOINT = "http://localhost:4318"
+
+_meter_provider: MeterProvider | None = None
+
+
+def is_telemetry_disabled() -> bool:
+    """Check if telemetry is disabled via environment variables.
+
+    Returns True if either ``CLP_DISABLE_TELEMETRY`` or ``DO_NOT_TRACK`` is
+    set to one of :data:`TELEMETRY_DISABLE_VALUES`.
+    """
+    disable_telemetry = os.getenv("CLP_DISABLE_TELEMETRY")
+    if disable_telemetry is not None and disable_telemetry.lower() in TELEMETRY_DISABLE_VALUES:
+        return True
+
+    do_not_track = os.getenv("DO_NOT_TRACK")
+    if do_not_track is not None and do_not_track.lower() in TELEMETRY_DISABLE_VALUES:
+        return True
+
+    return False
+
+
+def init_telemetry() -> None:
+    """Initialize OpenTelemetry metrics collection.
+
+    Sets the global :class:`MeterProvider` so that subsequent calls to
+    :func:`opentelemetry.api.metrics.get_meter` return a real meter.  When
+    telemetry is disabled (via env vars) this function is a no-op; callers
+    that create counters on the resulting no-op meter will silently do
+    nothing.
+
+    The collector endpoint is read from the standard
+    ``OTEL_EXPORTER_OTLP_ENDPOINT`` environment variable.  ``OTEL_SERVICE_NAME``
+    and ``OTEL_RESOURCE_ATTRIBUTES`` are picked up automatically by
+    :meth:`Resource.create`.
+    """
+    global _meter_provider
+
+    if is_telemetry_disabled():
+        return
+
+    endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", _DEFAULT_OTEL_ENDPOINT)
+    metrics_endpoint = f"{endpoint}/v1/metrics"
+
+    exporter = OTLPMetricExporter(endpoint=metrics_endpoint)
+    reader = PeriodicExportingMetricReader(exporter)
+    provider = MeterProvider(
+        resource=Resource.create(),
+        metric_readers=[reader],
+    )
+
+    set_meter_provider(provider)
+    _meter_provider = provider
+
+
+def shutdown_telemetry() -> None:
+    """Shutdown OpenTelemetry metrics, flushing any pending metric exports.
+
+    Safe to call even when telemetry was never initialized or was disabled.
+    """
+    global _meter_provider
+
+    if _meter_provider is not None:
+        _meter_provider.shutdown()
+        _meter_provider = None

components/job-orchestration/pyproject.toml

@@ -11,6 +11,9 @@ dependencies = [
     "mariadb>=1.1.14",
     "msgpack>=1.1.2",
     "mysql-connector-python>=9.5.0",
+    "opentelemetry-api>=1.20.0",
+    "opentelemetry-exporter-otlp-proto-http>=0.41b0",
+    "opentelemetry-sdk>=1.20.0",
     "pika>=1.3.2",
     "pydantic>=2.12.5",
     "pymongo>=4.16.0",