1from urllib.parse import urljoin
2
3from tornado import web
4from tornado.ioloop import IOLoop
5
6from .components.nvml import gpu_doc  # noqa: 1708
7from .components.nvml import NVML_ENABLED, gpu_memory_doc, gpu_utilization_doc
8from .components.scheduler import (
9    AggregateAction,
10    BandwidthTypes,
11    BandwidthWorkers,
12    ClusterMemory,
13    ComputePerKey,
14    CurrentLoad,
15    MemoryByKey,
16    Occupancy,
17    SystemMonitor,
18    SystemTimeseries,
19    TaskGraph,
20    TaskGroupGraph,
21    TaskProgress,
22    TaskStream,
23    WorkerNetworkBandwidth,
24    WorkersMemory,
25    WorkerTable,
26    events_doc,
27    graph_doc,
28    individual_doc,
29    individual_profile_doc,
30    individual_profile_server_doc,
31    profile_doc,
32    profile_server_doc,
33    status_doc,
34    stealing_doc,
35    systemmonitor_doc,
36    tasks_doc,
37    tg_graph_doc,
38    workers_doc,
39)
40from .core import BokehApplication
41from .worker import counters_doc
42
43applications = {
44    "/system": systemmonitor_doc,
45    "/stealing": stealing_doc,
46    "/workers": workers_doc,
47    "/events": events_doc,
48    "/counters": counters_doc,
49    "/tasks": tasks_doc,
50    "/status": status_doc,
51    "/profile": profile_doc,
52    "/profile-server": profile_server_doc,
53    "/graph": graph_doc,
54    "/groups": tg_graph_doc,
55    "/gpu": gpu_doc,
56    "/individual-task-stream": individual_doc(
57        TaskStream, 100, n_rectangles=1000, clear_interval="10s"
58    ),
59    "/individual-progress": individual_doc(TaskProgress, 100, height=160),
60    "/individual-graph": individual_doc(TaskGraph, 200),
61    "/individual-groups": individual_doc(TaskGroupGraph, 200),
62    "/individual-workers-memory": individual_doc(WorkersMemory, 100),
63    "/individual-cluster-memory": individual_doc(ClusterMemory, 100),
64    "/individual-cpu": individual_doc(CurrentLoad, 100, fig_attr="cpu_figure"),
65    "/individual-nprocessing": individual_doc(
66        CurrentLoad, 100, fig_attr="processing_figure"
67    ),
68    "/individual-occupancy": individual_doc(Occupancy, 100),
69    "/individual-workers": individual_doc(WorkerTable, 500),
70    "/individual-bandwidth-types": individual_doc(BandwidthTypes, 500),
71    "/individual-bandwidth-workers": individual_doc(BandwidthWorkers, 500),
72    "/individual-workers-network": individual_doc(
73        WorkerNetworkBandwidth, 500, fig_attr="bandwidth"
74    ),
75    "/individual-workers-disk": individual_doc(
76        WorkerNetworkBandwidth, 500, fig_attr="disk"
77    ),
78    "/individual-workers-network-timeseries": individual_doc(
79        SystemTimeseries, 500, fig_attr="bandwidth"
80    ),
81    "/individual-workers-cpu-timeseries": individual_doc(
82        SystemTimeseries, 500, fig_attr="cpu"
83    ),
84    "/individual-workers-memory-timeseries": individual_doc(
85        SystemTimeseries, 500, fig_attr="memory"
86    ),
87    "/individual-workers-disk-timeseries": individual_doc(
88        SystemTimeseries, 500, fig_attr="disk"
89    ),
90    "/individual-memory-by-key": individual_doc(MemoryByKey, 500),
91    "/individual-compute-time-per-key": individual_doc(ComputePerKey, 500),
92    "/individual-aggregate-time-per-action": individual_doc(AggregateAction, 500),
93    "/individual-scheduler-system": individual_doc(SystemMonitor, 500),
94    "/individual-profile": individual_profile_doc,
95    "/individual-profile-server": individual_profile_server_doc,
96    "/individual-gpu-memory": gpu_memory_doc,
97    "/individual-gpu-utilization": gpu_utilization_doc,
98}
99
100
101template_variables = {
102    "pages": [
103        "status",
104        "workers",
105        "tasks",
106        "system",
107        "profile",
108        "graph",
109        "groups",
110        "info",
111    ],
112    "plots": [x.replace("/", "") for x in applications if "individual" in x],
113}
114
115if NVML_ENABLED:
116    template_variables["pages"].insert(4, "gpu")
117
118
119def connect(application, http_server, scheduler, prefix=""):
120    bokeh_app = BokehApplication(
121        applications, scheduler, prefix=prefix, template_variables=template_variables
122    )
123    application.add_application(bokeh_app)
124    bokeh_app.initialize(IOLoop.current())
125
126    bokeh_app.add_handlers(
127        r".*",
128        [
129            (
130                r"/",
131                web.RedirectHandler,
132                {"url": urljoin((prefix or "").strip("/") + "/", r"status")},
133            )
134        ],
135    )
136