wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/dgl/graphbolt/internal/__init__.py
+4
View File
@@ -0,0 +1,4 @@
"""Utility functions for GraphBolt."""
from .utils import *
from .sample_utils import *
from .item_sampler_utils import *
python/dgl/graphbolt/internal/item_sampler_utils.py
+112
View File
@@ -0,0 +1,112 @@
"""Utility functions for DistributedItemSampler."""
def count_split(total, num_workers, worker_id, batch_size=1):
"""Calculate the number of assigned items after splitting them by batch
size evenly. It will return the number for this worker and also a sum of
previous workers.
"""
quotient, remainder = divmod(total, num_workers * batch_size)
if batch_size == 1:
assigned = quotient + (worker_id < remainder)
else:
batch_count, last_batch = divmod(remainder, batch_size)
assigned = quotient * batch_size + (
batch_size
if worker_id < batch_count
else (last_batch if worker_id == batch_count else 0)
)
prefix_sum = quotient * worker_id * batch_size + min(
worker_id * batch_size, remainder
)
return (assigned, prefix_sum)
def calculate_range(
distributed,
total,
num_replicas,
rank,
num_workers,
worker_id,
batch_size,
drop_last,
drop_uneven_inputs,
):
"""Calculates the range of items to be assigned to the current worker.
This function evenly distributes `total` items among multiple workers,
batching them using `batch_size`. Each replica has `num_workers` workers.
The batches generated by workers within the same replica are combined into
the replica`s output. The `drop_last` parameter determines whether
incomplete batches should be dropped. If `drop_last` is True, incomplete
batches are discarded. The `drop_uneven_inputs` parameter determines if the
number of batches assigned to each replica should be the same. If
`drop_uneven_inputs` is True, excessive batches for some replicas will be
dropped.
Args:
distributed (bool): Whether it's in distributed mode.
total (int): The total number of items.
num_replicas (int): The total number of replicas.
rank (int): The rank of the current replica.
num_workers (int): The number of workers per replica.
worker_id (int): The ID of the current worker.
batch_size (int): The desired batch size.
drop_last (bool): Whether to drop incomplete batches.
drop_uneven_inputs (bool): Whether to drop excessive batches for some
replicas.
Returns:
tuple: A tuple containing three numbers:
- start_offset (int): The starting offset of the range assigned to
the current worker.
- assigned_count (int): The length of the range assigned to the
current worker.
- output_count (int): The number of items that the current worker
will produce after dropping.
"""
# Check if it's distributed mode.
if not distributed:
if not drop_last:
return (0, total, total)
else:
return (0, total, total // batch_size * batch_size)
# First, equally distribute items into all replicas.
assigned_count, start_offset = count_split(
total, num_replicas, rank, batch_size
)
# Calculate the number of outputs when drop_uneven_inputs is True.
# `assigned_count` is the number of items distributed to the current
# process. `output_count` is the number of items should be output
# by this process after dropping.
if not drop_uneven_inputs:
if not drop_last:
output_count = assigned_count
else:
output_count = assigned_count // batch_size * batch_size
else:
if not drop_last:
min_item_count, _ = count_split(
total, num_replicas, num_replicas - 1, batch_size
)
min_batch_count = (min_item_count + batch_size - 1) // batch_size
output_count = min(min_batch_count * batch_size, assigned_count)
else:
output_count = total // (batch_size * num_replicas) * batch_size
# If there are multiple workers, equally distribute the batches to
# all workers.
if num_workers > 1:
# Equally distribute the dropped number too.
dropped_items, prev_dropped_items = count_split(
assigned_count - output_count, num_workers, worker_id
)
output_count, prev_output_count = count_split(
output_count,
num_workers,
worker_id,
batch_size,
)
assigned_count = output_count + dropped_items
start_offset += prev_output_count + prev_dropped_items
return (start_offset, assigned_count, output_count)
python/dgl/graphbolt/internal/sample_utils.py
+568
View File
@@ -0,0 +1,568 @@
"""Utility functions for sampling."""
from collections import defaultdict
from typing import Dict, List, Optional, Tuple, Union
import torch
from ..base import CSCFormatBase, etype_str_to_tuple, expand_indptr
def unique_and_compact(
nodes: Union[
List[torch.Tensor],
Dict[str, List[torch.Tensor]],
],
rank: int = 0,
world_size: int = 1,
async_op: bool = False,
):
"""
Compact a list of nodes tensor. The `rank` and `world_size` parameters are
relevant when using Cooperative Minibatching, which was initially proposed
in `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and
was later first fully described in
`Cooperative Minibatching in Graph Neural Networks
<https://arxiv.org/abs/2310.12403>`__.
Cooperation between the GPUs eliminates duplicate work performed across the
GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
performing GNN minibatching.
When `world_size` is greater than 1, then the given ids are partitioned
between the available ranks. The ids corresponding to the given rank are
guaranteed to come before the ids of other ranks. To do this, the
partitioned ids are rotated backwards by the given rank so that the ids are
ordered as: `[rank, rank + 1, world_size, 0, ..., rank - 1]`. This is
supported only for Volta and later generation NVIDIA GPUs.
Parameters
----------
nodes : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
List of nodes for compacting.
the unique_and_compact will be done per type
- If `nodes` is a list of tensor: All the tensors will do unique and
compact together, usually it is used for homogeneous graph.
- If `nodes` is a list of dictionary: The keys should be node type and
the values should be corresponding nodes, the unique and compact will
be done per type, usually it is used for heterogeneous graph.
rank : int
The rank of the current process.
world_size : int
The number of processes.
async_op: bool
Boolean indicating whether the call is asynchronous. If so, the result
can be obtained by calling wait on the returned future.
Returns
-------
Tuple[unique_nodes, compacted_node_list, unique_nodes_offsets]
The Unique nodes (per type) of all nodes in the input. And the compacted
nodes list, where IDs inside are replaced with compacted node IDs.
"Compacted node list" indicates that the node IDs in the input node
list are replaced with mapped node IDs, where each type of node is
mapped to a contiguous space of IDs ranging from 0 to N.
The unique nodes offsets tensor partitions the unique_nodes tensor. Has
size `world_size + 1` and `unique_nodes[offsets[i]: offsets[i + 1]]`
belongs to the rank `(rank + i) % world_size`.
"""
is_heterogeneous = isinstance(nodes, dict)
if not is_heterogeneous:
homo_ntype = "a"
nodes = {homo_ntype: nodes}
nums = {}
concat_nodes, empties = [], []
for ntype, nodes_of_type in nodes.items():
nums[ntype] = [node.size(0) for node in nodes_of_type]
concat_nodes.append(torch.cat(nodes_of_type))
empties.append(concat_nodes[-1].new_empty(0))
unique_fn = (
torch.ops.graphbolt.unique_and_compact_batched_async
if async_op
else torch.ops.graphbolt.unique_and_compact_batched
)
results = unique_fn(concat_nodes, empties, empties, rank, world_size)
class _Waiter:
def __init__(self, future, ntypes, nums):
self.future = future
self.ntypes = ntypes
self.nums = nums
def wait(self):
"""Returns the stored value when invoked."""
results = self.future.wait() if async_op else self.future
ntypes = self.ntypes
nums = self.nums
# Ensure there is no memory leak.
self.future = self.ntypes = self.nums = None
unique, compacted, offsets = {}, {}, {}
for ntype, result in zip(ntypes, results):
(
unique[ntype],
concat_compacted,
_,
offsets[ntype],
) = result
compacted[ntype] = list(concat_compacted.split(nums[ntype]))
if is_heterogeneous:
return unique, compacted, offsets
else:
return (
unique[homo_ntype],
compacted[homo_ntype],
offsets[homo_ntype],
)
post_processer = _Waiter(results, nodes.keys(), nums)
if async_op:
return post_processer
else:
return post_processer.wait()
def compact_temporal_nodes(nodes, nodes_timestamp):
"""Compact a list of temporal nodes without unique.
Note that since there is no unique, the nodes and nodes_timestamp are simply
concatenated. And the compacted nodes are consecutive numbers starting from
0.
Parameters
----------
nodes : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
List of nodes for compacting.
the compact operator will be done per type
- If `nodes` is a list of tensor: All the tensors will compact together,
usually it is used for homogeneous graph.
- If `nodes` is a list of dictionary: The keys should be node type and
the values should be corresponding nodes, the compact will be done per
type, usually it is used for heterogeneous graph.
nodes_timestamp : List[torch.Tensor] or Dict[str, List[torch.Tensor]]
List of timestamps for compacting.
Returns
-------
Tuple[nodes, nodes_timestamp, compacted_node_list]
The concatenated nodes and nodes_timestamp, and the compacted nodes list,
where IDs inside are replaced with compacted node IDs.
"""
def _compact_per_type(per_type_nodes, per_type_nodes_timestamp):
nums = [node.size(0) for node in per_type_nodes]
per_type_nodes = torch.cat(per_type_nodes)
per_type_nodes_timestamp = torch.cat(per_type_nodes_timestamp)
compacted_nodes = torch.arange(
0,
per_type_nodes.numel(),
dtype=per_type_nodes.dtype,
device=per_type_nodes.device,
)
compacted_nodes = list(compacted_nodes.split(nums))
return per_type_nodes, per_type_nodes_timestamp, compacted_nodes
if isinstance(nodes, dict):
ret_nodes, ret_timestamp, compacted = {}, {}, {}
for ntype, nodes_of_type in nodes.items():
(
ret_nodes[ntype],
ret_timestamp[ntype],
compacted[ntype],
) = _compact_per_type(nodes_of_type, nodes_timestamp[ntype])
return ret_nodes, ret_timestamp, compacted
else:
return _compact_per_type(nodes, nodes_timestamp)
def unique_and_compact_csc_formats(
csc_formats: Union[
Tuple[torch.Tensor, torch.Tensor],
Dict[str, Tuple[torch.Tensor, torch.Tensor]],
],
unique_dst_nodes: Union[
torch.Tensor,
Dict[str, torch.Tensor],
],
rank: int = 0,
world_size: int = 1,
async_op: bool = False,
):
"""
Compact csc formats and return unique nodes (per type). The `rank` and
`world_size` parameters are relevant when using Cooperative Minibatching,
which was initially proposed in
`Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__
and was later first fully described in
`Cooperative Minibatching in Graph Neural Networks
<https://arxiv.org/abs/2310.12403>`__.
Cooperation between the GPUs eliminates duplicate work performed across the
GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when
performing GNN minibatching.
When `world_size` is greater than 1, then the given ids are partitioned
between the available ranks. The ids corresponding to the given rank are
guaranteed to come before the ids of other ranks. To do this, the
partitioned ids are rotated backwards by the given rank so that the ids are
ordered as: `[rank, rank + 1, world_size, 0, ..., rank - 1]`. This is
supported only for Volta and later generation NVIDIA GPUs.
Parameters
----------
csc_formats : Union[CSCFormatBase, Dict(str, CSCFormatBase)]
CSC formats representing source-destination edges.
- If `csc_formats` is a CSCFormatBase: It means the graph is
homogeneous. Also, indptr and indice in it should be torch.tensor
representing source and destination pairs in csc format. And IDs inside
are homogeneous ids.
- If `csc_formats` is a Dict[str, CSCFormatBase]: The keys
should be edge type and the values should be csc format node pairs.
And IDs inside are heterogeneous ids.
unique_dst_nodes: torch.Tensor or Dict[str, torch.Tensor]
Unique nodes of all destination nodes in the node pairs.
- If `unique_dst_nodes` is a tensor: It means the graph is homogeneous.
- If `csc_formats` is a dictionary: The keys are node type and the
values are corresponding nodes. And IDs inside are heterogeneous ids.
rank : int
The rank of the current process.
world_size : int
The number of processes.
async_op: bool
Boolean indicating whether the call is asynchronous. If so, the result
can be obtained by calling wait on the returned future.
Returns
-------
Tuple[unique_nodes, csc_formats, unique_nodes_offsets]
The compacted csc formats, where node IDs are replaced with mapped node
IDs, and the unique nodes (per type).
"Compacted csc formats" indicates that the node IDs in the input node
pairs are replaced with mapped node IDs, where each type of node is
mapped to a contiguous space of IDs ranging from 0 to N. The unique
nodes offsets tensor partitions the unique_nodes tensor. Has size
`world_size + 1` and `unique_nodes[offsets[i]: offsets[i + 1]]` belongs
to the rank `(rank + i) % world_size`.
Examples
--------
>>> import dgl.graphbolt as gb
>>> N1 = torch.LongTensor([1, 2, 2])
>>> N2 = torch.LongTensor([5, 5, 6])
>>> unique_dst = {
... "n1": torch.LongTensor([1, 2]),
... "n2": torch.LongTensor([5, 6])}
>>> csc_formats = {
... "n1:e1:n2": gb.CSCFormatBase(indptr=torch.tensor([0, 2, 3]),indices=N1),
... "n2:e2:n1": gb.CSCFormatBase(indptr=torch.tensor([0, 1, 3]),indices=N2)}
>>> unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
... csc_formats, unique_dst
... )
>>> print(unique_nodes)
{'n1': tensor([1, 2]), 'n2': tensor([5, 6])}
>>> print(compacted_csc_formats)
{"n1:e1:n2": CSCFormatBase(indptr=torch.tensor([0, 2, 3]),
indices=torch.tensor([0, 1, 1])),
"n2:e2:n1": CSCFormatBase(indptr=torch.tensor([0, 1, 3]),
indices=torch.Longtensor([0, 0, 1]))}
"""
is_homogeneous = not isinstance(csc_formats, dict)
if is_homogeneous:
csc_formats = {"_N:_E:_N": csc_formats}
if unique_dst_nodes is not None:
assert isinstance(
unique_dst_nodes, torch.Tensor
), "Edge type not supported in homogeneous graph."
unique_dst_nodes = {"_N": unique_dst_nodes}
# Collect all source and destination nodes for each node type.
indices = defaultdict(list)
device = None
for etype, csc_format in csc_formats.items():
if device is None:
device = csc_format.indices.device
src_type, _, dst_type = etype_str_to_tuple(etype)
assert len(unique_dst_nodes.get(dst_type, [])) + 1 == len(
csc_format.indptr
), "The seed nodes should correspond to indptr."
indices[src_type].append(csc_format.indices)
indices = {ntype: torch.cat(nodes) for ntype, nodes in indices.items()}
ntypes = set(indices.keys())
dtype = list(indices.values())[0].dtype
default_tensor = torch.tensor([], dtype=dtype, device=device)
indice_list = []
unique_dst_list = []
for ntype in ntypes:
indice_list.append(indices.get(ntype, default_tensor))
unique_dst_list.append(unique_dst_nodes.get(ntype, default_tensor))
dst_list = [torch.tensor([], dtype=dtype, device=device)] * len(
unique_dst_list
)
uniq_fn = (
torch.ops.graphbolt.unique_and_compact_batched_async
if async_op
else torch.ops.graphbolt.unique_and_compact_batched
)
results = uniq_fn(indice_list, dst_list, unique_dst_list, rank, world_size)
class _Waiter:
def __init__(self, future, csc_formats):
self.future = future
self.csc_formats = csc_formats
def wait(self):
"""Returns the stored value when invoked."""
results = self.future.wait() if async_op else self.future
csc_formats = self.csc_formats
# Ensure there is no memory leak.
self.future = self.csc_formats = None
unique_nodes = {}
compacted_indices = {}
offsets = {}
for i, ntype in enumerate(ntypes):
(
unique_nodes[ntype],
compacted_indices[ntype],
_,
offsets[ntype],
) = results[i]
compacted_csc_formats = {}
# Map back with the same order.
for etype, csc_format in csc_formats.items():
num_elem = csc_format.indices.size(0)
src_type, _, _ = etype_str_to_tuple(etype)
indice = compacted_indices[src_type][:num_elem]
indptr = csc_format.indptr
compacted_csc_formats[etype] = CSCFormatBase(
indptr=indptr, indices=indice
)
compacted_indices[src_type] = compacted_indices[src_type][
num_elem:
]
# Return singleton for a homogeneous graph.
if is_homogeneous:
compacted_csc_formats = list(compacted_csc_formats.values())[0]
unique_nodes = list(unique_nodes.values())[0]
offsets = list(offsets.values())[0]
return unique_nodes, compacted_csc_formats, offsets
post_processer = _Waiter(results, csc_formats)
if async_op:
return post_processer
else:
return post_processer.wait()
def _broadcast_timestamps(csc, dst_timestamps):
"""Broadcast the timestamp of each destination node to its corresponding
source nodes."""
return expand_indptr(
csc.indptr, node_ids=dst_timestamps, output_size=len(csc.indices)
)
def compact_csc_format(
csc_formats: Union[CSCFormatBase, Dict[str, CSCFormatBase]],
dst_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]],
dst_timestamps: Optional[
Union[torch.Tensor, Dict[str, torch.Tensor]]
] = None,
):
"""
Relabel the row (source) IDs in the csc formats into a contiguous range from
0 and return the original row node IDs per type.
Note that
1. The column (destination) IDs are included in the relabeled row IDs.
2. If there are repeated row IDs, they would not be uniqued and will be
treated as different nodes.
3. If `dst_timestamps` is given, the timestamp of each destination node will
be broadcasted to its corresponding source nodes.
Parameters
----------
csc_formats: Union[CSCFormatBase, Dict[str, CSCFormatBase]]
CSC formats representing source-destination edges.
- If `csc_formats` is a CSCFormatBase: It means the graph is
homogeneous. Also, indptr and indice in it should be torch.tensor
representing source and destination pairs in csc format. And IDs inside
are homogeneous ids.
- If `csc_formats` is a Dict[str, CSCFormatBase]: The keys
should be edge type and the values should be csc format node pairs.
And IDs inside are heterogeneous ids.
dst_nodes: Union[torch.Tensor, Dict[str, torch.Tensor]]
Nodes of all destination nodes in the node pairs.
- If `dst_nodes` is a tensor: It means the graph is homogeneous.
- If `dst_nodes` is a dictionary: The keys are node type and the
values are corresponding nodes. And IDs inside are heterogeneous ids.
dst_timestamps: Optional[Union[torch.Tensor, Dict[str, torch.Tensor]]]
Timestamps of all destination nodes in the csc formats.
If given, the timestamp of each destination node will be broadcasted
to its corresponding source nodes.
Returns
-------
Tuple[original_row_node_ids, compacted_csc_formats, ...]
A tensor of original row node IDs (per type) of all nodes in the input.
The compacted CSC formats, where node IDs are replaced with mapped node
IDs ranging from 0 to N.
The source timestamps (per type) of all nodes in the input if
`dst_timestamps` is given.
Examples
--------
>>> import dgl.graphbolt as gb
>>> csc_formats = {
... "n2:e2:n1": gb.CSCFormatBase(
... indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([5, 4, 6])
... ),
... "n1:e1:n1": gb.CSCFormatBase(
... indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([1, 2, 3])
... ),
... }
>>> dst_nodes = {"n1": torch.LongTensor([2, 4])}
>>> original_row_node_ids, compacted_csc_formats = gb.compact_csc_format(
... csc_formats, dst_nodes
... )
>>> original_row_node_ids
{'n1': tensor([2, 4, 1, 2, 3]), 'n2': tensor([5, 4, 6])}
>>> compacted_csc_formats
{'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
indices=tensor([0, 1, 2]),
), 'n1:e1:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
indices=tensor([2, 3, 4]),
)}
>>> csc_formats = {
... "n2:e2:n1": gb.CSCFormatBase(
... indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([5, 4, 6])
... ),
... "n1:e1:n1": gb.CSCFormatBase(
... indptr=torch.tensor([0, 1, 3]), indices=torch.tensor([1, 2, 3])
... ),
... }
>>> dst_nodes = {"n1": torch.LongTensor([2, 4])}
>>> original_row_node_ids, compacted_csc_formats = gb.compact_csc_format(
... csc_formats, dst_nodes
... )
>>> original_row_node_ids
{'n1': tensor([2, 4, 1, 2, 3]), 'n2': tensor([5, 4, 6])}
>>> compacted_csc_formats
{'n2:e2:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
indices=tensor([0, 1, 2]),
), 'n1:e1:n1': CSCFormatBase(indptr=tensor([0, 1, 3]),
indices=tensor([2, 3, 4]),
)}
>>> dst_timestamps = {"n1": torch.LongTensor([10, 20])}
>>> (
... original_row_node_ids,
... compacted_csc_formats,
... src_timestamps,
... ) = gb.compact_csc_format(csc_formats, dst_nodes, dst_timestamps)
>>> src_timestamps
{'n1': tensor([10, 20, 10, 20, 20]), 'n2': tensor([10, 20, 20])}
"""
is_homogeneous = not isinstance(csc_formats, dict)
has_timestamp = dst_timestamps is not None
if is_homogeneous:
if dst_nodes is not None:
assert isinstance(
dst_nodes, torch.Tensor
), "Edge type not supported in homogeneous graph."
assert len(dst_nodes) + 1 == len(
csc_formats.indptr
), "The seed nodes should correspond to indptr."
offset = dst_nodes.size(0)
original_row_ids = torch.cat((dst_nodes, csc_formats.indices))
compacted_csc_formats = CSCFormatBase(
indptr=csc_formats.indptr,
indices=(
torch.arange(
0,
csc_formats.indices.size(0),
device=csc_formats.indices.device,
)
+ offset
),
)
src_timestamps = None
if has_timestamp:
src_timestamps = torch.cat(
[
dst_timestamps,
_broadcast_timestamps(
compacted_csc_formats, dst_timestamps
),
]
)
else:
compacted_csc_formats = {}
src_timestamps = None
original_row_ids = {key: val.clone() for key, val in dst_nodes.items()}
if has_timestamp:
src_timestamps = {
key: val.clone() for key, val in dst_timestamps.items()
}
for etype, csc_format in csc_formats.items():
src_type, _, dst_type = etype_str_to_tuple(etype)
assert len(dst_nodes.get(dst_type, [])) + 1 == len(
csc_format.indptr
), "The seed nodes should correspond to indptr."
device = csc_format.indices.device
offset = original_row_ids.get(
src_type, torch.tensor([], device=device)
).size(0)
original_row_ids[src_type] = torch.cat(
(
original_row_ids.get(
src_type,
torch.tensor(
[], dtype=csc_format.indices.dtype, device=device
),
),
csc_format.indices,
)
)
compacted_csc_formats[etype] = CSCFormatBase(
indptr=csc_format.indptr,
indices=(
torch.arange(
0,
csc_format.indices.size(0),
dtype=csc_format.indices.dtype,
device=device,
)
+ offset
),
)
if has_timestamp:
# If destination timestamps are given, broadcast them to the
# corresponding source nodes.
src_timestamps[src_type] = torch.cat(
(
src_timestamps.get(
src_type,
torch.tensor(
[],
dtype=dst_timestamps[dst_type].dtype,
device=device,
),
),
_broadcast_timestamps(
csc_format, dst_timestamps[dst_type]
),
)
)
if has_timestamp:
return original_row_ids, compacted_csc_formats, src_timestamps
return original_row_ids, compacted_csc_formats
python/dgl/graphbolt/internal/utils.py
+216
View File
@@ -0,0 +1,216 @@
"""Utility functions for GraphBolt."""
import hashlib
import json
import os
import shutil
from typing import List, Union
import numpy as np
import pandas as pd
import torch
from numpy.lib.format import read_array_header_1_0, read_array_header_2_0
def numpy_save_aligned(*args, **kwargs):
"""A wrapper for numpy.save(), ensures the array is stored 4KiB aligned."""
# https://github.com/numpy/numpy/blob/2093a6d5b933f812d15a3de0eafeeb23c61f948a/numpy/lib/format.py#L179
has_array_align = hasattr(np.lib.format, "ARRAY_ALIGN")
if has_array_align:
default_alignment = np.lib.format.ARRAY_ALIGN
# The maximum allowed alignment by the numpy code linked above is 4K.
# Most filesystems work with block sizes of 4K so in practice, the file
# size on the disk won't be larger.
np.lib.format.ARRAY_ALIGN = 4096
np.save(*args, **kwargs)
if has_array_align:
np.lib.format.ARRAY_ALIGN = default_alignment
def _read_torch_data(path):
return torch.load(path, weights_only=False)
def _read_numpy_data(path, in_memory=True):
if in_memory:
return torch.from_numpy(np.load(path))
return torch.as_tensor(np.load(path, mmap_mode="r+"))
def read_data(path, fmt, in_memory=True):
"""Read data from disk."""
if fmt == "torch":
return _read_torch_data(path)
elif fmt == "numpy":
return _read_numpy_data(path, in_memory=in_memory)
else:
raise RuntimeError(f"Unsupported format: {fmt}")
def save_data(data, path, fmt):
"""Save data into disk."""
# Make sure the directory exists.
os.makedirs(os.path.dirname(path), exist_ok=True)
if fmt not in ["numpy", "torch"]:
raise RuntimeError(f"Unsupported format: {fmt}")
# Perform necessary conversion.
if fmt == "numpy" and isinstance(data, torch.Tensor):
data = data.cpu().numpy()
elif fmt == "torch" and isinstance(data, np.ndarray):
data = torch.from_numpy(data).cpu()
# Save the data.
if fmt == "numpy":
if not data.flags["C_CONTIGUOUS"]:
Warning(
"The ndarray saved to disk is not contiguous, "
"so it will be copied to contiguous memory."
)
data = np.ascontiguousarray(data)
numpy_save_aligned(path, data)
elif fmt == "torch":
if not data.is_contiguous():
Warning(
"The tensor saved to disk is not contiguous, "
"so it will be copied to contiguous memory."
)
data = data.contiguous()
torch.save(data, path)
def get_npy_dim(npy_path):
"""Get the dim of numpy file."""
with open(npy_path, "rb") as f:
# For the read_array_header API provided by numpy will only read the
# length of the header, it will cause parsing failure and error if
# first 8 bytes which contains magin string and version are not read
# ahead of time. So, we need to make sure we have skipped these 8
# bytes.
f.seek(8, 0)
try:
shape, _, _ = read_array_header_1_0(f)
except ValueError:
try:
shape, _, _ = read_array_header_2_0(f)
except ValueError:
raise ValueError("Invalid file format")
return len(shape)
def _to_int32(data):
if isinstance(data, torch.Tensor):
return data.to(torch.int32)
elif isinstance(data, np.ndarray):
return data.astype(np.int32)
else:
raise TypeError(
"Unsupported input type. Please provide a torch tensor or numpy array."
)
def copy_or_convert_data(
input_path,
output_path,
input_format,
output_format="numpy",
in_memory=True,
is_feature=False,
within_int32=False,
):
"""Copy or convert the data from input_path to output_path."""
assert (
output_format == "numpy"
), "The output format of the data should be numpy."
os.makedirs(os.path.dirname(output_path), exist_ok=True)
# We read the data always in case we need to cast its type.
data = read_data(input_path, input_format, in_memory)
if within_int32:
data = _to_int32(data)
if input_format == "numpy":
# If dim of the data is 1, reshape it to n * 1 and save it to output_path.
if is_feature and get_npy_dim(input_path) == 1:
data = data.reshape(-1, 1)
# If the data does not need to be modified, just copy the file.
elif not within_int32 and data.numpy().flags["C_CONTIGUOUS"]:
shutil.copyfile(input_path, output_path)
return
else:
# If dim of the data is 1, reshape it to n * 1 and save it to output_path.
if is_feature and data.dim() == 1:
data = data.reshape(-1, 1)
save_data(data, output_path, output_format)
def read_edges(dataset_dir, edge_fmt, edge_path):
"""Read egde data from numpy or csv."""
assert edge_fmt in [
"numpy",
"csv",
], f"`numpy` or `csv` is expected when reading edges but got `{edge_fmt}`."
if edge_fmt == "numpy":
edge_data = read_data(
os.path.join(dataset_dir, edge_path),
edge_fmt,
)
assert (
edge_data.shape[0] == 2 and len(edge_data.shape) == 2
), f"The shape of edges should be (2, N), but got {edge_data.shape}."
src, dst = edge_data.numpy()
else:
edge_data = pd.read_csv(
os.path.join(dataset_dir, edge_path),
names=["src", "dst"],
)
src, dst = edge_data["src"].to_numpy(), edge_data["dst"].to_numpy()
return (src, dst)
def calculate_file_hash(file_path, hash_algo="md5"):
"""Calculate the hash value of a file."""
hash_algos = ["md5", "sha1", "sha224", "sha256", "sha384", "sha512"]
if hash_algo in hash_algos:
hash_obj = getattr(hashlib, hash_algo)()
else:
raise ValueError(
f"Hash algorithm must be one of: {hash_algos}, but got `{hash_algo}`."
)
with open(file_path, "rb") as file:
for chunk in iter(lambda: file.read(4096), b""):
hash_obj.update(chunk)
return hash_obj.hexdigest()
def calculate_dir_hash(
dir_path, hash_algo="md5", ignore: Union[str, List[str]] = None
):
"""Calculte the hash values of all files under the directory."""
hashes = {}
for dirpath, _, filenames in os.walk(dir_path):
for filename in filenames:
if ignore and filename in ignore:
continue
filepath = os.path.join(dirpath, filename)
file_hash = calculate_file_hash(filepath, hash_algo=hash_algo)
hashes[filepath] = file_hash
return hashes
def check_dataset_change(dataset_dir, processed_dir):
"""Check whether dataset has been changed by checking its hash value."""
hash_value_file = "dataset_hash_value.txt"
hash_value_file_path = os.path.join(
dataset_dir, processed_dir, hash_value_file
)
if not os.path.exists(hash_value_file_path):
return True
with open(hash_value_file_path, "r") as f:
oringinal_hash_value = json.load(f)
present_hash_value = calculate_dir_hash(dataset_dir, ignore=hash_value_file)
if oringinal_hash_value == present_hash_value:
force_preprocess = False
else:
force_preprocess = True
return force_preprocess