diff --git a/src/finn/util/vcd.py b/src/finn/util/vcd.py index d9e244422065314ceb790dc6719b57688ff76828..a4400f7bd7e75549189f081ce255fd67c49b3746 100644 --- a/src/finn/util/vcd.py +++ b/src/finn/util/vcd.py @@ -162,16 +162,23 @@ def _get_stats(x): return (x[0], get_stream_if_stats(x[1], x[0])) -def get_all_stream_if_stats(vcd_file, stream_ifs=None, sort_by="{'V': 1, 'R': 0}"): +def get_all_stream_if_stats(vcd_file, stream_ifs=None, sort_by="{'V': 1, 'R': 0}", num_workers=None): """Return a list of streaming interface stats, sorted by the percentage - for the given sort_by key. If stream_ifs is None, all streamin interface + for the given sort_by key. If stream_ifs is None, all streaming interface stats will be returned, otherwise treated as a list of interface names to - return the stats for.""" + return the stats for. + By default the number of parallel workers from the environment variable + NUM_DEFAULT_WORKERS will be used. This behavior can be changed on a per + call basis by supplying the optional parameter: num_workers + """ if stream_ifs is None: stream_ifs = list_stream_if(vcd_file) - with mp.Pool(get_num_default_workers()) as p: + if num_workers is None: + num_workers = get_num_default_workers() + + with mp.Pool(num_workers) as p: stream_ifs = map(lambda x: (x, vcd_file), stream_ifs) all_stats = p.map(_get_stats, stream_ifs)