Fix set_num_parlay_workers bug.

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,
+    ],
+)

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())

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"),

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)
 
 

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

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()