import copy import datetime import importlib import os import random import time import traceback from anydata.dataloaders.Base import InsufficientCamerasError from anydata.dataloaders.instantiate import dataset_from_config from anydata.utils.read import read_config import numpy as np from rich import print import torch from torch.utils.data import ConcatDataset, Dataset from custom.utils.logging import format_batch_summary from custom.dataset.camera_utils import reanchor_extrinsics from custom.utils.config import to_plain from custom.utils.constants import VAE_RATIO_THW from imaginaire.utils import log def weighted_choice(spec): ''' Pick one value from a temporal-aug spec. :param spec: int/scalar -> returned as-is; list -> uniform random (duplicates act as weights, e.g. [1,1,2,3]); dict -> weighted by its values/shares (e.g. {1: 0.5, 2: 0.25, 3: 0.25}, equivalently {1: 2, 2: 1, 3: 1}). NOTE: callers int()-coerce since OmegaConf stringifies keys. :return: one chosen element (or spec itself if scalar). ''' if isinstance(spec, dict): return random.choices(list(spec.keys()), weights=list(spec.values()), k=1)[0] if isinstance(spec, (list, tuple)): return random.choice(spec) return spec def _rewrite_s3_uri_to_sm_channel_path(uri, channel_specs): """Return the SageMaker input-channel path for an S3 URI, if it belongs to a mounted channel.""" uri = str(uri) for spec in channel_specs: s3_root = spec['s3_root'].rstrip('/') channel = spec['channel'] if uri == s3_root: return f'{channel}' if uri.startswith(f'{s3_root}/'): return f"{channel}/{uri[len(s3_root) + 1:]}" return None def _rewrite_config_for_sm_input_channels(config): """Rewrite AnyData S3 paths to SageMaker input-channel mount paths when channels are configured.""" count = int(os.environ.get('ANY4D_SM_INPUT_CHANNEL_COUNT', '0') or 0) if count <= 0: if os.environ.get('LOCAL_RANK', '0') == '0': log.info('FastFile mode disabled.') return config if os.environ.get('LOCAL_RANK', '0') == '0': input_source = os.environ.get('ANY4D_SM_INPUT_SOURCE', '') log.info(f'FastFile mode {"enabled" if input_source == "fastfile" else "disabled"}.') channel_specs = [] for i in range(count): suffix = f'{i:03d}' channel = os.environ.get(f'ANY4D_SM_INPUT_CHANNEL_{suffix}') s3_root = os.environ.get(f'ANY4D_SM_INPUT_S3_ROOT_{suffix}') if not channel or not s3_root: raise ValueError( f"ANY4D_SM_INPUT_CHANNEL_COUNT={count} but missing env vars for suffix={suffix}: " f"ANY4D_SM_INPUT_CHANNEL_{suffix}={channel!r}, ANY4D_SM_INPUT_S3_ROOT_{suffix}={s3_root!r}" ) channel_specs.append({'channel': channel, 's3_root': s3_root}) config = copy.deepcopy(config) path_prefix = str(config.get('path_prefix', '')).rstrip('/') rewritten_paths = [] for path_entry in config.get('path', []): if isinstance(path_entry, list): path_entry = list(path_entry) uri = str(path_entry[0]) full_uri = uri if uri.startswith('s3://') else f"{path_prefix}/{uri.lstrip('/')}" channel_path = _rewrite_s3_uri_to_sm_channel_path(full_uri, channel_specs) if channel_path is not None: path_entry[0] = channel_path rewritten_paths.append(path_entry) else: uri = str(path_entry) full_uri = uri if uri.startswith('s3://') else f"{path_prefix}/{uri.lstrip('/')}" channel_path = _rewrite_s3_uri_to_sm_channel_path(full_uri, channel_specs) rewritten_paths.append(channel_path if channel_path is not None else path_entry) config['path'] = rewritten_paths if any( str(path_entry[0] if isinstance(path_entry, list) else path_entry).startswith('anydata') for path_entry in rewritten_paths ): config['path_prefix'] = '/opt/ml/input/data' return config class AnyDataset(Dataset): def __init__( self, dataset_dir, num_frames=93, phase='train', vae_ratio=VAE_RATIO_THW, data_defaults=None, data_overrides=None, params=None, # unused, accepted for compatibility with VidarDataset single_gpu=False, # when True, webbed datasets won't split shards by rank shuffle=None, # for webdatasets: True = ResampledShards, False = SimpleShardList ): super().__init__() self.dataset_dir = dataset_dir self.num_frames = num_frames self.phase = phase self.vae_ratio = vae_ratio # to_plain: ${...}-interpolated overrides/defaults arrive as OmegaConf; normalize to plain. self.data_defaults = to_plain(data_defaults) or {} self.data_overrides = to_plain(data_overrides) or {} # Version-selectable robot-action parser module (default = current semantics). _afv = self.data_defaults.get('action_format_version', 'action_format_v1') self._action_parser = importlib.import_module(f'custom.dataset.{_afv}').parse_anydata_raw_robot_action self.params = params or {} self.single_gpu = single_gpu self.shuffle = shuffle config = read_config(dataset_dir)['dataset'] config = _rewrite_config_for_sm_input_channels(config) self.config = config # Log the resolved data sources so they show up in stdout / SageMaker logs. # Each path entry can be either `s3://...` or `[s3://..., weight]`. s3_paths = config.get('path', []) try: from megatron.core import parallel_state rank = parallel_state.get_data_parallel_rank() except (AssertionError, ImportError): rank = 0 path_lines = '\n '.join( f'{p[0]} (weight={p[1]})' if isinstance(p, (list, tuple)) and len(p) == 2 else str(p) for p in s3_paths ) log.info( f'[AnyDataset rank={rank}] yaml={dataset_dir} | phase={phase} | ' f'{len(s3_paths)} source(s):\n {path_lines}', rank0_only=False, ) # Per-dataset defaults from the custom section of the YAML config. # These are simple scalar values (e.g. frame_rate: 10.0, not sub-dicts). self.custom = to_plain(config.get('custom', config.get('added', {}))) or {} # Apply dataset-construction-time overrides from a4d_config. # These override top-level YAML keys before passing to anydata. if self.data_overrides: for key in ('length', 'single_sample', 'subsample', 'cameras_core_mode', 'cameras_aux_mode', 'multi_tarfiles'): val = self.data_overrides.get(key, None) if val is not None: config[key] = val # Nested override: augmentation.resize (e.g. [-16, 512]) resize = self.data_overrides.get('resize', None) if resize is not None: if 'augmentation' not in config: config['augmentation'] = {} config['augmentation']['resize'] = resize # Instantiate the anydata dataset. config.pop('custom', None) config.pop('added', None) if single_gpu: config['single_gpu'] = True # NOTE(bvh): for webdatasets, this controls shard resampling and sub-clip # offset shuffling (DistributedSampler has no effect on Webbed datasets). if shuffle and config.get('name', '') == 'Webbed': config['shard_type'] = 'resampled' config['shuffle_samples'] = True self.dataset, _ = dataset_from_config(config) self.success_count = 0 self.error_count = 0 self.camera_skip_count = 0 def __str__(self): return f'{len(self.dataset)} samples from {self.dataset_dir}' def __len__(self): return len(self.dataset) def _extract_prompt(self, sample): ''' Pull a prompt (str or list of candidates) from language metadata, or None if absent. ''' # TODO(bvh): confirm lang_try_order works for DROID: lang_try_order = ['prompt', 'task', 'omniworld', 'caption'] # NOTE(bvh): ^ e.g. DROID often has prompt, always has task, sometimes has omniworld, etc. for lang_key in ('language', 'lat_language'): if lang_key not in sample: continue # Language entries are keyed by (t, c) tuples; grab the first one. first_val = next(iter(sample[lang_key].values()), None) if not isinstance(first_val, dict): prompt = first_val else: prompt = None for sub_key in lang_try_order: cand = first_val.get(sub_key, None) while isinstance(cand, dict): # e.g. omniworld has further subkeys, one for each camera -> pick one cand = random.choice(list(cand.values())) if isinstance(cand, np.ndarray): cand = '\n'.join(str(c) for c in cand.ravel()) # Skip empty/whitespace candidates so the next key is tried instead # (e.g. YAMxdof has prompt='' on early frames but task is always set). if cand is not None and str(cand).strip(): prompt = str(cand) break while isinstance(prompt, dict): prompt = random.choice(list(prompt.values())) if isinstance(prompt, str): prompt = prompt.split("\n") # always becomes list of str! return prompt return None def supplement_sample(self, sample, index): ''' Resolve per-sample params; priority by merge order: defaults < custom (yaml) < sample < overrides. frame_stride / frame_trim / prompt may be int/list/dict specs and are randomly sampled here (pre-collate, so only the chosen scalar crosses collate). (clip/scale_depth: see a4d_model.) ''' # Per-sample metadata, keyed by config name (only keys this sample provides). meta = sample['metadata'] sample_params = {} if meta.get('framerate'): sample_params['frame_rate'] = list(meta['framerate'].values())[0] # Strided webbed splits (e.g. YAMxdof str3) keep the RAW framerate in metadata but # only store every stride-th frame; divide to get the effective rate of stored frames. web_stride = meta.get('stride', None) for _ in range(2): # unwrap {cam: val} and legacy {cam: {sub: val}} nesting if isinstance(web_stride, dict): web_stride = next(iter(web_stride.values()), None) if web_stride is not None and float(web_stride) > 1: sample_params['frame_rate'] = sample_params['frame_rate'] / float(web_stride) if 'scale_factor' in sample: sample_params['scale_factor'] = sample['scale_factor'] if 'reference_camera' in sample: sample_params['reference_camera'] = sample['reference_camera'] elif 'zero_origin' in sample: # legacy name sample_params['reference_camera'] = sample['zero_origin'] if 'stride' in sample: sample_params['frame_stride'] = sample['stride'] # This line is quite powerful! :-) p = {**self.data_defaults, **self.custom, **sample_params, **self.data_overrides} fps = p.get('frame_rate') scale_factor = p.get('scale_factor', 0.125) reference_camera = p.get('reference_camera', None) ref_cam_rel_per_timestep = p.get('ref_cam_rel_per_timestep', False) # Temporal augs: ALL randomness lives here. int/list/dict spec -> one chosen int (frame_trim: # False/None/<=0 -> 0 = off). frame_offset -> a fraction in [0,1) of the window slack # (0.0 = front / disabled); transforms maps it to an int once it knows the post-stride length. frame_stride = int(weighted_choice(p.get('frame_stride', 1))) frame_trim = weighted_choice(p.get('frame_trim', False)) frame_trim = max(int(frame_trim), 0) if frame_trim else 0 frame_offset = random.random() if weighted_choice(p.get('frame_offset', False)) else 0.0 # Prompt: bespoke language extraction, but an override wins; list of candidates -> pick one. if 'prompt' in self.data_overrides: prompt = self.data_overrides['prompt'] else: prompt = self._extract_prompt(sample) if prompt is None: prompt = p.get('prompt', 'do something') if not isinstance(prompt, str): # list of candidates -> pick one # Guard against all-empty candidates (would IndexError and kill the sample). candidates = [s.strip() for s in prompt if s.strip()] prompt = random.choice(candidates) if candidates else p.get('prompt', 'do something') return (fps, prompt, scale_factor, reference_camera, ref_cam_rel_per_timestep, frame_stride, frame_trim, frame_offset) def reset_iter(self): ''' Reset the wrapped webdataset iterator (no-op if the wrapped dataset is map-style or doesn't define reset_iter). Called by the trainer before each validation round so repeated passes re-consume the same leading subsequence of shards. ''' if hasattr(self.dataset, 'reset_iter'): self.dataset.reset_iter() def __getitem__(self, index): try: # Get anydata sample (can be from any dataset). # With trainer calling reset_iter() before each val round, the # wrapped webdataset is guaranteed to be at position 0 on entry; # StopIteration here means the dataset has no shards for this rank # (subsample < world_size), which is handled by the trainer via # the pre-flight broadcast path. fetch_start = time.perf_counter() anydata_dict = self.dataset[index] fetch_time = time.perf_counter() - fetch_start # ^ NOTE(bvh): Entries do not yet have batch dimension, # e.g. 'rgb' -> (3, H, W) tensors of float in [0, 1]. # Get short, meaningful dataset source identifier # Batch-level metadata has 'name' directly; file-level nests it under 'info'. meta = anydata_dict['metadata'] dset_name = meta.get('name') or meta.get('info', {}).get('name', 'unknown_name') # e.g. 'LBM' or 'HumanoidEveryday' # Get long, specific video identifier raw_id = anydata_dict['metadata'].get('raw_id', None) # ^ e.g. 'tri413_1403043282076853000_1403043293076853000' path = anydata_dict['metadata'].get('path', None) # ^ e.g. 's3://gaia-e2e-wfm-datasets/anydata/cv_webbed/videos/GAIA/split_all__IKE_32/fsdatasets/p4h_gaia_us/0.1.0/PARTITION=train/024604_tri413_1403043282076853000_1403043293076853000/%s/0028_0055.tar' storage = anydata_dict['metadata'].get('storage', 'frames') # e.g. 'videos' sample_id = f'{dset_name}_{index}' if raw_id is not None and '/' in raw_id: # raw_id presumably contains unified path (but not always!) sample_id = raw_id # ^ e.g. 'raw/ILIAD/success/...' elif path is not None: # "simulate" unified path by trimming 's3://.../GAIA/split_all__IKE_32/' prefix + '/%s/0028_0055.tar' suffix sample_id = path.split(f'/{storage}/', 1)[1].split('/', 2)[-1].split('/%s/', 1)[0] # ^ e.g. 'fsdatasets/p4h_gaia_us/0.1.0/PARTITION=train/024604_tri413_1403043282076853000_1403043293076853000' elif raw_id is not None: # fallback to slashless name (only available clue) sample_id = f'{index}_{raw_id}' # ^ e.g. '123_tri413_1403043282076853000_1403043293076853000' # NOTE(bvh): In Any4D, different viewpoints do NOT necessarily need to have the same # resolution or duration, and we eliminated the need to obtain representative values here. # Process raw actions if exist if 'action' in anydata_dict: anydata_dict['action_orig'] = copy.deepcopy(anydata_dict['action']) anydata_dict['action'] = self._action_parser(anydata_dict, dset_name) # Get per-sample metadata (fps, prompt, scale_factor, reference_camera, ref_cam_rel_per_timestep, frame_stride, frame_trim, frame_offset) = self.supplement_sample(anydata_dict, index) # Preserve original extrinsics and re-anchor if configured. # NOTE(bvh): this replaces the ego_camera system in vidar # NOTE(bvh): cams and cams_orig are constructed later in transforms.py:prepare_batch() if 'extrinsics' in anydata_dict: anydata_dict['extrinsics_orig'] = copy.deepcopy(anydata_dict['extrinsics']) if reference_camera is not False and reference_camera is not None: cam_names = meta.get('cameras', None) anydata_dict['extrinsics'] = reanchor_extrinsics( anydata_dict['extrinsics'], reference_camera, ref_cam_rel_per_timestep, camera_names=cam_names) # NOTE(bvh): FPS correction happens later in a4d_model because of stride alignment # + potential intra-batch differences. data_dict = { 'dl_idx': index, 'anydata': anydata_dict, 'dset_name': dset_name, 'sample_id': sample_id, #### 'fps': fps, 'fps_orig': fps, # preserved through stride / frame_rate overrides (cf. extrinsics_orig) 'prompt': prompt, 'scale_factor': scale_factor, 'reference_camera': reference_camera, 'frame_stride': frame_stride, 'frame_trim': frame_trim, 'frame_offset': frame_offset, #### 'is_preprocessed': True, 'is_anydata': True, 'sample_fetch_time': fetch_time, } self.success_count += 1 self.camera_skip_count = 0 return data_dict except StopIteration: # NOTE(bvh): Webdataset exhausted => raise (pass through) so the trainer's flexible # validation handling (batch broadcast / padding) can process it. # We should not swallow it ourselves because it will spam retry on # an already-exhausted iterator, and only crash once error_count >= 30 raise except InsufficientCamerasError as e: log.warning(f'Insufficient cameras (index {index}): {e}', rank0_only=False) self._log_error_inspection(index, e, locals().get('anydata_dict')) self.camera_skip_count += 1 # NOTE: short term counter, resets immediately upon success log.info(f'Number of consecutive camera skips so far: {self.camera_skip_count}', rank0_only=False) if self.camera_skip_count >= 20: log.error(f'Too many consecutive camera skips, crashing...', rank0_only=False) raise return self[np.random.randint(len(self))] except Exception as e: # TODO(bvh): unify this with the robustness logic in # custom/utils/utils.py:resilient_subroutine() log.warning(f'Invalid data encountered (index {index}), skipping...', rank0_only=False) log.warning('==== FULL TRACEBACK ====', rank0_only=False) log.warning(traceback.format_exc(), rank0_only=False) log.warning('========================', rank0_only=False) self._log_error_inspection(index, e, locals().get('anydata_dict')) self.error_count += 1 # NOTE: permanent counter, does not reset log.info(f'Number of errors so far: {self.error_count} (successes: {self.success_count})', rank0_only=False) time.sleep(1.0 + self.error_count / 10.0) # Crash above ~10% error rate, with a 10-error grace at startup. total_count = self.success_count + self.error_count crash_threshold = 10 + total_count * 0.1 if self.error_count >= crash_threshold: log.error(f'Too many errors ({self.error_count} / {total_count} total, ' f'threshold={crash_threshold:.1f}), crashing...', rank0_only=False) raise return self[np.random.randint(len(self))] def _log_error_inspection(self, index, exc, anydata_dict): ''' Dump per-rank inspection context for an error in __getitem__. Reads exc.webbed_context (attached by WebbedDataset.__getitem__) for shard-specific state. ''' try: ts = datetime.datetime.now().isoformat(timespec='seconds') wctx = getattr(exc, 'webbed_context', None) or {} # Strip large/dumped-separately fields so the header stays one line. wctx_head = {k: v for k, v in wctx.items() if k != 'super_sample'} rank_worker = wctx.get('rank_worker', '?') header = (f'time={ts} phase={self.phase} index={index} ' f'dataset_dir={self.dataset_dir!r} ' f'wrapped={type(self.dataset).__name__} ' f'successes={self.success_count} errors={self.error_count} ' f'camera_skips={self.camera_skip_count}\n' f'webbed_context={wctx_head!r}') log.warning(f'==== DATA ERROR INSPECTION ====\n{header}', rank0_only=False) if isinstance(anydata_dict, dict): log.warning( format_batch_summary(anydata_dict, label=f'anydata_dict at error (idx={index})'), rank0_only=False) elif wctx.get('super_sample'): # skip when empty (invalidated by create_sample) log.warning( format_batch_summary(wctx['super_sample'], label=f'Webbed.super_sample[{rank_worker}]'), rank0_only=False) log.warning('===============================', rank0_only=False) except Exception as fmt_e: log.warning(f'(error inspection dump itself failed: {fmt_e!r})', rank0_only=False)