File indexing completed on 2026-07-26 08:22:17
0001 import argparse
0002 import json
0003 import csv
0004 import pathlib
0005 import logging
0006 import functools
0007 import operator
0008 import os
0009 import shutil
0010 import itertools
0011 import random
0012 import tempfile
0013 import time
0014 import subprocess
0015 import re
0016
0017
0018 import traccc_bench_tools.parse_profile
0019 import traccc_bench_tools.types
0020
0021 log = logging.getLogger("traccc_cut_optimiser")
0022
0023
0024 def main():
0025 parser = argparse.ArgumentParser()
0026
0027 parser.add_argument(
0028 "exec",
0029 type=pathlib.Path,
0030 help="the traccc ",
0031 )
0032
0033 parser.add_argument(
0034 "db",
0035 type=pathlib.Path,
0036 help="the CSV database file",
0037 )
0038
0039 parser.add_argument(
0040 "parameters",
0041 type=pathlib.Path,
0042 help="the JSON parameter space file",
0043 )
0044
0045 parser.add_argument(
0046 "-v",
0047 "--verbose",
0048 help="enable verbose output",
0049 action="store_true",
0050 )
0051
0052 parser.add_argument(
0053 "--num-sm",
0054 help="number of SMs in the modelled GPU",
0055 type=int,
0056 required=True,
0057 dest="num_sm",
0058 )
0059
0060 parser.add_argument(
0061 "--num-threads-per-sm",
0062 help="number of thread slots per SM in the modelled GPU",
0063 type=int,
0064 required=True,
0065 dest="num_threads_per_sm",
0066 )
0067
0068 parser.add_argument(
0069 "--ncu-wrapper",
0070 help="wrapper to use around the ncu command",
0071 type=str,
0072 dest="ncu_wrapper",
0073 )
0074
0075 parser.add_argument(
0076 "--timeout",
0077 help="timeout in seconds for benchmarks",
0078 default=120,
0079 type=int,
0080 )
0081
0082 parser.add_argument(
0083 "--random",
0084 help="enable random search rather than grid search",
0085 action="store_true",
0086 )
0087
0088 args = parser.parse_args()
0089
0090 logging.basicConfig(
0091 level=logging.DEBUG if (args.verbose or False) else logging.INFO,
0092 format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
0093 )
0094
0095 log.info(
0096 "Using GPU with %d SMs and %d thread slots per SM (%d thread slots total)",
0097 getattr(args, "num_sm"),
0098 getattr(args, "num_threads_per_sm"),
0099 getattr(args, "num_sm") * getattr(args, "num_threads_per_sm"),
0100 )
0101
0102 gpu_spec = traccc_bench_tools.types.GpuSpec(
0103 getattr(args, "num_sm"), getattr(args, "num_threads_per_sm")
0104 )
0105
0106 with open(args.parameters, "r") as f:
0107 params = json.load(f)
0108
0109 parameter_space = params["parameters"]
0110 parameter_names = sorted(parameter_space)
0111
0112 log.info(
0113 "Running optimisation for %d parameters: %s",
0114 len(params["parameters"]),
0115 ", ".join(parameter_names),
0116 )
0117
0118 csv_field_names = parameter_names + [
0119 "success",
0120 "rec_throughput",
0121 "efficiency",
0122 "fake_rate",
0123 "duplicate_rate",
0124 ]
0125
0126 if args.db.is_file():
0127 log.info('Database file "%s" already exists; creating a backup', args.db)
0128 shutil.copy(str(args.db), str(args.db) + ".bak")
0129 results = {}
0130 with open(args.db, "r") as f:
0131 reader = csv.DictReader(f)
0132 for i in reader:
0133 key = (i[k] for k in parameter_names)
0134 if set(i.keys()) != set(csv_field_names):
0135 raise ValueError(
0136 "Input database has unexpected columns or is missing columns"
0137 )
0138 if i["success"] != 0:
0139 results[key] = {
0140 "rec_throughput": i["rec_throughput"],
0141 "efficiency": i["efficiency"],
0142 "fake_rate": i["fake_rate"],
0143 "duplicate_rate": i["duplicate_rate"],
0144 }
0145 else:
0146 results[key] = None
0147 log.info("Database contained %d pre-existing results", len(results))
0148 else:
0149 log.info('Database file "%s" does not exist; starting from scratch', args.db)
0150 results = {}
0151
0152 log.info("Starting optimisation with %d known results", len(results))
0153
0154 log.info(
0155 "Total parameter space has size %d",
0156 functools.reduce(operator.mul, [len(x) for x in parameter_space.values()], 1),
0157 )
0158
0159 if "TRACCC_TEST_DATA_DIR" not in os.environ:
0160 e = 'Environment variable "TRACCC_TEST_DATA_DIR" is not set; aborting!'
0161 log.fatal(e)
0162 raise RuntimeError(e)
0163
0164 for exec in ["ncu"]:
0165 if shutil.which(exec) is None:
0166 e = 'Executable "%s" is not available; aborting' % exec
0167 log.fatal(e)
0168 raise RuntimeError(e)
0169
0170 space_product = itertools.product(
0171 *[parameter_space[name] for name in parameter_names]
0172 )
0173 random_retry_count = 0
0174 it_count = 0
0175
0176 while True:
0177 if args.random:
0178 param_dict = {n: random.choice(vs) for n, vs in parameter_space.items()}
0179 param_list = tuple(param_dict.values())
0180
0181 if param_list in results:
0182 random_retry_count += 1
0183 if random_retry_count < 10:
0184 log.info(
0185 "Configuration is already known; continuing with random abort counter %d",
0186 random_retry_count,
0187 )
0188 continue
0189 else:
0190 log.info(
0191 "Configuration is already known and max retry counter is reached; aborting"
0192 )
0193 break
0194 else:
0195 random_retry_count = 0
0196 else:
0197 try:
0198 param_vector = next(space_product)
0199 assert len(param_vector) == len(parameter_names)
0200 param_dict = {parameter_names[i]: p for i, p in enumerate(param_vector)}
0201 param_list = tuple(param_dict.values())
0202 except StopIteration:
0203 log.info("Design space has been exhausted; exiting")
0204 break
0205 if param_list in results:
0206 log.info(
0207 "Configuration %s is already known; continuing", str(param_dict)
0208 )
0209 continue
0210
0211 try:
0212 log.info("Running benchmark for parameters %s", str(param_dict))
0213
0214 with tempfile.TemporaryDirectory() as tmpdirname:
0215 tmppath = pathlib.Path(tmpdirname)
0216
0217 log.info('Created temporary directory "%s"', str(tmppath))
0218
0219 log.info("Running benchmark step")
0220
0221 start_time = time.time()
0222
0223 profile_args = [
0224 "ncu",
0225 "--import-source",
0226 "no",
0227 "--section",
0228 "LaunchStats",
0229 "--section",
0230 "Occupancy",
0231 "--metrics",
0232 "gpu__time_duration.sum",
0233 "-f",
0234 "-o",
0235 "profile",
0236 str(args.exec.resolve()),
0237 "--input-directory=%s" % params["input"]["event_dir"],
0238 "--digitization-file=%s" % params["input"]["digitization_file"],
0239 "--conditions-file=%s" % params["input"]["conditions_file"],
0240 "--detector-file=%s" % params["input"]["detector_file"],
0241 "--grid-file=%s" % params["input"]["grid_file"],
0242 "--material-file=%s" % params["input"]["material_file"],
0243 "--input-events=1",
0244 "--check-performance",
0245 "--use-acts-geom-source=on",
0246 ]
0247
0248 for k, v in params["config"].items():
0249 profile_args.append("--%s=%s" % (k, str(v)))
0250
0251 for k, v in param_dict.items():
0252 profile_args.append("--%s=%s" % (k, str(v)))
0253
0254 if args.ncu_wrapper is not None:
0255 profile_args = args.ncu_wrapper.split() + profile_args
0256
0257 try:
0258 result = subprocess.run(
0259 profile_args,
0260 stdout=subprocess.PIPE,
0261 stderr=subprocess.STDOUT,
0262 cwd=tmppath,
0263 check=True,
0264 timeout=args.timeout,
0265 )
0266 except subprocess.CalledProcessError as e:
0267 log.warning("Process failed to execute; continuing")
0268 results[param_list] = None
0269 continue
0270 except subprocess.TimeoutExpired as e:
0271 log.warning("Process timed out; marking as failure and continuing")
0272 results[param_list] = None
0273 continue
0274
0275 stdout = result.stdout.decode("utf-8")
0276
0277 if match := re.search(
0278 r"Total track efficiency was (\d+(?:\.\d*)?)%", stdout
0279 ):
0280 result_efficiency = float(match.group(1)) / 100.0
0281 else:
0282 raise ValueError("Efficiency could not be parsed from stdout!")
0283
0284 if match := re.search(
0285 r"Total track duplicate rate was (\d+(?:\.\d*)?)", stdout
0286 ):
0287 result_duplicate_rate = float(match.group(1))
0288 else:
0289 raise ValueError("Dupliacate rate could not be parsed from stdout!")
0290
0291 if match := re.search(
0292 r"Total track fake rate was (\d+(?:\.\d*)?)", stdout
0293 ):
0294 result_fake_rate = float(match.group(1))
0295 else:
0296 raise ValueError("Fake rate could not be parsed from stdout!")
0297
0298 log.info(
0299 "Physics performance was %.1f%% efficiency, %.1f fake rate, %.1f duplicate rate",
0300 result_efficiency * 100.0,
0301 result_duplicate_rate,
0302 result_fake_rate,
0303 )
0304
0305 end_time = time.time()
0306
0307 log.info(
0308 "Completed benchmark step in %.1f seconds", end_time - start_time
0309 )
0310
0311 log.info("Running profile processing step")
0312
0313 start_time = time.time()
0314
0315 result_df = traccc_bench_tools.parse_profile.parse_profile_ncu(
0316 tmppath / "profile.ncu-rep",
0317 gpu_spec,
0318 event_marker_kernel="count_grid_capacities",
0319 )
0320
0321 total_rec_throughput = result_df["RecThroughputMean"].sum()
0322
0323 log.info(
0324 "Compute performance was a combined reciprocal throughput of %.1fms",
0325 1000.0 * total_rec_throughput,
0326 )
0327
0328 end_time = time.time()
0329
0330 log.info(
0331 "Completed profile processing step in %.1f seconds",
0332 end_time - start_time,
0333 )
0334
0335 results[param_list] = {
0336 "rec_throughput": total_rec_throughput,
0337 "efficiency": result_efficiency,
0338 "fake_rate": result_fake_rate,
0339 "duplicate_rate": result_duplicate_rate,
0340 }
0341
0342 except Exception as e:
0343 log.exception(e)
0344 results[param_list] = None
0345 except KeyboardInterrupt as e:
0346 log.info("Received keyboard interrupt; skipping to post-processing")
0347 break
0348
0349 log.info("Gathered a total of %d results (incl. pre-existing)", len(results))
0350
0351 log.info("Writing data to %s", args.db)
0352 with open(args.db, "w") as f:
0353 writer = csv.DictWriter(
0354 f,
0355 fieldnames=csv_field_names,
0356 )
0357 writer.writeheader()
0358
0359 for k, v in results.items():
0360
0361 param_dict = {parameter_names[i]: p for i, p in enumerate(k)}
0362 if v is None:
0363 writer.writerow(
0364 dict(
0365 **param_dict,
0366 success=0,
0367 rec_throughput=0.0,
0368 efficiency=0.0,
0369 fake_rate=0.0,
0370 duplicate_rate=0.0,
0371 )
0372 )
0373 else:
0374 writer.writerow(
0375 dict(
0376 **param_dict,
0377 **v,
0378 success=1,
0379 )
0380 )
0381
0382 log.info("Processing complete; goodbye!")
0383
0384
0385 if __name__ == "__main__":
0386 main()
