opengris-scaler 1.12.37__cp38-cp38-musllinux_1_2_x86_64.whl

scaler/ui/common/task_log.py

@@ -0,0 +1,111 @@
+import dataclasses
+from collections import deque
+from threading import Lock
+from typing import Deque, Dict, Optional
+
+from nicegui import ui
+
+from scaler.protocol.python.common import TaskState
+from scaler.protocol.python.message import StateTask, StateWorker
+from scaler.ui.common.utility import COMPLETED_TASK_STATUSES, display_capabilities
+from scaler.ui.util import NICEGUI_MAJOR_VERSION
+from scaler.utility.formatter import format_bytes
+from scaler.utility.metadata.profile_result import ProfileResult
+
+TASK_ID_HTML_TEMPLATE = (
+    "<span "
+    "style='display:inline-block;max-width:12rem;overflow:hidden;text-overflow:ellipsis;"
+    "white-space:nowrap;cursor:pointer;font:inherit;color:inherit' "
+    "title='{task}' onclick=\"navigator.clipboard.writeText('{task}')\">{task}</span>"
+)
+
+
+@dataclasses.dataclass
+class TaskData:
+    task: str = dataclasses.field(default="")
+    function: str = dataclasses.field(default="")
+    duration: str = dataclasses.field(default="")
+    peak_mem: str = dataclasses.field(default="")
+    status: str = dataclasses.field(default="")
+    capabilities: str = dataclasses.field(default="")
+
+    def populate(
+        self,
+        state: StateTask,
+        function_name: str,
+        profiling_data: Optional[ProfileResult],
+        task_capabilities: Dict[str, int],
+    ):
+        self.task = f"{state.task_id.hex()}"
+        self.function = function_name
+        self.status = state.state.name
+
+        if profiling_data:
+            duration = profiling_data.duration_s
+            mem = profiling_data.memory_peak
+            self.duration = f"{duration:.2f}s"
+            self.peak_mem = format_bytes(mem) if mem != 0 else "0"
+        else:
+            self.duration = "N/A"
+            self.peak_mem = "N/A"
+
+        self.capabilities = display_capabilities(set(task_capabilities.keys()))
+
+    def draw_row(self):
+        color = "color: green" if self.status == TaskState.Success.name else "color: red"
+        if NICEGUI_MAJOR_VERSION < 3:
+            ui.html(TASK_ID_HTML_TEMPLATE.format(task=self.task))
+        else:
+            ui.html(TASK_ID_HTML_TEMPLATE.format(task=self.task), sanitize=False)  # type: ignore[call-arg]
+        ui.label(self.function)
+        ui.label(self.duration)
+        ui.label(self.peak_mem)
+        ui.label(self.status).style(color)
+        ui.label(self.capabilities)
+
+    @staticmethod
+    def draw_titles():
+        ui.label("Task ID")
+        ui.label("Function")
+        ui.label("Duration")
+        ui.label("Peak mem")
+        ui.label("Status")
+        ui.label("Capabilities")
+
+
+class TaskLogTable:
+    def __init__(self):
+        self._task_log: Deque[TaskData] = deque(maxlen=100)
+        self._task_id_to_function_name: Dict[str, str] = {}
+        self._lock: Lock = Lock()
+
+    def handle_task_state(self, state_task: StateTask):
+        if state_task.function_name != b"" and state_task.task_id.hex() not in self._task_id_to_function_name:
+            self._task_id_to_function_name[state_task.task_id.hex()] = state_task.function_name.decode()
+
+        if state_task.state not in COMPLETED_TASK_STATUSES:
+            return
+
+        function_name = state_task.function_name.decode()
+        if function_name == "":
+            function_name = self._task_id_to_function_name.pop(state_task.task_id.hex(), "")
+
+        # Canceled/failed states don't have profiling metadata
+        profiling_data = ProfileResult.deserialize(state_task.metadata) if state_task.metadata != b"" else None
+
+        row = TaskData()
+        row.populate(state_task, function_name, profiling_data, state_task.capabilities)
+
+        with self._lock:
+            self._task_log.appendleft(row)
+
+    def handle_worker_state(self, _: StateWorker):
+        return
+
+    @ui.refreshable
+    def draw_section(self):
+        with self._lock:
+            with ui.card().classes("w-full q-mx-auto"), ui.grid(columns=6).classes("q-mx-auto"):
+                TaskData.draw_titles()
+                for task in self._task_log:
+                    task.draw_row()

scaler/ui/common/utility.py

@@ -0,0 +1,66 @@
+import datetime
+from typing import List, Set, Tuple
+
+from scaler.protocol.python.common import TaskState
+from scaler.ui.common.setting_page import Settings
+
+COMPLETED_TASK_STATUSES = {
+    TaskState.Success,
+    TaskState.Canceled,
+    TaskState.CanceledNotFound,
+    TaskState.Failed,
+    TaskState.FailedWorkerDied,
+}
+
+
+def format_timediff(a: datetime.datetime, b: datetime.datetime) -> float:
+    return (b - a).total_seconds()
+
+
+def format_worker_name(worker_name: str, cutoff: int = 15) -> str:
+    return worker_name[:cutoff] + "+" if len(worker_name) > cutoff else worker_name
+
+
+def get_bounds(now: datetime.datetime, start_time: datetime.datetime, settings: Settings) -> Tuple[int, int]:
+    upper_range = now - start_time
+    lower_range = upper_range - settings.stream_window
+
+    bound_upper_seconds = max(upper_range.seconds, settings.stream_window.seconds)
+    bound_lower_seconds = 0 if bound_upper_seconds == settings.stream_window.seconds else lower_range.seconds
+
+    return bound_lower_seconds, bound_upper_seconds
+
+
+def make_taskstream_ticks(lower_bound: int, upper_bound: int) -> List[int]:
+    distance = (upper_bound - lower_bound) // 6
+    return list(range(lower_bound, upper_bound + 1, distance))
+
+
+def make_memory_ticks(max_bytes: int) -> Tuple[List[int], List[str]]:
+    units = ["B", "KB", "MB", "GB", "TB"]
+    vals: List[int] = [0]
+    texts: List[str] = ["0"]
+    v = 1
+    i = 0
+    # ensure at least up to 1GB on empty data
+    target = max(1024 * 1024 * 1024, max_bytes)
+    while i < len(units) and v <= target:
+        vals.append(v)
+        texts.append(f"1{units[i]}")
+        v *= 1024
+        i += 1
+    return vals, texts
+
+
+def make_tick_text(window_length: int) -> List[int]:
+    upper = 0
+    lower = -1 * window_length
+    distance = (upper - lower) // 6
+    return list(range(lower, upper + 1, distance))
+
+
+def display_capabilities(capabilities: Set[str]) -> str:
+    if not capabilities or len(capabilities) == 0:
+        return "<no capabilities>"
+
+    return " & ".join(sorted(capabilities))

scaler/ui/common/webui.py

@@ -0,0 +1,80 @@
+import dataclasses
+import logging
+
+from scaler.protocol.python.message import StateBalanceAdvice, StateScheduler, StateTask, StateWorker
+from scaler.protocol.python.mixins import Message
+from scaler.ui.common.live_display import SchedulerSection, WorkersSection
+from scaler.ui.common.memory_window import MemoryChart
+from scaler.ui.common.setting_page import Settings
+from scaler.ui.common.task_graph import TaskStream
+from scaler.ui.common.task_log import TaskLogTable
+from scaler.ui.common.worker_processors import WorkerProcessors
+from scaler.utility.formatter import format_bytes, format_percentage
+
+
+@dataclasses.dataclass
+class Sections:
+    scheduler_section: SchedulerSection
+    workers_section: WorkersSection
+    task_stream_section: TaskStream
+    memory_usage_section: MemoryChart
+    tasklog_section: TaskLogTable
+    worker_processors: WorkerProcessors
+    settings_section: Settings
+
+
+def process_scheduler_message(status: Message, tables: Sections):
+    if isinstance(status, StateScheduler):
+        __update_scheduler_state(status, tables)
+        return
+
+    if isinstance(status, StateWorker):
+        logging.info(f"Received StateWorker update for worker {status.worker_id.decode()} with {status.state.name}")
+        tables.scheduler_section.handle_worker_state(status)
+        tables.workers_section.handle_worker_state(status)
+        tables.task_stream_section.handle_worker_state(status)
+        tables.memory_usage_section.handle_worker_state(status)
+        tables.tasklog_section.handle_worker_state(status)
+        tables.worker_processors.handle_worker_state(status)
+        tables.settings_section.handle_worker_state(status)
+        return
+
+    if isinstance(status, StateTask):
+        logging.debug(f"Received StateTask update for task {status.task_id.hex()} with {status.state.name}")
+        tables.scheduler_section.handle_task_state(status)
+        tables.workers_section.handle_task_state(status)
+        tables.task_stream_section.handle_task_state(status)
+        tables.memory_usage_section.handle_task_state(status)
+        tables.tasklog_section.handle_task_state(status)
+        tables.worker_processors.handle_task_state(status)
+        tables.settings_section.handle_task_state(status)
+        return
+
+    if isinstance(status, StateBalanceAdvice):
+        logging.debug(f"Received StateBalanceAdvice for {status.worker_id.decode()} with {len(status.task_ids)} tasks")
+        return
+
+    logging.info(f"Unhandled message received: {type(status)}")
+
+
+def __update_scheduler_state(data: StateScheduler, tables: Sections):
+    tables.scheduler_section.cpu = format_percentage(data.scheduler.cpu)
+    tables.scheduler_section.rss = format_bytes(data.scheduler.rss)
+    tables.scheduler_section.rss_free = format_bytes(data.rss_free)
+
+    previous_workers = set(tables.workers_section.workers.keys())
+    current_workers = set(worker_data.worker_id.decode() for worker_data in data.worker_manager.workers)
+
+    for worker_data in data.worker_manager.workers:
+        worker_name = worker_data.worker_id.decode()
+        tables.workers_section.workers[worker_name].populate(worker_data)
+
+    for died_worker in previous_workers - current_workers:
+        tables.workers_section.workers.pop(died_worker)
+        tables.worker_processors.remove_worker(died_worker)
+        tables.task_stream_section.mark_dead_worker(died_worker)
+
+    if previous_workers != current_workers:
+        tables.workers_section.draw_section.refresh()
+
+    tables.worker_processors.update_data(data.worker_manager.workers)

scaler/ui/common/worker_processors.py

@@ -0,0 +1,104 @@
+import dataclasses
+from threading import Lock
+from typing import Dict, List, Optional
+
+from nicegui import ui
+from nicegui.element import Element
+
+from scaler.protocol.python.common import WorkerState
+from scaler.protocol.python.message import StateTask, StateWorker
+from scaler.protocol.python.status import ProcessorStatus, WorkerStatus
+from scaler.ui.common.utility import format_worker_name
+
+
+@dataclasses.dataclass
+class WorkerProcessors:
+    workers: Dict[str, "WorkerProcessorTable"] = dataclasses.field(default_factory=dict)
+    _lock: Lock = Lock()
+
+    @ui.refreshable
+    def draw_section(self):
+        with self._lock:
+            for processor_table in self.workers.values():
+                processor_table.draw_table()
+
+    def update_data(self, data: List[WorkerStatus]):
+        with self._lock:
+            for worker in data:
+                worker_name = worker.worker_id.decode()
+                processor_table = self.workers.get(worker_name)
+
+                if processor_table is None:
+                    processor_table = WorkerProcessorTable(worker_name, 0, worker.rss_free, worker.processor_statuses)
+                    self.workers[worker_name] = processor_table
+                elif processor_table.processor_statuses != worker.processor_statuses:
+                    processor_table.processor_statuses = worker.processor_statuses
+
+    def remove_worker(self, dead_worker: str):
+        with self._lock:
+            self.workers.pop(dead_worker, None)
+
+    def handle_task_state(self, _: StateTask):
+        return
+
+    def handle_worker_state(self, state_worker: StateWorker):
+        worker_id = state_worker.worker_id.decode()
+        state = state_worker.state
+
+        if state == WorkerState.Disconnected:
+            self.remove_worker(worker_id)
+
+
+@dataclasses.dataclass
+class WorkerProcessorTable:
+    worker_name: str
+    rss_max: int
+    rss_free: int
+    processor_statuses: List[ProcessorStatus]
+
+    handler: Optional[Element] = dataclasses.field(default=None)
+
+    def draw_table(self):
+        formatted_worker_name = format_worker_name(self.worker_name)
+        with ui.card().classes("w-full") as handler:
+            self.handler = handler
+
+            ui.markdown(f"Worker **{formatted_worker_name}**").classes("text-xl")
+
+            with ui.grid(columns=7).classes("w-full"):
+                self.draw_titles()
+                for processor in sorted(self.processor_statuses, key=lambda x: x.pid):
+                    if processor.resource.rss > self.rss_max:
+                        self.rss_max = processor.resource.rss
+
+                    self.draw_row(processor, self.rss_free, self.rss_max)
+
+    @staticmethod
+    def draw_titles():
+        ui.label("Processor PID")
+        ui.label("CPU %")
+        ui.label("RSS (in MB)")
+        ui.label("Max RSS (in MB)")
+        ui.label("Initialized")
+        ui.label("Has Task")
+        ui.label("Suspended")
+
+    @staticmethod
+    def draw_row(processor_status: ProcessorStatus, rss_free: int, rss_max: int):
+        cpu = processor_status.resource.cpu / 10
+        rss = int(processor_status.resource.rss / 1e6)
+        rss_max = int(rss_max / 1e6)
+        rss_free = int(rss_free / 1e6)
+
+        ui.label(str(processor_status.pid))
+        ui.knob(value=cpu, track_color="grey-2", show_value=True, min=0, max=100)
+        ui.knob(value=rss, track_color="grey-2", show_value=True, min=0, max=rss + rss_free)
+        ui.knob(value=rss_max, track_color="grey-2", show_value=True, min=0, max=rss + rss_free)
+        ui.checkbox().bind_value_from(processor_status, "initialized")
+        ui.checkbox().bind_value_from(processor_status, "has_task")
+        ui.checkbox().bind_value_from(processor_status, "suspended")
+
+    def delete_row(self):
+        assert self.handler is not None
+        self.handler.clear()
+        self.handler.delete()

scaler/ui/v1.py

@@ -0,0 +1,76 @@
+import threading
+from functools import partial
+
+from nicegui import ui
+
+from scaler.config.section.webui import WebUIConfig
+from scaler.io.sync_subscriber import ZMQSyncSubscriber
+from scaler.ui.common.constants import (
+    MEMORY_USAGE_UPDATE_INTERVAL,
+    TASK_LOG_REFRESH_INTERVAL,
+    TASK_STREAM_UPDATE_INTERVAL,
+    WORKER_PROCESSORS_REFRESH_INTERVAL,
+)
+from scaler.ui.common.live_display import SchedulerSection, WorkersSection
+from scaler.ui.common.memory_window import MemoryChart
+from scaler.ui.common.setting_page import Settings
+from scaler.ui.common.task_graph import TaskStream
+from scaler.ui.common.task_log import TaskLogTable
+from scaler.ui.common.webui import Sections, process_scheduler_message
+from scaler.ui.common.worker_processors import WorkerProcessors
+
+
+def start_webui_v1(config: WebUIConfig):
+    tables = Sections(
+        scheduler_section=SchedulerSection(),
+        workers_section=WorkersSection(),
+        task_stream_section=TaskStream(),
+        memory_usage_section=MemoryChart(),
+        tasklog_section=TaskLogTable(),
+        worker_processors=WorkerProcessors(),
+        settings_section=Settings(),
+    )
+
+    with ui.tabs().classes("w-full h-full") as tabs:
+        live_tab = ui.tab("Live")
+        tasklog_tab = ui.tab("Task Log")
+        stream_tab = ui.tab("Worker Task Stream")
+        worker_processors_tab = ui.tab("Worker Processors")
+        settings_tab = ui.tab("Settings")
+
+    with ui.tab_panels(tabs, value=live_tab).classes("w-full"):
+        with ui.tab_panel(live_tab):
+            tables.scheduler_section.draw_section()
+            tables.workers_section.draw_section()  # type: ignore[call-arg]
+
+        with ui.tab_panel(tasklog_tab):
+            tables.tasklog_section.draw_section()  # type: ignore[call-arg]
+            ui.timer(TASK_LOG_REFRESH_INTERVAL, tables.tasklog_section.draw_section.refresh, active=True)
+
+        with ui.tab_panel(stream_tab):
+            tables.task_stream_section.setup_task_stream(tables.settings_section)
+            ui.timer(TASK_STREAM_UPDATE_INTERVAL, tables.task_stream_section.update_plot, active=True)
+
+            tables.memory_usage_section.setup_memory_chart(tables.settings_section)
+            ui.timer(MEMORY_USAGE_UPDATE_INTERVAL, tables.memory_usage_section.update_plot, active=True)
+
+        with ui.tab_panel(worker_processors_tab):
+            tables.worker_processors.draw_section()  # type: ignore[call-arg]
+            ui.timer(WORKER_PROCESSORS_REFRESH_INTERVAL, tables.worker_processors.draw_section.refresh, active=True)
+
+        with ui.tab_panel(settings_tab):
+            tables.settings_section.draw_section()
+
+    subscriber = ZMQSyncSubscriber(
+        address=config.monitor_address,
+        callback=partial(process_scheduler_message, tables=tables),
+        topic=b"",
+        timeout_seconds=-1,
+    )
+    subscriber.start()
+
+    ui_thread = threading.Thread(
+        target=partial(ui.run, host=config.web_host, port=config.web_port, reload=False), daemon=False
+    )
+    ui_thread.start()
+    ui_thread.join()

scaler/ui/v2.py

@@ -0,0 +1,102 @@
+import threading
+from typing import Optional
+
+from nicegui import Event, app, ui  # type: ignore[attr-defined]
+
+from scaler.config.section.webui import WebUIConfig
+from scaler.io.sync_subscriber import ZMQSyncSubscriber
+from scaler.protocol.python.mixins import Message
+from scaler.ui.common.constants import (
+    MEMORY_USAGE_UPDATE_INTERVAL,
+    TASK_LOG_REFRESH_INTERVAL,
+    TASK_STREAM_UPDATE_INTERVAL,
+    WORKER_PROCESSORS_REFRESH_INTERVAL,
+)
+from scaler.ui.common.live_display import SchedulerSection, WorkersSection
+from scaler.ui.common.memory_window import MemoryChart
+from scaler.ui.common.setting_page import Settings
+from scaler.ui.common.task_graph import TaskStream
+from scaler.ui.common.task_log import TaskLogTable
+from scaler.ui.common.webui import Sections, process_scheduler_message
+from scaler.ui.common.worker_processors import WorkerProcessors
+
+
+class WebUI:
+    def __init__(self) -> None:
+        self.scheduler_message = Event[Message]()
+        self.tables: Optional[Sections] = None
+
+    def start(self, host: str, port: int) -> None:
+        """Start the NiceGUI server in a separate thread."""
+        started = threading.Event()
+        app.on_startup(started.set)
+        thread = threading.Thread(
+            target=lambda: ui.run(self.root, host=host, port=port, reload=False),  # type: ignore[misc,arg-type]
+            daemon=True,
+        )
+        thread.start()
+        if not started.wait(timeout=3.0):
+            raise RuntimeError("NiceGUI did not start within 3 seconds.")
+
+    def root(self) -> None:
+        """Create the UI for each new visitor."""
+        self.scheduler_message.subscribe(self.handle_message)
+        tables = Sections(
+            scheduler_section=SchedulerSection(),
+            workers_section=WorkersSection(),
+            task_stream_section=TaskStream(),
+            memory_usage_section=MemoryChart(),
+            tasklog_section=TaskLogTable(),
+            worker_processors=WorkerProcessors(),
+            settings_section=Settings(),
+        )
+        self.tables = tables
+
+        with ui.tabs().classes("w-full h-full") as tabs:
+            live_tab = ui.tab("Live")
+            tasklog_tab = ui.tab("Task Log")
+            stream_tab = ui.tab("Worker Task Stream")
+            worker_processors_tab = ui.tab("Worker Processors")
+            settings_tab = ui.tab("Settings")
+
+        with ui.tab_panels(tabs, value=live_tab).classes("w-full"):
+            with ui.tab_panel(live_tab):
+                tables.scheduler_section.draw_section()
+                tables.workers_section.draw_section()  # type: ignore[call-arg]
+
+            with ui.tab_panel(tasklog_tab):
+                tables.tasklog_section.draw_section()  # type: ignore[call-arg]
+                ui.timer(TASK_LOG_REFRESH_INTERVAL, tables.tasklog_section.draw_section.refresh, active=True)
+
+            with ui.tab_panel(stream_tab):
+                tables.task_stream_section.setup_task_stream(tables.settings_section)
+                ui.timer(TASK_STREAM_UPDATE_INTERVAL, tables.task_stream_section.update_plot, active=True)
+
+                tables.memory_usage_section.setup_memory_chart(tables.settings_section)
+                ui.timer(MEMORY_USAGE_UPDATE_INTERVAL, tables.memory_usage_section.update_plot, active=True)
+
+            with ui.tab_panel(worker_processors_tab):
+                tables.worker_processors.draw_section()  # type: ignore[call-arg]
+                ui.timer(WORKER_PROCESSORS_REFRESH_INTERVAL, tables.worker_processors.draw_section.refresh, active=True)
+
+            with ui.tab_panel(settings_tab):
+                tables.settings_section.draw_section()
+
+    def new_message(self, status: Message):
+        self.scheduler_message.emit(status)
+
+    def handle_message(self, status: Message):
+        process_scheduler_message(status, self.tables)
+
+
+def start_webui_v2(config: WebUIConfig):
+    webui = WebUI()
+    webui.start(config.web_host, config.web_port)
+
+    subscriber = ZMQSyncSubscriber(
+        address=config.monitor_address, callback=webui.new_message, topic=b"", timeout_seconds=-1
+    )
+    subscriber.start()
+
+    while True:
+        pass

scaler/ui/webui.py

@@ -0,0 +1,21 @@
+import logging
+
+from scaler.config.section.webui import WebUIConfig
+from scaler.ui.util import NICEGUI_MAJOR_VERSION
+from scaler.utility.logging.utility import setup_logger
+
+
+def start_webui(config: WebUIConfig):
+
+    setup_logger(config.logging_config.paths, config.logging_config.config_file, config.logging_config.level)
+
+    if NICEGUI_MAJOR_VERSION < 3:
+        logging.info(f"Detected {NICEGUI_MAJOR_VERSION}. Using GUI v1.")
+        from scaler.ui.v1 import start_webui_v1
+
+        start_webui_v1(config)
+    else:
+        logging.info(f"Detected {NICEGUI_MAJOR_VERSION}. Using GUI v2.")
+        from scaler.ui.v2 import start_webui_v2
+
+        start_webui_v2(config)

scaler/utility/debug.py

@@ -0,0 +1,19 @@
+import functools
+import pdb
+import sys
+from typing import Callable
+
+
+def pdb_wrapped(func: Callable):
+    @functools.wraps(func)
+    def pdb_wrapper(*args, **kwargs):
+        try:
+            exit_code = func(*args, **kwargs)
+            sys.exit(exit_code)
+
+        except Exception:
+            ex_type, value, tb = sys.exc_info()
+            pdb.post_mortem(tb)
+            raise
+
+    return pdb_wrapper

scaler/utility/event_list.py

@@ -0,0 +1,63 @@
+import collections
+from typing import Callable
+
+
+class EventList(collections.UserList):
+    """A list that emits events when it is modified."""
+
+    def __init__(self, initlist=None):
+        super().__init__(initlist=initlist)
+        self._callbacks = []
+
+    def add_update_callback(self, callback: Callable[["EventList"], None]):
+        self._callbacks.append(callback)
+
+    def __setitem__(self, i, item):
+        super().__setitem__(i, item)
+        self._list_updated()
+
+    def __delitem__(self, i):
+        super().__delitem__(i)
+        self._list_updated()
+
+    def __add__(self, other):
+        super().__add__(other)
+        self._list_updated()
+
+    def __iadd__(self, other):
+        super().__iadd__(other)
+        self._list_updated()
+        return self
+
+    def append(self, item):
+        super().append(item)
+        self._list_updated()
+
+    def insert(self, i, item):
+        super().insert(i, item)
+        self._list_updated()
+
+    def pop(self, i: int = -1):
+        v = super().pop(i)
+        self._list_updated()
+        return v
+
+    def remove(self, item):
+        super().remove(item)
+        self._list_updated()
+
+    def clear(self) -> None:
+        super().clear()
+        self._list_updated()
+
+    def sort(self, /, *args, **kwargs):
+        super().sort(*args, **kwargs)
+        self._list_updated()
+
+    def extend(self, other) -> None:
+        super().extend(other)
+        self._list_updated()
+
+    def _list_updated(self):
+        for callback in self._callbacks:
+            callback(self)