diff --git a/.env.example b/.env.example index 22e8021..64f08f9 100644 --- a/.env.example +++ b/.env.example @@ -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 * * *" diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml index 0593968..8ff00b2 100644 --- a/.github/workflows/deploy.yml +++ b/.github/workflows/deploy.yml @@ -10,29 +10,13 @@ on: - 'feature/**' jobs: - get-version: - name: Determine Next Version - runs-on: ubuntu-latest - permissions: write-all - if: github.event.pull_request.merged == true - outputs: - version: ${{ steps.version.outputs.next_version }} - steps: - - uses: actions/checkout@v4 - - - name: Determine next version - id: version - uses: Aignosi/github_workflow_templates/.github/actions/determine-next-version@main - with: - branch_name: ${{ github.event.pull_request.head.ref }} - deploy: name: Deploy - needs: [get-version] + if: github.event.pull_request.merged == true permissions: write-all uses: Aignosi/github_workflow_templates/.github/workflows/reusable-deploy.yml@main with: - version: ${{ needs.get-version.outputs.version }} + branch_name: ${{ github.event.pull_request.head.ref }} project_type: 'python' image_name: 'sientia-dataops-model-manager' helm_chart_path: 'sientia-module' @@ -45,4 +29,5 @@ jobs: helm_repo_url: 'https://raw.githubusercontent.com/Aignosi/sientia-dataops-helm-repo/refs/heads/main/' helm_chart_version: '0.6.0' helm_values_file: './values.yaml' + use_vpn: true secrets: inherit diff --git a/README.md b/README.md index e27e636..60cd104 100644 --- a/README.md +++ b/README.md @@ -80,7 +80,7 @@ An enterprise-grade ML model training orchestration platform built on Temporal. ### Core Functionality - **ML Model Training Pipeline**: Complete training workflow from validation to deployment using MLFlow - **Polynomial Regression Support**: Configurable polynomial degree with interaction terms and mandatory scaler validation -- **Automated File Cleanup**: Scheduled cleanup of stale files from MinIO and local filesystem +- **Automated File Cleanup**: Scheduled cleanup of stale files from local filesystem - **Temporal Workflow Orchestration**: Robust workflow management with granular retry policies and fault tolerance - **Parameter Validation**: Defense-in-depth validation with business rules and type checking - **Experiment Tracking**: Comprehensive status tracking in PostgreSQL database @@ -184,14 +184,13 @@ The Model Manager system uses a Temporal-based workflow architecture with clear - No exception raising on failure - allows workflow to handle errors gracefully - Integration with TrainingRepository for business logic separation - MLFlow model saving and artifact management - - MinIO object storage operations - **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**: Local directory cleanup operations - `cleanup_temp_directories()`: Cleans local temporary directories - Configurable retention period (default: 24 hours) - Dry-run mode for testing + - No MinIO cleanup (files are managed by external processes) - **Key Features**: - Multiple inheritance pattern for unified activity interface - Parameter validation with business rules @@ -220,16 +219,15 @@ The Model Manager system uses a Temporal-based workflow architecture with clear #### **Model Training Pipeline** ``` -Training Request → Parameter Validation → MinIO Data Download → -Model Training → MLFlow Model Save → Resource Cleanup → Status Update +Training Request → Parameter Validation → Model Training → +MLFlow Model Save → Resource Cleanup → Status Update ``` **Key Stages:** 1. **Validation**: Experiment run ID and training parameters validation -2. **Data Acquisition**: Download training data from MinIO storage -3. **Training**: Execute ML model training with validated parameters -4. **Persistence**: Save trained model and artifacts to MLFlow -5. **Cleanup**: Remove temporary files and update experiment status +2. **Training**: Execute ML model training with validated parameters (data provided in request) +3. **Persistence**: Save trained model and artifacts to MLFlow +4. **Cleanup**: Remove temporary local directories and update experiment status ### Security Architecture @@ -267,10 +265,9 @@ The **TrainModel** workflow orchestrates the complete ML model training pipeline #### Execution Flow 1. **Validate Experiment Run ID**: Critical validation before any DB updates 2. **Validate Training Parameters**: Type checking + business rules validation -3. **Download Training Data**: Fetch file from MinIO storage -4. **Train Model**: Execute ML model training with validated parameters -5. **Save to MLFlow**: Save trained model and artifacts to MLFlow -6. **Cleanup Resources**: Delete temporary files and MinIO data +3. **Train Model**: Execute ML model training with validated parameters +4. **Save to MLFlow**: Save trained model and artifacts to MLFlow +5. **Cleanup Resources**: Delete temporary local directories #### Key Features - **Granular Retry Policies**: Different strategies for network, training, MLFlow, database, and filesystem operations @@ -317,18 +314,13 @@ The **TrainModel** workflow orchestrates the complete ML model training pipeline ```mermaid flowchart TD A[1. validate_experiment_run_id] --> B[2. validate_train_params] - B --> C[3. fetch_file_from_minio] - C --> D[4. train_model] - D --> E[5. save_model] - E --> F[6. cleanup_run_directory] - F --> G[7. delete_file_from_minio] + B --> C[3. train_model] + C --> D[4. cleanup_run_directory] B -.-> DB[(PostgreSQL)] - C -.-> MinIO[MinIO Storage] - D -.-> Training[ML Training] - E -.-> MLFlow[MLFlow] - F -.-> FS[Filesystem] - G -.-> MinIO + C -.-> Training[ML Training] + C -.-> MLFlow[MLFlow] + D -.-> FS[Filesystem] ``` #### Retry Strategies @@ -337,11 +329,9 @@ The workflow implements 5 different retry policies optimized for each operation | Operation Type | Initial Interval | Max Interval | Backoff | Max Attempts | Use Case | |---------------|------------------|--------------|---------|--------------|----------| -| **Network** | 1s | 10s | 2.0x | 5 | MinIO operations (transient network errors) | +| **Network** | 1s | 10s | 2.0x | 5 | Network operations (transient errors) | | **No Retry** | - | - | - | 1 | Training/Validation (permanent data errors) | -| **MLFlow** | 5s | 30s | 2.0x | 3 | MLFlow operations (API timeouts) | | **Database** | 2s | 20s | 2.0x | 5 | PostgreSQL updates (lock contention) | -| **Filesystem** | 2s | 10s | 1.5x | 3 | Cleanup operations (busy resources) | #### Business Validation Rules @@ -364,21 +354,20 @@ The workflow validates comprehensive business rules beyond type checking: ### Cleanup Files Workflow (`cleanup_files.py`) -The **CleanupFiles** workflow provides automated cleanup of stale files from MinIO storage and local temporary directories. It runs on a scheduled basis (default: daily at midnight UTC) to maintain storage hygiene. +The **CleanupFiles** workflow provides automated cleanup of stale local temporary directories. It runs on a scheduled basis (default: daily at midnight UTC) to maintain storage hygiene. #### Purpose -- **Storage Management**: Automatic removal of old files from MinIO and local filesystem +- **Storage Management**: Automatic removal of old temporary directories from local filesystem - **Retention Policy**: Configurable retention period (default: 24 hours) - **Scheduled Execution**: Cron-based scheduling for automated cleanup -- **Resource Optimization**: Prevents storage bloat and reduces costs +- **Resource Optimization**: Prevents storage bloat and reduces disk usage #### Execution Flow -1. **Cleanup MinIO Files**: Scan and delete files older than retention period from MinIO bucket -2. **Cleanup Local Directories**: Remove temporary directories older than retention period +1. **Cleanup Local Directories**: Remove temporary directories older than retention period #### Key Features -- **Timestamp-Based Cleanup**: Uses filename/directory timestamps for age determination -- **Pattern Matching**: Regex patterns for MinIO (`timestamp-filename`) and directories (`name_YYYYMMDD_HHMMSS_microseconds`) +- **Timestamp-Based Cleanup**: Uses directory timestamps for age determination +- **Pattern Matching**: Regex pattern for directories (`name_YYYYMMDD_HHMMSS_microseconds`) - **Configurable Retention**: Environment variable-based retention period - **Dry-Run Mode**: Test cleanup operations without actual deletion - **Idempotent**: Safe to run multiple times @@ -387,7 +376,7 @@ The **CleanupFiles** workflow provides automated cleanup of stale files from Min #### Input Parameters ```json { - "bucket_name": "model-training" // Optional, defaults to DEFAULT_CLEANUP_BUCKET env var + "temp_path": "model_manager/reports/temp" // Optional, defaults to 'model_manager/reports/temp' } ``` @@ -404,36 +393,28 @@ 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 | Operation Type | Initial Interval | Max Interval | Backoff | Max Attempts | Use Case | |---------------|------------------|--------------|---------|--------------|----------| -| **Network** | 1s | 10s | 2.0x | 5 | MinIO operations (transient network errors) | | **No Retry** | - | - | - | 1 | Local filesystem operations (permanent errors) | #### Cleanup Patterns -**MinIO Files:** -- Pattern: `{timestamp}-{filename}` where timestamp is milliseconds since epoch -- Example: `1638360000000-training_data.csv` -- Retention: Files older than `CLEANUP_RETENTION_HOURS` are deleted - **Local Directories:** - Pattern: `{name}_{YYYYMMDD}_{HHMMSS}_{microseconds}` - Example: `temp_20231201_143052_123456` - Retention: Directories older than `CLEANUP_RETENTION_HOURS` are deleted +- Location: `model_manager/reports/temp/` by default ## Installation & Setup @@ -1056,7 +1037,6 @@ The Model Manager system exposes comprehensive Prometheus metrics for operationa ### Cleanup Metrics - Cleanup execution success/failure rates -- Number of files deleted from MinIO - Number of directories cleaned from local filesystem - Cleanup duration and performance @@ -1101,8 +1081,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 | @@ -1120,14 +1098,12 @@ These timeouts control how long each activity in workflows can run before timing |----------|-------------|---------|-------------------| | `TIMEOUT_VALIDATE_PARAMS` | Parameter validation timeout | `30` | Fast operation, no I/O | | `TIMEOUT_TRAIN_MODEL` | Model training timeout | `2700` | Large dataset processing (45 min) | -| `TIMEOUT_DELETE_FILE` | Delete file from MinIO timeout | `120` | MinIO delete operation (2 min) | | `TIMEOUT_UPDATE_DATABASE` | Database update timeout | `30` | PostgreSQL update query (30 sec) | **Cleanup Workflow Timeouts:** | 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 +1317,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 diff --git a/model_manager/activities/cleanup.py b/model_manager/activities/cleanup.py index fcf22b0..81d85d8 100644 --- a/model_manager/activities/cleanup.py +++ b/model_manager/activities/cleanup.py @@ -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 @@ -27,7 +27,6 @@ with workflow.unsafe.imports_passed_through(): 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 +34,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, - minio_repository: MinioRepository, logger: Logger, notification_handler: NotificationHandler, metrics_controller: MetricsController, @@ -50,145 +47,21 @@ class Cleanup(SientiaMonitoring): Initialize Cleanup activity. Args: - minio_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.minio_repository = minio_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. MinioRepository applies BASE_PREFIX - # internally; we request all objects under that prefix for this bucket. - objects = await self.minio_repository.list_objects( - prefix='', - bucket=bucket_name, - recursive=True, - metadata=metadata, - ) - - for obj_key in objects: - files_scanned += 1 - - # Extract timestamp from the filename portion of the object key - filename = obj_key.split('/')[-1] - match = self.minio_timestamp_pattern.match(filename) - 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: - await self.minio_repository.delete_file( - object_name=obj_key, - bucket=None, - metadata=metadata, - ) - self.info(f'Deleted stale file: {obj_key}', metadata) - files_deleted += 1 - except Exception as e: # noqa: BLE001 - 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: """ diff --git a/model_manager/activities/training.py b/model_manager/activities/training.py index d86cdcf..4bc771f 100644 --- a/model_manager/activities/training.py +++ b/model_manager/activities/training.py @@ -24,7 +24,6 @@ with workflow.unsafe.imports_passed_through(): from sientia_model.wrappers.sientia_model import SientiaModel 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.data_manager_repository import DataManagerRepository @@ -261,11 +260,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() diff --git a/model_manager/utils/exceptions.py b/model_manager/utils/exceptions.py deleted file mode 100644 index 8f9c4a8..0000000 --- a/model_manager/utils/exceptions.py +++ /dev/null @@ -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) diff --git a/model_manager/utils/models/experiment_status.py b/model_manager/utils/models/experiment_status.py index 28c3b8b..1fb70fa 100644 --- a/model_manager/utils/models/experiment_status.py +++ b/model_manager/utils/models/experiment_status.py @@ -14,21 +14,13 @@ 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. """ ORCHESTRATOR_VALIDATION_ERROR = 'ORCHESTRATOR_VALIDATION_ERROR' 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' diff --git a/model_manager/worker/worker.py b/model_manager/worker/worker.py index 0729595..77bd453 100644 --- a/model_manager/worker/worker.py +++ b/model_manager/worker/worker.py @@ -190,7 +190,6 @@ async def main(): main_workflow=CleanupFiles, other_workflows=[], activities=[ - activities.cleanup_minio_files, activities.cleanup_temp_directories, ], temporal_client=temporal_client, diff --git a/model_manager/workflows/cleanup_files.py b/model_manager/workflows/cleanup_files.py index 685f946..de214e7 100644 --- a/model_manager/workflows/cleanup_files.py +++ b/model_manager/workflows/cleanup_files.py @@ -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, diff --git a/model_manager/workflows/train_model.py b/model_manager/workflows/train_model.py index 63273cb..9ff81d8 100644 --- a/model_manager/workflows/train_model.py +++ b/model_manager/workflows/train_model.py @@ -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 @@ -273,7 +272,7 @@ 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'), ) @@ -310,7 +309,7 @@ class TrainModel: metadata: dict[str, Any], ) -> None: """ - Cleanup resources and delete file from MinIO. + Cleanup resources. This method deletes the training file from MinIO. On success, updates DB status to FILE_DELETED. @@ -319,9 +318,6 @@ class TrainModel: Args: experiment_run_id: Validated experiment run ID metadata: Workflow execution metadata - - Raises: - Exception: If cleanup fails (after updating DB status) """ try: await workflow.execute_activity_method( diff --git a/pyproject.toml b/pyproject.toml index ea5d4cd..edf408a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "model-manager" -version = "1.1.2" +version = "1.2.0" description = "Sientia DataOps Model Manager - ML Model Orchestration System" readme = "README.md" requires-python = ">=3.11" diff --git a/scripts/run_cleanup_test.py b/scripts/run_cleanup_test.py index f27fe1f..612eb6b 100644 --- a/scripts/run_cleanup_test.py +++ b/scripts/run_cleanup_test.py @@ -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( diff --git a/tests/activities/test_cleanup.py b/tests/activities/test_cleanup.py index 2d8f7b3..5969939 100644 --- a/tests/activities/test_cleanup.py +++ b/tests/activities/test_cleanup.py @@ -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, diff --git a/tests/activities/test_training.py b/tests/activities/test_training.py index 7d20057..d3be47a 100644 --- a/tests/activities/test_training.py +++ b/tests/activities/test_training.py @@ -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('', '') diff --git a/tests/utils/models/test_experiment_status.py b/tests/utils/models/test_experiment_status.py index 44f494a..42682d9 100644 --- a/tests/utils/models/test_experiment_status.py +++ b/tests/utils/models/test_experiment_status.py @@ -5,18 +5,15 @@ from model_manager.utils.models.experiment_status import ExperimentStatus def test_experiment_status_values(): """Test that all expected status values exist.""" + assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR' 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 + """Test that enum has exactly 4 status values.""" + assert len(ExperimentStatus) == 4 def test_experiment_status_is_string(): @@ -28,30 +25,25 @@ def test_experiment_status_is_string(): def test_experiment_status_membership(): """Test membership checks for status values.""" + assert 'ORCHESTRATOR_VALIDATION_ERROR' in [s.value for s in ExperimentStatus] 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) == 4 + assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR in statuses 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 def test_experiment_status_comparison(): """Test that enum values can be compared with strings.""" + assert ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR == 'ORCHESTRATOR_VALIDATION_ERROR' assert ExperimentStatus.ORCHESTRATOR_WAITING_PROC == 'ORCHESTRATOR_WAITING_PROC' assert ExperimentStatus.TRAINING_SUCCESS == 'TRAINING_SUCCESS' assert ExperimentStatus.TRAINING_ERROR != 'TRAINING_SUCCESS' @@ -59,25 +51,25 @@ def test_experiment_status_comparison(): def test_experiment_status_access_by_name(): """Test accessing enum members by name.""" + assert ( + ExperimentStatus['ORCHESTRATOR_VALIDATION_ERROR'] + == ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR + ) assert ( ExperimentStatus['ORCHESTRATOR_WAITING_PROC'] == ExperimentStatus.ORCHESTRATOR_WAITING_PROC ) 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 def test_experiment_status_access_by_value(): """Test accessing enum members by value.""" + assert ( + ExperimentStatus('ORCHESTRATOR_VALIDATION_ERROR') + == ExperimentStatus.ORCHESTRATOR_VALIDATION_ERROR + ) assert ( ExperimentStatus('ORCHESTRATOR_WAITING_PROC') == ExperimentStatus.ORCHESTRATOR_WAITING_PROC ) 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 diff --git a/tests/utils/test_exceptions.py b/tests/utils/test_exceptions.py deleted file mode 100644 index 4444596..0000000 --- a/tests/utils/test_exceptions.py +++ /dev/null @@ -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) == '' diff --git a/tests/workflows/test_cleanup_files.py b/tests/workflows/test_cleanup_files.py index 28d02e5..d23bf75 100644 --- a/tests/workflows/test_cleanup_files.py +++ b/tests/workflows/test_cleanup_files.py @@ -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', diff --git a/tests/workflows/test_train_model.py b/tests/workflows/test_train_model.py index 2157bb2..23323d8 100644 --- a/tests/workflows/test_train_model.py +++ b/tests/workflows/test_train_model.py @@ -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,9 +513,8 @@ 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) + None, # update status (TRAINING_SUCCESS) 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