SIENTIAPDE-1646 Sync laborious/ from local release/SIENTIAPDE-1646 (272e02d)

laborious/worker/worker.py

@@ -1,78 +1,88 @@
 """
 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.
+Entry process that connects to Temporal, registers Laborious activities, and runs four workers in
+parallel. Each worker shares the same ``Activities`` instance (single Postgres pool, single MLflow
+repository, single PluginStore handle) but polls a different task queue.
 
-The worker supports two main task queues:
-- predictions_batch-queue: Handles batch prediction workflows
-- minimal_retrain-queue: Handles model retraining workflows
+Task queues (see ``sientia_do.temporal.worker.prepare_worker``):
+- ``predictions_batch-{runtime}-queue`` + sub-workflows on the same queue (ML-heavy path).
+- ``minimal_retrain-{runtime}-queue`` (retrain + promote + export).
+- ``drift-queue`` and ``simple_metrics-queue`` without a runtime suffix so existing schedulers
+  keep stable queue names.
 
-Key Features:
-- Automatic scaling with PollerBehaviorAutoscaling
-- Prometheus metrics integration
-- Comprehensive error handling and logging
-- Graceful shutdown with cleanup
-- Multiple worker instances for different workflow types
+Bootstrap order:
+1. Prometheus app metrics and Mongo-backed notification handler.
+2. ``RUNTIME`` validation and ``PluginStore.install_runtime`` so ``SientiaModel`` code is importable.
+3. ``Activities`` construction (builds ``SientiaMLflowRepository`` internally from env).
+4. OPC client initialization inside activities.
+5. Temporal ``Runtime`` with SDK Prometheus bind, client connect, then ``prepare_worker`` per workflow.
+
+Shutdown closes workers, notifications, activities (pools + OPC), and clears ``app_up``.
 
 Environment Variables:
-- 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)
+- RUNTIME: Required non-empty string passed to ``install_runtime``.
+- STORE_* / PYPI_*: Plugin store and private index (see ``build_plugin_store_config``).
+- TEMPORAL_HOST, TEMPORAL_NAMESPACE: Cluster connection.
+- POD_ID, HTTP_METRICS_PORT, HTTP_SDK_METRICS_PORT: Observability.
+- PROJECT_NAME, MONGODB_*: Notifications (via ``build_mongodb_config`` in handler).
+- POSTGRES_*, MINIO_*, OPC_*, PI_WEB_API_*, MLFLOW_*: Passed through ``Activities`` helpers.
 """
 
-from temporalio import workflow, client
-from temporalio.worker import Worker, PollerBehaviorAutoscaling
-from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
+from temporalio import client, workflow
+from temporalio.runtime import PrometheusConfig, Runtime, TelemetryConfig
 
 with workflow.unsafe.imports_passed_through():
+    import asyncio
     import os
     import sys
-    import asyncio
-    from laborious.workflows.minimal_retrain import MinimalRetrain
-    from laborious.workflows.predictions_batch import PredictionsBatch
-    from laborious.workflows.sub_workflows.prediction_process import PredictionProcess
-    from laborious.workflows.sub_workflows.format_and_export_prediction import \
-        FormatAndExportPrediction
-    from laborious.activities.activities import Activities
-    from laborious.utils.connectors_config import (
-        build_postgres_config,
-        build_mlflow_config,
-        build_opc_config,
-        build_mongodb_config
-    )
+
+    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 laborious import metrics
-    from prometheus_client import start_http_server
+    from sientia_do.observability.metrics_controller import MetricsController
+    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 sientia_model.model_repository.plugin_store import PluginStore
 
-POD_ID = os.getenv('POD_ID')
-SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
+    from laborious import metrics
+    from laborious.activities.activities import Activities
+    from laborious.utils.connectors_config import (
+        build_minio_config,
+        build_opc_config,
+        build_plugin_store_config,
+    )
+    from laborious.workflows.drift import Drift
+    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'))
 
 
 async def main():
     """
-    Main entry point for the Laborious worker application.
+    Run the full worker lifecycle: metrics, notifications, runtime install, workers, gather.
 
-    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.
+    Exits the process with code 0 on normal completion of all worker tasks, or 1 after logging
+    if any worker raises. ``finally`` always shuts down notifications and activities and sets
+    ``app_up`` to 0 before ``sys.exit``.
 
     Raises:
-        Exception: Any unhandled exception during worker execution
-        SystemExit: On graceful shutdown or error conditions
+        Exception: Propagated from ``asyncio.gather`` only before ``finally`` handling; typically
+            workers run until cancelled.
+
+    Return:
+        None (process terminates via ``sys.exit`` from the ``finally`` block).
     """
     host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
     logger = get_logger(__name__)
@@ -87,7 +97,7 @@ async def main():
 
     logger.custom_info(f'Starting Worker with POD_ID: {POD_ID}', metadata)
 
-    logger.custom_info("Starting prometheus client...", metadata)
+    logger.custom_info('Starting prometheus client...', metadata)
     start_prometheus_server()
 
     logger.custom_info('Starting Notification Handler...', metadata)
@@ -97,29 +107,77 @@ async def main():
         connection_string=mongo_config['connection_string'],
         database=mongo_config['database_name'],
         logger=logger,
-        project_name=os.getenv('PROJECT_NAME', 'laborious')
+        project_name=os.getenv('PROJECT_NAME', 'laborious'),
     )
 
+    metrics_controller = MetricsController(logger=logger)
+
+    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(f'Installing PluginStore runtime: {runtime}', metadata_runtime)
+
+    ps_cfg = build_plugin_store_config()
+    plugin_store = PluginStore(
+        base_url=ps_cfg['base_url'],
+        owner=ps_cfg['owner'],
+        repo=ps_cfg['repo'],
+        username=ps_cfg['username'],
+        password=ps_cfg['password'],
+        branch=ps_cfg['branch'],
+        cache_ttl_seconds=ps_cfg['cache_ttl_seconds'],
+        pypi_index_url=ps_cfg['pypi_index_url'],
+        pypi_username=ps_cfg['pypi_username'],
+        pypi_password=ps_cfg['pypi_password'],
+        logger=logger,
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
+    )
+
+    if runtime == 'legacy':
+        to_install_runtime = 'single'
+    else:
+        to_install_runtime = runtime
+
+    try:
+        await plugin_store.install_runtime(
+            runtime_name=to_install_runtime, metadata=metadata_runtime
+        )
+    except Exception as exc:
+        logger.custom_critical(
+            f'Failed to install runtime {to_install_runtime}: {exc}', metadata_runtime
+        )
+        metrics.APP_UP.labels(pod_id=POD_ID).set(0)
+        sys.exit(1)
+
     logger.custom_info('Starting Activities...', metadata)
 
     activities = Activities(
         postgres_config=build_postgres_config(),
-        mlflow_config=build_mlflow_config(),
+        plugin_store=plugin_store,
+        minio_config=build_minio_config(),
         opc_config=build_opc_config(),
+        pi_web_api_config=build_api_config(),
         logger=logger,
-        notification_handler=notification_handler
+        notification_handler=notification_handler,
+        metrics_controller=metrics_controller,
     )
 
     logger.custom_info('Initializing OPC...', metadata)
-    await activities.init_opc()
+    activities.init_opc()
 
-    logger.custom_info(
-        f'Starting SDK Metrics Server on port {SDK_METRICS_PORT}...', metadata)
+    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}")
+            metrics=PrometheusConfig(bind_address=f'0.0.0.0:{SDK_METRICS_PORT}')
         )
     )
 
@@ -128,34 +186,55 @@ async def main():
     temporal_client = await client.Client.connect(
         target_host=host,
         namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
-        runtime=new_runtime
+        runtime=new_runtime,
     )
 
     logger.custom_info('Starting Workers...', metadata)
 
     workers = [
-        Worker(
-            temporal_client,
-            task_queue='minimal_retrain-queue',
-            workflows=[MinimalRetrain],
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=MinimalRetrain,
+            other_workflows=[],
             activities=[
-                activities.load_custom_query,
+                activities.load_query_with_minio_offload,
                 activities.retrain_model,
                 activities.update_production_model,
-                activities.export_data_to_postgres
+                activities.format_retrain_report,
+                activities.export_data_to_postgres,
             ],
-            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()
+            logger=logger,
+            runtime=runtime,
         ),
-        Worker(
-            temporal_client,
-            task_queue='predictions_batch-queue',
-            workflows=[PredictionsBatch, PredictionProcess,
-                       FormatAndExportPrediction],
+        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='core',
+        ),
+        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='core',
+        ),
+        prepare_worker(
+            temporal_client=temporal_client,
+            main_workflow=PredictionsBatch,
+            other_workflows=[PredictionProcess, FormatAndExportPrediction],
             activities=[
                 # MLFlow
                 activities.request_predict,
@@ -164,24 +243,24 @@ async def main():
                 activities.input_gate,
                 activities.mlflow_response_gate,
                 activities.mlflow_content_gate,
+                activities.format_transformed_data,
                 activities.format_prediction,
                 activities.format_default_prediction,
-                activities.get_last_timestamp,
                 # OPC
                 activities.write_opc_data,
-                # Postgres
-                activities.load_custom_query,
+                # Postgres / MinIO offload
+                activities.load_query_with_minio_offload,
+                activities.cleanup_minio_objects_expired,
                 activities.repeat_last_prediction,
                 activities.export_data_to_postgres,
-                activities.write_metrics
+                activities.export_payload_to_postgres,
+                activities.write_metrics,
+                # Pi Web API
+                activities.write_pi_web_api_data,
             ],
-            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()
-        )
+            logger=logger,
+            runtime=runtime,
+        ),
     ]
 
     handlers = []
@@ -190,20 +269,17 @@ async def main():
 
     logger.custom_info('Workers started successfully', metadata)
 
+    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 BaseException as e:  # NOSONAR
-        logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
+        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()
-        # Exit with a non-zero status code to indicate failure to Kubernetes
+        notification_handler.shutdown()
+        activities.shutdown()
         metrics.APP_UP.labels(pod_id=POD_ID).set(0)  # Mark app as DOWN
-        sys.exit(1)
+        sys.exit(exit_code)
 
 
 def start_prometheus_server():
@@ -224,12 +300,12 @@ def start_prometheus_server():
         SystemExit: If the metrics server fails to start
     """
     try:
-        port = int(os.getenv("HTTP_METRICS_PORT", 9090))
+        port = int(os.getenv('HTTP_METRICS_PORT', 9090))
         start_http_server(port)
-        print(f"Prometheus server started on port {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}")
+        print(f'Failed to start Prometheus server: {e}')
         os._exit(1)