Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed_data / dataset_reader.py: 92%

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"""Worker options and collation for plan-aware sample loading.""" 

16# This distributed-data package is intentionally PyTorch-only. 

17 

18from __future__ import annotations 

19 

20from collections.abc import Mapping 

21from typing import Any 

22 

23import torch # pylint: disable=forbidden-backend-import 

24 

25_INTERNAL_DATALOADER_OPTIONS = frozenset({ 

26 "batch_sampler", 

27 "batch_size", 

28 "collate_fn", 

29 "dataset", 

30 "drop_last", 

31 "generator", 

32 "num_workers", 

33 "persistent_workers", 

34 "pin_memory", 

35 "prefetch_factor", 

36 "sampler", 

37 "shuffle", 

38}) 

39 

40 

41def _identity(value: Any) -> Any: 

42 return value 

43 

44 

45def _validate_worker_options(options: Mapping[str, Any]) -> None: 

46 """Validate worker settings shared by the public API and direct readers.""" 

47 num_workers = options["num_workers"] 

48 if not isinstance(num_workers, int) or isinstance(num_workers, bool) or num_workers < 0: 

49 raise ValueError(f"num_workers must be a non-negative integer, but got {num_workers!r}.") 

50 for name in ("pin_memory", "persistent_workers"): 

51 if not isinstance(options[name], bool): 

52 raise ValueError(f"{name} must be boolean, but got {options[name]!r}.") 

53 prefetch_factor = options["prefetch_factor"] 

54 if prefetch_factor is not None and ( 

55 not isinstance(prefetch_factor, int) or isinstance(prefetch_factor, bool) or prefetch_factor < 1 

56 ): 

57 raise ValueError("prefetch_factor must be a positive integer or None.") 

58 if num_workers == 0 and prefetch_factor is not None: 

59 raise ValueError("prefetch_factor requires num_workers > 0.") 

60 if options["persistent_workers"] and num_workers == 0: 

61 raise ValueError("persistent_workers=True requires num_workers > 0.") 

62 if options.get("in_order", True) is not True: 

63 raise ValueError("dataloader_kwargs in_order must remain True for deterministic sample routing.") 

64 

65 

66def _build_worker_options( 

67 *, 

68 num_workers: int, 

69 pin_memory: bool, 

70 prefetch_factor: int | None, 

71 persistent_workers: bool, 

72 generator: torch.Generator, 

73 dataloader_kwargs: Mapping[str, Any] | None, 

74) -> dict[str, Any]: 

75 """Merge additional DataLoader options without exposing owned semantics.""" 

76 try: 

77 options = dict({} if dataloader_kwargs is None else dataloader_kwargs) 

78 except Exception as exc: 

79 raise ValueError(f"dataloader_kwargs cannot be copied: {exc}") from exc 

80 conflicting_options = sorted(set(options) & _INTERNAL_DATALOADER_OPTIONS) 

81 if conflicting_options: 

82 raise ValueError(f"dataloader_kwargs contains internally managed options {conflicting_options}.") 

83 options.update({ 

84 "num_workers": num_workers, 

85 "pin_memory": pin_memory, 

86 "prefetch_factor": prefetch_factor, 

87 "persistent_workers": persistent_workers, 

88 "generator": generator, 

89 }) 

90 _validate_worker_options(options) 

91 if prefetch_factor is None: 

92 options.pop("prefetch_factor") 

93 return options