diff --git a/simulation_engine/multi_point.py b/simulation_engine/multi_point.py new file mode 100644 index 00000000..8f91fb4b --- /dev/null +++ b/simulation_engine/multi_point.py @@ -0,0 +1,135 @@ +from __future__ import annotations + +from collections.abc import Callable, Sequence +from concurrent.futures import Future, ProcessPoolExecutor, as_completed +from dataclasses import dataclass +import logging +import os +import random +import time +from typing import TypeVar + +import numpy as np + + +LOGGER = logging.getLogger(__name__) + +_ENV_PROCESSES = "SIM_MULTI_POINT_JOB_PROCESSES" +_MAX_NUMPY_SEED = (2**32) - 1 +_MISSING = object() + +P = TypeVar("P") +R = TypeVar("R") + + +@dataclass(frozen=True) +class PointWorkItem: + index: int + point: object + seed: int | None = None + + +def resolve_process_count(point_count: int) -> int: + if point_count <= 0: + return 1 + raw_value = os.getenv(_ENV_PROCESSES) + requested = 1 + if raw_value: + try: + requested = int(raw_value) + except ValueError: + requested = 1 + requested = max(1, requested) + cpu_count = os.cpu_count() or 1 + return max(1, min(requested, cpu_count, point_count)) + + +def seed_point_worker(seed: int | None) -> None: + if seed is None: + return + random.seed(seed) + np.random.seed(seed % _MAX_NUMPY_SEED) + + +def run_points( + *, + run_kind: str, + points: Sequence[P], + evaluate_point: Callable[[PointWorkItem], R], + parallel_allowed: bool = True, + on_complete: Callable[[int, int], None] | None = None, +) -> list[R]: + """Evaluate independent points and return results in input order.""" + point_count = len(points) + process_count = resolve_process_count(point_count) if parallel_allowed else 1 + started = time.monotonic() + try: + if process_count == 1: + results = _run_points_sequential(points, evaluate_point, on_complete) + else: + results = _run_points_parallel( + points, + evaluate_point, + process_count=process_count, + on_complete=on_complete, + ) + return results + finally: + duration_ms = int((time.monotonic() - started) * 1000) + LOGGER.info( + "multi_point_job_completed", + extra={ + "event": "multi_point_job_completed", + "run_kind": run_kind, + "point_count": point_count, + "process_count": process_count, + "duration_ms": duration_ms, + }, + ) + + +def _run_points_sequential( + points: Sequence[P], + evaluate_point: Callable[[PointWorkItem], R], + on_complete: Callable[[int, int], None] | None, +) -> list[R]: + total = len(points) + results: list[R] = [] + for index, point in enumerate(points): + results.append(evaluate_point(PointWorkItem(index=index, point=point))) + if on_complete is not None: + on_complete(index + 1, total) + return results + + +def _run_points_parallel( + points: Sequence[P], + evaluate_point: Callable[[PointWorkItem], R], + *, + process_count: int, + on_complete: Callable[[int, int], None] | None, +) -> list[R]: + total = len(points) + results: list[object] = [_MISSING] * total + seeds = [random.getrandbits(32) for _ in points] + executor = ProcessPoolExecutor(max_workers=process_count) + futures: dict[Future[R], int] = {} + try: + for index, point in enumerate(points): + work_item = PointWorkItem(index=index, point=point, seed=seeds[index]) + futures[executor.submit(evaluate_point, work_item)] = index + completed = 0 + for future in as_completed(futures): + index = futures[future] + results[index] = future.result() + completed += 1 + if on_complete is not None: + on_complete(completed, total) + except BaseException: + for future in futures: + future.cancel() + executor.shutdown(wait=False, cancel_futures=True) + raise + else: + executor.shutdown(wait=True) + return [result for result in results if result is not _MISSING] diff --git a/tests/test_multi_point_jobs.py b/tests/test_multi_point_jobs.py new file mode 100644 index 00000000..e39e7acb --- /dev/null +++ b/tests/test_multi_point_jobs.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import os +import random + +import pytest + +from simulation_engine import multi_point + + +class _InlineFuture: + def __init__(self, value): + self._value = value + self.cancelled = False + + def result(self): + return self._value + + def cancel(self): + self.cancelled = True + return True + + +class _InlineProcessPoolExecutor: + def __init__(self, max_workers): + self.max_workers = max_workers + self.shutdown_calls = [] + + def submit(self, fn, arg): + return _InlineFuture(fn(arg)) + + def shutdown(self, *, wait=True, cancel_futures=False): + self.shutdown_calls.append((wait, cancel_futures)) + + +def _install_inline_process_pool(monkeypatch): + monkeypatch.setattr( + multi_point, + "ProcessPoolExecutor", + _InlineProcessPoolExecutor, + ) + monkeypatch.setattr(multi_point, "as_completed", lambda futures: list(futures)) + + +def _return_indexed_value(work_item: multi_point.PointWorkItem): + multi_point.seed_point_worker(work_item.seed) + return { + "index": work_item.index, + "value": work_item.point, + "seed": work_item.seed, + "random": random.random(), + } + + +def test_resolve_process_count_defaults_to_one(monkeypatch): + monkeypatch.delenv("SIM_MULTI_POINT_JOB_PROCESSES", raising=False) + assert multi_point.resolve_process_count(10) == 1 + + +def test_resolve_process_count_caps_at_cpu_and_points(monkeypatch): + monkeypatch.setenv("SIM_MULTI_POINT_JOB_PROCESSES", "8") + monkeypatch.setattr(os, "cpu_count", lambda: 4) + assert multi_point.resolve_process_count(3) == 3 + assert multi_point.resolve_process_count(10) == 4 + + +def test_run_points_parallel_preserves_input_order_and_assigns_seeds(monkeypatch): + monkeypatch.setenv("SIM_MULTI_POINT_JOB_PROCESSES", "2") + _install_inline_process_pool(monkeypatch) + random.seed(12345) + first = multi_point.run_points( + run_kind="test", + points=["a", "b", "c", "d"], + evaluate_point=_return_indexed_value, + ) + random.seed(12345) + second = multi_point.run_points( + run_kind="test", + points=["a", "b", "c", "d"], + evaluate_point=_return_indexed_value, + ) + + assert [row["value"] for row in first] == ["a", "b", "c", "d"] + assert [row["index"] for row in first] == [0, 1, 2, 3] + assert [row["seed"] for row in first] == [row["seed"] for row in second] + assert [row["random"] for row in first] == [row["random"] for row in second] + assert all(row["seed"] is not None for row in first) + + +def test_run_points_sequential_does_not_seed_or_fork(monkeypatch): + monkeypatch.setenv("SIM_MULTI_POINT_JOB_PROCESSES", "2") + results = multi_point.run_points( + run_kind="test", + points=[1, 2], + evaluate_point=_return_indexed_value, + parallel_allowed=False, + ) + + assert [row["value"] for row in results] == [1, 2] + assert [row["seed"] for row in results] == [None, None] + + +def test_run_points_progress_callback_runs_for_completed_points(monkeypatch): + monkeypatch.setenv("SIM_MULTI_POINT_JOB_PROCESSES", "2") + _install_inline_process_pool(monkeypatch) + updates: list[tuple[int, int]] = [] + + multi_point.run_points( + run_kind="test", + points=[1, 2, 3], + evaluate_point=_return_indexed_value, + on_complete=lambda current, total: updates.append((current, total)), + ) + + assert sorted(updates) == [(1, 3), (2, 3), (3, 3)] + + +def test_run_points_cancels_pending_work_when_progress_raises(monkeypatch): + monkeypatch.setenv("SIM_MULTI_POINT_JOB_PROCESSES", "2") + _install_inline_process_pool(monkeypatch) + + def cancel_after_first(current: int, total: int) -> None: + raise RuntimeError(f"cancelled at {current}/{total}") + + with pytest.raises(RuntimeError, match="cancelled at"): + multi_point.run_points( + run_kind="test", + points=[1, 2, 3], + evaluate_point=_return_indexed_value, + on_complete=cancel_after_first, + )