bot_actions.py

import asyncio
import bot_checks
import datetime
import getpass
import shlex
import yaml

from termcolor import colored
from pygments import highlight
from pygments.lexers import YamlLexer
from pygments.formatters import TerminalFormatter

from db_connector import action

POWER_CYCLE_COMMAND = "sudo /global/home/groups/scs/sbin/ipmiwrapper.sh {node} chassis power cycle"
POWER_ON_COMMAND = "sudo /global/home/groups/scs/sbin/ipmiwrapper.sh {node} chassis power on"
POWER_OFF_COMMAND = "sudo /global/home/groups/scs/sbin/ipmiwrapper.sh {node} chassis power off"
SLURM_RESUME_COMMAND = "sudo scontrol update node={node} state=resume"
SYSTEMCTL_DAEMON_RELOAD_COMMAND = "sudo systemctl daemon-reload && sudo systemctl start slurmd"

def ssh(node):
    def ssh_command(command):
        return 'ssh -t {node} {command}'.format(node=shlex.quote(node),command=shlex.quote(command))
    return ssh_command

def power_cycle(node, state):
    return [
        POWER_CYCLE_COMMAND.format(node=node)
    ]

def power_on(node, state):
    return [
        POWER_ON_COMMAND.format(node=node)
    ]

def restart_slurm(node, state):
    return [
        ssh(node)(SYSTEMCTL_DAEMON_RELOAD_COMMAND)
    ]

def slurm_resume(node, state):
    return [
        SLURM_RESUME_COMMAND.format(node=node)
    ]

def nothing(node, state):
    return []

SUGGESTION = {
    'NODE_KILLED_IPMI_ON': power_cycle,
    'NODE_KILLED_IPMI_OFF': power_on,
    'SLURM_FAILED_USER_PROCESSES_ALIVE': power_cycle,
    'SLURM_FAILED_NO_USER_PROCESSES': restart_slurm,
    'NODE_WORKING': slurm_resume,
    'UNKNOWN': nothing
}

def suggest(node, state):
    return SUGGESTION[state](node, state)

def display_status(status):
    return highlight(yaml.dump(status, sort_keys=True, indent=2), YamlLexer(), TerminalFormatter())

def display_suggestion(index, suggestion):
    return str(index) + ') ' + colored(';'.join([ command for command in suggestion ]), attrs=['bold'])

async def interactive_suggest(suggestions, status):
    state_groups = {}
    count = 0
    for node, (suggestion, state) in suggestions.items():
        if not suggestion:
            continue
        key = state, status[node]['REASON']
        if key not in state_groups:
            state_groups[key] = []
        count += 1
        state_groups[key].append((node, suggestion))

    print('#!/bin/sh')
    print('#', 'Generated by bofhbot on', datetime.datetime.now(), 'by', getpass.getuser())
    print('#', '{}/{} nodes have suggestions'.format(count, len(status.keys())))
    print()
    for state, nodes_suggestions in state_groups.items():
        nodes = [ node for node, _ in nodes_suggestions ]
        suggestions = '\n'.join([ '; '.join(suggestion) for _, suggestion in nodes_suggestions ])
        print('#', state, ':', ','.join(nodes),'\n' + suggestions, '\n')