Coverage for mlos_bench/mlos_bench/tests/environments/local/local_env_vars_test.py: 93%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Unit tests for passing shell environment variables into LocalEnv scripts. 

7""" 

8import sys 

9 

10import pytest 

11 

12from mlos_bench.tunables.tunable_groups import TunableGroups 

13from mlos_bench.tests.environments import check_env_success 

14from mlos_bench.tests.environments.local import create_local_env 

15 

16 

17def _run_local_env(tunable_groups: TunableGroups, shell_subcmd: str, expected: dict) -> None: 

18 """ 

19 Check that LocalEnv can set shell environment variables. 

20 """ 

21 local_env = create_local_env(tunable_groups, { 

22 "const_args": { 

23 "const_arg": 111, # Passed into "shell_env_params" 

24 "other_arg": 222, # NOT passed into "shell_env_params" 

25 }, 

26 "tunable_params": ["kernel"], 

27 "shell_env_params": [ 

28 "const_arg", # From "const_arg" 

29 "kernel_sched_latency_ns", # From "tunable_params" 

30 ], 

31 "run": [ 

32 "echo const_arg,other_arg,unknown_arg,kernel_sched_latency_ns > output.csv", 

33 f"echo {shell_subcmd} >> output.csv", 

34 ], 

35 "read_results_file": "output.csv", 

36 }) 

37 

38 check_env_success(local_env, tunable_groups, expected, []) 

39 

40 

41@pytest.mark.skipif(sys.platform == 'win32', reason="sh-like shell only") 

42def test_local_env_vars_shell(tunable_groups: TunableGroups) -> None: 

43 """ 

44 Check that LocalEnv can set shell environment variables in sh-like shell. 

45 """ 

46 _run_local_env( 

47 tunable_groups, 

48 shell_subcmd="$const_arg,$other_arg,$unknown_arg,$kernel_sched_latency_ns", 

49 expected={ 

50 "const_arg": 111, # From "const_args" 

51 "other_arg": float("NaN"), # Not included in "shell_env_params" 

52 "unknown_arg": float("NaN"), # Unknown/undefined variable 

53 "kernel_sched_latency_ns": 2000000, # From "tunable_params" 

54 } 

55 ) 

56 

57 

58@pytest.mark.skipif(sys.platform != 'win32', reason="Windows only") 

59def test_local_env_vars_windows(tunable_groups: TunableGroups) -> None: 

60 """ 

61 Check that LocalEnv can set shell environment variables on Windows / cmd shell. 

62 """ 

63 _run_local_env( 

64 tunable_groups, 

65 shell_subcmd=r"%const_arg%,%other_arg%,%unknown_arg%,%kernel_sched_latency_ns%", 

66 expected={ 

67 "const_arg": 111, # From "const_args" 

68 "other_arg": r"%other_arg%", # Not included in "shell_env_params" 

69 "unknown_arg": r"%unknown_arg%", # Unknown/undefined variable 

70 "kernel_sched_latency_ns": 2000000, # From "tunable_params" 

71 } 

72 )