Coverage for mlos_bench/mlos_bench/services/types/local_exec_type.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Protocol interface for Service types that provide helper functions to run scripts and 

6commands locally on the scheduler side. 

7""" 

8 

9import contextlib 

10import tempfile 

11from typing import ( 

12 Iterable, 

13 Mapping, 

14 Optional, 

15 Protocol, 

16 Tuple, 

17 Union, 

18 runtime_checkable, 

19) 

20 

21from mlos_bench.tunables.tunable import TunableValue 

22 

23 

24@runtime_checkable 

25class SupportsLocalExec(Protocol): 

26 """ 

27 Protocol interface for a collection of methods to run scripts and commands in an 

28 external process on the node acting as the scheduler. 

29 

30 Can be useful for data processing due to reduced dependency management complications 

31 vs the target environment. Used in LocalEnv and provided by LocalExecService. 

32 """ 

33 

34 def local_exec( 

35 self, 

36 script_lines: Iterable[str], 

37 env: Optional[Mapping[str, TunableValue]] = None, 

38 cwd: Optional[str] = None, 

39 ) -> Tuple[int, str, str]: 

40 """ 

41 Execute the script lines from `script_lines` in a local process. 

42 

43 Parameters 

44 ---------- 

45 script_lines : Iterable[str] 

46 Lines of the script to run locally. 

47 Treat every line as a separate command to run. 

48 env : Mapping[str, Union[int, float, str]] 

49 Environment variables (optional). 

50 cwd : str 

51 Work directory to run the script at. 

52 If omitted, use `temp_dir` or create a temporary dir. 

53 

54 Returns 

55 ------- 

56 (return_code, stdout, stderr) : (int, str, str) 

57 A 3-tuple of return code, stdout, and stderr of the script process. 

58 """ 

59 

60 def temp_dir_context( 

61 self, 

62 path: Optional[str] = None, 

63 ) -> Union[tempfile.TemporaryDirectory, contextlib.nullcontext]: 

64 """ 

65 Create a temp directory or use the provided path. 

66 

67 Parameters 

68 ---------- 

69 path : str 

70 A path to the temporary directory. Create a new one if None. 

71 

72 Returns 

73 ------- 

74 temp_dir_context : TemporaryDirectory 

75 Temporary directory context to use in the `with` clause. 

76 """