SIENTIAPDE-1243: Refactor: Rename 'laborious' package to 'model_manager'
This commit renames the 'laborious' package to 'model_manager' across the entire project. This includes renaming directories, modules, references in code, configuration files, and documentation to reflect the new package name. This change improves clarity and consistency within the project.
This commit is contained in:
Bruno Domingues
0
model_manager/worker/__init__.py
Normal file
0
model_manager/worker/__init__.py
Normal file
237
model_manager/worker/worker.py
Normal file
237
model_manager/worker/worker.py
Normal file
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
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 two main task queues:
|
||||
- predictions_batch-queue: Handles batch prediction workflows
|
||||
- minimal_retrain-queue: Handles model retraining workflows
|
||||
|
||||
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
|
||||
|
||||
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)
|
||||
"""
|
||||
|
||||
from temporalio import workflow, client
|
||||
from temporalio.worker import Worker, PollerBehaviorAutoscaling
|
||||
from temporalio.runtime import Runtime, TelemetryConfig, PrometheusConfig
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import os
|
||||
import sys
|
||||
import asyncio
|
||||
from model_manager.workflows.minimal_retrain import MinimalRetrain
|
||||
from model_manager.workflows.predictions_batch import PredictionsBatch
|
||||
from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess
|
||||
from model_manager.workflows.sub_workflows.format_and_export_prediction import \
|
||||
FormatAndExportPrediction
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.utils.connectors_config import (
|
||||
build_postgres_config,
|
||||
build_mlflow_config,
|
||||
build_opc_config,
|
||||
build_mongodb_config
|
||||
)
|
||||
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
|
||||
from sientia_do.observability.logger import get_logger
|
||||
from model_manager import metrics
|
||||
from prometheus_client import start_http_server
|
||||
|
||||
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 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)
|
||||
|
||||
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', 'laborious')
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Activities...', metadata)
|
||||
|
||||
activities = Activities(
|
||||
postgres_config=build_postgres_config(),
|
||||
mlflow_config=build_mlflow_config(),
|
||||
opc_config=build_opc_config(),
|
||||
logger=logger,
|
||||
notification_handler=notification_handler
|
||||
)
|
||||
|
||||
logger.custom_info('Initializing OPC...', metadata)
|
||||
await activities.init_opc()
|
||||
|
||||
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', 'laborious'),
|
||||
runtime=new_runtime
|
||||
)
|
||||
|
||||
logger.custom_info('Starting Workers...', metadata)
|
||||
|
||||
workers = [
|
||||
Worker(
|
||||
temporal_client,
|
||||
task_queue='minimal_retrain-queue',
|
||||
workflows=[MinimalRetrain],
|
||||
activities=[
|
||||
activities.load_custom_query,
|
||||
activities.retrain_model,
|
||||
activities.update_production_model,
|
||||
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()
|
||||
),
|
||||
Worker(
|
||||
temporal_client,
|
||||
task_queue='predictions_batch-queue',
|
||||
workflows=[PredictionsBatch, PredictionProcess,
|
||||
FormatAndExportPrediction],
|
||||
activities=[
|
||||
# MLFlow
|
||||
activities.request_predict,
|
||||
activities.request_transform,
|
||||
# Gates
|
||||
activities.input_gate,
|
||||
activities.mlflow_response_gate,
|
||||
activities.mlflow_content_gate,
|
||||
activities.format_prediction,
|
||||
activities.format_default_prediction,
|
||||
activities.get_last_timestamp,
|
||||
# OPC
|
||||
activities.write_opc_data,
|
||||
# Postgres
|
||||
activities.load_custom_query,
|
||||
activities.repeat_last_prediction,
|
||||
activities.export_data_to_postgres,
|
||||
activities.write_metrics
|
||||
],
|
||||
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: # NOSONAR
|
||||
logger.custom_error(f"An unhandled exception occurred: {e}", metadata)
|
||||
finally:
|
||||
if notification_handler:
|
||||
notification_handler.shutdown()
|
||||
if activities:
|
||||
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 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())
|
||||
Reference in New Issue
Block a user