sientia-dataops-model-manager/model_manager/worker/worker.py

"""Model Manager Worker Module

This module provides the main worker implementation for the Sientia DataOps Model Manager system.
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
model training workflows.

The worker supports the train_model-queue task queue for ML model training workflows.

Key Features:
- Automatic scaling with PollerBehaviorAutoscaling
- Prometheus metrics integration
- Comprehensive error handling and logging
- Graceful shutdown with cleanup
- ML model training pipeline orchestration

Environment Variables:
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: model_manager)
- POD_ID: Kubernetes pod identifier for metrics
- HTTP_METRICS_PORT: Prometheus metrics server port (default: 9090)
- HTTP_SDK_METRICS_PORT: Temporal SDK metrics port (default: 9091)
- PROJECT_NAME: Project name for notifications (default: model_manager)
"""

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 model_manager import metrics
    from model_manager.activities.activities import Activities
    from model_manager.utils.connectors_config import (
        build_minio_config,
        build_mlflow_config,
        build_mongodb_config,
        build_postgres_config,
    )
    from model_manager.utils.logger_helper import get_logger
    from model_manager.workflows.train_model import TrainModel

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


async def main():
    """
    Main entry point for the Model Manager worker application.

    This function initializes and starts all components of the worker:
    1. Sets up logging and metadata
    2. Starts Prometheus metrics server
    3. Initializes notification handler
    4. Creates and configures activities
    5. Starts Temporal client and workers
    6. Manages worker lifecycle and graceful shutdown

    The function runs indefinitely until interrupted or an error occurs.
    On error, it performs cleanup and exits with a non-zero status code.

    Raises:
        Exception: Any unhandled exception during worker execution
        SystemExit: On graceful shutdown or error conditions
    """
    host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
    logger = get_logger(__name__)

    metadata = {
        'pod_id': POD_ID,
        'workflow_name': 'train_model',
    }

    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', 'model-manager'),
    )

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

    activities = Activities(
        postgres_config=build_postgres_config(),
        mlflow_config=build_mlflow_config(),
        minio_config=build_minio_config(),
        logger=logger,
        notification_handler=notification_handler,
    )

    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(f'Starting Temporal Client at {host}...', metadata)

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

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

    workers = [
        Worker(
            temporal_client,
            task_queue='train_model-queue',
            workflows=[TrainModel],
            activities=[
                activities.update_experiment_run,
                activities.validate_train_params,
                activities.train_model,
                activities.cleanup_resources,
            ],
            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:
        # This will run the workers and wait for them to complete.
        # If an exception occurs in any of the worker handlers, it will be propagated here.
        await asyncio.gather(*handlers)
    except BaseException as e:  # noqa: BLE001
        logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
    finally:
        notification_handler.shutdown()
        await 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():
    """
    Starts the Prometheus metrics server for monitoring and observability.

    This function initializes the Prometheus HTTP server on the configured port
    and sets the application health metric to indicate the service is running.

    The server exposes metrics that can be scraped by Prometheus for monitoring
    the health and performance of the Model Manager worker.

    Environment Variables:
        HTTP_METRICS_PORT: Port for the metrics server (default: 9090)
        POD_ID: Pod identifier for metrics labeling

    Raises:
        SystemExit: If the 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:  # noqa: BLE001
        print(f'Failed to start Prometheus server: {e}')
        os._exit(1)


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