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
    )

POD_ID = os.getenv("HOSTNAME", "localhost")
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))


async def main():
    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', 'laborious'),
        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():
    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())