sientia-dataops-model-manager/model_manager/activities/cleanup.py

"""
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
of the database, using timestamps embedded in filenames.
"""

from temporalio import activity, workflow

with workflow.unsafe.imports_passed_through():
    import os
    import re
    import shutil
    import traceback
    from datetime import datetime, timedelta
    from typing import Any

    from sientia_do.notifications.handlers import CoreNotificationHandler as NotificationHandler
    from sientia_do.notifications.models import NotificationLevel
    from sientia_do.observability.logger import Logger
    from sientia_do.observability.metrics_controller import MetricsController
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring

    from model_manager.metrics import ACTIVITY_EXECUTION_TOTAL, WORKFLOW_EXECUTION_TOTAL

    RETENTION_HOURS = int(os.getenv('CLEANUP_RETENTION_HOURS', '24'))
    DRY_RUN = os.getenv('CLEANUP_DRY_RUN', 'false').lower() == 'true'


class Cleanup(SientiaMonitoring):
    """
    Activity for cleaning up stale files and directories.

    This activity extends SientiaMonitoring and handles cleanup of:
    - Local temporary directories with timestamp suffixes
    """

    def __init__(
        self,
        logger: Logger,
        notification_handler: NotificationHandler,
        metrics_controller: MetricsController,
    ):
        """
        Initialize Cleanup activity.

        Args:
            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)

        # Configuration from environment variables
        self.retention_hours = RETENTION_HOURS
        self.dry_run = DRY_RUN

        # Regex patterns for timestamp extraction
        self.dir_timestamp_pattern = re.compile(
            r'^(.+)_(\d{8}_\d{6}_\d{6})$'
        )  # name_YYYYMMDD_HHMMSS_microseconds

    @activity.defn(name='cleanup_temp_directories')
    async def cleanup_temp_directories(self, input_data: dict[str, Any]) -> None:
        """
        Clean up stale temporary directories based on timestamp in directory name.

        This activity scans the reports/temp directory for subdirectories following
        the pattern '{name}_{timestamp}' where timestamp is in YYYYMMDD_HHMMSS_microseconds format.
        Directories older than the retention period are deleted.

        Args:
            input_data: Cleanup configuration containing:
                - metadata (dict): Workflow execution metadata
                - temp_path (str): Path to temp directory (optional, defaults to reports/temp)

        Returns:
            None: Results are logged and tracked via metrics

        Raises:
            Exception: If cleanup fails (after sending notification)
        """
        metadata = input_data.get('metadata', {})
        temp_path = input_data.get('temp_path', 'model_manager/reports/temp')
        metrics_status = 'success'

        cutoff_time = datetime.now() - timedelta(hours=self.retention_hours)

        try:
            self.info(
                f'Starting local directory cleanup - Path: {temp_path}, '
                f'Retention: {self.retention_hours}h, Dry run: {self.dry_run}',
                metadata,
            )

            if not os.path.exists(temp_path):
                self.warning(f'Temp directory does not exist: {temp_path}', metadata)
                return

            directories_scanned = 0
            directories_deleted = 0
            errors = []

            for item_name in os.listdir(temp_path):
                item_path = os.path.join(temp_path, item_name)

                if not os.path.isdir(item_path):
                    continue

                directories_scanned += 1

                # Extract timestamp from directory name
                match = self.dir_timestamp_pattern.match(item_name)
                if not match:
                    self.debug(
                        f'Skipping directory without timestamp pattern: {item_name}', metadata
                    )
                    continue

                timestamp_str = match.group(2)
                try:
                    # Parse YYYYMMDD_HHMMSS_microseconds
                    dir_time = datetime.strptime(timestamp_str, '%Y%m%d_%H%M%S_%f')

                    if dir_time < cutoff_time:
                        age_hours = (datetime.now() - dir_time).total_seconds() / 3600

                        if self.dry_run:
                            self.info(
                                f'[DRY RUN] Would delete directory: {item_name} (age: {age_hours:.1f}h)',
                                metadata,
                            )
                            directories_deleted += 1
                        else:
                            try:
                                shutil.rmtree(item_path)
                                self.info(
                                    f'Deleted stale directory: {item_name} (age: {age_hours:.1f}h)',
                                    metadata,
                                )
                                directories_deleted += 1
                            except OSError as e:
                                error_msg = f'Failed to delete directory {item_name}: {str(e)}'
                                errors.append(error_msg)
                                self.error(error_msg, metadata)
                    else:
                        age_hours = (datetime.now() - dir_time).total_seconds() / 3600
                        self.debug(
                            f'Keeping recent directory: {item_name} (age: {age_hours:.1f}h)',
                            metadata,
                        )

                except ValueError as e:
                    error_msg = f'Invalid timestamp format in directory {item_name}: {str(e)}'
                    errors.append(error_msg)
                    self.error(error_msg, metadata)

            self.info(
                f'Directory cleanup completed - Scanned: {directories_scanned}, '
                f'Deleted: {directories_deleted}, Errors: {len(errors)}',
                metadata,
            )

        except Exception as e:
            metrics_status = 'error'
            error_msg = f'Error in directory cleanup: {str(e)}'
            trace = traceback.format_exc()

            self.send_notification(
                metadata=metadata,
                notification_id='CLEANUP_DIRECTORIES_ERROR',
                message=error_msg,
                block='cleanup_temp_directories',
                level=NotificationLevel.ERROR,
                attachment_content=trace,
            )

            raise
        finally:
            await self._emit_metrics(
                metadata=metadata,
                metrics_status=metrics_status,
                activity_name='cleanup_temp_directories',
                emit_workflow_metric=True,
            )

    async def _emit_metrics(
        self,
        metadata: dict[str, Any],
        metrics_status: str,
        activity_name: str,
        emit_workflow_metric: bool,
    ) -> None:
        """
        Emit workflow and activity execution metrics.

        Args:
            metadata: Activity metadata containing pod_id and workflow_name
            metrics_status: Execution status ('success' or 'error')
            activity_name: Name of the activity being executed
        """
        if emit_workflow_metric:
            await self.emit_metric(
                metric_object=WORKFLOW_EXECUTION_TOTAL,
                tags={
                    'pod_id': metadata.get('pod_id'),
                    'workflow_name': metadata.get('workflow_name'),
                    'status': metrics_status,
                },
            )

        await self.emit_metric(
            metric_object=ACTIVITY_EXECUTION_TOTAL,
            tags={
                'pod_id': metadata.get('pod_id'),
                'activity_name': activity_name,
                'status': metrics_status,
            },
        )