Parallelize AntColonyMTSP's cluster solving

src/optimizers/combinatorial/mtsp.py

@@ -2,12 +2,13 @@
 from typing import Literal
 
 import numpy as np
+from joblib import Parallel, cpu_count, delayed
 from sklearn.cluster import KMeans, SpectralClustering
 
 from .base import TSPBase
 from .aco import AntColonyTSPConfig, AntColonyTSP
 from ..core.base import OptimizerResult, create_from_dict, literal_options
-from ..core.random import get_seed
+from ..core.random import get_seed, spawn_stream_roots, use_stream_root
 from ..core.types import AF
 
 # NOTE: "FCM" is accepted but always raises NotImplementedError (see
@@ -17,6 +18,31 @@
 ClusterMethod = Literal["kmeans", "spectral", "FCM"]
 
 
+def _solve_cluster(
+    cluster_id: int,
+    cluster: list[int],
+    city_locations: AF,
+    base_config: "AntColonyMTSPConfig",
+    n_jobs: int,
+    stream_root: np.random.SeedSequence,
+) -> OptimizerResult:
+    """Solve one cluster's independent ACO tour. Module-level so it's picklable
+    for joblib's ``processes`` backend; ``stream_root`` isolates this cluster's
+    internal RNG draws from every other concurrently-running cluster."""
+    with use_stream_root(stream_root):
+        cluster_cities = city_locations[cluster, :]
+        cluster_config = create_from_dict(base_config.__dict__, AntColonyTSPConfig)
+        cluster_config.name = f"{base_config.name}-{cluster_id + 1}"
+        cluster_config.n_jobs = n_jobs
+        tsp_solve = AntColonyTSP(config=cluster_config, city_locations=cluster_cities)
+        cluster_result = tsp_solve.solve()
+        # Map cluster indices back to original indices
+        cluster_result.solution_vector = np.array(
+            [cluster[i] for i in cluster_result.solution_vector]
+        )
+        return cluster_result
+
+
 @dataclass
 class AntColonyMTSPConfig(AntColonyTSPConfig):
     n_clusters: int = 10
@@ -33,33 +59,39 @@ def __init__(self, *, config: AntColonyMTSPConfig, city_locations: AF):
         super().__init__(config=config, city_locations=city_locations)
 
     def solve(self, *, preserve_percent: float = 0.0) -> OptimizerResult:
-        # Each cluster's ACO run is independent and could in principle run in
-        # parallel, but every solver here draws from the seeded global RNG via
-        # core.random.spawn_streams()/rng() -- and spawn_streams' docstring is
-        # explicit that it must be called from a single thread (its counter
-        # isn't synchronized). Running clusters concurrently would race on that
-        # counter and break the reproducibility guarantee the RNG-determinism
-        # work (core/random.py) established. Parallelizing this safely needs a
-        # dedicated per-cluster stream handed in up front (spawned once, here,
-        # before dispatch) rather than each cluster spawning its own -- left
-        # sequential until that's done.
         clusters = self.do_clustering()
+        n_clusters = len(clusters)
 
-        results = []
-        for cluster_id, cluster in enumerate(clusters):
-            cluster_cities = self.city_locations[cluster, :]
-            cluster_config = create_from_dict(self.config.__dict__, AntColonyTSPConfig)
-            cluster_config.name = f"{self.config.name}-{cluster_id + 1}"
-            tsp_solve = AntColonyTSP(
-                config=cluster_config, city_locations=cluster_cities
-            )
-            cluster_result = tsp_solve.solve()
-            # Map cluster indices back to original indices
-            cluster_result.solution_vector = np.array(
-                [cluster[i] for i in cluster_result.solution_vector]
-            )
+        # Split the configured processor budget between cluster-level and
+        # per-cluster ant-level parallelism instead of giving every cluster
+        # the full budget (which would oversubscribe once clusters run
+        # concurrently) -- this is what the old TODO here ("handle the number
+        # of processors based upon parallel clusters") was asking for.
+        total_jobs = self.config.n_jobs if self.config.n_jobs > 0 else cpu_count() - 1
+        outer_jobs = max(1, min(n_clusters, total_jobs))
+        inner_jobs = max(1, total_jobs // outer_jobs)
+
+        # Each cluster's ACO run draws from the seeded RNG via
+        # core.random.spawn_streams()/rng() internally, and spawn_streams'
+        # counter isn't itself synchronized across threads. Spawn one
+        # independent stream *root* per cluster up front, here, from the
+        # single (calling) thread, then have each cluster task stand in its
+        # own root for the duration of its run (use_stream_root) so its
+        # internal spawn_streams() calls draw from that root instead of
+        # racing on the shared global one -- see core/random.py.
+        stream_roots = spawn_stream_roots(n_clusters)
 
-            results.append(cluster_result)
+        results = Parallel(n_jobs=outer_jobs, prefer=self.config.joblib_prefer)(
+            delayed(_solve_cluster)(
+                cluster_id,
+                cluster,
+                self.city_locations,
+                self.config,
+                inner_jobs,
+                stream_roots[cluster_id],
+            )
+            for cluster_id, cluster in enumerate(clusters)
+        )
 
         optimal_paths = [result.solution_vector for result in results]
 

src/optimizers/core/random.py

@@ -79,6 +79,18 @@ def rng() -> np.random.Generator:
     return _global_rng
 
 
+def _current_worker_root() -> np.random.SeedSequence:
+    local_root: np.random.SeedSequence | None = getattr(
+        _thread_local, "worker_sequence", None
+    )
+    if local_root is not None:
+        return local_root
+    if _worker_sequence is None:
+        set_seed(None)
+    assert _worker_sequence is not None
+    return _worker_sequence
+
+
 def spawn_streams(n: int) -> list[np.random.Generator]:
     """``n`` independent Generators, derived deterministically from the seed.
 
@@ -92,12 +104,30 @@ def spawn_streams(n: int) -> list[np.random.Generator]:
     Successive calls return successive families, so a per-generation call gives
     every generation fresh numbers while remaining a pure function of the seed and
     the call order. Call it from one thread -- the counter it advances is not
-    itself synchronized.
+    itself synchronized -- unless that thread is inside a :func:`use_stream_root`
+    scope, in which case it draws from that scope's own independent root instead
+    of the shared global one (see :func:`use_stream_root`).
     """
-    if _worker_sequence is None:
-        set_seed(None)
-    assert _worker_sequence is not None
-    return [np.random.default_rng(child) for child in _worker_sequence.spawn(n)]
+    root = _current_worker_root()
+    return [np.random.default_rng(child) for child in root.spawn(n)]
+
+
+def spawn_stream_roots(n: int) -> list[np.random.SeedSequence]:
+    """``n`` independent ``SeedSequence`` roots, for a *nested* level of parallelism.
+
+    :func:`spawn_streams` hands out ready-to-use Generators for leaf tasks. This
+    instead returns the ``SeedSequence`` objects themselves, for a task that is
+    itself going to dispatch further parallel work which internally calls
+    :func:`spawn_streams` many times (e.g. one independent multi-generation
+    solver run per cluster). Call it once, from the single thread doing the
+    dispatching, before launching the nested tasks -- the same one-thread
+    requirement as :func:`spawn_streams`, and for the same reason (the spawn
+    counter it advances isn't synchronized). Each task then wraps its own work in
+    ``with use_stream_root(roots[i]):`` so its internal ``spawn_streams()`` calls
+    draw from its own independent sub-tree instead of racing on the shared root.
+    """
+    root = _current_worker_root()
+    return list(root.spawn(n))
 
 
 @contextmanager
@@ -113,3 +143,25 @@ def use_stream(generator: np.random.Generator) -> Iterator[None]:
         yield
     finally:
         _thread_local.rng = previous
+
+
+@contextmanager
+def use_stream_root(seed_sequence: np.random.SeedSequence) -> Iterator[None]:
+    """Make ``seed_sequence`` the root :func:`spawn_streams` draws from, for a task.
+
+    Companion to :func:`use_stream`: where ``use_stream`` overrides what plain
+    ``rng()`` calls see on this thread, this overrides what *nested*
+    ``spawn_streams()``/``spawn_stream_roots()`` calls see, so a task that
+    itself dispatches further parallel work draws from its own independent
+    sub-tree instead of racing on the shared global root's spawn counter.
+    Restores whatever was in place on exit, so nesting is safe and a worker
+    thread reused by a later task never inherits the previous task's root.
+    """
+    previous: np.random.SeedSequence | None = getattr(
+        _thread_local, "worker_sequence", None
+    )
+    _thread_local.worker_sequence = seed_sequence
+    try:
+        yield
+    finally:
+        _thread_local.worker_sequence = previous

tests/test_determinism.py

@@ -33,8 +33,16 @@
     ParticleSwarmOptimizer,
     ParticleSwarmOptimizerConfig,
 )
+from optimizers.combinatorial.mtsp import AntColonyMTSPConfig, AntColonyMTSP
 from optimizers.continuous.variables import InputContinuousVariable
-from optimizers.core.random import rng, set_seed, spawn_streams, use_stream
+from optimizers.core.random import (
+    rng,
+    set_seed,
+    spawn_stream_roots,
+    spawn_streams,
+    use_stream,
+    use_stream_root,
+)
 
 OPTIMIZERS = {
     "ga": (GeneticAlgorithmOptimizer, GeneticAlgorithmOptimizerConfig),
@@ -170,3 +178,109 @@ def test_worker_streams_do_not_disturb_the_main_stream():
     for generator in spawn_streams(4):
         generator.random(10)
     assert rng().random(5).tolist() == expected
+
+
+# ---------------------------------------------------------------------------
+# Nested parallelism: spawn_stream_roots / use_stream_root
+#
+# AntColonyMTSP dispatches one independent, itself-multi-generation ACO run
+# per cluster in parallel. Each of those runs calls spawn_streams() many times
+# internally, so giving the clusters plain spawned Generators (like leaf
+# tasks get) isn't enough -- two clusters running concurrently would each be
+# advancing spawn_streams()'s single shared counter from a different thread.
+# spawn_stream_roots()/use_stream_root() give each such task its own
+# independent root to spawn *from*, so its internal spawn_streams() calls
+# can't collide with any other concurrently-running task's.
+# ---------------------------------------------------------------------------
+
+
+def test_stream_roots_are_deterministic_and_distinct():
+    set_seed(7)
+    first = [np.random.default_rng(r).random(4).tolist() for r in spawn_stream_roots(3)]
+    set_seed(7)
+    second = [
+        np.random.default_rng(r).random(4).tolist() for r in spawn_stream_roots(3)
+    ]
+
+    assert first == second, "the same seed must produce the same root family"
+    assert first[0] != first[1] != first[2], "roots must not collide"
+
+
+def test_use_stream_root_isolates_nested_spawn_streams():
+    """Two tasks standing in different roots must not see each other's
+    nested spawn_streams() numbers, even though both call it the same way."""
+    set_seed(7)
+    root_a, root_b = spawn_stream_roots(2)
+
+    with use_stream_root(root_a):
+        a_children = [g.random(4).tolist() for g in spawn_streams(2)]
+    with use_stream_root(root_b):
+        b_children = [g.random(4).tolist() for g in spawn_streams(2)]
+
+    assert a_children != b_children
+
+
+def test_use_stream_root_scopes_to_the_calling_thread_and_restores():
+    """Nested spawn_streams() calls inside a use_stream_root scope must not
+    advance the shared global root's own spawn counter -- so code outside the
+    scope sees exactly the numbers it would have if the scope never ran."""
+    set_seed(7)
+    (nested_root,) = spawn_stream_roots(1)
+    with use_stream_root(nested_root):
+        spawn_streams(5)  # busywork inside the nested root
+    outside = [g.random(4).tolist() for g in spawn_streams(2)]
+
+    set_seed(7)
+    spawn_stream_roots(1)  # same single spawn from the global root as above
+    expected_outside = [g.random(4).tolist() for g in spawn_streams(2)]
+
+    assert outside == expected_outside
+
+
+def test_use_stream_root_nests():
+    set_seed(7)
+    first, second = spawn_stream_roots(2)
+    with use_stream_root(first):
+        with use_stream_root(second):
+            inner = spawn_streams(1)[0].random(4).tolist()
+        middle = spawn_streams(1)[0].random(4).tolist()
+    assert inner != middle, "nested and restored scopes must draw from different roots"
+
+
+# ---------------------------------------------------------------------------
+# AntColonyMTSP: clusters solve concurrently via joblib; each must be immune
+# to how the others happen to be scheduled.
+# ---------------------------------------------------------------------------
+
+
+def solve_mtsp_once(seed, *, n_jobs=4):
+    """One seeded AntColonyMTSP run over a fixed city layout."""
+    city_locations = np.random.default_rng(42).uniform(0.0, 10.0, size=(24, 2))
+    set_seed(seed)
+    config = AntColonyMTSPConfig(
+        name="determinism-mtsp",
+        num_generations=3,
+        population_size=8,
+        n_clusters=3,
+        clustering_method="kmeans",
+        stop_after_iterations=8,
+        n_jobs=n_jobs,
+        joblib_prefer="threads",
+    )
+    with (
+        contextlib.redirect_stdout(io.StringIO()),
+        contextlib.redirect_stderr(io.StringIO()),
+    ):
+        result = AntColonyMTSP(config=config, city_locations=city_locations).solve()
+    return round(float(result.solution_score), 9)
+
+
+def test_mtsp_parallel_clusters_are_reproducible():
+    """Clusters run concurrently (n_clusters=3, n_jobs=4); a seeded run must
+    still be a pure function of the seed regardless of scheduling."""
+    runs = [solve_mtsp_once(3) for _ in range(3)]
+    assert len(set(runs)) == 1, f"mtsp varied across runs: {runs}"
+
+
+def test_mtsp_different_seeds_still_differ():
+    assert solve_mtsp_once(3) != solve_mtsp_once(4)