Coverage for src/tinycta/hyper/_setup.py: 100%

1"""Experiment setup helpers: logger configuration.""" 

2 

3import os 

4from pathlib import Path 

5from typing import Any, NamedTuple 

6 

7import yaml 

8from loguru import logger 

9 

10_FILE_SINKS: dict[str, int] = {} 

11 

12 

13class ExperimentConfig(NamedTuple): 

14 """Resources bundled for a notebook experiment run. 

15 

16 Example: 

17 >>> from tinycta.hyper import ExperimentConfig 

18 >>> cfg = ExperimentConfig(name="momentum", logger=None, params={"fast": 16, "slow": 64}) 

19 >>> cfg.name 

20 'momentum' 

21 >>> cfg.params 

22 {'fast': 16, 'slow': 64} 

23 

24 The three config sections are optional and default to ``None``, so a 

25 config file that omits one is not an error: 

26 

27 >>> cfg.optuna is None and cfg.data is None 

28 True 

29 

30 Being a :class:`~typing.NamedTuple`, it also unpacks positionally: 

31 

32 >>> name, _logger, params, _optuna, _data = cfg 

33 >>> name, sorted(params) 

34 ('momentum', ['fast', 'slow']) 

35 """ 

36 

37 name: str 

38 logger: Any 

39 params: dict[str, Any] | None = None 

40 optuna: dict[str, Any] | None = None 

41 data: dict[str, Any] | None = None 

42 

43 

44def _load_yaml(path: Path) -> dict[str, Any]: 

45 """Load a YAML file, returning an empty dict if the file does not exist.""" 

46 if not path.exists(): 

47 return {} 

48 with open(path) as f: 

49 return yaml.safe_load(f) or {} 

50 

51 

52def _resolve_base(config_path: Path) -> Path: 

53 """Return the notebooks directory that paths are resolved relative to. 

54 

55 This is the grandparent of ``config_path`` when it lives inside a ``config/`` 

56 subdirectory, otherwise its direct parent. 

57 """ 

58 return config_path.parent.parent if config_path.parent.name == "config" else config_path.parent 

59 

60 

61def _merge_sections( 

62 cfg: dict[str, Any], sibling: dict[str, Any] 

63) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]: 

64 """Merge the shared ``config.yml`` with the experiment-specific sibling. 

65 

66 Each of the ``data``, ``params`` and ``optuna`` sections is taken from 

67 ``cfg`` when present and truthy, else from ``sibling``, else an empty dict. 

68 

69 Returns: 

70 tuple: ``(data, params, optuna)`` section dicts. 

71 """ 

72 

73 def pick(section: str) -> dict[str, Any]: 

74 """Return ``section`` from ``cfg`` if truthy, else ``sibling``, else an empty dict.""" 

75 return cfg.get(section) or sibling.get(section) or {} 

76 

77 return pick("data"), pick("params"), pick("optuna") 

78 

79 

80def _output_dir(base: Path, data: dict[str, Any], name: str) -> Path: 

81 """Resolve the experiment output directory and create it. 

82 

83 ``NOTEBOOK_OUTPUT_FOLDER`` is a deliberate operator override and is used 

84 verbatim. Otherwise the directory is ``base / output_path / name`` and is 

85 confined under ``base`` so an untrusted ``output_path`` (e.g. ``../../etc`` 

86 or an absolute path) cannot escape the notebooks directory. 

87 

88 Raises: 

89 ValueError: When the config-derived path escapes ``base``. 

90 """ 

91 env_folder = os.environ.get("NOTEBOOK_OUTPUT_FOLDER") 

92 if env_folder: 

93 output_dir = Path(env_folder) 

94 else: 

95 folder = data.get("output_path", "output") 

96 base_resolved = base.resolve() 

97 output_dir = (base_resolved / folder / name).resolve() 

98 if not output_dir.is_relative_to(base_resolved): 

99 msg = f"output_path {folder!r} escapes the notebooks directory {base_resolved}" 

100 raise ValueError(msg) 

101 output_dir.mkdir(parents=True, exist_ok=True) 

102 return output_dir 

103 

104 

105def _ensure_sink(log_path: Path) -> None: 

106 """Register a loguru file sink for ``log_path`` once, keyed by its resolved path.""" 

107 key = str(log_path.resolve()) 

108 if key not in _FILE_SINKS: 

109 _FILE_SINKS[key] = logger.add(log_path) 

110 

111 

112def get_config(name: str, config_path: Path | str | None = None) -> ExperimentConfig: 

113 """Return logger and config sections for an experiment. 

114 

115 Accepts either a shared ``config.yml`` or an experiment-specific 

116 ``config/{name}.yml``. Paths in the config are resolved relative to the 

117 notebooks directory (one level above any ``config/`` subdirectory). 

118 ``NOTEBOOK_OUTPUT_FOLDER`` env var overrides the output directory used for 

119 the log file sink; otherwise the config-derived output directory is confined 

120 under the notebooks directory. 

121 

122 Example: 

123 >>> import tempfile 

124 >>> from pathlib import Path 

125 >>> import yaml 

126 >>> from tinycta.hyper import get_config 

127 

128 A shared ``config.yml`` supplies the sections directly: 

129 

130 >>> with tempfile.TemporaryDirectory() as tmp: 

131 ... config_path = Path(tmp) / "config.yml" 

132 ... _ = config_path.write_text( 

133 ... yaml.safe_dump({"params": {"fast": 16, "slow": 64}, "data": {"output_path": "out"}}) 

134 ... ) 

135 ... cfg = get_config("momentum", config_path=config_path) 

136 ... output_exists = (Path(tmp) / "out" / "momentum").is_dir() 

137 >>> cfg.name 

138 'momentum' 

139 >>> cfg.params 

140 {'fast': 16, 'slow': 64} 

141 

142 The output directory is created as a side effect, ready for the run's 

143 artefacts and its ``output.log`` sink: 

144 

145 >>> output_exists 

146 True 

147 

148 An ``output_path`` that would escape the notebooks directory is rejected 

149 rather than followed: 

150 

151 >>> with tempfile.TemporaryDirectory() as tmp: 

152 ... config_path = Path(tmp) / "config.yml" 

153 ... _ = config_path.write_text(yaml.safe_dump({"data": {"output_path": "../../etc"}})) 

154 ... try: 

155 ... get_config("momentum", config_path=config_path) 

156 ... except ValueError: 

157 ... print("escaping output_path rejected") 

158 escaping output_path rejected 

159 """ 

160 config_path = Path(config_path) if config_path else Path.cwd() / "config.yml" 

161 cfg = _load_yaml(config_path) 

162 base = _resolve_base(config_path) 

163 sibling = _load_yaml(base / "config" / f"{name}.yml") 

164 

165 data, params, optuna_cfg = _merge_sections(cfg, sibling) 

166 

167 output_dir = _output_dir(base, data, name) 

168 _ensure_sink(output_dir / "output.log") 

169 logger.info(f"Writing output to: {output_dir}\nCurrent working directory: {os.getcwd()}") 

170 

171 return ExperimentConfig( 

172 name=name, 

173 logger=logger, 

174 params=params, 

175 optuna=optuna_cfg, 

176 data=data, 

177 )