SIENTIAPDE-1253: Enforce TrainModelParams object in Training activity and tests, removing dict conversion.
This commit is contained in:
@@ -71,7 +71,7 @@ class Training(BaseActivity):
|
||||
Required keys:
|
||||
- metadata (dict): Workflow execution metadata
|
||||
- uploaded_file (BytesIO): Training data file (already downloaded from MinIO)
|
||||
- train_params (dict): Training parameters (converted to TrainModelParams)
|
||||
- train_params (TrainModelParams): Training parameters object
|
||||
|
||||
Returns:
|
||||
dict: Training result with the following structure:
|
||||
@@ -86,15 +86,7 @@ class Training(BaseActivity):
|
||||
result = await train_model({
|
||||
'metadata': {'workflow_id': 'train-123', 'experiment_run_id': 456},
|
||||
'uploaded_file': BytesIO(csv_data),
|
||||
'train_params': {
|
||||
'experiment_run_id': 456,
|
||||
'target_variable': 'price',
|
||||
'variable_columns': ['feature1', 'feature2'],
|
||||
'train_size': 80,
|
||||
'shuffle': True,
|
||||
'use_scaler': True,
|
||||
# ... other TrainModelParams fields
|
||||
}
|
||||
'train_params': TrainModelParams(...) # Already converted object
|
||||
})
|
||||
# Returns: {'success': True, 'result': TrainModelResult(...), 'error_message': None}
|
||||
|
||||
@@ -103,21 +95,22 @@ class Training(BaseActivity):
|
||||
"""
|
||||
metadata = input_data.get('metadata', {})
|
||||
uploaded_file = input_data['uploaded_file']
|
||||
train_params_dict = input_data['train_params']
|
||||
train_params = input_data['train_params']
|
||||
|
||||
try:
|
||||
self.info(
|
||||
f'Starting model training for target: {train_params_dict.get("target_variable")}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Convert dict to TrainModelParams
|
||||
train_params = TrainModelParams.from_dict(train_params_dict)
|
||||
|
||||
# Validate uploaded_file is BytesIO
|
||||
if not isinstance(uploaded_file, BytesIO):
|
||||
raise ValueError(f'uploaded_file must be BytesIO, got {type(uploaded_file)}')
|
||||
|
||||
# Validate train_params is TrainModelParams
|
||||
if not isinstance(train_params, TrainModelParams):
|
||||
raise ValueError(f'train_params must be TrainModelParams, got {type(train_params)}')
|
||||
|
||||
self.info(
|
||||
f'Starting model training for target: {train_params.target_variable}',
|
||||
metadata,
|
||||
)
|
||||
|
||||
# Step 1: Train the model
|
||||
self.info('Training model with TrainingRepository', metadata)
|
||||
train_result = self.training_repository.train(uploaded_file, train_params)
|
||||
@@ -141,7 +134,12 @@ class Training(BaseActivity):
|
||||
}
|
||||
|
||||
except Exception as e: # noqa: BLE001
|
||||
error_msg = f'Error training model - Target: {train_params_dict.get("target_variable", "unknown")}, Error: {str(e)}'
|
||||
target = (
|
||||
train_params.target_variable
|
||||
if hasattr(train_params, 'target_variable')
|
||||
else 'unknown'
|
||||
)
|
||||
error_msg = f'Error training model - Target: {target}, Error: {str(e)}'
|
||||
trace = traceback.format_exc()
|
||||
|
||||
# Send notification (MongoDB)
|
||||
|
||||
Reference in New Issue
Block a user