Source code for gigl.common.data.load_torch_tensors

import time
import traceback
from dataclasses import dataclass, replace
from typing import MutableMapping, Optional, Union, cast

import torch
import torch.multiprocessing as mp
from graphlearn_torch.distributed.rpc import barrier, rpc_is_initialized
from torch.multiprocessing import Manager

from gigl.common.data.dataloaders import (
    SerializedTFRecordInfo,
    TFDatasetOptions,
    TFRecordDataLoader,
)
from gigl.common.logger import Logger
from gigl.src.common.types.graph_data import EdgeType, NodeType
from gigl.types.graph import (
    DEFAULT_HOMOGENEOUS_EDGE_TYPE,
    DEFAULT_HOMOGENEOUS_NODE_TYPE,
    FeatureQuantizationMetadata,
    LoadedGraphTensors,
)
from gigl.utils.share_memory import share_memory



logger = Logger()


_ID_FMT = "{entity}_ids"
_FEATURE_FMT = "{entity}_features"
_PACKED_FEATURE_FMT = "{entity}_packed_features"
_LABEL_FMT = "{entity}_labels"
_EDGE_WEIGHTS_KEY = "edge_weights"
_NODE_KEY = "node"


def _extract_weight_col(
    feat_tensor: torch.Tensor,
    feature_keys: list[str],
    feature_spec: dict,
    col_name: str,
    edge_type: EdgeType,
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
    """Slice a named weight column out of a feature tensor.

    Accounts for multi-dim features: each feature key may contribute more than one column
    to ``feat_tensor`` (e.g. ``FixedLenFeature(shape=[16])`` contributes 16 columns).
    The weight feature must be a scalar (width 1).

    Args:
        feat_tensor: Edge feature tensor of shape ``[num_edges, total_feature_cols]``.
        feature_keys: Ordered list of feature names matching the columns of ``feat_tensor``.
        feature_spec: Feature spec dict mapping feature name to its TF feature spec (used to
            determine per-key column widths).
        col_name: Name of the column to extract as weights.
        edge_type: Edge type (used only in error messages).

    Returns:
        A tuple ``(weights, trimmed_features)`` where ``weights`` is a 1-D tensor of shape
        ``[num_edges]`` and ``trimmed_features`` is ``feat_tensor`` with the weight column
        removed.

    Raises:
        ValueError: If ``col_name`` is not in ``feature_keys`` or the weight feature is not
            width 1.
    """
    if col_name not in feature_keys:
        raise ValueError(
            f"weight_edge_feat_name '{col_name}' not found in edge feature keys "
            f"for edge type {edge_type}: {feature_keys}"
        )
    key_idx = feature_keys.index(col_name)
    col_widths = []
    for key in feature_keys:
        spec = feature_spec[key]
        col_widths.append(spec.shape[-1] if spec.shape else 1)
    weight_width = col_widths[key_idx]
    if weight_width != 1:
        raise ValueError(
            f"weight_edge_feat_name '{col_name}' for edge type {edge_type} must be a scalar "
            f"feature (width 1), but has width {weight_width}."
        )
    col_offset = sum(col_widths[:key_idx])
    weights = feat_tensor[:, col_offset]
    keep_cols = [i for i in range(feat_tensor.shape[1]) if i != col_offset]
    trimmed = feat_tensor[:, keep_cols] if keep_cols else None
    return weights, trimmed


_EDGE_KEY = "edge"
_POSITIVE_LABEL_KEY = "positive_label"
_NEGATIVE_LABEL_KEY = "negative_label"


@dataclass(frozen=True)


class SerializedGraphMetadata:
    """
    Stores information for all entities. If homogeneous, all types are of type SerializedTFRecordInfo. Otherwise, they are dictionaries with the corresponding mapping.
    """

    # Node Entity Info for loading node tensors, a SerializedTFRecordInfo for homogeneous and dict[NodeType, SerializedTFRecordInfo] for heterogeneous cases


    node_entity_info: Union[
        SerializedTFRecordInfo, dict[NodeType, SerializedTFRecordInfo]
    ]

    # Edge Entity Info for loading edge tensors, a SerializedTFRecordInfo for homogeneous and dict[EdgeType, SerializedTFRecordInfo] for heterogeneous cases


    edge_entity_info: Union[
        SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]
    ]

    # Positive Label Entity Info, if present, a SerializedTFRecordInfo for homogeneous and dict[EdgeType, SerializedTFRecordInfo] for heterogeneous cases.
    # If there are no positive labels for any edge type, this value is None


    positive_label_entity_info: Optional[
        Union[SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]]
    ] = None

    # Negative Label Entity Info, if present, a SerializedTFRecordInfo for homogeneous and dict[EdgeType, SerializedTFRecordInfo] for heterogeneous cases.
    # If there are no negative labels for any edge type, this value is None.


    negative_label_entity_info: Optional[
        Union[SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]]
    ] = None

    # Optional node quantization metadata.


    node_quantization_metadata: Optional[
        Union[FeatureQuantizationMetadata, dict[NodeType, FeatureQuantizationMetadata]]
    ] = None



    edge_quantization_metadata: Optional[
        Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]]
    ] = None




def _validate_weight_edge_feature_name(
    edge_entity_info: Union[
        SerializedTFRecordInfo, dict[EdgeType, SerializedTFRecordInfo]
    ],
    weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]],
) -> None:
    if weight_edge_feat_name is None:
        return

    configured_weights: list[tuple[EdgeType, str, SerializedTFRecordInfo]]
    if isinstance(edge_entity_info, SerializedTFRecordInfo):
        if not isinstance(weight_edge_feat_name, str):
            raise ValueError("weight_edge_feat_name must be str for homogeneous graph")
        edge_type = DEFAULT_HOMOGENEOUS_EDGE_TYPE
        configured_weights = [(edge_type, weight_edge_feat_name, edge_entity_info)]
    else:
        if isinstance(weight_edge_feat_name, str):
            if len(edge_entity_info) != 1:
                raise ValueError(
                    "weight_edge_feat_name must be dict[EdgeType, str] for heterogeneous graph with multiple edge types"
                )
            edge_type, serialized_info = next(iter(edge_entity_info.items()))
            configured_weights = [(edge_type, weight_edge_feat_name, serialized_info)]
        else:
            unknown_edge_types = set(weight_edge_feat_name) - set(edge_entity_info)
            if unknown_edge_types:
                raise ValueError(
                    f"weight_edge_feat_name contains unknown edge types: {unknown_edge_types}"
                )
            configured_weights = [
                (edge_type, feature_name, edge_entity_info[edge_type])
                for edge_type, feature_name in weight_edge_feat_name.items()
            ]

    for edge_type, feature_name, serialized_info in configured_weights:
        if feature_name not in serialized_info.feature_keys:
            raise ValueError(
                f"Sampling-weight field '{feature_name}' for edge type {edge_type} must be an unquantized raw edge feature."
            )




def remove_sampling_weight_from_edge_quantization_metadata(
    serialized_graph_metadata: SerializedGraphMetadata,
    weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]],
) -> Optional[
    Union[FeatureQuantizationMetadata, dict[EdgeType, FeatureQuantizationMetadata]]
]:
    """Remove separately stored sampling weights from edge reconstruction metadata.

    TFRecord loading removes the sampling-weight column from raw edge features
    before registering it with the weighted sampler. The resulting metadata
    must describe the remaining model features so batch reconstruction scatters
    raw and dequantized columns into the correct positions.

    Args:
        serialized_graph_metadata: Serialized edge schema and quantization metadata.
        weight_edge_feat_name: Raw scalar feature configured as sampling weights.

    Returns:
        Quantization metadata for the model-facing edge features.
    """
    quantization_metadata = serialized_graph_metadata.edge_quantization_metadata
    if quantization_metadata is None or weight_edge_feat_name is None:
        return quantization_metadata

    if isinstance(serialized_graph_metadata.edge_entity_info, SerializedTFRecordInfo):
        assert isinstance(quantization_metadata, FeatureQuantizationMetadata)
        assert isinstance(weight_edge_feat_name, str)
        edge_info_by_type: dict[EdgeType, SerializedTFRecordInfo] = {
            DEFAULT_HOMOGENEOUS_EDGE_TYPE: serialized_graph_metadata.edge_entity_info
        }
        metadata_by_type: dict[EdgeType, FeatureQuantizationMetadata] = {
            DEFAULT_HOMOGENEOUS_EDGE_TYPE: quantization_metadata
        }
        weight_by_type: dict[EdgeType, str] = {
            DEFAULT_HOMOGENEOUS_EDGE_TYPE: weight_edge_feat_name
        }
        is_homogeneous = True
    else:
        assert isinstance(quantization_metadata, dict)
        edge_info_by_type: dict[EdgeType, SerializedTFRecordInfo] = (
            serialized_graph_metadata.edge_entity_info
        )
        metadata_by_type: dict[EdgeType, FeatureQuantizationMetadata] = cast(
            dict[EdgeType, FeatureQuantizationMetadata], quantization_metadata
        )
        if isinstance(weight_edge_feat_name, str):
            if len(edge_info_by_type) != 1:
                raise ValueError(
                    "weight_edge_feat_name must be dict[EdgeType, str] for "
                    "heterogeneous graph with multiple edge types"
                )
            edge_type = next(iter(edge_info_by_type))
            weight_by_type: dict[EdgeType, str] = {edge_type: weight_edge_feat_name}
        else:
            weight_by_type: dict[EdgeType, str] = weight_edge_feat_name
        is_homogeneous = False

    adjusted_metadata: dict[EdgeType, FeatureQuantizationMetadata] = {}
    for edge_type, metadata in metadata_by_type.items():
        weight_feature_name = weight_by_type.get(edge_type)
        if weight_feature_name is None:
            adjusted_metadata[edge_type] = metadata
            continue

        edge_info = edge_info_by_type[edge_type]
        raw_column_offset = 0
        for feature_name in edge_info.feature_keys:
            if feature_name == weight_feature_name:
                break
            feature_spec = edge_info.feature_spec[feature_name]
            raw_column_offset += feature_spec.shape[-1] if feature_spec.shape else 1
        # The sampling-weight column is removed before reconstruction, shifting
        # every following logical feature index one position to the left.
        weight_logical_index = metadata.raw_feature_indices[raw_column_offset]
        adjusted_quantized_feature_indices = tuple(
            quantized_feature_index - 1
            if quantized_feature_index > weight_logical_index
            else quantized_feature_index
            for quantized_feature_index in metadata.quantized_feature_indices
        )
        adjusted_metadata[edge_type] = replace(
            metadata,
            feature_dim=metadata.feature_dim - 1,
            quantized_feature_indices=adjusted_quantized_feature_indices,
        )

    if is_homogeneous:
        return adjusted_metadata[DEFAULT_HOMOGENEOUS_EDGE_TYPE]
    return adjusted_metadata



def _data_loading_process(
    tf_record_dataloader: TFRecordDataLoader,
    output_dict: MutableMapping[
        str, Union[torch.Tensor, dict[Union[NodeType, EdgeType], torch.Tensor]]
    ],
    error_dict: MutableMapping[str, str],
    entity_type: str,
    serialized_tf_record_info: Union[
        SerializedTFRecordInfo,
        dict[Union[NodeType, EdgeType], SerializedTFRecordInfo],
    ],
    rank: int,
    tf_dataset_options: TFDatasetOptions = TFDatasetOptions(),
    weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]] = None,
) -> None:
    """
    Spawned multiprocessing.Process which loads homogeneous or heterogeneous information for a specific entity type [node, edge, positive_label, negative_label]
    and moves to shared memory. Also logs timing information for duration of loading. If an exception is thrown, its traceback will be stored in
    the error_dict "error" field, since exceptions for spawned processes won't properly be raised to the parent process.

    Args:
        tf_record_dataloader (TFRecordDataLoader): TFRecordDataloader used for loading tensors from serialized tfrecords
        output_dict (MutableMapping[str, Union[torch.Tensor, dict[Union[NodeType, EdgeType], torch.Tensor]]]):
            Dictionary initialized by mp.Manager().dict() in which outputs of tensor loading will be written to
        error_dict (MutableMapping[str, str]): Dictionary initialized by mp.Manager().dict() in which error of errors in current process will be written to
        entity_type (str): Entity type to prefix ids, features, and error keys with when
            writing to the output_dict and error_dict fields
        serialized_tf_record_info (Union[SerializedTFRecordInfo, dict[NodeType, SerializedTFRecordInfo], dict[EdgeType, SerializedTFRecordInfo]]):
            Serialized information for current entity
        rank (int): Rank of the current machine
        tf_dataset_options (TFDatasetOptions): The options to use when building the dataset.
        weight_edge_feat_name (Optional[Union[str, dict[EdgeType, str]]]): Only used when
            ``entity_type == _EDGE_KEY``. Name of the edge feature column to extract as
            sampling weights. Ignored for node, positive_label, and negative_label entities.
            Supply a single string for homogeneous graphs or a per-edge-type dict for
            heterogeneous graphs.
    """
    # We add a try - except clause here to ensure that exceptions are properly circulated back to the parent process
    try:
        # To simplify the logic to proceed on a singular code path, we convert homogeneous inputs to heterogeneous just within the scope of this function
        if isinstance(serialized_tf_record_info, SerializedTFRecordInfo):
            serialized_tf_record_info = (
                {DEFAULT_HOMOGENEOUS_NODE_TYPE: serialized_tf_record_info}
                if serialized_tf_record_info.is_node_entity
                else {DEFAULT_HOMOGENEOUS_EDGE_TYPE: serialized_tf_record_info}
            )
            is_input_homogeneous = True
        else:
            is_input_homogeneous = False

        all_tf_record_uris = [
            serialized_entity.tfrecord_uri_prefix.uri
            for serialized_entity in serialized_tf_record_info.values()
        ]

        start_time = time.time()

        logger.info(
            f"Rank {rank} has begun to load data from tfrecord directories: {all_tf_record_uris}"
        )

        ids: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
        features: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
        quantized_features: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
        labels: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
        weights: dict[Union[NodeType, EdgeType], torch.Tensor] = {}
        for (
            graph_type,
            serialized_entity_tf_record_info,
        ) in serialized_tf_record_info.items():
            # We currently do not support training with labels for edge entities
            if (
                serialized_entity_tf_record_info.label_keys
                and not serialized_entity_tf_record_info.is_node_entity
            ):
                raise NotImplementedError(
                    "Label keys are not supported for edge entities"
                )
            loaded_entity = tf_record_dataloader.load_as_torch_tensors(
                serialized_tf_record_info=serialized_entity_tf_record_info,
                tf_dataset_options=tf_dataset_options,
            )
            entity_ids = loaded_entity.ids
            entity_features = loaded_entity.features
            entity_quantized_features = loaded_entity.quantized_features
            entity_labels = loaded_entity.labels
            ids[graph_type] = entity_ids
            logger.info(
                f"Rank {rank} finished loading {entity_type} ids of shape {entity_ids.shape} for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
            )
            if entity_features is not None:
                features[graph_type] = entity_features
                logger.info(
                    f"Rank {rank} finished loading {entity_type} features of shape {entity_features.shape} for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                )
            else:
                logger.info(
                    f"Rank {rank} did not detect {entity_type} features for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                )

            if entity_quantized_features is not None:
                quantized_features[graph_type] = entity_quantized_features
                logger.info(
                    f"Rank {rank} finished loading {entity_type} quantized features of shape {entity_quantized_features.shape} for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                )
            else:
                logger.info(
                    f"Rank {rank} did not detect {entity_type} quantized features for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                )

            if entity_labels is not None:
                labels[graph_type] = entity_labels
                logger.info(
                    f"Rank {rank} finished loading {entity_type} labels of shape {entity_labels.shape} for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                )
            else:
                logger.info(
                    f"Rank {rank} did not detect {entity_type} labels for graph type {graph_type} from {serialized_entity_tf_record_info.tfrecord_uri_prefix.uri}"
                )

        # Extract weight column from edge features when weight_edge_feat_name is set.
        # The weight column is sliced out of each edge type's feature tensor and stored
        # separately so it is not duplicated in the feature matrix.
        if weight_edge_feat_name is not None and entity_type == _EDGE_KEY:
            if isinstance(weight_edge_feat_name, str):
                if len(serialized_tf_record_info) != 1 or len(features) != 1:
                    raise ValueError(
                        f"weight_edge_feat_name must be a dict[EdgeType, str] for heterogeneous "
                        f"graphs with multiple edge types ({sorted(serialized_tf_record_info)}). "
                        "Provide an explicit per-edge-type mapping instead of a single string."
                    )
                col_name = weight_edge_feat_name
                edge_type, feat_tensor = next(iter(features.items()))
                assert isinstance(edge_type, EdgeType)
                feature_keys = list(serialized_tf_record_info[edge_type].feature_keys)
                weights[edge_type], trimmed = _extract_weight_col(
                    feat_tensor,
                    feature_keys,
                    serialized_tf_record_info[edge_type].feature_spec,
                    col_name,
                    edge_type,
                )
                if trimmed is not None:
                    features[edge_type] = trimmed
                else:
                    del features[edge_type]
                logger.info(
                    f"Rank {rank} extracted weight column '{col_name}' "
                    f"from {entity_type} features for type {edge_type}"
                )
            else:
                # Iterate the EdgeType-keyed dict directly to stay within EdgeType.
                for edge_type, col_name in weight_edge_feat_name.items():
                    if edge_type not in features:
                        continue
                    feat_tensor = features[edge_type]
                    feature_keys = list(
                        serialized_tf_record_info[edge_type].feature_keys
                    )
                    weights[edge_type], trimmed = _extract_weight_col(
                        feat_tensor,
                        feature_keys,
                        serialized_tf_record_info[edge_type].feature_spec,
                        col_name,
                        edge_type,
                    )
                    if trimmed is not None:
                        features[edge_type] = trimmed
                    else:
                        del features[edge_type]
                    logger.info(
                        f"Rank {rank} extracted weight column '{col_name}' "
                        f"from {entity_type} features for type {edge_type}"
                    )

        logger.info(
            f"Rank {rank} is attempting to share {entity_type} id memory for tfrecord directories: {all_tf_record_uris}"
        )
        share_memory(ids)
        # We convert the ids back to homogeneous from the default heterogeneous setup if our provided input was homogeneous

        if features:
            logger.info(
                f"Rank {rank} is attempting to share {entity_type} feature memory for tfrecord directories: {all_tf_record_uris}"
            )
            share_memory(features)
            # We convert the features back to homogeneous from the default heterogeneous setup if our provided input was homogeneous

        if quantized_features:
            logger.info(
                f"Rank {rank} is attempting to share {entity_type} quantized feature memory for tfrecord directories: {all_tf_record_uris}"
            )
            share_memory(quantized_features)

        if labels:
            logger.info(
                f"Rank {rank} is attempting to share {entity_type} label memory for tfrecord directories: {all_tf_record_uris}"
            )
            share_memory(labels)

        if weights:
            logger.info(
                f"Rank {rank} is attempting to share {entity_type} weight memory for tfrecord directories: {all_tf_record_uris}"
            )
            share_memory(weights)

        output_dict[_ID_FMT.format(entity=entity_type)] = (
            list(ids.values())[0] if is_input_homogeneous else ids
        )
        if features:
            output_dict[_FEATURE_FMT.format(entity=entity_type)] = (
                list(features.values())[0] if is_input_homogeneous else features
            )
        if quantized_features:
            output_dict[_PACKED_FEATURE_FMT.format(entity=entity_type)] = (
                list(quantized_features.values())[0]
                if is_input_homogeneous
                else quantized_features
            )
        if labels:
            output_dict[_LABEL_FMT.format(entity=entity_type)] = (
                list(labels.values())[0] if is_input_homogeneous else labels
            )
        if weights:
            output_dict[_EDGE_WEIGHTS_KEY] = (
                list(weights.values())[0] if is_input_homogeneous else weights
            )

        logger.info(
            f"Rank {rank} has finished loading {entity_type} data from tfrecord directories: {all_tf_record_uris}, elapsed time: {time.time() - start_time:.2f} seconds"
        )

    except Exception:
        error_dict[entity_type] = traceback.format_exc()




def load_torch_tensors_from_tf_record(
    tf_record_dataloader: TFRecordDataLoader,
    serialized_graph_metadata: SerializedGraphMetadata,
    should_load_tensors_in_parallel: bool,
    rank: int = 0,
    node_tf_dataset_options: TFDatasetOptions = TFDatasetOptions(),
    edge_tf_dataset_options: TFDatasetOptions = TFDatasetOptions(),
    weight_edge_feat_name: Optional[Union[str, dict[EdgeType, str]]] = None,
) -> LoadedGraphTensors:
    """
    Loads all torch tensors from a SerializedGraphMetadata object for all entity [node, edge, positive_label, negative_label] and edge / node types.

    Running these processes in parallel slows the runtime of each individual process, but may still result in a net speedup across all entity types. As a result,
    there is a tradeoff that needs to be made between parallel and sequential tensor loading, which is why we don't parallelize across node and edge types. We enable
    the `should_load_tensors_in_parallel` to allow some customization for loading strategies based on the input data.

    Args:
        tf_record_dataloader (TFRecordDataLoader): TFRecordDataloader used for loading tensors from serialized tfrecords
        serialized_graph_metadata (SerializedGraphMetadata): Serialized graph metadata contained serialized information for loading tfrecords across node and edge types
        should_load_tensors_in_parallel (bool): Whether tensors should be loaded from serialized information in parallel or in sequence across the [node, edge, pos_label, neg_label] entity types.
        rank (int): Rank on current machine
        node_tf_dataset_options (TFDatasetOptions): The options to use for nodes when building the dataset.
        edge_tf_dataset_options (TFDatasetOptions): The options to use for edges when building the dataset.
        weight_edge_feat_name (Optional[Union[str, dict[EdgeType, str]]]): Name of the edge feature column to extract
            as sampling weights. The column is removed from the edge feature matrix and returned separately via
            ``LoadedGraphTensors.edge_weights``. Supply a single string for homogeneous graphs or a per-edge-type
            dict for heterogeneous graphs.
    Returns:
        loaded_graph_tensors (LoadedGraphTensors): Unpartitioned Graph Tensors
    """

    _validate_weight_edge_feature_name(
        edge_entity_info=serialized_graph_metadata.edge_entity_info,
        weight_edge_feat_name=weight_edge_feat_name,
    )

    logger.info(f"Rank {rank} starting loading torch tensors from serialized info ...")
    start_time = time.time()

    manager = Manager()

    # By default, torch processes are created using the `fork` method, which makes a copy of the entire process. This can be problematic in multi-threaded settings,
    # especially when working with TensorFlow, since this includes all threads, which can lead to deadlocks or other synchronization issues. As a result, we set the
    # start method to spawn, which creates a new Python interpreter process and is much safer with multi-threading applications.
    ctx = mp.get_context("spawn")

    node_output_dict: MutableMapping[
        str, Union[torch.Tensor, dict[NodeType, torch.Tensor]]
    ] = manager.dict()

    edge_output_dict: MutableMapping[
        str, Union[torch.Tensor, dict[EdgeType, torch.Tensor]]
    ] = manager.dict()

    error_dict: MutableMapping[str, str] = manager.dict()

    node_data_loading_process = ctx.Process(
        target=_data_loading_process,
        kwargs={
            "tf_record_dataloader": tf_record_dataloader,
            "output_dict": node_output_dict,
            "error_dict": error_dict,
            "entity_type": _NODE_KEY,
            "serialized_tf_record_info": serialized_graph_metadata.node_entity_info,
            "rank": rank,
            "tf_dataset_options": node_tf_dataset_options,
        },
    )

    edge_data_loading_process = ctx.Process(
        target=_data_loading_process,
        kwargs={
            "tf_record_dataloader": tf_record_dataloader,
            "output_dict": edge_output_dict,
            "error_dict": error_dict,
            "entity_type": _EDGE_KEY,
            "serialized_tf_record_info": serialized_graph_metadata.edge_entity_info,
            "rank": rank,
            "tf_dataset_options": edge_tf_dataset_options,
            "weight_edge_feat_name": weight_edge_feat_name,
        },
    )

    if serialized_graph_metadata.positive_label_entity_info is not None:
        positive_label_data_loading_process = ctx.Process(
            target=_data_loading_process,
            kwargs={
                "tf_record_dataloader": tf_record_dataloader,
                "output_dict": edge_output_dict,
                "error_dict": error_dict,
                "entity_type": _POSITIVE_LABEL_KEY,
                "serialized_tf_record_info": serialized_graph_metadata.positive_label_entity_info,
                "rank": rank,
            },
        )
    else:
        logger.info(f"No positive labels detected from input data")

    if serialized_graph_metadata.negative_label_entity_info is not None:
        negative_label_data_loading_process = ctx.Process(
            target=_data_loading_process,
            kwargs={
                "tf_record_dataloader": tf_record_dataloader,
                "output_dict": edge_output_dict,
                "error_dict": error_dict,
                "entity_type": _NEGATIVE_LABEL_KEY,
                "serialized_tf_record_info": serialized_graph_metadata.negative_label_entity_info,
                "rank": rank,
            },
        )
    else:
        logger.info(f"No negative labels detected from input data")

    if should_load_tensors_in_parallel:
        # In this setting, we start all the processes at once and join them at the end to achieve parallelized tensor loading
        logger.info("Loading Serialized TFRecord Data in Parallel ...")
        node_data_loading_process.start()
        edge_data_loading_process.start()
        if serialized_graph_metadata.positive_label_entity_info is not None:
            positive_label_data_loading_process.start()
        if serialized_graph_metadata.negative_label_entity_info is not None:
            negative_label_data_loading_process.start()

        node_data_loading_process.join()
        edge_data_loading_process.join()
        if serialized_graph_metadata.positive_label_entity_info is not None:
            positive_label_data_loading_process.join()
        if serialized_graph_metadata.negative_label_entity_info is not None:
            negative_label_data_loading_process.join()
    else:
        # In this setting, we start and join each process one-at-a-time in order to achieve sequential tensor loading
        logger.info("Loading Serialized TFRecord Data in Sequence ...")
        # Here we launch edge loading process first since experimentally
        # we have found that, since edge data is larger than node data,
        # launching edge before launching node reduces peak memory usage
        # compared with first launching node loading and then launching edge loading.
        edge_data_loading_process.start()
        edge_data_loading_process.join()
        node_data_loading_process.start()
        node_data_loading_process.join()
        if serialized_graph_metadata.positive_label_entity_info is not None:
            positive_label_data_loading_process.start()
            positive_label_data_loading_process.join()
        if serialized_graph_metadata.negative_label_entity_info is not None:
            negative_label_data_loading_process.start()
            negative_label_data_loading_process.join()

    if error_dict:
        for entity_type, traceback in error_dict.items():
            logger.error(
                f"Identified error in {entity_type} data loading process: \n{traceback}"
            )
        raise ValueError(
            f"Raised error in data loading processes for entity types {error_dict.keys()}."
        )

    node_ids = node_output_dict[_ID_FMT.format(entity=_NODE_KEY)]
    node_features = node_output_dict.get(_FEATURE_FMT.format(entity=_NODE_KEY), None)
    node_quantized_features = node_output_dict.get(
        _PACKED_FEATURE_FMT.format(entity=_NODE_KEY), None
    )
    node_labels = node_output_dict.get(_LABEL_FMT.format(entity=_NODE_KEY), None)

    edge_index = edge_output_dict[_ID_FMT.format(entity=_EDGE_KEY)]
    edge_features = edge_output_dict.get(_FEATURE_FMT.format(entity=_EDGE_KEY), None)
    edge_quantized_features = edge_output_dict.get(
        _PACKED_FEATURE_FMT.format(entity=_EDGE_KEY), None
    )
    edge_weights = edge_output_dict.get(_EDGE_WEIGHTS_KEY, None)

    positive_labels = edge_output_dict.get(
        _ID_FMT.format(entity=_POSITIVE_LABEL_KEY), None
    )

    negative_labels = edge_output_dict.get(
        _ID_FMT.format(entity=_NEGATIVE_LABEL_KEY), None
    )

    if rpc_is_initialized():
        logger.info(
            f"Rank {rank} has finished loading data in {time.time() - start_time:.2f} seconds. Wait for other ranks to finish loading data from tfrecords"
        )
        barrier()

    logger.info(
        f"All ranks have finished loading data from tfrecords, rank {rank} finished in {time.time() - start_time:.2f} seconds"
    )

    return LoadedGraphTensors(
        node_ids=node_ids,
        node_features=node_features,
        node_quantized_features=node_quantized_features,
        node_labels=node_labels,
        edge_index=edge_index,
        edge_features=edge_features,
        edge_quantized_features=edge_quantized_features,
        positive_label=positive_labels,
        negative_label=negative_labels,
        edge_weights=edge_weights,
    )