Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / config / parser.py: 88%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""Training configuration entry point: a YAML file plus dotted CLI overrides.""" 

16 

17from __future__ import annotations 

18 

19__all__ = [ 

20 "parse_training_args", 

21] 

22 

23import argparse 

24from pathlib import Path 

25from typing import Sequence 

26 

27import yaml 

28 

29from hyper_parallel.trainer.config.resolver import ( 

30 ConfigResolutionError, 

31 replace_override_path, 

32 resolve_config, 

33) 

34from hyper_parallel.trainer.config.trainer import TrainerConfig 

35 

36 

37def _parse_override_tokens(tokens: Sequence[str]) -> list[tuple[str, str]]: 

38 """Split ``--field=value`` CLI tokens into validated paths and raw values.""" 

39 parsed = [] 

40 for token in tokens: 

41 if not token.startswith("--") or token == "--": 

42 raise ConfigResolutionError( 

43 "CLI", f"expected a dotted override in '--field=value' form, got {token!r}" 

44 ) 

45 option = token[2:] 

46 if "=" not in option: 

47 raise ConfigResolutionError( 

48 "CLI", f"expected a dotted override in '--field=value' form, got {token!r}" 

49 ) 

50 path, raw_value = option.split("=", 1) 

51 if not path or any(not part for part in path.split(".")): 

52 raise ConfigResolutionError("CLI", f"invalid override path {path!r}") 

53 parsed.append((path, raw_value)) 

54 return parsed 

55 

56 

57def _apply_typed_overrides( 

58 config: TrainerConfig, 

59 overrides: Sequence[str], 

60) -> TrainerConfig: 

61 """Parse CLI override values and apply them to the resolved configuration.""" 

62 config_with_overrides = config 

63 for dot_path, raw_value in _parse_override_tokens(overrides): 

64 try: 

65 parsed_value = yaml.safe_load(raw_value) 

66 except yaml.YAMLError as exc: 

67 raise ConfigResolutionError( 

68 f"CLI.{dot_path}", f"invalid YAML value {raw_value!r}: {exc}" 

69 ) from exc 

70 config_with_overrides = replace_override_path( 

71 config_with_overrides, dot_path.split("."), parsed_value, path="" 

72 ) 

73 return config_with_overrides 

74 

75 

76def _load_training_config( 

77 config_file: str | Path, 

78 cli_overrides: Sequence[str] = (), 

79) -> TrainerConfig: 

80 """Load a YAML configuration and apply CLI overrides.""" 

81 config_path = Path(config_file) 

82 try: 

83 raw = yaml.safe_load(config_path.read_text(encoding="utf-8")) 

84 except OSError as exc: 

85 raise ConfigResolutionError(config_path, f"could not read config file: {exc}") from exc 

86 except yaml.YAMLError as exc: 

87 raise ConfigResolutionError(config_path, f"invalid YAML: {exc}") from exc 

88 

89 config = resolve_config(raw) 

90 return _apply_typed_overrides(config, cli_overrides) 

91 

92 

93def parse_training_args(argv: Sequence[str] | None = None) -> TrainerConfig: 

94 """Parse a YAML configuration path and CLI overrides into a ``TrainerConfig``. 

95 

96 Args: 

97 argv: Argument tokens excluding the program name. Uses ``sys.argv[1:]`` 

98 when ``None``. The configuration path precedes any 

99 ``--field=value`` overrides. 

100 

101 Returns: 

102 The resolved configuration with CLI overrides applied. 

103 

104 Raises: 

105 ConfigResolutionError: The file cannot be read, or the YAML 

106 configuration or CLI overrides are invalid. 

107 SystemExit: Help is requested or required command-line arguments 

108 are missing. 

109 

110 Example: 

111 config = parse_training_args(["train.yaml", "--accelerator.tp_size=4"]) 

112 """ 

113 parser = argparse.ArgumentParser(description="HyperParallel training config") 

114 parser.add_argument("config_file", help="Path to the YAML training config") 

115 args, overrides = parser.parse_known_args(argv) 

116 return load_training_config(args.config_file, overrides) 

117 

118 

119def load_training_config( 

120 config_file: str | Path, 

121 cli_overrides: Sequence[str] = (), 

122) -> TrainerConfig: 

123 """Load a training YAML into ``TrainerConfig``. 

124 

125 Args: 

126 config_file: Path to the YAML training config. 

127 cli_overrides: Canonical ``--field=value`` dotted overrides. 

128 

129 Returns: 

130 Resolved training configuration, carrying the resolved YAML path for 

131 callers that need the source of the configuration (codegen artifact 

132 layout, signature). 

133 """ 

134 config = _load_training_config(config_file, cli_overrides) 

135 config._yaml_path = str(Path(config_file).resolve()) # pylint: disable=protected-access 

136 return config