Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / registry.py: 62%

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"""Dataset registry — single ``data.type`` name → builder dispatch. 

16 

17The registry is the only piece of dataset-related code aware of the full 

18catalogue of supported formats. Each format lives in its own module and 

19registers a builder via ``@DATASET_REGISTRY.register("<name>")``; the trainer 

20then calls :func:`build_dataset` with the ``data.type`` from the YAML config 

21and never needs to know about specific formats. New formats plug in by adding 

22one module-level registration instead of modifying trainer code. 

23""" 

24from typing import Any, Callable, Dict, Iterator, List 

25 

26 

27BuilderFn = Callable[..., Any] 

28 

29 

30class DatasetRegistry: 

31 """Map ``data.type`` names to builder callables. 

32 

33 Each builder takes keyword arguments ``base`` (the :class:`BaseTrainer` 

34 instance), ``args`` (the parsed config) plus any extras the builder 

35 needs (``tokenizer``, ``data_transform`` ...). The registry never holds 

36 state about a specific run — builders are looked up fresh each time. 

37 

38 Example: 

39 >>> @DATASET_REGISTRY.register("dummy") 

40 ... def build_dummy(*, base, args, **_): 

41 ... return _DummyDataset(...) 

42 """ 

43 

44 def __init__(self) -> None: 

45 self._builders: Dict[str, BuilderFn] = {} 

46 

47 def register(self, name: str) -> Callable[[BuilderFn], BuilderFn]: 

48 """Decorator that registers ``fn`` under ``name``. 

49 

50 Raises: 

51 ValueError: When ``name`` is already registered. Re-registration 

52 is treated as a programming error (typo / duplicate import) 

53 rather than silently overwritten. 

54 """ 

55 if not name: 

56 raise ValueError("Dataset registry name must be a non-empty string") 

57 

58 def decorator(fn: BuilderFn) -> BuilderFn: 

59 if name in self._builders: 

60 raise ValueError( 

61 f"Dataset type '{name}' is already registered; " 

62 f"check for duplicate registrations / circular imports" 

63 ) 

64 self._builders[name] = fn 

65 return fn 

66 

67 return decorator 

68 

69 def get(self, name: str) -> BuilderFn: 

70 """Return the builder for ``name`` or raise with a helpful message.""" 

71 try: 

72 return self._builders[name] 

73 except KeyError as exc: 

74 available = sorted(self._builders) 

75 raise ValueError( 

76 f"Unknown data.type '{name}'. Available: {available}. " 

77 f"Add a new format by registering a builder under " 

78 f"hyper_parallel.data." 

79 ) from exc 

80 

81 def names(self) -> List[str]: 

82 """Return registered names in deterministic (sorted) order.""" 

83 return sorted(self._builders) 

84 

85 def __contains__(self, name: str) -> bool: 

86 return name in self._builders 

87 

88 def __iter__(self) -> Iterator[str]: 

89 return iter(self.names()) 

90 

91 

92DATASET_REGISTRY = DatasetRegistry() 

93 

94 

95def build_dataset(name: str, **kwargs: Any) -> Any: 

96 """Build a dataset by ``data.type`` name. 

97 

98 Args: 

99 name: Dataset type identifier — must match a registered builder. 

100 **kwargs: Forwarded to the builder. The trainer passes ``base`` 

101 (the :class:`BaseTrainer` instance) and ``args`` (the parsed 

102 config) plus, when available, ``tokenizer`` / ``data_transform``. 

103 

104 Returns: 

105 A ``torch.utils.data.Dataset``-compatible object suitable for the 

106 trainer's distributed sampler + ``StatefulDataLoader``. 

107 """ 

108 return DATASET_REGISTRY.get(name)(**kwargs)