base/eco.py - dev - kotkov/eco

base/eco.py
119 строк5 KB
Arseny Kotkov
Add multithreading logging to file with log.ecolog
04 мар 2025, 08:01
04 мар 2025, 08:0103de008
import argparse
from threading import Thread

from SCHEDULER.scheduler import Scheduler
from SLURM.slurm import Slurm


from AAA.config import IP, PORT
from log import ecolog
from utils import Socket


def load_basic_configurations(args):
    # Определяем тип агента
    if args.type:
        if args.type.lower() == "slurm":
            agent_cls = Slurm
        if args.type.lower() == "scheduler":
            agent_cls = Scheduler

    # Инициализируем агента выбранного типа
    agent = agent_cls()

    # Если хост и порт явно указаны - используем их
    if args.name and args.ip and args.port:
        name, ip, port = str(args.name), str(args.ip), int(args.port)
        # Обновляем конфиг соответствующего агента
        if isinstance(agent, Slurm):
            with open('SLURM/config.py', 'w', encoding='utf-8') as file:
                file.write(
                    f"NAME = '{name}'\nIP = '{ip}'\nPORT = {port}\n")
            ecolog("Конфиг SLURM обновлён")
        elif isinstance(agent, Scheduler):
            with open('SCHEDULER/config.py', 'w', encoding='utf-8') as file:
                file.write(
                    f"NAME = '{name}'\nIP = '{ip}'\nPORT = {port}\n")
            ecolog("Конфиг SCHEDULER обновлён")
    # Иначе, если регистрируется агент типа Slurm
    elif isinstance(agent, Slurm):
        import SLURM.config
        name, ip, port = SLURM.config.NAME, SLURM.config.IP, SLURM.config.PORT
        ecolog("Использован конфиг SLURM")
    # Иначе, если регистрируется агент типа Scheduler
    elif isinstance(agent, Scheduler):
        import SCHEDULER.config
        name, ip, port = SCHEDULER.config.NAME, SCHEDULER.config.IP, SCHEDULER.config.PORT
        ecolog("Использован конфиг SCHEDULER")
    return agent, name, ip, port

def handle_startagent(args):
    if args.type:
        if args.type.lower() == "slurm":
            agent_cls = Slurm
        if args.type.lower() == "scheduler":
            agent_cls = Scheduler

    # Инициализируем агента выбранного типа
    agent = agent_cls()

    if args.type == 'SLURM':
        import SLURM.config
        name, ip, port = SLURM.config.NAME, SLURM.config.IP, SLURM.config.PORT
        ecolog("Использован конфиг SLURM")
    elif args.type == 'SCHEDULER':
        import SCHEDULER.config
        name, ip, port = SCHEDULER.config.NAME, SCHEDULER.config.IP, SCHEDULER.config.PORT
        ecolog("Использован конфиг SCHEDULER")

    # Запускаем агента на указанном порту, поскольку IP и порт агенту
    # присваиваются только в момент запуска метода serve().
    kwargs = {'blocking', True}
    # t0 = Thread(target=agent.serve, args=((ip, port)), kwargs=kwargs)
    # Запускаем в режиме демона, чтобы когда утилита регистрации отработала
    # агент бы выключился
    # t0.daemon = True
    # t0.start()
    agent.serve((ip,port), blocking=True)


def handle_registeragent(args):
    agent, name, ip, port = load_basic_configurations(args)

    # Запускаем агента на указанном порту, поскольку IP и порт агенту
    # присваиваются только в момент запуска метода serve().
    t0 = Thread(target=agent.serve, args=((ip, port),))
    # Запускаем в режиме демона, чтобы когда утилита регистрации отработала
    # агент бы выключился
    t0.daemon = True
    t0.start()
    agent.aaa = Socket(IP, PORT)
    agent.name = name
    # Выполняем RPC-запрос к AAA
    ecolog(agent.register())


def handle_command(args):
    if args.command == "startagent":
        handle_startagent(args)
    elif args.command == "registeragent":
        handle_registeragent(args)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", help="Доступные команды")

    # Регистрация агента
    registeragent_parser = subparsers.add_parser("registeragent", help="Зарегистрировать агента")
    registeragent_parser.add_argument("type", help="Тип агента")
    registeragent_parser.add_argument("-n", "--name", help="Имя агента")
    registeragent_parser.add_argument("-p", "--port", help="Порт агента")
    registeragent_parser.add_argument("-i", "--ip", help="IP-адрес агента")

    # Запуск агента
    startagent_parser = subparsers.add_parser("startagent", help="Запустить агента")
    startagent_parser.add_argument("type", help="Тип агента")

    parsed_args = parser.parse_args()
    handle_command(parsed_args)