"""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-queue: For ML model training workflows - cleanup-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) - TRAIN_TASK_QUEUE: Task queue for training workflows (default: train_model-queue) - CLEANUP_TASK_QUEUE: Task queue for cleanup workflows (default: cleanup-queue) - 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.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')) TRAIN_TASK_QUEUE = os.getenv('TRAIN_TASK_QUEUE', 'train_model-queue') CLEANUP_TASK_QUEUE = os.getenv('CLEANUP_TASK_QUEUE', 'cleanup-queue') 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 """ if not RUNTIME: raise ValueError('RUNTIME environment variable is required') 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) logger.custom_info(f'Installing runtime {RUNTIME}', metadata) 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, ) await plugin_store.install_runtime(runtime_name=RUNTIME) 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) temporal_client = await client.Client.connect( target_host=host, namespace=os.getenv('TEMPORAL_NAMESPACE', 'model-manager'), runtime=new_runtime, tls=use_tls, ) logger.custom_info(f'Temporal client initialized at {host}', 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 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) logger.custom_info(f'Prometheus server initialized on port {port}.', metadata) metrics.APP_UP.labels(pod_id=POD_ID).set(1) # Mark app as UP 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())