gisaf-backend/src/gisaf/scheduler.py

328 lines
10 KiB
Python
Raw Normal View History

#!/usr/bin/env python
"""
Gisaf task scheduler, orchestrating the background tasks
like remote device data collection, etc.
"""
2024-05-09 00:04:53 +02:00
import os
2024-05-09 00:04:53 +02:00
import psutil
import logging
import sys
import asyncio
from json import dumps
from datetime import datetime
from importlib.metadata import entry_points
2024-05-09 01:42:11 +02:00
from typing import Any, List
from fastapi import FastAPI
from pydantic_settings import BaseSettings, SettingsConfigDict
# from apscheduler import SchedulerStarted
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from apscheduler.triggers.date import DateTrigger
# from gisaf.ipynb_tools import Gisaf
formatter = logging.Formatter(
2024-05-09 00:04:53 +02:00
"%(asctime)s:%(levelname)s:%(name)s:%(message)s", "%Y-%m-%d %H:%M:%S"
)
for handler in logging.root.handlers:
handler.setFormatter(formatter)
2024-05-09 00:04:53 +02:00
logger = logging.getLogger("gisaf.scheduler")
class Settings(BaseSettings):
2024-05-09 00:04:53 +02:00
model_config = SettingsConfigDict(env_prefix="gisaf_scheduler_")
app_name: str = "Gisaf scheduler"
job_names: List[str] = []
exclude_job_names: List[str] = []
list: bool = False
class JobBaseClass:
"""
Base class for all the jobs.
"""
2024-05-09 00:04:53 +02:00
task_id = None
interval = None
cron = None
enabled = True
2024-05-09 00:04:53 +02:00
type = "" ## interval, cron or longrun
sched_params = ""
name = "<unnammed task>"
features = None
2024-05-09 00:04:53 +02:00
def __init__(self):
self.last_run = None
self.current_run = None
async def get_feature_ids(self):
"""
Subclasses might define a get_features function to inform the
front-ends about the map features it works on.
The scheduler runs this on startup.
"""
return []
async def run(self):
"""
Subclasses should define a run async function to run
"""
2024-05-09 00:04:53 +02:00
logger.info(f"Noop defined for {self.name}")
class JobScheduler:
# gs: Gisaf
jobs: dict[str, Any]
2024-05-09 00:04:53 +02:00
# tasks: dict[str, Any]
wss: dict[str, Any]
subscribers: set[Any]
scheduler: AsyncIOScheduler
2024-05-09 00:04:53 +02:00
def __init__(self):
2024-05-09 00:04:53 +02:00
# self.redis_store = gs.app['store']
self.jobs = {}
2024-05-09 00:04:53 +02:00
# self.tasks = {}
self.wss = {}
self.subscribers = set()
self.scheduler = AsyncIOScheduler()
def start(self):
self.scheduler.start()
2024-05-09 00:04:53 +02:00
# def scheduler_event_listener(self, event):
# asyncio.create_task(self.scheduler_event_alistener(event))
2024-05-09 00:04:53 +02:00
# async def scheduler_event_alistener(self, event):
# if isinstance(event, SchedulerStarted):
# pid = os.getpid()
# logger.debug(f'Scheduler started, pid={pid}')
# #await self.gs.app['store'].pub.set('_scheduler/pid', pid)
async def job_event_added(self, event):
task = await self.scheduler.data_store.get_task(event.task_id)
2024-05-09 00:04:53 +02:00
schedules = [
ss
for ss in await self.scheduler.get_schedules()
if ss.task_id == event.task_id
]
if len(schedules) > 1:
2024-05-09 00:04:53 +02:00
logger.warning(f"More than 1 schedule matching task {event.task_id}")
return
else:
schedule = schedules[0]
async def job_acquired(self, event):
pass
async def job_cancelled(self, event):
pass
async def job_released(self, event):
pass
# task = self.tasks.get(event.job_id)
# if not task:
# breakpoint()
# logger.warning(f'Got an event {event} for unregistered task {event.task_id}')
# return
# if isinstance(event, apscheduler.JobCancelled): #events.EVENT_JOB_ERROR:
# msg = f'"{task.name}" cancelled ({task.task_id})'
# task.last_run = event
# task.current_run = None
# logger.warning(msg)
# ## TODO: try to restart the task
# elif isinstance(event, apscheduler.JobAcquired): #events.EVENT_JOB_SUBMITTED:
# ## XXX: should be task.last_run = None
# task.last_run = event
# task.current_run = event
# msg = f'"{task.name}" started ({task.task_id})'
# elif isinstance(event, apscheduler.JobReleased): #events.EVENT_JOB_EXECUTED:
# task.last_run = event
# task.current_run = None
# msg = f'"{task.name}" worked ({task.task_id})'
# else:
# logger.info(f'*********** Unhandled event: {event}')
# pass
# #await self.send_to_redis_store(task, event, msg)
# ## Send to notification graphql websockets subscribers
# for queue in self.subscribers:
# queue.put_nowait((task, event))
# ## Send raw messages through websockets
# await self.send_to_websockets(task, event, msg)
async def send_to_redis_store(self, job, event, msg):
"""
Send to Redis store
"""
try:
2024-05-09 00:04:53 +02:00
self.gs.app["store"].pub.publish(
"admin:scheduler:json", dumps({"msg": msg})
)
except Exception as err:
logger.warning(f'Cannot publish updates for "{job.name}" to Redis: {err}')
logger.exception(err)
async def send_to_websockets(self, job, event, msg):
"""
Send to all connected websockets
"""
for ws in self.wss.values():
2024-05-09 00:04:53 +02:00
asyncio.create_task(ws.send_json({"msg": msg}))
def add_subscription(self, ws):
self.wss[id(ws)] = ws
def delete_subscription(self, ws):
del self.wss[id(ws)]
def get_available_jobs(self):
return [
entry_point.name
2024-05-09 00:04:53 +02:00
for entry_point in entry_points().select(group="gisaf_jobs")
]
async def setup(self, job_names=None, exclude_job_names=None):
if job_names is None:
job_names = []
if exclude_job_names is None:
exclude_job_names = []
## Go through entry points and define the tasks
2024-05-09 00:04:53 +02:00
for entry_point in entry_points().select(group="gisaf_jobs"):
## Eventually skip task according to arguments of the command line
2024-05-09 00:04:53 +02:00
if (entry_point.name in exclude_job_names) or (
(len(job_names) > 0) and entry_point.name not in job_names
):
logger.info(f"Skip task {entry_point.name}")
continue
try:
task_class = entry_point.load()
except Exception as err:
2024-05-09 00:04:53 +02:00
logger.error(f"Task {entry_point.name} skipped cannot be loaded: {err}")
continue
## Create the task instance
try:
2024-05-09 00:04:53 +02:00
task = task_class()
except Exception as err:
2024-05-09 00:04:53 +02:00
logger.error(f"Task {entry_point.name} cannot be instanciated: {err}")
continue
task.name = entry_point.name
if not task.enabled:
logger.debug(f'Job "{entry_point.name}" disabled')
continue
logger.debug(f'Add task "{entry_point.name}"')
2024-05-09 00:04:53 +02:00
if not hasattr(task, "run"):
logger.error(f"Task {entry_point.name} skipped: no run method")
continue
task.features = await task.get_feature_ids()
2024-05-09 00:04:53 +02:00
kwargs: dict[str, Any] = {
# 'tags': [entry_point.name],
}
if isinstance(task.interval, dict):
2024-05-09 00:04:53 +02:00
kwargs["trigger"] = IntervalTrigger(**task.interval)
kwargs["next_run_time"] = datetime.now()
task.type = "interval"
## TODO: format user friendly text for interval
task.sched_params = get_pretty_format_interval(task.interval)
elif isinstance(task.cron, dict):
## FIXME: CronTrigger
2024-05-09 00:04:53 +02:00
kwargs["trigger"] = CronTrigger(**task.cron)
kwargs.update(task.cron)
2024-05-09 00:04:53 +02:00
task.type = "cron"
## TODO: format user friendly text for cron
task.sched_params = get_pretty_format_cron(task.cron)
else:
2024-05-09 00:04:53 +02:00
task.type = "longrun"
task.sched_params = "always running"
kwargs["trigger"] = DateTrigger(datetime.now())
# task.task_id = await self.scheduler.add_job(task.run, **kwargs)
# self.tasks[task.task_id] = task
# continue
## Create the APScheduler task
try:
2024-05-09 00:04:53 +02:00
task.task_id = self.scheduler.add_job(
task.run,
name=entry_point.name,
**kwargs,
)
except Exception as err:
2024-05-09 00:04:53 +02:00
logger.warning(f"Cannot add task {entry_point.name}: {err}")
logger.exception(err)
else:
2024-05-09 00:04:53 +02:00
logger.info(f'Job "{entry_point.name}" added ({task.task_id.id})')
# self.tasks[task.task_id] = task
## Subscribe to all events
# self.scheduler.subscribe(self.job_acquired, JobAcquired)
# self.scheduler.subscribe(self.job_cancelled, JobCancelled)
# self.scheduler.subscribe(self.job_released, JobReleased)
# self.scheduler.subscribe(self.job_event_added, JobAdded)
# self.scheduler.subscribe(self.scheduler_event_listener, SchedulerEvent)
class GSFastAPI(FastAPI):
js: JobScheduler
2024-05-09 00:04:53 +02:00
allowed_interval_params = set(("seconds", "minutes", "hours", "days", "weeks"))
def get_pretty_format_interval(params):
"""
Return a format for describing interval
"""
2024-05-09 00:04:53 +02:00
return str({k: v for k, v in params.items() if k in allowed_interval_params})
def get_pretty_format_cron(params):
"""
Return a format for describing cron
"""
return str(params)
async def startup(settings):
if settings.list:
## Just print avalable jobs and exit
jobs = js.get_available_jobs()
2024-05-09 00:04:53 +02:00
print(" ".join(jobs))
# sys.exit(0)
parent_pid = os.getpid()
parent = psutil.Process(parent_pid)
for child in parent.children(
recursive=True
): # or parent.children() for recursive=False
child.kill()
parent.kill()
# try:
# await js.gs.setup()
# await js.gs.make_models()
# except Exception as err:
# logger.error('Cannot setup Gisaf')
# logger.exception(err)
# sys.exit(1)
try:
2024-05-09 00:04:53 +02:00
await js.setup(
job_names=settings.job_names, exclude_job_names=settings.exclude_job_names
)
except Exception as err:
2024-05-09 00:04:53 +02:00
logger.error("Cannot setup scheduler")
logger.exception(err)
sys.exit(1)
2024-05-09 00:04:53 +02:00
js = JobScheduler()