diff --git a/dgf/src/gbbs/BUILD b/dgf/src/gbbs/BUILD index b905d17..ad94ea7 100644 --- a/dgf/src/gbbs/BUILD +++ b/dgf/src/gbbs/BUILD @@ -125,3 +125,16 @@ py_test( # numpy dep, ], ) + +# Unlike the tests above, this target deliberately does not set +# PARLAY_NUM_THREADS: the variable takes precedence over +# set_num_parlay_workers(), so the worker count could not be exercised with it +# set. This test only manipulates the scheduler and runs no graph algorithms. +py_test( + name = "parlay_workers_test", + srcs = ["parlay_workers_test.py"], + deps = [ + ":loader", + # absl/testing:absltest dep, + ], +) diff --git a/dgf/src/gbbs/connected_components_test.py b/dgf/src/gbbs/connected_components_test.py index 602dc8b..85fddc5 100644 --- a/dgf/src/gbbs/connected_components_test.py +++ b/dgf/src/gbbs/connected_components_test.py @@ -73,10 +73,6 @@ def _make_directed_cycle_with_tail() -> loader.GbbsGraphHandle: class ValidateGraphParamsTest(parameterized.TestCase): - def setUp(self): - super().setUp() - loader.set_num_parlay_workers(1) - @parameterized.named_parameters( dict( testcase_name="simple_union_async", @@ -165,10 +161,6 @@ def test_symmetric_only_params_reject_asymmetric_graph( class ConnectedComponentsTest(parameterized.TestCase): """Tests for the connected_components() dispatch function.""" - def setUp(self): - super().setUp() - loader.set_num_parlay_workers(1) - @parameterized.named_parameters( dict( testcase_name="default_params", @@ -320,10 +312,6 @@ def test_progress_flag_does_not_change_result(self): class StronglyConnectedComponentsTest(parameterized.TestCase): """Tests for the SCC algorithm via connected_components().""" - def setUp(self): - super().setUp() - loader.set_num_parlay_workers(1) - def test_scc_finds_cycle_and_tail(self): """Directed cycle 0→1→2→0 with tail 2→3 yields 2 SCCs.""" graph = _make_directed_cycle_with_tail() @@ -404,10 +392,6 @@ def test_scc_labels_length_matches_num_nodes(self): class IsSymmetricTest(absltest.TestCase): """Tests for the GbbsGraphHandle.is_symmetric() nanobind binding.""" - def setUp(self): - super().setUp() - loader.set_num_parlay_workers(1) - def test_symmetric_graph_reports_symmetric(self): graph = _make_chain_graph(num_nodes=3, symmetric=True) self.assertTrue(graph.is_symmetric()) diff --git a/dgf/src/gbbs/gbbs_ext.cc b/dgf/src/gbbs/gbbs_ext.cc index 29fa1d5..ecc4438 100644 --- a/dgf/src/gbbs/gbbs_ext.cc +++ b/dgf/src/gbbs/gbbs_ext.cc @@ -37,7 +37,11 @@ NB_MODULE(_gbbs_ext, m) { m.def( "set_num_parlay_workers", [](unsigned int num_workers) { - // Technically this resets the global scheduler with new threads. + // Drop the running scheduler first. `initialize_scheduler` returns the + // existing scheduler and ignores `num_workers` if one is already + // running, and `global_scheduler` keeps the scheduler created at module + // import alive for the lifetime of the module. + global_scheduler.shared.reset(); global_scheduler = parlay::initialize_scheduler(num_workers); }, nb::arg("num_workers"), diff --git a/dgf/src/gbbs/loader.py b/dgf/src/gbbs/loader.py index b304f5d..a8fe34e 100644 --- a/dgf/src/gbbs/loader.py +++ b/dgf/src/gbbs/loader.py @@ -74,11 +74,28 @@ def set_num_parlay_workers(num_workers: int) -> None: Calling this function is an optional user-exposed knob to customize or throttle thread concurrency. Typically, this value is set once per program - before executing graph operations. + before executing graph operations. Each call tears down the running thread + pool and starts a new one, so it must not be called while another thread is + executing a Parlay operation. + + `PARLAY_NUM_THREADS` takes precedence over this function: Parlay reads it + every time a scheduler is created, so the worker count cannot be changed + programmatically while it is set. Args: num_workers: The desired number of worker threads. + + Raises: + RuntimeError: If `PARLAY_NUM_THREADS` is set, which would make this call + silently ineffective. """ + env_num_threads = os.environ.get("PARLAY_NUM_THREADS") + if env_num_threads is not None: + raise RuntimeError( + f"PARLAY_NUM_THREADS is set to {env_num_threads!r}, which takes" + " precedence over set_num_parlay_workers(). Unset it to set the worker" + " count programmatically." + ) _gbbs_ext.set_num_parlay_workers(num_workers) diff --git a/dgf/src/gbbs/loader_test.py b/dgf/src/gbbs/loader_test.py index 37c9f9f..866d900 100644 --- a/dgf/src/gbbs/loader_test.py +++ b/dgf/src/gbbs/loader_test.py @@ -25,7 +25,6 @@ class LoaderTest(parameterized.TestCase): def test_read_gbbs_graph_from_parquet(self): - loader.set_num_parlay_workers(1) work_dir = self.create_tempdir().full_path # Create a small homogeneous graph schema diff --git a/dgf/src/gbbs/parlay_workers_test.py b/dgf/src/gbbs/parlay_workers_test.py new file mode 100644 index 0000000..2375162 --- /dev/null +++ b/dgf/src/gbbs/parlay_workers_test.py @@ -0,0 +1,40 @@ +# Copyright 2022 Google LLC. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from unittest import mock + +from absl.testing import absltest +from dgf.src.gbbs import loader + + +class ParlayWorkersTest(absltest.TestCase): + + def test_set_num_parlay_workers_takes_effect(self): + # The scheduler is already running by the time this test starts: it is + # created when the extension module is imported. + loader.set_num_parlay_workers(2) + self.assertEqual(loader.num_parlay_workers(), 2) + + loader.set_num_parlay_workers(1) + self.assertEqual(loader.num_parlay_workers(), 1) + + def test_set_num_parlay_workers_raises_when_env_var_is_set(self): + with mock.patch.dict(os.environ, {"PARLAY_NUM_THREADS": "4"}): + with self.assertRaisesRegex(RuntimeError, "PARLAY_NUM_THREADS is set"): + loader.set_num_parlay_workers(2) + + +if __name__ == "__main__": + absltest.main()