Coverage for mlos_bench/mlos_bench/services/local/temp_dir_context.py: 100%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""Helper functions to work with temp files locally on the scheduler side.""" 

6 

7import abc 

8import logging 

9import os 

10from contextlib import nullcontext 

11from string import Template 

12from tempfile import TemporaryDirectory 

13from typing import Any, Callable, Dict, List, Optional, Union 

14 

15from mlos_bench.services.base_service import Service 

16 

17_LOG = logging.getLogger(__name__) 

18 

19 

20class TempDirContextService(Service, metaclass=abc.ABCMeta): 

21 """ 

22 A *base* service class that provides a method to create a temporary directory 

23 context for local scripts. 

24 

25 It is inherited by LocalExecService and MockLocalExecService. This class is not 

26 supposed to be used as a standalone service. 

27 """ 

28 

29 def __init__( 

30 self, 

31 config: Optional[Dict[str, Any]] = None, 

32 global_config: Optional[Dict[str, Any]] = None, 

33 parent: Optional[Service] = None, 

34 methods: Union[Dict[str, Callable], List[Callable], None] = None, 

35 ): 

36 """ 

37 Create a new instance of a service that provides temporary directory context for 

38 local exec service. 

39 

40 Parameters 

41 ---------- 

42 config : dict 

43 Free-format dictionary that contains parameters for the service. 

44 (E.g., root path for config files, etc.) 

45 global_config : dict 

46 Free-format dictionary of global parameters. 

47 parent : Service 

48 An optional parent service that can provide mixin functions. 

49 methods : Union[Dict[str, Callable], List[Callable], None] 

50 New methods to register with the service. 

51 """ 

52 super().__init__( 

53 config, 

54 global_config, 

55 parent, 

56 self.merge_methods(methods, [self.temp_dir_context]), 

57 ) 

58 self._temp_dir = self.config.get("temp_dir") 

59 if self._temp_dir: 

60 # expand globals 

61 self._temp_dir = Template(self._temp_dir).safe_substitute(global_config or {}) 

62 # and resolve the path to absolute path 

63 self._temp_dir = self._config_loader_service.resolve_path(self._temp_dir) 

64 _LOG.info("%s: temp dir: %s", self, self._temp_dir) 

65 

66 def temp_dir_context( 

67 self, 

68 path: Optional[str] = None, 

69 ) -> Union[TemporaryDirectory, nullcontext]: 

70 """ 

71 Create a temp directory or use the provided path. 

72 

73 Parameters 

74 ---------- 

75 path : str 

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

77 

78 Returns 

79 ------- 

80 temp_dir_context : tempfile.TemporaryDirectory 

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

82 """ 

83 temp_dir = path or self._temp_dir 

84 if temp_dir is None: 

85 return TemporaryDirectory() 

86 os.makedirs(temp_dir, exist_ok=True) 

87 return nullcontext(temp_dir)