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,