Zordon/dashboard.py

#!/usr/bin/env python
import datetime
import os.path
import socket
import threading
import time
import traceback

import requests
from rich import box
from rich.console import Console
from rich.layout import Layout
from rich.live import Live
from rich.panel import Panel
from rich.table import Column
from rich.table import Table
from rich.text import Text
from rich.tree import Tree

import zordon_server
from zordon_server import RenderStatus, string_to_status

"""
The RenderDashboard is designed to be run on a remote machine or on the local server
This provides a detailed status of all jobs running on the server
"""

status_colors = {RenderStatus.ERROR: "red", RenderStatus.CANCELLED: 'orange1', RenderStatus.COMPLETED: 'green',
                 RenderStatus.NOT_STARTED: "yellow", RenderStatus.SCHEDULED: 'purple',
                 RenderStatus.RUNNING: 'cyan'}

categories = [RenderStatus.RUNNING, RenderStatus.ERROR, RenderStatus.NOT_STARTED, RenderStatus.SCHEDULED,
              RenderStatus.COMPLETED, RenderStatus.CANCELLED]

renderer_colors = {'ffmpeg': '[magenta]', 'Blender': '[orange1]', 'aerender':'[purple'}

local_hostname = socket.gethostname()


def status_string_to_color(status_string):
    job_status = string_to_status(status_string)
    job_color = '[{}]'.format(status_colors[job_status])
    return job_color


def sorted_jobs(all_jobs):

    sort_by_date = True
    if not sort_by_date:
        sorted_job_list = []
        if all_jobs:
            for status_category in categories:
                found_jobs = [x for x in all_jobs if x['status'] == status_category.value]
                if found_jobs:
                    sorted_found_jobs = sorted(found_jobs, key=lambda d: datetime.datetime.fromisoformat(d['date_created']), reverse=True)
                    sorted_job_list.extend(sorted_found_jobs)
    else:
        sorted_job_list = sorted(all_jobs, key=lambda d: datetime.datetime.fromisoformat(d['date_created']), reverse=True)
    return sorted_job_list


def create_node_tree(all_server_data) -> Tree:
    main_tree = Tree("[magenta]Server Cluster")

    for server_host, server_data in all_server_data['servers'].items():

        node_title_local = f"[cyan bold]{server_host}[/] [yellow](This Computer)[default]"
        node_title_remote = f"[cyan]{server_host} [magenta](Remote)[default]"
        node_tree_text = node_title_local if (server_host == local_hostname) else node_title_remote

        if server_data.get('is_online', False):

            node_tree_text = node_tree_text + " - [green]Running"
            node_tree = Tree(node_tree_text)

            stats_text = f"CPU: [yellow]{server_data['status']['cpu_percent']}% [default]| RAM: " \
                         f"[yellow]{server_data['status']['memory_percent']}% [default]| Cores: " \
                         f"[yellow]{server_data['status']['cpu_count']} [default]| " \
                         f"{server_data['status']['platform'].split('-')[0]}"

            node_tree.add(Tree(stats_text))

            running_jobs = [job for job in server_data['jobs'] if job['status'] == 'running']
            not_started = [job for job in server_data['jobs'] if job['status'] == 'not_started']
            scheduled = [job for job in server_data['jobs'] if job['status'] == 'scheduled']
            jobs_to_display = running_jobs + not_started + scheduled

            jobs_tree = Tree(f"Running: [green]{len(running_jobs)} [default]| Queued: [cyan]{len(not_started)}"
                             f"[default] | Scheduled: [cyan]{len(scheduled)}")

            for job in jobs_to_display:
                renderer = f"{renderer_colors[job['renderer']]}{job['renderer']}[default]"
                filename = os.path.basename(job['render']['input']).split('.')[0]
                if job['status'] == 'running':
                    jobs_tree.add(f"[bold]{renderer} {filename} ({job['id']}) - {status_string_to_color(job['status'])}{(float(job['render']['percent_complete']) * 100):.1f}%")
                else:
                    jobs_tree.add(f"{filename} ({job['id']}) - {status_string_to_color(job['status'])}{job['status'].title()}")

            if not jobs_to_display:
                jobs_tree.add("[italic]No running jobs")

            node_tree.add(jobs_tree)
            main_tree.add(node_tree)
        else:
            # if server is offline
            node_tree_text = node_tree_text + " - [red]Offline"
            node_tree = Tree(node_tree_text)
            main_tree.add(node_tree)
    return main_tree


def create_jobs_table(all_server_data) -> Table:
    table = Table("ID", "Project", "Output", "Renderer", Column(header="Priority", justify="center"),
                  Column(header="Status", justify="center"), Column(header="Time Elapsed", justify="right"),
                  Column(header="# Frames", justify="right"), "Client", show_lines=True,
                  box=box.HEAVY_HEAD)

    all_jobs = []
    for server_name, server_data in all_server_data['servers'].items():
        for job in server_data['jobs']:
            #todo: clean this up
            all_jobs.append(job)

    all_jobs = sorted_jobs(all_jobs)

    for job in all_jobs:

        job_status = string_to_status(job['status'])
        job_color = '[{}]'.format(status_colors[job_status])
        job_text = f"{job_color}" + job_status.value.title()

        if job_status == RenderStatus.ERROR and job['render']['errors']:
            job_text = job_text + "\n" + "\n".join(job['render']['errors'])

        elapsed_time = job['render'].get('time_elapsed', 'unknown')

        # Project name
        project_name = job_color + os.path.basename(job['render']['input'])
        project_name = project_name.replace(".", "[default].")

        if job_status == RenderStatus.RUNNING:
            job_text = f"{job_color}[bold]Running - {float(job['render']['percent_complete']) * 100:.1f}%"
            delta = datetime.datetime.now() - datetime.datetime.fromisoformat(job['render']['start_time'])
            elapsed_time = "[bold]" + str(delta)
            project_name = "[bold]" + project_name
        elif job_status == RenderStatus.CANCELLED or job_status == RenderStatus.ERROR:
            project_name = "[strike]" + project_name

        # Priority
        priority_color = ["red", "yellow", "cyan"][(job['priority'] - 1)]

        client_name = job['client'] or 'unknown'
        client_colors = {'unknown': '[red]', local_hostname: '[yellow]'}
        client_title = client_colors.get(client_name, '[magenta]') + client_name

        table.add_row(
            job['id'],
            project_name,
            os.path.basename(job['render']['output']),
            renderer_colors.get(job['renderer'], '[cyan]') + job['renderer'] + '[default]-' + job['render']['renderer_version'],
            f"[{priority_color}]{job['priority']}",
            job_text,
            elapsed_time,
            str(max(int(job['render']['total_frames']), 1)),
            client_title
        )

    return table


def create_status_panel(all_server_data):
    for key, value in all_server_data['servers'].items():
        if key == local_hostname:
            return str(value['status'])
    return "no status"


class RenderDashboard:

    def __init__(self, server_ip=None, server_port="8080"):
        self.server_ip = server_ip
        self.server_port = server_port
        self.local_hostname = local_hostname

    def connect(self):
        status = self.request_data('status')
        return status

    def request_data(self, payload, timeout=2):
        try:
            req = requests.get(f'http://{self.server_ip}:{self.server_port}/{payload}', timeout=timeout)
            if req.ok:
                return req.json()
        except Exception as e:
            pass
        return None

    def get_jobs(self):
        all_jobs = self.request_data('jobs')
        sorted_jobs = []
        if all_jobs:
            for status_category in categories:
                found_jobs = [x for x in all_jobs if x['status'] == status_category.value]
                if found_jobs:
                    sorted_jobs.extend(found_jobs)
        return sorted_jobs

    def get_data(self):
        all_data = self.request_data('full_status')
        return all_data


class KeyboardThread(threading.Thread):

    def __init__(self, input_cbk = None, name='keyboard-input-thread'):
        self.input_cbk = input_cbk
        super(KeyboardThread, self).__init__(name=name)
        self.start()

    def run(self):
        while True:
            self.input_cbk(input()) #waits to get input + Return


def my_callback(inp):
    #evaluate the keyboard input
    print('You Entered:', inp)


if __name__ == '__main__':

    get_server_ip = input("Enter server IP or None for local: ") or local_hostname

    client = RenderDashboard(get_server_ip, "8080")

    if not client.connect():
        if client.server_ip == local_hostname:
            start_server = input("Local server not running. Start server? (y/n) ")
            if start_server and start_server[0].lower() == "y":
                # Startup the local server
                zordon_server.RenderServer.start(background_thread=True)
                test = client.connect()
                print(f"connected? {test}")
        else:
            print(f"\nUnable to connect to server: {client.server_ip}")
            print("\nVerify IP address is correct and server is running")
            exit(1)

    # start the Keyboard thread
    kthread = KeyboardThread(my_callback)

    # Console Layout
    console = Console()
    layout = Layout()

    # Divide the "screen" in to three parts
    layout.split(
        Layout(name="header", size=3),
        Layout(ratio=1, name="main"),
        Layout(size=10, name="footer"),
    )
    # Divide the "main" layout in to "side" and "body"
    layout["main"].split_row(
        Layout(name="side"),
        Layout(name="body",
               ratio=3))
    # Divide the "side" layout in to two
    layout["side"].split(Layout(name="side_top"), Layout(name="side_bottom"))

    # Server connection header
    header_text = Text(f"Connected to server: ")
    header_text.append(f"{client.server_ip} ", style="green")
    if client.server_ip == local_hostname:
        header_text.append("(This Computer)", style="yellow")
    else:
        header_text.append("(Remote)", style="magenta")

    layout["header"].update(Panel(Text("Zordon Render Client - Version 0.0.1 alpha", justify="center")))

    with Live(console=console, screen=False, refresh_per_second=1, transient=True) as live:
        while True:

            server_data = client.get_data()
            try:
                if server_data:
                    layout["body"].update(create_jobs_table(server_data))
                    layout["side_top"].update(Panel(create_node_tree(server_data)))
                    layout["side_bottom"].update(Panel(create_status_panel(server_data)))
                    live.update(layout, refresh=False)
            except Exception as e:
                print(f"Exception updating table: {e}")
                traceback.print_exception(e)
            time.sleep(1)
    #     #     # todo: Add input prompt to manage running jobs (ie add, cancel, get info, etc)