opengris-scaler 1.12.28__cp313-cp313-musllinux_1_2_x86_64.whl

scaler/ui/task_log.py

@@ -0,0 +1,107 @@
+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.utility import COMPLETED_TASK_STATUSES, display_capabilities
+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"
+        ui.html(TASK_ID_HTML_TEMPLATE.format(task=self.task))
+        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/utility.py

@@ -0,0 +1,66 @@
+import datetime
+from typing import List, Set, Tuple
+
+from scaler.protocol.python.common import TaskState
+from scaler.ui.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/webui.py

@@ -0,0 +1,147 @@
+import dataclasses
+import logging
+import threading
+from functools import partial
+from typing import Optional, Tuple
+
+from nicegui import ui
+
+from scaler.config.types.zmq import ZMQConfig
+from scaler.io.sync_subscriber import ZMQSyncSubscriber
+from scaler.protocol.python.message import StateBalanceAdvice, StateScheduler, StateTask, StateWorker
+from scaler.protocol.python.mixins import Message
+from scaler.ui.constants import (
+    MEMORY_USAGE_UPDATE_INTERVAL,
+    TASK_LOG_REFRESH_INTERVAL,
+    TASK_STREAM_UPDATE_INTERVAL,
+    WORKER_PROCESSORS_REFRESH_INTERVAL,
+)
+from scaler.ui.live_display import SchedulerSection, WorkersSection
+from scaler.ui.memory_window import MemoryChart
+from scaler.ui.setting_page import Settings
+from scaler.ui.task_graph import TaskStream
+from scaler.ui.task_log import TaskLogTable
+from scaler.ui.worker_processors import WorkerProcessors
+from scaler.utility.formatter import format_bytes, format_percentage
+from scaler.utility.logging.utility import setup_logger
+
+
+@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 start_webui(
+    address: str,
+    host: str,
+    port: int,
+    logging_paths: Tuple[str, ...],
+    logging_config_file: Optional[str],
+    logging_level: str,
+):
+
+    setup_logger(logging_paths, logging_config_file, logging_level)
+
+    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=ZMQConfig.from_string(address),
+        callback=partial(__show_status, tables=tables),
+        topic=b"",
+        timeout_seconds=-1,
+    )
+    subscriber.start()
+
+    ui_thread = threading.Thread(target=partial(ui.run, host=host, port=port, reload=False), daemon=False)
+    ui_thread.start()
+    ui_thread.join()
+
+
+def __show_status(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)
+
+    for worker_data in data.worker_manager.workers:
+        worker_name = worker_data.worker_id.decode()
+        tables.workers_section.workers[worker_name].populate(worker_data)
+
+    tables.worker_processors.update_data(data.worker_manager.workers)

scaler/ui/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.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/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)

scaler/utility/event_loop.py

@@ -0,0 +1,58 @@
+import asyncio
+import enum
+import logging
+from typing import Awaitable, Callable
+
+
+class EventLoopType(enum.Enum):
+    builtin = enum.auto()
+    uvloop = enum.auto()
+
+    @staticmethod
+    def allowed_types():
+        return {m.name for m in EventLoopType}
+
+
+def register_event_loop(event_loop_type: str):
+    if event_loop_type not in EventLoopType.allowed_types():
+        raise TypeError(f"allowed event loop types are: {EventLoopType.allowed_types()}")
+
+    event_loop_type_enum = EventLoopType[event_loop_type]
+    if event_loop_type_enum == EventLoopType.uvloop:
+        try:
+            import uvloop  # noqa
+        except ImportError:
+            raise ImportError("please use pip install uvloop if try to use uvloop as event loop")
+
+        uvloop.install()
+
+    assert event_loop_type in EventLoopType.allowed_types()
+
+    logging.info(f"use event loop: {event_loop_type}")
+
+
+def create_async_loop_routine(routine: Callable[[], Awaitable], seconds: int):
+    """create async loop routine,
+
+    - if seconds is negative, means disable
+    - 0 means looping without any wait, as fast as possible
+    - positive number means execute routine every positive seconds, if passing 1 means run once every 1 seconds"""
+
+    async def loop():
+        if seconds < 0:
+            logging.info(f"{routine.__self__.__class__.__name__}: disabled")  # type: ignore[attr-defined]
+            return
+
+        logging.info(f"{routine.__self__.__class__.__name__}: started")  # type: ignore[attr-defined]
+        try:
+            while True:
+                await routine()
+                await asyncio.sleep(seconds)
+        except asyncio.CancelledError:
+            pass
+        except KeyboardInterrupt:
+            pass
+
+        logging.info(f"{routine.__self__.__class__.__name__}: exited")  # type: ignore[attr-defined]
+
+    return loop()

scaler/utility/exceptions.py

@@ -0,0 +1,42 @@
+class TaskNotFoundError(Exception):
+    pass
+
+
+class WorkerDiedError(Exception):
+    pass
+
+
+class NoWorkerError(Exception):
+    pass
+
+
+class DisconnectedError(Exception):
+    pass
+
+
+class ProcessorDiedError(Exception):
+    pass
+
+
+class DeserializeObjectError(Exception):
+    pass
+
+
+class MissingObjects(Exception):
+    pass
+
+
+class ClientCancelledException(Exception):
+    pass
+
+
+class ClientShutdownException(Exception):
+    pass
+
+
+class ClientQuitException(Exception):
+    pass
+
+
+class ObjectStorageException(Exception):
+    pass

scaler/utility/formatter.py

@@ -0,0 +1,44 @@
+STORAGE_SIZE_MODULUS = 1024.0
+TIME_MODULUS = 1000
+
+
+def format_bytes(number) -> str:
+    for unit in ["B", "K", "M", "G", "T"]:
+        if number >= STORAGE_SIZE_MODULUS:
+            number /= STORAGE_SIZE_MODULUS
+            continue
+
+        if unit in {"B", "K"}:
+            return f"{int(number)}{unit}"
+
+        return f"{number:.1f}{unit}"
+
+    raise ValueError("This should not happen")
+
+
+def format_integer(number):
+    return f"{number:,}"
+
+
+def format_percentage(number: int):
+    return f"{(number/1000):.1%}"
+
+
+def format_microseconds(number: int):
+    for unit in ["us", "ms", "s"]:
+        if number >= TIME_MODULUS:
+            number = int(number / TIME_MODULUS)
+            continue
+
+        if unit == "us":
+            return f"{number/TIME_MODULUS:.1f}ms"
+
+        too_big_sign = "+" if unit == "s" and number > TIME_MODULUS else ""
+        return f"{int(number)}{too_big_sign}{unit}"
+
+
+def format_seconds(number: int):
+    if number > 60:
+        return "60+s"
+
+    return f"{number}s"

scaler/utility/graph/optimization.py

@@ -0,0 +1,27 @@
+from collections import deque
+from typing import Any, Callable, Dict, List, Tuple, Union
+
+
+def cull_graph(
+    graph: Dict[str, Tuple[Union[Callable, Any], ...]], keys: List[str]
+) -> Dict[str, Tuple[Union[Callable, Any], ...]]:
+    queue = deque(keys)
+    visited = set()
+    for target_key in keys:
+        visited.add(target_key)
+
+    while queue:
+        key = queue.popleft()
+
+        task = graph[key]
+        if not (isinstance(task, tuple) and task and callable(task[0])):
+            continue
+
+        dependencies = set(task[1:])
+        for predecessor_key in dependencies:
+            if predecessor_key in visited:
+                continue
+            visited.add(predecessor_key)
+            queue.append(predecessor_key)
+
+    return {key: graph[key] for key in visited}