from dataclasses import dataclass from typing import Callable, List, Optional, Tuple, Dict import functools import torch import tensorflow as tf from tml.common.batch import DataclassBatch from tml.projects.home.recap.data.config import RecapDataConfig, TaskData from tml.projects.home.recap.data import preprocessors from tml.projects.home.recap.config import JobMode from tml.projects.home.recap.data.tfe_parsing import get_seg_dense_parse_fn from tml.projects.home.recap.data.util import ( keyed_jagged_tensor_from_tensors_dict, sparse_or_dense_tf_to_torch, ) from absl import logging import torch.distributed as dist @dataclass class RecapBatch(DataclassBatch): """Holds features and labels from the Recap dataset.""" continuous_features: torch.Tensor binary_features: torch.Tensor discrete_features: torch.Tensor sparse_features: "KeyedJaggedTensor" # type: ignore[name-defined] # noqa: F821 labels: torch.Tensor user_embedding: torch.Tensor = None user_eng_embedding: torch.Tensor = None author_embedding: torch.Tensor = None weights: torch.Tensor = None def __post_init__(self): if self.weights is None: self.weights = torch.ones_like(self.labels) for feature_name, feature_value in self.as_dict().items(): if ("embedding" in feature_name) and (feature_value is None): setattr(self, feature_name, torch.empty([0, 0])) def to_batch(x, sparse_feature_names: Optional[List[str]] = None) -> RecapBatch: """Converts a torch data loader output into `RecapBatch`.""" x = tf.nest.map_structure(functools.partial( sparse_or_dense_tf_to_torch, pin_memory=False), x) try: features_in, labels = x except ValueError: # For Mode.INFERENCE, we do not expect to recieve labels as part of the input tuple features_in, labels = x, None sparse_features = keyed_jagged_tensor_from_tensors_dict({}) if sparse_feature_names: sparse_features = keyed_jagged_tensor_from_tensors_dict( {embedding_name: features_in[embedding_name] for embedding_name in sparse_feature_names} ) user_embedding, user_eng_embedding, author_embedding = None, None, None if "user_embedding" in features_in: if sparse_feature_names and "meta__user_id" in sparse_feature_names: raise ValueError( "Only one source of embedding for user is supported") else: user_embedding = features_in["user_embedding"] if "user_eng_embedding" in features_in: if sparse_feature_names and "meta__user_eng_id" in sparse_feature_names: raise ValueError( "Only one source of embedding for user is supported") else: user_eng_embedding = features_in["user_eng_embedding"] if "author_embedding" in features_in: if sparse_feature_names and "meta__author_id" in sparse_feature_names: raise ValueError( "Only one source of embedding for user is supported") else: author_embedding = features_in["author_embedding"] return RecapBatch( continuous_features=features_in["continuous"], binary_features=features_in["binary"], discrete_features=features_in["discrete"], sparse_features=sparse_features, user_embedding=user_embedding, user_eng_embedding=user_eng_embedding, author_embedding=author_embedding, labels=labels, # Defaults to torch.ones_like(labels) weights=features_in.get("weights", None), ) def _chain(param, f1, f2): """ Reduce multiple functions into one chained function _chain(x, f1, f2) -> f2(f1(x)) """ output = param fns = [f1, f2] for f in fns: output = f(output) return output def _add_weights(inputs, tasks: Dict[str, TaskData]): """Adds weights based on label sampling for positive and negatives. This is useful for numeric calibration etc. This mutates inputs. Args: inputs: A dictionary of strings to tensor-like structures. tasks: A dict of string (label) to `TaskData` specifying inputs. Returns: A tuple of features and labels; weights are added to features. """ weights = [] for key, task in tasks.items(): label = inputs[key] float_label = tf.cast(label, tf.float32) weights.append( float_label / task.pos_downsampling_rate + (1.0 - float_label) / task.neg_downsampling_rate ) # Ensure we are batch-major (assumes we batch before this call). inputs["weights"] = tf.squeeze(tf.transpose( tf.convert_to_tensor(weights)), axis=0) return inputs def get_datetimes(explicit_datetime_inputs): """Compute list datetime strings for train/validation data.""" datetime_format = "%Y/%m/%d/%H" end = datetime.strptime( explicit_datetime_inputs.end_datetime, datetime_format) dates = sorted( [ (end - timedelta(hours=i + 1)).strftime(datetime_format) for i in range(int(explicit_datetime_inputs.hours)) ] ) return dates def get_explicit_datetime_inputs_files(explicit_datetime_inputs): """ Compile list of files for training/validation. Used with DataConfigs that use the `explicit_datetime_inputs` format to specify data. For each hour of data, if the directory is missing or empty, we increment a counter to keep track of the number of missing data hours. Returns only files with a `.gz` extension. Args: explicit_datetime_inputs: An `ExplicitDatetimeInputs` object within a `datasets.DataConfig` object Returns: data_files: Sorted list of files to read corresponding to data at the desired datetimes num_hours_missing: Number of hours that we are missing data """ datetimes = get_datetimes(explicit_datetime_inputs) folders = [os.path.join(explicit_datetime_inputs.data_root, datetime) for datetime in datetimes] data_files = [] num_hours_missing = 0 for folder in folders: try: files = tf.io.gfile.listdir(folder) if not files: logging.warning(f"{folder} contained no data files") num_hours_missing += 1 data_files.extend( [ os.path.join(folder, filename) for filename in files if filename.rsplit(".", 1)[-1].lower() == "gz" ] ) except tf.errors.NotFoundError as e: num_hours_missing += 1 logging.warning( f"Cannot find directory {folder}. Missing one hour of data. Error: \n {e}") return sorted(data_files), num_hours_missing def _map_output_for_inference( inputs, tasks: Dict[str, TaskData], preprocessor: tf.keras.Model = None, add_weights: bool = False ): """ Map the input data for inference. Args: inputs (dict): Input data dictionary. tasks (Dict[str, TaskData]): Dictionary of task data configurations. preprocessor (tf.keras.Model, optional): Preprocessor model for input data. Defaults to None. add_weights (bool, optional): Whether to add weights. Defaults to False. Returns: dict: Modified input data dictionary for inference. """ if preprocessor: raise ValueError("No preprocessor should be used at inference time.") if add_weights: raise NotImplementedError() # Add zero weights. inputs["weights"] = tf.zeros_like( tf.expand_dims(inputs["continuous"][:, 0], -1)) for label in tasks: del inputs[label] return inputs def _map_output_for_train_eval( inputs, tasks: Dict[str, TaskData], preprocessor: tf.keras.Model = None, add_weights: bool = False ): """ Map the input data for training and evaluation. Args: inputs (dict): Input data dictionary. tasks (Dict[str, TaskData]): Dictionary of task data configurations. preprocessor (tf.keras.Model, optional): Preprocessor model for input data. Defaults to None. add_weights (bool, optional): Whether to add weights. Defaults to False. Returns: Tuple[dict, tf.Tensor]: Modified input data dictionary and label values for training and evaluation. """ if add_weights: inputs = _add_weights_based_on_sampling_rates(inputs, tasks) # Warning this has to happen first as it changes the input if preprocessor: inputs = preprocessor(inputs) label_values = tf.squeeze( tf.stack([inputs[label] for label in tasks], axis=1), axis=[-1]) for label in tasks: del inputs[label] return inputs, label_values def _add_weights_based_on_sampling_rates(inputs, tasks: Dict[str, TaskData]): """Adds weights based on label sampling for positive and negatives. This is useful for numeric calibration etc. This mutates inputs. Args: inputs: A dictionary of strings to tensor-like structures. tasks: A dict of string (label) to `TaskData` specifying inputs. Returns: A tuple of features and labels; weights are added to features. """ weights = [] for key, task in tasks.items(): label = inputs[key] float_label = tf.cast(label, tf.float32) weights.append( float_label / task.pos_downsampling_rate + (1.0 - float_label) / task.neg_downsampling_rate ) # Ensure we are batch-major (assumes we batch before this call). inputs["weights"] = tf.squeeze(tf.transpose( tf.convert_to_tensor(weights)), axis=0) return inputs class RecapDataset(torch.utils.data.IterableDataset): def __init__( self, data_config: RecapDataConfig, dataset_service: Optional[str] = None, mode: JobMode = JobMode.TRAIN, compression: Optional[str] = "AUTO", repeat: bool = False, vocab_mapper: tf.keras.Model = None, ): """ Create a RecapDataset for training or inference. Args: data_config (RecapDataConfig): Data configuration. dataset_service (str, optional): Dataset service identifier. Defaults to None. mode (JobMode, optional): Job mode (TRAIN or INFERENCE). Defaults to JobMode.TRAIN. compression (str, optional): Compression type. Defaults to "AUTO". repeat (bool, optional): Whether to repeat the dataset. Defaults to False. vocab_mapper (tf.keras.Model, optional): Vocabulary mapper. Defaults to None. """ logging.info("***** Labels *****") logging.info(list(data_config.tasks.keys())) self._data_config = data_config self._parse_fn = get_seg_dense_parse_fn(data_config) self._mode = mode self._repeat = repeat self._num_concurrent_iterators = 1 self._vocab_mapper = vocab_mapper self.dataset_service = dataset_service preprocessor = None self._batch_size_multiplier = 1 if data_config.preprocess: preprocessor = preprocessors.build_preprocess( data_config.preprocess, mode=mode) if data_config.preprocess.downsample_negatives: self._batch_size_multiplier = data_config.preprocess.downsample_negatives.batch_multiplier self._preprocessor = preprocessor if mode == JobMode.INFERENCE: if preprocessor is not None: raise ValueError("Expect no preprocessor at inference time.") should_add_weights = False output_map_fn = _map_output_for_inference # (features,) else: # Only add weights if there is a reason to! If all weights will # be equal to 1.0, save bandwidth between DDS and Chief by simply # relying on the fact that weights default to 1.0 in `RecapBatch` # WARNING: Weights may still be added as a side effect of a preprocessor # such as `DownsampleNegatives`. should_add_weights = any( [ task_cfg.pos_downsampling_rate != 1.0 or task_cfg.neg_downsampling_rate != 1.0 for task_cfg in data_config.tasks.values() ] ) output_map_fn = _map_output_for_train_eval # (features, labels) self._output_map_fn = functools.partial( output_map_fn, tasks=data_config.tasks, preprocessor=preprocessor, add_weights=should_add_weights, ) sparse_feature_names = list( vocab_mapper.vocabs.keys()) if vocab_mapper else None self._tf_dataset = self._create_tf_dataset() self._init_tensor_spec() def _init_tensor_spec(self): """ Initialize the tensor specification for the dataset. """ def _tensor_spec_to_torch_shape(spec): if spec.shape is None: return None shape = [x if x is not None else -1 for x in spec.shape] return torch.Size(shape) self.torch_element_spec = tf.nest.map_structure( _tensor_spec_to_torch_shape, self._tf_dataset.element_spec ) def _create_tf_dataset(self): """ Create a TensorFlow dataset from the data files. Returns: tf.data.Dataset: TensorFlow dataset. """ if hasattr(self, "_tf_dataset"): raise ValueError( "Do not call `_create_tf_dataset` more than once.") world_size = dist.get_world_size() if dist.is_initialized() else 1 per_replica_bsz = ( self._batch_size_multiplier * self._data_config.global_batch_size // world_size ) dataset: tf.data.Dataset = self._create_base_tf_dataset( batch_size=per_replica_bsz, ) if self._repeat: logging.info("Repeating dataset") dataset = dataset.repeat() if self.dataset_service: if self._num_concurrent_iterators > 1: if not self.machines_config: raise ValueError( "Must supply a machine_config for autotuning in order to use >1 concurrent iterators" ) dataset = dataset_lib.with_auto_tune_budget( dataset, machine_config=self.machines_config.chief, num_concurrent_iterators=self.num_concurrent_iterators, on_chief=False, ) self.dataset_id, self.job_name = register_dataset( dataset=dataset, dataset_service=self.dataset_service, compression=self.compression ) dataset = distribute_from_dataset_id( dataset_id=self.dataset_id, # type: ignore[arg-type] job_name=self.job_name, dataset_service=self.dataset_service, compression=self.compression, ) elif self._num_concurrent_iterators > 1: if not self.machines_config: raise ValueError( "Must supply a machine_config for autotuning in order to use >1 concurrent iterators" ) dataset = dataset_lib.with_auto_tune_budget( dataset, machine_config=self.machines_config.chief, num_concurrent_iterators=self._num_concurrent_iterators, on_chief=True, ) # Vocabulary mapping happens on the training node, not in dds because of size. if self._vocab_mapper: dataset = dataset.map(self._vocab_mapper) return dataset.prefetch(world_size * 2) def _create_base_tf_dataset(self, batch_size: int): """ Create the base TensorFlow dataset. Args: batch_size (int): Batch size. Returns: tf.data.Dataset: Base TensorFlow dataset. """ if self._data_config.inputs: glob = self._data_config.inputs filenames = sorted(tf.io.gfile.glob(glob)) elif self._data_config.explicit_datetime_inputs: num_missing_hours_tol = self._data_config.explicit_datetime_inputs.num_missing_hours_tol filenames, num_hours_missing = get_explicit_datetime_inputs_files( self._data_config.explicit_datetime_inputs, increment="hourly", ) if num_hours_missing > num_missing_hours_tol: raise ValueError( f"We are missing {num_hours_missing} hours of data" f"more than tolerance {num_missing_hours_tol}." ) elif self._data_config.explicit_date_inputs: num_missing_days_tol = self._data_config.explicit_date_inputs.num_missing_days_tol filenames, num_days_missing = get_explicit_datetime_inputs_files( self._data_config.explicit_date_inputs, increment="daily", ) if num_days_missing > num_missing_days_tol: raise ValueError( f"We are missing {num_days_missing} days of data" f"more than tolerance {num_missing_days_tol}." ) else: raise ValueError( "Must specifiy either `inputs`, `explicit_datetime_inputs`, or `explicit_date_inputs` in data_config" ) num_files = len(filenames) logging.info(f"Found {num_files} data files") if num_files < 1: raise ValueError("No data files found") if self._data_config.num_files_to_keep is not None: filenames = filenames[: self._data_config.num_files_to_keep] logging.info(f"Retaining only {len(filenames)} files.") filenames_ds = ( tf.data.Dataset.from_tensor_slices( filenames).shuffle(len(filenames)) # Because of drop_remainder, if our dataset does not fill # up a batch, it will emit nothing without this repeat. .repeat(-1) ) if self._data_config.file_batch_size: filenames_ds = filenames_ds.batch( self._data_config.file_batch_size) def per_shard_dataset(filename): """ Create a TensorFlow dataset for a single shard file. Args: filename (str): Path to the shard file. Returns: tf.data.Dataset: TensorFlow dataset for the shard file. """ ds = tf.data.TFRecordDataset([filename], compression_type="GZIP") return ds.prefetch(4) ds = filenames_ds.interleave( per_shard_dataset, block_length=4, deterministic=False, num_parallel_calls=self._data_config.interleave_num_parallel_calls or tf.data.experimental.AUTOTUNE, ) # Combine functions into one map call to reduce overhead. map_fn = functools.partial( _chain, f1=self._parse_fn, f2=self._output_map_fn, ) # Shuffle -> Batch -> Parse is the correct ordering # Shuffling needs to be performed before batching otherwise there is not much point # Batching happens before parsing because tf.Example parsing is actually vectorized # and works much faster overall on batches of data. ds = ( # DANGER DANGER: there is a default shuffle size here. ds.shuffle(self._data_config.examples_shuffle_buffer_size) .batch(batch_size=batch_size, drop_remainder=True) .map( map_fn, num_parallel_calls=self._data_config.map_num_parallel_calls or tf.data.experimental.AUTOTUNE, ) ) if self._data_config.cache: ds = ds.cache() if self._data_config.ignore_data_errors: ds = ds.apply(tf.data.experimental.ignore_errors()) options = tf.data.Options() options.experimental_deterministic = False ds = ds.with_options(options) return ds def _gen(self): for x in self._tf_dataset: yield to_batch(x) def to_dataloader(self) -> Dict[str, torch.Tensor]: """ Convert the RecapDataset to a PyTorch DataLoader. Returns: torch.utils.data.DataLoader: PyTorch DataLoader for the dataset. """ return torch.utils.data.DataLoader(self, batch_size=None) def __iter__(self): return iter(self._gen())