Repository: dmlc/dgl
Pull request: [Graphbolt] Add pickle serialization support for Graphbolt SampledSubgraph
## Description
Implemented `__setstate__` and `__getstate__` function for `SampledSubgraph` which enables pickle to dump it.
Resolve #5953.
[Updated] Added unit test.

## Checklist
Please feel free to remove inapplicable items for your PR.
- [x] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [x] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code.
- [ ] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
- [x] All changes have test coverage
- [ ] Code is well-documented
- [ ] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
- [x] Related issue is referred in this PR
- [ ] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md).

## Changes
- Modified `graphbolt/include/graphbolt/sampled_subgraph.h` and `graphbolt/src/sampled_subgraph.cc` to add `SetState()` and `GetState()` functions. It serializes the graph into a vector of tensors.
- Modified `graphbolt/src/python_bindings.cc`. Use `m.class_<>.def_pickle()` to bind them onto python `__setstate__` and `__getstate__` functions. (Ref: [pytorch docs](https://pytorch.org/tutorials/advanced/torch_script_custom_classes.html#defining-serialization-deserialization-methods-for-custom-c-classes))
- If in the future there is `Dict` type attributes, `SetState()` and `GetState()` could be easily modified to support it.
- Added unit test in `tests/python/pytorch/graphbolt/impl/test_c_sampled_subgraph.py`


Problem statement:
[GraphBolt] GraphBolt objects cannot be serialized during multiprocessing
## 🐛 Bug

It seems that GraphBolt objects like `SampledSubgraph` cannot be passed around in multiprocessing due to lack of serialization method.

## To Reproduce

```python
import dgl.graphbolt
import torch
import multiprocessing as mp
import scipy.sparse as sp

def rand_csc_graph(N, density):
    adj = sp.random(N, N, density)
    adj = adj + adj.T
    adj = adj.tocsc()

    indptr = torch.LongTensor(adj.indptr)
    indices = torch.LongTensor(adj.indices)

    graph = dgl.graphbolt.from_csc(indptr, indices)

    return graph

def _entry(graph, q):
    adjs = []
    seeds = torch.arange(5)

    for hop in range(2):
        sg = graph.sample_neighbors(seeds, torch.LongTensor([2]))
        seeds = sg.indices
        adjs.insert(0, sg)

    print('Putting')
    q.put(adjs)
    print('Done put')

if __name__ == '__main__':
    graph = rand_csc_graph(200, 0.15)
    q = mp.Queue()
    proc = mp.Process(target=_entry, args=(graph, q))
    proc.start()
    print('Getting')
    item = q.get()
    print('Done get')
    proc.join()
    print(item)
```

This code snippet hangs indefinitely, either silently or throws the following error:
```
Traceback (most recent call last):
  File "/home/ubuntu/miniconda3/lib/python3.9/multiprocessing/queues.py", line 245, in _feed
    obj = _ForkingPickler.dumps(obj)
  File "/home/ubuntu/miniconda3/lib/python3.9/multiprocessing/reduction.py", line 51, in dumps
    cls(buf, protocol).dump(obj)
RuntimeError: Tried to serialize object __torch__.torch.classes.graphbolt.SampledSubgraph which does not have a __getstate__ method defined!
```

## Expected behavior

The data should be successfully printed.  Replacing `q.put(adjs)` to something else like `q.put(42)` works.  So it should be something explicit to GraphBolt.

## Environment

 - DGL Version (e.g., 1.0): master
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:

Hints from issue discussion:

Commit message: Fix lint issue

Patch to review:
diff --git a/graphbolt/include/graphbolt/sampled_subgraph.h b/graphbolt/include/graphbolt/sampled_subgraph.h
index 49ebb4651f88..ab958545e9b0 100644
--- a/graphbolt/include/graphbolt/sampled_subgraph.h
+++ b/graphbolt/include/graphbolt/sampled_subgraph.h
@@ -113,6 +113,24 @@ struct SampledSubgraph : torch::CustomClassHolder {
    * subgraph.
    */
   torch::optional<torch::Tensor> type_per_edge;
+
+  /**
+   * @brief Version number to indicate graph version in serialization and
+   * deserialization.
+   */
+  static constexpr int64_t kSampledSubgraphSerializeVersionNumber = 1;
+
+  /**
+   * @brief Get graph state (for pickle serialization).
+   * @return A vector of Tensors.
+   */
+  std::vector<torch::Tensor> GetState();
+
+  /**
+   * @brief Set graph state (for pickle deserialization).
+   * @param state A vector of Tensors.
+   */
+  void SetState(std::vector<torch::Tensor>& state);
 };
 
 }  // namespace sampling
diff --git a/graphbolt/src/python_binding.cc b/graphbolt/src/python_binding.cc
index a307fd63e41a..231d5f50eb14 100644
--- a/graphbolt/src/python_binding.cc
+++ b/graphbolt/src/python_binding.cc
@@ -20,7 +20,18 @@ TORCH_LIBRARY(graphbolt, m) {
       .def_readwrite(
           "reverse_column_node_ids", &SampledSubgraph::reverse_column_node_ids)
       .def_readwrite("reverse_edge_ids", &SampledSubgraph::reverse_edge_ids)
-      .def_readwrite("type_per_edge", &SampledSubgraph::type_per_edge);
+      .def_readwrite("type_per_edge", &SampledSubgraph::type_per_edge)
+      .def_pickle(
+          // __getstate__
+          [](const c10::intrusive_ptr<SampledSubgraph>& self)
+              -> std::vector<torch::Tensor> { return self->GetState(); },
+          // __setstate__
+          [](std::vector<torch::Tensor> state)
+              -> c10::intrusive_ptr<SampledSubgraph> {
+            auto g = c10::make_intrusive<SampledSubgraph>();
+            g->SetState(state);
+            return g;
+          });
   m.class_<CSCSamplingGraph>("CSCSamplingGraph")
       .def("num_nodes", &CSCSamplingGraph::NumNodes)
       .def("num_edges", &CSCSamplingGraph::NumEdges)
diff --git a/graphbolt/src/sampled_subgraph.cc b/graphbolt/src/sampled_subgraph.cc
new file mode 100644
index 000000000000..8a6422d467fa
--- /dev/null
+++ b/graphbolt/src/sampled_subgraph.cc
@@ -0,0 +1,92 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file sampled_subgraph.cc
+ * @brief Source file of sampled subgraph.
+ */
+
+#include <graphbolt/sampled_subgraph.h>
+#include <graphbolt/serialize.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+namespace graphbolt {
+namespace sampling {
+
+std::vector<torch::Tensor> SampledSubgraph::GetState() {
+  std::vector<torch::Tensor> state;
+
+  // Version number.
+  torch::Tensor version_num_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt64)) *
+      kSampledSubgraphSerializeVersionNumber;
+  state.push_back(version_num_tensor);
+
+  // Tensors.
+  state.push_back(indptr);
+  state.push_back(indices);
+  state.push_back(reverse_column_node_ids);
+
+  // Optional tensors.
+  static torch::Tensor true_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt32));
+  static torch::Tensor false_tensor =
+      torch::zeros(1, torch::TensorOptions().dtype(torch::kInt32));
+  if (reverse_row_node_ids.has_value()) {
+    state.push_back(true_tensor);
+    state.push_back(reverse_row_node_ids.value());
+  } else {
+    state.push_back(false_tensor);
+  }
+  if (reverse_edge_ids.has_value()) {
+    state.push_back(true_tensor);
+    state.push_back(reverse_edge_ids.value());
+  } else {
+    state.push_back(false_tensor);
+  }
+  if (type_per_edge.has_value()) {
+    state.push_back(true_tensor);
+    state.push_back(type_per_edge.value());
+  } else {
+    state.push_back(false_tensor);
+  }
+
+  return state;
+}
+
+void SampledSubgraph::SetState(std::vector<torch::Tensor>& state) {
+  // Iterator.
+  uint32_t i = 0;
+
+  // Version number.
+  torch::Tensor version_num_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt64)) *
+      SampledSubgraph::kSampledSubgraphSerializeVersionNumber;
+  TORCH_CHECK(
+      version_num_tensor.equal(state[i++]),
+      "Version number mismatch when deserializing SampledSubgraph.");
+
+  // Tensors.
+  indptr = state[i++];
+  indices = state[i++];
+  reverse_column_node_ids = state[i++];
+
+  // Optional tensors.
+  static torch::Tensor true_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt32));
+  reverse_row_node_ids = torch::nullopt;
+  reverse_edge_ids = torch::nullopt;
+  type_per_edge = torch::nullopt;
+  if (state[i++].equal(true_tensor)) {
+    reverse_row_node_ids = state[i++];
+  }
+  if (state[i++].equal(true_tensor)) {
+    reverse_edge_ids = state[i++];
+  }
+  if (state[i++].equal(true_tensor)) {
+    type_per_edge = state[i++];
+  }
+}
+
+}  // namespace sampling
+}  // namespace graphbolt
diff --git a/tests/python/pytorch/graphbolt/impl/test_c_sampled_subgraph.py b/tests/python/pytorch/graphbolt/impl/test_c_sampled_subgraph.py
new file mode 100644
index 000000000000..2d77d831743d
--- /dev/null
+++ b/tests/python/pytorch/graphbolt/impl/test_c_sampled_subgraph.py
@@ -0,0 +1,93 @@
+import multiprocessing as mp
+import os
+import tempfile
+import time
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.graphbolt as gb
+import gb_test_utils as gbt
+
+import pytest
+import torch
+
+
+def subprocess_entry(q):
+    num_nodes = 5
+    num_edges = 12
+    indptr = torch.LongTensor([0, 3, 5, 7, 9, 12])
+    indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 1, 1, 2, 0, 3, 4])
+    type_per_edge = torch.LongTensor([0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1])
+    assert indptr[-1] == num_edges
+    assert indptr[-1] == len(indices)
+    ntypes = {"n1": 0, "n2": 1, "n3": 2}
+    etypes = {("n1", "e1", "n2"): 0, ("n1", "e2", "n3"): 1}
+    metadata = gb.GraphMetadata(ntypes, etypes)
+
+    # Construct CSCSamplingGraph.
+    graph = gb.from_csc(
+        indptr, indices, type_per_edge=type_per_edge, metadata=metadata
+    )
+
+    adjs = []
+    seeds = torch.arange(5)
+
+    # Sampling.
+    for hop in range(2):
+        sg = graph.sample_neighbors(seeds, torch.LongTensor([2]))
+        seeds = sg.indices
+        adjs.append(sg)
+
+    # 1. Put.
+    q.put(adjs)
+    time.sleep(2)
+    # 4. Get.
+    result = q.get()
+
+    # Verification.
+    for hop in range(2):
+        # Tensors.
+        assert torch.equal(adjs[hop].indptr, result[hop].indptr)
+        assert torch.equal(adjs[hop].indices, result[hop].indices)
+        assert torch.equal(
+            adjs[hop].reverse_column_node_ids,
+            result[hop].reverse_column_node_ids,
+        )
+        # Optional tensors.
+        assert (
+            adjs[hop].reverse_row_node_ids is None
+            and adjs[hop].reverse_row_node_ids is None
+        ) or torch.equal(
+            adjs[hop].reverse_row_node_ids, result[hop].reverse_row_node_ids
+        )
+        assert (
+            adjs[hop].reverse_edge_ids is None
+            and result[hop].reverse_edge_ids is None
+        ) or torch.equal(
+            adjs[hop].reverse_edge_ids, result[hop].reverse_edge_ids
+        )
+        assert (
+            adjs[hop].type_per_edge is None
+            and result[hop].type_per_edge is None
+        ) or torch.equal(adjs[hop].type_per_edge, result[hop].type_per_edge)
+
+
+@unittest.skipIf(
+    F._default_context_str == "gpu",
+    reason="Graph is CPU only at present.",
+)
+def test_subgraph_serialization():
+    # Create a sub-process.
+    q = mp.Queue()
+    proc = mp.Process(target=subprocess_entry, args=(q,))
+    proc.start()
+
+    # 2. Get.
+    items = q.get()
+    # 3. Put again.
+    q.put(items)
+
+    time.sleep(1)
+    proc.join()


Code review comments:
diff_hunk: @@ -113,6 +113,24 @@ struct SampledSubgraph : torch::CustomClassHolder {
    * subgraph.
    */
   torch::optional<torch::Tensor> type_per_edge;
+
+  /**
+   * @brief Version number to indicate graph version in serialization and
+   * deserialization.
+   */
+  static constexpr int64_t kSampledSubgraphSerializeVersionNumber = 1;
path: graphbolt/include/graphbolt/sampled_subgraph.h
text: @user1:
Does external class need to access this field?

@author:
No I think. Should it be set to private and non-static?

@user1:
move it to sampled_subgraph.cc within the namespace should be fine.

diff_hunk: @@ -0,0 +1,93 @@
+import multiprocessing as mp
+import os
+import tempfile
+import time
+import unittest
+
+import backend as F
+
+import dgl
+import dgl.graphbolt as gb
+import gb_test_utils as gbt
+
+import pytest
+import torch
+
+
+def subprocess_entry(q):
+    num_nodes = 5
+    num_edges = 12
+    indptr = torch.LongTensor([0, 3, 5, 7, 9, 12])
+    indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 1, 1, 2, 0, 3, 4])
+    type_per_edge = torch.LongTensor([0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1])
+    assert indptr[-1] == num_edges
+    assert indptr[-1] == len(indices)
+    ntypes = {"n1": 0, "n2": 1, "n3": 2}
+    etypes = {("n1", "e1", "n2"): 0, ("n1", "e2", "n3"): 1}
+    metadata = gb.GraphMetadata(ntypes, etypes)
+
+    # Construct CSCSamplingGraph.
+    graph = gb.from_csc(
+        indptr, indices, type_per_edge=type_per_edge, metadata=metadata
+    )
+
+    adjs = []
+    seeds = torch.arange(5)
+
+    # Sampling.
+    for hop in range(2):
+        sg = graph.sample_neighbors(seeds, torch.LongTensor([2]))
+        seeds = sg.indices
+        adjs.append(sg)
+
+    # 1. Put.
path: tests/python/pytorch/graphbolt/impl/test_c_sampled_subgraph.py
text: @user1:
The 1, 2, 3, 4 which representing the order of the operators is not very clear, and also how do you guarantee the #2 happens after #1?

diff_hunk: @@ -0,0 +1,92 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file sampled_subgraph.cc
+ * @brief Source file of sampled subgraph.
+ */
+
+#include <graphbolt/sampled_subgraph.h>
+#include <graphbolt/serialize.h>
+#include <torch/torch.h>
+
+#include <vector>
+
+namespace graphbolt {
+namespace sampling {
+
+std::vector<torch::Tensor> SampledSubgraph::GetState() {
+  std::vector<torch::Tensor> state;
+
+  // Version number.
+  torch::Tensor version_num_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt64)) *
+      kSampledSubgraphSerializeVersionNumber;
+  state.push_back(version_num_tensor);
+
+  // Tensors.
+  state.push_back(indptr);
+  state.push_back(indices);
+  state.push_back(reverse_column_node_ids);
+
+  // Optional tensors.
+  static torch::Tensor true_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt32));
+  static torch::Tensor false_tensor =
+      torch::zeros(1, torch::TensorOptions().dtype(torch::kInt32));
+  if (reverse_row_node_ids.has_value()) {
+    state.push_back(true_tensor);
+    state.push_back(reverse_row_node_ids.value());
+  } else {
+    state.push_back(false_tensor);
+  }
+  if (reverse_edge_ids.has_value()) {
+    state.push_back(true_tensor);
+    state.push_back(reverse_edge_ids.value());
+  } else {
+    state.push_back(false_tensor);
+  }
+  if (type_per_edge.has_value()) {
+    state.push_back(true_tensor);
+    state.push_back(type_per_edge.value());
+  } else {
+    state.push_back(false_tensor);
+  }
+
+  return state;
+}
+
+void SampledSubgraph::SetState(std::vector<torch::Tensor>& state) {
+  // Iterator.
+  uint32_t i = 0;
+
+  // Version number.
+  torch::Tensor version_num_tensor =
+      torch::ones(1, torch::TensorOptions().dtype(torch::kInt64)) *
+      SampledSubgraph::kSampledSubgraphSerializeVersionNumber;
+  TORCH_CHECK(
+      version_num_tensor.equal(state[i++]),
path: graphbolt/src/sampled_subgraph.cc
text: @user1:
Avoid change value in a MARCO statement, which is usually error-prone.

One option:
```
  TORCH_CHECK(
      version_num_tensor.equal(state[0),
...

  // Iterator.
  uint32_t i = 1;
```

Or:

```
  torch::Tensor actual_version = state[i++];
  torch::Tensor expected_version = torch::ones(...)
  TORCH_CHECK(
      actual_version.equal(expected_version)...
...

  // Iterator.
  uint32_t i = 1;
```