diff --git a/src/optimizers/combinatorial/mtsp.py b/src/optimizers/combinatorial/mtsp.py index f45a442..6830263 100644 --- a/src/optimizers/combinatorial/mtsp.py +++ b/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] diff --git a/src/optimizers/core/random.py b/src/optimizers/core/random.py index 9fbc7e4..5037300 100644 --- a/src/optimizers/core/random.py +++ b/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 diff --git a/tests/test_determinism.py b/tests/test_determinism.py index 9ed18e4..3aa5d4a 100644 --- a/tests/test_determinism.py +++ b/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)