Coverage for mlos_bench/mlos_bench/config/environments/apps/fio/scripts/local/process_fio_results.py: 44%

1#!/usr/bin/env python3 

2# 

3# Copyright (c) Microsoft Corporation. 

4# Licensed under the MIT License. 

5# 

6"""Script for post-processing FIO results for mlos_bench.""" 

7 

8import argparse 

9import itertools 

10import json 

11from collections.abc import Iterator 

12from typing import Any 

13 

14import pandas 

15 

16 

17def _flat_dict(data: Any, path: str) -> Iterator[tuple[str, Any]]: 

18 """Flatten every dict in the hierarchy and rename the keys with the dict path.""" 

19 if isinstance(data, dict): 

20 for key, val in data.items(): 

21 yield from _flat_dict(val, f"{path}.{key}") 

22 else: 

23 yield (path, data) 

24 

25 

26def _main(input_file: str, output_file: str, prefix: str) -> None: 

27 """Convert FIO read data from JSON to tall CSV.""" 

28 with open(input_file, encoding="utf-8") as fh_input: 

29 json_data = json.load(fh_input) 

30 

31 data = list( 

32 itertools.chain( 

33 _flat_dict(json_data["jobs"][0], prefix), 

34 _flat_dict(json_data["disk_util"][0], f"{prefix}.disk_util"), 

35 ) 

36 ) 

37 

38 tall_df = pandas.DataFrame(data, columns=["metric", "value"]) 

39 tall_df.to_csv(output_file, index=False) 

40 print(f"Converted: {input_file} -> {output_file}") 

41 # print(tall_df) 

42 

43 

44if __name__ == "__main__": 

45 

46 parser = argparse.ArgumentParser(description="Post-process FIO benchmark results.") 

47 

48 parser.add_argument( 

49 "input", 

50 help="FIO benchmark results in JSON format (downloaded from a remote VM).", 

51 ) 

52 parser.add_argument( 

53 "output", 

54 help="Converted FIO benchmark data (CSV, to be consumed by mlos_bench).", 

55 ) 

56 parser.add_argument( 

57 "--prefix", 

58 default="fio", 

59 help="Prefix of the metric IDs (default 'fio')", 

60 ) 

61 

62 args = parser.parse_args() 

63 _main(args.input, args.output, args.prefix)