o 2wim@s(ddlZdddZdddZddZdS) NcCsd}d}d}d}dtjvr!dtjvr!ttjd}ttjd}n/z%ddl}|jrE|jrE|p6|jjj}|jj |d}|jj |d}Wn t yOYnwdtjvridtjvrittjd}ttjd}n zddl }|j j}|dur~|j}|j}Wn t yYnw||||fS) z Return node and worker info for PyTorch and some distributed environments. This function is copied from WebDataset: https://github.com/webdataset/webdataset rRANK WORLD_SIZENgroupWORKER NUM_WORKERS)osenvironinttorch.distributed distributed is_availableis_initializedrWORLDget_rankget_world_sizeModuleNotFoundErrortorch.utils.datautilsdataget_worker_infoid num_workers)rrank world_sizeworkerrtorch worker_inforV/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/lhotse/shar/readers/utils.pypytorch_worker_infos>    r!cCs t|d\}}}}||d|S)u This function is copied from WebDataset: https://github.com/webdataset/webdataset and adapted to lists. rNr!)srcrrrrrrrr split_by_node+sr%cCst\}}}}||d|S)r"Nr#)r$rrrrrrr split_by_worker4sr&)N)r r!r%r&rrrr s  '