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 and cleanup workflows.

The worker supports two task queues:
- train_model-<runtime>-queue: For ML model training workflows
- cleanup_files-<runtime>-queue: For file cleanup workflows

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

Environment Variables:
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: model-manager)
- TEMPORAL_USE_TLS: Enable TLS for Temporal connection (default: false)
- RUNTIME: Runtime identifier used in queue naming (default: single)
- 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

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 Logger as SientiaLogger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_model.model_repository.plugin_store import PluginStore

    from model_manager import metrics
    from model_manager.activities.activities import Activities
    from model_manager.runtime_paths import ensure_runtime_directories
    from model_manager.schedules.cleanup_schedule import create_cleanup_schedule
    from model_manager.utils.connectors_config import (
        build_minio_config,
        build_mlflow_config,
        build_mongodb_config,
        build_plugin_store_config,
        build_postgres_config,
    )
    from model_manager.utils.logger_helper import get_logger
    from model_manager.worker.prepare_worker import prepare_worker
    from model_manager.workflows.cleanup_files import CleanupFiles
    from model_manager.workflows.train_model import TrainModel

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


def _get_runtime(runtime: str | None) -> str:
    """
    Resolve runtime using fallback when missing.

    Args:
     - runtime: str | None, runtime value from environment

    Return:
     str: normalized runtime value
    """
    normalized_runtime = runtime.strip() if runtime else ''
    return normalized_runtime or 'single'


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 and Plugin Store client, then sets ``app_up``
       (what the probes check) to 1. The runtime libraries are already installed: the
       sientia-module entrypoint merges them into the resolved requirements before boot
    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
    """

    runtime = _get_runtime(RUNTIME)

    ensure_runtime_directories()

    host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
    use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
    logger = get_logger(__name__)

    metadata = {
        'pod_id': POD_ID,
        'runtime': runtime,
    }

    start_prometheus_server(logger, 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(f'MongoDB client initialized at {mongo_config["uri"]}', metadata)

    logger.custom_info('Initializing metrics controller', metadata)

    metrics_controller = MetricsController(logger=logger)

    plugin_store_parameters = build_plugin_store_config()
    plugin_store = PluginStore(
        base_url=plugin_store_parameters['base_url'],
        owner=plugin_store_parameters['owner'],
        repo=plugin_store_parameters['repo'],
        username=plugin_store_parameters['username'],
        password=plugin_store_parameters['password'],
        branch=plugin_store_parameters['branch'],
        cache_ttl_seconds=plugin_store_parameters['cache_ttl_seconds'],
        pypi_index_url=plugin_store_parameters['pypi_index_url'],
        pypi_username=plugin_store_parameters['pypi_username'],
        pypi_password=plugin_store_parameters['pypi_password'],
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )

    # The probes only check app_up, so it goes to 1 only after initialization.
    metrics.APP_UP.labels(pod_id=POD_ID).set(1)

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

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

    logger.custom_info(f'SDK metrics server initialized on port {SDK_METRICS_PORT}', metadata)

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

    logger.custom_info(f'Temporal client initialized at {host}/{namespace}', metadata)

    # Create cleanup schedule (idempotent - only creates if doesn't exist)
    try:
        await create_cleanup_schedule(temporal_client, logger, metadata)
    except Exception as e:  # noqa: BLE001
        logger.custom_error(f'Failed to configure cleanup schedule: {e}', metadata)
        # Don't fail the worker startup if schedule creation fails
        # The schedule can be created manually if needed

    workers = [
        prepare_worker(
            main_workflow=TrainModel,
            other_workflows=[],
            activities=[
                activities.update_experiment_run,
                activities.load_model_metadata,
                activities.validate_train_params,
                activities.train_model,
                activities.cleanup_resources,
            ],
            temporal_client=temporal_client,
            logger=logger,
            runtime=runtime,
        ),
        prepare_worker(
            main_workflow=CleanupFiles,
            other_workflows=[],
            activities=[
                activities.cleanup_temp_directories,
            ],
            temporal_client=temporal_client,
            logger=logger,
            runtime=runtime,
        ),
    ]

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

    logger.custom_info('Model manager workers initialized', 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()
        logger.custom_info('MongoDB client closed', metadata)
        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(logger: SientiaLogger, metadata: dict[str, str | None]):
    """
    Starts the Prometheus metrics server for monitoring and observability.

    This function initializes the Prometheus HTTP server on the configured port.
    ``main`` sets ``app_up`` to 1 later, once initialization is done.

    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)
        logger.custom_info(f'Prometheus server initialized on port {port}.', metadata)
    except Exception as e:  # noqa: BLE001
        logger.custom_critical(f'Failed to start Prometheus server: {e}', metadata)
        os._exit(1)


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