Code import - branch release/SIENTIAPDE-1645

model_manager/schedules/cleanup_schedule.py

@@ -0,0 +1,160 @@
+"""Schedule configuration for cleanup workflow."""
+
+import os
+from datetime import timedelta
+from typing import Any
+
+from sientia_do.observability.logger import Logger as SientiaLogger
+from temporalio.client import (
+    Client,
+    Schedule,
+    ScheduleActionStartWorkflow,
+    ScheduleSpec,
+)
+
+from model_manager.worker.prepare_worker import build_queue_name
+
+# Schedule configuration from environment variables
+CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 0 * * *')  # Default: midnight UTC
+CLEANUP_TIMEZONE = os.getenv('CLEANUP_TIMEZONE', 'UTC')
+CLEANUP_EXECUTION_TIMEOUT_HOURS = int(os.getenv('CLEANUP_EXECUTION_TIMEOUT_HOURS', '1'))
+
+
+def build_cleanup_schedule_id(runtime: str | None) -> str:
+    """
+    Build cleanup schedule ID using runtime-derived naming.
+
+    Args:
+        runtime: Runtime suffix used by workers
+
+    Return:
+        str: Cleanup schedule ID
+    """
+    normalized_runtime = runtime.strip() if runtime else ''
+    return f'cleanup-files-{normalized_runtime or "single"}-daily'
+
+
+async def _needs_schedule_reconcile(
+    schedule_handle: Any,
+    cleanup_task_queue: str,
+    logger: SientiaLogger,
+    metadata: dict[str, str | None],
+) -> bool:
+    """
+    Compare configured cleanup schedule against current expected values.
+
+    Args:
+        schedule_handle: Temporal schedule handle for current schedule ID
+        cleanup_task_queue: Expected cleanup task queue
+        logger: Logger instance for errors
+        metadata: Metadata dictionary for logging context
+
+    Return:
+        bool: True when schedule should be recreated to apply current config
+    """
+    try:
+        schedule_description = await schedule_handle.describe()
+        schedule = getattr(schedule_description, 'schedule', None)
+        action = getattr(schedule, 'action', None)
+        spec = getattr(schedule, 'spec', None)
+
+        current_task_queue = getattr(action, 'task_queue', None)
+        current_execution_timeout = getattr(action, 'execution_timeout', None)
+        current_cron = getattr(spec, 'cron_expressions', None)
+        current_timezone = getattr(spec, 'time_zone_name', None)
+
+        return (
+            current_task_queue != cleanup_task_queue
+            or current_execution_timeout != timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS)
+            or current_cron != [CLEANUP_CRON]
+            or current_timezone != CLEANUP_TIMEZONE
+        )
+    except Exception as e:  # noqa: BLE001
+        logger.custom_error(f'Error describing cleanup schedule for reconcile: {e}', metadata)
+        return True
+
+
+async def schedule_exists(
+    client: Client, schedule_id: str, logger: SientiaLogger, metadata: dict[str, str | None]
+) -> bool:
+    """
+    Check if a schedule already exists.
+
+    Args:
+        client: Temporal client instance
+        schedule_id: ID of the schedule to check
+        logger: Logger instance for error logging
+        metadata: Metadata dictionary for logging context
+
+    Returns:
+        True if schedule exists, False otherwise
+    """
+    try:
+        async for schedule in await client.list_schedules():
+            if schedule.id == schedule_id:
+                return True
+        return False
+    except Exception as e:  # noqa: BLE001
+        logger.custom_error(f'Error checking if schedule exists: {e}', metadata)
+        return False
+
+
+async def create_cleanup_schedule(
+    client: Client, logger: SientiaLogger, metadata: dict[str, str | None]
+) -> None:
+    """
+    Create or update the cleanup files schedule.
+
+    This function is idempotent and can be called multiple times safely.
+    It will only create the schedule if it doesn't already exist.
+
+    Args:
+        client: Temporal client instance
+        logger: Logger instance for logging schedule operations
+        metadata: Metadata dictionary for logging context
+    """
+    runtime = (os.getenv('RUNTIME') or 'single').strip()
+    cleanup_task_queue = build_queue_name('CleanupFiles', runtime or 'single')
+    schedule_id = build_cleanup_schedule_id(runtime)
+    updated = False
+
+    if await schedule_exists(client, schedule_id, logger, metadata):
+        handle = client.get_schedule_handle(schedule_id)
+        if await _needs_schedule_reconcile(handle, cleanup_task_queue, logger, metadata):
+            await handle.delete()
+            updated = True
+        else:
+            logger.custom_info(
+                f"Schedule '{schedule_id}' is already up to date, no-op reconcile",
+                metadata,
+            )
+            return
+
+    await client.create_schedule(
+        schedule_id,
+        Schedule(
+            action=ScheduleActionStartWorkflow(
+                'cleanup_files',
+                {},  # Empty input, will use default bucket from environment
+                id=f'cleanup-files-scheduled-{schedule_id}',
+                task_queue=cleanup_task_queue,
+                execution_timeout=timedelta(hours=CLEANUP_EXECUTION_TIMEOUT_HOURS),
+            ),
+            spec=ScheduleSpec(
+                cron_expressions=[CLEANUP_CRON],
+                time_zone_name=CLEANUP_TIMEZONE,
+            ),
+        ),
+    )
+    if updated:
+        logger.custom_info(
+            f"Schedule '{schedule_id}' reconciled successfully. "
+            f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
+            metadata,
+        )
+    else:
+        logger.custom_info(
+            f"Schedule '{schedule_id}' created successfully. "
+            f'Cleanup will run at: {CLEANUP_CRON} ({CLEANUP_TIMEZONE})',
+            metadata,
+        )