vitor-aignosi
7512963e19
Enhance MinIO integration and update environment configurations - Added MinIO configuration parameters to .env.example and values.yaml for improved storage management. - Updated requirements.txt to include necessary libraries for MinIO support. - Refactored Activities class to utilize Storage for MinIO interactions. - Enhanced MLFlow class to integrate MinIO for data retrieval during model retraining. - Introduced build_minio_config function to streamline MinIO configuration setup. - Updated minimal_retrain workflow to support data storage in MinIO.
293 lines
10 KiB
Python
293 lines
10 KiB
Python
"""
|
|
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
|
|
from temporalio.converter import PayloadCodec, DataConverter
|
|
from temporalio.api.common.v1 import Payload
|
|
|
|
with workflow.unsafe.imports_passed_through():
|
|
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_minio_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 laborious import metrics
|
|
from prometheus_client import start_http_server
|
|
import lzma
|
|
import dataclasses
|
|
|
|
|
|
POD_ID = os.getenv('POD_ID')
|
|
SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', "9091"))
|
|
LZMA_MIN_MB = float(os.getenv('LZMA_MIN_MB', "1.5"))
|
|
|
|
|
|
class LzmaPayloadCodec(PayloadCodec):
|
|
async def encode(self, payloads):
|
|
out = []
|
|
for p in payloads:
|
|
if p.data:
|
|
old_len = len(p.data) / 1000000
|
|
# Only compress payloads larger than 1.5 MB
|
|
if old_len > LZMA_MIN_MB:
|
|
compressed_data = lzma.compress(p.data)
|
|
new_len = len(compressed_data) / 1000000
|
|
ratio = new_len / old_len if old_len else 0
|
|
print(
|
|
f"[codec] encode lzma: {old_len} MB -> {new_len} MB ({ratio:.2f}x)")
|
|
meta = dict(p.metadata or {})
|
|
meta[b"codec"] = b"lzma"
|
|
out.append(Payload(metadata=meta, data=compressed_data))
|
|
else:
|
|
out.append(p)
|
|
|
|
else:
|
|
out.append(p)
|
|
return out
|
|
|
|
async def decode(self, payloads):
|
|
out = []
|
|
for p in payloads:
|
|
if p.data and (p.metadata or {}).get(b"codec") == b"lzma":
|
|
# comp_len = len(p.data)
|
|
decomp = lzma.decompress(p.data)
|
|
# decomp_len = len(decomp)
|
|
# ratio = (decomp_len / comp_len) if comp_len else 0
|
|
# print(
|
|
# f"[codec] decode lzma: {comp_len} B -> {decomp_len} B ({ratio:.2f}x)")
|
|
meta = dict(p.metadata or {})
|
|
meta.pop(b"codec", None)
|
|
out.append(Payload(metadata=meta, data=decomp))
|
|
else:
|
|
out.append(p)
|
|
return out
|
|
|
|
|
|
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(),
|
|
minio_config=build_minio_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)
|
|
|
|
codec_dc = DataConverter(payload_codec=LzmaPayloadCodec())
|
|
|
|
temporal_client = await client.Client.connect(
|
|
target_host=host,
|
|
namespace=os.getenv('TEMPORAL_NAMESPACE', 'laborious'),
|
|
data_converter=codec_dc,
|
|
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.query_to_minio,
|
|
activities.retrain_model,
|
|
activities.update_production_model,
|
|
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()
|
|
),
|
|
Worker(
|
|
temporal_client,
|
|
task_queue='predictions_batch-queue',
|
|
workflows=[PredictionsBatch, PredictionProcess,
|
|
FormatAndExportPrediction],
|
|
activities=[
|
|
# MLFlow
|
|
activities.request_predict,
|
|
activities.request_transform,
|
|
activities.query_to_minio,
|
|
# 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())
|