gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

SIENTIAPDE-1241: refactor train_model workflow due to I/O errors.

Browse Source
This commit is contained in:
Bruno Domingues
2025-10-22 15:37:56 -03:00
parent f2a1c88ff3
commit 5789a13023
31 changed files with 37878 additions and 6097 deletions
Download Patch File Download Diff File Expand all files Collapse all files

330
model_manager/workflows/train_model.py
View File

@@ -20,18 +20,15 @@ with workflow.unsafe.imports_passed_through():
from model_manager.activities.activities import Activities
from model_manager.activities.experiment_tracking import UpdateType
from model_manager.utils.exceptions import ModelTrainingError
from model_manager.utils.models.experiment_status import ExperimentStatus
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.models.train_model_result import TrainModelResult
# Activity Timeouts (in seconds) - Configurable via environment variables
# Defaults are designed to handle large files (up to 200MB)
TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
TIMEOUT_DOWNLOAD_FILE = int(os.getenv('TIMEOUT_DOWNLOAD_FILE', '600'))
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '1800'))
TIMEOUT_SAVE_MODEL = int(os.getenv('TIMEOUT_SAVE_MODEL', '300'))
TIMEOUT_CLEANUP_DIRECTORY = int(os.getenv('TIMEOUT_CLEANUP_DIRECTORY', '60'))
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '60'))
TIMEOUT_TRAIN_MODEL = int(os.getenv('TIMEOUT_TRAIN_MODEL', '2700'))
TIMEOUT_DELETE_FILE = int(os.getenv('TIMEOUT_DELETE_FILE', '120'))
TIMEOUT_UPDATE_DATABASE = int(os.getenv('TIMEOUT_UPDATE_DATABASE', '30'))
# Retry Policies - Granular strategies for different operation types
@@ -48,14 +45,6 @@ with workflow.unsafe.imports_passed_through():
maximum_attempts=1,
)
# Moderate retry with backoff for MLFlow operations
mlflow_retry_policy = RetryPolicy(
initial_interval=timedelta(seconds=5),
maximum_interval=timedelta(seconds=30),
backoff_coefficient=2.0,
maximum_attempts=3,
)
# Database retry with exponential backoff
database_retry_policy = RetryPolicy(
initial_interval=timedelta(seconds=2),
@@ -64,14 +53,6 @@ with workflow.unsafe.imports_passed_through():
maximum_attempts=5,
)
# Filesystem retry for cleanup operations
filesystem_retry_policy = RetryPolicy(
initial_interval=timedelta(seconds=2),
maximum_interval=timedelta(seconds=10),
backoff_coefficient=1.5,
maximum_attempts=3,
)
@workflow.defn(name='train_model')
class TrainModel:
@@ -113,8 +94,6 @@ class TrainModel:
Raises:
ValueError: If experiment_run_id is missing or invalid
"""
# CRITICAL: Validate experiment_run_id first
# Without it, we cannot update database status, so fail immediately
experiment_run_id = self._validate_experiment_run_id(input_data)
metadata = {
@@ -124,33 +103,21 @@ class TrainModel:
}
}
# Step 1: Validate and convert training parameters
train_params = await self._validate_training_parameters(
input_data, experiment_run_id, metadata
)
# Step 2 & 3: Download from MinIO and Train model
# The _download_and_train_model method handles both steps:
# - Downloads file from MinIO (returns BytesIO)
# - Trains model with the downloaded file
# Any error (download OR training) = TRAINING_ERROR
train_result = await self._download_and_train_model(
train_result = await self._train_model(
train_params=train_params,
experiment_run_id=experiment_run_id,
metadata=metadata,
)
# Step 4: Save model to MLFlow
saved_result = await self._save_model_to_mlflow(
train_result=train_result,
experiment_run_id=experiment_run_id,
metadata=metadata,
)
# Step 5: Cleanup resources and delete file from MinIO
await self._cleanup_resources(
saved_result=saved_result,
experiment_run_id=experiment_run_id,
run_dir=(train_result.get('run_dir') or ''),
bucket_name=train_params.bucket_name,
file_name=train_params.file_name,
metadata=metadata,
)
@@ -174,16 +141,12 @@ class TrainModel:
experiment_run_id = input_data.get('experiment_run_id')
if experiment_run_id is None:
error_msg = 'experiment_run_id is required but was not provided'
workflow.logger.error(error_msg)
raise ValueError(error_msg)
raise ValueError('experiment_run_id is required but was not provided')
if not isinstance(experiment_run_id, int):
error_msg = (
raise ValueError(
f'experiment_run_id must be an integer, got {type(experiment_run_id).__name__}'
)
workflow.logger.error(error_msg)
raise ValueError(error_msg)
return experiment_run_id
@@ -211,21 +174,17 @@ class TrainModel:
Raises:
Exception: If validation fails (after updating DB status)
"""
validation_input = {
**metadata,
**input_data, # All training params at same level
}
try:
# Execute validation activity
train_params = await workflow.execute_activity_method(
Activities.validate_train_params,
validation_input,
retry_policy=no_retry_policy, # Validation errors are permanent
{
**metadata,
**input_data,
},
retry_policy=no_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
)
# Validation succeeded: Update status
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
@@ -234,41 +193,23 @@ class TrainModel:
)
return train_params
except Exception as e:
# Validation failed: Update status with error
error_message = str(e)
error_type = type(e).__name__
# Log with rich context for debugging
workflow.logger.error(
f'[VALIDATION_ERROR] Experiment {experiment_run_id} validation failed',
extra={
'step': 'validate_training_parameters',
'experiment_run_id': experiment_run_id,
'error_type': error_type,
'error_message': error_message,
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
},
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR,
error_message=error_message,
error_message=self._extract_error_message(e),
)
# Re-raise exception to stop workflow
raise
async def _download_and_train_model(
async def _train_model(
self,
train_params: TrainModelParams,
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelResult:
) -> dict[str, str | None]:
"""
Download file from MinIO and train model.
@@ -287,167 +228,53 @@ class TrainModel:
Raises:
Exception: If download or training fails (after updating DB status)
"""
uploaded_file = None
try:
# Step 1: Download file from MinIO
download_input = {
**metadata,
'bucket_name': train_params.bucket_name,
'file_name': train_params.file_name,
}
uploaded_file = await workflow.execute_activity_method(
Activities.fetch_file_from_minio,
download_input,
retry_policy=network_retry_policy, # Fast retry for network issues
start_to_close_timeout=timedelta(seconds=TIMEOUT_DOWNLOAD_FILE),
)
# Step 2: Train model with downloaded file
train_input = {
**metadata,
'uploaded_file': uploaded_file,
'train_params': train_params,
}
train_result = await workflow.execute_activity_method(
Activities.train_model,
train_input,
retry_policy=no_retry_policy, # Training errors are permanent (bad data)
{
**metadata,
'train_params': train_params,
},
retry_policy=no_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
)
# Training succeeded: Update status
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS,
status=ExperimentStatus.TRAINING_SUCCESS,
)
return train_result
except Exception as e:
# Download or training failed: Update status with error
error_message = str(e)
error_type = type(e).__name__
# Log with rich context for debugging
workflow.logger.error(
f'[TRAINING_ERROR] Experiment {experiment_run_id} training failed',
extra={
'step': 'download_and_train_model',
'experiment_run_id': experiment_run_id,
'experiment_name': train_params.experiment_name,
'bucket_name': train_params.bucket_name,
'file_name': train_params.file_name,
'error_type': error_type,
'error_message': error_message,
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
},
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.TRAINING_ERROR,
error_message=error_message,
)
# Re-raise exception to stop workflow
raise
finally:
# Ensure BytesIO is closed
if uploaded_file and hasattr(uploaded_file, 'close'):
uploaded_file.close()
async def _save_model_to_mlflow(
self,
train_result: TrainModelResult,
experiment_run_id: int,
metadata: dict[str, Any],
) -> TrainModelResult:
"""
Save trained model to MLFlow.
This method calls the save_model activity to save the trained model and
its artifacts to MLFlow. On success, updates DB status to MLFLOW_SENT
with MODEL_SAVED type and run_name. On error, updates DB status to
MLFLOW_SEND_ERROR.
Args:
train_result: TrainModelResult from training step
experiment_run_id: Validated experiment run ID
metadata: Workflow execution metadata
Returns:
TrainModelResult: Updated training result with MLFlow run name
Raises:
Exception: If model saving fails (after updating DB status)
"""
try:
# Save model to MLFlow
save_input = {
**metadata,
'train_result': train_result,
}
saved_result = await workflow.execute_activity_method(
Activities.save_model,
save_input,
retry_policy=mlflow_retry_policy, # Retry MLFlow with backoff
start_to_close_timeout=timedelta(seconds=TIMEOUT_SAVE_MODEL),
)
# Model saved successfully: Update status to MLFLOW_SENT with run_name
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.MLFLOW_SENT,
run_name=saved_result.run_name,
run_name=train_result.get('run_name'),
)
return saved_result
return train_result
except Exception as e:
# Model saving failed: Update status with error
error_message = str(e)
error_type = type(e).__name__
# Mapear flags -> status
# False/False: erro no treino
# True/False: erro ao salvar (MLflow)
# False/True: estado inconsistente, tratar como erro de treino
# True/True: não deveria cair aqui; tratar como erro genérico de treino
status = ExperimentStatus.TRAINING_ERROR
# Log with rich context for debugging
workflow.logger.error(
f'[MLFLOW_ERROR] Experiment {experiment_run_id} model save failed',
extra={
'step': 'save_model_to_mlflow',
'experiment_run_id': experiment_run_id,
'experiment_name': train_result.params.experiment_name,
'run_dir': train_result.run_dir if hasattr(train_result, 'run_dir') else None,
'error_type': error_type,
'error_message': error_message,
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
},
)
if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
status = ExperimentStatus.MLFLOW_SEND_ERROR
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.MLFLOW_SEND_ERROR,
error_message=error_message,
status=status,
error_message=self._extract_error_message(e),
)
# Re-raise exception to stop workflow
raise
async def _cleanup_resources(
self,
saved_result: TrainModelResult,
experiment_run_id: int,
run_dir: str,
bucket_name: str,
file_name: str,
metadata: dict[str, Any],
) -> None:
"""
@@ -466,75 +293,33 @@ class TrainModel:
Exception: If cleanup fails (after updating DB status)
"""
try:
# Step 1: Remove temporary run directory via activity (deterministic)
if hasattr(saved_result, 'run_dir') and saved_result.run_dir:
cleanup_input = {
**metadata,
'run_dir': saved_result.run_dir,
}
await workflow.execute_activity_method(
Activities.cleanup_run_directory,
cleanup_input,
retry_policy=filesystem_retry_policy, # Retry filesystem operations
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_DIRECTORY),
)
workflow.logger.info(
f'Run directory cleanup completed for experiment {experiment_run_id}'
)
# Step 2: Delete file from MinIO
delete_input = {
**metadata,
'bucket_name': saved_result.params.bucket_name,
'file_name': saved_result.params.file_name,
}
await workflow.execute_activity_method(
Activities.delete_file_from_minio,
delete_input,
retry_policy=network_retry_policy, # Fast retry for network issues
Activities.cleanup_resources,
{
**metadata,
'run_dir': run_dir,
'bucket_name': bucket_name,
'file_name': file_name,
},
retry_policy=network_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
)
# Cleanup succeeded: Update status to FILE_DELETED
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS,
status=ExperimentStatus.FILE_DELETED,
)
except Exception as e:
# Cleanup failed: Update status with error
error_message = str(e)
error_type = type(e).__name__
# Log with rich context for debugging
workflow.logger.error(
f'[CLEANUP_ERROR] Experiment {experiment_run_id} cleanup failed',
extra={
'step': 'cleanup_resources',
'experiment_run_id': experiment_run_id,
'bucket_name': saved_result.params.bucket_name,
'file_name': saved_result.params.file_name,
'run_dir': saved_result.run_dir if hasattr(saved_result, 'run_dir') else None,
'error_type': error_type,
'error_message': error_message,
'workflow_id': metadata.get('metadata', {}).get('workflow_id'),
},
)
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=ExperimentStatus.FILE_DELETE_ERROR,
error_message=error_message,
error_message=self._extract_error_message(e),
)
# Re-raise exception to stop workflow
raise
async def _update_experiment_run(
@@ -567,17 +352,34 @@ class TrainModel:
'status': status,
}
# Add error_message only if provided
if error_message is not None:
update_input['error_message'] = error_message
# Add run_name only if provided
if run_name is not None:
update_input['run_name'] = run_name
await workflow.execute_activity_method(
Activities.update_experiment_run,
update_input,
retry_policy=database_retry_policy, # Retry DB with exponential backoff
retry_policy=database_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_UPDATE_DATABASE),
)
def _extract_error_message(self, exc: Exception) -> str:
message_parts: list[str] = []
seen: set[int] = set()
current: Exception | None = exc
while current and id(current) not in seen:
seen.add(id(current))
text = str(current).strip()
if text and text not in message_parts:
message_parts.append(text)
current = getattr(current, 'cause', None)
if not message_parts:
return repr(exc)
return ' | '.join(message_parts)