Various improvements of the in-process sampler.

benchmark/BUILD

@@ -44,6 +44,7 @@ py_library(
         ":utils",
         "//dgf",
         "//dgf/src/util:log",
+        # numpy dep,
     ],
 )
 

benchmark/in_process_sampling.py

@@ -21,6 +21,7 @@
 import dgf
 from dgf.benchmark import utils as benchmark_utils
 from dgf.src.util import log
+import numpy as np
 
 
 class OutputFormat(enum.Enum):
@@ -58,6 +59,7 @@ def __init__(
       output_format: OutputFormat = OutputFormat.NUMPY,
       edgeset_to_mask: Optional[str] = None,
       with_replacement: bool = False,
+      multi_visit: bool = True,
   ):
     self.seed_nodeset = seed_nodeset
     self.extract_features = extract_features
@@ -70,6 +72,7 @@ def __init__(
     self.batch_size = 12
     self.edgeset_to_mask = edgeset_to_mask
     self.set_unit_multiplicator(self.batch_size)
+    self.multi_visit = multi_visit
 
     self.sum_sampled_nodes = 0
     self.num_samples = 0
@@ -84,6 +87,7 @@ def setup(self):
         num_hops=self.num_hops,
         hop_width=self.hop_width,
         with_replacement=self.with_replacement,
+        multi_visit=self.multi_visit,
     )
     sampling_plan = dgf.sampling.simple_sampling_config_to_sampling_plan(
         self.sampling_config,
@@ -134,12 +138,12 @@ def output_fn(
     self.output_fn = output_fn
 
   def run_unit(self):
-    seed_node_idxs = [
-        random.randrange(0, self.num_nodes) for _ in range(self.batch_size)
-    ]
+    seed_node_idxs = np.random.randint(
+        0, self.num_nodes, size=self.batch_size, dtype=np.int64
+    )
     if self.edgeset_to_mask is not None:
       # Pass dummy masked edge indices (e.g. all 0).
-      masked_edge_idxs = [0 for _ in range(self.batch_size)]
+      masked_edge_idxs = np.zeros(self.batch_size, dtype=np.int64)
       samples = self.sampler.sample(
           seed_node_idxs, masked_edge_idxs=masked_edge_idxs
       )
@@ -154,14 +158,15 @@ def run_unit(self):
 
   def details(self) -> str:
     return (
-        f"num_hops={self.sampling_config.num_hops}"
-        f" hop_width={self.sampling_config.hop_width}"
-        f" extract_features={self.extract_features}"
-        f" output_format={self.output_format.value}"
-        f" with_replacement={self.with_replacement}"
-        f" batch_size={self.batch_size}"
-        f" edgeset_to_mask={self.edgeset_to_mask}"
-        f" nodes_per_sample={self.sum_sampled_nodes / self.num_samples}"
+        f"hops={self.sampling_config.num_hops}"
+        f" width={self.sampling_config.hop_width}"
+        f" feat.={self.extract_features}"
+        f" format={self.output_format.value}"
+        f" with_rep.={int(self.with_replacement)}"
+        f" batch={self.batch_size}"
+        f" mask={self.edgeset_to_mask}"
+        f" nodes/spl.={int(self.sum_sampled_nodes / self.num_samples)}"
+        f" multi_visit={self.multi_visit}"
     )
 
 
@@ -261,31 +266,50 @@ def in_process_sampling(
   edgeset_to_mask = edgeset_names[0] if edgeset_names else None
 
   for num_hops in list_num_hops:
-    for extract_features in [False, True]:
-      benchmarker.run(
-          GenGraphSubsets(
-              graph=graph,
-              schema=schema,
-              seed_nodeset=seed_nodeset,
-              extract_features=extract_features,
-              num_hops=num_hops,
-          ),
-          repetitions=1,
-          warmup_repetitions=1,
-      )
+    for extract_features in [True, False]:
+      for with_replacement in [True, False]:
+        benchmarker.run(
+            GenGraphSubsets(
+                graph=graph,
+                schema=schema,
+                seed_nodeset=seed_nodeset,
+                extract_features=extract_features,
+                num_hops=num_hops,
+            ),
+            repetitions=1,
+            warmup_repetitions=1,
+        )
 
-      benchmarker.run(
-          GenGraphSamples(
-              num_hops=num_hops,
-              graph=graph,
-              schema=schema,
-              seed_nodeset=seed_nodeset,
-              extract_features=extract_features,
-              output_format=OutputFormat.NUMPY,
-          ),
-          repetitions=1,
-          warmup_repetitions=1,
-      )
+        benchmarker.run(
+            GenGraphSamples(
+                num_hops=num_hops,
+                graph=graph,
+                schema=schema,
+                seed_nodeset=seed_nodeset,
+                extract_features=extract_features,
+                output_format=OutputFormat.NUMPY,
+                with_replacement=with_replacement,
+                multi_visit=True,
+            ),
+            repetitions=1,
+            warmup_repetitions=1,
+        )
+
+        if not with_replacement:
+          benchmarker.run(
+              GenGraphSamples(
+                  num_hops=num_hops,
+                  graph=graph,
+                  schema=schema,
+                  seed_nodeset=seed_nodeset,
+                  extract_features=extract_features,
+                  output_format=OutputFormat.NUMPY,
+                  with_replacement=with_replacement,
+                  multi_visit=False,
+              ),
+              repetitions=1,
+              warmup_repetitions=1,
+          )
 
       if edgeset_to_mask is not None:
         benchmarker.run(

benchmark/in_process_sampling_main.py

@@ -27,7 +27,20 @@
   --graph_path=/cns/iz-d/home/research-graph/public/graphflow_datasets/fetch_repo/ogb_arxiv\
   --seed_nodeset=nodes
 
+Profiling Instructions:
+1. Build the binary in opt mode:
+   blaze build -c opt --cpu=haswell //third_party/py/dgf/benchmark:in_process_sampling_main
+2. Run with CPUPROFILE enabled:
+   CPUPROFILE=/tmp/prof.out ./blaze-bin/third_party/py/dgf/benchmark/in_process_sampling_main \
+     --work_dir=/tmp/gf_benchmark \
+     --graph_path=/cns/iz-d/home/research-graph/public/graphflow_datasets/fetch_repo/ogb_arxiv \
+    --seed_nodeset=nodes
 
+3. View the profiling results:
+   # Top functions:
+   pprof --text ./blaze-bin/third_party/py/dgf/benchmark/in_process_sampling_main /tmp/prof.out | head -n 40
+   # Interactive Web UI:
+   pprof -http=localhost:8080 ./blaze-bin/third_party/py/dgf/benchmark/in_process_sampling_main /tmp/prof.out
 
 # Note: For very large datasets like `papers100` (1.6B edges), avoid using
 # `--work_dir` to prevent local caching, as the dataset is very large.
@@ -58,9 +71,9 @@
 )
 _LIST_NUM_HOPS = flags.DEFINE_list(
     "list_num_hops",
-    "2,3,4",
+    "3",
     "A comma-separated list of integers representing the number of hops to"
-    " sample.",
+    " sample. Can also be a single value.",
 )
 _BENCHMARK_OUTPUT_FORMATS = flags.DEFINE_bool(
     "benchmark_output_formats",

dgf/src/sampling/config.py

@@ -69,6 +69,7 @@ class SimpleSamplingConfig:
   reverse: bool = True
   with_replacement: bool = False
   temporal_sampling: bool = False
+  multi_visit: bool = True
 
 
 @dataclasses.dataclass
@@ -117,6 +118,7 @@ class SamplingPlan:
   root: PlanNode
   with_replacement: bool = False
   temporal_sampling: bool = False
+  multi_visit: bool = True
   edgeset_timestamp_features: Dict[str, str] = dataclasses.field(
       default_factory=dict
   )
@@ -168,14 +170,12 @@ def rec_build(nodeset: str, depth: int) -> PlanNode:
 
   edgeset_ts_features = {}
   if src.temporal_sampling:
-    edgeset_ts_features = temporal_util.edgeset_timestamp_features(
-        schema
-    )
+    edgeset_ts_features = temporal_util.edgeset_timestamp_features(schema)
 
   return SamplingPlan(
       root=rec_build(src.seed_nodeset, depth=0),
       with_replacement=src.with_replacement,
       temporal_sampling=src.temporal_sampling,
+      multi_visit=src.multi_visit,
       edgeset_timestamp_features=edgeset_ts_features,
   )
-

dgf/src/sampling/in_memory_sampler.h

@@ -214,6 +214,7 @@ struct SamplingPlan {
 
   std::unique_ptr<Node> root;
   bool with_replacement;
+  bool multi_visit;
 
   // Total number of steps. "step_idx" in "Node" are in [0, num_steps).
   size_t num_steps;

dgf/src/sampling/in_memory_sampler.py

@@ -335,15 +335,18 @@ def add_features_to_samples(
   if return_features or not return_node_idxs:
     # Extract feature values.
     # TODO(gbm): Do this in C++.
+    full_node_sets = full_graph.node_sets
     for sample in samples:
       for node_set_name, node_set in sample.node_sets.items():
         node_idxs = node_set.features["#idx"]
         if not return_node_idxs:
           del node_set.features["#idx"]
         if return_features:
-          features = full_graph.node_sets[node_set_name].features
+          features = full_node_sets[node_set_name].features
           for feature_name, full_feature_value in features.items():
-            node_set.features[feature_name] = full_feature_value[node_idxs]
+            node_set.features[feature_name] = np.take(
+                full_feature_value, node_idxs, axis=0
+            )
 
 
 def create_sampler(