sientia-dataops-scouter_temporal/scouter/worker/worker.py

from temporalio import client, workflow
from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
from temporalio.worker import PollerBehaviorAutoscaling, Worker

with workflow.unsafe.imports_passed_through():
    import asyncio
    import os
    import sys

    from prometheus_client import start_http_server
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import get_logger

    from scouter import metrics
    from scouter.activities.activities import Activities
    from scouter.utils.connectors_config import (
        build_mongodb_config,
        build_postgres_config,
        build_redis_config,
    )
    from scouter.workflow.scouter import Scouter
    from scouter.workflow.sub_workflows.core_scouter import CoreScouter

# Environment configuration
POD_ID = os.getenv('HOSTNAME', 'localhost')
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091'))


async def main():
    """
    Main entry point for the Scouter Temporal worker.

    This function initializes and starts all required services:
    - Prometheus metrics server
    - Notification handler for MongoDB
    - Activity implementations for data processing
    - Temporal client and workers
    - Multiple task queues for different workflow types

    The worker supports two main task queues:
    - scouter-queue: Main data processing workflows
    - fake_data-queue: Test data generation workflows

    Returns:
        None

    Raises:
        SystemExit: If worker initialization or execution fails
    """
    host = os.getenv('TEMPORAL_HOST', 'localhost:7233')

    logger = get_logger(__name__)

    metadata = {
        'pod_id': POD_ID,
        'model_name': '-',
        'model_id': '-',
        'workflow_name': '-',
        'schedule_name': '-',
    }

    logger.custom_info(f'Starting Worker with pod_id: {POD_ID}', metadata)

    logger.custom_info('Starting prometheus client...', metadata)
    start_prometheus_server()

    logger.custom_info('Starting Notification Handler...', metadata)

    mongo_config = build_mongodb_config()
    notification_handler = NotificationHandler(
        connection_string=mongo_config['connection_string'],
        database=mongo_config['database_name'],
        logger=logger,
        project_name=os.getenv('PROJECT_NAME', 'scouter'),
    )

    logger.custom_info('Starting Activities...', metadata)

    activities = Activities(
        logger=logger,
        notification_handler=notification_handler,
        postgres_config=build_postgres_config(),
        redis_config=build_redis_config(),
        mongodb_config=build_mongodb_config(),
    )

    logger.custom_info(f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)

    new_runtime = Runtime(
        telemetry=TelemetryConfig(
            metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
        )
    )

    logger.custom_info('Starting Temporal Client...', metadata)

    temporal_client = await client.Client.connect(
        target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'scouter'), runtime=new_runtime
    )

    logger.custom_info('Starting Workers...', metadata)

    workers = [
        Worker(
            temporal_client,
            task_queue='scouter-queue',
            workflows=[Scouter, CoreScouter],
            activities=[
                activities.load_latest_data,
                activities.get_last_data_timestamp,
                activities.put_last_data_timestamp,
                activities.data_quality_gate,
                activities.aggregate_data,
                activities.group_and_hold_data,
                activities.export_data_to_postgres,
                activities.write_metrics,
                activities.store_data_package,
            ],
            max_concurrent_workflow_tasks=50,
            max_concurrent_activities=50,
            max_concurrent_local_activities=50,
            max_cached_workflows=200,
            workflow_task_poller_behavior=PollerBehaviorAutoscaling(),
            activity_task_poller_behavior=PollerBehaviorAutoscaling(),
        )
    ]

    handlers = []
    for w in workers:
        handlers.append(w.run())

    logger.custom_info('Workers started successfully', metadata)

    try:
        await asyncio.gather(*handlers)

    except BaseException as e:  # NOSONAR
        logger.custom_error(
            'An unhandled exception occurred: %s', metadata=metadata
        )
    finally:
        if notification_handler:
            notification_handler.shutdown()
        if activities:
            activities.shutdown()
        # Exit with a non-zero status code to indicate failure to Kubernetes
        metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
        sys.exit(1)


def start_prometheus_server():
    """
    Start the Prometheus metrics HTTP server.

    This function initializes the Prometheus metrics server on the configured
    port and sets the application health status. It's essential for
    monitoring and observability of the Scouter system.

    Returns:
        None

    Raises:
        SystemExit: If metrics server fails to start
    """
    try:
        port = int(os.getenv('HTTP_METRICS_PORT', 9090))
        start_http_server(port)
        print(f'Prometheus server started on port {port}.')
        metrics.APP_UP.labels(pod_id=POD_ID).set(1)  # Mark app as UP
    except Exception as e:
        print(f'Failed to start Prometheus server: {e}')
        os._exit(1)


if __name__ == '__main__':
    asyncio.run(main())