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 name: str 

17 logger: Any 

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

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

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

21 

22 

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

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

25 if not path.exists(): 

26 return {} 

27 with open(path) as f: 

28 return yaml.safe_load(f) or {} 

29 

30 

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

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

33 

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

35 subdirectory, otherwise its direct parent. 

36 """ 

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

38 

39 

40def _merge_sections( 

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

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

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

44 

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

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

47 

48 Returns: 

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

50 """ 

51 

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

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

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

55 

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

57 

58 

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

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

61 

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

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

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

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

66 

67 Raises: 

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

69 """ 

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

71 if env_folder: 

72 output_dir = Path(env_folder) 

73 else: 

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

75 base_resolved = base.resolve() 

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

77 if not output_dir.is_relative_to(base_resolved): 

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

79 raise ValueError(msg) 

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

81 return output_dir 

82 

83 

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

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

86 key = str(log_path.resolve()) 

87 if key not in _FILE_SINKS: 

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

89 

90 

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

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

93 

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

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

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

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

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

99 under the notebooks directory. 

100 """ 

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

102 cfg = _load_yaml(config_path) 

103 base = _resolve_base(config_path) 

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

105 

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

107 

108 output_dir = _output_dir(base, data, name) 

109 _ensure_sink(output_dir / "output.log") 

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

111 

112 return ExperimentConfig( 

113 name=name, 

114 logger=logger, 

115 params=params, 

116 optuna=optuna_cfg, 

117 data=data, 

118 )