@@ -7,36 +7,38 @@ def cli():
77 cli .add_argument ('-c' ,'--clara' , metavar = 'DIR' ,help = 'CLARA_HOME path (default=$CLARA_HOME)' ,default = os .getenv ('CLARA_HOME' ,None ))
88 cli .add_argument ('-t' ,'--threads' ,metavar = '#' ,help = 'threads (default=8,16,24,34,44)' ,default = '8,16,24,34,44' )
99 cli .add_argument ('-e' ,'--events' , metavar = '#' ,help = 'events per thread (default=555)' ,default = 555 ,type = int )
10- cli .add_argument ('-N' ,'--numa' , metavar = '#' ,help = 'NUMA socket (default=None, choices=[0,1])' ,default = None ,type = int ,choices = [0 ,1 ])
10+ cli .add_argument ('-N' ,'--numa' , metavar = '#' ,help = 'restrict to a NUMA socket (choices=[0,1])' ,choices = [0 ,1 ])
11+ cli .add_argument ('-C' ,'--cpus' , metavar = '#' ,help = 'restrict to a CPU list (e.g. 0,1,2,3)' )
12+ cli .add_argument ('-T' ,'--cputhr' , metavar = '#' ,help = 'restrict to same number of CPUs as threads' )
1113 cli .add_argument ('datafile' , help = 'input EVIO/HIPO data file' )
1214 cfg = cli .parse_args ()
1315 cfg .threads = cfg .threads .split (',' )
1416 if not cfg .clara :
1517 cli .error ('cannot find CLARA installation via -c or $CLARA_HOME' )
16- cfg .run_clara = find_run_clara (cfg .clara )
17- if not cfg .run_clara :
18- cli .error ('cannot find run-clara in $PATH or CLARA installation' )
19- if cfg .numa is not None :
18+ elif shutil .which ('run-clara' ):
19+ cfg .run_clara = shutil .which ('run-clara' )
20+ elif os .path .exists (clara_home + '/plugins/clas12/bin/run-clara' ):
21+ cfg .run_clara = clara_home + '/plugins/clas12/bin/run-clara'
22+ else :
23+ cli .error ('cannot find run-clara in $PATH or $CLARA_HOME' )
24+ if cfg .cpus :
25+ try :
26+ cfg .cpus = ',' .join ([str (int (i )) for i in cfg .cpus .split (',' )])
27+ except Exception as e :
28+ cli .error ('invalid --cpus argument: {cfg.cpus}' )
29+ elif cfg .numa or cfg .cputhr :
30+ if not cfg .numa :
31+ cfg .numa = 0
2032 if not shutil .which ('numactl' ):
2133 cli .error ('numactl is not in $PATH, --numa option not supported' )
22- if len (list (get_numa_cpus (cfg .numa ))) == 0 :
23- cli .error ('invalid --numa node: {cfg.numa}' )
24- cfg .numa = ',' .join (list (get_numa_cpus (cfg .numa )))
34+ for x in run (['numactl' ,'-H' ]):
35+ if x .startswith (f'node { cfg .numa } cpus:' ):
36+ cfg .cpus = ',' .join ([str (int (i )) for i in x .strip ().split ()[3 :]])
37+ break
38+ if not cfg .cpus :
39+ cli .error ('error determining cpu list for -C or -T' )
2540 return cfg
2641
27- def find_run_clara (clara_home ):
28- import os ,shutil
29- if shutil .which ('run-clara' ):
30- return shutil .which ('run-clara' )
31- elif os .path .exists (clara_home + '/plugins/clas12/bin/run-clara' ):
32- return clara_home + '/plugins/clas12/bin/run-clara'
33-
34- def get_numa_cpus (node ):
35- for line in run (['numactl' ,'-H' ]):
36- if line .startswith (f'node { node } cpus:' ):
37- for col in line .strip ().split ()[3 :]:
38- yield col
39-
4042def run (cmd ):
4143 import subprocess
4244 print ('run >>> ' + ' ' .join (cmd ))
@@ -46,15 +48,16 @@ def run(cmd):
4648 if len (line ) > 0 :
4749 yield line
4850 p .wait ()
49- if p .returncode != 0 :
50- pass
5151
5252def benchmark (cfg , threads , log ):
5353 import collections
5454 cmd = []
55- # use taskset:
56- if cfg .numa is not None :
57- cmd .extend (['taskset' ,'-c' ,cfg .numa ])
55+ # use taskset for thread affinity:
56+ if cfg .cpus :
57+ if cfg .cputhr :
58+ cmd .extend (['taskset' ,'-c' ,',' .join (cfg .cpus .split (',' )[:threads ])])
59+ else :
60+ cmd .extend (['taskset' ,'-c' ,cfg .cpus ])
5861 # add the run-clara command:
5962 cmd .extend ([cfg .run_clara ,
6063 '-c' ,cfg .clara ,
0 commit comments