cf5111e520
- Added PluginStore integration for model management. - Replaced StorageRepository with MinIORepository in Activities, Cleanup, and Training classes. - Updated training logic to handle validation files and improved data management. - Enhanced configuration for MinIO and PluginStore in connectors. - Removed deprecated model repository and storage repository files. - Updated environment variable handling for new configurations.
250 lines
9.1 KiB
Python
250 lines
9.1 KiB
Python
"""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
|
|
from temporalio.worker import PollerBehaviorAutoscaling, Worker
|
|
|
|
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_model.model_repository.plugin_store import PluginStore
|
|
from sientia_do.observability.metrics_controller import MetricsController
|
|
from sientia_do.temporal.worker.prepare_worker import prepare_worker
|
|
|
|
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.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(f'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.validate_train_params,
|
|
activities.train_model,
|
|
activities.cleanup_resources,
|
|
],
|
|
temporal_client=temporal_client,
|
|
logger=logger,
|
|
),
|
|
prepare_worker(
|
|
main_workflow=CleanupFiles,
|
|
other_workflows=[],
|
|
activities=[
|
|
activities.cleanup_minio_files,
|
|
activities.cleanup_temp_directories,
|
|
],
|
|
temporal_client=temporal_client,
|
|
logger=logger,
|
|
),
|
|
]
|
|
|
|
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())
|