import time
import traceback
from dataclasses import dataclass, replace
from typing import (
Generic,
Mapping,
MutableMapping,
NamedTuple,
Optional,
TypeVar,
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.common.utils.feature_quantization.lossless_compression import (
FeatureCompression,
compress_loaded_features,
resolve_feature_compression,
)
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
_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)
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 _shift_indices_past_removed_column(
indices: tuple[int, ...], removed_index: int
) -> tuple[int, ...]:
"""Shifts every index after ``removed_index`` one position to the left."""
shifted_indices: list[int] = []
for index in indices:
if index > removed_index:
shifted_indices.append(index - 1)
else:
shifted_indices.append(index)
return tuple(shifted_indices)
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] = {}
packed_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_packed_features = loaded_entity.packed_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_packed_features is not None:
packed_features[graph_type] = entity_packed_features
logger.info(
f"Rank {rank} finished loading {entity_type} packed features of shape {entity_packed_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} packed 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 packed_features:
logger.info(
f"Rank {rank} is attempting to share {entity_type} packed feature memory for tfrecord directories: {all_tf_record_uris}"
)
share_memory(packed_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 packed_features:
output_dict[_PACKED_FEATURE_FMT.format(entity=entity_type)] = (
list(packed_features.values())[0]
if is_input_homogeneous
else packed_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()
# Node features are keyed by NodeType and edge features by EdgeType. A constrained TypeVar keeps that key type
# through the helper: ``dict`` is invariant in its key, so a dict keyed by NodeType is not a dict keyed by
# ``Union[NodeType, EdgeType]``.
_EntityKey = TypeVar("_EntityKey", NodeType, EdgeType)
class _EntityFeatures(NamedTuple, Generic[_EntityKey]):
"""Raw features, packed features and quantization metadata of one entity type."""
features: Optional[Union[torch.Tensor, dict[_EntityKey, torch.Tensor]]]
packed_features: Optional[Union[torch.Tensor, dict[_EntityKey, torch.Tensor]]]
metadata: Optional[
Union[
FeatureQuantizationMetadata,
dict[_EntityKey, FeatureQuantizationMetadata],
]
]
def _read_and_compress_features(
output_dict: Mapping[str, Union[torch.Tensor, dict[_EntityKey, torch.Tensor]]],
error_dict: Mapping[str, str],
entity_type: str,
metadata: Optional[
Union[
FeatureQuantizationMetadata,
dict[_EntityKey, FeatureQuantizationMetadata],
]
],
should_compress: bool,
is_homogeneous: bool,
default_key: _EntityKey,
) -> _EntityFeatures[_EntityKey]:
"""Reads the raw and packed features of one entity type once its loading process has joined.
Compresses them with ``compress_loaded_features`` if ``should_compress`` is set and this entity type's own
loading process did not fail.
If this entity type's load failed, compression is skipped so the caller raises the stored load error once all
loading processes have joined. The other ranks wait until the RPC timeout either way, as at the load barrier.
Args:
output_dict: Outputs of the loading process of ``entity_type``.
error_dict: Tracebacks of the loading processes that failed, keyed by entity type.
entity_type: Entity type whose features are read.
metadata: Quantization metadata describing the packed features.
should_compress: Whether to compress the features.
is_homogeneous: Whether the graph is homogeneous, which decides the form of ``metadata``.
default_key: Key the homogeneous graph's single node or edge type uses when tensors are keyed.
Returns:
The raw features, packed features and metadata, after compression if it ran.
"""
features = output_dict.get(_FEATURE_FMT.format(entity=entity_type), None)
packed_features = output_dict.get(
_PACKED_FEATURE_FMT.format(entity=entity_type), None
)
# A failed load wrote no output, so compressing would raise a second, confusing error instead of the stored
# load error that the caller raises after all joins. Only this entity's failure is checked: another entity's
# process can fail at any time while this one compresses.
if not should_compress or entity_type in error_dict:
return _EntityFeatures(features, packed_features, metadata)
features, packed_features, metadata = compress_loaded_features(
features=features,
quantized_features=packed_features,
metadata=metadata,
is_homogeneous=is_homogeneous,
default_key=default_key,
)
return _EntityFeatures(features, packed_features, metadata)
[docs]
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.
If RPC is initialized, node and edge feature columns are losslessly compressed, each right after its own loading
process joins, unless the ``GIGL_FEATURE_COMPRESSION`` environment variable is ``off``
(see :mod:`gigl.common.utils.feature_quantization.lossless_compression`). When loading in parallel, node
compression runs while the edge loading process is still reading. The sampling-weight edge feature is removed
from the edge features before they are compressed, so it is never compressed.
With compression on and RPC initialized, every rank must call this function, with the same
``should_load_tensors_in_parallel``, because compression joins a collective RPC ``all_gather``.
Without RPC, features are returned uncompressed.
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, with the quantization metadata that
describes their packed features.
Raises:
ValueError: If a loading process failed, or if ``GIGL_FEATURE_COMPRESSION`` holds an unknown value while RPC
is initialized.
"""
_validate_weight_edge_feature_name(
edge_entity_info=serialized_graph_metadata.edge_entity_info,
weight_edge_feat_name=weight_edge_feat_name,
)
should_compress = (
rpc_is_initialized()
and resolve_feature_compression() == FeatureCompression.AUTO
)
node_is_homogeneous = isinstance(
serialized_graph_metadata.node_entity_info, SerializedTFRecordInfo
)
edge_is_homogeneous = isinstance(
serialized_graph_metadata.edge_entity_info, SerializedTFRecordInfo
)
node_quantization_metadata = serialized_graph_metadata.node_quantization_metadata
edge_quantization_metadata = remove_sampling_weight_from_edge_quantization_metadata(
serialized_graph_metadata=serialized_graph_metadata,
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()
# We do this because in practice nodes get loaded quicker, and we can get a head start on the compression
# while the edge loading process is still reading.
node_features, node_packed_features, node_quantization_metadata = (
_read_and_compress_features(
output_dict=node_output_dict,
error_dict=error_dict,
entity_type=_NODE_KEY,
metadata=node_quantization_metadata,
should_compress=should_compress,
is_homogeneous=node_is_homogeneous,
default_key=DEFAULT_HOMOGENEOUS_NODE_TYPE,
)
)
edge_data_loading_process.join()
edge_features, edge_packed_features, edge_quantization_metadata = (
_read_and_compress_features(
output_dict=edge_output_dict,
error_dict=error_dict,
entity_type=_EDGE_KEY,
metadata=edge_quantization_metadata,
should_compress=should_compress,
is_homogeneous=edge_is_homogeneous,
default_key=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
)
)
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()
edge_features, edge_packed_features, edge_quantization_metadata = (
_read_and_compress_features(
output_dict=edge_output_dict,
error_dict=error_dict,
entity_type=_EDGE_KEY,
metadata=edge_quantization_metadata,
should_compress=should_compress,
is_homogeneous=edge_is_homogeneous,
default_key=DEFAULT_HOMOGENEOUS_EDGE_TYPE,
)
)
node_data_loading_process.start()
node_data_loading_process.join()
node_features, node_packed_features, node_quantization_metadata = (
_read_and_compress_features(
output_dict=node_output_dict,
error_dict=error_dict,
entity_type=_NODE_KEY,
metadata=node_quantization_metadata,
should_compress=should_compress,
is_homogeneous=node_is_homogeneous,
default_key=DEFAULT_HOMOGENEOUS_NODE_TYPE,
)
)
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_labels = node_output_dict.get(_LABEL_FMT.format(entity=_NODE_KEY), None)
edge_index = edge_output_dict[_ID_FMT.format(entity=_EDGE_KEY)]
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_packed_features=node_packed_features,
node_labels=node_labels,
edge_index=edge_index,
edge_features=edge_features,
edge_packed_features=edge_packed_features,
positive_label=positive_labels,
negative_label=negative_labels,
edge_weights=edge_weights,
node_quantization_metadata=node_quantization_metadata,
edge_quantization_metadata=edge_quantization_metadata,
)
