SIENTIAPDE-1645: code snapshot (part 1)
This commit is contained in:
vitor-aignosi
0
model_manager/workflows/__init__.py
Normal file
0
model_manager/workflows/__init__.py
Normal file
64
model_manager/workflows/cleanup_files.py
Normal file
64
model_manager/workflows/cleanup_files.py
Normal file
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Cleanup workflow for removing local filesystem.
|
||||
|
||||
This module provides a Temporal cron workflow that runs daily to clean up
|
||||
temporary files and directories older than the configured retention period.
|
||||
"""
|
||||
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.runtime_paths import REPORTS_TEMP_DIR
|
||||
from model_manager.workflows.train_model import no_retry_policy
|
||||
|
||||
TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120'))
|
||||
POD_ID = os.getenv('POD_ID')
|
||||
|
||||
|
||||
@workflow.defn(name='cleanup_files')
|
||||
class CleanupFiles:
|
||||
"""
|
||||
Cleanup workflow for removing stale files.
|
||||
|
||||
This workflow cleans up:
|
||||
- Local temporary directories with timestamp suffixes
|
||||
|
||||
The workflow is designed to be simple and robust, with error handling
|
||||
delegated to the individual activities.
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any] | None = None) -> None:
|
||||
"""
|
||||
Execute the cleanup workflow.
|
||||
|
||||
This method orchestrates the cleanup of local directories
|
||||
in sequence. No exception handling is needed as activities handle their
|
||||
own errors and notifications.
|
||||
"""
|
||||
payload = input_data or {}
|
||||
temp_path = payload.get('temp_path') or REPORTS_TEMP_DIR
|
||||
|
||||
# Metadata for tracking
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'pod_id': POD_ID,
|
||||
'workflow_name': 'cleanup_files',
|
||||
}
|
||||
}
|
||||
|
||||
# Execute local directory cleanup
|
||||
await workflow.execute_activity_method(
|
||||
Activities.cleanup_temp_directories,
|
||||
{
|
||||
**metadata,
|
||||
'temp_path': temp_path,
|
||||
},
|
||||
retry_policy=no_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_LOCAL),
|
||||
)
|
||||
400
model_manager/workflows/train_model.py
Normal file
400
model_manager/workflows/train_model.py
Normal file
@@ -0,0 +1,400 @@
|
||||
"""
|
||||
Train Model Workflow for ML model training pipeline.
|
||||
|
||||
This workflow orchestrates the complete ML model training process, including:
|
||||
- Parameter validation and conversion
|
||||
- Data download from MinIO
|
||||
- Model training
|
||||
- Model saving to MLFlow
|
||||
- Experiment tracking and status updates
|
||||
"""
|
||||
|
||||
from temporalio import workflow
|
||||
|
||||
with workflow.unsafe.imports_passed_through():
|
||||
import os
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from temporalio.common import RetryPolicy
|
||||
from temporalio.exceptions import ApplicationError
|
||||
|
||||
from model_manager.activities.activities import Activities
|
||||
from model_manager.activities.experiment_tracking import UpdateType
|
||||
from model_manager.utils.models.experiment_status import ExperimentStatus
|
||||
|
||||
# Activity timeouts (seconds). Tune per environment (large uploads, long training).
|
||||
# Training uses no_retry_policy: extend TIMEOUT_TRAIN_MODEL instead of adding retries
|
||||
# to avoid duplicate MLflow side effects. Cleanup/delete uses network_retry_policy.
|
||||
TIMEOUT_VALIDATE_PARAMS = int(os.getenv('TIMEOUT_VALIDATE_PARAMS', '30'))
|
||||
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
|
||||
# Fast retry for transient network errors (MinIO operations)
|
||||
network_retry_policy = RetryPolicy(
|
||||
initial_interval=timedelta(seconds=1),
|
||||
maximum_interval=timedelta(seconds=10),
|
||||
backoff_coefficient=2.0,
|
||||
maximum_attempts=5,
|
||||
)
|
||||
|
||||
# No retry for training - data errors are permanent
|
||||
no_retry_policy = RetryPolicy(
|
||||
maximum_attempts=1,
|
||||
)
|
||||
|
||||
# Database retry with exponential backoff
|
||||
database_retry_policy = RetryPolicy(
|
||||
initial_interval=timedelta(seconds=2),
|
||||
maximum_interval=timedelta(seconds=20),
|
||||
backoff_coefficient=2.0,
|
||||
maximum_attempts=5,
|
||||
)
|
||||
|
||||
|
||||
@workflow.defn(name='train_model')
|
||||
class TrainModel:
|
||||
"""
|
||||
Complete ML model training workflow.
|
||||
|
||||
This workflow implements the full training pipeline from parameter validation
|
||||
through model training and saving to MLFlow. It provides comprehensive error
|
||||
handling with database status updates at each stage.
|
||||
|
||||
The workflow ensures:
|
||||
- Proper parameter validation before training starts
|
||||
- Status tracking in database for monitoring
|
||||
- Error handling with detailed error messages
|
||||
- Cleanup and proper resource management
|
||||
"""
|
||||
|
||||
@workflow.run
|
||||
async def run(self, input_data: dict[str, Any]) -> dict[str, Any] | None:
|
||||
"""
|
||||
Execute the complete model training workflow.
|
||||
|
||||
This method orchestrates all steps of the training pipeline:
|
||||
1. Validate and convert training parameters
|
||||
2. Download training data from MinIO
|
||||
3. Train the model
|
||||
4. Save model to MLFlow
|
||||
5. Update experiment tracking status
|
||||
|
||||
Args:
|
||||
input_data: Complete configuration for the training workflow
|
||||
All TrainModelParams fields at the same level:
|
||||
- experiment_run_id (int): Unique identifier for the experiment run (REQUIRED)
|
||||
- target_variable (str): Target variable to predict
|
||||
- variable_columns (list[str]): Feature columns
|
||||
- train_size (int): Training data percentage
|
||||
- ... (all other TrainModelParams fields)
|
||||
|
||||
Raises:
|
||||
ValueError: If experiment_run_id is missing or invalid
|
||||
"""
|
||||
workflow.logger.info(f'Starting train_model workflow for {input_data}')
|
||||
|
||||
try:
|
||||
experiment_run_id = self._validate_experiment_run_id(input_data)
|
||||
except ValueError as exc:
|
||||
# Prevent workflow-task retries on deterministic input contract violations.
|
||||
raise ApplicationError(str(exc), non_retryable=True) from exc
|
||||
input_data = {**input_data, 'experiment_run_id': experiment_run_id}
|
||||
|
||||
model_name = input_data.get('model_name')
|
||||
model_id = input_data.get('model_id')
|
||||
|
||||
metadata = {
|
||||
'metadata': {
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'workflow_name': 'train_model',
|
||||
'model_name': model_name,
|
||||
'model_id': model_id,
|
||||
}
|
||||
}
|
||||
|
||||
train_params = await self._validate_training_parameters(
|
||||
input_data, experiment_run_id, metadata
|
||||
)
|
||||
|
||||
training_succeeded = False
|
||||
train_result: dict[str, Any] | None = None
|
||||
|
||||
try:
|
||||
train_result = await self._train_model(
|
||||
train_params=train_params,
|
||||
experiment_run_id=experiment_run_id,
|
||||
metadata=metadata,
|
||||
)
|
||||
training_succeeded = True
|
||||
finally:
|
||||
try:
|
||||
if train_result is not None:
|
||||
await self._cleanup_resources(
|
||||
run_dir=train_result.get('run_dir'),
|
||||
metadata=metadata,
|
||||
)
|
||||
else:
|
||||
pass
|
||||
except Exception: # noqa: BLE001
|
||||
# If cleanup fails after training failed, there is nothing extra to log (DB not committed).
|
||||
if training_succeeded: # pragma: no branch
|
||||
workflow.logger.warning(
|
||||
'cleanup_resources failed after successful training; model and DB status '
|
||||
'are already committed. Temp files may remain until scheduled cleanup.',
|
||||
)
|
||||
return train_result
|
||||
|
||||
def _validate_experiment_run_id(self, input_data: dict[str, Any]) -> int:
|
||||
"""
|
||||
Validate experiment_run_id from input data.
|
||||
|
||||
This method ensures that experiment_run_id is present and valid.
|
||||
Without a valid experiment_run_id, we cannot update database status,
|
||||
so this validation must happen before any other operation.
|
||||
|
||||
Args:
|
||||
input_data: Input data dictionary containing experiment_run_id
|
||||
|
||||
Returns:
|
||||
int: Validated experiment_run_id
|
||||
|
||||
Raises:
|
||||
ValueError: If experiment_run_id is missing or not an integer
|
||||
"""
|
||||
experiment_run_id = input_data.get('experiment_run_id')
|
||||
|
||||
if experiment_run_id is None:
|
||||
raise ValueError('experiment_run_id is required but was not provided')
|
||||
|
||||
if isinstance(experiment_run_id, int):
|
||||
return experiment_run_id
|
||||
|
||||
if isinstance(experiment_run_id, str) and experiment_run_id.strip().isdigit():
|
||||
return int(experiment_run_id.strip())
|
||||
|
||||
raise ValueError(
|
||||
f'experiment_run_id must be an integer or numeric string, got {type(experiment_run_id).__name__}'
|
||||
)
|
||||
|
||||
async def _validate_training_parameters(
|
||||
self,
|
||||
input_data: dict[str, Any],
|
||||
experiment_run_id: int,
|
||||
metadata: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Validate and convert training parameters from dict to TrainModelParams.
|
||||
|
||||
This method calls the validate_train_params activity to convert and validate
|
||||
the input parameters. On success, updates DB status to ORCHESTRATOR_WAITING_PROC.
|
||||
On error, updates DB status to ORCHESTRATOR_VALIDATION_ERROR.
|
||||
|
||||
Args:
|
||||
input_data: Input data dictionary containing all training parameters
|
||||
experiment_run_id: Validated experiment run ID
|
||||
metadata: Workflow execution metadata
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Validated training parameters
|
||||
|
||||
Raises:
|
||||
Exception: If validation fails (after updating DB status)
|
||||
"""
|
||||
try:
|
||||
input_data = await workflow.execute_activity_method(
|
||||
Activities.load_model_metadata,
|
||||
{
|
||||
**input_data,
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=no_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
|
||||
)
|
||||
|
||||
train_params = await workflow.execute_activity_method(
|
||||
Activities.validate_train_params,
|
||||
{
|
||||
**input_data,
|
||||
**metadata,
|
||||
},
|
||||
retry_policy=no_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_VALIDATE_PARAMS),
|
||||
)
|
||||
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.STATUS,
|
||||
status=ExperimentStatus.ORCHESTRATOR_WAITING_PROC,
|
||||
)
|
||||
|
||||
return train_params
|
||||
except Exception as e:
|
||||
try:
|
||||
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=self._extract_error_message(e),
|
||||
)
|
||||
except Exception as secondary: # noqa: BLE001
|
||||
workflow.logger.warning(
|
||||
'Failed to persist ORCHESTRATOR_VALIDATION_ERROR to experiment_run: %s',
|
||||
secondary,
|
||||
)
|
||||
raise
|
||||
|
||||
async def _train_model(
|
||||
self,
|
||||
train_params: dict[str, Any],
|
||||
experiment_run_id: int,
|
||||
metadata: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Download file from MinIO and train model.
|
||||
|
||||
This method orchestrates the download and training steps using proper
|
||||
resource management with try/catch/finally. On success, updates DB status
|
||||
to TRAINING_SUCCESS. On error, updates DB status to TRAINING_ERROR.
|
||||
|
||||
Args:
|
||||
train_params: TrainModelParams object with training configuration
|
||||
experiment_run_id: Validated experiment run ID
|
||||
metadata: Workflow execution metadata
|
||||
|
||||
Returns:
|
||||
dict[str, Any]: Serializable training summary from the train_model activity
|
||||
|
||||
Raises:
|
||||
Exception: If download or training fails (after updating DB status)
|
||||
"""
|
||||
try:
|
||||
train_result = await workflow.execute_activity_method(
|
||||
Activities.train_model,
|
||||
{
|
||||
**metadata,
|
||||
'train_params': train_params,
|
||||
},
|
||||
retry_policy=no_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_TRAIN_MODEL),
|
||||
)
|
||||
|
||||
await self._update_experiment_run(
|
||||
metadata=metadata,
|
||||
experiment_run_id=experiment_run_id,
|
||||
update_type=UpdateType.MODEL_SAVED,
|
||||
status=ExperimentStatus.TRAINING_SUCCESS,
|
||||
run_name=train_result.get('run_name'),
|
||||
)
|
||||
|
||||
return train_result
|
||||
except Exception as e:
|
||||
try:
|
||||
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=self._extract_error_message(e),
|
||||
)
|
||||
except Exception as secondary: # noqa: BLE001
|
||||
workflow.logger.warning(
|
||||
'Failed to persist TRAINING_ERROR status to experiment_run: %s',
|
||||
secondary,
|
||||
)
|
||||
raise
|
||||
|
||||
async def _cleanup_resources(
|
||||
self,
|
||||
run_dir: str | None,
|
||||
metadata: dict[str, Any],
|
||||
) -> None:
|
||||
"""
|
||||
Cleanup resources.
|
||||
|
||||
This method removes the temporary run directory via activity.
|
||||
|
||||
Args:
|
||||
run_dir: Temporary directory to remove
|
||||
metadata: Workflow execution metadata
|
||||
"""
|
||||
if run_dir is None:
|
||||
return
|
||||
|
||||
await workflow.execute_activity_method(
|
||||
Activities.cleanup_resources,
|
||||
{
|
||||
**metadata,
|
||||
'run_dir': run_dir,
|
||||
},
|
||||
retry_policy=network_retry_policy,
|
||||
start_to_close_timeout=timedelta(seconds=TIMEOUT_DELETE_FILE),
|
||||
)
|
||||
|
||||
async def _update_experiment_run(
|
||||
self,
|
||||
metadata: dict[str, Any],
|
||||
experiment_run_id: int,
|
||||
update_type: UpdateType,
|
||||
status: ExperimentStatus,
|
||||
error_message: str | None = None,
|
||||
run_name: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Update experiment run status in the database.
|
||||
|
||||
This is a helper method to simplify calls to the update_experiment_run activity.
|
||||
It handles both success and error status updates.
|
||||
|
||||
Args:
|
||||
metadata: Workflow execution metadata
|
||||
experiment_run_id: Unique identifier for the experiment run
|
||||
update_type: Type of update (STATUS, STATUS_WITH_ERROR, or MODEL_SAVED)
|
||||
status: Status to set in the database
|
||||
error_message: Error message (required if update_type is STATUS_WITH_ERROR)
|
||||
run_name: MLFlow run name (required if update_type is MODEL_SAVED)
|
||||
"""
|
||||
update_input = {
|
||||
**metadata,
|
||||
'experiment_run_id': experiment_run_id,
|
||||
'update_type': update_type,
|
||||
'status': status,
|
||||
}
|
||||
|
||||
if error_message is not None:
|
||||
update_input['error_message'] = error_message
|
||||
|
||||
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,
|
||||
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)
|
||||
|
||||
cause = getattr(current, '__cause__', None)
|
||||
context = getattr(current, '__context__', None)
|
||||
current = cause if isinstance(cause, Exception) else context
|
||||
|
||||
if not message_parts:
|
||||
return repr(exc)
|
||||
|
||||
return ' | '.join(message_parts)
|
||||
Reference in New Issue
Block a user