remote_run_callbacks.py

import os
import re
import threading

from socket import error
from uuid import uuid4

import dearpygui.dearpygui as dpg
from paramiko.ssh_exception import BadHostKeyException, AuthenticationException, UnableToAuthenticate, NoValidConnectionsError, SSHException

from utils import show_message, get_config_file_path, parse_dirs, construct_folder_structure, get_config, update_config_values, get_run_window_value, set_run_window_value, delete_item_children_and_clear_aliases
from remote_run import Connection, BuildingSession, RunningSession
from remote_run_constants import *
from style_constants import FONT_SIZE
from dynamic_ui import construct_model_output_structure_ui
from config import dump_config, load_config
from style import get_accent_color

def add_running_session_ui(running_session_id, machine, model):
    with dpg.table_row(parent='models_table', tag=running_session_id):
        dpg.add_text(MODEL_TO_LABEL[model], tag=f'{running_session_id}_model')
        dpg.add_text(MACHINE_TO_LABEL[machine], tag=f'{running_session_id}_machine')
        dpg.add_text('building', color=list(get_accent_color()[:3]) + [255], tag=f'{running_session_id}_status')
        with dpg.group(horizontal=False):
            dpg.add_spacer(height=0)
            dpg.add_progress_bar(label='progress', default_value=0.0, tag=f'{running_session_id}_progress_bar', height=FONT_SIZE, width=150)
        dpg.add_button(label='download output', enabled=False, tag=f'{running_session_id}_download_button', callback=show_download_output_window_clb, user_data=running_session_id)
    dpg.bind_item_theme(f'{running_session_id}_progress_bar', 'progress_bar_theme')

def update_run_model_window():
    builds = get_run_window_value('builds')
    connections = get_run_window_value('connections')

    machine = LABEL_TO_MACHINE[dpg.get_value('machine')]
    model = LABEL_TO_MODEL[dpg.get_value('model')]

    if (machine, model) in builds:
        dpg.hide_item('credentials')
    elif machine in connections:
        dpg.show_item('gitlab_group')
        dpg.hide_item('lab_server_group')
        dpg.hide_item('lomonosov_server_group')
    else:
        if machine == LAB:
            dpg.show_item('gitlab_group')
            dpg.show_item('lab_server_group')
            dpg.hide_item('lomonosov_server_group')
        elif machine == LOMONOSOV:
            dpg.show_item('gitlab_group')
            dpg.hide_item('lab_server_group')
            dpg.show_item('lomonosov_server_group')

def save_running_session(running_session_id):
    with open('sessions.txt', 'a+', encoding='utf-8') as log:
        model = dpg.get_value(f'{running_session_id}_model')
        machine = dpg.get_value(f'{running_session_id}_machine')
        log.write(f'{running_session_id};{machine};{model}\n')

def write_log(output):
    if not output:
        return

    with open('log.txt', 'a+', encoding='utf-8') as log:
        log.write(output + '\n')

def load_running_sessions():
    if not os.path.exists('sessions.txt'):
        return

    with open('sessions.txt', 'r', encoding='utf-8') as file:
        for line in file:
            running_session_id, machine, model = line.strip('\r\n').split(';')
            add_running_session_ui(running_session_id, LABEL_TO_MACHINE[machine], LABEL_TO_MODEL[model])
            dpg.set_value(f'{running_session_id}_status', 'exited long ago')
            dpg.hide_item(f'{running_session_id}_progress_bar')
            dpg.enable_item(f'{running_session_id}_download_button')

def update_run_model_window_clb(s, a, u):
    update_run_model_window()

def show_run_model_window_clb(s, a, u):
    if get_config_file_path():
        update_config_values()
        dump_config(get_config(), get_config_file_path())

    dpg.set_value('config_file_path', get_config_file_path())

    update_run_model_window()

    dpg.show_item('run_model_window')

def make_execute_command():
    np = int(dpg.get_value('-np'))
    arch = dpg.get_value('-arch')
    udump = dpg.get_value('-udump')
    model_output = dpg.get_value('-model-output')

    command = f'mpirun -np {np} ./nsenx -arch {arch}'
    if dpg.get_value('use-udump'):
        command += f' -udump {dpg.get_value('-udump')}'
    if dpg.get_value('use-model-output'):
        command += f' -model-output {dpg.get_value('-model-output')}'

    return command

def make_run_commands(running_session_id, executable_file_path, machine, model):
    run_commands = []

    folders = ['chem-forcing', 'chem-init', 'drag-configs', 'drag-forcing', 'meteo-forcing', 'meteo-init']
    prefix = executable_file_path.split('/')[0] + '/nse-gabls1-urban-les'

    for folder in folders:
        path = f'{prefix}/{folder}'
        run_commands.append(f'cp -r ../{path} {folder}')

    if machine == LOMONOSOV:
        run_commands += RUN_COMMANDS[(machine, model)] # + make_run_sh_for_slurm() + ['sbatch ./run.sh']
    else:
        run_commands += RUN_COMMANDS[(machine, model)] + [make_execute_command()]

    return run_commands

def update_progress(running_session_id, output):
    if not re.findall(r'([0-9]+)%', output):
        progress = dpg.get_value(f'{running_session_id}_progress_bar')
    else:
        progress = int(re.findall(r'([0-9]+)%', output)[-1]) / 100
    dpg.set_value(f'{running_session_id}_progress_bar', progress)

def run_model():
    dpg.hide_item('run_model_window')

    config_file_path = dpg.get_value('config_file_path')

    if not os.path.exists(config_file_path):
        show_message(f'no {config_file_path} file :(')
        return

    builds = get_run_window_value('builds')
    connections = get_run_window_value('connections')

    machine = LABEL_TO_MACHINE[dpg.get_value('machine')]
    model = LABEL_TO_MODEL[dpg.get_value('model')]

    if machine not in connections:
        if machine == LAB:
            username = dpg.get_value('server_username')
            password = dpg.get_value('server_password')
        elif machine == LOMONOSOV:
            username = dpg.get_value('lomonosov_username')
            if 'ID_RSA_PATH' in os.environ:
                id_rsa_path = os.environ['ID_RSA_PATH']
            else:
                id_rsa_path = None
        try:
            if machine == LAB:
                conn = Connection(MACHINE_TO_HOST[machine], username, password)
            elif machine == LOMONOSOV:
                if id_rsa_path:
                    conn = Connection(MACHINE_TO_HOST[machine], username, id_rsa_path=id_rsa_path)
                else:
                    conn = Connection(MACHINE_TO_HOST[machine], username)
        except BadHostKeyException:
            show_message('server’s host key could not be verified\ntry again')
            return
        except AuthenticationException:
            show_message('wrong username or password')
            return
        except error:
            show_message('check your internet connection :)')
            return
        except NoValidConnectionsError:
            show_message('something wrong with the server ¯\\_(ツ)_/¯')
            return
        except SSHException:
            show_message('i have no idea what happend\nplease contact @maryshca')
            return
        connections[machine] = conn
        set_run_window_value('connections', connections)
    else:
        conn = connections[machine]

    running_session_id = str(uuid4())

    if machine == LOMONOSOV:
        running_session_id = '~/_scratch/' + running_session_id

    add_running_session_ui(running_session_id, machine, model)

    save_running_session(running_session_id)

    if (machine, model) not in builds:
        dpg.hide_item(f'{running_session_id}_progress_bar')

        gitlab_username = dpg.get_value('gitlab_username')
        gitlab_password = dpg.get_value('gitlab_password')
        build_commands = BUILD_COMMANDS[(machine, model)]
        exec_file_path = EXECUTABLE_FILE_PATH[(machine, model)]
        building_session = BuildingSession(conn, build_commands, exec_file_path, gitlab_username, gitlab_password)

        while building_session.building():
            output = building_session.output()
            update_progress(running_session_id, output)

        executable_file_path = building_session.executable_file_path

        builds[(machine, model)] = executable_file_path
        set_run_window_value('builds', builds)

        dpg.show_item(f'{running_session_id}_progress_bar')
    else:
        executable_file_path = builds[(machine, model)]

    dpg.set_value(f'{running_session_id}_status', 'running')

    run_commands = make_run_commands(running_session_id, executable_file_path, machine, model)

    upload_paths = dict()

    if get_config_file_path():
        config = get_config()

        if 'topography.filename' in config:
            upload_paths = {config['topography.filename'].value: config['topography.filename'].value}
    else:
        try:
            config = load_config(config_file_path)
            if 'topography.filename' in config:
                upload_paths = {config['topography.filename'].value: config['topography.filename'].value}
        except Exception as e:
            show_message('bad file format')
            return

    running_session = RunningSession(conn, config_file_path, run_commands, executable_file_path, running_session_id, upload_paths=upload_paths)

    if machine == LOMONOSOV:
        output = running_session.output()
        sbatch_id = re.findall(r'([0-1]+)', output)[0]
        return

    while running_session.running():
        output = running_session.output()
        write_log(output)
        update_progress(running_session_id, output)

    dpg.set_value(f'{running_session_id}_status', 'successful run')
    dpg.set_value(f'{running_session_id}_progress_bar', 1.0)
    dpg.enable_item(f'{running_session_id}_download_button')

def run_model_clb():
    threading.Thread(target=run_model).start()

def show_download_output_window_clb(s, a, running_session_id):
    delete_item_children_and_clear_aliases('output_header')

    machine = LABEL_TO_MACHINE[dpg.get_value(f'{running_session_id}_machine')]

    conn = get_run_window_value('connections')[machine]
    stdout = conn.exec_command(f'cd ./{running_session_id}/output; find . -type f')
    output = stdout.read().decode()
    dirs, filepaths = parse_dirs(output)
    construct_model_output_structure_ui(filepaths, dirs)

    dpg.set_item_user_data('model_output_window', (dirs, filepaths, running_session_id))

    dpg.show_item('model_output_window')

def open_download_to_folder_dialog_clb():
    dpg.show_item('download_to_folder_dialog')

def open_config_file_path_dialog_clb():
    dpg.show_item('config_file_path_dialog')

def set_download_output_to_folder_clb(s, a, u):
    dpg.set_value('download_output_to_folder', a['file_path_name'])

def set_config_file_path_clb(s, a, u):
    dpg.set_value('config_file_path', a['file_path_name'])

def download_output():
    dpg.hide_item('model_output_window')
    dirs, filepaths, running_session_id = dpg.get_item_user_data('model_output_window')

    download = set()
    download_folders = set()
    for filepath in filepaths:
        if dpg.get_value(filepath):
            download.add(filepath)
            download_folders.add('/'.join(filepath.split('/')[:-1]))

    machine = LABEL_TO_MACHINE[dpg.get_value(f'{running_session_id}_machine')]

    download_to_folder = dpg.get_value('download_output_to_folder')

    dpg.set_value(f'{running_session_id}_status', 'downloading')

    prefix = os.path.join(download_to_folder)
    
    construct_folder_structure(download_folders, prefix=prefix)

    conn = get_run_window_value('connections')[machine]
    
    for filepath in download:
        conn.download(f'{running_session_id}/output/{filepath}', os.path.join(prefix, filepath))

    conn.download(f'{running_session_id}/map.txt', os.path.join(prefix, 'map.txt'))

    dpg.set_value(f'{running_session_id}_status', 'downloaded')

def download_output_clb(s, a, u):
    threading.Thread(target=download_output).start()