pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__py3-none-any.whl

torch_geometric/datasets/web_qsp_dataset.py

@@ -1,241 +1,342 @@
 # Code adapted from the G-Retriever paper: https://arxiv.org/abs/2402.07630
-from typing import Any, Dict, List, Tuple, no_type_check
+import gc
+import os
+from itertools import chain
+from typing import Any, Dict, Iterator, List, Optional
 
-import numpy as np
 import torch
-from torch import Tensor
 from tqdm import tqdm
 
-from torch_geometric.data import Data, InMemoryDataset
-from torch_geometric.nn.nlp import SentenceTransformer
-
-
-@no_type_check
-def retrieval_via_pcst(
-    data: Data,
-    q_emb: Tensor,
-    textual_nodes: Any,
-    textual_edges: Any,
-    topk: int = 3,
-    topk_e: int = 3,
-    cost_e: float = 0.5,
-) -> Tuple[Data, str]:
-    c = 0.01
-    if len(textual_nodes) == 0 or len(textual_edges) == 0:
-        desc = textual_nodes.to_csv(index=False) + "\n" + textual_edges.to_csv(
-            index=False,
-            columns=["src", "edge_attr", "dst"],
-        )
-        return data, desc
-
-    from pcst_fast import pcst_fast
-
-    root = -1
-    num_clusters = 1
-    pruning = 'gw'
-    verbosity_level = 0
-    if topk > 0:
-        n_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.x)
-        topk = min(topk, data.num_nodes)
-        _, topk_n_indices = torch.topk(n_prizes, topk, largest=True)
-
-        n_prizes = torch.zeros_like(n_prizes)
-        n_prizes[topk_n_indices] = torch.arange(topk, 0, -1).float()
-    else:
-        n_prizes = torch.zeros(data.num_nodes)
-
-    if topk_e > 0:
-        e_prizes = torch.nn.CosineSimilarity(dim=-1)(q_emb, data.edge_attr)
-        topk_e = min(topk_e, e_prizes.unique().size(0))
-
-        topk_e_values, _ = torch.topk(e_prizes.unique(), topk_e, largest=True)
-        e_prizes[e_prizes < topk_e_values[-1]] = 0.0
-        last_topk_e_value = topk_e
-        for k in range(topk_e):
-            indices = e_prizes == topk_e_values[k]
-            value = min((topk_e - k) / sum(indices), last_topk_e_value - c)
-            e_prizes[indices] = value
-            last_topk_e_value = value * (1 - c)
-        # reduce the cost of the edges such that at least one edge is selected
-        cost_e = min(cost_e, e_prizes.max().item() * (1 - c / 2))
-    else:
-        e_prizes = torch.zeros(data.num_edges)
-
-    costs = []
-    edges = []
-    virtual_n_prizes = []
-    virtual_edges = []
-    virtual_costs = []
-    mapping_n = {}
-    mapping_e = {}
-    for i, (src, dst) in enumerate(data.edge_index.t().numpy()):
-        prize_e = e_prizes[i]
-        if prize_e <= cost_e:
-            mapping_e[len(edges)] = i
-            edges.append((src, dst))
-            costs.append(cost_e - prize_e)
-        else:
-            virtual_node_id = data.num_nodes + len(virtual_n_prizes)
-            mapping_n[virtual_node_id] = i
-            virtual_edges.append((src, virtual_node_id))
-            virtual_edges.append((virtual_node_id, dst))
-            virtual_costs.append(0)
-            virtual_costs.append(0)
-            virtual_n_prizes.append(prize_e - cost_e)
-
-    prizes = np.concatenate([n_prizes, np.array(virtual_n_prizes)])
-    num_edges = len(edges)
-    if len(virtual_costs) > 0:
-        costs = np.array(costs + virtual_costs)
-        edges = np.array(edges + virtual_edges)
-
-    vertices, edges = pcst_fast(edges, prizes, costs, root, num_clusters,
-                                pruning, verbosity_level)
-
-    selected_nodes = vertices[vertices < data.num_nodes]
-    selected_edges = [mapping_e[e] for e in edges if e < num_edges]
-    virtual_vertices = vertices[vertices >= data.num_nodes]
-    if len(virtual_vertices) > 0:
-        virtual_vertices = vertices[vertices >= data.num_nodes]
-        virtual_edges = [mapping_n[i] for i in virtual_vertices]
-        selected_edges = np.array(selected_edges + virtual_edges)
-
-    edge_index = data.edge_index[:, selected_edges]
-    selected_nodes = np.unique(
-        np.concatenate(
-            [selected_nodes, edge_index[0].numpy(), edge_index[1].numpy()]))
-
-    n = textual_nodes.iloc[selected_nodes]
-    e = textual_edges.iloc[selected_edges]
-    desc = n.to_csv(index=False) + '\n' + e.to_csv(
-        index=False, columns=['src', 'edge_attr', 'dst'])
-
-    mapping = {n: i for i, n in enumerate(selected_nodes.tolist())}
-    src = [mapping[i] for i in edge_index[0].tolist()]
-    dst = [mapping[i] for i in edge_index[1].tolist()]
-
-    data = Data(
-        x=data.x[selected_nodes],
-        edge_index=torch.tensor([src, dst]),
-        edge_attr=data.edge_attr[selected_edges],
-    )
-
-    return data, desc
-
-
-class WebQSPDataset(InMemoryDataset):
-    r"""The WebQuestionsSP dataset of the `"The Value of Semantic Parse
-    Labeling for Knowledge Base Question Answering"
-    <https://aclanthology.org/P16-2033/>`_ paper.
+from torch_geometric.data import InMemoryDataset
+from torch_geometric.llm.large_graph_indexer import (
+    EDGE_RELATION,
+    LargeGraphIndexer,
+    TripletLike,
+    get_features_for_triplets_groups,
+)
+from torch_geometric.llm.models import SentenceTransformer
+from torch_geometric.llm.utils.backend_utils import (
+    preprocess_triplet,
+    retrieval_via_pcst,
+)
+
+
+class KGQABaseDataset(InMemoryDataset):
+    r"""Base class for the 2 KGQA datasets used in `"Reasoning on Graphs:
+    Faithful and Interpretable Large Language Model Reasoning"
+    <https://arxiv.org/pdf/2310.01061>`_ paper.
 
     Args:
+        dataset_name (str): HuggingFace `dataset` name.
         root (str): Root directory where the dataset should be saved.
         split (str, optional): If :obj:`"train"`, loads the training dataset.
             If :obj:`"val"`, loads the validation dataset.
             If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
         force_reload (bool, optional): Whether to re-process the dataset.
             (default: :obj:`False`)
+        verbose (bool, optional): Whether to print output. Defaults to False.
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+        load_dataset_kwargs (dict, optional):
+            Keyword arguments for the `datasets.load_dataset` function.
+            (default: :obj:`{}`)
+        retrieval_kwargs (dict, optional):
+            Keyword arguments for the
+            `get_features_for_triplets_groups` function.
+            (default: :obj:`{}`)
     """
     def __init__(
         self,
+        dataset_name: str,
         root: str,
         split: str = "train",
         force_reload: bool = False,
+        verbose: bool = False,
+        use_pcst: bool = True,
+        load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        retrieval_kwargs: Optional[Dict[str, Any]] = None,
     ) -> None:
+        self.split = split
+        self.dataset_name = dataset_name
+        self.use_pcst = use_pcst
+        self.load_dataset_kwargs = load_dataset_kwargs or {}
+        """
+        NOTE: If running into memory issues,
+        try reducing this batch size for the LargeGraphIndexer
+        used to build our KG.
+        Example: self.retrieval_kwargs = {"batch_size": 64}
+        """
+        self.retrieval_kwargs = retrieval_kwargs or {}
+
+        # Caching custom subsets of the dataset results in unsupported behavior
+        if 'split' in self.load_dataset_kwargs:
+            print("WARNING: Caching custom subsets of the dataset \
+                results in unsupported behavior.\
+                Please specify a separate root directory for each split,\
+                or set force_reload=True on subsequent instantiations\
+                of the dataset.")
+
+        self.required_splits = ['train', 'validation', 'test']
+
+        self.verbose = verbose
+        self.force_reload = force_reload
         super().__init__(root, force_reload=force_reload)
-
-        if split not in {'train', 'val', 'test'}:
+        """
+        NOTE: Current behavior is to process the entire dataset,
+        and only return the split specified by the user.
+        """
+        if f'{split}_data.pt' not in set(self.processed_file_names):
             raise ValueError(f"Invalid 'split' argument (got {split})")
+        if split == 'val':
+            split = 'validation'
 
-        path = self.processed_paths[['train', 'val', 'test'].index(split)]
-        self.load(path)
+        self.load(self.processed_paths[self.required_splits.index(split)])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ["raw.pt"]
 
     @property
     def processed_file_names(self) -> List[str]:
-        return ['train_data.pt', 'val_data.pt', 'test_data.pt']
+        return ["train_data.pt", "val_data.pt", "test_data.pt"]
 
-    def process(self) -> None:
+    def download(self) -> None:
         import datasets
-        import pandas as pd
 
-        datasets = datasets.load_dataset('rmanluo/RoG-webqsp')
+        # HF Load Dataset by dataset name if no path is specified
+        self.load_dataset_kwargs['path'] = self.load_dataset_kwargs.get(
+            'path', self.dataset_name)
+        raw_dataset = datasets.load_dataset(**self.load_dataset_kwargs)
+
+        # Assert that the dataset contains the required splits
+        assert all(split in raw_dataset for split in self.required_splits), \
+            f"Dataset '{self.dataset_name}' is missing required splits: \
+            {self.required_splits}"
+
+        raw_dataset.save_to_disk(self.raw_paths[0])
+
+    def _get_trips(self) -> Iterator[TripletLike]:
+        # Iterate over each element's graph in each split of the dataset
+        # Using chain to lazily iterate without storing all trips in memory
+        split_iterators = []
+
+        for split in self.required_splits:
+            # Create an iterator for each element's graph in the current split
+            split_graphs = (element['graph']
+                            for element in self.raw_dataset[split])
+            split_iterators.append(chain.from_iterable(split_graphs))
+
+        # Chain all split iterators together
+        return chain.from_iterable(split_iterators)
+
+    def _build_graph(self) -> None:
+        print("Encoding graph...")
+        trips = self._get_trips()
+        self.indexer: LargeGraphIndexer = LargeGraphIndexer.from_triplets(
+            trips, pre_transform=preprocess_triplet)
+
+        # Nodes:
+        print("\tEncoding nodes...")
+        nodes = self.indexer.get_unique_node_features()
+        x = self.model.encode(nodes, batch_size=256, output_device='cpu')
+        self.indexer.add_node_feature(new_feature_name="x", new_feature_vals=x)
+
+        # Edges:
+        print("\tEncoding edges...")
+        edges = self.indexer.get_unique_edge_features(
+            feature_name=EDGE_RELATION)
+        edge_attr = self.model.encode(edges, batch_size=256,
+                                      output_device='cpu')
+        self.indexer.add_edge_feature(
+            new_feature_name="edge_attr",
+            new_feature_vals=edge_attr,
+            map_from_feature=EDGE_RELATION,
+        )
+
+        print("\tSaving graph...")
+        self.indexer.save(self.indexer_path)
+
+    def _retrieve_subgraphs(self) -> None:
+        raw_splits = [
+            self.raw_dataset[split] for split in self.required_splits
+        ]
+        zipped = zip(
+            self.required_splits,
+            raw_splits,  # noqa
+            self.processed_paths,
+        )
+        for split_name, dataset, path in zipped:
+            print(f"Processing {split_name} split...")
+
+            print("\tEncoding questions...")
+            split_questions = [str(element['question']) for element in dataset]
+            split_q_embs = self.model.encode(split_questions, batch_size=256,
+                                             output_device='cpu')
+
+            print("\tRetrieving subgraphs...")
+            results_graphs = []
+            retrieval_kwargs = {
+                **self.retrieval_kwargs,
+                **{
+                    'pre_transform': preprocess_triplet,
+                    'verbose': self.verbose,
+                }
+            }
+            graph_gen = get_features_for_triplets_groups(
+                self.indexer, (element['graph'] for element in dataset),
+                **retrieval_kwargs)
+
+            for index in tqdm(range(len(dataset)), disable=not self.verbose):
+                data_i = dataset[index]
+                graph = next(graph_gen)
+                textual_nodes = self.textual_nodes.iloc[
+                    graph["node_idx"]].reset_index()
+                textual_edges = self.textual_edges.iloc[
+                    graph["edge_idx"]].reset_index()
+                if self.use_pcst and len(textual_nodes) > 0 and len(
+                        textual_edges) > 0:
+                    subgraph, desc = retrieval_via_pcst(
+                        graph,
+                        split_q_embs[index],
+                        textual_nodes,
+                        textual_edges,
+                    )
+                else:
+                    desc = textual_nodes.to_csv(
+                        index=False) + "\n" + textual_edges.to_csv(
+                            index=False,
+                            columns=["src", "edge_attr", "dst"],
+                        )
+                    subgraph = graph
+                question = f"Question: {data_i['question']}\nAnswer: "
+                label = ("|").join(data_i["answer"]).lower()
+
+                subgraph["question"] = question
+                subgraph["label"] = label
+                subgraph["desc"] = desc
+                results_graphs.append(subgraph.to("cpu"))
+            print("\tSaving subgraphs...")
+            self.save(results_graphs, path)
+
+    def process(self) -> None:
+        import datasets
+        from pandas import DataFrame
+        self.raw_dataset = datasets.load_from_disk(self.raw_paths[0])
 
         device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         model_name = 'sentence-transformers/all-roberta-large-v1'
-        model = SentenceTransformer(model_name).to(device)
-        model.eval()
-
-        for dataset, path in zip(
-            [datasets['train'], datasets['validation'], datasets['test']],
-                self.processed_paths,
-        ):
-            questions = [example["question"] for example in dataset]
-            question_embs = model.encode(
-                questions,
-                batch_size=256,
-                output_device='cpu',
-            )
-
-            data_list = []
-            for i, example in enumerate(tqdm(dataset)):
-                raw_nodes: Dict[str, int] = {}
-                raw_edges = []
-                for tri in example["graph"]:
-                    h, r, t = tri
-                    h = h.lower()
-                    t = t.lower()
-                    if h not in raw_nodes:
-                        raw_nodes[h] = len(raw_nodes)
-                    if t not in raw_nodes:
-                        raw_nodes[t] = len(raw_nodes)
-                    raw_edges.append({
-                        "src": raw_nodes[h],
-                        "edge_attr": r,
-                        "dst": raw_nodes[t]
-                    })
-                nodes = pd.DataFrame([{
-                    "node_id": v,
-                    "node_attr": k,
-                } for k, v in raw_nodes.items()],
-                                     columns=["node_id", "node_attr"])
-                edges = pd.DataFrame(raw_edges,
-                                     columns=["src", "edge_attr", "dst"])
-
-                nodes.node_attr = nodes.node_attr.fillna("")
-                x = model.encode(
-                    nodes.node_attr.tolist(),
-                    batch_size=256,
-                    output_device='cpu',
-                )
-                edge_attr = model.encode(
-                    edges.edge_attr.tolist(),
-                    batch_size=256,
-                    output_device='cpu',
-                )
-                edge_index = torch.tensor([
-                    edges.src.tolist(),
-                    edges.dst.tolist(),
-                ], dtype=torch.long)
-
-                question = f"Question: {example['question']}\nAnswer: "
-                label = ('|').join(example['answer']).lower()
-                data = Data(
-                    x=x,
-                    edge_index=edge_index,
-                    edge_attr=edge_attr,
-                )
-                data, desc = retrieval_via_pcst(
-                    data,
-                    question_embs[i],
-                    nodes,
-                    edges,
-                    topk=3,
-                    topk_e=5,
-                    cost_e=0.5,
-                )
-                data.question = question
-                data.label = label
-                data.desc = desc
-                data_list.append(data)
-
-            self.save(data_list, path)
+        self.model: SentenceTransformer = SentenceTransformer(model_name).to(
+            device)
+        self.model.eval()
+        self.indexer_path = os.path.join(self.processed_dir,
+                                         "large_graph_indexer")
+        if self.force_reload or not os.path.exists(self.indexer_path):
+            self._build_graph()
+        else:
+            print("Loading graph...")
+            self.indexer = LargeGraphIndexer.from_disk(self.indexer_path)
+        self.textual_nodes = DataFrame.from_dict(
+            {"node_attr": self.indexer.get_node_features()})
+        self.textual_nodes["node_id"] = self.textual_nodes.index
+        self.textual_nodes = self.textual_nodes[["node_id", "node_attr"]]
+        self.textual_edges = DataFrame(self.indexer.get_edge_features(),
+                                       columns=["src", "edge_attr", "dst"])
+        self.textual_edges["src"] = [
+            self.indexer._nodes[h] for h in self.textual_edges["src"]
+        ]
+        self.textual_edges["dst"] = [
+            self.indexer._nodes[h] for h in self.textual_edges["dst"]
+        ]
+        self._retrieve_subgraphs()
+
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+class WebQSPDataset(KGQABaseDataset):
+    r"""The WebQuestionsSP dataset of the `"The Value of Semantic Parse
+    Labeling for Knowledge Base Question Answering"
+    <https://aclanthology.org/P16-2033/>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+        verbose (bool, optional): Whether to print output. Defaults to False.
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+        load_dataset_kwargs (dict, optional):
+            Keyword arguments for the `datasets.load_dataset` function.
+            (default: :obj:`{}`)
+        retrieval_kwargs (dict, optional):
+            Keyword arguments for the
+            `get_features_for_triplets_groups` function.
+            (default: :obj:`{}`)
+    """
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+        verbose: bool = False,
+        use_pcst: bool = True,
+        load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        retrieval_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        load_dataset_kwargs = load_dataset_kwargs or {}
+        retrieval_kwargs = retrieval_kwargs or {}
+        # Modify these paramters if running into memory/compute issues
+        default_retrieval_kwargs = {
+            'max_batch_size': 250,  # Lower batch size to reduce memory usage
+            'num_workers':
+            None,  # Use all available workers, or set to number of threads
+        }
+        retrieval_kwargs = {**default_retrieval_kwargs, **retrieval_kwargs}
+        dataset_name = 'rmanluo/RoG-webqsp'
+        super().__init__(dataset_name, root, split, force_reload, verbose,
+                         use_pcst, load_dataset_kwargs=load_dataset_kwargs,
+                         retrieval_kwargs=retrieval_kwargs)
+
+
+class CWQDataset(KGQABaseDataset):
+    r"""The ComplexWebQuestions (CWQ) dataset of the `"The Web as a
+    Knowledge-base forAnswering Complex Questions"
+    <https://arxiv.org/pdf/1803.06643>`_ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        split (str, optional): If :obj:`"train"`, loads the training dataset.
+            If :obj:`"val"`, loads the validation dataset.
+            If :obj:`"test"`, loads the test dataset. (default: :obj:`"train"`)
+        force_reload (bool, optional): Whether to re-process the dataset.
+            (default: :obj:`False`)
+        verbose (bool, optional): Whether to print output. Defaults to False.
+        use_pcst (bool, optional): Whether to preprocess the dataset's graph
+            with PCST or return the full graphs. (default: :obj:`True`)
+        load_dataset_kwargs (dict, optional):
+            Keyword arguments for the `datasets.load_dataset` function.
+            (default: :obj:`{}`)
+        retrieval_kwargs (dict, optional):
+            Keyword arguments for the
+            `get_features_for_triplets_groups` function.
+            (default: :obj:`{}`)
+    """
+    def __init__(
+        self,
+        root: str,
+        split: str = "train",
+        force_reload: bool = False,
+        verbose: bool = False,
+        use_pcst: bool = True,
+        load_dataset_kwargs: Optional[Dict[str, Any]] = None,
+        retrieval_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        load_dataset_kwargs = load_dataset_kwargs or {}
+        retrieval_kwargs = retrieval_kwargs or {}
+        dataset_name = 'rmanluo/RoG-cwq'
+        super().__init__(dataset_name, root, split, force_reload, verbose,
+                         use_pcst, load_dataset_kwargs=load_dataset_kwargs,
+                         retrieval_kwargs=retrieval_kwargs)

torch_geometric/datasets/wikics.py

@@ -45,7 +45,8 @@ class WikiCS(InMemoryDataset):
             warnings.warn(
                 f"The {self.__class__.__name__} dataset now returns an "
                 f"undirected graph by default. Please explicitly specify "
-                f"'is_undirected=False' to restore the old behavior.")
+                f"'is_undirected=False' to restore the old behavior.",
+                stacklevel=2)
             is_undirected = True
         self.is_undirected = is_undirected
         super().__init__(root, transform, pre_transform,

torch_geometric/deprecation.py

@@ -23,7 +23,7 @@ def deprecated(
             out = f"'{name}' is deprecated"
             if details is not None:
                 out += f", {details}"
-            warnings.warn(out)
+            warnings.warn(out, stacklevel=2)
             return func(*args, **kwargs)
 
         return wrapper

torch_geometric/distributed/__init__.py

@@ -1,3 +1,5 @@
+from warnings import warn
+
 from .dist_context import DistContext
 from .local_feature_store import LocalFeatureStore
 from .local_graph_store import LocalGraphStore
@@ -7,6 +9,17 @@ from .dist_loader import DistLoader
 from .dist_neighbor_loader import DistNeighborLoader
 from .dist_link_neighbor_loader import DistLinkNeighborLoader
 
+warn(
+    "`torch_geometric.distributed` has been deprecated since 2.7.0 and will "
+    "no longer be maintained. For distributed training, refer to our "
+    "tutorials on distributed training at "
+    "https://pytorch-geometric.readthedocs.io/en/latest/tutorial/distributed.html "  # noqa: E501
+    "or cuGraph examples at "
+    "https://github.com/rapidsai/cugraph-gnn/tree/main/python/cugraph-pyg/cugraph_pyg/examples",  # noqa: E501
+    stacklevel=2,
+    category=DeprecationWarning,
+)
+
 __all__ = classes = [
     'DistContext',
     'LocalFeatureStore',

torch_geometric/distributed/dist_loader.py

@@ -138,9 +138,9 @@ class DistLoader:
             # close RPC & worker group at exit:
             atexit.register(shutdown_rpc, self.current_ctx_worker.worker_name)
 
-        except RuntimeError:
+        except RuntimeError as e:
             raise RuntimeError(f"`{self}.init_fn()` could not initialize the "
-                               f"worker loop of the neighbor sampler")
+                               f"worker loop of the neighbor sampler") from e
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}(pid={self.pid})'

torch_geometric/distributed/partition.py

@@ -304,7 +304,7 @@ class Partitioner:
                 elif self.is_node_level_time:
                     node_time = data.time
 
-                # Sort by column to avoid keeping track of permuations in
+                # Sort by column to avoid keeping track of permutations in
                 # `NeighborSampler` when converting to CSC format:
                 global_row, global_col, perm = sort_csc(
                     global_row, global_col, node_time, edge_time)
@@ -361,7 +361,7 @@ class Partitioner:
             'edge_types': self.edge_types,
             'node_offset': list(node_offset.values()) if node_offset else None,
             'is_hetero': self.is_hetero,
-            'is_sorted': True,  # Based on colum/destination.
+            'is_sorted': True,  # Based on column/destination.
         }
         with open(osp.join(self.root, 'META.json'), 'w') as f:
             json.dump(meta, f)

torch_geometric/distributed/rpc.py

@@ -92,7 +92,7 @@ def shutdown_rpc(id: str = None, graceful: bool = True,
 class RPCRouter:
     r"""A router to get the worker based on the partition ID."""
     def __init__(self, partition_to_workers: List[List[str]]):
-        for pid, rpc_worker_list in enumerate(partition_to_workers):
+        for rpc_worker_list in partition_to_workers:
             if len(rpc_worker_list) == 0:
                 raise ValueError('No RPC worker is in worker list')
         self.partition_to_workers = partition_to_workers
@@ -120,7 +120,7 @@ def rpc_partition_to_workers(
     partition_to_workers = [[] for _ in range(num_partitions)]
     gathered_results = global_all_gather(
         (ctx.role, num_partitions, current_partition_idx))
-    for worker_name, (role, nparts, idx) in gathered_results.items():
+    for worker_name, (_, _, idx) in gathered_results.items():
         partition_to_workers[idx].append(worker_name)
     return partition_to_workers
 
@@ -144,7 +144,7 @@ _rpc_call_pool: Dict[int, RPCCallBase] = {}
 @rpc_require_initialized
 def rpc_register(call: RPCCallBase) -> int:
     r"""Registers a call for RPC requests."""
-    global _rpc_call_id, _rpc_call_pool
+    global _rpc_call_id
 
     with _rpc_call_lock:
         call_id = _rpc_call_id