Add laborious
This commit is contained in:
vitor-aignosi
0
laborious/worker/__init__.py
Normal file
0
laborious/worker/__init__.py
Normal file
336
laborious/worker/worker.py
Normal file
336
laborious/worker/worker.py
Normal file
@@ -0,0 +1,336 @@
|
||||
"""
|
||||
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_*: import configuration, see
|
||||
``laborious.utils.connectors_config``
|
||||
"""
|
||||
|
||||
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'))
|
||||
|
||||
|
||||
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 os.getenv('IMPORT_MODEL_ENABLED', '').strip().lower() in ('1', 'true', 'yes', 'on'):
|
||||
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,
|
||||
# Bundle
|
||||
activities.download_import_bundle,
|
||||
activities.open_import_bundle,
|
||||
# MLflow provisioning
|
||||
activities.create_import_experiment,
|
||||
activities.upload_import_artifacts,
|
||||
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 BaseException as e: # NOSONAR
|
||||
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())
|
||||
Reference in New Issue
Block a user