Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / parallel / dataloader_parallel.py: 75%
92 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
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"""DataLoader rank ownership, Dataset cache synchronization, and DP topology."""
17from __future__ import annotations
19from collections.abc import Callable
20from dataclasses import dataclass
21from functools import partial
22from typing import Any
24import torch.distributed as dist
26from hyper_parallel.data.dataset_logging import get_dataset_logger
28logger = get_dataset_logger(__name__)
31def _always_build() -> bool:
32 """Enable Dataset construction in a standalone process."""
33 return True
36def _no_barrier() -> None:
37 """Provide a no-op barrier for standalone Dataset construction."""
40@dataclass(frozen=True)
41class DataLoaderParallelContext:
42 """Describe DataLoader ownership and its DP/TP/CP topology.
44 ``dp_rank`` and ``dp_world_size`` are consumed only by online iterable
45 sources. Mapping sources leave DP sample ownership to their batch sampler.
46 Native loading reads on ``tp_source_rank`` at every CP coordinate.
47 With ``tp_source_rank=None``, the loader distributes batches itself and
48 every model rank advances it; no additional TP broadcast is needed.
49 """
51 build_on_rank: Callable[[], bool] = _always_build
52 build_cache_on_rank: Callable[[], bool] = _always_build
53 barrier: Callable[[], None] = _no_barrier
54 distributed_enabled: bool = False
55 data_index_cache: bool = False
56 tp_source_rank: int | None = 0
57 dp_rank: int = 0
58 dp_world_size: int = 1
59 tp_rank: int = 0
60 tp_world_size: int = 1
61 tp_group: Any = None
62 cp_rank: int = 0
63 cp_world_size: int = 1
64 cp_mesh: Any = None
67def _is_dataloader_rank(mesh_context: Any, source_rank: int) -> bool:
68 """Return whether this CP coordinate owns its DataLoader input."""
69 tp_rank = int(getattr(mesh_context, "tp_rank", 0))
70 owns_dataloader = tp_rank == source_rank
71 return owns_dataloader
74def _is_global_cache_builder_rank(shared_storage: bool) -> bool:
75 """Select global rank zero as writer when Dataset caches are shared."""
76 if not shared_storage:
77 return True
79 try:
80 global_rank = int(dist.get_rank())
81 except (RuntimeError, ValueError):
82 global_rank = 0
83 builds_shared_cache = global_rank == 0
84 return builds_shared_cache
87def create_dataloader_parallel_context(
88 mesh_context: Any,
89 *,
90 data_index_cache: bool = False,
91 tp_source_rank: int | None = 0,
92 shared_storage: bool = True,
93 barrier: Callable[[], None] | None = None,
94) -> DataLoaderParallelContext:
95 """Create DataLoader ownership and synchronization policy from the Trainer mesh.
97 Args:
98 mesh_context: Trainer mesh state that provides TP, CP, and DP ranks.
99 data_index_cache: Whether every Dataset rank may consume an existing index cache.
100 tp_source_rank: TP-local source rank, or ``None`` when the loader already
101 distributes batches to all model ranks.
102 shared_storage: Whether Dataset index caches are visible to every process.
103 barrier: Optional Dataset-specific long-wait synchronization callback.
105 Returns:
106 DataLoader rank ownership, cache ownership, and synchronization callbacks.
107 """
108 device_mesh = getattr(mesh_context, "device_mesh", None)
109 try:
110 world_size = int(dist.get_world_size())
111 except (RuntimeError, ValueError):
112 world_size = 1
113 distributed_enabled = device_mesh is not None and world_size > 1
114 if not distributed_enabled:
115 dataloader_context = DataLoaderParallelContext(
116 data_index_cache=data_index_cache,
117 tp_source_rank=tp_source_rank,
118 )
119 logger.debug("Created standalone DataLoader context: data_index_cache=%s", data_index_cache)
120 return dataloader_context
122 tp_size = int(getattr(mesh_context, "tp_size", 1))
123 if tp_source_rank is not None and not 0 <= tp_source_rank < tp_size:
124 raise ValueError("tp_source_rank must be a rank within the TP group or None")
125 build_on_rank = _always_build if tp_source_rank is None else partial(
126 _is_dataloader_rank, mesh_context, tp_source_rank,
127 )
128 build_cache_on_rank = partial(_is_global_cache_builder_rank, shared_storage)
129 mesh_dim_names = tuple(getattr(device_mesh, "mesh_dim_names", ()) or ())
130 tp_mesh = device_mesh["tp"] if "tp" in mesh_dim_names else None
131 cp_mesh = device_mesh["cp"] if "cp" in mesh_dim_names else None
132 dataloader_context = DataLoaderParallelContext(
133 build_on_rank=build_on_rank,
134 build_cache_on_rank=build_cache_on_rank,
135 barrier=barrier or dist.barrier,
136 distributed_enabled=True,
137 data_index_cache=data_index_cache,
138 tp_source_rank=tp_source_rank,
139 dp_rank=int(getattr(mesh_context, "dp_rank", 0)),
140 dp_world_size=int(getattr(mesh_context, "dp_size", 1)),
141 tp_rank=int(getattr(mesh_context, "tp_rank", 0)),
142 tp_world_size=int(getattr(mesh_context, "tp_size", 1)),
143 tp_group=tp_mesh.get_group() if tp_mesh is not None else None,
144 cp_rank=int(getattr(mesh_context, "cp_rank", 0)),
145 cp_world_size=int(getattr(mesh_context, "cp_size", 1)),
146 cp_mesh=cp_mesh,
147 )
148 logger.debug(
149 "Created distributed DataLoader context: dp_rank=%d, dp_world_size=%d, "
150 "tp_rank=%d, tp_world_size=%d, cp_rank=%d, cp_world_size=%d, "
151 "data_index_cache=%s, shared_storage=%s",
152 dataloader_context.dp_rank,
153 dataloader_context.dp_world_size,
154 dataloader_context.tp_rank,
155 dataloader_context.tp_world_size,
156 dataloader_context.cp_rank,
157 dataloader_context.cp_world_size,
158 data_index_cache,
159 shared_storage,
160 )
161 return dataloader_context
164def split_iterable_dataset_by_dp(
165 dataset: Any,
166 dataloader_context: DataLoaderParallelContext,
167) -> Any:
168 """Split a Hugging Face iterable Dataset across data-parallel ranks."""
169 if dataloader_context.dp_world_size <= 1:
170 split_dataset = dataset
171 return split_dataset
173 from datasets.distributed import split_dataset_by_node # pylint: disable=C0415
175 split_dataset = split_dataset_by_node(
176 dataset,
177 rank=dataloader_context.dp_rank,
178 world_size=dataloader_context.dp_world_size,
179 )
180 return split_dataset
183def build_dataset_for_dataloader(
184 dataset_factory: Callable[[], Any],
185 dataloader_context: DataLoaderParallelContext,
186 *,
187 barrier_needed: bool,
188) -> Any | None:
189 """Build a Dataset on the ranks that own DataLoader input.
191 Args:
192 dataset_factory: Deferred Dataset construction callable.
193 dataloader_context: DataLoader ownership and cache synchronization policy.
194 barrier_needed: Whether cache writers and readers must synchronize.
196 Returns:
197 Dataset used by the local DataLoader, or ``None`` on non-owning ranks.
199 Note:
200 With shared storage, the cache-builder rank calls ``dataset_factory``
201 first. Other DataLoader ranks wait for the cache and then reopen it.
202 """
203 if not dataloader_context.distributed_enabled:
204 local_dataset = dataset_factory()
205 return local_dataset
207 local_dataset = None
208 builds_cache_first = dataloader_context.build_cache_on_rank()
209 owns_dataset = (
210 dataloader_context.data_index_cache
211 or dataloader_context.build_on_rank()
212 )
213 logger.debug(
214 "DataLoader Dataset synchronization: owns_dataset=%s, builds_cache_first=%s",
215 owns_dataset,
216 builds_cache_first,
217 enabled=owns_dataset,
218 )
220 if builds_cache_first and owns_dataset:
221 local_dataset = dataset_factory()
223 if barrier_needed:
224 try:
225 dataloader_context.barrier()
226 finally:
227 # Dataset-specific barriers may own a temporary auxiliary group.
228 # Close it after all ranks have crossed the barrier so repeated
229 # dataset builds do not retain process groups and file handles.
230 barrier = dataloader_context.barrier
231 barrier_owner = getattr(barrier, "__self__", None)
232 if barrier_owner is None and hasattr(barrier, "close"):
233 barrier_owner = barrier
234 close = getattr(barrier_owner, "close", None)
235 if close is not None:
236 close()
238 if not builds_cache_first and owns_dataset:
239 local_dataset = dataset_factory()
241 return local_dataset