-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
118 lines (96 loc) · 3.79 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import argparse
import logging
import os
import signal
import sys
import threading
import time
from core.logger import setup_logger, LoggerWriter
from core.metrics import Metric, MetricList
from core.runner import ScriptRunner
from core.visual import MetricsDisplay
logger = setup_logger(f"smaug_{os.getpid()}")
sys.stderr = LoggerWriter(logger, logging.ERROR)
class App:
def __init__(self, script_file: str, num: int, use_buffer: bool):
os.makedirs("logs", exist_ok=True)
self.runner = ScriptRunner(script_file, use_buffer)
self.runner.run(num)
self.monitor = self.runner.monitor
self.display = MetricsDisplay()
self.stop_flag = False
self.collect_data_thread = threading.Thread(target=self._collect_data)
self.collect_data_thread.start()
signal.signal(signal.SIGINT, self.signal_handler)
self._wait_scripts()
def signal_handler(self, signal, frame) -> None:
self.stop()
def stop(self) -> None:
self.stop_flag = True
self.runner.stop()
def _collect_data(self) -> None:
collect_interval = 0.3 # 0.3 seconds
while not self.stop_flag:
epoch_now = int(time.time())
cpu_metrics = self.monitor.cpu_monitor.record_stats()
cpu_metrics.append(
Metric('cpu average', self.monitor.cpu_monitor.get_average(), epoch=epoch_now)
)
memory_metrics = self.monitor.memory_monitor.record_stats()
for key, value in self.monitor.memory_monitor.get_average().items():
memory_metrics.append(Metric(f'{key} average', value, epoch=epoch_now))
disk_metrics = self.monitor.disk_monitor.record_stats()
disk_metrics.append(
Metric('disk usage difference',
self.monitor.disk_monitor.get_diff(),
epoch=epoch_now)
)
process_metrics = MetricList([
Metric('execution time',
self.monitor.process_monitor.get_execution_time(),
epoch=epoch_now),
Metric('total thread usage',
self.monitor.process_monitor.get_total_thread_usage(),
epoch=epoch_now),
])
thread_usage_diff = self.monitor.process_monitor.get_diff()['total thread usage diff']
process_metrics.append(
Metric('total thread usage difference', thread_usage_diff, epoch=epoch_now)
)
app_size_metric = Metric('app size',
self.monitor.get_app_size(), epoch=epoch_now)
metrics = (cpu_metrics + memory_metrics + disk_metrics
+ MetricList([app_size_metric]) + process_metrics)
self.display.update(metrics)
time.sleep(collect_interval)
def _wait_scripts(self) -> None:
for process in self.runner.processes:
process.wait()
self.stop()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Run the application with a specified main file."
)
parser.add_argument(
"-mf", "--main-file", type=str, help="The main file to run (path)"
)
parser.add_argument(
"-n",
"--num",
type=int,
default=1,
help="The number of times to run your script. Default is 1",
)
parser.add_argument(
"-ub",
"--use-buffer",
type=lambda x: (str(x).lower() in ["true", "1", "yes"]),
default=True,
help="Use buffer for the script output. Default is True",
)
args = parser.parse_args()
if args.main_file:
num = args.num
main_file = args.main_file
use_buffer = args.use_buffer
app = App(main_file, num, use_buffer)