Skip to content

Commit 01d1503

Browse files
authored
Support config v2 for nemesis and stress tool (#20688)
1 parent f01a94b commit 01d1503

File tree

4 files changed

+99
-45
lines changed

4 files changed

+99
-45
lines changed

ydb/tests/library/harness/kikimr_cluster.py

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -29,26 +29,33 @@ def kikimr_cluster_factory(configurator=None, config_path=None):
2929
class ExternalKiKiMRCluster(KiKiMRClusterInterface):
3030
def __init__(
3131
self,
32-
config_path,
32+
cluster_template,
3333
kikimr_configure_binary_path,
3434
kikimr_path,
3535
kikimr_next_path=None,
3636
ssh_username=None,
3737
deploy_cluster=False,
38-
):
39-
self.__config_path = config_path
40-
with open(config_path, 'r') as r:
41-
self.__yaml_config = yaml.safe_load(r.read())
38+
yaml_config=None):
39+
with open(cluster_template, 'r') as r:
40+
self.__cluster_template = yaml.safe_load(r.read())
41+
if yaml_config is not None:
42+
with open(yaml_config, 'r') as r:
43+
self.__yaml_config = yaml.safe_load(r.read())
4244
self.__kikimr_configure_binary_path = kikimr_configure_binary_path
43-
self.__hosts = [host.get('name', host.get('host')) for host in self.__yaml_config.get('hosts', [])]
4445
self._slots = None
4546
self.__kikimr_path = kikimr_path
4647
self.__kikimr_next_path = kikimr_next_path
4748
self.__ssh_username = ssh_username
4849
self.__deploy_cluster = deploy_cluster
49-
self.__slot_count = 0
5050

51-
for domain in self.__yaml_config['domains']:
51+
if yaml_config is not None:
52+
self.__hosts = [host.get('name', host.get('host')) for host in self.__yaml_config.get('config', {}).get('hosts')]
53+
else:
54+
# Backward compatibility for cluster_template
55+
self.__hosts = [host.get('name', host.get('host')) for host in self.__cluster_template.get('hosts')]
56+
57+
self.__slot_count = 0
58+
for domain in self.__cluster_template['domains']:
5259
self.__slot_count = max(self.__slot_count, domain['dynamic_slots'])
5360

5461
super(ExternalKiKiMRCluster, self).__init__()
@@ -133,7 +140,7 @@ def _prepare_cluster(self):
133140
self._run_on(
134141
inst_set,
135142
lambda x: x.prepare_artifacts(
136-
self.__config_path
143+
self.__cluster_template
137144
)
138145
)
139146

ydb/tests/library/harness/kikimr_runner.py

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1022,8 +1022,9 @@ def switch_version(self):
10221022
self.update_binary_links()
10231023

10241024
def prepare_artifacts(self, cluster_yml):
1025-
self.copy_file_or_dir(
1026-
self.__kikimr_configure_binary_path, self.kikimr_configure_binary_deploy_path)
1025+
if self.__kikimr_configure_binary_path is not None:
1026+
self.copy_file_or_dir(
1027+
self.__kikimr_configure_binary_path, self.kikimr_configure_binary_deploy_path)
10271028

10281029
for version, local_driver in zip(self.versions, self.local_drivers_path):
10291030
self.ssh_command("sudo rm -rf %s" % version)
@@ -1033,14 +1034,15 @@ def prepare_artifacts(self, cluster_yml):
10331034
self.ssh_command("sudo /sbin/setcap 'CAP_SYS_RAWIO,CAP_SYS_NICE=ep' %s" % version)
10341035

10351036
self.update_binary_links()
1036-
self.ssh_command("sudo mkdir -p %s" % self.kikimr_configuration_deploy_path)
1037-
self.copy_file_or_dir(cluster_yml, self.kikimr_cluster_yaml_deploy_path)
1038-
self.ssh_command(self.__generate_configs_cmd())
1039-
self.ssh_command(
1040-
self.__generate_configs_cmd(
1041-
"--dynamic"
1037+
if self.__kikimr_configure_binary_path is not None:
1038+
self.ssh_command("sudo mkdir -p %s" % self.kikimr_configuration_deploy_path)
1039+
self.copy_file_or_dir(cluster_yml, self.kikimr_cluster_yaml_deploy_path)
1040+
self.ssh_command(self.__generate_configs_cmd())
1041+
self.ssh_command(
1042+
self.__generate_configs_cmd(
1043+
"--dynamic"
1044+
)
10421045
)
1043-
)
10441046

10451047
def format_pdisk(self, pdisk_id):
10461048
pass

ydb/tests/stability/tool/__main__.py

Lines changed: 46 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -224,13 +224,14 @@ class bcolors:
224224

225225

226226
class StabilityCluster:
227-
def __init__(self, ssh_username, cluster_path, ydbd_path=None, ydbd_next_path=None):
227+
def __init__(self, ssh_username, cluster_path, ydbd_path=None, ydbd_next_path=None, yaml_config=None):
228228
self.working_dir = os.path.join(tempfile.gettempdir(), "ydb_stability")
229229
os.makedirs(self.working_dir, exist_ok=True)
230230
self.ssh_username = ssh_username
231231
self.slice_directory = cluster_path
232232
self.ydbd_path = ydbd_path
233233
self.ydbd_next_path = ydbd_next_path
234+
self.yaml_config = yaml_config
234235

235236
self.artifacts = (
236237
self._unpack_resource('nemesis'),
@@ -244,14 +245,24 @@ def __init__(self, ssh_username, cluster_path, ydbd_path=None, ydbd_next_path=No
244245
self._unpack_resource('ydb_cli'),
245246
)
246247

247-
self.kikimr_cluster = ExternalKiKiMRCluster(
248-
config_path=self.slice_directory,
249-
kikimr_configure_binary_path=self._unpack_resource("cfg"),
250-
kikimr_path=self.ydbd_path,
251-
kikimr_next_path=self.ydbd_next_path,
252-
ssh_username=self.ssh_username,
253-
deploy_cluster=True,
254-
)
248+
if self.yaml_config is None:
249+
self.kikimr_cluster = ExternalKiKiMRCluster(
250+
cluster_template=self.slice_directory,
251+
kikimr_configure_binary_path=self._unpack_resource("cfg"),
252+
kikimr_path=self.ydbd_path,
253+
kikimr_next_path=self.ydbd_next_path,
254+
ssh_username=self.ssh_username,
255+
deploy_cluster=True,
256+
)
257+
else:
258+
self.kikimr_cluster = ExternalKiKiMRCluster(
259+
cluster_template=self.slice_directory,
260+
kikimr_configure_binary_path=None,
261+
kikimr_path=self.ydbd_path,
262+
kikimr_next_path=self.ydbd_next_path,
263+
ssh_username=self.ssh_username,
264+
yaml_config=self.yaml_config,
265+
)
255266

256267
def _unpack_resource(self, name):
257268
res = resource.find(name)
@@ -401,7 +412,7 @@ def perform_checks(self):
401412
print(f' {node}: {minidumps_search_results[node]}')
402413

403414
def start_nemesis(self):
404-
self.prepare_cluster_yaml()
415+
self.prepare_config_files()
405416
with ThreadPoolExecutor() as pool:
406417
pool.map(lambda node: node.ssh_command(DICT_OF_SERVICES['nemesis']['start_command'], raise_on_error=True), self.kikimr_cluster.nodes.values())
407418

@@ -603,17 +614,27 @@ def deploy_node_tools(self, node):
603614
)
604615
node.ssh_command(f"sudo chmod 777 {node_artifact_path}", raise_on_error=False)
605616

606-
def prepare_cluster_yaml(self):
617+
def prepare_config_files(self):
607618
with ThreadPoolExecutor() as pool:
608-
pool.map(lambda node: node.copy_file_or_dir(
609-
self.slice_directory,
610-
'/Berkanavt/kikimr/cfg/cluster.yaml'
611-
), self.kikimr_cluster.nodes.values())
619+
if self.yaml_config is None:
620+
pool.map(lambda node: node.copy_file_or_dir(
621+
self.slice_directory,
622+
'/Berkanavt/kikimr/cfg/cluster.yaml'
623+
), self.kikimr_cluster.nodes.values())
624+
else:
625+
pool.map(lambda node: node.copy_file_or_dir(
626+
self.slice_directory,
627+
'/Berkanavt/kikimr/cfg/databases.yaml'
628+
), self.kikimr_cluster.nodes.values())
629+
pool.map(lambda node: node.copy_file_or_dir(
630+
self.yaml_config,
631+
'/Berkanavt/kikimr/cfg/config.yaml'
632+
), self.kikimr_cluster.nodes.values())
612633

613634
def deploy_tools(self):
614635
with ThreadPoolExecutor(len(self.kikimr_cluster.nodes)) as pool:
615636
pool.map(self.deploy_node_tools, self.kikimr_cluster.nodes.values())
616-
self.prepare_cluster_yaml()
637+
self.prepare_config_files()
617638

618639
def get_workload_outputs(self, mode='err', last_n_lines=10):
619640
"""Capture last N lines of output from all running workload screens."""
@@ -1040,6 +1061,13 @@ def parse_args():
10401061
type=path_type,
10411062
help="Path to next ydbd version binary (for cross-version testing)",
10421063
)
1064+
parser.add_argument(
1065+
"--yaml-config",
1066+
required=False,
1067+
default=None,
1068+
type=path_type,
1069+
help="Path to Yandex DB configuration v2",
1070+
)
10431071
parser.add_argument(
10441072
"--ssh_user",
10451073
required=False,
@@ -1184,12 +1212,14 @@ def parse_args():
11841212
def main():
11851213
args = parse_args()
11861214
ssh_username = args.ssh_user
1215+
yaml_config = args.yaml_config
11871216
print('Initing cluster info')
11881217
stability_cluster = StabilityCluster(
11891218
ssh_username=ssh_username,
11901219
cluster_path=args.cluster_path,
11911220
ydbd_path=args.ydbd_path,
11921221
ydbd_next_path=args.next_ydbd_path,
1222+
yaml_config=yaml_config,
11931223
)
11941224

11951225
for action in args.actions:

ydb/tests/tools/nemesis/driver/__main__.py

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -128,25 +128,40 @@ def __exit__(self, exc_type, exc_val, exc_tb):
128128
def nemesis_logic(arguments):
129129
logging.config.dictConfig(setup_logging_config(arguments.log_file))
130130
ssh_username = os.getenv('NEMESIS_USER', 'robot-nemesis')
131-
nemesis = catalog.nemesis_factory(
132-
ExternalKiKiMRCluster(
133-
arguments.ydb_cluster_template,
134-
kikimr_configure_binary_path=None,
135-
kikimr_path=arguments.ydb_binary_path,
131+
yaml_config = arguments.yaml_config
132+
if yaml_config is not None:
133+
nemesis = catalog.nemesis_factory(
134+
ExternalKiKiMRCluster(
135+
cluster_template=arguments.ydb_cluster_template,
136+
kikimr_configure_binary_path=None,
137+
kikimr_path=arguments.ydb_binary_path,
138+
ssh_username=ssh_username,
139+
yaml_config=yaml_config,
140+
),
136141
ssh_username=ssh_username,
137-
),
138-
ssh_username=ssh_username,
139-
enable_nemesis_list_filter_by_hostname=arguments.enable_nemesis_list_filter_by_hostname,
140-
)
142+
enable_nemesis_list_filter_by_hostname=arguments.enable_nemesis_list_filter_by_hostname,
143+
)
144+
else:
145+
nemesis = catalog.nemesis_factory(
146+
ExternalKiKiMRCluster(
147+
cluster_template=arguments.ydb_cluster_template,
148+
kikimr_configure_binary_path=None,
149+
kikimr_path=arguments.ydb_binary_path,
150+
ssh_username=ssh_username,
151+
),
152+
ssh_username=ssh_username,
153+
enable_nemesis_list_filter_by_hostname=arguments.enable_nemesis_list_filter_by_hostname,
154+
)
141155
nemesis.start()
142156
monitor.setup_page(arguments.mon_host, arguments.mon_port)
143157
nemesis.stop()
144158

145159

146160
def main():
147161
parser = argparse.ArgumentParser(formatter_class=argparse.RawDescriptionHelpFormatter)
148-
parser.add_argument('--ydb-cluster-template', required=True, help='Path to the Yandex DB cluster template')
149-
parser.add_argument('--ydb-binary-path', required=True, help='Path to the Yandex DB binary')
162+
parser.add_argument('--ydb-cluster-template', required=True, help='Path to the YDB cluster template')
163+
parser.add_argument('--ydb-binary-path', required=True, help='Path to the YDB binary')
164+
parser.add_argument('--yaml-config', required=False, default=None, help='Path to the YDB configuration v2')
150165
parser.add_argument('--private-key-file', default='')
151166
parser.add_argument('--log-file', default=None)
152167
parser.add_argument('--mon-port', default=8666, type=lambda x: int(x))

0 commit comments

Comments
 (0)