competitions/api.py

import asyncio
import os
import signal
import sqlite3
from contextlib import asynccontextmanager

import psutil
from fastapi import FastAPI
from loguru import logger

from competitions.utils import run_evaluation


def get_process_status(pid):
    try:
        process = psutil.Process(pid)
        proc_status = process.status()
        return proc_status
    except psutil.NoSuchProcess:
        logger.info(f"No process found with PID: {pid}")
        return "Completed"


def kill_process_by_pid(pid):
    """Kill process by PID."""
    os.kill(pid, signal.SIGTERM)


class JobDB:
    def __init__(self, db_path):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        self.c = self.conn.cursor()
        self.create_jobs_table()

    def create_jobs_table(self):
        self.c.execute(
            """CREATE TABLE IF NOT EXISTS jobs
            (id INTEGER PRIMARY KEY, pid INTEGER)"""
        )
        self.conn.commit()

    def add_job(self, pid):
        sql = f"INSERT INTO jobs (pid) VALUES ({pid})"
        self.c.execute(sql)
        self.conn.commit()

    def get_running_jobs(self):
        self.c.execute("""SELECT pid FROM jobs""")
        running_pids = self.c.fetchall()
        running_pids = [pid[0] for pid in running_pids]
        return running_pids

    def delete_job(self, pid):
        sql = f"DELETE FROM jobs WHERE pid={pid}"
        self.c.execute(sql)
        self.conn.commit()


PARAMS = os.environ.get("PARAMS")
DB = JobDB("job.db")


class BackgroundRunner:
    async def run_main(self):
        while True:
            running_jobs = DB.get_running_jobs()
            if running_jobs:
                for _pid in running_jobs:
                    proc_status = get_process_status(_pid)
                    proc_status = proc_status.strip().lower()
                    if proc_status in ("completed", "error", "zombie"):
                        logger.info(f"Process {_pid} is already completed. Skipping...")
                        try:
                            kill_process_by_pid(_pid)
                        except Exception as e:
                            logger.info(f"Error while killing process: {e}")
                        DB.delete_job(_pid)

            running_jobs = DB.get_running_jobs()
            if not running_jobs:
                logger.info("No running jobs found. Shutting down the server.")
                os.kill(os.getpid(), signal.SIGINT)
            await asyncio.sleep(30)


runner = BackgroundRunner()


@asynccontextmanager
async def lifespan(app: FastAPI):
    process_pid = run_evaluation(params=PARAMS)
    logger.info(f"Started training with PID {process_pid}")
    DB.add_job(process_pid)
    asyncio.create_task(runner.run_main())
    yield


api = FastAPI(lifespan=lifespan)


@api.get("/")
async def root():
    return "Your model is being evaluated..."


@api.get("/health")
async def health():
    return "OK"