From 7abd95180694fd9c1e529b3fde24f2dfe7712e92 Mon Sep 17 00:00:00 2001 From: Bruno Domingues Date: Thu, 16 Oct 2025 18:15:52 -0300 Subject: [PATCH] SIENTIAPDE-1255: Refactor worker to support only the train_model-queue and remove prediction workflows. --- model_manager/worker/worker.py | 64 ++++++++++------------------------ tests/worker/test_worker.py | 21 +++++------ 2 files changed, 26 insertions(+), 59 deletions(-) diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index 7a8b771..a9cf7db 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -1,20 +1,17 @@ -""" -Model Manager Worker Module +"""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 -prediction and retraining workflows. +model training workflows. -The worker supports two main task queues: -- predictions_batch-queue: Handles batch prediction workflows -- minimal_retrain-queue: Handles model retraining workflows +The worker supports the train_model-queue task queue for ML model training 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 +- ML model training pipeline orchestration Environment Variables: - TEMPORAL_HOST: Temporal server address (default: localhost:7233) @@ -46,12 +43,7 @@ with workflow.unsafe.imports_passed_through(): build_mongodb_config, build_postgres_config, ) - from model_manager.workflows.minimal_retrain import MinimalRetrain - from model_manager.workflows.predictions_batch import PredictionsBatch - from model_manager.workflows.sub_workflows.format_and_export_prediction import ( - FormatAndExportPrediction, - ) - from model_manager.workflows.sub_workflows.prediction_process import PredictionProcess + from model_manager.workflows.train_model import TrainModel POD_ID = os.getenv('POD_ID') SDK_METRICS_PORT = int(os.getenv('HTTP_SDK_METRICS_PORT', '9091')) @@ -133,41 +125,21 @@ async def main(): workers = [ Worker( temporal_client, - task_queue='minimal_retrain-queue', - workflows=[MinimalRetrain], - activities=[ - activities.load_custom_query, - activities.retrain_model, - activities.update_production_model, - 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], + task_queue='train_model-queue', + workflows=[TrainModel], activities=[ + # Training & Validation + activities.validate_train_params, + activities.train_model, # MLFlow - activities.request_predict, - activities.request_transform, - # Gates - activities.input_gate, - activities.mlflow_response_gate, - activities.mlflow_content_gate, - activities.format_prediction, - activities.format_default_prediction, - activities.get_last_timestamp, - # Postgres - activities.load_custom_query, - activities.repeat_last_prediction, - activities.export_data_to_postgres, - activities.write_metrics, + activities.save_model, + # MinIO + activities.fetch_file_from_minio, + activities.delete_file_from_minio, + # Filesystem + activities.cleanup_run_directory, + # Database + activities.update_experiment_run, ], max_concurrent_workflow_tasks=50, max_concurrent_activities=50, diff --git a/tests/worker/test_worker.py b/tests/worker/test_worker.py index 0c43222..cefbdd9 100644 --- a/tests/worker/test_worker.py +++ b/tests/worker/test_worker.py @@ -125,7 +125,7 @@ async def test_main_success( namespace='test-namespace', runtime=ANY, ) - assert mock_worker.call_count == 2 # Two workers created + assert mock_worker.call_count == 1 # Only one worker created mock_gather.assert_called_once() mock_sys_exit.assert_called_once_with(1) @@ -201,7 +201,7 @@ async def test_main_exception_handling( @patch('model_manager.worker.worker.NotificationHandler') @patch('model_manager.worker.worker.start_prometheus_server') @patch('model_manager.worker.worker.get_logger') -async def test_main_creates_two_workers( +async def test_main_creates_only_one_worker( mock_get_logger, mock_start_prometheus, mock_notification_handler, @@ -213,7 +213,7 @@ async def test_main_creates_two_workers( mock_sys_exit, mock_env_vars, ): - """Test that main creates two workers with correct configurations.""" + """Test that main creates only one worker with correct configurations.""" # Arrange mock_logger = MagicMock() mock_get_logger.return_value = mock_logger @@ -240,18 +240,13 @@ async def test_main_creates_two_workers( # Act await main() - # Assert - Verify two workers were created - assert mock_worker.call_count == 2 + # Assert - Verify only one worker was created + assert mock_worker.call_count == 1 - # Verify first worker (minimal_retrain-queue) + # Verify worker (train_model-queue) first_call = mock_worker.call_args_list[0] - assert first_call[1]['task_queue'] == 'minimal_retrain-queue' - assert 'MinimalRetrain' in str(first_call[1]['workflows']) - - # Verify second worker (predictions_batch-queue) - second_call = mock_worker.call_args_list[1] - assert second_call[1]['task_queue'] == 'predictions_batch-queue' - assert 'PredictionsBatch' in str(second_call[1]['workflows']) + assert first_call[1]['task_queue'] == 'train_model-queue' + assert 'TrainModel' in str(first_call[1]['workflows']) @pytest.mark.asyncio