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

SIENTIAPDE-1717: Remove MinIO cleanup functionality and associated components. This change streamlines the cleanup workflow to focus solely on local temporary directories, removes the ModelTrainingError exception, and updates related configurations, documentation, and tests.

Browse Source
This commit is contained in:
Bruno Domingues
2026-03-30 14:14:10 -03:00
parent 63a94dae0a
commit 7a0961f29d
20 changed files with 56 additions and 778 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
.env.example
View File

@@ -44,10 +44,7 @@ TIMEOUT_UPDATE_DATABASE=30
CLEANUP_RETENTION_HOURS=24
CLEANUP_DRY_RUN=false
TIMEOUT_CLEANUP_MINIO=300
TIMEOUT_CLEANUP_LOCAL=120
MAX_KEYS_CLEANUP=1000
DEFAULT_CLEANUP_BUCKET=model-training
CLEANUP_SCHEDULE_ID=cleanup-files-daily
CLEANUP_CRON="0 0 * * *"

14
README.md
View File

@@ -188,7 +188,6 @@ The Model Manager system uses a Temporal-based workflow architecture with clear
- **Polynomial Regression**: Support for configurable degree and interaction terms
- **Training Predictions**: Calculates y_train_pred before denormalization for accurate metrics
- **Cleanup**: File and directory cleanup operations
- `cleanup_minio_files()`: Removes stale files from MinIO based on timestamp prefixes
- `cleanup_temp_directories()`: Cleans local temporary directories
- Configurable retention period (default: 24 hours)
- Dry-run mode for testing
@@ -404,16 +403,13 @@ The cleanup schedule is automatically created when the worker starts:
| **Execution Timeout** | `CLEANUP_EXECUTION_TIMEOUT_HOURS` | `1` | Maximum execution time (hours) |
| **Retention Period** | `CLEANUP_RETENTION_HOURS` | `24` | Files older than this are deleted |
| **Dry Run** | `CLEANUP_DRY_RUN` | `false` | Test mode without actual deletion |
| **Max Keys** | `MAX_KEYS_CLEANUP` | `1000` | MinIO list operation page size |
#### Architecture Diagram
```mermaid
flowchart TD
A[Scheduled Trigger] --> B[cleanup_minio_files]
B --> C[cleanup_temp_directories]
B -.-> MinIO[MinIO Storage]
C -.-> FS[Local Filesystem]
A[Scheduled Trigger] --> B[cleanup_temp_directories]
B -.-> FS[Local Filesystem]
```
#### Retry Strategies
@@ -1101,8 +1097,6 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa
| `CLEANUP_EXECUTION_TIMEOUT_HOURS` | Cleanup execution timeout | `1` | No |
| `CLEANUP_RETENTION_HOURS` | File retention period (hours) | `24` | No |
| `CLEANUP_DRY_RUN` | Dry-run mode (no actual deletion) | `false` | No |
| `MAX_KEYS_CLEANUP` | MinIO list operation page size | `1000` | No |
| `DEFAULT_CLEANUP_BUCKET` | Default bucket for cleanup | `model-training` | No |
| `LOG_LEVEL` | Application log level | `INFO` | No |
| `PROJECT_NAME` | Project name for metrics | `model-manager` | No |
| `HTTP_METRICS_PORT` | Prometheus metrics port | `9090` | No |
@@ -1127,7 +1121,6 @@ These timeouts control how long each activity in workflows can run before timing
| Variable | Description | Default | Calculation Basis |
|----------|-------------|---------|-------------------|
| `TIMEOUT_CLEANUP_MINIO` | MinIO cleanup timeout | `300` | Scan and delete multiple files (5 min) |
| `TIMEOUT_CLEANUP_LOCAL` | Local cleanup timeout | `120` | Scan and delete directories (2 min) |
**Note**: These timeouts can be adjusted based on your infrastructure performance and file sizes. If you're processing files larger than 200MB or have slower network/compute resources, increase these values accordingly.
@@ -1341,7 +1334,6 @@ export LOG_LEVEL=DEBUG
- **Connection Pools**: Optimize database connection pool sizes
- **Model Retention**: Configure MLFlow model retention based on requirements
- **Batch Sizes**: Adjust data processing batch sizes for optimal throughput
- **Cleanup Performance**: Tune `MAX_KEYS_CLEANUP` for MinIO list operation page size
### Scaling Considerations

1
model_manager/activities/activities.py
View File

@@ -111,7 +111,6 @@ class Activities(ExperimentTracking, Training, Cleanup):
Cleanup.__init__(
self,
storage_repository=self.storage_repository,
logger=logger,
notification_handler=notification_handler,
metrics_controller=metrics_controller,

122
model_manager/activities/cleanup.py
View File

@@ -1,5 +1,5 @@
"""
Cleanup activities for removing stale files from MinIO and local filesystem.
Cleanup activities for removing stale files from local filesystem.
This module provides activities for cleaning up temporary files and directories
that are older than the configured retention period. It operates independently
@@ -13,7 +13,7 @@ with workflow.unsafe.imports_passed_through():
import re
import shutil
import traceback
from datetime import UTC, datetime, timedelta
from datetime import datetime, timedelta
from typing import Any
from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
@@ -23,11 +23,9 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
from model_manager.utils.repository.storage_repository import StorageRepository
RETENTION_HOURS = int(os.getenv('CLEANUP_RETENTION_HOURS', '24'))
DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true'
MAX_KEYS_CLEANUP = int(os.getenv('MAX_KEYS_CLEANUP', '1000'))
class Cleanup(SientiaMonitoring):
@@ -35,13 +33,11 @@ class Cleanup(SientiaMonitoring):
Activity for cleaning up stale files and directories.
This activity extends SientiaMonitoring and handles cleanup of:
- MinIO files with timestamp prefixes (timestamp-filename pattern)
- Local temporary directories with timestamp suffixes
"""
def __init__(
self,
storage_repository: StorageRepository,
logger: Logger,
notification_handler: NotificationHandler,
metrics_controller: MetricsController,
@@ -50,135 +46,21 @@ class Cleanup(SientiaMonitoring):
Initialize Cleanup activity.
Args:
storage_repository: Repository for MinIO operations
logger: Logger instance for observability
notification_handler: Handler for sending notifications
metrics_controller: Controller for metrics emission
"""
SientiaMonitoring.__init__(self, logger, notification_handler, metrics_controller)
self.storage_repository = storage_repository
# Configuration from environment variables
self.retention_hours = RETENTION_HOURS
self.dry_run = DRY_RUN
# MinIO list operation page size
self.max_keys_cleanup = MAX_KEYS_CLEANUP
# Regex patterns for timestamp extraction
self.minio_timestamp_pattern = re.compile(r'^(\d{13})-(.+)') # timestamp-filename
self.dir_timestamp_pattern = re.compile(
r'^(.+)_(\d{8}_\d{6}_\d{6})$'
) # name_YYYYMMDD_HHMMSS_microseconds
@activity.defn(name='cleanup_minio_files')
async def cleanup_minio_files(self, input_data: dict[str, Any]) -> None:
"""
Clean up stale files from MinIO based on timestamp in filename.
This activity scans a single MinIO bucket for files following the pattern
'{timestamp}-{filename}' where timestamp is milliseconds since epoch.
Files older than the retention period are deleted.
Args:
input_data: Cleanup configuration containing:
- metadata (dict): Workflow execution metadata
- bucket_name (str): Name of the bucket to scan
Returns:
None: Results are logged and tracked via metrics
Raises:
Exception: If cleanup fails (after sending notification)
"""
metadata = input_data.get('metadata', {})
bucket_name = input_data.get('bucket_name')
metrics_status = 'success'
if not bucket_name:
raise ValueError('bucket_name must be provided')
cutoff_time = datetime.now(UTC) - timedelta(hours=self.retention_hours)
cutoff_timestamp_ms = int(cutoff_time.timestamp() * 1000)
try:
self.info(
f'Starting MinIO cleanup - Bucket: {bucket_name}, '
f'Retention: {self.retention_hours}h, Dry run: {self.dry_run}, '
f'Cutoff: {cutoff_time.isoformat()}',
metadata,
)
files_scanned = 0
files_deleted = 0
errors = []
# List objects in the specified bucket
max_keys = self.max_keys_cleanup # Use environment variable for page size
objects = self.storage_repository.list_bucket_objects(bucket_name, max_keys)
for obj_key in objects:
files_scanned += 1
# Extract timestamp from filename
match = self.minio_timestamp_pattern.match(obj_key)
if not match:
self.debug(f'Skipping file without timestamp pattern: {obj_key}', metadata)
continue
file_timestamp_ms = int(match.group(1))
if file_timestamp_ms < cutoff_timestamp_ms:
if self.dry_run:
self.info(
f'[DRY RUN] Would delete: {obj_key} (age: {(cutoff_time.timestamp() - file_timestamp_ms / 1000) / 3600:.1f}h)',
metadata,
)
files_deleted += 1
else:
try:
self.storage_repository.delete_file(bucket_name, obj_key)
self.info(f'Deleted stale file: {obj_key}', metadata)
files_deleted += 1
except OSError as e:
error_msg = f'Failed to delete {obj_key}: {str(e)}'
errors.append(error_msg)
self.error(error_msg, metadata)
else:
self.debug(
f'Keeping recent file: {obj_key} (age: {(cutoff_time.timestamp() - file_timestamp_ms / 1000) / 3600:.1f}h)',
metadata,
)
self.info(
f'MinIO cleanup completed - Bucket: {bucket_name}, '
f'Scanned: {files_scanned}, Deleted: {files_deleted}, Errors: {len(errors)}',
metadata,
)
except Exception as e:
metrics_status = 'error'
error_msg = f'Error in MinIO cleanup: {str(e)}'
trace = traceback.format_exc()
self.send_notification(
metadata=metadata,
notification_id='CLEANUP_MINIO_ERROR',
message=error_msg,
block='cleanup_minio_files',
level=NotificationLevel.ERROR,
attachment_content=trace,
)
raise
finally:
await self._emit_metrics(
metadata=metadata,
metrics_status=metrics_status,
activity_name='cleanup_minio_files',
emit_workflow_metric=(metrics_status == 'error'),
)
@activity.defn(name='cleanup_temp_directories')
async def cleanup_temp_directories(self, input_data: dict[str, Any]) -> None:
"""

26
model_manager/activities/training.py
View File

@@ -19,7 +19,6 @@ with workflow.unsafe.imports_passed_through():
from sientia_do.observability.sientia_monitoring import SientiaMonitoring
from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL
from model_manager.utils.exceptions import ModelTrainingError
from model_manager.utils.models.train_model_params import TrainModelParams
from model_manager.utils.repository.model_repository import ModelRepository
from model_manager.utils.repository.storage_repository import StorageRepository
@@ -143,8 +142,6 @@ class Training(SientiaMonitoring):
train_params = TrainModelParams.from_dict(train_params)
# type: ignore[assignment]
model_trained = False
model_saved = False
metrics_status = 'success'
try:
@@ -157,9 +154,7 @@ class Training(SientiaMonitoring):
train_params, train_result
)
model_trained = True
train_result = self.model_repository.save_model(train_result)
model_saved = True
return {
'run_name': train_result.run_name,
@@ -168,11 +163,7 @@ class Training(SientiaMonitoring):
except Exception as e: # noqa: BLE001
metrics_status = 'error'
error_msg = (
'Error training model - '
f'model_trained={model_trained}, model_saved={model_saved}, '
f'error: {str(e)}'
)
error_msg = f'Error training model - error: {str(e)}'
trace = traceback.format_exc()
@@ -185,10 +176,7 @@ class Training(SientiaMonitoring):
attachment_content=trace,
)
raise ModelTrainingError(
model_trained=model_trained,
model_saved=model_saved,
) from e
raise e
finally:
await self._emit_metrics(
metadata=metadata,
@@ -206,28 +194,20 @@ class Training(SientiaMonitoring):
input_data: Cleanup configuration containing:
- metadata (dict): Workflow execution metadata.
- run_dir (str): Temporary directory to remove.
- bucket_name (str): MinIO bucket of the uploaded file.
- file_name (str): MinIO object key to delete.
Raises:
Exception: If cleanup fails (after sending notification).
"""
metadata = input_data.get('metadata', {})
run_dir = input_data.get('run_dir', '')
bucket_name = input_data.get('bucket_name', '')
file_name = input_data.get('file_name', '')
metrics_status = 'success'
try:
self.model_repository.cleanup_run_directory(run_dir)
self.storage_repository.delete_file(bucket_name, file_name)
except Exception as e: # noqa: BLE001
metrics_status = 'error'
error_msg = (
f'Error cleaning up resources - Run directory: {run_dir}, '
f'File: {bucket_name}/{file_name}, Error: {str(e)}'
)
error_msg = f'Error cleaning up resources - Run directory: {run_dir}, Error: {str(e)}'
trace = traceback.format_exc()

38
model_manager/utils/exceptions.py
View File

@@ -1,38 +0,0 @@
"""
Custom exception types for the Model Manager.
This module defines domain-specific exceptions used across the training
workflow to convey additional context (e.g., flags indicating which steps
completed successfully) without altering control flow semantics.
"""
class ModelTrainingError(Exception):
"""
Exception raised when the model training workflow fails.
This exception carries flags indicating whether the model was trained
and/or saved successfully, enabling the workflow to map errors to
appropriate experiment statuses.
"""
def __init__(self, model_trained: bool, model_saved: bool, message: str | None = None):
"""
Initialize ModelTrainingError with training state flags.
Args:
model_trained: True if the training step completed successfully.
model_saved: True if the model saving step completed successfully.
message: Optional custom error message. If None, a default message
including the state flags is generated.
"""
self.model_trained = model_trained
self.model_saved = model_saved
if message is None:
message = (
'Model training workflow failed '
f'(model_trained={model_trained}, model_saved={model_saved})'
)
super().__init__(message)

8
model_manager/utils/models/experiment_status.py
View File

@@ -14,13 +14,10 @@ class ExperimentStatus(StrEnum):
to maintain compatibility with existing database records and monitoring systems.
Attributes:
ORCHESTRATOR_VALIDATION_ERROR: Error in the parameters validation.
ORCHESTRATOR_WAITING_PROC: Initial status indicating experiment is registered and waiting for processing.
ORCHESTRATOR_VALIDATION_ERROR: Error in the parameters validation.
TRAINING_SUCCESS: Training completed successfully with model and metrics calculated.
TRAINING_ERROR: Training failed due to data issues, model errors, or other exceptions.
TRACKING_SENT: Model successfully saved to MLFlow.
TRACKING_SEND_ERROR: Model saving to MLFlow failed due to connection or serialization errors.
FILE_DELETED: Cleanup completed successfully with all artifacts removed.
FILE_DELETE_ERROR: Cleanup failed due to file system or MinIO errors.
"""
@@ -28,7 +25,4 @@ class ExperimentStatus(StrEnum):
ORCHESTRATOR_WAITING_PROC = 'ORCHESTRATOR_WAITING_PROC'
TRAINING_SUCCESS = 'TRAINING_SUCCESS'
TRAINING_ERROR = 'TRAINING_ERROR'
TRACKING_SENT = 'TRACKING_SENT'
TRACKING_SEND_ERROR = 'TRACKING_SEND_ERROR'
FILE_DELETED = 'FILE_DELETED'
FILE_DELETE_ERROR = 'FILE_DELETE_ERROR'

11
model_manager/utils/repository/storage_repository.py
View File

@@ -119,17 +119,6 @@ class StorageRepository:
return BytesIO(file_content)
def delete_file(self, bucket_name: str, file_name: str) -> None:
"""
Remove an object from MinIO storage.
Args:
bucket_name: Bucket that contains the object.
file_name: Object key to delete.
"""
self.minio_client.delete_object(Bucket=bucket_name, Key=file_name)
self.logger.info(f'File deleted successfully: {bucket_name}/{file_name}')
def list_bucket_objects(self, bucket_name: str, max_keys: int = 1000) -> list[str]:
"""
List objects in a MinIO bucket.

1
model_manager/worker/worker.py
View File

@@ -155,7 +155,6 @@ async def main():
task_queue=CLEANUP_TASK_QUEUE,
workflows=[CleanupFiles],
activities=[
activities.cleanup_minio_files,
activities.cleanup_temp_directories,
],
max_concurrent_workflow_tasks=20,

27
model_manager/workflows/cleanup_files.py
View File

@@ -1,5 +1,5 @@
"""
Cleanup workflow for removing stale files from MinIO and local filesystem.
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.
@@ -13,11 +13,9 @@ with workflow.unsafe.imports_passed_through():
from typing import Any
from model_manager.activities.activities import Activities
from model_manager.workflows.train_model import POD_ID, network_retry_policy, no_retry_policy
from model_manager.workflows.train_model import POD_ID, no_retry_policy
TIMEOUT_CLEANUP_MINIO = int(os.getenv('TIMEOUT_CLEANUP_MINIO', '300'))
TIMEOUT_CLEANUP_LOCAL = int(os.getenv('TIMEOUT_CLEANUP_LOCAL', '120'))
DEFAULT_CLEANUP_BUCKET = os.getenv('DEFAULT_CLEANUP_BUCKET', 'model-training')
@workflow.defn(name='cleanup_files')
@@ -26,7 +24,6 @@ class CleanupFiles:
Cleanup workflow for removing stale files.
This workflow cleans up:
- MinIO files with timestamp prefixes
- Local temporary directories with timestamp suffixes
The workflow is designed to be simple and robust, with error handling
@@ -38,17 +35,10 @@ class CleanupFiles:
"""
Execute the cleanup workflow.
This method orchestrates the cleanup of MinIO files and local directories
This method orchestrates the cleanup of local directories
in sequence. No exception handling is needed as activities handle their
own errors and notifications.
Args:
input_data: Workflow configuration containing optional:
- bucket_name (str): Bucket to clean (defaults to environment variable)
"""
# Get bucket name from input or environment
bucket_name = input_data.get('bucket_name', DEFAULT_CLEANUP_BUCKET)
# Default temp path for local cleanup
temp_path = 'model_manager/reports/temp'
@@ -60,17 +50,6 @@ class CleanupFiles:
}
}
# Execute MinIO cleanup
await workflow.execute_activity_method(
Activities.cleanup_minio_files,
{
**metadata,
'bucket_name': bucket_name,
},
retry_policy=network_retry_policy,
start_to_close_timeout=timedelta(seconds=TIMEOUT_CLEANUP_MINIO),
)
# Execute local directory cleanup
await workflow.execute_activity_method(
Activities.cleanup_temp_directories,

71
model_manager/workflows/train_model.py
View File

@@ -20,7 +20,6 @@ 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
@@ -117,10 +116,7 @@ class TrainModel:
)
await self._cleanup_resources(
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,
)
@@ -246,27 +242,17 @@ class TrainModel:
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.TRACKING_SENT,
status=ExperimentStatus.TRAINING_SUCCESS,
run_name=train_result.get('run_name'),
)
return train_result
except Exception as e:
# 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
if isinstance(e, ModelTrainingError) and (e.model_trained and not e.model_saved):
status = ExperimentStatus.TRACKING_SEND_ERROR
await self._update_experiment_run(
metadata=metadata,
experiment_run_id=experiment_run_id,
update_type=UpdateType.STATUS_WITH_ERROR,
status=status,
status=ExperimentStatus.TRAINING_ERROR,
error_message=self._extract_error_message(e),
)
@@ -274,56 +260,27 @@ class TrainModel:
async def _cleanup_resources(
self,
experiment_run_id: int,
run_dir: str,
bucket_name: str,
file_name: str,
metadata: dict[str, Any],
) -> None:
"""
Cleanup resources and delete file from MinIO.
Cleanup resources.
This method removes the temporary run directory via activity and deletes
the training file from MinIO. On success, updates DB status to FILE_DELETED.
On error, updates DB status to FILE_DELETE_ERROR.
This method removes the temporary run directory via activity.
Args:
saved_result: TrainModelResult with run_dir and params information
experiment_run_id: Validated experiment run ID
run_dir: Temporary directory to remove
metadata: Workflow execution metadata
Raises:
Exception: If cleanup fails (after updating DB status)
"""
try:
await workflow.execute_activity_method(
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),
)
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:
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=self._extract_error_message(e),
)
raise
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,

10
scripts/run_cleanup_test.py
View File

@@ -45,14 +45,8 @@ async def main(argv: list[str]) -> None:
temporal_host = os.getenv('TEMPORAL_HOST')
temporal_namespace = os.getenv('TEMPORAL_NAMESPACE')
task_queue = os.getenv('CLEANUP_TASK_QUEUE')
default_bucket = os.getenv('DEFAULT_CLEANUP_BUCKET')
use_tls = os.getenv('TEMPORAL_USE_TLS', 'false').lower() == 'true'
# Optional CLI: bucket name override
bucket_name = default_bucket
if argv:
bucket_name = argv[0]
print(f'Connecting to Temporal at {temporal_host} (namespace={temporal_namespace})...')
client = await Client.connect(
target_host=temporal_host,
@@ -61,7 +55,6 @@ async def main(argv: list[str]) -> None:
)
input_data: dict[str, Any] = {
'bucket_name': bucket_name,
}
workflow_id = f'cleanup-files-manual-{int(asyncio.get_event_loop().time())}'
@@ -69,8 +62,7 @@ async def main(argv: list[str]) -> None:
print(
f'Starting cleanup_files workflow once...\n'
f' workflow_id = {workflow_id}\n'
f' task_queue = {task_queue}\n'
f' bucket_name = {bucket_name}'
f' task_queue = {task_queue}'
)
handle = await client.start_workflow(

202
tests/activities/test_cleanup.py
View File

@@ -4,7 +4,7 @@ import asyncio
import os
import shutil
import tempfile
from datetime import UTC, datetime, timedelta
from datetime import datetime, timedelta
from importlib import reload
from unittest.mock import AsyncMock, MagicMock, patch
@@ -34,15 +34,6 @@ def mock_metrics_controller():
return controller
@pytest.fixture
def mock_storage_repository():
"""Fixture for a mock storage repository."""
repo = MagicMock()
repo.delete_file = MagicMock()
repo.list_bucket_objects = MagicMock()
return repo
@pytest.fixture
def temp_dir():
"""Fixture to create and clean up a temporary directory."""
@@ -59,11 +50,9 @@ def temp_dir():
{
'CLEANUP_RETENTION_HOURS': '24',
'CLEANUP_DRY_RUN': 'false',
'MAX_KEYS_CLEANUP': '1000',
},
)
def test_cleanup_init_default_values(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -75,7 +64,6 @@ def test_cleanup_init_default_values(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -83,11 +71,9 @@ def test_cleanup_init_default_values(
assert cleanup.retention_hours == 24
assert cleanup.dry_run is False
assert cleanup.max_keys_cleanup == 1000
def test_cleanup_init_custom_env_values(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -98,7 +84,6 @@ def test_cleanup_init_custom_env_values(
{
'CLEANUP_RETENTION_HOURS': '48',
'CLEANUP_DRY_RUN': 'true',
'MAX_KEYS_CLEANUP': '500',
},
):
import model_manager.activities.cleanup
@@ -106,16 +91,14 @@ def test_cleanup_init_custom_env_values(
reload(model_manager.activities.cleanup)
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup = Cleanup(
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
assert cleanup.retention_hours == 48
assert cleanup.dry_run is True
assert cleanup.max_keys_cleanup == 500
assert cleanup.retention_hours == 48
assert cleanup.dry_run is True
@patch.dict(os.environ, {'CLEANUP_RETENTION_HOURS': 'invalid'})
@@ -127,162 +110,10 @@ def test_cleanup_init_invalid_env_value_raises_error():
reload(model_manager.activities.cleanup)
# --- MinIO Cleanup Tests ---
def test_cleanup_minio_files_missing_bucket_name(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test cleanup_minio_files raises ValueError if bucket_name is missing."""
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
with pytest.raises(ValueError, match='bucket_name must be provided'):
asyncio.run(cleanup.cleanup_minio_files({'metadata': {}}))
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
def test_cleanup_minio_files_success_with_deletions(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test successful deletion of old files from MinIO."""
import model_manager.activities.cleanup
reload(model_manager.activities.cleanup)
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
old_ts = int((datetime.now(UTC) - timedelta(hours=48)).timestamp() * 1000)
recent_ts = int((datetime.now(UTC) - timedelta(hours=1)).timestamp() * 1000)
mock_storage_repository.list_bucket_objects.return_value = [
f'{old_ts}-old-file.txt',
f'{recent_ts}-recent-file.txt',
'no-timestamp-file.txt',
]
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
mock_storage_repository.delete_file.assert_called_once_with(
'test-bucket', f'{old_ts}-old-file.txt'
)
cleanup._emit_metrics.assert_called_once()
def test_cleanup_minio_files_dry_run(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test MinIO cleanup in dry_run mode does not delete files."""
with patch.dict(os.environ, {'CLEANUP_DRY_RUN': 'true'}):
import model_manager.activities.cleanup
reload(model_manager.activities.cleanup)
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
old_ts = int((datetime.now(UTC) - timedelta(hours=48)).timestamp() * 1000)
mock_storage_repository.list_bucket_objects.return_value = [f'{old_ts}-old-file.txt']
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
mock_storage_repository.delete_file.assert_not_called()
cleanup._emit_metrics.assert_called_once()
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
def test_cleanup_minio_files_delete_error(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test error during MinIO file deletion is handled gracefully."""
import model_manager.activities.cleanup
reload(model_manager.activities.cleanup)
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup._emit_metrics = AsyncMock()
cleanup.error = MagicMock()
old_ts = int((datetime.now(UTC) - timedelta(hours=48)).timestamp() * 1000)
mock_storage_repository.list_bucket_objects.return_value = [f'{old_ts}-old-file.txt']
mock_storage_repository.delete_file.side_effect = OSError('Permission Denied')
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
cleanup.error.assert_called_once()
cleanup._emit_metrics.assert_called_once()
def test_cleanup_minio_files_exception_handling(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
):
"""Test exception during MinIO cleanup triggers notification and metrics."""
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
)
cleanup.send_notification = MagicMock()
cleanup._emit_metrics = AsyncMock()
mock_storage_repository.list_bucket_objects.side_effect = Exception('Connection Error')
with pytest.raises(Exception, match='Connection Error'):
asyncio.run(cleanup.cleanup_minio_files({'bucket_name': 'test-bucket', 'metadata': {}}))
cleanup.send_notification.assert_called_once()
cleanup._emit_metrics.assert_called_once()
# --- Temp Directory Cleanup Tests ---
def test_cleanup_temp_directories_nonexistent_path(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -291,7 +122,6 @@ def test_cleanup_temp_directories_nonexistent_path(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -310,7 +140,6 @@ def test_cleanup_temp_directories_nonexistent_path(
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
def test_cleanup_temp_directories_success_with_deletions(
temp_dir,
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -322,7 +151,6 @@ def test_cleanup_temp_directories_success_with_deletions(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -347,7 +175,6 @@ def test_cleanup_temp_directories_success_with_deletions(
@patch.dict(os.environ, {'CLEANUP_DRY_RUN': 'true'})
def test_cleanup_temp_directories_dry_run(
temp_dir,
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -359,7 +186,6 @@ def test_cleanup_temp_directories_dry_run(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -379,7 +205,6 @@ def test_cleanup_temp_directories_dry_run(
@patch.dict('model_manager.activities.cleanup.os.environ', {'CLEANUP_DRY_RUN': 'false'})
def test_cleanup_temp_directories_delete_error(
temp_dir,
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -391,7 +216,6 @@ def test_cleanup_temp_directories_delete_error(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -414,7 +238,6 @@ def test_cleanup_temp_directories_delete_error(
def test_emit_metrics(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -423,7 +246,6 @@ def test_emit_metrics(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -444,7 +266,6 @@ def test_emit_metrics(
def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
temp_dir,
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -456,7 +277,6 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -480,7 +300,6 @@ def test_cleanup_temp_directories_with_files_and_unmatched_dirs(
def test_cleanup_temp_directories_invalid_timestamp_format(
temp_dir,
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -492,7 +311,6 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -512,7 +330,6 @@ def test_cleanup_temp_directories_invalid_timestamp_format(
def test_cleanup_temp_directories_generic_exception(
temp_dir,
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -524,7 +341,6 @@ def test_cleanup_temp_directories_generic_exception(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,
@@ -541,7 +357,6 @@ def test_cleanup_temp_directories_generic_exception(
def test_emit_metrics_activity_only(
mock_storage_repository,
mock_logger,
mock_notification_handler,
mock_metrics_controller,
@@ -550,7 +365,6 @@ def test_emit_metrics_activity_only(
from model_manager.activities.cleanup import Cleanup
cleanup = Cleanup(
storage_repository=mock_storage_repository,
logger=mock_logger,
notification_handler=mock_notification_handler,
metrics_controller=mock_metrics_controller,

16
tests/activities/test_training.py
View File

@@ -342,7 +342,6 @@ def test_train_model_training_fails(
):
"""Test train_model when training fails."""
from model_manager.activities.training import Training
from model_manager.utils.exceptions import ModelTrainingError
training = Training(
model_repository=mock_model_repository,
@@ -362,11 +361,10 @@ def test_train_model_training_fails(
'train_params': mock_train_params,
}
with pytest.raises(ModelTrainingError) as exc_info:
with pytest.raises(RuntimeError) as exc_info:
asyncio.run(training.train_model(input_data))
assert exc_info.value.model_trained is False
assert exc_info.value.model_saved is False
assert str(exc_info.value) == 'Training failed'
training.send_notification.assert_called_once()
@@ -382,7 +380,6 @@ def test_train_model_save_fails(
):
"""Test train_model when model saving fails."""
from model_manager.activities.training import Training
from model_manager.utils.exceptions import ModelTrainingError
training = Training(
model_repository=mock_model_repository,
@@ -406,11 +403,10 @@ def test_train_model_save_fails(
'train_params': mock_train_params,
}
with pytest.raises(ModelTrainingError) as exc_info:
with pytest.raises(RuntimeError) as exc_info:
asyncio.run(training.train_model(input_data))
assert exc_info.value.model_trained is True
assert exc_info.value.model_saved is False
assert str(exc_info.value) == 'Save failed'
training.send_notification.assert_called_once()
@@ -437,14 +433,11 @@ def test_cleanup_resources_success(
input_data = {
'metadata': {'workflow_id': 'test-123'},
'run_dir': '/tmp/run_001', # noqa: S108
'bucket_name': 'test-bucket',
'file_name': 'test-file.csv',
}
asyncio.run(training.cleanup_resources(input_data))
mock_model_repository.cleanup_run_directory.assert_called_once_with('/tmp/run_001') # noqa: S108
mock_storage_repository.delete_file.assert_called_once_with('test-bucket', 'test-file.csv')
@patch('model_manager.activities.training.TrainingRepository')
@@ -510,4 +503,3 @@ def test_cleanup_resources_with_empty_values(
asyncio.run(training.cleanup_resources(input_data))
mock_model_repository.cleanup_run_directory.assert_called_once_with('')
mock_storage_repository.delete_file.assert_called_once_with('', '')

19
tests/utils/models/test_experiment_status.py
View File

@@ -8,15 +8,12 @@ def test_experiment_status_values():
assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC'
assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS'
assert ExperimentStatus.TRAINING_ERROR == 'TRAINING_ERROR'
assert ExperimentStatus.TRACKING_SENT == 'TRACKING_SENT'
assert ExperimentStatus.TRACKING_SEND_ERROR == 'TRACKING_SEND_ERROR'
assert ExperimentStatus.FILE_DELETED == 'FILE_DELETED'
assert ExperimentStatus.FILE_DELETE_ERROR == 'FILE_DELETE_ERROR'
def test_experiment_status_count():
"""Test that enum has exactly 8 status values."""
assert len(ExperimentStatus) == 8
assert len(ExperimentStatus) == 5
def test_experiment_status_is_string():
@@ -31,22 +28,16 @@ def test_experiment_status_membership():
assert 'ORCHESTRATOR_WAITING_PROC' in [s.value for s in ExperimentStatus]
assert 'TRAINING_SUCCESS' in [s.value for s in ExperimentStatus]
assert 'TRAINING_ERROR' in [s.value for s in ExperimentStatus]
assert 'TRACKING_SENT' in [s.value for s in ExperimentStatus]
assert 'TRACKING_SEND_ERROR' in [s.value for s in ExperimentStatus]
assert 'FILE_DELETED' in [s.value for s in ExperimentStatus]
assert 'FILE_DELETE_ERROR' in [s.value for s in ExperimentStatus]
def test_experiment_status_iteration():
"""Test that enum can be iterated."""
statuses = list(ExperimentStatus)
assert len(statuses) == 8
assert len(statuses) == 5
assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC in statuses
assert ExperimentStatus.TRAINING_SUCCESS in statuses
assert ExperimentStatus.TRAINING_ERROR in statuses
assert ExperimentStatus.TRACKING_SENT in statuses
assert ExperimentStatus.TRACKING_SEND_ERROR in statuses
assert ExperimentStatus.FILE_DELETED in statuses
assert ExperimentStatus.FILE_DELETE_ERROR in statuses
@@ -64,9 +55,6 @@ def test_experiment_status_access_by_name():
)
assert ExperimentStatus['TRAINING_SUCCESS'] == ExperimentStatus.TRAINING_SUCCESS
assert ExperimentStatus['TRAINING_ERROR'] == ExperimentStatus.TRAINING_ERROR
assert ExperimentStatus['TRACKING_SENT'] == ExperimentStatus.TRACKING_SENT
assert ExperimentStatus['TRACKING_SEND_ERROR'] == ExperimentStatus.TRACKING_SEND_ERROR
assert ExperimentStatus['FILE_DELETED'] == ExperimentStatus.FILE_DELETED
assert ExperimentStatus['FILE_DELETE_ERROR'] == ExperimentStatus.FILE_DELETE_ERROR
@@ -77,7 +65,4 @@ def test_experiment_status_access_by_value():
)
assert ExperimentStatus('TRAINING_SUCCESS') == ExperimentStatus.TRAINING_SUCCESS
assert ExperimentStatus('TRAINING_ERROR') == ExperimentStatus.TRAINING_ERROR
assert ExperimentStatus('TRACKING_SENT') == ExperimentStatus.TRACKING_SENT
assert ExperimentStatus('TRACKING_SEND_ERROR') == ExperimentStatus.TRACKING_SEND_ERROR
assert ExperimentStatus('FILE_DELETED') == ExperimentStatus.FILE_DELETED
assert ExperimentStatus('FILE_DELETE_ERROR') == ExperimentStatus.FILE_DELETE_ERROR

95
tests/utils/repository/test_storage_repository.py
View File

@@ -237,83 +237,6 @@ def test_fetch_file_network_error(mock_boto3, mock_logger, storage_config):
repo.fetch_file('test-bucket', 'test-file.csv')
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_delete_file_success(mock_boto3, mock_logger, storage_config):
"""Test successful file deletion from MinIO."""
from model_manager.utils.repository.storage_repository import StorageRepository
mock_s3_client = Mock()
mock_s3_client.delete_object.return_value = {}
mock_boto3.client.return_value = mock_s3_client
repo = StorageRepository(logger=mock_logger, **storage_config)
# Delete file
repo.delete_file('test-bucket', 'test-file.csv')
# Verify delete_object was called correctly
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key='test-file.csv')
# Verify logging
assert mock_logger.info.call_count >= 2 # Init + delete
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_delete_file_non_existent(mock_boto3, mock_logger, storage_config):
"""Test delete file that doesn't exist (should succeed silently in S3)."""
from model_manager.utils.repository.storage_repository import StorageRepository
# S3/MinIO delete is idempotent - deleting non-existent file succeeds
mock_s3_client = Mock()
mock_s3_client.delete_object.return_value = {}
mock_boto3.client.return_value = mock_s3_client
repo = StorageRepository(logger=mock_logger, **storage_config)
# Delete non-existent file (should succeed)
repo.delete_file('test-bucket', 'non-existent.csv')
mock_s3_client.delete_object.assert_called_once()
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_delete_file_access_denied(mock_boto3, mock_logger, storage_config):
"""Test delete file when access is denied."""
from model_manager.utils.repository.storage_repository import StorageRepository
# Setup mock to raise AccessDenied error
mock_s3_client = Mock()
mock_s3_client.delete_object.side_effect = ClientError(
{'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}, 'DeleteObject'
)
mock_boto3.client.return_value = mock_s3_client
repo = StorageRepository(logger=mock_logger, **storage_config)
# Attempt to delete file without permissions
with pytest.raises(ClientError) as exc_info:
repo.delete_file('test-bucket', 'protected-file.csv')
assert exc_info.value.response['Error']['Code'] == 'AccessDenied'
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_delete_file_network_error(mock_boto3, mock_logger, storage_config):
"""Test delete file when network error occurs."""
from model_manager.utils.repository.storage_repository import StorageRepository
# Setup mock to raise network error
mock_s3_client = Mock()
mock_s3_client.delete_object.side_effect = ConnectionError('Network unreachable')
mock_boto3.client.return_value = mock_s3_client
repo = StorageRepository(logger=mock_logger, **storage_config)
# Attempt to delete file with network error
with pytest.raises(ConnectionError):
repo.delete_file('test-bucket', 'test-file.csv')
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_storage_repository_with_ssl(mock_boto3, mock_logger, storage_config):
"""Test StorageRepository initialization with SSL enabled."""
@@ -404,24 +327,6 @@ def test_fetch_file_with_special_characters(mock_boto3, mock_logger, storage_con
mock_s3_client.get_object.assert_called_once_with(Bucket='test-bucket', Key=special_filename)
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_delete_file_with_path_separators(mock_boto3, mock_logger, storage_config):
"""Test delete file with path separators in object key."""
from model_manager.utils.repository.storage_repository import StorageRepository
mock_s3_client = Mock()
mock_s3_client.delete_object.return_value = {}
mock_boto3.client.return_value = mock_s3_client
repo = StorageRepository(logger=mock_logger, **storage_config)
# Delete file with path separators
file_path = 'data/2023/01/test-file.csv'
repo.delete_file('test-bucket', file_path)
mock_s3_client.delete_object.assert_called_once_with(Bucket='test-bucket', Key=file_path)
@patch('model_manager.utils.repository.storage_repository.boto3')
def test_storage_repository_different_regions(mock_boto3, mock_logger, storage_config):
"""Test StorageRepository with different AWS regions."""

79
tests/utils/test_exceptions.py
View File

@@ -1,79 +0,0 @@
"""Unit tests for custom exceptions with 100% coverage."""
import pytest
def test_model_training_error_with_default_message():
"""Test ModelTrainingError with default message."""
from model_manager.utils.exceptions import ModelTrainingError
error = ModelTrainingError(model_trained=True, model_saved=False)
assert error.model_trained is True
assert error.model_saved is False
assert str(error) == 'Model training workflow failed (model_trained=True, model_saved=False)'
def test_model_training_error_with_custom_message():
"""Test ModelTrainingError with custom message."""
from model_manager.utils.exceptions import ModelTrainingError
custom_msg = 'Custom error occurred during training'
error = ModelTrainingError(model_trained=False, model_saved=False, message=custom_msg)
assert error.model_trained is False
assert error.model_saved is False
assert str(error) == custom_msg
def test_model_training_error_both_true():
"""Test ModelTrainingError when both flags are True."""
from model_manager.utils.exceptions import ModelTrainingError
error = ModelTrainingError(model_trained=True, model_saved=True)
assert error.model_trained is True
assert error.model_saved is True
assert str(error) == 'Model training workflow failed (model_trained=True, model_saved=True)'
def test_model_training_error_both_false():
"""Test ModelTrainingError when both flags are False."""
from model_manager.utils.exceptions import ModelTrainingError
error = ModelTrainingError(model_trained=False, model_saved=False)
assert error.model_trained is False
assert error.model_saved is False
assert str(error) == 'Model training workflow failed (model_trained=False, model_saved=False)'
def test_model_training_error_is_exception():
"""Test ModelTrainingError is an Exception subclass."""
from model_manager.utils.exceptions import ModelTrainingError
error = ModelTrainingError(model_trained=True, model_saved=False)
assert isinstance(error, Exception)
def test_model_training_error_can_be_raised():
"""Test ModelTrainingError can be raised and caught."""
from model_manager.utils.exceptions import ModelTrainingError
with pytest.raises(ModelTrainingError) as exc_info:
raise ModelTrainingError(model_trained=True, model_saved=False)
assert exc_info.value.model_trained is True
assert exc_info.value.model_saved is False
def test_model_training_error_with_empty_message():
"""Test ModelTrainingError with empty string message."""
from model_manager.utils.exceptions import ModelTrainingError
error = ModelTrainingError(model_trained=True, model_saved=True, message='')
assert error.model_trained is True
assert error.model_saved is True
assert str(error) == ''

54
tests/workflows/test_cleanup_files.py
View File

@@ -8,8 +8,8 @@ import pytest
@pytest.mark.asyncio
@patch('model_manager.workflows.cleanup_files.workflow')
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
"""Test the CleanupFiles workflow when bucket_name is provided in the input."""
async def test_cleanup_files_workflow(mock_workflow_module):
"""Test the CleanupFiles workflow."""
from model_manager.workflows.cleanup_files import CleanupFiles
# Mock execute_activity_method
@@ -17,58 +17,14 @@ async def test_cleanup_files_workflow_with_input_bucket(mock_workflow_module):
# Instantiate and run the workflow
workflow_instance = CleanupFiles()
await workflow_instance.run({'bucket_name': 'input-bucket'})
await workflow_instance.run({})
# Verify that the activities were called with the correct parameters
calls = mock_workflow_module.execute_activity_method.call_args_list
assert len(calls) == 2
# Check cleanup_minio_files call
minio_call_args = calls[0][0][1]
assert minio_call_args['bucket_name'] == 'input-bucket'
assert minio_call_args['metadata'] == {
'pod_id': 'temporal-pod',
'workflow_name': 'cleanup_files',
}
assert len(calls) == 1
# Check cleanup_temp_directories call
local_call_args = calls[1][0][1]
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
assert local_call_args['metadata'] == {
'pod_id': 'temporal-pod',
'workflow_name': 'cleanup_files',
}
@pytest.mark.asyncio
@patch('model_manager.workflows.cleanup_files.workflow')
@patch('model_manager.workflows.cleanup_files.POD_ID', 'temporal-pod')
@patch('model_manager.workflows.cleanup_files.DEFAULT_CLEANUP_BUCKET', 'env-var-bucket')
async def test_cleanup_files_workflow_with_default_bucket(mock_workflow_module):
"""Test the CleanupFiles workflow when using the default bucket from environment variables."""
from model_manager.workflows.cleanup_files import CleanupFiles
# Mock execute_activity_method
mock_workflow_module.execute_activity_method = AsyncMock()
# Instantiate and run the workflow
workflow_instance = CleanupFiles()
await workflow_instance.run({}) # Empty input
# Verify that the activities were called
calls = mock_workflow_module.execute_activity_method.call_args_list
assert len(calls) == 2
# Check cleanup_minio_files call
minio_call_args = calls[0][0][1]
assert minio_call_args['bucket_name'] == 'env-var-bucket'
assert minio_call_args['metadata'] == {
'pod_id': 'temporal-pod',
'workflow_name': 'cleanup_files',
}
# Check cleanup_temp_directories call
local_call_args = calls[1][0][1]
local_call_args = calls[0][0][1]
assert local_call_args['temp_path'] == 'model_manager/reports/temp'
assert local_call_args['metadata'] == {
'pod_id': 'temporal-pod',

31
tests/workflows/test_train_model.py
View File

@@ -4,7 +4,6 @@ from unittest.mock import AsyncMock, Mock, patch
import pytest
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
@@ -275,20 +274,18 @@ async def test_train_model_mlflow_error(mock_workflow_module, mock_train_params)
from model_manager.workflows.train_model import TrainModel
# Setup mocks - MLflow save fails
mlflow_error = ModelTrainingError(
model_trained=True, model_saved=False, message='MLflow save failed'
)
mlflow_error = RuntimeError('MLflow save failed')
mock_workflow_module.execute_activity_method = AsyncMock(side_effect=[mlflow_error, None])
workflow_instance = TrainModel()
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
with pytest.raises(ModelTrainingError):
with pytest.raises(RuntimeError):
await workflow_instance._train_model(mock_train_params, 123, metadata)
# Verify TRACKING_SEND_ERROR status was set
# Verify TRAINING_ERROR status was set
call_args = mock_workflow_module.execute_activity_method.call_args_list[1]
assert call_args[0][1]['status'] == ExperimentStatus.TRACKING_SEND_ERROR
assert call_args[0][1]['status'] == ExperimentStatus.TRAINING_ERROR
@pytest.mark.asyncio
@@ -306,14 +303,11 @@ async def test_cleanup_resources_success(mock_workflow_module):
metadata = {'metadata': {'pod_id': 'test-pod', 'experiment_run_id': 123}}
await workflow_instance._cleanup_resources(
experiment_run_id=123,
run_dir='/tmp/test-run', # noqa: S108
bucket_name='test-bucket',
file_name='test-file.csv',
metadata=metadata,
)
assert mock_workflow_module.execute_activity_method.call_count == 2
assert mock_workflow_module.execute_activity_method.call_count == 1
@pytest.mark.asyncio
@@ -332,15 +326,12 @@ async def test_cleanup_resources_failure(mock_workflow_module):
with pytest.raises(RuntimeError, match='Cleanup failed'):
await workflow_instance._cleanup_resources(
experiment_run_id=123,
run_dir='/tmp/test-run', # noqa: S108
bucket_name='test-bucket',
file_name='test-file.csv',
metadata=metadata,
)
# Verify error status update was called
assert mock_workflow_module.execute_activity_method.call_count == 2
assert mock_workflow_module.execute_activity_method.call_count == 1
@pytest.mark.asyncio
@@ -411,7 +402,7 @@ async def test_update_experiment_run_with_run_name(mock_workflow_module):
metadata=metadata,
experiment_run_id=123,
update_type=UpdateType.MODEL_SAVED,
status=ExperimentStatus.TRACKING_SENT,
status=ExperimentStatus.TRAINING_SUCCESS,
run_name='test-run-123',
)
@@ -440,9 +431,8 @@ async def test_run_complete_workflow_success(
mock_train_params, # validate_train_params
None, # update status (ORCHESTRATOR_WAITING_PROC)
train_result, # train_model
None, # update status (TRACKING_SENT)
None, # update status (TRAINING_SUCCESS)
None, # cleanup_resources
None, # update status (FILE_DELETED)
]
)
@@ -452,7 +442,7 @@ async def test_run_complete_workflow_success(
await workflow_instance.run(sample_input_data)
# Verify all activities were called
assert mock_workflow_module.execute_activity_method.call_count == 6
assert mock_workflow_module.execute_activity_method.call_count == 5
@pytest.mark.asyncio
@@ -523,7 +513,6 @@ async def test_run_workflow_cleanup_error(
mock_train_params, # validate_train_params
None, # update status (ORCHESTRATOR_WAITING_PROC)
train_result, # train_model
None, # update status (TRACKING_SENT)
RuntimeError('Cleanup failed'), # cleanup_resources fails
None, # update status (FILE_DELETE_ERROR)
]
@@ -534,7 +523,7 @@ async def test_run_workflow_cleanup_error(
with pytest.raises(RuntimeError, match='Cleanup failed'):
await workflow_instance.run(sample_input_data)
assert mock_workflow_module.execute_activity_method.call_count == 6
assert mock_workflow_module.execute_activity_method.call_count == 5
@pytest.mark.asyncio

6
values.yaml
View File

@@ -237,14 +237,8 @@ env:
value: "24"
- name: CLEANUP_DRY_RUN
value: "false"
- name: TIMEOUT_CLEANUP_MINIO
value: "300"
- name: TIMEOUT_CLEANUP_LOCAL
value: "120"
- name: MAX_KEYS_CLEANUP
value: "1000"
- name: DEFAULT_CLEANUP_BUCKET
value: "model-training"
# Cleanup Schedule Configuration
- name: CLEANUP_SCHEDULE_ID