Coverage for mlos_bench/mlos_bench/tests/config/environments/local/scripts/bench_setup.py: 25%

1#!/usr/bin/env python3 

2# 

3# Copyright (c) Microsoft Corporation. 

4# Licensed under the MIT License. 

5# 

6""" 

7Helper script to update the environment parameters from JSON. 

8 

9This is a sample script that demonstrates how to read the tunable parameters 

10and metadata from JSON and produce some kind of a configuration file for the 

11application that is being benchmarked or optimized. 

12 

13THIS IS A TOY EXAMPLE. The script does not have any actual effect on the system. 

14Please copy and extend it to suit your needs. 

15 

16Run: 

17 `./bench_setup.py ./input-params.json ./input-params-meta.json` 

18""" 

19 

20import argparse 

21import json 

22import os 

23 

24 

25def _main(fname_input: str, fname_meta: str, fname_output: str) -> None: 

26 

27 # In addition to the input JSON files, 

28 # MLOS can pass parameters through the OS environment: 

29 print(f'# RUN: {os.environ["experiment_id"]}:{os.environ["trial_id"]}') 

30 

31 # Key-value pairs of tunable parameters, e.g., 

32 # {"shared_buffers": "128", ...} 

33 with open(fname_input, "rt", encoding="utf-8") as fh_tunables: 

34 tunables_data = json.load(fh_tunables) 

35 

36 # Optional free-format metadata for tunable parameters, e.g. 

37 # {"shared_buffers": {"suffix": "MB"}, ...} 

38 with open(fname_meta, "rt", encoding="utf-8") as fh_meta: 

39 tunables_meta = json.load(fh_meta) 

40 

41 # Pretend that we are generating a PG config file with lines like: 

42 # shared_buffers = 128MB 

43 with open(fname_output, "wt", encoding="utf-8", newline="") as fh_config: 

44 for key, val in tunables_data.items(): 

45 meta = tunables_meta.get(key, {}) 

46 suffix = meta.get("suffix", "") 

47 line = f"{key} = {val}{suffix}" 

48 fh_config.write(line + "\n") 

49 print(line) 

50 

51 

52if __name__ == "__main__": 

53 

54 parser = argparse.ArgumentParser(description="Update the environment parameters from JSON.") 

55 

56 parser.add_argument("input", help="JSON file with tunable parameters.") 

57 parser.add_argument("meta", help="JSON file with tunable parameters metadata.") 

58 parser.add_argument("output", help="Output config file.") 

59 

60 args = parser.parse_args() 

61 

62 _main(args.input, args.meta, args.output)