#!/usr/bin/env python3 """ Swarm Scheduler - Story 4 ========================= Time-window based scheduler for the capability discovery swarm. Manages: - 8pm-4am AEST execution window - Agent scheduling and coordination - Graceful shutdown at window end """ import asyncio from datetime import datetime, timedelta from typing import Callable, Optional, List, Dict, Any from dataclasses import dataclass from enum import Enum import pytz import logging class SchedulerState(Enum): STOPPED = "stopped" WAITING = "waiting" RUNNING = "running" SHUTTING_DOWN = "shutting_down" @dataclass class ScheduleConfig: """Scheduler configuration.""" start_hour: int = 20 # 8pm end_hour: int = 4 # 4am timezone: str = "Australia/Sydney" check_interval_seconds: int = 60 graceful_shutdown_minutes: int = 15 class SwarmScheduler: """ Time-window scheduler for the capability swarm. Usage: scheduler = SwarmScheduler() async def my_swarm_run(): # Swarm execution logic pass # Run the scheduler (blocks until stopped) await scheduler.run(my_swarm_run) # Or check if in window manually if scheduler.is_in_window(): await my_swarm_run() """ def __init__(self, config: Optional[ScheduleConfig] = None): self.config = config or ScheduleConfig() self.state = SchedulerState.STOPPED self.tz = pytz.timezone(self.config.timezone) self.logger = logging.getLogger("swarm.scheduler") self._shutdown_event = asyncio.Event() self._current_task: Optional[asyncio.Task] = None def get_aest_time(self) -> datetime: """Get current time in AEST.""" return datetime.now(self.tz) def is_in_window(self) -> bool: """Check if current time is within the execution window.""" now = self.get_aest_time() hour = now.hour # Window: 20:00 (8pm) to 04:00 (4am) return hour >= self.config.start_hour or hour < self.config.end_hour def is_near_window_end(self) -> bool: """Check if we're within graceful shutdown period.""" now = self.get_aest_time() end_time = now.replace(hour=self.config.end_hour, minute=0, second=0, microsecond=0) # Handle day boundary if now.hour >= self.config.start_hour: end_time += timedelta(days=1) time_to_end = (end_time - now).total_seconds() / 60 return time_to_end <= self.config.graceful_shutdown_minutes def time_until_window_start(self) -> timedelta: """Get time until next window start.""" now = self.get_aest_time() start_time = now.replace(hour=self.config.start_hour, minute=0, second=0, microsecond=0) if now.hour >= self.config.start_hour or now.hour < self.config.end_hour: # Already in window or past today's start if now.hour < self.config.end_hour: # It's between midnight and 4am, window started yesterday return timedelta(0) return timedelta(0) # Before today's window if now.hour < self.config.start_hour: return start_time - now return timedelta(hours=24) - (now - start_time) def time_until_window_end(self) -> timedelta: """Get time until current window ends.""" if not self.is_in_window(): return timedelta(0) now = self.get_aest_time() end_time = now.replace(hour=self.config.end_hour, minute=0, second=0, microsecond=0) if now.hour >= self.config.start_hour: end_time += timedelta(days=1) return end_time - now async def run(self, swarm_callback: Callable[[], Any], on_window_start: Optional[Callable] = None, on_window_end: Optional[Callable] = None): """ Run the scheduler continuously. Args: swarm_callback: Async function to call during window on_window_start: Optional callback when window opens on_window_end: Optional callback when window closes """ self.state = SchedulerState.WAITING self.logger.info(f"Scheduler started. Window: {self.config.start_hour}:00 - {self.config.end_hour}:00 {self.config.timezone}") while not self._shutdown_event.is_set(): if self.is_in_window(): if self.state == SchedulerState.WAITING: self.state = SchedulerState.RUNNING self.logger.info("Entering execution window") if on_window_start: await self._safe_call(on_window_start) # Check for graceful shutdown if self.is_near_window_end(): if self.state != SchedulerState.SHUTTING_DOWN: self.state = SchedulerState.SHUTTING_DOWN self.logger.info("Entering graceful shutdown period") else: # Run swarm await self._safe_call(swarm_callback) else: if self.state == SchedulerState.RUNNING or self.state == SchedulerState.SHUTTING_DOWN: self.state = SchedulerState.WAITING self.logger.info("Execution window closed") if on_window_end: await self._safe_call(on_window_end) # Wait until next window wait_time = min( self.time_until_window_start().total_seconds(), self.config.check_interval_seconds ) self.logger.debug(f"Outside window. Waiting {wait_time:.0f}s") await asyncio.sleep(self.config.check_interval_seconds) self.state = SchedulerState.STOPPED self.logger.info("Scheduler stopped") async def _safe_call(self, callback: Callable): """Safely call a callback with error handling.""" try: result = callback() if asyncio.iscoroutine(result): await result except Exception as e: self.logger.exception(f"Callback error: {e}") def stop(self): """Signal the scheduler to stop.""" self._shutdown_event.set() def get_status(self) -> Dict[str, Any]: """Get scheduler status.""" now = self.get_aest_time() return { "state": self.state.value, "current_time_aest": now.isoformat(), "in_window": self.is_in_window(), "near_window_end": self.is_near_window_end(), "time_until_window_start": str(self.time_until_window_start()), "time_until_window_end": str(self.time_until_window_end()) if self.is_in_window() else "N/A", "config": { "start_hour": self.config.start_hour, "end_hour": self.config.end_hour, "timezone": self.config.timezone } } # Quick test if __name__ == "__main__": import json scheduler = SwarmScheduler() status = scheduler.get_status() print("Scheduler Status:") print(json.dumps(status, indent=2, default=str)) print(f"\nIn window: {scheduler.is_in_window()}") print(f"Time until window: {scheduler.time_until_window_start()}")