SIENTIAPDE-1255: Refactor worker to support only the train_model-queue and remove prediction workflows.

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,

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