Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / parallel / dataloader_parallel.py: 34%
88 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 Every CP coordinate uses its TP rank zero as the DataLoader source rank.
47 """
49 build_on_rank: Callable[[], bool] = _always_build
50 build_cache_on_rank: Callable[[], bool] = _always_build
51 barrier: Callable[[], None] = _no_barrier
52 distributed_enabled: bool = False
53 data_index_cache: bool = False
54 dp_rank: int = 0
55 dp_world_size: int = 1
56 tp_rank: int = 0
57 tp_world_size: int = 1
58 tp_group: Any = None
59 cp_rank: int = 0
60 cp_world_size: int = 1
61 cp_mesh: Any = None
64def _is_dataloader_rank(mesh_context: Any) -> bool:
65 """Return whether this CP coordinate owns its DataLoader input."""
66 tp_rank = int(getattr(mesh_context, "tp_rank", 0))
67 owns_dataloader = tp_rank == 0
68 return owns_dataloader
71def _is_global_cache_builder_rank(shared_storage: bool) -> bool:
72 """Select global rank zero as writer when Dataset caches are shared."""
73 if not shared_storage:
74 return True
76 try:
77 global_rank = int(dist.get_rank())
78 except (RuntimeError, ValueError):
79 global_rank = 0
80 builds_shared_cache = global_rank == 0
81 return builds_shared_cache
84def create_dataloader_parallel_context(
85 mesh_context: Any,
86 *,
87 data_index_cache: bool = False,
88 shared_storage: bool = True,
89 barrier: Callable[[], None] | None = None,
90) -> DataLoaderParallelContext:
91 """Create DataLoader ownership and synchronization policy from the Trainer mesh.
93 Args:
94 mesh_context: Trainer mesh state that provides TP, CP, and DP ranks.
95 data_index_cache: Whether every Dataset rank may consume an existing index cache.
96 shared_storage: Whether Dataset index caches are visible to every process.
97 barrier: Optional Dataset-specific long-wait synchronization callback.
99 Returns:
100 DataLoader rank ownership, cache ownership, and synchronization callbacks.
101 """
102 device_mesh = getattr(mesh_context, "device_mesh", None)
103 try:
104 world_size = int(dist.get_world_size())
105 except (RuntimeError, ValueError):
106 world_size = 1
107 distributed_enabled = device_mesh is not None and world_size > 1
108 if not distributed_enabled:
109 dataloader_context = DataLoaderParallelContext(data_index_cache=data_index_cache)
110 logger.debug("Created standalone DataLoader context: data_index_cache=%s", data_index_cache)
111 return dataloader_context
113 build_on_rank = partial(_is_dataloader_rank, mesh_context)
114 build_cache_on_rank = partial(_is_global_cache_builder_rank, shared_storage)
115 mesh_dim_names = tuple(getattr(device_mesh, "mesh_dim_names", ()) or ())
116 tp_mesh = device_mesh["tp"] if "tp" in mesh_dim_names else None
117 cp_mesh = device_mesh["cp"] if "cp" in mesh_dim_names else None
118 dataloader_context = DataLoaderParallelContext(
119 build_on_rank=build_on_rank,
120 build_cache_on_rank=build_cache_on_rank,
121 barrier=barrier or dist.barrier,
122 distributed_enabled=True,
123 data_index_cache=data_index_cache,
124 dp_rank=int(getattr(mesh_context, "dp_rank", 0)),
125 dp_world_size=int(getattr(mesh_context, "dp_size", 1)),
126 tp_rank=int(getattr(mesh_context, "tp_rank", 0)),
127 tp_world_size=int(getattr(mesh_context, "tp_size", 1)),
128 tp_group=tp_mesh.get_group() if tp_mesh is not None else None,
129 cp_rank=int(getattr(mesh_context, "cp_rank", 0)),
130 cp_world_size=int(getattr(mesh_context, "cp_size", 1)),
131 cp_mesh=cp_mesh,
132 )
133 logger.debug(
134 "Created distributed DataLoader context: dp_rank=%d, dp_world_size=%d, "
135 "tp_rank=%d, tp_world_size=%d, cp_rank=%d, cp_world_size=%d, "
136 "data_index_cache=%s, shared_storage=%s",
137 dataloader_context.dp_rank,
138 dataloader_context.dp_world_size,
139 dataloader_context.tp_rank,
140 dataloader_context.tp_world_size,
141 dataloader_context.cp_rank,
142 dataloader_context.cp_world_size,
143 data_index_cache,
144 shared_storage,
145 )
146 return dataloader_context
149def split_iterable_dataset_by_dp(
150 dataset: Any,
151 dataloader_context: DataLoaderParallelContext,
152) -> Any:
153 """Split a Hugging Face iterable Dataset across data-parallel ranks."""
154 if dataloader_context.dp_world_size <= 1:
155 split_dataset = dataset
156 return split_dataset
158 from datasets.distributed import split_dataset_by_node # pylint: disable=C0415
160 split_dataset = split_dataset_by_node(
161 dataset,
162 rank=dataloader_context.dp_rank,
163 world_size=dataloader_context.dp_world_size,
164 )
165 return split_dataset
168def build_dataset_for_dataloader(
169 dataset_factory: Callable[[], Any],
170 dataloader_context: DataLoaderParallelContext,
171 *,
172 barrier_needed: bool,
173) -> Any | None:
174 """Build a Dataset on the ranks that own DataLoader input.
176 Args:
177 dataset_factory: Deferred Dataset construction callable.
178 dataloader_context: DataLoader ownership and cache synchronization policy.
179 barrier_needed: Whether cache writers and readers must synchronize.
181 Returns:
182 Dataset used by the local DataLoader, or ``None`` on non-owning ranks.
184 Note:
185 With shared storage, the cache-builder rank calls ``dataset_factory``
186 first. Other DataLoader ranks wait for the cache and then reopen it.
187 """
188 if not dataloader_context.distributed_enabled:
189 local_dataset = dataset_factory()
190 return local_dataset
192 local_dataset = None
193 builds_cache_first = dataloader_context.build_cache_on_rank()
194 owns_dataset = dataloader_context.data_index_cache or dataloader_context.build_on_rank()
195 logger.debug(
196 "DataLoader Dataset synchronization: owns_dataset=%s, builds_cache_first=%s",
197 owns_dataset,
198 builds_cache_first,
199 enabled=owns_dataset,
200 )
202 if builds_cache_first and owns_dataset:
203 local_dataset = dataset_factory()
205 if barrier_needed:
206 try:
207 dataloader_context.barrier()
208 finally:
209 # Dataset-specific barriers may own a temporary auxiliary group.
210 # Close it after all ranks have crossed the barrier so repeated
211 # dataset builds do not retain process groups and file handles.
212 barrier = dataloader_context.barrier
213 barrier_owner = getattr(barrier, "__self__", None)
214 if barrier_owner is None and hasattr(barrier, "close"):
215 barrier_owner = barrier
216 close = getattr(barrier_owner, "close", None)
217 if close is not None:
218 close()
220 if not builds_cache_first and owns_dataset:
221 local_dataset = dataset_factory()
223 return local_dataset