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 collections.abc import Callable 

11from contextlib import nullcontext 

12from string import Template 

13from tempfile import TemporaryDirectory 

14from typing import Any 

15 

16from mlos_bench.services.base_service import Service 

17 

18_LOG = logging.getLogger(__name__) 

19 

20 

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

22 """ 

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

24 context for local scripts. 

25 

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

27 supposed to be used as a standalone service. 

28 """ 

29 

30 def __init__( 

31 self, 

32 config: dict[str, Any] | None = None, 

33 global_config: dict[str, Any] | None = None, 

34 parent: Service | None = None, 

35 methods: dict[str, Callable] | list[Callable] | None = None, 

36 ): 

37 """ 

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

39 local exec service. 

40 

41 Parameters 

42 ---------- 

43 config : dict 

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

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

46 global_config : dict 

47 Free-format dictionary of global parameters. 

48 parent : Service 

49 An optional parent service that can provide mixin functions. 

50 methods : Union[dict[str, Callable], list[Callable], None] 

51 New methods to register with the service. 

52 """ 

53 super().__init__( 

54 config, 

55 global_config, 

56 parent, 

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

58 ) 

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

60 if self._temp_dir: 

61 # expand globals 

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

63 # and resolve the path to absolute path 

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

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

66 

67 def temp_dir_context( 

68 self, 

69 path: str | None = None, 

70 ) -> TemporaryDirectory | nullcontext: 

71 """ 

72 Create a temp directory or use the provided path. 

73 

74 Parameters 

75 ---------- 

76 path : str 

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

78 

79 Returns 

80 ------- 

81 temp_dir_context : tempfile.TemporaryDirectory 

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

83 """ 

84 temp_dir = path or self._temp_dir 

85 if temp_dir is None: 

86 return TemporaryDirectory() 

87 os.makedirs(temp_dir, exist_ok=True) 

88 return nullcontext(temp_dir)