import os import random import io import datetime from glob import glob from pathlib import Path import s3fs from time import perf_counter import boto3 from boto3.s3.transfer import TransferConfig import logging logging.getLogger('botocore.credentials').setLevel(logging.WARNING) import time import torch from torch.distributed import init_process_group, destroy_process_group, is_initialized, barrier import uuid import random import subprocess import argparse import copy from typing import Dict, Iterable, Tuple import importlib from anydata.utils.read import read_config from anydata.utils.types import is_str, is_tensor from anydata.utils.data import flatten from anydata.dataloaders.instantiate import augmentations from anydata.augmentations.tensor import to_tensor_sample from anydata.webdataset.misc.utils_legacy import prep_buffer, save_data_single, get_name, make_cfg_name, unwrap_cfg import numpy as np import multiprocessing from tqdm import tqdm from webdataset import TarWriter import torchvision.transforms as T import traceback SAGEMAKER = os.environ.get('SAGEMAKER') == 'enabled' import shutil import json def have_s5cmd(): return shutil.which("s5cmd") is not None def run_cp(local_src, s3_dst, acl=True, quiet=True): """Upload local→s3, using s5cmd if enabled, else boto3/aws cli""" if args.s5cmd and have_s5cmd(): # print("using s5cmd to copy ", local_src, " to ", s3_dst) cmd = ["s5cmd", "cp"] if acl: cmd += ["--acl", "bucket-owner-full-control"] cmd += [local_src, s3_dst] subprocess.run(cmd, check=True) else: # boto3 path (keeps your old behavior) s3_bucket = s3_dst.strip('s3://').split('/')[0] s3_key = "/".join(s3_dst.strip('s3://').split('/')[1:]) boto3.client('s3').upload_file( Filename=local_src, Bucket=s3_bucket, Key=s3_key, ExtraArgs={'ACL': 'bucket-owner-full-control'} if acl else {}, Config=TransferConfig(max_concurrency=20) ) def run_cp_recursive(src, dst, acl=False, concurrency=64): """Recursive copy src→dst (downloads/uploads)""" if args.s5cmd and have_s5cmd(): print("using s5cmd to copy recursively ", src, " to ", dst) if src.startswith("s3://") and not src.endswith("**"): if not src.endswith("/"): src += "/" src += "**" elif not src.startswith("s3://") and not src.endswith("/"): src += "/**" cmd = ["s5cmd", "cp", f"--concurrency={concurrency}"] if acl: cmd += ["--acl", "bucket-owner-full-control"] cmd += [src, dst] subprocess.run(cmd, check=True) else: cmd = ["aws", "s3", "cp", "--recursive", src, dst] if acl: cmd += ["--acl", "bucket-owner-full-control"] subprocess.run(cmd, check=True) def write_empty_txt(filename, name, folder, message=None): folder = f'{folder}/{filename.replace("/", "__")}' os.makedirs(folder, exist_ok=True) if name is None: folder = f'{folder}.txt' else: folder = f'{folder}/{name}.txt' with open(folder, 'w') as f: if message is not None: f.write(f'{message}\n') def arg_parser(): parser = argparse.ArgumentParser() parser.add_argument('--datasets', type=str, nargs='+') parser.add_argument('--split', type=str, default=None) parser.add_argument("--split_json", type=str, default=None) parser.add_argument('--multi_bar', type=lambda x: x.lower() == 'true', default=False) parser.add_argument('--src_path', type=str, default='/data/cv_unified', help='Source path of data to be processed; can be a local/S3 folder') parser.add_argument('--start_idx', type=int, default=None, help="Start index of all task suffixes to process") parser.add_argument('--end_idx', type=int, default=None, help="End index of all task suffixes to process") parser.add_argument('--local_path', type=str, default='/data/cv_webdatasets', help="Where the processed data will be stored locally") parser.add_argument('--s3_path', type=str, default='s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_datasets/webdatasets', help="Where processed data will be uploaded & stored on S3") parser.add_argument('--s5cmd', type=lambda x: x.lower() == 'true', default=False, help="Use s5cmd for uploads/downloads if installed, else fall back to aws/boto3") parser.add_argument('--producer', type=lambda x: x.lower() == 'true', default=False, help="Stream from S3 with producer/consumer (True) or pre-download dataset to /tmp like DGX (False)") parser.add_argument('--delete', type=lambda x: x.lower() == 'true', default=False) parser.add_argument('--merge', type=lambda x: x.lower() == 'true', default=False) parser.add_argument('--seed', type=int, default=42) parser.add_argument('--fill_depth', type=lambda x: x.lower() == 'true', default=False) parser.add_argument('--resize', type=int, nargs='+', default=None) parser.add_argument('--resolution', type=int, nargs='+', default=None) parser.add_argument('--cameras', type=str, nargs='+', default=None) parser.add_argument('--cameras_context_sample', type=int, default=None) parser.add_argument('--context', type=int, nargs='+', default=None) parser.add_argument('--context_sample', type=str, default=None) parser.add_argument('--repeat', type=int, default=None) parser.add_argument('--no_resize_supervision', type=lambda x: x.lower() == 'true', default=False) parser.add_argument('--min_context', type=int, default=None) parser.add_argument('--buffer', type=str, default=None) parser.add_argument('--suffix', type=str, default=None) parser.add_argument('--zone', type=str, default='west') parser.add_argument('--cfg_path', type=str, default=f'scripts/datasets/webdataset/configs') parser.add_argument('--multi_sample', type=int, default=None) parser.add_argument("--local-rank", "--local_rank", type=int, default=0) parser.add_argument('--parallelize', type=str, default='sequence', choices=['task','sequence']) parser.add_argument('--source', type=str, default='lbm', choices=['lbm','custom']) parser.add_argument('--subset', type=str, default=None) parser.add_argument('--mode', type=str, default='latents') parser.add_argument('--data_path', type=str, default=None) parser.add_argument('--filter', type=str, default='') parser.add_argument('--check_resolution', type=str, default=None) parser.add_argument('--resize_ratio', type=str, default='') # Optional: pre-download unified dataset using repo helper before processing parser.add_argument("--download_data", type=str, default=None, help="If set, run anydata/sync/download_uni.py before processing. " "Accepts a split json path or a dataset folder (e.g. spartan_tiny or spartan_tiny/split_all.json).") parser.add_argument("--download_local_bucket", type=str, default=os.environ.get('ANYDATA_LOCAL_ROOT', '/data'), help="Local root for download_uni.py (e.g. /tmp or /data)") parser.add_argument("--download_s3_bucket", type=str, default="s3://tri-ml-sandbox-16011-us-west-2-datasets", help="S3 bucket root used by download_uni.py") parser.add_argument("--download_data_folder", type=str, default="cv_unified", help="Data folder under bucket used by download_uni.py (e.g. cv_unified)") parser.add_argument("--download_num_procs", type=int, default=16, help="Number of processes for download_uni.py") parser.add_argument("--input_split", type=str, default="split_all.json", help="Split json filename to use when download_data is a folder (default: split_all.json)") args = parser.parse_args() args = prep_buffer(args) if args.check_resolution == '': args.check_resolution = None if args.check_resolution is not None: args.check_resolution = [int(v) for v in args.check_resolution.split('_')] args.local_path += f'/{args.mode}' args.s3_path += f'/{args.mode}' if args.parallelize == 'sequence': args.merge = True # args.src_path = args.src_path.replace('fm_datasets', f'{args.data_path}_datasets') dataset = args.datasets[0].split('/')[0] args.source = 'custom' for d in ['LBM']: if d in dataset: args.source = 'lbm' # args.src_path = f"{args.src_path}/{dataset}" # for d in ['LBM','RoboCasa']: # if d in dataset: # args.src_path = args.src_path.replace('fm_datasets', 'cv_datasets') return args def upload_single(local_path_i, s3_path_i, sample): # NOTE (Dian): create file handle inside the subprocess to avoid undefined/corrupted states try: file_i = TarWriter(local_path_i, compress=False, encoder=True, keep_meta=True) file_i.write(sample) file_i.close() except Exception as e: print(f"TarWriter failed for {local_path_i}: {repr(e)}") try: for k, v in sample.items(): print(f" key={k} type={type(v)}") except Exception: pass raise uploaded = True try: if args.s5cmd and have_s5cmd(): print("using s5cmd to upload ", local_path_i, " to ", s3_path_i) cmd = ["s5cmd", "cp", "--acl", "bucket-owner-full-control", local_path_i, s3_path_i] subprocess.run(cmd, check=True) else: print("using boto3 to upload ", local_path_i, " to ", s3_path_i) s3_bucket = s3_path_i.strip('s3://').split('/')[0] s3_key = "/".join(s3_path_i.strip('s3://').split('/')[1:]) s3_bucket = s3_bucket.replace('//', '/') s3_key = s3_key.replace('//', '/') s3 = boto3.client('s3') s3.upload_file( Filename=local_path_i, Bucket=s3_bucket, Key=s3_key, ExtraArgs={'ACL': 'bucket-owner-full-control'}, Config=TransferConfig(max_concurrency=20) ) print("upload with s3 complete ", local_path_i, " to ", s3_path_i) except subprocess.CalledProcessError as e: uploaded = False print(f"Upload failed: {local_path_i} -> {s3_path_i} ({repr(e)})") raise except Exception as e: uploaded = False print(f"Upload failed: {local_path_i} -> {s3_path_i} ({repr(e)})") raise if uploaded: print(f"Removing local temp file {local_path_i}") os.remove(local_path_i) def save_data(data, n, prefix, file=None, repeat_idx=None, tokenizers=None, error_data=None, dataset_tag=None, dataset_idx=None, sequence_name=None): if dataset_tag is None: dataset_tag = data.get("tag", None) if dataset_tag is None and "metadata" in data: dataset_tag = data["metadata"].get("tag", "Unified") if dataset_tag is None: dataset_tag = "Unified" if dataset_idx is None: dataset_idx = data.get("idx", None) if dataset_idx is None and "metadata" in data: dataset_idx = data["metadata"].get("idx", 0) prefix = f'{prefix.split("/")[-1]}/{n:012}_{int(dataset_idx):012}_{dataset_tag}' if 'extrinsics' in data and 'pose' not in data: data['pose'] = data['extrinsics'] def filter_cam(data_in, cam_id): out = {} for key, val in data_in.items(): if isinstance(val, dict): # Keep shared language tokens unfiltered if key in ['language', 'lat_language']: out[key] = val continue if len(val) == 0: out[key] = val continue first_key = next(iter(val.keys())) if not isinstance(first_key, tuple): out[key] = val else: out[key] = {k: v for k, v in val.items() if isinstance(k, tuple) and len(k) == 2 and k[1] == cam_id} else: out[key] = val return out cam_ids = sorted(list(set(key[1] for key in data['rgb'].keys()))) # Prefer metadata camera names from Unified samples when available meta_cams = None if isinstance(data.get("metadata", None), dict): meta_cams = data["metadata"].get("cameras", None) if meta_cams and len(meta_cams) > 0: cam_labels = {cam: (meta_cams[cam] if cam < len(meta_cams) else f"cam{cam}") for cam in cam_ids} else: # Map numeric ids to names if provided (e.g., args.cameras = ['front', 'rear', ...]) cam_labels = {cam: (args.cameras[cam] if args.cameras and cam < len(args.cameras) else f"cam{cam}") for cam in cam_ids} print(f"cameras labels are: {cam_labels}") # No buffering: write one entry per cam into the same tar writer. if not args.buffer: size_total = 0 for cam in cam_ids: data_cam = filter_cam(data, cam) valid, sample, _ = save_data_single( args, data_cam, n, f"{prefix}_{cam_labels[cam]}", file, store_raw=args.mode == 'frames' or args.mode == 'both', store_latents=args.mode == 'latents' or args.mode == 'both', ) if valid: file.write(sample[1]) size_total += 1 return size_total time = sorted(list(set(key[0] for key in data['rgb'].keys()))) cam = cam_ids total_length = len(time) max_length = args.buffer overlap = args.buffer_overlap stride = args.buffer_stride pad_end = args.buffer_pad fixed_length = args.buffer_fixed if not pad_end and total_length <= max_length: return 0 # if not pad_end and max_length > total_length: # max_length = total_length if stride is not None: intervals = [[0, max_length]] while intervals[-1][0] < total_length - stride[1]: skip = random.randint(*stride[0]) if isinstance(stride[0], list) else stride[0] next_interval = intervals[-1][0] + skip intervals.append([next_interval, next_interval + max_length]) else: available = total_length - max_length skip = int(np.round(available / overlap)) + 1 interval = np.linspace(0, available, skip) interval = [int(i) for i in interval] intervals = [[0, max_length]] for i in range(1, len(interval)): intervals.append([interval[i], interval[i] + max_length]) if len(intervals) > 1: for i in range(0, len(interval) - 1): intervals[i][1] += (intervals[i+1][0] - intervals[i][0]) intervals = intervals[:-1] else: intervals[0][1] = total_length if args.buffer_random is not None: intervals = random.sample(intervals, args.buffer_random) intervals.sort(key=lambda x: x[0]) local_path = file # os.path.dirname(local_path) s3_path = local_path.replace(args.local_path, args.s3_path) + '/tarfiles' s3_path = s3_path.replace('__', '/') print(f"local_path: {local_path}, s3_path: {s3_path}") filename = _trim_sequence_name(sequence_name) print(f"filename after trimming: {filename}") if repeat_idx is not None: filename += '___%02d' % repeat_idx # if pad_end: # n = n - max_length // 2 # else: # n = n - max_length s3_path = s3_path.replace('//tarfiles', '/tarfiles') expected = len(intervals) cam_skip = {} for cam_id in cam_ids: s3_path_cam = f'{s3_path}/{cam_labels[cam_id]}/{filename}' proc = subprocess.Popen(f'aws s3 ls {s3_path_cam}/', stdout=subprocess.PIPE, shell=True) folders = proc.communicate()[0].decode().replace('\n', ' ').split(' ') folders = [f for f in folders if f.endswith('.tar') or f.endswith('/')] proc.kill() existing = len([f for f in folders if f.endswith('.tar')]) if existing == expected: cam_skip[cam_id] = True else: cam_skip[cam_id] = False if len(folders) > 0: print(f'!!!! INCOMPLETE {s3_path_cam}') os.makedirs(local_path, exist_ok=True) # print(intervals, len(folders), n) lat_language = {} print(f"n is {n}, total intervals: {len(intervals)}, for file: {filename}") # NOTE (Dian): use pool manager to avoid I/O handles exceeding system limit upload_list = [] with multiprocessing.Pool(processes=64) as pool: progress = range(len(intervals)) # progress = tqdm(range(n), ncols=128) if global_rank == 0 else range(n) for i in progress: # for i in range(n): time_st, time_fn = intervals[i] for cam_id in cam_ids: if cam_skip.get(cam_id, False): continue # name_i = f'seq-%06d-{time_st}_{time_fn}_{cam_labels[cam_id]}.tar' % i name_i = f'seq-%06d-%d_%d.tar' % (i, time_st, time_fn) temp_id = uuid.uuid4().hex local_path_i = f'/tmp/wds_{temp_id}.tar' # local_path_i = f'{local_path}/{cam_labels[cam_id]}/{filename}/{name_i}' s3_path_i = f'{s3_path}/{cam_labels[cam_id]}/{filename}/{name_i}' print(f"cam_label: {cam_labels[cam_id]}, filename: {filename}, name_i: {name_i} local_path_i: {local_path_i} s3_path_i: {s3_path_i}") os.makedirs(os.path.dirname(local_path_i), exist_ok=True) try: file_i = None data_cam = filter_cam(data, cam_id) size, sample, lat_language = save_data_single( args, data_cam, n, f"{prefix}_{cam_labels[cam_id]}", file_i, time_st, time_fn, total_length, tokenizers, lat_language, store_raw=args.mode == 'frames' or args.mode == 'both', store_latents=args.mode == 'latents' or args.mode == 'both', ) sample = sample[1] except Exception as e: print(f'!!!! ERROR {local_path_i}', e) traceback.print_exc() s3_path_error, dataset_error, idx_error = error_data split = s3_path_error.split('/') # TODO @vitor inline hardcoded path may be confusing invalid_folder = f'/data/cv_datasets/invalids/{split[-2]}/{split[-1]}' filename = dataset_error.get_filename(idx_error) write_empty_txt(filename, name_i, invalid_folder, e) size = 0 if size > 0: # NOTE (Dian): use another process to upload to unblock the following processing print("uploading ", local_path_i, " to ", s3_path_i) result = pool.apply_async(upload_single, args=(local_path_i, s3_path_i, sample)) upload_list.append(result) # if global_rank == 0: # start = perf_counter() # print(f"Reached progress end ---------------------------") if len(upload_list) > 0: for res in upload_list: try: res.get() # raises if upload failed except Exception as e: print(f"UPLOAD ERROR: {repr(e)}") # if global_rank == 0: # end = perf_counter() # print(f"Main process took {end - start}s to join subprocesses") create_indices(s3_path) return 0 def process_batch(dataset, n, splits, indices, shuffled_indices, prefix, s3_path, files, repeat_idx, tokenizers): print(f'Processing batch {n} with {len(files)} files') progress1 = range(len(files)) global_rank = int(os.environ.get('RANK', '0')) for i in progress1: progress2 = list(range(splits[i], splits[i+1])) if global_rank == 0: progress2 = tqdm(progress2, ncols=128, leave=False) size = 0 n_tries = 2 for j in progress2: for attempt in range(n_tries): try: dataset_idx = shuffled_indices[j] data = dataset[dataset_idx] sequence_name = dataset.get_filename(dataset_idx) size += save_data( data, indices[j], prefix, files[i], repeat_idx, tokenizers, (s3_path, dataset, dataset_idx), dataset_tag=getattr(dataset, "tag", None), dataset_idx=dataset_idx, sequence_name=sequence_name, ) break except Exception as e: split = s3_path.split('/') # TODO @vitor inline hardcoded path may be confusing invalid_folder = f'/data/cv_datasets/invalids/{split[-2]}/{split[-1]}' filename = dataset.get_filename(shuffled_indices[j]) print('!!!! ERROR') print(f"Reason: {repr(e)}") traceback.print_exc() write_empty_txt(filename, None, invalid_folder, e) if not args.buffer: name_base = files[i].own_fileobj.name files[i].close() name = name_base.replace('.tar', '_%04d.tar' % size) os.rename(name_base, name) if size > 0: uploaded = True try: if args.s5cmd and have_s5cmd(): cmd = ["s5cmd", "cp", "--acl", "bucket-owner-full-control", name, f"{s3_path}/{os.path.basename(name)}"] subprocess.run(cmd, check=True) else: result = subprocess.run( f'aws s3 cp {name} {s3_path}/{os.path.basename(name)} --acl bucket-owner-full-control --quiet', shell=True, capture_output=True, text=True, ) if result.returncode != 0: uploaded = False print(f"Upload failed: {name} -> {s3_path}\n{result.stderr.strip()}") except subprocess.CalledProcessError as e: uploaded = False print(f"Upload failed: {name} -> {s3_path} ({e})") if uploaded: os.remove(name) else: os.remove(name) def process_dataset(s3_path, dataset, local_path, repeat_idx=None, tokenizers=None): size = len(dataset) # NOTE (Dian): clean up some logic for multi-gpu processing indices = list(range(size)) shuffled_indices = [i for i in indices] if args.suffix is not None: local_path = f'{local_path}/{args.suffix.replace("/","__")}' if os.path.exists(local_path) and not args.merge: print(f'### DELETING LOCAL: {local_path}') shutil.rmtree(local_path) if args.parallelize == 'task': all_splits, all_sizes = [[0, size]], [[size]] all_sizes_flat = flatten(all_sizes) elif args.parallelize == 'sequence': local_rank = int(os.environ.get('LOCAL_RANK', '0')) local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', '1')) chunks = np.linspace(0, size, local_world_size + 1) chunks = [int(v) for v in chunks] all_splits, all_sizes = [[chunks[local_rank], chunks[local_rank+1]]], [[chunks[local_rank+1]]] all_sizes_flat = flatten(all_sizes) if args.subset is not None and args.subset != '': a, b = [int(v) for v in args.subset.split('/')] chunks = np.linspace(all_splits[0][0], all_splits[0][1], a + 1) chunks = [int(v) for v in chunks] all_splits = [[chunks[b], chunks[b + 1]]] else: raise ValueError(f'invalid parallelize {args.parallelize}') os.makedirs(local_path, exist_ok=True) if not args.buffer: tar_files = [TarWriter( f'{local_path}/%06d.tar' % i, compress=False, encoder=True, keep_meta=True, ) for i, s in enumerate(all_sizes_flat)] else: tar_files = [local_path for i, s in enumerate(all_sizes_flat)] fn = 0 for i in range(len(all_sizes)): st = fn fn = st + len(all_sizes[i]) process_batch( dataset, i, all_splits[i], indices, shuffled_indices, local_path, s3_path, tar_files[st:fn], repeat_idx, tokenizers, ) def create_indices(s3_path): s3_path = s3_path.replace('//tarfiles', '/tarfiles') if s3_path[-1] == '/': s3_path = s3_path[:-1] proc = subprocess.Popen(f'aws s3 ls {s3_path}/', stdout=subprocess.PIPE, shell=True) folders = proc.communicate()[0].decode().replace('\n', ' ').split(' ') folders = [f for f in folders if f.endswith('.tar') or f.endswith('/')] proc.kill() tars = [f for f in folders if f.endswith('tar')] if len(tars) == 0: for f in folders: create_indices(f'{s3_path}/{f[:-1]}') else: name = s3_path.split('/') idx = name.index('tarfiles') - 1 # name = '__'.join(name[idx:]) + '.txt' import hashlib full_key = '__'.join(name[idx:]) short = hashlib.md5(full_key.encode()).hexdigest()[:12] name = f'index_{short}.txt' with open(name, 'w') as file: for t in tars: file.write(f'{s3_path}/{t}\n') s3_path_indices = s3_path.replace('tarfiles', 'indices') subprocess.run( f'aws s3 cp {name} {s3_path_indices}/filenames.txt --quiet', shell=True, capture_output=True, text=True, check=True, ) os.remove(name) def create_all_indices(s3_path, all_tars): proc = subprocess.Popen(f'aws s3 ls {s3_path}/', stdout=subprocess.PIPE, shell=True) folders = proc.communicate()[0].decode().replace('\n', ' ').split(' ') folders = [f for f in folders if f.endswith('.tar') or f.endswith('/')] proc.kill() tars = [f for f in folders if f.endswith('tar')] if len(tars) == 0: for f in folders: create_all_indices(f'{s3_path}/{f[:-1]}', all_tars) else: tars = [f'{s3_path}/{t}' for t in tars] all_tars.extend(tars) def process_all_indices(s3_path_i, filename): print('####################################################### CREATING INDICES', filename) all_tars = [] create_all_indices(s3_path_i + '/tarfiles', all_tars) with open(filename, 'w') as file: for t in all_tars: file.write(f'{t}\n') subprocess.run( f'aws s3 cp {filename} {s3_path_i}/indices/filenames.txt --quiet', shell=True, capture_output=True, text=True, check=True, ) os.remove(filename) def get_all_suffixes(path, start=None, end=None): """ List all task suffixes to process under a local/S3 path """ task_suffixes = [] # if given S3 path if path.startswith('s3://'): path = path.replace('s3://', '') s3 = s3fs.S3FileSystem() folders1 = s3.glob(str(Path(path) / '*') + '/') for folder1 in folders1: folders2 = s3.glob(str(Path(folder1) / '*') + '/') for folder2 in folders2: folders3 = s3.glob(str(Path(folder2) / '*') + '/') task_suffixes.extend(folders3) # if given local path else: if args.source == 'lbm': folders1 = glob(str(Path(path) / '*') + '/') for folder1 in folders1: folders2 = glob(str(Path(folder1) / '*') + '/') for folder2 in folders2: folders3 = glob(str(Path(folder2) / '*') + '/') task_suffixes.extend(folders3) elif args.source == 'custom': task_suffixes = glob(str(Path(path)) + '/') else: raise ValueError(f'Invalid source {args.source}') for i in range(len(task_suffixes)): task_suffixes[i] = task_suffixes[i].replace(path, '').strip('/') print(f'Total number of suffixes: {len(task_suffixes)}') # ensure deterministic order task_suffixes = sorted(task_suffixes) if start is None: start = 0 if end is None: end = len(task_suffixes) task_suffixes = task_suffixes[start:end] return task_suffixes def download_producer(task_queue, src_path, task_suffixes): global_rank = int(os.environ.get('RANK', '0')) for ti, task_suffix in enumerate(task_suffixes): print(f"On global rank {global_rank}, {task_queue.qsize()} downloaded tasks in the queue") try: download_task(src_path, task_suffix) except: task_queue.put(None) break task_queue.put((ti, task_suffix)) # manually signal the end of tasks task_queue.put(None) def copy_subfolder(info): s3_path, local_path, subdir = info src = s3_path + subdir dst = os.path.join(local_path, subdir) os.makedirs(dst, exist_ok=True) try: if args.s5cmd and have_s5cmd(): # print("using s5cmd to copy ", src, " to ", dst) cmd = ["s5cmd", "cp", "--concurrency=32", src + "**", dst] else: cmd = ["aws", "s3", "sync", src, dst, "--only-show-errors"] subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True) except subprocess.CalledProcessError as e: print(f"Failed to copy {src} -> {dst}") print(f"Return code: {e.returncode}") print("STDOUT:", e.stdout) print("STDERR:", e.stderr) def download_task(s3_folder, suffix): global_rank = int(os.environ.get('RANK', '0')) local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', '1')) node_rank = global_rank // local_world_size s3_path = f"{s3_folder.strip('/')}/{suffix}/" local_path = os.path.join(LOCAL_ROOT, suffix) # print(f"Downloading {s3_path} to {local_path} on node {node_rank}, global rank {global_rank}") os.makedirs(local_path, exist_ok=True) # cmd = [ # "aws", "s3", "ls", s3_path, "|", # "awk", "'{print $2}'", "|", # "xargs", "-P60", "-I{}", # "aws", "s3", "cp", "--recursive", "--quiet", s3_path + "{}", local_path + "{}" # ] try: start = perf_counter() result = subprocess.run(["aws", "s3", "ls", s3_path], capture_output=True, text=True, check=True) # dirs = [line.split()[1] for line in result.stdout.strip().splitlines()] dirs = [] for line in result.stdout.splitlines(): line = line.strip() if line.startswith("PRE "): d = line.split()[-1] # e.g., 'lowdim/' if not d.endswith("/"): d += "/" dirs.append(d) infos = [(s3_path, local_path, dir_) for dir_ in dirs] with multiprocessing.Pool(processes=4) as pool: pool.map(copy_subfolder, infos) end = perf_counter() # print(f"Download {s3_path} finished in {end - start} seconds.") # subprocess.run(" ".join(cmd), shell=True, check=True, capture_output=True, text=True) except subprocess.CalledProcessError as e: print('Return code: ', e.returncode) print('STD output: ', e.stdout) print('Error message: ', e.stderr) def delete_task(suffix): global_rank = int(os.environ.get('RANK', '0')) local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', '1')) node_rank = global_rank // local_world_size local_path = os.path.join(LOCAL_ROOT, suffix) print(f"Deleting local data {local_path} on node {node_rank}, global rank {global_rank}") try: subprocess.run( ["rm", "-r", local_path], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) except subprocess.CalledProcessError as e: print(f"Return code: {e.returncode}") print("STDOUT:", e.stdout) print("STDERR:", e.stderr) def get_from_cfg_list(cfg, key, idx): if key not in cfg: return None data = cfg[key] if not isinstance(data, list): return data # For camera lists written as [cam0, cam1, ...], treat as a single dataset entry. if key in ['cameras', 'cameras_context'] and len(data) > 0 and not isinstance(data[0], list): return data return data[idx] if len(data) > 1 else data[0] def _trim_sequence_name(sequence_name): """ Normalize sequence identifier to old high-level structure by removing trailing rgb/camera/frame components. """ if sequence_name is None: return sequence_name if "/" in sequence_name: parts = sequence_name.split("/") else: parts = sequence_name.split("__") if len(parts) <= 3: return sequence_name return "/".join(parts[:-3]) def setup_dataset(cfg, verbose=False): cfg_dict = unwrap_cfg(cfg) if not isinstance(cfg, dict) else cfg shared_keys = ['context', 'labels', 'labels_context'] # Normalize cameras to list-of-lists so it doesn't look like multi-dataset. if 'cameras' in cfg_dict and isinstance(cfg_dict['cameras'], list) and cfg_dict['cameras']: if not isinstance(cfg_dict['cameras'][0], list): cfg_dict['cameras'] = [cfg_dict['cameras']] num_datasets = 0 for key, val in cfg_dict.items(): if key not in shared_keys and isinstance(val, list): if len(val) > 1: num_datasets = max(num_datasets, len(val)) if num_datasets == 0: num_datasets = 1 datasets, datasets_cfg = [], [] for i in range(num_datasets): cfg_args = {} for key, val in cfg_dict.items(): if key in shared_keys: cfg_args[key] = val else: cfg_args[key] = get_from_cfg_list(cfg_dict, key, i) if cfg_args.get('augmentation', None) is not None: cfg_args['data_transform'] = lambda sample, cfg=cfg_args['augmentation']: augmentations(sample, cfg=cfg) else: cfg_args['data_transform'] = to_tensor_sample name = get_from_cfg_list(cfg_dict, 'name', i) if isinstance(name, list): name = name[i] if len(name) > i else name[0] module = importlib.import_module(f'anydata.dataloaders.{name}') dataset_cls = getattr(module, name + 'Dataset', None) dataset = dataset_cls(**{k: v for k, v in cfg_args.items() if k != 'name'}) if verbose: string = f'######### {name}: {len(dataset)} samples' context = cfg_dict.get('context', []) cameras = cfg_dict.get('cameras', []) labels = cfg_dict.get('labels', []) suffix = cfg_dict.get('suffix', []) stride = cfg_dict.get('context_stride', []) string += f' | context {context}'.replace(', ', ',') if stride: string += f' | stride {stride}'.replace(', ', ',') if cameras: string += f' | cameras {cameras}'.replace(', ', ',').replace("'", "") if labels: string += f' | labels {labels}'.replace(', ', ',').replace("'", "") if suffix: string += f' | suffix {suffix}'.replace(', ', ',').replace("'", "") print(string) datasets.append(dataset) cfg_for_name = copy.deepcopy(cfg_args) if 'path' in cfg_for_name and not isinstance(cfg_for_name['path'], list): cfg_for_name['path'] = [cfg_for_name['path']] datasets_cfg.append(make_cfg_name(args, cfg_for_name)) return datasets, datasets_cfg def set_random_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) def process_task(cfg_name, tokenizers): """ Each task can either be sim or real, with suffix like: BimanualPutRedBellPepperInBin/riverway/sim """ cfg_dict = unwrap_cfg(cfg_name) if not isinstance(cfg_name, dict) else cfg_name if args.buffer_pad_last is not None: if args.buffer_pad_last == 'pad': cfg_dict['pad_last'] = [int(args.buffer * 0.8)] else: cfg_dict['pad_last'] = [args.buffer_pad_last] if args.buffer_pad_first is not None: cfg_dict['pad_first'] = [args.buffer_pad_first] if args.buffer_crop_first is not None: cfg_dict['crop_first'] = [args.buffer_crop_first] if args.buffer_crop_last is not None: cfg_dict['crop_last'] = [args.buffer_crop_last] datasets, cfgs = setup_dataset(cfg_dict, verbose=True) for dataset_i, cfg_i in zip(datasets, cfgs): name_i = get_name(dataset_name, cfg_i, dataset_i, args) local_path_i = f'{args.local_path}/{args.cfg_file}/{name_i}' s3_path_i = local_path_i.replace(args.local_path, args.s3_path) s3_base_i = f'{args.s3_path}/{args.cfg_file}' if args.suffix is not None: s3_path_i = f'{s3_path_i}/{args.suffix}' print(f's3_path_i: {s3_path_i}') yaml_local_path = args.cfg_yaml_path yaml_s3_path = f"{s3_path_i}/{args.cfg_file}.yaml" run_cp(yaml_local_path, yaml_s3_path) print(f"Uploading YAML config to: {yaml_s3_path}") dirname = os.path.dirname(s3_path_i) basename = os.path.basename(s3_path_i) idx = -1 if args.suffix is None else -3 filename = f"{'__'.join(s3_path_i.split('/')[idx:])}.txt" proc = subprocess.Popen(f'aws s3 ls {dirname}/', stdout=subprocess.PIPE, shell=True) done = proc.communicate()[0].decode().replace('\n', ' ').split(' ') done = [d[:-1] for d in done if d != '' and d != 'PRE'] proc.kill() if basename in done: if not args.delete and not args.merge: print(f'### ALREADY PROCESSED: {s3_path_i}') else: if args.delete: print(f'### DELETING AWS: {s3_path_i}') subprocess.run( f'aws s3 rm {s3_path_i} --recursive --quiet', shell=True, capture_output=True, text=True, check=True, ) # elif args.merge: # print(f'### MERGING AWS: {s3_path_i}') if args.multi_sample is not None: for i in range(args.multi_sample): set_random_seed(i) process_dataset(s3_path_i, dataset_i, local_path_i, repeat_idx=i, tokenizers=tokenizers) else: process_dataset(s3_path_i, dataset_i, local_path_i, tokenizers=tokenizers) if args.parallelize == 'task': process_all_indices(s3_path_i, filename) else: if args.multi_sample is not None: for i in range(args.multi_sample): set_random_seed(i) process_dataset(s3_path_i, dataset_i, local_path_i, repeat_idx=i, tokenizers=tokenizers) else: process_dataset(s3_path_i, dataset_i, local_path_i, tokenizers=tokenizers) if args.parallelize == 'task': process_all_indices(s3_path_i, filename) def greedy_partition(lengths, task_suffixes, n_parts): """ Partition a decending list in to N parts such that the sum of each sublist is as close as possible. This is for scheduling all tasks across N ranks so that they finish all tasks within roughly the same time, and therefore optimize the GPU compute time by avoiding certain ranks waiting too long for the others. """ items_with_meta = list(zip(lengths, task_suffixes)) items_with_meta = sorted(items_with_meta, key=lambda x: x[0], reverse=True) partitions = [[] for _ in range(n_parts)] totals = [0] * n_parts for value, meta in items_with_meta: idx = totals.index(min(totals)) partitions[idx].append((value, meta)) # Keep both value and meta totals[idx] += value return partitions if __name__ == "__main__": args = arg_parser() print(f"Arguments: {args}") print(f"have s5cmd: {have_s5cmd()}") # Optional pre-download using repo helper if args.download_data: global_rank_pre = int(os.environ.get("RANK", "0")) world_size_pre = int(os.environ.get("WORLD_SIZE", "1")) split_rel = args.download_data if not split_rel.endswith(".json"): split_rel = f"{split_rel.rstrip('/')}/{args.input_split}" split_dir = os.path.dirname(split_rel) args._download_local_root = os.path.join(args.download_local_bucket, args.download_data_folder, split_dir) args._download_local_split = os.path.join(args.download_local_bucket, args.download_data_folder, split_rel) done_marker = f"{args._download_local_split}.download_done" cmd = [ "python3", "anydata/sync/download_uni.py", split_rel, "--local_bucket", args.download_local_bucket, "--s3_bucket", args.download_s3_bucket, "--data_folder", args.download_data_folder, "--num_procs", str(args.download_num_procs), ] # In distributed launch, only one rank should extract to the shared local folder. if world_size_pre == 1 or global_rank_pre == 0: print(f"Running download_uni.py: {' '.join(cmd)}") subprocess.run(cmd, check=True) with open(done_marker, "w") as f: f.write("done\n") else: print( f"Rank {global_rank_pre} waiting for download to finish at {done_marker}" ) waited = 0 while not os.path.exists(done_marker): time.sleep(5) waited += 5 if waited % 60 == 0: print(f"Rank {global_rank_pre} waiting... {waited}s") if waited >= 24 * 60 * 60: raise TimeoutError(f"Timed out waiting for download marker: {done_marker}") args.src_path = args._download_local_root print(f"Source path: {args.src_path}") if args.src_path.startswith("s3://"): # Example: s3://tri-ml-sandbox-16011-us-west-2-datasets/cv_unified/DDAD LOCAL_ROOT = args.src_path.replace("s3://tri-ml-sandbox-16011-us-west-2-datasets", "/tmp") else: # TODO @vitor inline hardcoded path may be confusing LOCAL_ROOT = "/data/cv_unified" # ---------------- GPU setup ---------------- local_rank = int(os.environ.get('LOCAL_RANK', '0')) global_rank = int(os.environ.get('RANK', '0')) world_size = int(os.environ.get('WORLD_SIZE', '1')) local_world_size = int(os.environ.get('LOCAL_WORLD_SIZE', '1')) node_rank = global_rank // local_world_size args.multi_gpu = world_size > 1 if args.multi_gpu: torch.cuda.set_device(local_rank) init_process_group(backend="nccl", timeout=datetime.timedelta(hours=6)) if args.seed > 0: set_random_seed(args.seed) # ---------------- Tokenizer init ---------------- if args.mode == "frames": tokenizers = None else: from custom.vae.a4d import VAE tokenizers = VAE( f"cuda:{local_rank}", vae_path="s3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints/nvidia/Cosmos-Predict2-2B-Video2World/tokenizer/tokenizer.pth", text_encoder_path="s3://tri-ml-sandbox-16011-us-west-2-datasets/cosmos-predict-2/checkpoints/google-t5/t5-11b", ) print(f"===============================Initialized model on local rank {local_rank}, node {node_rank}") # ---------------- Process each dataset ---------------- for dataset_name in args.datasets: print("################################################# PROCESSING", dataset_name) args.cfg_file, dataset_name = dataset_name.split("/") cfg_file = f"{args.cfg_path}/{args.cfg_file}.yaml" args.cfg_yaml_path = cfg_file data_cfg = read_config(cfg_file) datasets_cfg = data_cfg.datasets if hasattr(data_cfg, "datasets") else data_cfg.get("datasets", {}) dataset_cfg = datasets_cfg.__dict__.get(dataset_name, None) if hasattr(datasets_cfg, "__dict__") else datasets_cfg.get(dataset_name, None) if dataset_cfg is None: raise ValueError(f"Dataset {dataset_name} not found in {cfg_file}") dataset_cfg = unwrap_cfg(dataset_cfg) if not isinstance(dataset_cfg, dict) else dataset_cfg # If we pre-downloaded, override dataset path to local split file if args.download_data: dataset_cfg['path'] = [[args._download_local_split]] # Populate args.cameras from dataset cfg if not set in args_file/CLI if not args.cameras: cams = dataset_cfg.get("cameras", None) if cams and len(cams) > 0: args.cameras = list(cams[0]) print(f"args.cameras: {args.cameras}") cfg_name = make_cfg_name(args, dataset_cfg) print(f"cfg name: {cfg_name}") # args.src_path = data_cfg.datasets.__dict__[dataset_name].path[0] # --- Handle remote S3 datasets --- if args.src_path.startswith('s3://'): if args.producer: task_suffixes = get_all_suffixes(args.src_path, args.start_idx, args.end_idx) task_queue = multiprocessing.Queue(maxsize=3) downloader_p = multiprocessing.Process( target=download_producer, args=(task_queue, args.src_path, task_suffixes) ) downloader_p.start() while True: task_info = task_queue.get() if task_info is None: break ti, task_suffix = task_info print(f"Working on task ({ti + 1}/{len(task_suffixes)}) {task_suffix} on", f"node {node_rank}, global rank {global_rank}......") args.suffix = task_suffix os.makedirs(LOCAL_ROOT, exist_ok=True) # cfg_name.path = [f"{LOCAL_ROOT}/{args.suffix}"] process_task(cfg_name, tokenizers) # delete_task(task_suffix) downloader_p.join() else: LOCAL_ROOT = f"/tmp/cv_datasets/unified/{args.cfg_file}" os.makedirs(LOCAL_ROOT, exist_ok=True) dataset_cfg = unwrap_cfg(dataset_cfg) if args.split_json: dataset_cfg['path'] = [args.split_json] cfg_name = make_cfg_name(args, dataset_cfg) cfg_paths = dataset_cfg.get('path', []) flat_paths = flatten(cfg_paths) if isinstance(cfg_paths, (list, tuple)) else [cfg_paths] split_paths = [Path(p) for p in flat_paths if str(p).endswith('.json')] print("split paths:", split_paths) if split_paths: download_list = [] collected_keys = [] for split_path in split_paths: local_split_path = Path(split_path) if not local_split_path.exists(): try: rel = local_split_path.relative_to(LOCAL_ROOT) except ValueError: rel = None if rel is not None: s3_split_path = f"{args.src_path.rstrip('/')}/{rel.as_posix()}" local_split_path.parent.mkdir(parents=True, exist_ok=True) if args.s5cmd and have_s5cmd(): cmd = ["s5cmd", "cp", s3_split_path, str(local_split_path)] else: cmd = ["aws", "s3", "cp", s3_split_path, str(local_split_path)] try: subprocess.run(cmd, check=True) except subprocess.CalledProcessError as e: print(f"Failed to download split file {s3_split_path}: {e}") continue if not local_split_path.exists(): print(f"Split file {local_split_path} not found; skipping") continue try: split_data = json.loads(local_split_path.read_text()) except Exception as e: print(f"Failed to read split file {local_split_path}: {e}") continue try: rel_base = local_split_path.parent.relative_to(LOCAL_ROOT).as_posix() except ValueError: rel_base = "" if rel_base == ".": # split file at LOCAL_ROOT rel_base = "" if isinstance(split_data, dict): for key in split_data.keys(): collected_keys.append(key) suffix = f"{rel_base}/{key}" if rel_base else key download_list.append(suffix) else: print(f"Split file {local_split_path} must contain a JSON object mapping scene ids") if collected_keys: split_keys = sorted(set(collected_keys)) print(f"Total {len(split_keys)} sequences found in split files, and example keys: {split_keys[:5]}") if download_list: for idx, suffix in enumerate(download_list, start=1): print(f"[{idx}/{len(download_list)}] Syncing sequence {suffix}") download_task(args.src_path, suffix) # Ensure all ranks see a fully-synced local dataset before processing if args.multi_gpu: if is_initialized(): barrier() else: start_time = time.time() print(f"Downloading entire dataset {args.src_path} -> {LOCAL_ROOT}") run_cp_recursive(args.src_path, LOCAL_ROOT) # 🔹 download whole dataset once print(f"download finisehed using {time.time() - start_time} seconds") args.src_path = dataset_cfg.get('path', [None])[0] print("process task") process_task(cfg_name, tokenizers) else: # iterate over tasks stored locally args.src_path = dataset_cfg.get('path', [None])[0] process_task(cfg_name, tokenizers) print("####################################################### DONE", dataset_name) if args.multi_gpu: destroy_process_group() print("######################################## DONE!!!")