From 9a0bc7ba3e655da3bdff11f707b787d02779bbe5 Mon Sep 17 00:00:00 2001 From: Mathieu Guillame-Bert Date: Mon, 24 Aug 2026 04:35:44 -0700 Subject: [PATCH] Various improvements of the in-process sampler. - Reduce the number of concurent mutex lockings. - Replace a treeset by a vector+sorting+dedup for the edges. - Add option to disallow multi visit of nodes (multi_visit=False). PiperOrigin-RevId: 969769906 --- benchmark/BUILD | 1 + benchmark/in_process_sampling.py | 96 ++++++++------ benchmark/in_process_sampling_main.py | 17 ++- dgf/src/sampling/config.py | 8 +- dgf/src/sampling/in_memory_sampler.h | 1 + dgf/src/sampling/in_memory_sampler.py | 7 +- dgf/src/sampling/in_memory_sampler_ext.cc | 141 +++++++++++---------- dgf/src/sampling/in_memory_sampler_nb.cc | 2 + dgf/src/sampling/in_memory_sampler_test.py | 61 +++++++++ dgf/src/util/concurrency.h | 11 +- 10 files changed, 227 insertions(+), 118 deletions(-) diff --git a/benchmark/BUILD b/benchmark/BUILD index b4a906f..8081ee9 100644 --- a/benchmark/BUILD +++ b/benchmark/BUILD @@ -44,6 +44,7 @@ py_library( ":utils", "//dgf", "//dgf/src/util:log", + # numpy dep, ], ) diff --git a/benchmark/in_process_sampling.py b/benchmark/in_process_sampling.py index 90122c8..f39b623 100644 --- a/benchmark/in_process_sampling.py +++ b/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( diff --git a/benchmark/in_process_sampling_main.py b/benchmark/in_process_sampling_main.py index d08322a..6a2f7e3 100644 --- a/benchmark/in_process_sampling_main.py +++ b/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", diff --git a/dgf/src/sampling/config.py b/dgf/src/sampling/config.py index c57d712..e5d75d0 100644 --- a/dgf/src/sampling/config.py +++ b/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, ) - diff --git a/dgf/src/sampling/in_memory_sampler.h b/dgf/src/sampling/in_memory_sampler.h index 250acf4..8fd574d 100644 --- a/dgf/src/sampling/in_memory_sampler.h +++ b/dgf/src/sampling/in_memory_sampler.h @@ -214,6 +214,7 @@ struct SamplingPlan { std::unique_ptr root; bool with_replacement; + bool multi_visit; // Total number of steps. "step_idx" in "Node" are in [0, num_steps). size_t num_steps; diff --git a/dgf/src/sampling/in_memory_sampler.py b/dgf/src/sampling/in_memory_sampler.py index 98dc08f..5c9f72f 100644 --- a/dgf/src/sampling/in_memory_sampler.py +++ b/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( diff --git a/dgf/src/sampling/in_memory_sampler_ext.cc b/dgf/src/sampling/in_memory_sampler_ext.cc index 3ddd5f3..bb8daad 100644 --- a/dgf/src/sampling/in_memory_sampler_ext.cc +++ b/dgf/src/sampling/in_memory_sampler_ext.cc @@ -14,9 +14,7 @@ #include #include #include -#include #include -#include #include #include #include @@ -51,8 +49,6 @@ #include "dgf/src/util/status_caster.h" #include "dgf/src/util/util.h" -namespace nb = nanobind; - namespace dgf::sampling::in_memory_sampler { struct SampleBuilder; // Forward declaration @@ -110,10 +106,7 @@ struct Sampler { // Pool of available SampleBuilders to reuse them across seeds and calls. std::vector> sample_builder_pool_; - std::mutex sample_builder_pool_mutex_; - - std::unique_ptr AcquireSampleBuilder(); - void ReleaseSampleBuilder(std::unique_ptr builder); + util::concurrency::Mutex sample_builder_pool_mutex_; Sampler(int num_threads) : thread_pool(num_threads) {} @@ -201,7 +194,7 @@ struct Sampler { struct SampleBuilder { struct EdgeSet { // The pair of sampled edges (i.e., src/target node index). - absl::btree_set> edges; + std::vector> edges; }; struct NodeSet { @@ -209,7 +202,7 @@ struct SampleBuilder { std::vector sampled_node_idx_to_node_idx; // Inverse mapping of "sampled_node_idx_to_node_idx". // Only used when sampling without replacement. - absl::btree_map node_idx_to_sampled_node_idx; + absl::flat_hash_map node_idx_to_sampled_node_idx; }; // Per edge-set information indexed by edgeset index. @@ -278,9 +271,20 @@ struct SampleBuilder { // Recursively sample the sub-nodes. auto& sample_target_nodeset = nodesets[plan_edge.node->nodeset_idx]; auto& sample_edgeset = edgesets[plan_edge.edgeset_idx]; + + const auto add_edge = [&](SampleIdx src, SampleIdx trg) { + if (!plan_edge.reversed) { + sample_edgeset.edges.push_back({src, trg}); + } else { + sample_edgeset.edges.push_back({trg, src}); + } + }; + for (const auto target_node : cache_node_idxs) { // Record the target node and create/reuse a sampled node. + bool recuse = true; + SampleIdx target_sampled_node; if (!sampler.plan_.with_replacement) { // Sampling without replacement. @@ -289,21 +293,15 @@ struct SampleBuilder { target_node, sample_target_nodeset.node_idx_to_sampled_node_idx.size()); target_sampled_node = target_sampled_node_it->second; + add_edge(source_sampled_node, target_sampled_node); if (inserted) { sample_target_nodeset.sampled_node_idx_to_node_idx.push_back( target_node); - DGF_STATUS_CHECK( - sample_target_nodeset.sampled_node_idx_to_node_idx.size() == - sample_target_nodeset.node_idx_to_sampled_node_idx.size()); - } - // Record the edge, if it is a new edge. - if (!plan_edge.reversed) { - sample_edgeset.edges.insert( - {source_sampled_node, target_sampled_node}); - } else { - sample_edgeset.edges.insert( - {target_sampled_node, source_sampled_node}); + DCHECK(sample_target_nodeset.sampled_node_idx_to_node_idx.size() == + sample_target_nodeset.node_idx_to_sampled_node_idx.size()); + } else if (!sampler.plan_.multi_visit) { + recuse = false; } } else { // Sampling with replacement. @@ -311,20 +309,15 @@ struct SampleBuilder { sample_target_nodeset.sampled_node_idx_to_node_idx.size(); sample_target_nodeset.sampled_node_idx_to_node_idx.push_back( target_node); - // Record the edge, if it is a new edge. - if (!plan_edge.reversed) { - sample_edgeset.edges.insert( - {source_sampled_node, target_sampled_node}); - } else { - sample_edgeset.edges.insert( - {target_sampled_node, source_sampled_node}); - } + add_edge(source_sampled_node, target_sampled_node); } // Build the sub-sample. - DGF_RETURN_IF_ERROR(RecursiveGrow(sampler, *plan_edge.node, target_node, - target_sampled_node, seed_timestamp, - masked_edge_idx, depth + 1)); + if (recuse) { + DGF_RETURN_IF_ERROR(RecursiveGrow( + sampler, *plan_edge.node, target_node, target_sampled_node, + seed_timestamp, masked_edge_idx, depth + 1)); + } } } return absl::OkStatus(); @@ -364,9 +357,26 @@ struct SampleBuilder { } // Recursive transversal. - return RecursiveGrow(sampler, *sampler.plan_.root, seed_node_idx, - sample_seed_node_idx, seed_timestamp, masked_edge_idx, - /*depth=*/0); + DGF_RETURN_IF_ERROR(RecursiveGrow(sampler, *sampler.plan_.root, + seed_node_idx, sample_seed_node_idx, + seed_timestamp, masked_edge_idx, + /*depth=*/0)); + + // Deduplicate the edges. + // Note: We also sort edges in debug mode. + const bool dedup_edges = + !sampler.plan_.with_replacement && sampler.plan_.multi_visit; + if (dedup_edges || sampler.debug_sampling_) { + for (auto& sampled_edgeset : edgesets) { + std::sort(sampled_edgeset.edges.begin(), sampled_edgeset.edges.end()); + if (dedup_edges) { + sampled_edgeset.edges.erase(std::unique(sampled_edgeset.edges.begin(), + sampled_edgeset.edges.end()), + sampled_edgeset.edges.end()); + } + } + } + return absl::OkStatus(); } absl::StatusOr ExportToInMemoryGraph(const Sampler& sampler) { @@ -397,21 +407,6 @@ struct SampleBuilder { } }; -std::unique_ptr Sampler::AcquireSampleBuilder() { - std::lock_guard lock(sample_builder_pool_mutex_); - if (sample_builder_pool_.empty()) { - return std::make_unique(rng_()); - } - auto builder = std::move(sample_builder_pool_.back()); - sample_builder_pool_.pop_back(); - return builder; -} - -void Sampler::ReleaseSampleBuilder(std::unique_ptr builder) { - std::lock_guard lock(sample_builder_pool_mutex_); - sample_builder_pool_.push_back(std::move(builder)); -} - // Creates a graph sample starting from a given seed node. // The returned `nb::object` is an instance of `InMemoryGraph` // containing the sampled subgraph. @@ -446,15 +441,28 @@ absl::StatusOr Sampler::Sample( } } - // Pre-allocate pool if needed. + std::vector> active_builders(num_seeds); + + // Pre-allocate builders + grab the ones we need. { - std::lock_guard lock(sample_builder_pool_mutex_); - while (sample_builder_pool_.size() < thread_pool.num_threads()) { - sample_builder_pool_.push_back(std::make_unique(rng_())); + util::concurrency::MutexLock lock(sample_builder_pool_mutex_); + for (size_t seed_idx = 0; seed_idx < num_seeds; seed_idx++) { + if (!sample_builder_pool_.empty()) { + active_builders[seed_idx] = std::move(sample_builder_pool_.back()); + sample_builder_pool_.pop_back(); + } else { + active_builders[seed_idx] = std::make_unique(rng_()); + } } } - std::vector> active_builders(num_seeds); + // Release builders back to pool even on failure. + const auto release_builders = [&]() { + util::concurrency::MutexLock lock(sample_builder_pool_mutex_); + for (auto& builder : active_builders) { + sample_builder_pool_.push_back(std::move(builder)); + } + }; // Create samples { @@ -469,7 +477,7 @@ absl::StatusOr Sampler::Sample( // Start the sampling. absl::Status global_status; - std::mutex global_status_mutex; + util::concurrency::Mutex global_status_mutex; std::latch latch(num_seeds); for (size_t seed_idx = 0; seed_idx < num_seeds; seed_idx++) { @@ -484,16 +492,15 @@ absl::StatusOr Sampler::Sample( masked_edge_idx = masked_edge_idxs->view()(seed_idx); } const uint64_t seed = seeds[seed_idx]; - thread_pool.Schedule([this, seed_idx, seed, &active_builders, - seed_node_idx, seed_timestamp, masked_edge_idx, - &latch, &global_status_mutex, &global_status]() { - std::unique_ptr sample_builder = AcquireSampleBuilder(); + SampleBuilder* sample_builder = active_builders[seed_idx].get(); + thread_pool.Schedule([this, sample_builder, seed, seed_node_idx, + seed_timestamp, masked_edge_idx, &latch, + &global_status_mutex, &global_status]() { sample_builder->rng.seed(seed); const auto status = sample_builder->Grow( *this, seed_node_idx, seed_timestamp, masked_edge_idx); - active_builders[seed_idx] = std::move(sample_builder); latch.count_down(); // If the sampling failed, record the failure. @@ -509,12 +516,7 @@ absl::StatusOr Sampler::Sample( // Return an error if any of the samplers failed. if (!global_status.ok()) { - // Release builders back to pool even on failure. - for (auto& builder : active_builders) { - if (builder != nullptr) { - ReleaseSampleBuilder(std::move(builder)); - } - } + release_builders(); return global_status; } } @@ -526,9 +528,10 @@ absl::StatusOr Sampler::Sample( DGF_ASSIGN_OR_RETURN(auto graph, sample_builder->ExportToInMemoryGraph(*this)); graphs.append(graph); - // Release builder back to pool. - ReleaseSampleBuilder(std::move(sample_builder)); } + + // Release builders back to pool in bulk. + release_builders(); return graphs; } diff --git a/dgf/src/sampling/in_memory_sampler_nb.cc b/dgf/src/sampling/in_memory_sampler_nb.cc index 5ecc9d5..ea8fbb6 100644 --- a/dgf/src/sampling/in_memory_sampler_nb.cc +++ b/dgf/src/sampling/in_memory_sampler_nb.cc @@ -39,8 +39,10 @@ absl::StatusOr CreateSamplingPlan( DGF_GET_ATTR_OR_RETURN(nb::object, py_root, py_plan, "root"); DGF_GET_ATTR_OR_RETURN(bool, with_replacement, py_plan, "with_replacement"); + DGF_GET_ATTR_OR_RETURN(bool, multi_visit, py_plan, "multi_visit"); plan.with_replacement = with_replacement; + plan.multi_visit = multi_visit; std::function>( const nb::object&)> diff --git a/dgf/src/sampling/in_memory_sampler_test.py b/dgf/src/sampling/in_memory_sampler_test.py index d9b259e..d75c343 100644 --- a/dgf/src/sampling/in_memory_sampler_test.py +++ b/dgf/src/sampling/in_memory_sampler_test.py @@ -1518,5 +1518,66 @@ def test_slice_timeseries_by_seed_validation_errors(self): ) +class InMemorySamplerDiamon(parameterized.TestCase): + + @classmethod + def setUpClass(cls): + super().setUpClass() + cls.schema = schema_lib.GraphSchema( + node_sets={ + "n1": schema_lib.NodeSchema( + features={ + "f1": schema_lib.FeatureSchema( + format=schema_lib.FeatureFormat.INTEGER_64 + ) + } + ), + }, + edge_sets={ + "e11": schema_lib.EdgeSchema(source="n1", target="n1"), + }, + ) + # Create a diamon followed by 3 nodes. + cls.graph = in_memory_graph_lib.InMemoryGraph( + node_sets={ + "n1": in_memory_graph_lib.InMemoryNodeSet( + features={"f1": np.array([10, 11, 12, 13, 14, 15, 16])}, + num_nodes=7, + ), + }, + edge_sets={ + "e11": in_memory_graph_lib.InMemoryEdgeSet( + adjacency=np.array( + [[0, 0, 1, 2, 3, 3, 3], [1, 2, 3, 3, 4, 5, 6]] + ) + ), + }, + ) + + def test_sample(self): + plan = config_lib.simple_sampling_config_to_sampling_plan( + config_lib.SimpleSamplingConfig( + seed_nodeset="n1", + num_hops=5, + hop_width=2, + reverse=False, + multi_visit=False, + ), + self.schema, + ) + sampler = in_memory_sampler_lib.create_sampler( + self.graph, + plan, + self.schema, + return_features=True, + return_node_idxs=False, + batch_size=5, + ) + sample = sampler.sample(0) + self.assertIsNotNone(sample) + # The diamon and exactly 2 or the last 3 nodes are sampled. + self.assertEqual(sample.node_sets["n1"].num_nodes, 6) + + if __name__ == "__main__": absltest.main() diff --git a/dgf/src/util/concurrency.h b/dgf/src/util/concurrency.h index cede52e..ebf0ae7 100644 --- a/dgf/src/util/concurrency.h +++ b/dgf/src/util/concurrency.h @@ -21,7 +21,8 @@ namespace dgf::util::concurrency { // TODO(gbm): Use XYZ's thread. -typedef std::unique_lock MutexLock; +typedef std::lock_guard MutexLock; +typedef std::unique_lock UniqueMutexLock; typedef std::thread Thread; typedef std::mutex Mutex; @@ -40,7 +41,7 @@ class Channel { // Close the channel. No new items can be push in the channel. Any "Pop" will // return immediately with an empty optional. void Close() { - MutexLock l(mutex_); + UniqueMutexLock l(mutex_); if (channel_closed_) { return; } @@ -51,7 +52,7 @@ class Channel { // Push an item in the channel. void Push(Input item) { - MutexLock l(mutex_); + UniqueMutexLock l(mutex_); if (max_items_.has_value()) { while (content_.size() >= *max_items_ && !channel_closed_) { not_full_.wait(l); @@ -68,7 +69,7 @@ class Channel { // {}. If the channel is empty but not closed, blocks. If the channel is not // empty, returns the first added element. std::optional Pop() { - MutexLock l(mutex_); + UniqueMutexLock l(mutex_); while (content_.empty() && !channel_closed_) { not_empty_.wait(l); } @@ -88,7 +89,7 @@ class Channel { std::atomic channel_closed_ = false; std::condition_variable not_empty_; // Signaled when an item is added. std::condition_variable not_full_; // Signaled when an item is removed. - std::mutex mutex_; + Mutex mutex_; std::optional max_items_; };