Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,804 Bytes
e6af450 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
# Copyright 2025 Bytedance Ltd. and/or its affiliates.
# SPDX-License-Identifier: Apache-2.0
import os
import xml.etree.ElementTree as ET
import subprocess
import logging
import pyarrow.fs as pf
import torch.distributed as dist
logger = logging.getLogger(__name__)
def get_parquet_data_paths(data_dir_list, num_sampled_data_paths, rank=0, world_size=1):
num_data_dirs = len(data_dir_list)
if world_size > 1:
chunk_size = (num_data_dirs + world_size - 1) // world_size
start_idx = rank * chunk_size
end_idx = min(start_idx + chunk_size, num_data_dirs)
local_data_dir_list = data_dir_list[start_idx:end_idx]
local_num_sampled_data_paths = num_sampled_data_paths[start_idx:end_idx]
else:
local_data_dir_list = data_dir_list
local_num_sampled_data_paths = num_sampled_data_paths
local_data_paths = []
for data_dir, num_data_path in zip(local_data_dir_list, local_num_sampled_data_paths):
if data_dir.startswith("hdfs://"):
files = hdfs_ls_cmd(data_dir)
data_paths_per_dir = [
file for file in files if file.endswith(".parquet")
]
else:
files = os.listdir(data_dir)
data_paths_per_dir = [
os.path.join(data_dir, name)
for name in files
if name.endswith(".parquet")
]
repeat = num_data_path // len(data_paths_per_dir)
data_paths_per_dir = data_paths_per_dir * (repeat + 1)
local_data_paths.extend(data_paths_per_dir[:num_data_path])
if world_size > 1:
gather_list = [None] * world_size
dist.all_gather_object(gather_list, local_data_paths)
combined_chunks = []
for chunk_list in gather_list:
if chunk_list is not None:
combined_chunks.extend(chunk_list)
else:
combined_chunks = local_data_paths
return combined_chunks
# NOTE: cumtomize this function for your cluster
def get_hdfs_host():
return "hdfs://xxx"
# NOTE: cumtomize this function for your cluster
def get_hdfs_block_size():
return 134217728
# NOTE: cumtomize this function for your cluster
def get_hdfs_extra_conf():
return None
def init_arrow_pf_fs(parquet_file_path):
if parquet_file_path.startswith("hdfs://"):
fs = pf.HadoopFileSystem(
host=get_hdfs_host(),
port=0,
buffer_size=get_hdfs_block_size(),
extra_conf=get_hdfs_extra_conf(),
)
else:
fs = pf.LocalFileSystem()
return fs
def hdfs_ls_cmd(dir):
result = subprocess.run(["hdfs", "dfs", "ls", dir], capture_output=True, text=True).stdout
return ['hdfs://' + i.split('hdfs://')[-1].strip() for i in result.split('\n') if 'hdfs://' in i]
|