sientia-dataops-laborious-legacy/laborious/worker/worker.py

"""
Laborious Worker Module

This module provides the main worker implementation for the Sientia DataOps Laborious system.
It orchestrates Temporal workers, manages task queues, and handles the lifecycle of
prediction and retraining workflows.

The worker supports multiple runtime-scoped task queues (via ``sientia_do.temporal.worker.prepare_worker``):
- predictions_batch-{runtime}-queue: Batch prediction workflows (heavy workload)
- minimal_retrain-{runtime}-queue: Model retraining workflows
- drift-{runtime}-queue: Drift detection workflows
- simple_metrics-{runtime}-queue: Simple metrics workflows
- import_model-{runtime}-queue: `.sientia` model import — registered **only** when
  ``IMPORT_MODEL_ENABLED`` is truthy, so the containment boundary stays a deployment decision

``RUNTIME`` must be set; it is passed to every ``prepare_worker`` call. Schedulers must use the
same queue names (breaking change vs legacy ``drift-queue`` / ``simple_metrics-queue``).

Key Features:
- Resource-based scaling with WorkerTuner (CPU and memory aware)
- Automatic polling scaling with PollerBehaviorAutoscaling
- Prometheus metrics integration
- Comprehensive error handling and logging
- Graceful shutdown with cleanup
- Multiple worker instances for different workflow types

Environment Variables:
- RUNTIME: Required non-empty string; suffix for all task queue names
- TEMPORAL_HOST: Temporal server address (default: localhost:7233)
- TEMPORAL_NAMESPACE: Temporal namespace (default: laborious)
- 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: laborious)
- IMPORT_MODEL_ENABLED: register the fifth (import) worker; unset by default
- IMPORT_PASSWORD_KEY / IMPORT_BUNDLE_* / IMPORT_STATUS_DB_NAME: import configuration, see
  ``laborious.utils.connectors_config`` (the import log's server is POSTGRES_*)
"""

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 get_logger
    from sientia_do.temporal.worker.prepare_worker import prepare_worker
    from sientia_do.utils.connectors_config import (
        build_api_config,
        build_mongodb_config,
        build_postgres_config,
    )

    from laborious import metrics
    from laborious.activities.activities import Activities
    from laborious.utils.connectors_config import (
        build_import_config,
        build_import_status_config,
        build_minio_config,
        build_mlflow_config,
        build_opc_config,
    )
    from laborious.workflows.drift import Drift
    from laborious.workflows.import_model import ImportModel
    from laborious.workflows.minimal_retrain import MinimalRetrain
    from laborious.workflows.predictions_batch import PredictionsBatch
    from laborious.workflows.simple_metrics import SimpleMetrics
    from laborious.workflows.sub_workflows.format_and_export_prediction import (
        FormatAndExportPrediction,
    )
    from laborious.workflows.sub_workflows.prediction_process import PredictionProcess

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

# The values that turn the import worker on. Anything else — unset, empty, `0`, `false`, a typo —
# leaves the four default workers alone and polls no import queue at all.
IMPORT_ENABLED_VALUES = ('1', 'true', 'yes', 'on')


def import_worker_enabled() -> bool:
    """Whether this runtime registers the fifth (import) worker.

    Read at worker start, not per workflow: containment (QTZPOC-19) wants the importer in its own
    runtime with its own credentials, and this flag is what makes that a deployment decision.
    """
    return os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in IMPORT_ENABLED_VALUES


async def main():
    """
    Main entry point for the Laborious 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. Initializes OPC connections
    6. Starts Temporal client and workers
    7. 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,
        'model_name': '-',
        'model_id': '-',
        'workflow_name': '-',
        'schedule_name': '-',
    }

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

    runtime = os.getenv('RUNTIME', '').strip()
    if not runtime:
        logger.custom_critical(
            'RUNTIME environment variable is required and must be non-empty',
            metadata,
        )
        metrics.APP_UP.labels(pod_id=POD_ID).set(0)
        sys.exit(1)

    metadata_runtime = {**metadata, 'runtime': runtime}

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

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

    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', 'laborious'),
    )

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

    activities = Activities(
        postgres_config=build_postgres_config(),
        mlflow_config=build_mlflow_config(),
        minio_config=build_minio_config(),
        opc_config=build_opc_config(),
        pi_web_api_config=build_api_config(),
        logger=logger,
        notification_handler=notification_handler,
        mongo_config=mongo_config,
        import_config=build_import_config(),
        import_status_config=build_import_status_config(),
    )

    logger.custom_info('Initializing OPC...', metadata_runtime)
    await activities.init_opc()

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

    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_runtime)

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

    logger.custom_info(f'Starting Workers (runtime={runtime})...', metadata_runtime)

    workers = [
        prepare_worker(
            temporal_client=temporal_client,
            main_workflow=MinimalRetrain,
            other_workflows=[],
            activities=[
                activities.load_query_with_minio_offload,
                activities.retrain_model,
                activities.update_production_model,
                activities.format_retrain_report,
                activities.export_data_to_postgres,
            ],
            logger=logger,
            runtime=runtime,
        ),
        prepare_worker(
            temporal_client=temporal_client,
            main_workflow=SimpleMetrics,
            other_workflows=[],
            activities=[
                activities.load_custom_query,
                activities.calculate_simple_metrics,
                activities.export_data_to_postgres,
            ],
            logger=logger,
            runtime=runtime,
        ),
        prepare_worker(
            temporal_client=temporal_client,
            main_workflow=Drift,
            other_workflows=[],
            activities=[
                activities.load_custom_query,
                activities.get_reference_data,
                activities.calculate_drift,
                activities.export_data_to_postgres,
            ],
            logger=logger,
            runtime=runtime,
        ),
        prepare_worker(
            temporal_client=temporal_client,
            main_workflow=PredictionsBatch,
            other_workflows=[PredictionProcess, FormatAndExportPrediction],
            activities=[
                # MLFlow
                activities.request_predict,
                activities.request_transform,
                # Gates
                activities.input_gate,
                activities.mlflow_response_gate,
                activities.mlflow_content_gate,
                activities.format_transformed_data,
                activities.format_prediction,
                activities.format_default_prediction,
                # OPC
                activities.write_opc_data,
                # Postgres / MinIO offload
                activities.load_query_with_minio_offload,
                activities.cleanup_minio_objects_expired,
                activities.repeat_last_prediction,
                activities.export_data_to_postgres,
                activities.export_payload_to_postgres,
                activities.write_metrics,
                # Pi Web API
                activities.write_pi_web_api_data,
            ],
            logger=logger,
            runtime=runtime,
        ),
    ]

    # The fifth worker is opt-in: containment (QTZPOC-19) wants the importer in its own runtime with
    # its own credentials, and a flag lets that be a deployment decision instead of forcing every
    # Laborious runtime to poll the import queue today. Unset, the four workers above are unaffected
    # and no import queue is polled at all.
    if import_worker_enabled():
        logger.custom_info('Import model worker enabled', metadata_runtime)
        workers.append(
            prepare_worker(
                temporal_client=temporal_client,
                main_workflow=ImportModel,
                other_workflows=[],
                activities=[
                    # Import log row (the BFF database, update-only)
                    activities.claim_import_status,
                    activities.record_import_status,
                    activities.record_import_names,
                    activities.record_import_terminal_status,
                    activities.report_import_status_write_failure,
                    activities.report_import_pipeline_failure,
                    # Bundle download/open, MLflow experiment creation and artifact upload — one
                    # fused activity, so they never split across worker processes.
                    activities.prepare_and_upload_import_bundle,
                    # MLflow provisioning
                    activities.register_import_model_version,
                    # The model listing
                    activities.write_import_model_document,
                    # Cleanup and retention
                    activities.cleanup_import_files,
                    activities.ensure_import_bundle_retention,
                ],
                logger=logger,
                runtime=runtime,
            )
        )

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

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

    exit_code = 0
    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 (SystemExit, KeyboardInterrupt) as e:
        # Deliberate shutdown (Ctrl+C, a handler's own sys.exit) — not a failure to log as one.
        # `finally` below calls sys.exit(exit_code) unconditionally, which supersedes this raise
        # and becomes the exception that actually propagates — the raise is here so this branch
        # still honors the reraise-on-SystemExit contract, not because it changes what exits.
        exit_code = e.code if isinstance(e, SystemExit) and isinstance(e.code, int) else 0
        raise
    except Exception as e:
        logger.custom_error(f'An unhandled exception occurred: {e}', metadata)
        exit_code = 1
    finally:
        if notification_handler:
            notification_handler.shutdown()
        if activities:
            await activities.shutdown()
        metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
        sys.exit(exit_code)


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 Laborious 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:
        print(f'Failed to start Prometheus server: {e}')
        os._exit(1)


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