Initial commit of integrated agent system
This commit is contained in:
cillin
362
backend/core/scheduler.py
Normal file
362
backend/core/scheduler.py
Normal file
@@ -0,0 +1,362 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Scheduler - Market-aware trigger system for trading cycles
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from datetime import datetime, time, timedelta
|
||||
from typing import Any, Callable, Optional
|
||||
from zoneinfo import ZoneInfo
|
||||
|
||||
import pandas_market_calendars as mcal
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# NYSE timezone for US stock trading
|
||||
NYSE_TZ = ZoneInfo("America/New_York")
|
||||
NYSE_CALENDAR = mcal.get_calendar("NYSE")
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""
|
||||
Market-aware scheduler for live trading.
|
||||
Uses NYSE timezone and trading calendar.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mode: str = "daily",
|
||||
trigger_time: Optional[str] = None,
|
||||
interval_minutes: Optional[int] = None,
|
||||
heartbeat_interval: Optional[int] = None,
|
||||
config: Optional[dict] = None,
|
||||
):
|
||||
self.mode = mode
|
||||
self.trigger_time = trigger_time or "09:30" # NYSE timezone
|
||||
self.trigger_now = self.trigger_time == "now"
|
||||
self.interval_minutes = interval_minutes or 60
|
||||
self.heartbeat_interval = heartbeat_interval # e.g. 3600 = 1 hour
|
||||
self.config = config or {}
|
||||
|
||||
self.running = False
|
||||
self._task: Optional[asyncio.Task] = None
|
||||
self._heartbeat_task: Optional[asyncio.Task] = None
|
||||
self._callback: Optional[Callable] = None
|
||||
self._heartbeat_callback: Optional[Callable] = None
|
||||
|
||||
def _now_nyse(self) -> datetime:
|
||||
"""Get current time in NYSE timezone"""
|
||||
return datetime.now(NYSE_TZ)
|
||||
|
||||
def _is_trading_day(self, date: datetime) -> bool:
|
||||
"""Check if date is a NYSE trading day"""
|
||||
date_str = date.strftime("%Y-%m-%d")
|
||||
valid_days = NYSE_CALENDAR.valid_days(
|
||||
start_date=date_str,
|
||||
end_date=date_str,
|
||||
)
|
||||
return len(valid_days) > 0
|
||||
|
||||
def _is_trading_hours(self, now: datetime) -> bool:
|
||||
"""Check if current time is within NYSE trading hours (9:30-16:00 ET)."""
|
||||
market_time = now.time()
|
||||
return time(9, 30) <= market_time <= time(16, 0)
|
||||
|
||||
def set_heartbeat_callback(self, callback: Callable) -> None:
|
||||
"""Register callback for heartbeat triggers."""
|
||||
self._heartbeat_callback = callback
|
||||
|
||||
def _next_trading_day(self, from_date: datetime) -> datetime:
|
||||
"""Find the next trading day from given date"""
|
||||
check_date = from_date
|
||||
for _ in range(10): # Max 10 days ahead (handles holidays)
|
||||
if self._is_trading_day(check_date):
|
||||
return check_date
|
||||
check_date += timedelta(days=1)
|
||||
return check_date
|
||||
|
||||
async def start(self, callback: Callable):
|
||||
"""Start scheduler"""
|
||||
if self.running:
|
||||
logger.warning("Scheduler already running")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
self._callback = callback
|
||||
self._schedule_task()
|
||||
|
||||
# Start heartbeat loop if configured
|
||||
if self.heartbeat_interval and self._heartbeat_callback:
|
||||
self._heartbeat_task = asyncio.create_task(self._run_heartbeat_loop())
|
||||
logger.info(
|
||||
f"Heartbeat loop started: interval={self.heartbeat_interval}s",
|
||||
)
|
||||
|
||||
logger.info(
|
||||
f"Scheduler started: mode={self.mode}, timezone=America/New_York",
|
||||
)
|
||||
|
||||
def _schedule_task(self):
|
||||
"""Create the active scheduler task for the current mode."""
|
||||
if not self._callback:
|
||||
raise ValueError("Scheduler callback is not set")
|
||||
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
self._task = None
|
||||
|
||||
if self.mode == "daily":
|
||||
self._task = asyncio.create_task(self._run_daily(self._callback))
|
||||
elif self.mode == "intraday":
|
||||
self._task = asyncio.create_task(
|
||||
self._run_intraday(self._callback),
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown scheduler mode: {self.mode}")
|
||||
|
||||
def reconfigure(
|
||||
self,
|
||||
*,
|
||||
mode: Optional[str] = None,
|
||||
trigger_time: Optional[str] = None,
|
||||
interval_minutes: Optional[int] = None,
|
||||
) -> bool:
|
||||
"""Update scheduler parameters in-place and restart its timing loop."""
|
||||
changed = False
|
||||
|
||||
if mode and mode != self.mode:
|
||||
self.mode = mode
|
||||
changed = True
|
||||
|
||||
if trigger_time and trigger_time != self.trigger_time:
|
||||
self.trigger_time = trigger_time
|
||||
self.trigger_now = self.trigger_time == "now"
|
||||
changed = True
|
||||
|
||||
if (
|
||||
interval_minutes is not None
|
||||
and interval_minutes > 0
|
||||
and interval_minutes != self.interval_minutes
|
||||
):
|
||||
self.interval_minutes = interval_minutes
|
||||
changed = True
|
||||
|
||||
if changed and self.running and self._callback:
|
||||
self._schedule_task()
|
||||
logger.info(
|
||||
"Scheduler reconfigured: mode=%s, trigger_time=%s, interval_minutes=%s",
|
||||
self.mode,
|
||||
self.trigger_time,
|
||||
self.interval_minutes,
|
||||
)
|
||||
|
||||
return changed
|
||||
|
||||
async def _run_heartbeat_loop(self):
|
||||
"""Run heartbeat checks on a separate interval during trading hours."""
|
||||
while self.running:
|
||||
now = self._now_nyse()
|
||||
if self._is_trading_day(now) and self._is_trading_hours(now):
|
||||
if self._heartbeat_callback:
|
||||
try:
|
||||
current_date = now.strftime("%Y-%m-%d")
|
||||
logger.debug(
|
||||
f"[Heartbeat] Triggering heartbeat check for {current_date}",
|
||||
)
|
||||
await self._heartbeat_callback(date=current_date)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"[Heartbeat] Callback failed: {e}",
|
||||
exc_info=True,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"[Heartbeat] Callback not set, skipping heartbeat",
|
||||
)
|
||||
|
||||
await asyncio.sleep(self.heartbeat_interval)
|
||||
|
||||
async def _run_daily(self, callback: Callable):
|
||||
"""Run once per trading day at specified time (NYSE timezone)"""
|
||||
first_run = True
|
||||
|
||||
while self.running:
|
||||
now = self._now_nyse()
|
||||
|
||||
# Handle "now" trigger - run immediately on first iteration
|
||||
if self.trigger_now and first_run:
|
||||
first_run = False
|
||||
current_date = now.strftime("%Y-%m-%d")
|
||||
logger.info(f"Triggering immediately for {current_date}")
|
||||
await callback(date=current_date)
|
||||
# After immediate run, stop (one-shot mode)
|
||||
self.running = False
|
||||
break
|
||||
|
||||
target_time = datetime.strptime(self.trigger_time, "%H:%M").time()
|
||||
|
||||
# Calculate next trigger datetime
|
||||
if now.time() < target_time:
|
||||
next_run = now.replace(
|
||||
hour=target_time.hour,
|
||||
minute=target_time.minute,
|
||||
second=0,
|
||||
microsecond=0,
|
||||
)
|
||||
else:
|
||||
next_run = (now + timedelta(days=1)).replace(
|
||||
hour=target_time.hour,
|
||||
minute=target_time.minute,
|
||||
second=0,
|
||||
microsecond=0,
|
||||
)
|
||||
|
||||
# Skip to next trading day
|
||||
next_run = self._next_trading_day(next_run)
|
||||
next_run = next_run.replace(
|
||||
hour=target_time.hour,
|
||||
minute=target_time.minute,
|
||||
second=0,
|
||||
microsecond=0,
|
||||
)
|
||||
|
||||
wait_seconds = (next_run - now).total_seconds()
|
||||
logger.info(
|
||||
f"Next trigger: {next_run.strftime('%Y-%m-%d %H:%M %Z')} "
|
||||
f"(in {wait_seconds/3600:.1f} hours)",
|
||||
)
|
||||
|
||||
await asyncio.sleep(wait_seconds)
|
||||
|
||||
current_date = self._now_nyse().strftime("%Y-%m-%d")
|
||||
logger.info(f"Triggering daily cycle for {current_date}")
|
||||
await callback(date=current_date)
|
||||
|
||||
async def _run_intraday(self, callback: Callable):
|
||||
"""Run every N minutes (for future use)"""
|
||||
while self.running:
|
||||
now = self._now_nyse()
|
||||
current_date = now.strftime("%Y-%m-%d")
|
||||
|
||||
if self._is_trading_day(now):
|
||||
logger.info(f"Triggering intraday cycle for {current_date}")
|
||||
await callback(date=current_date)
|
||||
|
||||
await asyncio.sleep(self.interval_minutes * 60)
|
||||
|
||||
def stop(self):
|
||||
"""Stop scheduler"""
|
||||
self.running = False
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
self._task = None
|
||||
if self._heartbeat_task:
|
||||
self._heartbeat_task.cancel()
|
||||
self._heartbeat_task = None
|
||||
logger.info("Scheduler stopped")
|
||||
|
||||
|
||||
class BacktestScheduler:
|
||||
"""Backtest Scheduler - Runs through historical trading dates"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
start_date: str,
|
||||
end_date: str,
|
||||
trading_calendar: Optional[Any] = None,
|
||||
delay_between_days: float = 0.1,
|
||||
):
|
||||
self.start_date = start_date
|
||||
self.end_date = end_date
|
||||
self.trading_calendar = trading_calendar
|
||||
self.delay_between_days = delay_between_days
|
||||
|
||||
self.running = False
|
||||
self._task: Optional[asyncio.Task] = None
|
||||
self._dates: list = []
|
||||
|
||||
def get_trading_dates(self) -> list:
|
||||
"""Get list of trading dates in the backtest period"""
|
||||
import pandas as pd
|
||||
|
||||
start = pd.to_datetime(self.start_date)
|
||||
end = pd.to_datetime(self.end_date)
|
||||
|
||||
if self.trading_calendar:
|
||||
calendar = mcal.get_calendar(self.trading_calendar)
|
||||
trading_dates = calendar.valid_days(
|
||||
start_date=self.start_date,
|
||||
end_date=self.end_date,
|
||||
)
|
||||
dates = [d.strftime("%Y-%m-%d") for d in trading_dates]
|
||||
else:
|
||||
all_dates = pd.date_range(start=start, end=end, freq="D")
|
||||
dates = [
|
||||
d.strftime("%Y-%m-%d") for d in all_dates if d.weekday() < 5
|
||||
]
|
||||
|
||||
self._dates = dates
|
||||
return dates
|
||||
|
||||
async def start(self, callback: Callable):
|
||||
"""Start async backtest scheduler"""
|
||||
if self.running:
|
||||
logger.warning("Backtest scheduler already running")
|
||||
return
|
||||
|
||||
self.running = True
|
||||
dates = self.get_trading_dates()
|
||||
|
||||
logger.info(
|
||||
f"Starting backtest: {self.start_date} to {self.end_date} "
|
||||
f"({len(dates)} trading days)",
|
||||
)
|
||||
|
||||
self._task = asyncio.create_task(self._run_async(callback, dates))
|
||||
|
||||
async def _run_async(self, callback: Callable, dates: list):
|
||||
"""Run backtest asynchronously"""
|
||||
for i, date in enumerate(dates, 1):
|
||||
if not self.running:
|
||||
break
|
||||
|
||||
logger.info(f"[{i}/{len(dates)}] Processing {date}")
|
||||
await callback(date=date)
|
||||
|
||||
if self.delay_between_days > 0:
|
||||
await asyncio.sleep(self.delay_between_days)
|
||||
|
||||
logger.info("Backtest complete")
|
||||
self.running = False
|
||||
|
||||
def run(self, callback: Callable, **kwargs):
|
||||
"""Run backtest synchronously through all trading dates"""
|
||||
dates = self.get_trading_dates()
|
||||
results = []
|
||||
|
||||
logger.info(
|
||||
f"Starting backtest: {self.start_date} to {self.end_date} "
|
||||
f"({len(dates)} trading days)",
|
||||
)
|
||||
|
||||
for i, date in enumerate(dates, 1):
|
||||
logger.info(f"[{i}/{len(dates)}] Processing {date}")
|
||||
result = callback(date=date, **kwargs)
|
||||
results.append({"date": date, "result": result})
|
||||
|
||||
logger.info("Backtest complete")
|
||||
return results
|
||||
|
||||
def stop(self):
|
||||
"""Stop backtest scheduler"""
|
||||
self.running = False
|
||||
if self._task:
|
||||
self._task.cancel()
|
||||
self._task = None
|
||||
logger.info("Backtest scheduler stopped")
|
||||
|
||||
def get_total_days(self) -> int:
|
||||
"""Get total number of trading days"""
|
||||
if not self._dates:
|
||||
self.get_trading_dates()
|
||||
return len(self._dates)
|
||||
Reference in New Issue
Block a user