ASE/elab_orchestrator.py

#!.venv/bin/python

# Import necessary libraries
import logging
import asyncio
import os
import aiomysql
import contextvars


# Import custom modules for configuration and database connection
from utils.config import loader_matlab_elab as setting
from utils.database import DATA_LOADED
from utils.database.matlab_query import get_matlab_command
from utils.csv.loaders import get_next_csv_atomic


# Crea una context variable per identificare il worker
worker_context = contextvars.ContextVar('worker_id', default='^-^')

# Formatter personalizzato che include il worker_id
class WorkerFormatter(logging.Formatter):
    """Formatter personalizzato che include l'ID del worker nei log."""
    def format(self, record: str) -> str:
        """Formatta il record di log includendo l'ID del worker.

        Args:
            record (str): Il record di log da formattare.

        Returns:
            La stringa formattata del record di log.
        """
        record.worker_id = worker_context.get()
        return super().format(record)

# Initialize the logger for this module
logger = logging.getLogger()

# Delay tra un processamento CSV e il successivo (in secondi)
ELAB_PROCESSING_DELAY = 0.2
# Tempo di attesa se non ci sono record da elaborare
NO_RECORD_SLEEP = 60

async def worker(worker_id: int, cfg: object, pool: object) -> None:
    """Esegue il ciclo di lavoro per l'elaborazione dei dati caricati.

    Il worker preleva un record dal database che indica dati pronti per
    l'elaborazione, esegue un comando Matlab associato e attende
    prima di iniziare un nuovo ciclo.

    Args:
        worker_id (int): L'ID univoco del worker.
        cfg (object): L'oggetto di configurazione.
        pool (object): Il pool di connessioni al database.
    """
    # Imposta il context per questo worker
    worker_context.set(f"W{worker_id:02d}")

    debug_mode = (logging.getLogger().getEffectiveLevel() == logging.DEBUG)
    logger.info("Avviato")

    while True:
        try:
            logger.info("Inizio elaborazione")

            record = await get_next_csv_atomic(pool, cfg.dbrectable, DATA_LOADED)

            if record:
                id, unit_type, tool_type, unit_name, tool_name = [x.lower().replace(" ", "_") if isinstance(x, str) else x for x in record]
                matlab_info = await get_matlab_command(cfg, tool_name, unit_name, pool)
                if matlab_info:
                    matlab_cmd = f"timeout {cfg.matlab_timeout} ./run_{matlab_info['matcall']}.sh {cfg.matlab_runtime} {unit_name} {tool_name}"

                    # matlab_error_filename = f'{cfg.matlab_error_path}{unit_name}{tool_name}_output_error.txt'

                    proc = await asyncio.create_subprocess_shell(
                        matlab_cmd,
                        cwd=cfg.matlab_func_path,
                        stdout=asyncio.subprocess.PIPE,
                        stderr=asyncio.subprocess.PIPE
                    )

                    stdout, stderr = await proc.communicate()

                    if proc.returncode != 0:
                        logger.error("Errore durante l'elaborazione")
                        logger.error(stderr.decode().strip())
                    logger.info(stdout.decode().strip())
                    await asyncio.sleep(ELAB_PROCESSING_DELAY)
            else:
                logger.info("Nessun record disponibile")
                await asyncio.sleep(NO_RECORD_SLEEP)

        except Exception as e:
            logger.error(f"Errore durante l'esecuzione: {e}", exc_info=debug_mode)
            await asyncio.sleep(1)


async def main():
    """Funzione principale che inizializza e avvia il sistema di elaborazione.

    Questa funzione si occupa di:
    - Caricare la configurazione.
    - Impostare il logging.
    - Creare un pool di connessioni al database.
    - Avviare i worker concorrenti per l'elaborazione.
    - Gestire l'arresto controllato del sistema.
    """
    logger.info("Avvio del sistema...")

    cfg = setting.Config()
    logger.info("Configurazione caricata correttamente")

    try:
        # Configura il logging globale
        log_level = os.getenv("LOG_LEVEL", "INFO").upper()
        debug_mode = (logging.getLogger().getEffectiveLevel() == logging.DEBUG)

        # Configura il logging con il formatter personalizzato
        handler = logging.FileHandler(cfg.logfilename)
        formatter = WorkerFormatter(
            "%(asctime)s - PID: %(process)d.Worker-%(worker_id)s.%(name)s.%(funcName)s.%(levelname)s: %(message)s"
        )
        handler.setFormatter(formatter)

        # Rimuovi eventuali handler esistenti e aggiungi il nostro
        logger.handlers.clear()
        logger.addHandler(handler)
        logger.setLevel(getattr(logging, log_level))

        logger.info("Logging configurato correttamente")

        # Numero massimo di worker concorrenti
        logger.info(f"Avvio di {cfg.max_threads} worker concorrenti")

        pool = await aiomysql.create_pool(
            host=cfg.dbhost,
            user=cfg.dbuser,
            password=cfg.dbpass,
            db=cfg.dbname,
            minsize=cfg.max_threads,
            maxsize=cfg.max_threads*4,
            pool_recycle=3600
        )

        # Avvia i worker
        workers = [
            asyncio.create_task(worker(i, cfg, pool))
            for i in range(cfg.max_threads)
        ]

        logger.info("Sistema avviato correttamente. In attesa di nuovi task...")

        try:
            await asyncio.gather(*workers, return_exceptions=debug_mode)
        finally:
            pool.close()
            await pool.wait_closed()

    except KeyboardInterrupt:
        logger.info("Info: Shutdown richiesto... chiusura in corso")

    except Exception as e:
        logger.error(f"Errore principale: {e}", exc_info=debug_mode)

if __name__ == "__main__":
    asyncio.run(main())