(
self, path, process_fn, seed, *, meta_names=[], nshards=sys.maxsize, shuffle_buffer=1000, include_dirs=None
)
| 250 | """ |
| 251 | |
| 252 | def __init__( |
| 253 | self, path, process_fn, seed, *, meta_names=[], nshards=sys.maxsize, shuffle_buffer=1000, include_dirs=None |
| 254 | ): |
| 255 | # os.environ['WDS_SHOW_SEED'] = '1' |
| 256 | import torch |
| 257 | |
| 258 | if torch.distributed.get_rank() == 0: |
| 259 | if include_dirs is not None: # /webdatasets/A,/webdatasets/C |
| 260 | other_paths = [] |
| 261 | include_dirs = include_dirs.split(",") |
| 262 | for include_dir in include_dirs: |
| 263 | if "*" in include_dir: |
| 264 | include_dir, n = include_dir.split("*") |
| 265 | n = int(n) |
| 266 | else: |
| 267 | n = 1 |
| 268 | for cur_dir, dirs, files in os.walk(include_dir): |
| 269 | for f in files: |
| 270 | if f.endswith("tar") and os.path.getsize(os.path.join(cur_dir, f)) > 0: |
| 271 | # other_paths.append(os.path.join(cur_dir,f)) |
| 272 | other_paths.extend([os.path.join(cur_dir, f)] * n) |
| 273 | # print(f'Adding dataset paths {",".join(other_paths)}') |
| 274 | from braceexpand import braceexpand |
| 275 | |
| 276 | if len(path) > 0: # not "" |
| 277 | path = list(braceexpand(path)) + other_paths |
| 278 | else: |
| 279 | path = other_paths |
| 280 | path = [path] |
| 281 | else: |
| 282 | path = [ |
| 283 | None, |
| 284 | ] |
| 285 | torch.distributed.broadcast_object_list(path, src=0) |
| 286 | path = path[0] |
| 287 | |
| 288 | tarfile_samples = partial(tarfile_samples_with_meta, meta_names=meta_names) |
| 289 | tarfile_to_samples = pipelinefilter(tarfile_samples) |
| 290 | |
| 291 | # if model parallel, shuffle_buffer should be 1 to disable shuffling |
| 292 | try: |
| 293 | from sat.mpu import get_model_parallel_world_size |
| 294 | |
| 295 | if get_model_parallel_world_size() > 1: |
| 296 | shuffle_buffer = 1 |
| 297 | except Exception: |
| 298 | pass |
| 299 | |
| 300 | super().__init__( |
| 301 | ConfiguredResampledShards(path, seed, nshards=nshards), |
| 302 | tarfile_to_samples(), |
| 303 | wds.shuffle(shuffle_buffer), |
| 304 | process_fn, |
| 305 | ) |
| 306 | |
| 307 | |
| 308 | # rclone support |
nothing calls this directly
no test coverage detected