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

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

with workflow.unsafe.imports_passed_through():
    import sys
    import os
    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.observability.logger import get_logger
    from scouter.activities.activities import Activities
    from scouter.workflow.scouter import Scouter
    from scouter.workflow.sub_workflows.core_scouter import CoreScouter
    from scouter.workflow.fake_data import FakeData
    from scouter.activities.faker import Faker
    import asyncio
    from prometheus_client import start_http_server
    from scouter import metrics
    from scouter.utils.connectors_config import (
        build_postgres_config,
        build_redis_config,
        build_mongodb_config
    )

# 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('Starting Faker Activities...', metadata)

    faker_activities = Faker(
        logger=logger,
        notification_handler=notification_handler,
        bootstrap_servers=os.getenv(
            'KAFKA_BOOTSTRAP_SERVERS', 'localhost:9092')
    )

    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()
        ),
        Worker(
            temporal_client,
            task_queue='fake_data-queue',
            workflows=[FakeData],
            activities=[
                faker_activities.generate_and_send_data,
            ]
        )
    ]

    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",
                            e, exc_info=True, 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())