Skip to content

Commit 3f631ff

Browse files
phvalguimaphvalguima
authored
Merge pull request #1 from canonical/DPE-5997-py-wrapper
[DPE-6253] Add benchmark wrapper logic
2 parents 894773b + 464a6a8 commit 3f631ff

File tree

10 files changed

+3332
-0
lines changed

10 files changed

+3332
-0
lines changed

‎poetry.lock

+2,333
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎pyproject.toml

+129
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
# Copyright 2023 Canonical Ltd.
2+
# See LICENSE file for licensing details.
3+
4+
[tool.poetry]
5+
package-mode = false
6+
requires-poetry = ">=2.0.0"
7+
8+
[tool.poetry.dependencies]
9+
python = "^3.10"
10+
ops = "^2.17.0"
11+
tenacity = "^9.0.0"
12+
jinja2 = "^3.1.4"
13+
overrides = "7.7.0"
14+
requests = "2.32.3"
15+
shortuuid = "1.0.13"
16+
cryptography = "^43.0.1"
17+
jsonschema = "^4.23.0"
18+
prometheus-client = ">=0.19.0"
19+
pydantic = "^1.10.18, <2"
20+
# pydantic = ">=2.0,<3.0"
21+
fastapi = ">=0.115.0"
22+
uvicorn = ">0.11.5"
23+
kafka-python = ">=2.0"
24+
25+
[tool.poetry.group.charm-libs.dependencies]
26+
# data_platform_libs/v0/data_interfaces.py
27+
ops = "^2.17"
28+
# data_platform_libs/v0/upgrade.py
29+
# grafana_agent/v0/cos_agent.py requires pydantic <2
30+
pydantic = "^1.10, <2"
31+
# tls_certificates_interface/v1/tls_certificates.py
32+
cryptography = "^43.0.0"
33+
jsonschema = "^4.23.0"
34+
# grafana_agent/v0/cos_agent.py
35+
cosl = "^0.0.41"
36+
bcrypt = "^4.1.3"
37+
38+
[tool.poetry.group.format]
39+
optional = true
40+
41+
[tool.poetry.group.format.dependencies]
42+
ruff = "^0.6.8"
43+
44+
[tool.poetry.group.lint]
45+
optional = true
46+
47+
[tool.poetry.group.lint.dependencies]
48+
codespell = "^2.3.0"
49+
shellcheck-py = "^0.10.0.1"
50+
51+
[tool.poetry.group.unit.dependencies]
52+
pytest = "^8.3.3"
53+
pytest-asyncio = "^0.21.2"
54+
coverage = {extras = ["toml"], version = "^7.6.1"}
55+
parameterized = "^0.9.0"
56+
57+
[tool.poetry.group.integration.dependencies]
58+
pytest = "^8.3.3"
59+
pytest-github-secrets = {git = "https://github.com/canonical/data-platform-workflows", tag = "v24.0.6", subdirectory = "python/pytest_plugins/github_secrets"}
60+
pytest-operator = "^0.37.0"
61+
pytest-operator-cache = {git = "https://github.com/canonical/data-platform-workflows", tag = "v24.0.6", subdirectory = "python/pytest_plugins/pytest_operator_cache"}
62+
pytest-operator-groups = {git = "https://github.com/canonical/data-platform-workflows", tag = "v24.0.6", subdirectory = "python/pytest_plugins/pytest_operator_groups"}
63+
pytest-microceph = {git = "https://github.com/canonical/data-platform-workflows", tag = "v24.0.6", subdirectory = "python/pytest_plugins/microceph"}
64+
juju = "^3.5.2"
65+
ops = "^2.17.0"
66+
tenacity = "^9.0.0"
67+
pyyaml = "^6.0.2"
68+
urllib3 = "^2.2.3"
69+
protobuf = "5.28.2"
70+
71+
[tool.coverage.run]
72+
branch = true
73+
74+
[tool.coverage.report]
75+
show_missing = true
76+
77+
[tool.pytest.ini_options]
78+
minversion = "6.0"
79+
log_cli_level = "INFO"
80+
markers = ["unstable"]
81+
asyncio_mode = "auto"
82+
83+
# Formatting tools configuration
84+
[tool.black]
85+
line-length = 99
86+
target-version = ["py310"]
87+
88+
# Linting tools configuration
89+
[tool.ruff]
90+
# preview and explicit preview are enabled for CPY001
91+
preview = true
92+
target-version = "py310"
93+
src = ["src", "."]
94+
line-length = 99
95+
96+
[tool.ruff.lint]
97+
explicit-preview-rules = true
98+
select = ["A", "E", "W", "F", "C", "N", "D", "I001", "CPY001"]
99+
extend-ignore = [
100+
"D203",
101+
"D204",
102+
"D213",
103+
"D215",
104+
"D400",
105+
"D404",
106+
"D406",
107+
"D407",
108+
"D408",
109+
"D409",
110+
"D413",
111+
]
112+
# Ignore E501 because using black creates errors with this
113+
# Ignore D107 Missing docstring in __init__
114+
ignore = ["E501", "D107"]
115+
116+
[tool.ruff.lint.per-file-ignores]
117+
"tests/*" = ["D100", "D101", "D102", "D103", "D104"]
118+
119+
[tool.ruff.lint.flake8-copyright]
120+
# Check for properly formatted copyright header in each file
121+
author = "Canonical Ltd."
122+
notice-rgx = "Copyright\\s\\d{4}([-,]\\d{4})*\\s+"
123+
min-file-size = 1
124+
125+
[tool.ruff.lint.mccabe]
126+
max-complexity = 10
127+
128+
[tool.ruff.lint.pydocstyle]
129+
convention = "google"

‎src/benchmark/literals.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# Copyright 2024 Canonical Ltd.
2+
# See LICENSE file for licensing details.
3+
4+
"""This module contains the constants and models used by the sysbench charm."""
5+
6+
BENCHMARK_WORKLOAD_PATH = "/root/.benchmark/charmed_parameters"

‎src/benchmark/wrapper/core.py

+133
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
# Copyright 2024 Canonical Ltd.
2+
# See LICENSE file for licensing details.
3+
4+
"""The core models for the wrapper script."""
5+
6+
from enum import Enum
7+
8+
from prometheus_client import Gauge
9+
from pydantic import BaseModel
10+
11+
12+
class BenchmarkCommand(str, Enum):
13+
"""Enum to hold the benchmark phase."""
14+
15+
PREPARE = "prepare"
16+
RUN = "run"
17+
STOP = "stop"
18+
COLLECT = "collect"
19+
UPLOAD = "upload"
20+
CLEANUP = "cleanup"
21+
22+
23+
class ProcessStatus(str, Enum):
24+
"""Enum to hold the process status."""
25+
26+
RUNNING = "running"
27+
STOPPED = "stopped"
28+
ERROR = "error"
29+
TO_START = "to_start"
30+
31+
32+
class ProcessModel(BaseModel):
33+
"""Model to hold the process information."""
34+
35+
cmd: str
36+
pid: int = -1
37+
status: str = ProcessStatus.TO_START
38+
user: str | None = None
39+
group: str | None = None
40+
cwd: str | None = None
41+
42+
43+
class MetricOptionsModel(BaseModel):
44+
"""Model to hold the metrics."""
45+
46+
label: str | None = None
47+
extra_labels: list[str] = []
48+
description: str | None = None
49+
50+
51+
class WorkloadCLIArgsModel(BaseModel):
52+
"""Model to hold the workload options."""
53+
54+
test_name: str
55+
command: BenchmarkCommand
56+
workload: str
57+
threads: int
58+
parallel_processes: int
59+
duration: int
60+
run_count: int
61+
report_interval: int
62+
extra_labels: str
63+
peers: str
64+
log_file: str = "/var/log/dpe_benchmark_workload.log"
65+
is_coordinator: bool
66+
67+
68+
class BenchmarkMetrics:
69+
"""Class to hold the benchmark metrics."""
70+
71+
def __init__(
72+
self,
73+
options: MetricOptionsModel,
74+
):
75+
self.options = options
76+
self.metrics = {}
77+
78+
def add(self, sample: BaseModel):
79+
"""Add the benchmark to the prometheus metric."""
80+
for key, value in sample.dict().items():
81+
if f"{self.options.label}_{key}" not in self.metrics:
82+
self.metrics[f"{self.options.label}_{key}"] = Gauge(
83+
f"{self.options.label}_{key}",
84+
f"{self.options.description} {key}",
85+
["model", "unit"],
86+
)
87+
self.metrics[f"{self.options.label}_{key}"].labels(*self.options.extra_labels).set(
88+
value
89+
)
90+
91+
92+
class KafkaBenchmarkSample(BaseModel):
93+
"""Sample from the benchmark tool."""
94+
95+
produce_rate: float # in msgs / s
96+
produce_throughput: float # in MB/s
97+
produce_error_rate: float # in err/s
98+
99+
produce_latency_avg: float # in (ms)
100+
produce_latency_50: float
101+
produce_latency_99: float
102+
produce_latency_99_9: float
103+
produce_latency_max: float
104+
105+
produce_delay_latency_avg: float # in (us)
106+
produce_delay_latency_50: float
107+
produce_delay_latency_99: float
108+
produce_delay_latency_99_9: float
109+
produce_delay_latency_max: float
110+
111+
consume_rate: float # in msgs / s
112+
consume_throughput: float # in MB/s
113+
consume_backlog: float # in KB
114+
115+
116+
class KafkaBenchmarkSampleMatcher(Enum):
117+
"""Hard-coded regexes to process the benchmark sample."""
118+
119+
produce_rate: str = r"Pub rate\s+(.*?)\s+msg/s"
120+
produce_throughput: str = r"Pub rate\s+\d+.\d+\s+msg/s\s+/\s+(.*?)\s+MB/s"
121+
produce_error_rate: str = r"Pub err\s+(.*?)\s+err/s"
122+
produce_latency_avg: str = r"Pub Latency \(ms\) avg:\s+(.*?)\s+"
123+
# Match: Pub Latency (ms) avg: 1478.1 - 50%: 1312.6 - 99%: 4981.5 - 99.9%: 5104.7 - Max: 5110.5
124+
# Generates: [('1478.1', '1312.6', '4981.5', '5104.7', '5110.5')]
125+
produce_latency_percentiles: str = r"Pub Latency \(ms\) avg:\s+(.*?)\s+- 50%:\s+(.*?)\s+- 99%:\s+(.*?)\s+- 99.9%:\s+(.*?)\s+- Max:\s+(.*?)\s+"
126+
127+
# Pub Delay Latency (us) avg: 21603452.9 - 50%: 21861759.0 - 99%: 23621631.0 - 99.9%: 24160895.0 - Max: 24163839.0
128+
# Generates: [('21603452.9', '21861759.0', '23621631.0', '24160895.0', '24163839.0')]
129+
produce_latency_delay_percentiles: str = r"Pub Delay Latency \(us\) avg:\s+(.*?)\s+- 50%:\s+(.*?)\s+- 99%:\s+(.*?)\s+- 99.9%:\s+(.*?)\s+- Max:\s+(\d+\.\d+)"
130+
131+
consume_rate: str = r"Cons rate\s+(.*?)\s+msg/s"
132+
consume_throughput: str = r"Cons rate\s+\d+.\d+\s+msg/s\s+/\s+(.*?)\s+MB/s"
133+
consume_backlog: str = r"Backlog:\s+(.*?)\s+K"

‎src/benchmark/wrapper/main.py

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
#!/usr/bin/python3
2+
# Copyright 2024 Canonical Ltd.
3+
# See LICENSE file for licensing details.
4+
5+
"""This script runs the benchmark tool, collects its output and forwards to prometheus."""
6+
7+
import logging
8+
import signal
9+
10+
from core import WorkloadCLIArgsModel
11+
from process import WorkloadToProcessMapping
12+
from prometheus_client import start_http_server
13+
14+
15+
class MainWrapper:
16+
"""Main class to manage the benchmark tool."""
17+
18+
mapping: WorkloadToProcessMapping
19+
20+
def __init__(self, args: WorkloadCLIArgsModel):
21+
self.args = args
22+
23+
def run(self):
24+
"""Prepares the workload and runs the benchmark."""
25+
manager, _ = self.mapping.map(self.args.command)
26+
27+
logging.basicConfig(filename=self.args.log_file, encoding="utf-8", level=logging.INFO)
28+
29+
def _exit(*args, **kwargs):
30+
manager.stop()
31+
32+
signal.signal(signal.SIGINT, _exit)
33+
signal.signal(signal.SIGTERM, _exit)
34+
start_http_server(8008)
35+
36+
# Start the manager and process the output
37+
manager.start()
38+
# Now, start the event loop to monitor the processes:
39+
manager.run()
40+
41+
42+
# EXAMPLE
43+
# The code below is an example usage of the main function + the wrapper classes
44+
#
45+
# if __name__ == "__main__":
46+
# parser = argparse.ArgumentParser(
47+
# prog="wrapper", description="Runs the benchmark command as an argument."
48+
# )
49+
# parser.add_argument("--test_name", type=str, help="Test name to be used")
50+
# parser.add_argument("--command", type=str, help="Command to be executed", default="run")
51+
# parser.add_argument(
52+
# "--workload", type=str, help="Name of the workload to be executed", default="default"
53+
# )
54+
# parser.add_argument("--report_interval", type=int, default=10)
55+
# parser.add_argument("--parallel_processes", type=int, default=1)
56+
# parser.add_argument("--threads", type=int, default=1)
57+
# parser.add_argument("--duration", type=int, default=0)
58+
# parser.add_argument("--run_count", type=int, default=1)
59+
# parser.add_argument(
60+
# "--target_hosts", type=str, default="", help="comma-separated list of target hosts"
61+
# )
62+
# parser.add_argument(
63+
# "--log_file", type=str, default="/var/log/dpe_benchmark_workload.log", help="Log file for all threads"
64+
# )
65+
# parser.add_argument(
66+
# "--extra_labels",
67+
# type=str,
68+
# help="comma-separated list of extra labels to be used.",
69+
# default="",
70+
# )
71+
# # Parse the arguments as dictionary, using the same logic as:
72+
# # https://github.com/python/cpython/blob/ \
73+
# # 47c5a0f307cff3ed477528536e8de095c0752efa/Lib/argparse.py#L134
74+
# args = parser.parse_args().__dict__ | {"command": BenchmarkCommand(parser.parse_args().command)}
75+
# MainWrapper(WorkloadCLIArgsModel.parse_obj(args)).run()

0 commit comments

Comments
 (0)