Merge branch 'main' into release/SIENTIAPDE-1645
This commit is contained in:
vitor-aignosi
@@ -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 * * *"
|
||||
|
||||
21
.github/workflows/deploy.yml
vendored
21
.github/workflows/deploy.yml
vendored
@@ -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
|
||||
|
||||
79
README.md
79
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]
|
||||
A[Scheduled Trigger] --> B[cleanup_temp_directories]
|
||||
|
||||
B -.-> MinIO[MinIO Storage]
|
||||
C -.-> FS[Local Filesystem]
|
||||
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
|
||||
|
||||
|
||||
@@ -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:
|
||||
"""
|
||||
|
||||
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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'
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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('', '')
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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) == ''
|
||||
@@ -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',
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user