Coverage for mlos_bench/mlos_bench/services/local/temp_dir_context.py: 100%
24 statements
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-06 00:35 +0000
« prev ^ index » next coverage.py v7.5.1, created at 2024-05-06 00:35 +0000
1#
2# Copyright (c) Microsoft Corporation.
3# Licensed under the MIT License.
4#
5"""
6Helper functions to work with temp files locally on the scheduler side.
7"""
9import abc
10import logging
11import os
12from contextlib import nullcontext
13from string import Template
14from tempfile import TemporaryDirectory
15from typing import Any, Callable, Dict, List, Optional, Union
17from mlos_bench.services.base_service import Service
19_LOG = logging.getLogger(__name__)
22class TempDirContextService(Service, metaclass=abc.ABCMeta):
23 """
24 A *base* service class that provides a method to create a temporary
25 directory context for local scripts.
27 It is inherited by LocalExecService and MockLocalExecService.
28 This class is not supposed to be used as a standalone service.
29 """
31 def __init__(self,
32 config: Optional[Dict[str, Any]] = None,
33 global_config: Optional[Dict[str, Any]] = None,
34 parent: Optional[Service] = None,
35 methods: Union[Dict[str, Callable], List[Callable], None] = None):
36 """
37 Create a new instance of a service that provides temporary directory context
38 for local exec service.
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, global_config, parent,
54 self.merge_methods(methods, [self.temp_dir_context])
55 )
56 self._temp_dir = self.config.get("temp_dir")
57 if self._temp_dir:
58 # expand globals
59 self._temp_dir = Template(self._temp_dir).safe_substitute(global_config or {})
60 # and resolve the path to absolute path
61 self._temp_dir = self._config_loader_service.resolve_path(self._temp_dir)
62 _LOG.info("%s: temp dir: %s", self, self._temp_dir)
64 def temp_dir_context(self, path: Optional[str] = None) -> Union[TemporaryDirectory, nullcontext]:
65 """
66 Create a temp directory or use the provided path.
68 Parameters
69 ----------
70 path : str
71 A path to the temporary directory. Create a new one if None.
73 Returns
74 -------
75 temp_dir_context : TemporaryDirectory
76 Temporary directory context to use in the `with` clause.
77 """
78 temp_dir = path or self._temp_dir
79 if temp_dir is None:
80 return TemporaryDirectory()
81 os.makedirs(temp_dir, exist_ok=True)
82 return nullcontext(temp_dir)