dfbdd4b605
The probes only check app_up, which was set to 1 when the metrics server started, before install_runtime. A pod that would fail the guard could turn Ready first, so a rollout would remove the old pod. app_up now goes to 1 right after the guard. Also: copy sys.modules keys before iterating, raise the sientia_model floor to the runtime's >=0.14.0, pin the same runtime versions in requirements-local.txt, test that the snapshot is taken before the install and the real importlib.metadata path, and document the boot failure in README and CLAUDE.md.
307 lines
11 KiB
Python
307 lines
11 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-<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 importlib.metadata
|
|
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'
|
|
|
|
|
|
def _loaded_distribution_versions() -> dict[str, str]:
|
|
"""Version on disk of every installed distribution whose top-level module is imported."""
|
|
# list() copies the keys: another thread importing while we iterate would raise RuntimeError.
|
|
loaded = {name.partition('.')[0] for name in list(sys.modules)}
|
|
return {
|
|
dist: importlib.metadata.version(dist)
|
|
for module, dists in importlib.metadata.packages_distributions().items()
|
|
if module in loaded
|
|
for dist in dists
|
|
}
|
|
|
|
|
|
def _runtime_changed_loaded_packages(before: dict[str, str]) -> dict[str, str]:
|
|
"""
|
|
Packages that were imported before ``install_runtime`` and whose version changed on disk.
|
|
|
|
The process keeps the old code in memory, so a later lazy import mixes two versions
|
|
(SIENTIAPDE-2214: sklearn 1.9.1 in memory, 1.8.0 on disk, ``_cyutility`` import error).
|
|
"""
|
|
changed = {}
|
|
for dist, old in before.items():
|
|
try:
|
|
new = importlib.metadata.version(dist)
|
|
except importlib.metadata.PackageNotFoundError:
|
|
new = 'removed'
|
|
if new != old:
|
|
changed[dist] = f'{old} -> {new}'
|
|
return changed
|
|
|
|
|
|
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, then installs the runtime; exits if the install
|
|
changed the version of a package the worker already imported, and only then sets
|
|
``app_up`` (what the probes check) to 1
|
|
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,
|
|
)
|
|
|
|
loaded_before_install = _loaded_distribution_versions()
|
|
await plugin_store.install_runtime(runtime_name=runtime)
|
|
|
|
changed = _runtime_changed_loaded_packages(loaded_before_install)
|
|
if changed:
|
|
logger.custom_critical(
|
|
f'Runtime {runtime} changed packages already imported by the worker: {changed}. '
|
|
'Pin them in requirements.txt to the runtime versions.',
|
|
metadata,
|
|
)
|
|
metrics.APP_UP.labels(pod_id=POD_ID).set(0)
|
|
sys.exit(1)
|
|
|
|
# The probes only check app_up, so it goes to 1 after the guard: a pod that would fail it
|
|
# never turns Ready, and a rollout keeps the old pod.
|
|
metrics.APP_UP.labels(pod_id=POD_ID).set(1)
|
|
|
|
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.
|
|
``main`` sets ``app_up`` to 1 later, once the runtime install passed the version guard.
|
|
|
|
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)
|
|
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())
|