[llm bench] Move calculation of memory consumption to memory_monitor tool

Author: sbalandi

tools/llm_bench/benchmark.py

@@ -9,7 +9,7 @@
 from openvino import get_version
 import torch
 import traceback
-from llm_bench_utils.memory_profile import MemConsumption
+from llm_bench_utils.memory_profile import MemMonitorWrapper
 import llm_bench_utils.output_csv
 import llm_bench_utils.output_json
 import task.visual_language_generation as bench_vlm
@@ -19,7 +19,7 @@
 import task.speech_to_text_generation as bench_speech
 
 DEFAULT_TORCH_THREAD_NUMS = 16
-mem_consumption = MemConsumption()
+memory_monitor = MemMonitorWrapper()
 
 
 def num_iters_type(x):
@@ -85,6 +85,14 @@ def get_argprser():
         help='if the value is 1, output the maximum memory consumption in warm-up iterations. If the value is 2,'
         ' output the maximum memory consumption in all iterations.',
     )
+    parser.add_argument(
+        '-mc_dir',
+        '--memory_consumption_dir',
+        default=None,
+        required=False,
+        type=str,
+        help='Path to store memory consamption logs and chart.',
+    )
     parser.add_argument('-bs', '--batch_size', type=int, default=1, required=False, help='Batch size value')
     parser.add_argument('--num_beams', type=int, default=1, help='Number of beams in the decoding strategy, activates beam_search if greater than 1')
     parser.add_argument(
@@ -221,21 +229,23 @@ def main():
                     if half_nums_of_torch_threads > DEFAULT_TORCH_THREAD_NUMS:
                         torch.set_num_threads(DEFAULT_TORCH_THREAD_NUMS)
                     else:
+                        half_nums_of_torch_threads = int(half_nums_of_torch_threads) if int(half_nums_of_torch_threads) else 1
                         torch.set_num_threads(int(half_nums_of_torch_threads))
             log.info(f"The num_beams is {model_args['num_beams']}, update Torch thread num from "
                      f'{original_torch_thread_nums} to {torch.get_num_threads()}, avoid to use the CPU cores for OpenVINO inference.')
     log.info(out_str)
     if args.memory_consumption:
-        mem_consumption.start_collect_mem_consumption_thread()
+        memory_monitor.create_monitors()
+        if args.memory_consumption_dir:
+            memory_monitor.set_dir(args.memory_consumption_dir)
     try:
         if model_args['use_case'] in ['text_gen', 'code_gen']:
             iter_data_list, pretrain_time, iter_timestamp = CASE_TO_BENCH[model_args['use_case']](
                 model_path, framework, args.device, args.tokens_len, args.streaming, model_args,
-                args.num_iters, mem_consumption)
+                args.num_iters, memory_monitor)
         else:
             iter_data_list, pretrain_time, iter_timestamp = CASE_TO_BENCH[model_args['use_case']](
-                model_path, framework, args.device, model_args, args.num_iters,
-                mem_consumption)
+                model_path, framework, args.device, model_args, args.num_iters, memory_monitor)
         if args.report is not None or args.report_json is not None:
             model_precision = ''
             if framework == 'ov':
@@ -276,7 +286,7 @@ def main():
         exit(1)
     finally:
         if args.memory_consumption:
-            mem_consumption.end_collect_mem_consumption_thread()
+            memory_monitor.stop()
 
 
 if __name__ == '__main__':

tools/llm_bench/llm_bench_utils/gen_output_data.py

@@ -12,8 +12,7 @@ def gen_iterate_data(
     latency='',
     res_md5='',
     max_rss_mem='',
-    max_shared_mem='',
-    max_uss_mem='',
+    max_sys_mem='',
     prompt_idx='',
     tokenization_time=[],
     mm_embeddings_preparation_time=''
@@ -31,8 +30,7 @@ def gen_iterate_data(
     iter_data['first_token_infer_latency'] = -1
     iter_data['other_tokens_infer_avg_latency'] = -1
     iter_data['max_rss_mem_consumption'] = max_rss_mem
-    iter_data['max_shared_mem_consumption'] = max_shared_mem
-    iter_data['max_uss_mem_consumption'] = max_uss_mem
+    iter_data['max_sys_mem_consumption'] = max_sys_mem
     iter_data['prompt_idx'] = prompt_idx
     iter_data['tokenization_time'] = tokenization_time[0] if len(tokenization_time) > 0 else ''
     iter_data['detokenization_time'] = tokenization_time[1] if len(tokenization_time) > 1 else ''