Code import - branch release/SIENTIAPDE-1645
This commit is contained in:
0
model_manager/worker/__init__.py
Normal file
0
model_manager/worker/__init__.py
Normal file
119
model_manager/worker/prepare_worker.py
Normal file
119
model_manager/worker/prepare_worker.py
Normal file
@@ -0,0 +1,119 @@
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
from sientia_do.observability.logger import Logger
|
||||
from temporalio.client import Client
|
||||
from temporalio.worker import PollerBehaviorAutoscaling, Worker
|
||||
|
||||
# Worker configuration parameters with default values.
|
||||
parameters = [
|
||||
('MAX_CONCURRENT_WORKFLOW_TASKS', '200'),
|
||||
('MAX_CONCURRENT_ACTIVITIES', '200'),
|
||||
('MAX_CONCURRENT_LOCAL_ACTIVITIES', '200'),
|
||||
('MAX_CACHED_WORKFLOWS', '200'),
|
||||
('WORKFLOW_POLLER_BEHAVIOUR_MINIMUM', '10'),
|
||||
('WORKFLOW_POLLER_BEHAVIOUR_INITIAL', '100'),
|
||||
('WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM', '200'),
|
||||
('ACTIVITY_POLLER_BEHAVIOUR_MINIMUM', '10'),
|
||||
('ACTIVITY_POLLER_BEHAVIOUR_INITIAL', '100'),
|
||||
('ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM', '200'),
|
||||
('ACTIVITY_EXECUTOR_MAX_WORKERS', '200'),
|
||||
]
|
||||
|
||||
|
||||
def camel_to_snake(text: str) -> str:
|
||||
"""
|
||||
Convert a CamelCase or camelCase string into snake_case.
|
||||
|
||||
Args:
|
||||
- text: str, original string in CamelCase or camelCase format
|
||||
|
||||
Return:
|
||||
str: converted string in snake_case format
|
||||
"""
|
||||
text = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', text)
|
||||
text = re.sub('([a-z0-9])([A-Z])', r'\1_\2', text)
|
||||
return text.lower()
|
||||
|
||||
|
||||
def build_queue_name(workflow_name: str, runtime: str | None = None) -> str:
|
||||
"""
|
||||
Build Temporal queue name from workflow name and runtime.
|
||||
|
||||
Args:
|
||||
- workflow_name: str, workflow class name in CamelCase format
|
||||
- runtime: str | None, runtime suffix for environment-specific queues
|
||||
|
||||
Return:
|
||||
str: queue name in the format <workflow>-<runtime>-queue or <workflow>-queue
|
||||
"""
|
||||
snake_workflow_name = camel_to_snake(workflow_name)
|
||||
if runtime:
|
||||
return f'{snake_workflow_name}-{runtime}-queue'
|
||||
return f'{snake_workflow_name}-queue'
|
||||
|
||||
|
||||
def prepare_worker(
|
||||
main_workflow: type,
|
||||
other_workflows: Sequence[type],
|
||||
activities: Sequence[Any],
|
||||
temporal_client: Client,
|
||||
logger: Logger,
|
||||
runtime: str | None = None,
|
||||
) -> Worker:
|
||||
"""
|
||||
Build and configure a Temporal worker for the given workflow and activities.
|
||||
|
||||
Args:
|
||||
- main_workflow: type, main workflow class used as worker entry point
|
||||
- other_workflows: Sequence[type], additional workflows in the same worker
|
||||
- activities: Sequence[Any], activity callables registered in this worker
|
||||
- temporal_client: Client, Temporal client used by the worker
|
||||
- logger: Logger, logger instance used during worker preparation
|
||||
- runtime: str | None, runtime suffix appended to queue name when present
|
||||
|
||||
Return:
|
||||
Worker: fully configured Temporal worker instance ready to run
|
||||
"""
|
||||
main_workflow_name = main_workflow.__name__.upper()
|
||||
queue_name = build_queue_name(main_workflow.__name__, runtime)
|
||||
|
||||
local_workflow_parameters: dict[str, int] = {}
|
||||
for parameter_name, default_value in parameters:
|
||||
local_workflow_parameters[parameter_name] = int(
|
||||
os.getenv(f'{main_workflow_name}_{parameter_name}', default_value)
|
||||
)
|
||||
|
||||
logger.info(f'Preparing worker for {main_workflow_name} with queue {queue_name}')
|
||||
|
||||
activity_executor = ThreadPoolExecutor(
|
||||
max_workers=local_workflow_parameters['ACTIVITY_EXECUTOR_MAX_WORKERS'],
|
||||
thread_name_prefix=f'{queue_name}-activity',
|
||||
)
|
||||
|
||||
return Worker(
|
||||
temporal_client,
|
||||
task_queue=queue_name,
|
||||
workflows=[main_workflow, *other_workflows],
|
||||
activities=[*activities],
|
||||
activity_executor=activity_executor,
|
||||
max_concurrent_workflow_tasks=local_workflow_parameters['MAX_CONCURRENT_WORKFLOW_TASKS'],
|
||||
max_concurrent_activities=local_workflow_parameters['MAX_CONCURRENT_ACTIVITIES'],
|
||||
max_concurrent_local_activities=local_workflow_parameters[
|
||||
'MAX_CONCURRENT_LOCAL_ACTIVITIES'
|
||||
],
|
||||
max_cached_workflows=local_workflow_parameters['MAX_CACHED_WORKFLOWS'],
|
||||
workflow_task_poller_behavior=PollerBehaviorAutoscaling(
|
||||
minimum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MINIMUM'],
|
||||
initial=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_INITIAL'],
|
||||
maximum=local_workflow_parameters['WORKFLOW_POLLER_BEHAVIOUR_MAXIMUM'],
|
||||
),
|
||||
activity_task_poller_behavior=PollerBehaviorAutoscaling(
|
||||
minimum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MINIMUM'],
|
||||
initial=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_INITIAL'],
|
||||
maximum=local_workflow_parameters['ACTIVITY_POLLER_BEHAVIOUR_MAXIMUM'],
|
||||
),
|
||||
)
|
||||
259
model_manager/worker/worker.py
Normal file
259
model_manager/worker/worker.py
Normal file
@@ -0,0 +1,259 @@
|
||||
"""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-<runtime>-queue: For ML model training workflows
|
||||
- cleanup_files-<runtime>-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)
|
||||
- RUNTIME: Runtime identifier used in queue naming (default: single)
|
||||
- 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.runtime_paths import ensure_runtime_directories
|
||||
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'))
|
||||
|
||||
|
||||
def _get_runtime(runtime: str | None) -> str:
|
||||
"""
|
||||
Resolve runtime using fallback when missing.
|
||||
|
||||
Args:
|
||||
- runtime: str | None, runtime value from environment
|
||||
|
||||
Return:
|
||||
str: normalized runtime value
|
||||
"""
|
||||
normalized_runtime = runtime.strip() if runtime else ''
|
||||
return normalized_runtime or 'single'
|
||||
|
||||
|
||||
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
|
||||
"""
|
||||
|
||||
runtime = _get_runtime(RUNTIME)
|
||||
|
||||
ensure_runtime_directories()
|
||||
|
||||
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)
|
||||
|
||||
namespace = os.getenv('TEMPORAL_NAMESPACE', 'model-manager')
|
||||
temporal_client = await client.Client.connect(
|
||||
target_host=host,
|
||||
namespace=namespace,
|
||||
runtime=new_runtime,
|
||||
tls=use_tls,
|
||||
)
|
||||
|
||||
logger.custom_info(f'Temporal client initialized at {host}/{namespace}', 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())
|
||||
Reference in New Issue
Block a user