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"""YAML configuration parser for PP optimization."""
16
17from __future__ import annotations
18
19import math
20from dataclasses import dataclass
21from typing import Any, Optional
22
23import yaml
24
25
26def _validate_int(raw: Any, field_name: str) -> int:
27 """Convert *raw* to int, rejecting floats and bools.
28
29 ``int(1.9)`` silently truncates to ``1``. Even ``int(3.0)`` accepts
30 a float where an integer was intended. ``int(True)`` yields ``1``,
31 which is almost never the intent for a pipeline-parallel integer
32 config field. This helper raises ``ValueError`` when the input is
33 a float or bool.
34
35 Args:
36 raw: Value from parsed YAML.
37 field_name: Human-readable field name for error messages.
38
39 Returns:
40 Integer value.
41
42 Raises:
43 ValueError: If *raw* is a float or bool.
44 """
45 if isinstance(raw, bool):
46 raise ValueError(
47 f"{field_name} must be an integer, got bool {raw}"
48 )
49 if isinstance(raw, float):
50 raise ValueError(
51 f"{field_name} must be an integer, got {raw}"
52 )
53 return int(raw)
54
55
56def _to_bool(value: Any, field_name: str = "field") -> bool:
57 """Convert a YAML-parsed value to bool with strict string handling.
58
59 ``bool("false")`` is ``True`` in Python because any non-empty string
60 is truthy. This helper uses a whitelist of recognized boolean
61 representations and raises ``ValueError`` for ambiguous inputs.
62
63 Accepted values:
64
65 * ``bool`` — ``True`` / ``False`` directly.
66 * ``int`` — only ``0`` (False) or ``1`` (True); other integers
67 are rejected.
68 * ``float`` — only ``0.0`` (False) or ``1.0`` (True); other
69 floats are rejected.
70 * ``str`` — ``"true"``/``"yes"``/``"1"`` → True,
71 ``"false"``/``"no"``/``"0"`` → False; anything else is
72 rejected.
73
74 Args:
75 value: Value from ``yaml.safe_load`` — typically ``bool``,
76 ``int``, ``float``, or ``str``.
77 field_name: Human-readable field name for error messages.
78
79 Returns:
80 Boolean interpretation of *value*.
81
82 Raises:
83 ValueError: If *value* is not a recognized boolean
84 representation (e.g. int other than 0/1, ambiguous
85 string, etc.).
86 """
87 if isinstance(value, bool):
88 return value
89 if isinstance(value, int) and not isinstance(value, bool):
90 if value == 0:
91 return False
92 if value == 1:
93 return True
94 raise ValueError(
95 f"{field_name}: cannot interpret {value} as boolean; "
96 f"expected 0 or 1"
97 )
98 if isinstance(value, float) and value.is_integer():
99 int_val = int(value)
100 if int_val == 0:
101 return False
102 if int_val == 1:
103 return True
104 raise ValueError(
105 f"{field_name}: cannot interpret {value} as boolean; "
106 f"expected 0.0 or 1.0"
107 )
108 if isinstance(value, str):
109 low = value.strip().lower()
110 if low in ("true", "yes", "1"):
111 return True
112 if low in ("false", "no", "0"):
113 return False
114 raise ValueError(
115 f"{field_name}: cannot interpret '{value}' as boolean; "
116 f"expected true/false/yes/no/1/0"
117 )
118 raise ValueError(
119 f"{field_name}: cannot interpret {value!r} as boolean"
120 )
121
122
123@dataclass
124class YamlOptimizationConfig:
125 """Parsed YAML configuration for PP optimization.
126
127 Attributes:
128 pp_degree: Number of pipeline stages. Parsed from the YAML key
129 ``pipeline_num``; the field name uses ILP domain terminology.
130 micro_batch_num: Number of micro batches.
131 num_layer: Number of homogeneous body layers. Optional; when
132 ``None``, actual layer counts come from the JSON profile
133 via :func:`generate_layers_list`.
134 num_of_interleave: VPP interleaving factor.
135 vpp_less_memory: Use the less-memory VPP schedule (``vpp2``).
136 optimization_level: ILP optimization level (0-2).
137 memory_limit: Per-stage memory limit in MB. Must be positive
138 for ILP load balancing.
139 constant_memory: Constant memory per stage in MB.
140 enable_simulation: Whether to run the pipeline simulator after
141 ILP solving.
142 sim_comm_time: P2P communication time between adjacent stages
143 in ms. Used only by the pipeline simulator; does NOT
144 affect ILP optimization.
145 use_backward_time: Whether to use real backward times from the
146 JSON profile for simulation. When ``False`` (default), the
147 simulator derives backward time from forward time using
148 ``backward_ratio`` (original behaviour). When ``True``,
149 the simulator uses actual backward times from profiling.
150 """
151
152 pp_degree: int
153 micro_batch_num: int
154 num_layer: Optional[int] = None
155 num_of_interleave: int = 1
156 vpp_less_memory: bool = False
157 optimization_level: int = 1
158 memory_limit: int = 0
159 constant_memory: int = 0
160 enable_simulation: bool = True
161 sim_comm_time: float = 0.0
162 use_backward_time: bool = False
163
164 def _validate_field_types(self) -> None:
165 """Check that every config field has the expected Python type.
166
167 Integer fields reject ``bool`` and ``float``; boolean fields
168 must be genuine ``bool`` instances; ``sim_comm_time`` must be
169 a finite ``int`` or ``float`` (not ``bool``).
170
171 Raises:
172 ValueError: If any field has an unexpected type.
173 """
174 int_fields = (
175 "pp_degree", "micro_batch_num", "num_of_interleave",
176 "optimization_level", "memory_limit", "constant_memory",
177 )
178 for name in int_fields:
179 val = getattr(self, name)
180 if isinstance(val, bool) or not isinstance(val, int):
181 raise ValueError(
182 f"{name} must be an integer, "
183 f"got {type(val).__name__} {val!r}"
184 )
185 if self.num_layer is not None:
186 if (isinstance(self.num_layer, bool)
187 or not isinstance(self.num_layer, int)):
188 raise ValueError(
189 f"num_layer must be an integer or None, "
190 f"got {type(self.num_layer).__name__} "
191 f"{self.num_layer!r}"
192 )
193 for name in ("vpp_less_memory", "enable_simulation", "use_backward_time"):
194 if not isinstance(getattr(self, name), bool):
195 raise ValueError(
196 f"{name} must be a boolean, "
197 f"got {type(getattr(self, name)).__name__}"
198 )
199 if (isinstance(self.sim_comm_time, bool)
200 or not isinstance(self.sim_comm_time, (int, float))
201 or not math.isfinite(self.sim_comm_time)):
202 raise ValueError(
203 f"sim_comm_time must be a finite number, "
204 f"got {self.sim_comm_time}"
205 )
206
207 def validate(self) -> None:
208 """Validate all pipeline configuration fields.
209
210 Checks both types (via :meth:`_validate_field_types`) and
211 value ranges.
212
213 Raises:
214 ValueError: If any field has an invalid type or value.
215 """
216 self._validate_field_types()
217 if self.pp_degree <= 0:
218 raise ValueError(
219 f"pp_degree must be positive, got {self.pp_degree}"
220 )
221 if self.micro_batch_num <= 0:
222 raise ValueError(
223 f"micro_batch_num must be positive, "
224 f"got {self.micro_batch_num}"
225 )
226 if self.num_layer is not None and self.num_layer <= 0:
227 raise ValueError(
228 f"num_layer must be positive when provided, "
229 f"got {self.num_layer}"
230 )
231 if self.num_of_interleave <= 0:
232 raise ValueError(
233 f"num_of_interleave must be positive, "
234 f"got {self.num_of_interleave}"
235 )
236 if self.optimization_level not in (0, 1, 2):
237 raise ValueError(
238 f"optimization_level must be 0, 1, or 2, "
239 f"got {self.optimization_level}"
240 )
241 if self.memory_limit < 0:
242 raise ValueError(
243 f"memory_limit must be non-negative, "
244 f"got {self.memory_limit}"
245 )
246 if self.constant_memory < 0:
247 raise ValueError(
248 f"constant_memory must be non-negative, "
249 f"got {self.constant_memory}"
250 )
251 if self.sim_comm_time < 0.0:
252 raise ValueError(
253 f"sim_comm_time must be non-negative, "
254 f"got {self.sim_comm_time}"
255 )
256
257
258def _extract_required_fields(
259 pipeline_cfg: dict,
260 yaml_path: str,
261) -> tuple[int, Optional[int], int]:
262 """Extract and validate required pipeline topology fields.
263
264 Args:
265 pipeline_cfg: The ``pipeline_config`` mapping from the YAML file.
266 yaml_path: Path to the YAML file (for error messages).
267
268 Returns:
269 ``(pp_degree, num_layer, micro_batch_num)`` tuple.
270
271 Raises:
272 ValueError: If a required field is missing or invalid.
273 """
274 pipeline_num = pipeline_cfg.get("pipeline_num")
275 if pipeline_num is None:
276 raise ValueError(
277 f"{yaml_path}: pipeline_config.pipeline_num is required"
278 )
279 pp_degree = _validate_int(pipeline_num, "pipeline_config.pipeline_num")
280 if pp_degree <= 0:
281 raise ValueError(
282 f"{yaml_path}: pipeline_config.pipeline_num must be "
283 f"positive, got {pp_degree}"
284 )
285
286 num_layer_raw = pipeline_cfg.get("num_layer")
287 num_layer: Optional[int] = None
288 if num_layer_raw is not None:
289 num_layer = _validate_int(num_layer_raw, "pipeline_config.num_layer")
290 if num_layer <= 0:
291 raise ValueError(
292 f"pipeline_config.num_layer must be positive, got {num_layer}"
293 )
294
295 micro_batch_num = pipeline_cfg.get("micro_batch_num")
296 if micro_batch_num is None:
297 raise ValueError("YAML pipeline_config.micro_batch_num is required")
298 micro_batch_num = _validate_int(
299 micro_batch_num, "pipeline_config.micro_batch_num",
300 )
301 if micro_batch_num <= 0:
302 raise ValueError(
303 f"pipeline_config.micro_batch_num must be positive, "
304 f"got {micro_batch_num}"
305 )
306
307 return pp_degree, num_layer, micro_batch_num
308
309
310def _extract_optional_fields(pipeline_cfg: dict) -> tuple[int, bool, int, int, int, bool, float, bool]:
311 """Extract and type-convert optional pipeline configuration fields.
312
313 Range validation is delegated to :meth:`YamlOptimizationConfig.validate`
314 so that direct construction also benefits from the same checks.
315
316 Args:
317 pipeline_cfg: The ``pipeline_config`` mapping from the YAML file.
318
319 Returns:
320 ``(num_of_interleave, vpp_less_memory, optimization_level,
321 memory_limit, constant_memory, enable_simulation,
322 sim_comm_time, use_backward_time)`` tuple.
323
324 Raises:
325 ValueError: If a field cannot be converted to the expected type
326 (e.g. float where int is required, ambiguous boolean string).
327 """
328 num_of_interleave = _validate_int(
329 pipeline_cfg.get("num_of_interleave", 1),
330 "pipeline_config.num_of_interleave",
331 )
332
333 vpp_less_memory = _to_bool(
334 pipeline_cfg.get("vpp_less_memory", False),
335 "pipeline_config.vpp_less_memory",
336 )
337
338 optimization_level = _validate_int(
339 pipeline_cfg.get("optimization_level", 1),
340 "pipeline_config.optimization_level",
341 )
342
343 memory_limit = _validate_int(
344 pipeline_cfg.get("memory_limit", 0),
345 "pipeline_config.memory_limit",
346 )
347
348 constant_memory = _validate_int(
349 pipeline_cfg.get("constant_memory", 0),
350 "pipeline_config.constant_memory",
351 )
352
353 enable_simulation = _to_bool(
354 pipeline_cfg.get("enable_simulation", True),
355 "pipeline_config.enable_simulation",
356 )
357
358 sim_comm_time = float(pipeline_cfg.get("sim_comm_time", 0.0))
359
360 use_backward_time = _to_bool(
361 pipeline_cfg.get("use_backward_time", False),
362 "pipeline_config.use_backward_time",
363 )
364
365 return (
366 num_of_interleave, vpp_less_memory, optimization_level,
367 memory_limit, constant_memory, enable_simulation, sim_comm_time,
368 use_backward_time,
369 )
370
371
372def parse_yaml_for_optimization(yaml_path: str) -> YamlOptimizationConfig:
373 """Parse a YAML configuration file for PP optimization.
374
375 The YAML must contain a ``pipeline_config`` section with
376 ``pipeline_num`` and ``micro_batch_num``.
377 ``num_layer`` is optional; when omitted it defaults to ``None``
378 and the layer count is derived from the JSON profile.
379 ``num_of_interleave`` is optional and defaults to 1.
380 ``memory_limit``, ``constant_memory``, ``enable_simulation``,
381 ``sim_comm_time`` are optional.
382
383 Args:
384 yaml_path: Path to the YAML configuration file.
385
386 Returns:
387 :class:`YamlOptimizationConfig` with all required fields.
388
389 Raises:
390 ValueError: If required fields are missing or values are invalid.
391 FileNotFoundError: If the YAML file does not exist.
392
393 Example:
394 >>> config = parse_yaml_for_optimization("pp_config.yaml")
395 >>> config.pp_degree
396 4
397 >>> config.num_layer
398 32
399 """
400 with open(yaml_path, encoding="utf-8") as fp:
401 cfg = yaml.safe_load(fp)
402
403 if not isinstance(cfg, dict):
404 raise ValueError(
405 f"YAML file {yaml_path} must contain a top-level mapping, "
406 f"got {type(cfg).__name__}"
407 )
408
409 pipeline_cfg: dict = cfg.get("pipeline_config", {})
410 if not isinstance(pipeline_cfg, dict):
411 raise ValueError(
412 f"YAML file {yaml_path} must contain a 'pipeline_config' section"
413 )
414
415 pp_degree, num_layer, micro_batch_num = _extract_required_fields(
416 pipeline_cfg, yaml_path,
417 )
418 (
419 num_of_interleave, vpp_less_memory, optimization_level,
420 memory_limit, constant_memory, enable_simulation, sim_comm_time,
421 use_backward_time,
422 ) = _extract_optional_fields(pipeline_cfg)
423
424 config = YamlOptimizationConfig(
425 pp_degree=pp_degree,
426 num_layer=num_layer,
427 micro_batch_num=micro_batch_num,
428 num_of_interleave=num_of_interleave,
429 vpp_less_memory=vpp_less_memory,
430 optimization_level=optimization_level,
431 memory_limit=memory_limit,
432 constant_memory=constant_memory,
433 enable_simulation=enable_simulation,
434 sim_comm_time=sim_comm_time,
435 use_backward_time=use_backward_time,
436 )
437 config.validate()
438 return config