k3-node 1.0.0__py3-none-any.whl

k3_node/loader/link_neighbor_loader.py

@@ -0,0 +1,202 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+
+try:
+    import torch
+    from torch import Tensor
+except ImportError:
+    torch = None
+    Tensor = type(None)
+
+from k3_node.data import Data, HeteroData
+from k3_node.loader.link_loader import EdgeSamplerInput, LinkLoader
+from k3_node.loader.node_loader import HeteroSamplerOutput, SamplerOutput
+from k3_node.loader.sampler_utils import FastGraph, sample_neighbors_hetero, sample_neighbors_homo
+
+
+class InternalLinkNeighborSampler:
+    r"""Pure Python/NumPy link neighborhood sampler."""
+    def __init__(
+        self,
+        data: Union[Data, HeteroData],
+        num_neighbors: Union[List[int], Dict[Tuple[str, str, str], List[int]]],
+        replace: bool = False,
+        subgraph_type: str = 'directional',
+        disjoint: bool = False,
+        neg_sampling_ratio: float = 0.0,
+    ):
+        self.data = data
+        self.num_neighbors = num_neighbors
+        self.replace = replace
+        self.subgraph_type = subgraph_type
+        self.disjoint = disjoint
+        self.neg_sampling_ratio = neg_sampling_ratio
+        self.edge_permutation = None
+
+    def _cached_graph(self):
+        # The CSR index depends only on the graph; build it once instead of per batch.
+        if getattr(self, "_graph", None) is None:
+            self._graph = FastGraph(self.data.edge_index, num_nodes=self.data.num_nodes)
+        return self._graph
+
+    def sample_from_edges(self, input_data: EdgeSamplerInput) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        is_torch = torch is not None and isinstance(input_data.row, Tensor)
+
+        pos_row = input_data.row
+        pos_col = input_data.col
+        pos_len = pos_row.size(0) if is_torch else len(pos_row)
+
+        if self.neg_sampling_ratio > 0:
+            num_neg = round(self.neg_sampling_ratio * pos_len)
+            num_total_nodes = self.data.num_nodes
+            if is_torch:
+                # As in PyG, both endpoints of a negative are random nodes
+                neg_row = torch.randint(0, num_total_nodes, (num_neg,), dtype=torch.long, device=pos_row.device)
+                neg_col = torch.randint(0, num_total_nodes, (num_neg,), dtype=torch.long, device=pos_row.device)
+                total_row = torch.cat([pos_row, neg_row], dim=0)
+                total_col = torch.cat([pos_col, neg_col], dim=0)
+                edge_label = torch.cat([torch.ones(pos_len, dtype=torch.float, device=pos_row.device),
+                                        torch.zeros(num_neg, dtype=torch.float, device=pos_row.device)], dim=0)
+            else:
+                neg_row = np.random.randint(0, num_total_nodes, size=(num_neg,), dtype=np.int64)
+                neg_col = np.random.randint(0, num_total_nodes, size=(num_neg,), dtype=np.int64)
+                total_row = np.concatenate([pos_row, neg_row], axis=0)
+                total_col = np.concatenate([pos_col, neg_col], axis=0)
+                edge_label = np.concatenate([np.ones(pos_len, dtype=np.float32), np.zeros(num_neg, dtype=np.float32)], axis=0)
+        else:
+            total_row = pos_row
+            total_col = pos_col
+            edge_label = input_data.label
+
+        # The sampled subgraph starts with the (sorted, unique) seed nodes, so `edge_label_index`
+        # is relabeled to their positions, i.e. to local node ids, as in PyG.
+        def unique_inverse(values):
+            if is_torch:
+                return torch.unique(values, return_inverse=True)
+            return np.unique(np.asarray(values), return_inverse=True)
+
+        def stack(a, b):
+            return torch.stack([a, b], dim=0) if is_torch else np.stack([a, b], axis=0)
+
+        if is_torch:
+            seed_nodes, inverse = unique_inverse(torch.cat([total_row, total_col], dim=0))
+        else:
+            seed_nodes, inverse = unique_inverse(np.concatenate([total_row, total_col], axis=0))
+        edge_label_index = stack(inverse[:len(total_row)], inverse[len(total_row):])
+
+        if isinstance(self.data, Data):
+            node, row, col, edge, n_counts, e_counts = sample_neighbors_homo(
+                edge_index=self.data.edge_index,
+                seed_nodes=seed_nodes,
+                num_neighbors=self.num_neighbors,
+                num_nodes=self.data.num_nodes,
+                graph=self._cached_graph(),
+                replace=self.replace,
+                subgraph_type=self.subgraph_type,
+                disjoint=self.disjoint,
+            )
+            return SamplerOutput(
+                node=node,
+                row=row,
+                col=col,
+                edge=edge,
+                num_sampled_nodes=n_counts,
+                num_sampled_edges=e_counts,
+                metadata=(input_data.input_id, edge_label_index, edge_label),
+            )
+        elif isinstance(self.data, HeteroData):
+            edge_index_dict = {k: self.data[k].edge_index for k in self.data.edge_types}
+            input_type = input_data.input_type or self.data.edge_types[0]
+            src_type, _, dst_type = input_type
+
+            seed_dict = {k: None for k in self.data.node_types}
+            if src_type == dst_type:
+                seed_dict[src_type] = seed_nodes
+            else:
+                seed_dict[src_type], src_inverse = unique_inverse(total_row)
+                seed_dict[dst_type], dst_inverse = unique_inverse(total_col)
+                edge_label_index = stack(src_inverse, dst_inverse)
+
+            norm_num_neighbors = self.num_neighbors
+            if isinstance(norm_num_neighbors, dict):
+                norm_num_neighbors = {self.data._to_canonical(*k): v for k, v in norm_num_neighbors.items()}
+
+            node_dict, row_dict, col_dict, edge_dict, n_counts, e_counts = sample_neighbors_hetero(
+                edge_index_dict=edge_index_dict,
+                seed_nodes_dict=seed_dict,
+                num_neighbors=norm_num_neighbors,
+                replace=self.replace,
+                subgraph_type=self.subgraph_type,
+            )
+            return HeteroSamplerOutput(
+                node=node_dict,
+                row=row_dict,
+                col=col_dict,
+                edge=edge_dict,
+                num_sampled_nodes=n_counts,
+                num_sampled_edges=e_counts,
+                metadata=(input_data.input_id, edge_label_index, edge_label),
+            )
+
+        raise TypeError(f"Invalid data type: {type(self.data)}")
+
+
+class LinkNeighborLoader(LinkLoader):
+    r"""A link-based data loader derived as an extension of NeighborLoader."""
+    def __init__(
+        self,
+        data: Union[Data, HeteroData],
+        num_neighbors: Union[List[int], Dict[Tuple[str, str, str], List[int]]],
+        edge_label_index: Any = None,
+        edge_label: Optional[Any] = None,
+        edge_label_time: Optional[Any] = None,
+        replace: bool = False,
+        subgraph_type: str = 'directional',
+        disjoint: bool = False,
+        temporal_strategy: str = 'uniform',
+        neg_sampling: Optional[Any] = None,
+        neg_sampling_ratio: Optional[Union[int, float]] = None,
+        time_attr: Optional[str] = None,
+        weight_attr: Optional[str] = None,
+        transform: Optional[Callable] = None,
+        transform_sampler_output: Optional[Callable] = None,
+        is_sorted: bool = False,
+        filter_per_worker: Optional[bool] = None,
+        neighbor_sampler: Optional[Any] = None,
+        directed: bool = True,
+        **kwargs,
+    ):
+        if not directed:
+            subgraph_type = 'induced'
+
+        ratio = 0.0
+        if neg_sampling_ratio is not None:
+            ratio = float(neg_sampling_ratio)
+        elif neg_sampling is not None:
+            ratio = float(getattr(neg_sampling, 'amount', 1.0)) if hasattr(neg_sampling, 'amount') else 1.0
+
+        if neighbor_sampler is None:
+            neighbor_sampler = InternalLinkNeighborSampler(
+                data,
+                num_neighbors=num_neighbors,
+                replace=replace,
+                subgraph_type=str(getattr(subgraph_type, 'value', subgraph_type)),
+                disjoint=disjoint,
+                neg_sampling_ratio=ratio,
+            )
+
+        super().__init__(
+            data=data,
+            link_sampler=neighbor_sampler,
+            edge_label_index=edge_label_index,
+            edge_label=edge_label,
+            edge_label_time=edge_label_time,
+            neg_sampling=neg_sampling,
+            neg_sampling_ratio=neg_sampling_ratio,
+            transform=transform,
+            transform_sampler_output=transform_sampler_output,
+            filter_per_worker=filter_per_worker,
+            **kwargs,
+        )
+

k3_node/loader/mixin.py

@@ -0,0 +1,190 @@
+import glob
+import logging
+import os
+import os.path as osp
+import warnings
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, List, Optional, Union
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+
+def get_numa_nodes_cores() -> Dict[str, Any]:
+    """Parses numa nodes information into a dictionary."""
+    numa_node_paths = glob.glob('/sys/devices/system/node/node[0-9]*')
+    if not numa_node_paths:
+        return {}
+
+    nodes = {}
+    try:
+        for node_path in numa_node_paths:
+            numa_node_id = int(osp.basename(node_path)[4:])
+            thread_siblings = {}
+            for cpu_dir in glob.glob(osp.join(node_path, 'cpu[0-9]*')):
+                cpu_id = int(osp.basename(cpu_dir)[3:])
+                if cpu_id > 0:
+                    with open(osp.join(cpu_dir, 'online')) as core_online_file:
+                        core_online = int(core_online_file.read().splitlines()[0])
+                else:
+                    core_online = 1  # cpu0 is always online
+                if core_online == 1:
+                    with open(osp.join(cpu_dir, 'topology', 'core_id')) as core_id_file:
+                        core_id = int(core_id_file.read().strip())
+                        if core_id in thread_siblings:
+                            thread_siblings[core_id].append(cpu_id)
+                        else:
+                            thread_siblings[core_id] = [cpu_id]
+
+            nodes[numa_node_id] = sorted([(k, sorted(v)) for k, v in thread_siblings.items()])
+    except (OSError, ValueError, IndexError):
+        warnings.warn('Failed to read NUMA info')
+        return {}
+
+    return nodes
+
+
+class WorkerInitWrapper:
+    r"""Wraps the :attr:`worker_init_fn` argument for DataLoader workers."""
+    def __init__(self, func: Optional[Callable]) -> None:
+        self.func = func
+
+    def __call__(self, worker_id: int) -> None:
+        if self.func is not None:
+            self.func(worker_id)
+
+
+class LogMemoryMixin:
+    r"""A context manager to enable logging of memory consumption in
+    DataLoader workers.
+    """
+    def _mem_init_fn(self, worker_id: int) -> None:
+        if psutil is not None:
+            proc = psutil.Process(os.getpid())
+            memory = proc.memory_info().rss / (1024 * 1024)
+            logging.debug(f"Worker {worker_id} @ PID {proc.pid}: {memory:.2f} MB")
+        self._old_worker_init_fn(worker_id)
+
+    @contextmanager
+    def enable_memory_log(self):
+        self._old_worker_init_fn = WorkerInitWrapper(getattr(self, 'worker_init_fn', None))
+        try:
+            self.worker_init_fn = self._mem_init_fn
+            yield
+        finally:
+            self.worker_init_fn = self._old_worker_init_fn
+
+
+class MultithreadingMixin:
+    r"""A context manager to enable multi-threading in DataLoader workers."""
+    def _mt_init_fn(self, worker_id: int) -> None:
+        if torch is not None:
+            try:
+                torch.set_num_threads(int(self._worker_threads))
+            except IndexError as e:
+                raise ValueError(f"Cannot set {self._worker_threads} threads in worker {worker_id}") from e
+        self._old_worker_init_fn(worker_id)
+
+    @contextmanager
+    def enable_multithreading(self, worker_threads: Optional[int] = None):
+        num_workers = getattr(self, 'num_workers', 0)
+        if not num_workers > 0:
+            raise ValueError(f"'enable_multithreading' needs to be performed with at least one worker (got {num_workers})")
+
+        if torch is not None:
+            if worker_threads is None:
+                worker_threads = torch.get_num_threads() // num_workers
+            if worker_threads > torch.get_num_threads():
+                raise ValueError(
+                    f"'worker_threads' should be smaller than total available threads {torch.get_num_threads()} (got {worker_threads})"
+                )
+            context = torch.multiprocessing.get_context()._name
+            if context != 'spawn':
+                raise ValueError(f"'enable_multithreading' can only be used with 'spawn' multiprocessing context (got {context})")
+        else:
+            if worker_threads is None:
+                worker_threads = 1
+
+        self._worker_threads = worker_threads
+        self._old_worker_init_fn = WorkerInitWrapper(getattr(self, 'worker_init_fn', None))
+        try:
+            logging.debug(f"Using {worker_threads} threads in each worker")
+            self.worker_init_fn = self._mt_init_fn
+            yield
+        finally:
+            self.worker_init_fn = self._old_worker_init_fn
+
+
+class AffinityMixin:
+    r"""A context manager to enable CPU affinity for data loader workers."""
+    def _aff_init_fn(self, worker_id: int) -> None:
+        try:
+            worker_cores = self.loader_cores[worker_id]
+            if not isinstance(worker_cores, list):
+                worker_cores = [worker_cores]
+
+            if torch is not None and torch.multiprocessing.get_context()._name == 'spawn':
+                torch.set_num_threads(len(worker_cores))
+
+            if psutil is not None:
+                psutil.Process().cpu_affinity(worker_cores)
+        except IndexError as e:
+            raise ValueError(f"Cannot use CPU affinity for worker ID {worker_id} on CPU {self.loader_cores}") from e
+
+        self._old_worker_init_fn(worker_id)
+
+    @contextmanager
+    def enable_cpu_affinity(self, loader_cores: Optional[Union[List[List[int]], List[int]]] = None):
+        num_workers = getattr(self, 'num_workers', 0)
+        if not num_workers > 0:
+            raise ValueError(f"'enable_cpu_affinity' should be used with at least one worker (got {num_workers})")
+        if loader_cores and len(loader_cores) != num_workers:
+            raise ValueError(
+                f"The number of loader cores ({len(loader_cores)}) in 'enable_cpu_affinity' should match number of workers ({num_workers})"
+            )
+
+        from k3_node.data import HeteroData
+        if hasattr(self, 'data') and isinstance(self.data, HeteroData):
+            warnings.warn(
+                "Due to conflicting parallelization methods it is not advised to use affinitization with 'HeteroData' datasets.",
+                stacklevel=2,
+            )
+
+        self.loader_cores = loader_cores[:] if loader_cores else None
+        if self.loader_cores is None:
+            numa_info = get_numa_nodes_cores()
+            if numa_info and len(numa_info.get(0, [])) > num_workers:
+                node0_cores = [cpus[0] for core_id, cpus in numa_info[0]]
+                node0_cores.sort()
+            elif psutil is not None:
+                node0_cores = list(range(psutil.cpu_count(logical=False) or 1))
+            else:
+                node0_cores = list(range(os.cpu_count() or 1))
+
+            if len(node0_cores) < num_workers:
+                raise ValueError(f"More workers ({num_workers}) than available cores ({len(node0_cores)})")
+
+            if torch is not None and torch.multiprocessing.get_context()._name == 'spawn':
+                work_thread_pool = int(len(node0_cores) / num_workers)
+                self.loader_cores = [
+                    list(range(work_thread_pool * i, work_thread_pool * (i + 1)))
+                    for i in range(num_workers)
+                ]
+            else:
+                self.loader_cores = node0_cores[:num_workers]
+
+        self._old_worker_init_fn = WorkerInitWrapper(getattr(self, 'worker_init_fn', None))
+        try:
+            logging.debug(f"{num_workers} data loader workers assigned to CPUs {self.loader_cores}")
+            self.worker_init_fn = self._aff_init_fn
+            yield
+        finally:
+            self.worker_init_fn = self._old_worker_init_fn
+

k3_node/loader/neighbor_loader.py

@@ -0,0 +1,159 @@
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from k3_node.data import Data, HeteroData
+from k3_node.loader.node_loader import HeteroSamplerOutput, NodeLoader, NodeSamplerInput, SamplerOutput
+import numpy as np
+
+from k3_node.loader.sampler_utils import (FastGraph, sample_neighbors_disjoint, sample_neighbors_hetero,
+                                       sample_neighbors_homo)
+
+
+class InternalNeighborSampler:
+    r"""Pure Python/NumPy neighborhood sampling engine."""
+    def __init__(
+        self,
+        data: Union[Data, HeteroData],
+        num_neighbors: Union[List[int], Dict[Tuple[str, str, str], List[int]]],
+        replace: bool = False,
+        subgraph_type: str = 'directional',
+        disjoint: bool = False,
+    ):
+        self.data = data
+        self.num_neighbors = num_neighbors
+        self.replace = replace
+        self.subgraph_type = subgraph_type
+        self.disjoint = disjoint
+        self.edge_permutation = None
+        self._graph = None  # CSR index, built once and reused for every batch
+
+    def sample_from_nodes(self, input_data: NodeSamplerInput) -> Union[SamplerOutput, HeteroSamplerOutput]:
+        if isinstance(self.data, Data):
+            if self._graph is None:
+                self._graph = FastGraph(self.data.edge_index, num_nodes=self.data.num_nodes)
+            if self.disjoint:  # one separate subgraph per seed node
+                seeds = input_data.node
+                is_torch = hasattr(seeds, "detach")
+                out = sample_neighbors_disjoint(self._graph, seeds.detach().cpu().numpy() if is_torch else seeds,
+                                                self.num_neighbors, replace=self.replace)
+                node, row, col, edge, batch, n_counts, e_counts = out
+                if is_torch:
+                    import torch
+                    node, row, col, edge, batch = (torch.from_numpy(a) for a in (node, row, col, edge, batch))
+                return SamplerOutput(node=node, row=row, col=col, edge=edge, batch=batch,
+                                     num_sampled_nodes=n_counts, num_sampled_edges=e_counts,
+                                     metadata=(input_data.input_id, input_data.time))
+            node, row, col, edge, n_counts, e_counts = sample_neighbors_homo(
+                edge_index=self.data.edge_index,
+                seed_nodes=input_data.node,
+                num_neighbors=self.num_neighbors,
+                num_nodes=self.data.num_nodes,
+                replace=self.replace,
+                subgraph_type=self.subgraph_type,
+                disjoint=self.disjoint,
+                graph=self._graph,
+            )
+            return SamplerOutput(
+                node=node,
+                row=row,
+                col=col,
+                edge=edge,
+                num_sampled_nodes=n_counts,
+                num_sampled_edges=e_counts,
+                metadata=(input_data.input_id, input_data.time),
+            )
+        elif isinstance(self.data, HeteroData):
+            edge_index_dict = {}
+            for edge_type in self.data.edge_types:
+                canonical = self.data._to_canonical(*edge_type) if hasattr(self.data, '_to_canonical') else edge_type
+                edge_index_dict[canonical] = self.data[edge_type].edge_index
+
+            node_type = input_data.input_type or self.data.node_types[0]
+            seed_dict = {k: None for k in self.data.node_types}
+            seed_dict[node_type] = input_data.node
+
+            # Normalize num_neighbors for hetero
+            if isinstance(self.num_neighbors, dict):
+                norm_num_neighbors = {}
+                for k, v in self.num_neighbors.items():
+                    can = self.data._to_canonical(*k) if hasattr(self.data, '_to_canonical') else k
+                    norm_num_neighbors[can] = v
+            else:
+                norm_num_neighbors = self.num_neighbors
+
+            node_dict, row_dict, col_dict, edge_dict, n_counts, e_counts = sample_neighbors_hetero(
+                edge_index_dict=edge_index_dict,
+                seed_nodes_dict=seed_dict,
+                num_neighbors=norm_num_neighbors,
+                replace=self.replace,
+                subgraph_type=self.subgraph_type,
+            )
+            return HeteroSamplerOutput(
+                node=node_dict,
+                row=row_dict,
+                col=col_dict,
+                edge=edge_dict,
+                num_sampled_nodes=n_counts,
+                num_sampled_edges=e_counts,
+                metadata=(input_data.input_id, input_data.time),
+            )
+
+        raise TypeError(f"Invalid data type for sampling: {type(self.data)}")
+
+
+class NeighborLoader(NodeLoader):
+    r"""A data loader that performs neighbor sampling as introduced in
+    "Inductive Representation Learning on Large Graphs".
+
+    Args:
+        data (Data or HeteroData): The graph data object.
+        num_neighbors (List[int] or Dict[EdgeType, List[int]]): Number of neighbors to sample per iteration.
+        input_nodes (Tensor or str or Tuple[str, Tensor], optional): Seed nodes. (default: :obj:`None`)
+        replace (bool, optional): Sample with replacement. (default: :obj:`False`)
+        subgraph_type (str, optional): :obj:`"directional"`, :obj:`"bidirectional"`, or :obj:`"induced"`.
+            (default: :obj:`"directional"`)
+        disjoint (bool, optional): If :obj:`True`, creates disjoint subgraphs per seed node. (default: :obj:`False`)
+        **kwargs (optional): Additional arguments of :class:`torch.utils.data.DataLoader`.
+    """
+    def __init__(
+        self,
+        data: Union[Data, HeteroData],
+        num_neighbors: Union[List[int], Dict[Tuple[str, str, str], List[int]]],
+        input_nodes: Any = None,
+        input_time: Optional[Any] = None,
+        replace: bool = False,
+        subgraph_type: str = 'directional',
+        disjoint: bool = False,
+        temporal_strategy: str = 'uniform',
+        time_attr: Optional[str] = None,
+        weight_attr: Optional[str] = None,
+        transform: Optional[Callable] = None,
+        transform_sampler_output: Optional[Callable] = None,
+        is_sorted: bool = False,
+        filter_per_worker: Optional[bool] = None,
+        neighbor_sampler: Optional[Any] = None,
+        directed: bool = True,
+        **kwargs,
+    ):
+        if not directed:
+            subgraph_type = 'induced'
+
+        if neighbor_sampler is None:
+            neighbor_sampler = InternalNeighborSampler(
+                data,
+                num_neighbors=num_neighbors,
+                replace=replace,
+                subgraph_type=str(getattr(subgraph_type, 'value', subgraph_type)),
+                disjoint=disjoint,
+            )
+
+        super().__init__(
+            data=data,
+            node_sampler=neighbor_sampler,
+            input_nodes=input_nodes,
+            input_time=input_time,
+            transform=transform,
+            transform_sampler_output=transform_sampler_output,
+            filter_per_worker=filter_per_worker,
+            **kwargs,
+        )
+