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

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"""Resolve target-selected YAML groups into a typed trainer config tree.""" 

16 

17from __future__ import annotations 

18 

19__all__ = [ 

20 "ConfigResolutionError", 

21 "replace_override_path", 

22 "resolve_config", 

23] 

24 

25import dataclasses 

26import difflib 

27import importlib 

28import inspect 

29import types 

30from collections.abc import Iterable, Mapping 

31from dataclasses import MISSING, fields, is_dataclass, replace 

32from typing import Any, Literal, Union, get_args, get_origin, get_type_hints 

33 

34from hyper_parallel.trainer.config.data import ( 

35 DataLoaderConfig, 

36 DatasetConfig, 

37 ModelAssetsConfig, 

38) 

39from hyper_parallel.trainer.config.optimization import ( 

40 OptimizerConfig, 

41 OptimizerSwapConfig, 

42) 

43from hyper_parallel.trainer.config.target import Target 

44from hyper_parallel.components.modules.family import get_module_family 

45from hyper_parallel.trainer.config.replacement import ReplacementTarget 

46from hyper_parallel.trainer.config.trainer import TrainerConfig 

47 

48 

49_NONE_TYPE: type = types.NoneType # pylint: disable=no-member 

50_UNION_TYPE: type = types.UnionType # pylint: disable=no-member 

51 

52 

53class ConfigResolutionError(ValueError): 

54 """A target or typed configuration value is invalid.""" 

55 

56 def __init__(self, location: str, message: str) -> None: 

57 """Prefix the error message with its configuration location.""" 

58 super().__init__(f"{location}: {message}") 

59 

60 

61def import_target(target_path: str, *, location: str) -> object: 

62 """Import a callable from a dotted Python path. 

63 

64 Args: 

65 target_path: Module path followed by callable attributes. 

66 location: Configuration location used in error messages. 

67 

68 Returns: 

69 The imported callable, without invoking it. 

70 

71 Raises: 

72 ConfigResolutionError: The path is invalid, cannot be imported, or 

73 does not identify a callable. 

74 """ 

75 if not isinstance(target_path, str) or not target_path.strip(): 

76 raise ConfigResolutionError(location, "_target_ must be a non-empty dotted path") 

77 

78 parts = target_path.split(".") 

79 if any(not part for part in parts): 

80 raise ConfigResolutionError(location, f"invalid target path {target_path!r}") 

81 

82 for split_at in range(len(parts), 0, -1): 

83 module_name = ".".join(parts[:split_at]) 

84 try: 

85 target = importlib.import_module(module_name) 

86 except ModuleNotFoundError as exc: 

87 if exc.name == module_name or module_name.startswith(f"{exc.name}."): 

88 continue 

89 raise ConfigResolutionError( 

90 location, 

91 f"target {target_path!r} failed while importing dependency {exc.name!r}", 

92 ) from exc 

93 except ImportError as exc: 

94 raise ConfigResolutionError(location, f"target {target_path!r} could not be imported: {exc}") from exc 

95 

96 for attribute in parts[split_at:]: 

97 if not hasattr(target, attribute): 

98 raise ConfigResolutionError( 

99 location, 

100 f"target {target_path!r} has no attribute {attribute!r}", 

101 ) 

102 target = getattr(target, attribute) 

103 

104 if not callable(target): 

105 raise ConfigResolutionError(location, f"target {target_path!r} is not callable") 

106 return target 

107 

108 raise ConfigResolutionError(location, f"target {target_path!r} could not be imported") 

109 

110 

111def _is_union(annotation: object) -> bool: 

112 """Return whether the annotation is a ``Union`` or a PEP 604 union.""" 

113 return get_origin(annotation) in (Union, _UNION_TYPE) 

114 

115 

116def _annotation_name(annotation: object) -> str: 

117 """Format an annotation for configuration error messages.""" 

118 if annotation is Any: 

119 return "Any" 

120 if isinstance(annotation, type): 

121 return annotation.__qualname__ 

122 return str(annotation).replace("typing.", "") 

123 

124 

125def _target_hints(target: object, *, path: str) -> dict[str, object]: 

126 """Resolve annotations for a target function or class constructor.""" 

127 hint_source = target.__init__ if inspect.isclass(target) else target 

128 try: 

129 return get_type_hints(hint_source) 

130 except (NameError, TypeError) as exc: 

131 raise ConfigResolutionError(path, f"could not resolve target type annotations: {exc}") from exc 

132 

133 

134# Annotations whose payload is a container of other values. 

135_COLLECTION_TYPES: tuple[type, ...] = (list, tuple, dict, Mapping) 

136 

137 

138def _normalize_sequence( 

139 value: object, 

140 item_types: tuple, 

141 *, 

142 uniform: bool, 

143 kind: str, 

144 path: str, 

145) -> list | tuple: 

146 """Normalize sequence items against their declared types. 

147 

148 With ``uniform=True``, one item type applies to every element. Otherwise, 

149 ``item_types`` must contain one type per element. Empty ``item_types`` 

150 produces a tuple without normalizing its items. 

151 """ 

152 if not isinstance(value, (list, tuple)): 

153 raise ConfigResolutionError(path, f"expected {kind}, got {type(value).__name__}") 

154 if not item_types: 

155 return tuple(value) 

156 if uniform: 

157 item_types = item_types * len(value) 

158 if len(item_types) != len(value): 

159 raise ConfigResolutionError( 

160 path, 

161 f"expected {kind} of length {len(item_types)}, got {len(value)}", 

162 ) 

163 normalized = [ 

164 normalize_value(item, item_type, path=f"{path}[{index}]") 

165 for index, (item, item_type) in enumerate(zip(value, item_types)) 

166 ] 

167 return tuple(normalized) if kind == "tuple" else normalized 

168 

169 

170def _require_none_allowed(annotation: object, *, path: str) -> None: 

171 """Accept ``None`` only when the annotation permits it.""" 

172 if annotation is _NONE_TYPE or ( 

173 _is_union(annotation) and _NONE_TYPE in get_args(annotation) 

174 ): 

175 return None 

176 raise ConfigResolutionError(path, f"expected {_annotation_name(annotation)}, got None") 

177 

178 

179def _normalize_union(value: object, annotation: object, *, path: str) -> object: 

180 """Normalize a value against the first compatible union member.""" 

181 members = get_args(annotation) 

182 non_none_members = tuple(member for member in members if member is not _NONE_TYPE) 

183 if len(non_none_members) == 1 and len(non_none_members) != len(members): 

184 return normalize_value(value, non_none_members[0], path=path) 

185 

186 for member in members: 

187 try: 

188 return normalize_value(value, member, path=path) 

189 except ConfigResolutionError: 

190 continue 

191 raise ConfigResolutionError( 

192 path, 

193 f"expected {_annotation_name(annotation)}, got {type(value).__name__}", 

194 ) 

195 

196 

197def _normalize_scalar(value: object, annotation: object, *, path: str) -> object: 

198 """Validate a scalar's type, allowing integers to normalize to floats.""" 

199 if annotation is bool and isinstance(value, bool): 

200 return value 

201 if annotation is int and isinstance(value, int) and not isinstance(value, bool): 

202 return value 

203 if annotation is float and isinstance(value, (int, float)) and not isinstance(value, bool): 

204 return float(value) 

205 if annotation is str and isinstance(value, str): 

206 return value 

207 raise ConfigResolutionError( 

208 path, 

209 f"expected {_annotation_name(annotation)}, got {type(value).__name__}", 

210 ) 

211 

212 

213def _normalize_literal(value: object, choices: tuple, *, path: str) -> object: 

214 """Validate a literal value, restoring YAML boolean aliases when allowed.""" 

215 # PyYAML 1.1 parses unquoted on/off/yes/no/true/false scalars as bool. 

216 # Map the bool back to the matching word when it is one of the choices. 

217 if isinstance(value, bool): 

218 words = ("on", "yes", "true") if value else ("off", "no", "false") 

219 for word in words: 

220 if word in choices: 

221 return word 

222 if any(type(value) is type(choice) and value == choice for choice in choices): 

223 return value 

224 expected = ", ".join(repr(choice) for choice in choices) 

225 raise ConfigResolutionError(path, f"expected one of ({expected}), got {value!r}") 

226 

227 

228def _normalize_collection( 

229 value: object, 

230 annotation: object, 

231 origin: object, 

232 args: tuple, 

233 *, 

234 path: str, 

235) -> object: 

236 """Normalize a value against a list, tuple, or mapping annotation.""" 

237 if origin is list or annotation is list: 

238 return _normalize_sequence( 

239 value, 

240 (args[0] if args else Any,), 

241 uniform=True, 

242 kind="list", 

243 path=path, 

244 ) 

245 if origin is tuple or annotation is tuple: 

246 if len(args) == 2 and args[1] is Ellipsis: 

247 return _normalize_sequence(value, args[:1], uniform=True, kind="tuple", path=path) 

248 return _normalize_sequence(value, args, uniform=False, kind="tuple", path=path) 

249 if not isinstance(value, Mapping): 

250 raise ConfigResolutionError(path, f"expected mapping, got {type(value).__name__}") 

251 return dict(value) 

252 

253 

254def _normalize_class(value: object, annotation: object, *, path: str) -> object: 

255 """Validate that a value is an instance of the annotated class.""" 

256 if not isinstance(annotation, type): 

257 raise ConfigResolutionError(path, f"unsupported type annotation {_annotation_name(annotation)}") 

258 if isinstance(value, annotation): 

259 return value 

260 raise ConfigResolutionError( 

261 path, 

262 f"expected {_annotation_name(annotation)}, got {type(value).__name__}", 

263 ) 

264 

265 

266def _normalize_override_scalar(value: object, annotation: object) -> object: 

267 """Convert a CLI string to the numeric type declared by its annotation. 

268 

269 Supports ``int``, ``float``, and their optional forms. YAML may leave 

270 numeric CLI input such as ``1e-4`` as a string. Failed conversions leave 

271 the value unchanged for subsequent validation by ``normalize_value``. 

272 """ 

273 if not isinstance(value, str): 

274 return value 

275 target = annotation 

276 if get_origin(target) in (Union, types.UnionType): 

277 members = [member for member in get_args(target) if member is not types.NoneType] 

278 if len(members) != 1: 

279 return value 

280 target = members[0] 

281 if target is int: 

282 try: 

283 return int(value) 

284 except ValueError: 

285 return value 

286 if target is float: 

287 try: 

288 return float(value) 

289 except ValueError: 

290 return value 

291 return value 

292 

293 

294def normalize_value(value: object, annotation: object, *, path: str) -> object: 

295 """Validate and normalize a parsed value against its annotation. 

296 

297 Args: 

298 value: Parsed value to validate or convert. 

299 annotation: Type annotation defining the accepted value. 

300 path: Configuration path used in error messages. 

301 

302 Returns: 

303 The validated or converted value. 

304 

305 Raises: 

306 ConfigResolutionError: The value does not match the annotation or the 

307 annotation is unsupported. 

308 """ 

309 # A nested target is a deferred instance of the annotated runtime type. 

310 # Its constructor arguments are validated when its own node is resolved; 

311 # the parent Target materializes it immediately before invocation. 

312 if isinstance(value, Target): 

313 return value 

314 if annotation in (Any, object): 

315 return value 

316 if isinstance(annotation, dataclasses.InitVar): 

317 return normalize_value(value, annotation.type, path=path) 

318 if value is None: 

319 return _require_none_allowed(annotation, path=path) 

320 if _is_union(annotation): 

321 return _normalize_union(value, annotation, path=path) 

322 if annotation in (bool, int, float, str): 

323 return _normalize_scalar(value, annotation, path=path) 

324 

325 origin = get_origin(annotation) 

326 args = get_args(annotation) 

327 if origin is Literal: 

328 return _normalize_literal(value, args, path=path) 

329 if origin in _COLLECTION_TYPES or annotation in _COLLECTION_TYPES: 

330 return _normalize_collection(value, annotation, origin, args, path=path) 

331 if isinstance(annotation, type) and dataclasses.is_dataclass(annotation): 

332 # Nested dataclass items resolve from mappings in the same way as 

333 # top-level dataclass components. 

334 return _resolve_dataclass(value, annotation, path=path) 

335 return _normalize_class(value, annotation, path=path) 

336 

337 

338def _suggestion(name: str, candidates: Iterable[str]) -> str: 

339 """Return a spelling suggestion, or an empty string if no name is close.""" 

340 matches = difflib.get_close_matches(name, candidates, n=1) 

341 return f"; did you mean {matches[0]!r}?" if matches else "" 

342 

343 

344def _replace_family_selector( 

345 config: ReplacementTarget, parts: list[str], value: object, *, path: str, 

346) -> ReplacementTarget: 

347 """Validate one terminal selector and report reselection errors at its path.""" 

348 if len(parts) != 1: 

349 raise ConfigResolutionError(f"CLI.{path}", "a scalar selector cannot have child fields") 

350 try: 

351 return config.replace(**{parts[0]: value}) 

352 except (TypeError, ValueError) as exc: 

353 raise ConfigResolutionError(f"CLI.{path}", str(exc)) from exc 

354 

355 

356def _replace_target_path( 

357 config: Target[Any], 

358 parts: list[str], 

359 value: object, 

360 *, 

361 path: str, 

362) -> Target[Any]: 

363 """Return a ``Target`` copy with an argument or nested value replaced.""" 

364 name = parts[0] 

365 full_path = f"{path}.{name}" if path else name 

366 if name == "_target_": 

367 raise ConfigResolutionError( 

368 f"CLI.{full_path}", "changing _target_ through an override is not supported" 

369 ) 

370 

371 selector_spec = config.get_family_selector_spec(name) if isinstance(config, ReplacementTarget) else None 

372 if selector_spec is not None: 

373 return _replace_family_selector(config, parts, value, path=full_path) 

374 

375 signature = inspect.signature(config.callable) 

376 parameter = signature.parameters.get(name) 

377 has_var_kwargs = any( 

378 item.kind is inspect.Parameter.VAR_KEYWORD 

379 for item in signature.parameters.values() 

380 ) 

381 if parameter is None and not has_var_kwargs: 

382 candidates = [ 

383 item.name 

384 for item in signature.parameters.values() 

385 if item.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY) 

386 ] 

387 raise ConfigResolutionError( 

388 f"CLI.{full_path}", 

389 f"unknown target argument {name!r}{_suggestion(name, candidates)}", 

390 ) 

391 

392 if len(parts) > 1: 

393 try: 

394 child = getattr(config, name) 

395 except AttributeError as exc: 

396 raise ConfigResolutionError( 

397 f"CLI.{full_path}", "target argument is not configured" 

398 ) from exc 

399 return config.replace(**{name: replace_override_path(child, parts[1:], value, path=full_path)}) 

400 

401 normalized = value 

402 if parameter is not None: 

403 annotation = _target_hints(config.callable, path=path).get( 

404 name, 

405 parameter.annotation, 

406 ) 

407 if annotation not in (Any, object, inspect.Signature.empty): 

408 normalized = normalize_value( 

409 _normalize_override_scalar(value, annotation), 

410 annotation, 

411 path=f"CLI.{full_path}", 

412 ) 

413 return config.replace(**{name: normalized}) 

414 

415 

416def replace_override_path(config: object, parts: list[str], value: object, *, path: str) -> object: 

417 """Return a configuration copy with a dotted-path override applied. 

418 

419 Args: 

420 config: Resolved configuration object or nested value to update. 

421 parts: Nonempty sequence of remaining dotted-path components. 

422 value: Parsed override value to validate and assign. 

423 path: Path already traversed, or an empty string at the root. 

424 

425 Returns: 

426 A copy with the override applied; the input is not modified. 

427 

428 Raises: 

429 ConfigResolutionError: The path cannot be updated or the value does 

430 not match the selected field's annotation. 

431 """ 

432 if isinstance(config, Target): 

433 return _replace_target_path(config, parts, value, path=path) 

434 

435 if isinstance(config, Mapping): 

436 name = parts[0] 

437 full_path = f"{path}.{name}" if path else name 

438 if name not in config: 

439 raise ConfigResolutionError(f"CLI.{path}", f"unknown mapping key {name!r}") 

440 if len(parts) == 1: 

441 return {**config, name: value} 

442 return {**config, name: replace_override_path(config[name], parts[1:], value, path=full_path)} 

443 

444 location = f"CLI.{path}" if path else "CLI" 

445 if not is_dataclass(config): 

446 raise ConfigResolutionError( 

447 location, f"value of type {type(config).__name__} has no configurable fields" 

448 ) 

449 

450 config_fields = {field.name: field for field in fields(config)} 

451 name = parts[0] 

452 if name not in config_fields: 

453 target = getattr(config, "target", None) 

454 if isinstance(target, Target): 

455 return replace(config, target=_replace_target_path(target, parts, value, path=path)) 

456 raise ConfigResolutionError( 

457 location, 

458 f"unknown field {name!r} on {type(config).__name__}" 

459 f"{_suggestion(name, config_fields)}", 

460 ) 

461 

462 full_path = f"{path}.{name}" if path else name 

463 if len(parts) == 1: 

464 annotation = get_type_hints(type(config))[name] 

465 normalized = normalize_value( 

466 _normalize_override_scalar(value, annotation), annotation, path=f"CLI.{full_path}" 

467 ) 

468 return replace(config, **{name: normalized}) 

469 

470 child = getattr(config, name) 

471 if child is None: 

472 raise ConfigResolutionError( 

473 f"CLI.{full_path}", "component was not selected by the YAML" 

474 ) 

475 return replace(config, **{name: replace_override_path(child, parts[1:], value, path=full_path)}) 

476 

477 

478def _resolve_union(node: object, annotation: object, *, path: str) -> object: 

479 """Resolve a YAML value against a compatible non-``None`` union member.""" 

480 non_none_members = [ 

481 member for member in get_args(annotation) if member is not _NONE_TYPE 

482 ] 

483 if isinstance(node, Mapping) and "_target_" in node: 

484 target_member = next( 

485 (member for member in non_none_members 

486 if get_origin(member) is Target or member is Target), 

487 None, 

488 ) 

489 if target_member is not None: 

490 # A _target_ mapping is explicit; preserve its import/signature error 

491 # instead of masking it behind a generic union mismatch. 

492 return resolve_component(node, annotation=target_member, path=path) 

493 if len(non_none_members) == 1: 

494 # Single-member union (e.g. Optional[Target]): resolve directly so the 

495 # member's specific error (e.g. an unexpected target argument) reaches 

496 # the user instead of being swallowed by the generic union failure. 

497 return resolve_component(node, annotation=non_none_members[0], path=path) 

498 

499 last_error: ConfigResolutionError | None = None 

500 for member in non_none_members: 

501 try: 

502 return resolve_component(node, annotation=member, path=path) 

503 except ConfigResolutionError as exc: 

504 last_error = exc 

505 raise ConfigResolutionError( 

506 path, 

507 f"expected {_annotation_name(annotation)}, got {type(node).__name__}", 

508 ) from last_error 

509 

510 

511def _resolve_dataclass(node: object, config_type: type, *, path: str) -> object: 

512 """Resolve YAML fields and construct a configuration dataclass.""" 

513 if not isinstance(node, Mapping): 

514 raise ConfigResolutionError(path, "configuration section must be a YAML mapping") 

515 

516 config_fields = {field.name: field for field in fields(config_type)} 

517 unknown = sorted(set(node) - set(config_fields)) 

518 if unknown: 

519 raise ConfigResolutionError(path, f"unknown configuration fields: {unknown}") 

520 

521 required = [ 

522 field.name 

523 for field in config_fields.values() 

524 if field.default is MISSING and field.default_factory is MISSING 

525 ] 

526 missing = [name for name in required if name not in node] 

527 if missing: 

528 raise ConfigResolutionError(path, f"missing required configuration fields: {missing}") 

529 

530 try: 

531 hints = get_type_hints(config_type) 

532 except (NameError, TypeError) as exc: 

533 raise ConfigResolutionError(path, f"could not resolve configuration type annotations: {exc}") from exc 

534 

535 resolved = { 

536 name: resolve_component( 

537 value, 

538 annotation=hints[name], 

539 path=f"{path}.{name}", 

540 ) 

541 for name, value in node.items() 

542 } 

543 try: 

544 return config_type(**resolved) 

545 except TypeError as exc: 

546 raise ConfigResolutionError(path, f"could not construct {config_type.__name__}: {exc}") from exc 

547 

548 

549def _resolve_nested_target_nodes(value: object, *, path: str) -> object: 

550 """Resolve reserved ``_target_`` nodes inside a target argument tree.""" 

551 if isinstance(value, Mapping): 

552 if "_target_" in value: 

553 return _resolve_target(value, path=path) 

554 return { 

555 key: _resolve_nested_target_nodes(item, path=f"{path}.{key}") 

556 for key, item in value.items() 

557 } 

558 if isinstance(value, list): 

559 return [ 

560 _resolve_nested_target_nodes(item, path=f"{path}[{index}]") 

561 for index, item in enumerate(value) 

562 ] 

563 if isinstance(value, tuple): 

564 return tuple( 

565 _resolve_nested_target_nodes(item, path=f"{path}[{index}]") 

566 for index, item in enumerate(value) 

567 ) 

568 return value 

569 

570 

571def _resolve_target_args( 

572 raw_args: Mapping[str, object], 

573 signature: inspect.Signature, 

574 hints: Mapping[str, object], 

575 *, 

576 path: str, 

577) -> dict[str, object]: 

578 """Validate target arguments and fill omitted keyword-callable defaults. 

579 

580 Positional-only defaults are not added because ``Target`` invokes its 

581 callable with keyword arguments. 

582 """ 

583 try: 

584 signature.bind_partial(**raw_args) 

585 except TypeError as exc: 

586 raise ConfigResolutionError(path, f"target arguments are invalid: {exc}") from exc 

587 

588 normalized = {} 

589 for name, value in raw_args.items(): 

590 parameter = signature.parameters.get(name) 

591 value = _resolve_nested_target_nodes(value, path=f"{path}.{name}") 

592 if parameter is None: 

593 normalized[name] = value 

594 continue 

595 

596 annotation = hints.get(name, parameter.annotation) 

597 if annotation in (Any, object, inspect.Signature.empty): 

598 normalized[name] = value 

599 else: 

600 normalized[name] = normalize_value( 

601 value, 

602 annotation, 

603 path=f"{path}.{name}", 

604 ) 

605 

606 for name, parameter in signature.parameters.items(): 

607 if ( 

608 name in normalized 

609 or parameter.default is inspect.Signature.empty 

610 or parameter.kind is inspect.Parameter.POSITIONAL_ONLY 

611 ): 

612 continue 

613 normalized[name] = parameter.default 

614 

615 return normalized 

616 

617 

618def _resolve_target(node: object, *, path: str) -> Target[Any]: 

619 """Resolve a YAML mapping into a ``Target`` without invoking its callable.""" 

620 if not isinstance(node, Mapping): 

621 raise ConfigResolutionError(path, "target section must be a YAML mapping") 

622 if "_target_" not in node: 

623 raise ConfigResolutionError(path, "target section is missing required _target_") 

624 

625 target_path = node["_target_"] 

626 target = import_target(target_path, location=f"{path}._target_") 

627 try: 

628 signature = inspect.signature(target) 

629 except (TypeError, ValueError) as exc: 

630 raise ConfigResolutionError(path, f"target signature is unavailable: {exc}") from exc 

631 

632 for parameter in signature.parameters.values(): 

633 if ( 

634 parameter.kind is inspect.Parameter.POSITIONAL_ONLY 

635 and parameter.default is inspect.Signature.empty 

636 ): 

637 raise ConfigResolutionError( 

638 path, 

639 f"target parameter {parameter.name!r} must be callable by keyword", 

640 ) 

641 

642 raw_args = {key: value for key, value in node.items() if key != "_target_"} 

643 normalized_args = _resolve_target_args( 

644 raw_args, 

645 signature, 

646 _target_hints(target, path=path), 

647 path=path, 

648 ) 

649 return Target( 

650 target, 

651 target_path=target_path, 

652 **normalized_args, 

653 ) 

654 

655 

656def _resolve_replacement_target(node: object, *, path: str) -> Target[Any]: 

657 """Apply Family selection only to registered baselines in replacement entries.""" 

658 if not isinstance(node, Mapping) or "_target_" not in node: 

659 raise ConfigResolutionError(path, "replacement requires a mapping with _target_") 

660 target_path = node["_target_"] 

661 requested = import_target(target_path, location=f"{path}._target_") 

662 if get_module_family(requested) is None: 

663 return _resolve_target(node, path=path) 

664 try: 

665 selected = ReplacementTarget( 

666 requested, target_path=target_path, 

667 **{key: value for key, value in node.items() if key != "_target_"}, 

668 ) 

669 if selected.low_precision_dtype_scheme is not None: 

670 return selected 

671 selected_signature = inspect.signature(selected.callable) 

672 for parameter in selected_signature.parameters.values(): 

673 if parameter.kind is inspect.Parameter.POSITIONAL_ONLY and parameter.default is inspect.Parameter.empty: 

674 raise TypeError(f"replacement parameter {parameter.name!r} must be callable by keyword") 

675 # Normalize raw arguments; to_dict would serialize configured objects and callables. 

676 normalized = _resolve_target_args( 

677 selected._kwargs, selected_signature, 

678 _target_hints(selected.callable, path=path), path=path, 

679 ) 

680 return selected.replace(**normalized) 

681 except (ValueError, TypeError, ImportError, AttributeError) as exc: 

682 raise ConfigResolutionError(path, str(exc)) from exc 

683 

684 

685def _resolve_dataloader_config(node: object, *, path: str) -> DataLoaderConfig: 

686 """Resolve a DataLoader target with nested collator and batch runtime.""" 

687 if not isinstance(node, Mapping): 

688 raise ConfigResolutionError(path, "DataLoader configuration must be a YAML mapping") 

689 

690 target_node = dict(node) 

691 if "dataloader_type" in target_node: 

692 raise ConfigResolutionError(f"{path}.dataloader_type", "renamed to sampler_type") 

693 

694 if "batch_adapter" in target_node: 

695 raise ConfigResolutionError( 

696 f"{path}.batch_adapter", 

697 "removed by the Omni data lifecycle; use dataset.data_transform for " 

698 "sample/batch encoding and dataloader.get_batch.runtime_input_adapter " 

699 "for model-owned forward metadata", 

700 ) 

701 

702 collate_node = target_node.pop("collate_fn", None) 

703 get_batch_node = target_node.pop("get_batch", None) 

704 policy_nodes = {name: target_node.pop(name, None) for name in ("cost_model", "balancing_algorithm")} 

705 load_balance = None 

706 if "load_balance" in target_node: 

707 balance_node = target_node.pop("load_balance") 

708 load_balance = normalize_value( 

709 {} if balance_node is None else balance_node, dict[str, Any], path=f"{path}.load_balance", 

710 ) 

711 sampler_type = normalize_value( 

712 target_node.pop("sampler_type", "single"), 

713 Literal["single", "cyclic"], 

714 path=f"{path}.sampler_type", 

715 ) 

716 data_rearrange_map = target_node.pop("data_rearrange_map", None) 

717 data_sharding = normalize_value( 

718 target_node.pop("data_sharding", False), 

719 bool, 

720 path=f"{path}.data_sharding", 

721 ) 

722 use_background_prefetcher = normalize_value( 

723 target_node.pop("use_background_prefetcher", False), 

724 bool, 

725 path=f"{path}.use_background_prefetcher", 

726 ) 

727 target = _resolve_target(target_node, path=path) 

728 collate_fn = ( 

729 None 

730 if collate_node is None 

731 else _resolve_target(collate_node, path=f"{path}.collate_fn") 

732 ) 

733 get_batch = ( 

734 None 

735 if get_batch_node is None 

736 else _resolve_target(get_batch_node, path=f"{path}.get_batch") 

737 ) 

738 return DataLoaderConfig( 

739 target=target, 

740 collate_fn=collate_fn, 

741 get_batch=get_batch, 

742 load_balance=load_balance, 

743 **{ 

744 name: None if value is None else _resolve_target(value, path=f"{path}.{name}") 

745 for name, value in policy_nodes.items() 

746 }, 

747 sampler_type=sampler_type, 

748 data_rearrange_map=data_rearrange_map, 

749 data_sharding=data_sharding, 

750 use_background_prefetcher=use_background_prefetcher, 

751 ) 

752 

753 

754def _resolve_dataset_config(node: object, *, path: str) -> DatasetConfig: 

755 """Resolve a ``DatasetConfig`` with its model assets and sample transform.""" 

756 if not isinstance(node, Mapping): 

757 raise ConfigResolutionError(path, "Dataset configuration must be a YAML mapping") 

758 

759 target_node = dict(node) 

760 model_assets_node = target_node.pop("model_assets", {}) 

761 data_transform_node = target_node.pop("data_transform", None) 

762 training_mode_node = target_node.pop("training_mode", None) 

763 training_mode = None if training_mode_node is None else normalize_value( 

764 training_mode_node, Literal["sft", "pretrain"], path=f"{path}.training_mode" 

765 ) 

766 # Dataset classes need path/split arguments supplied later by the indexed builder. 

767 data_options = dict(target_node.get("data_config", {})) 

768 dataset_class = data_options.pop("low_level_dataset", None) 

769 if dataset_class is not None: 

770 target_node["data_config"] = data_options 

771 target = _resolve_target(target_node, path=path) 

772 if dataset_class is not None: 

773 target = target.replace(data_config={**target.data_config, "low_level_dataset": dataset_class}) 

774 if isinstance(model_assets_node, Mapping) and "_target_" in model_assets_node: 

775 model_assets = _resolve_target( 

776 model_assets_node, 

777 path=f"{path}.model_assets", 

778 ) 

779 else: 

780 model_assets = resolve_component( 

781 model_assets_node, 

782 annotation=ModelAssetsConfig, 

783 path=f"{path}.model_assets", 

784 ) 

785 data_transform = ( 

786 None 

787 if data_transform_node is None