Coverage for /home/jenkins/.local/lib/python3.10/site-packages/hyper_parallel/data/indexed/indexed_data

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

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

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 indexed data paths and normalize GPT Dataset build options."""

17from __future__ import annotations

19import os

20import re

21from collections.abc import Mapping, Sequence

22from dataclasses import dataclass, field

23from importlib import import_module

24from typing import Any

26import numpy as np

28import torch.distributed as dist

30from hyper_parallel.data.dataset_logging import get_dataset_logger

32logger = get_dataset_logger(__name__)

35@dataclass

36class BlendedMegatronDatasetConfig:

37 """Configuration shared by blended indexed datasets."""

39 random_seed: int

40 sequence_length: int

41 blend: list[str] | None = None

42 blend_per_split: list[list[str] | None] | None = None

43 split: str | None = None

44 split_matrix: list[tuple[float, float] | None] | None = field(init=False, default=None)

45 path_to_cache: str | None = None

46 mmap_bin_files: bool = True

47 mock: bool = False

48 tokenizer: Any = None

50 def __post_init__(self) -> None:

51 """Validate the configured data distribution and derive split ranges."""

52 if self.mock:

53 return

55 if self.blend_per_split is not None and any(self.blend_per_split):

56 if self.blend is not None:

57 raise ValueError("blend and blend_per_split are incompatible")

58 if self.split is not None:

59 raise ValueError("split and blend_per_split are incompatible")

60 if len(self.blend_per_split) != 3:

61 raise ValueError("blend_per_split must contain train, validation, and test blends")

62 return

64 if self.blend is None:

65 raise ValueError("one of either blend or blend_per_split must be provided")

66 if self.split is None:

67 raise ValueError("both blend and split must be provided")

69 # Parse at most three train/validation/test ratios, padding omitted

70 # trailing splits with zero before normalization.

71 split_values = [float(value) for value in re.findall(r"[.0-9]+", self.split)]

72 split_values.extend([0.0] * (3 - len(split_values)))

73 if len(split_values) != 3:

74 raise ValueError(f"split length {len(split_values)} does not match expected 3")

76 if not all(value >= 0.0 for value in split_values):

77 raise ValueError("all split values must be non-negative")

79 normalized_values = (np.asarray(split_values, dtype=np.float64) / np.sum(split_values)).tolist()

80 canonical_values = []

81 for value in normalized_values:

82 rounded_value = round(value, 12)

83 precision = 2 if rounded_value == round(rounded_value, 2) else 12

84 canonical_values.append(f"{rounded_value:.{precision}f}")

85 self.split = ", ".join(canonical_values)

87 # Convert normalized ratios into non-overlapping ranges. A zero-sized

88 # split is represented by None.

89 boundaries = [0.0]

90 for value in normalized_values:

91 boundaries.append(boundaries[-1] + value)

92 self.split_matrix = []

93 for begin, end in zip(boundaries[:-1], boundaries[1:]):

94 self.split_matrix.append(None if end <= begin else (begin, end))

97@dataclass

98class GPTDatasetConfig(BlendedMegatronDatasetConfig):

99 """Configuration required to build indexed GPT Dataset splits."""

100

101 reset_position_ids: bool = False

102 reset_attention_mask: bool = False

103 eod_mask_loss: bool = False

104 create_ltor_fields: bool = False

105 drop_last_partial_validation_sequence: bool = True

106 add_extra_token_to_sequence: bool = True

107 allow_ambiguous_pad_tokens: bool = False

108 dataset_margin: float = 1.005

109 skip_data_check: bool = False

110 reuse_idx: bool = False

111 data_lazy_load: bool = False

112

113 low_level_dataset: type | None = None

114 simple_blend: str | None = None

115 dataset_options: Mapping[str, Any] = field(default_factory=dict)

116

117 def __post_init__(self) -> None:

118 """Validate the blended Dataset fields followed by GPT-specific fields."""

119 super().__post_init__()

120

121 if self.low_level_dataset is None:

122 raise ValueError("low_level_dataset must be explicitly configured")

123 if self.simple_blend is None:

124 raise ValueError("simple_blend must be explicitly configured")

125 if self.sequence_length <= 0:

126 raise ValueError("sequence_length must be positive")

127 if self.tokenizer is None and not self.mock:

128 raise ValueError("Attribute 'tokenizer' must not be None")

129

130

131def _discover_indexed_paths(

132 current_path: str,

133 rank: int,

134 world_size: int,

135 distributed_walk: bool,

136 file_suffix: str,

137) -> list[str]:

138 """Discover indexed prefixes under one directory."""

139 if not distributed_walk:

140 return _walk_data_directories([current_path], file_suffix)

141

142 root_layout = None

143 if rank == 0:

144 subdirectories = []

145 top_files = []

146 for name in os.listdir(current_path):

147 path = os.path.join(current_path, name)

148 if os.path.isdir(path):

149 subdirectories.append(path)

150 elif name.endswith(file_suffix):

151 top_files.append(os.path.splitext(path)[0] if file_suffix == ".bin" else path)

152 root_layout = (subdirectories, top_files)

153

154 gathered_layouts: list[tuple[list[str], list[str]] | None] = [None] * world_size

155 dist.all_gather_object(gathered_layouts, root_layout)

156 rank_zero_layout = gathered_layouts[0]

157 if rank_zero_layout is None:

158 raise ValueError("Rank 0 did not provide the indexed directory layout")

159 subdirectories, file_paths = rank_zero_layout

160 local_files = _walk_data_directories(subdirectories[rank::world_size], file_suffix)

161 gathered_files: list[list[str] | None] = [None] * world_size

162 dist.all_gather_object(gathered_files, local_files)

163 for rank_files in gathered_files:

164 if rank_files:

165 file_paths.extend(rank_files)

166 return file_paths

167

168

169def resolve_data_paths(

170 data_path: str | Sequence[str],

171 *,

172 distributed_walk: bool = True,

173 file_suffix: str = ".bin",

174) -> list[str]:

175 """Discover independent data files and return source paths with optional weights.

176

177 Args:

178 data_path: One indexed prefix, Arrow file, or sequence of directories/source paths.

179 distributed_walk: Whether initialized ranks divide directory walking.

180 file_suffix: Dataset storage suffix: .bin returns prefixes, .arrow keeps full file paths.

181

182 Returns:

183 One prefix directly, or alternating weight/prefix values for multiple sources. If ``data_path`` is not a

184 collection of discoverable indexed paths, return it unchanged so an explicit blend can be parsed.

185 """

186 # 1. Normalize one path and multiple paths to the same traversal form.

187 data_paths = [data_path] if isinstance(data_path, str) else list(data_path)

188

189 # 2. Use distributed walking only after the process group is available.

190 # Before initialization, one process walks all configured directories.

191 try:

192 rank = int(dist.get_rank())

193 world_size = int(dist.get_world_size())

194 except (RuntimeError, ValueError):

195 rank, world_size = 0, 1

196 use_distributed_walk = distributed_walk and world_size > 1

197

198 # 3. Accept a direct prefix, or recursively discover prefixes in a directory.

199 file_paths: list[str] = []

200 for current_path in data_paths:

201 if file_suffix == ".arrow" and current_path.endswith(file_suffix) and os.path.isfile(current_path):

202 file_paths.append(current_path)

203 continue

204 if file_suffix == ".bin" and os.path.isfile(current_path + ".bin"):

205 file_paths.append(current_path)

206 continue

207 if not os.path.isdir(current_path):

208 # Keep args.data_path unchanged when filesystem discovery

209 # returns no files. This preserves explicit forms such as

210 # ["30", prefix_a, "70", prefix_b] for the blend parser.

211 original_paths = list(data_paths)

212 return original_paths

213

214 file_paths.extend(_discover_indexed_paths(current_path, rank, world_size, use_distributed_walk, file_suffix))

215

216 if not file_paths:

217 raise ValueError(f"data_path {data_paths!r} has no data file ending with {file_suffix}")

218

219 if len(file_paths) == 1:

220 return file_paths

221

222 if file_suffix == ".arrow":

223 # Every rank must blend files in the same order after directory discovery.

224 file_paths.sort()

225 else:

226 # 4. Prefer the numeric shard order encoded by the corpus filenames.

227 # Unknown naming schemes remain usable in their original walk order.

228 try:

229 def numeric_shard_key(path: str) -> tuple[int, int]:

230 """Extract numeric corpus and shard identifiers from a path."""

231 path_parts = os.path.basename(path).split("_")

232 return int(path_parts[2]), int(path_parts[-1])

233

234 file_paths.sort(key=numeric_shard_key)

235 except (IndexError, TypeError, ValueError):

236 logger.warning("Cannot sort indexed files with the numeric filename rule; keeping walk order")

237

238 # 5. Automatically discovered sources use equal blend weights.

239 weighted_paths = [value for path in file_paths for value in ("1", path)]

240 return weighted_paths

241

242

243def build_gpt_dataset_config(

244 data_paths: Sequence[str],

245 data_config: Mapping[str, Any],

246) -> GPTDatasetConfig:

247 """Build a GPT Dataset config from provider fields.

248

249 Args:

250 data_paths: Weighted prefixes returned by :func:`resolve_data_paths`.

251 data_config: Indexed pretraining configuration supplied by the caller.

252

253 Returns:

254 Configuration required by GPT Dataset and split construction.

255

256 Raises:

257 ValueError: If a required build field is missing.

258 """

259 # The source argument parser owns defaults and validation. This stage only

260 # renames fields where the Dataset implementation expects another name.

261 blend_per_split = [

262 data_config.get("train_data_path"),

263 data_config.get("valid_data_path"),

264 data_config.get("test_data_path"),

265 ]

266 has_independent_blends = any(blend_per_split)

267 try:

268 mock_data = bool(data_config.get("mock_data", False))

269 tokenizer = data_config.get("tokenizer")

270 if tokenizer is None and not mock_data:

271 raise ValueError("tokenizer is required when mock_data is false")

272

273 config = GPTDatasetConfig(

274 random_seed=data_config["random_seed"],

275 sequence_length=data_config["seq_length"],

276 tokenizer=tokenizer,

277 mock=mock_data,

278 blend=None if has_independent_blends else list(data_paths),

279 blend_per_split=blend_per_split if has_independent_blends else None,

280 split=None if has_independent_blends else data_config.get("split", "1, 0, 0"),

281 path_to_cache=data_config.get("data_cache_path"),

282 mmap_bin_files=data_config.get("mmap_bin_files", True),

283 reuse_idx=data_config.get("reuse_idx", False),

284 data_lazy_load=data_config.get("data_lazy_load", False),

285 reset_position_ids=data_config.get("reset_position_ids", False),

286 reset_attention_mask=data_config.get("reset_attention_mask", False),

287 eod_mask_loss=data_config.get("eod_mask_loss", False),

288 create_ltor_fields=data_config.get("create_ltor_fields_in_dataloader", False),

289 drop_last_partial_validation_sequence=data_config.get("drop_last_partial_validation_sequence", True),

290 add_extra_token_to_sequence=data_config.get("add_extra_token_to_sequence", True),

291 allow_ambiguous_pad_tokens=data_config.get("allow_ambiguous_pad_tokens", False),

292 dataset_margin=data_config.get("dataset_margin", 1.005),

293 skip_data_check=data_config.get("skip_data_check", False),

294 low_level_dataset=resolve_dataset_type(data_config["low_level_dataset"]),

295 dataset_options=(

296 {key: value for key, value in data_config["low_level_dataset"].items() if key != "_target_"}

297 if isinstance(data_config["low_level_dataset"], Mapping) else {}

298 ),

299 simple_blend=data_config["simple_blend"],

300 )

301 except KeyError as error:

302 raise ValueError(f"Missing indexed Dataset build option: {error.args[0]}") from error

303 source_paths = list(data_paths) if len(data_paths) <= 1 else list(data_paths[1::2])

304 logger.debug("Resolved indexed sources=%d, first four=%s", len(source_paths), source_paths[:4])

305 return config

306

307

308def _walk_data_directories(directories: Sequence[str], file_suffix: str) -> list[str]:

309 """Recursively find source files below the assigned directories."""

310 file_paths = []

311 for directory in directories:

312 for root, _, filenames in os.walk(directory):

313 for filename in filenames:

314 if filename.endswith(file_suffix):

315 path = os.path.abspath(os.path.join(root, filename))

316 file_paths.append(os.path.splitext(path)[0] if file_suffix == ".bin" else path)

317 return file_paths

318

319

320def resolve_dataset_type(spec: Mapping[str, Any] | type) -> type:

321 """Resolve an explicitly selected class without constructing a Dataset."""

322 if isinstance(spec, type):

323 return spec

324 if not isinstance(spec, Mapping) or "_target_" not in spec:

325 raise ValueError("low_level_dataset requires an explicit Dataset _target_")

326 module, _, name = spec["_target_"].rpartition(".")

327 dataset_type = getattr(import_module(module), name)

328 if not isinstance(dataset_type, type):

329 raise ValueError("low_level_dataset._target_ must identify a Dataset class")

330 return dataset_type

Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / indexed / indexed_data_config.py: 86%

178 statements