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"""Unit tests for passing shell environment variables into LocalEnv scripts.""" 

6import sys 

7 

8import pytest 

9 

10from mlos_bench.tests.environments import check_env_success 

11from mlos_bench.tests.environments.local import create_local_env 

12from mlos_bench.tunables.tunable_groups import TunableGroups 

13 

14 

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

16 """Check that LocalEnv can set shell environment variables.""" 

17 local_env = create_local_env( 

18 tunable_groups, 

19 { 

20 "const_args": { 

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

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

23 }, 

24 "tunable_params": ["kernel"], 

25 "shell_env_params": [ 

26 "const_arg", # From "const_arg" 

27 "kernel_sched_latency_ns", # From "tunable_params" 

28 ], 

29 "run": [ 

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

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

32 ], 

33 "read_results_file": "output.csv", 

34 }, 

35 ) 

36 

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

38 

39 

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

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

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

43 _run_local_env( 

44 tunable_groups, 

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

46 expected={ 

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

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

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

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

51 }, 

52 ) 

53 

54 

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

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

57 """Check that LocalEnv can set shell environment variables on Windows / cmd 

58 shell. 

59 """ 

60 _run_local_env( 

61 tunable_groups, 

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

63 expected={ 

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

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

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

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

68 }, 

69 )