diff --git a/docs/how_to_guide/21_running_simulations_in_parallel.ipynb b/docs/how_to_guide/21_running_simulations_in_parallel.ipynb new file mode 100644 index 000000000..c38c70c6d --- /dev/null +++ b/docs/how_to_guide/21_running_simulations_in_parallel.ipynb @@ -0,0 +1,248 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# How to run simulations in parallel" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Many simulators are computationally expensive. When performing Simulation-Based Inference, we often need thousands of simulations. If your simulator takes even a fraction of a second, running them sequentially can be prohibitively slow.\n", + "\n", + "**Vectorization (or batching)** is one strategy to improve speed, where the simulator processes multiple parameters at once in a single call. However, this is not always straightforward to implement (e.g., rewriting legacy code) and can lead to high memory consumption if the batches are too large.\n", + "\n", + "**Parallelization** is an alternative that allows us to run multiple simulations at the same time in separate processes, reducing the total wall-clock time without requiring the simulator to be vectorized.\n", + "\n", + "This guide illustrates how to parallelize simulations using `sbi`'s `simulate_from_theta` utility. We will use a toy simulator that returns a file path—a common scenario where the simulator is an external executable that writes results to disk, making standard in-memory batching impossible or impractical." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "from typing import Any\n", + "\n", + "import numpy as np\n", + "\n", + "\n", + "def simulator(theta: Any) -> str:\n", + " \"\"\"\n", + " A toy simulator that sleeps and returns a dummy file path.\n", + " \"\"\"\n", + " time.sleep(0.1) # Simulate expensive computation\n", + "\n", + " # Format the parameter into a string suitable for a filename\n", + " # We use 4 decimal places and replace the dot with a hyphen\n", + " # e.g., 0.1234 -> 0-1234\n", + " theta_id = f\"{theta[0]:.4f}\".replace(\".\", \"_\")\n", + "\n", + " return f\"/path/to/simulation/output_{theta_id}.npy\"\n", + "\n", + "thetas = np.random.uniform(0, 1, size=(100, 1)) # 100 parameter sets" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Naive for loop\n", + "First, let's establish a baseline by running the simulations sequentially in a simple for loop. This represents the wall-clock time it takes without any parallelization optimization." + ] + }, + { + "cell_type": "code", + "execution_count": 25, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Elapsed time: 10.022978331000104 seconds\n", + "['/path/to/simulation/output_0-5192.npy', '/path/to/simulation/output_0-0983.npy', '/path/to/simulation/output_0-6823.npy', '/path/to/simulation/output_0-9948.npy', '/path/to/simulation/output_0-3902.npy', '/path/to/simulation/output_0-4715.npy', '/path/to/simulation/output_0-0382.npy', '/path/to/simulation/output_0-9635.npy', '/path/to/simulation/output_0-8807.npy', '/path/to/simulation/output_0-2155.npy']\n" + ] + } + ], + "source": [ + "start = time.perf_counter()\n", + "x = [simulator(np.array(theta)) for theta in thetas]\n", + "end = time.perf_counter()\n", + "print(f\"Elapsed time: {end - start} seconds\")\n", + "print(x[:10])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Parallel execution\n", + "\n", + "`sbi` provides the `simulate_from_theta` utility to easily parallelize simulations. It uses `joblib` under the hood.\n", + "\n", + "We can use `joblib.parallel_config` context manager to specify the number of workers (`n_jobs`)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "3f2e3e21b42048b7b727bfc75756f4ba", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + " 0%| | 0/100 [00:00 -from typing import Any, Callable, Optional, Tuple, Union +import warnings +from typing import Any, Callable, List, Optional, Tuple, Union, overload import numpy as np import torch -from joblib import Parallel, delayed -from numpy import ndarray +from joblib import Parallel, delayed, parallel_config from torch import Tensor, float32 from tqdm.auto import tqdm from sbi.utils.sbiutils import seed_all_backends +X = Tensor | np.ndarray | List[str] | str +Theta = Tensor | np.ndarray | List[Any] + # Refactoring following #1175. tl:dr: letting joblib iterate over numpy arrays # allows for a roughly 10x performance gain. The resulting casting necessity @@ -26,7 +29,7 @@ def simulate_for_sbi( simulation_batch_size: Union[int, None] = 1, seed: Optional[int] = None, show_progress_bar: bool = True, -) -> Tuple[Tensor, Tensor]: +) -> Tuple[Tensor, Tensor | List[str] | str]: r"""Returns pairs :math:`(\theta, x)` by sampling proposal and running simulations. This function performs two steps: @@ -60,61 +63,226 @@ def simulate_for_sbi( """ if num_simulations == 0: - theta = torch.tensor([], dtype=float32) - x = torch.tensor([], dtype=float32) + return torch.tensor([], dtype=float32), torch.tensor([], dtype=float32) + + seed_all_backends(seed) + theta = proposal.sample((num_simulations,)) + # Cast to numpy for joblib efficiency + theta_numpy = theta.cpu().numpy() + if simulation_batch_size is None: + simulation_batch_size = num_simulations else: - # Cast theta to numpy for better joblib performance (seee #1175) - seed_all_backends(seed) - theta = proposal.sample((num_simulations,)) - - # Parse the simulation_batch_size logic - if simulation_batch_size is None: - simulation_batch_size = num_simulations - else: - simulation_batch_size = min(simulation_batch_size, num_simulations) - - if num_workers != 1: - # For multiprocessing, we want to switch to numpy arrays. - # The batch size will be an approximation, since np.array_split does - # not take as argument the size of the batch but their total. - num_batches = num_simulations // simulation_batch_size - batches = np.array_split(theta.cpu().numpy(), num_batches, axis=0) - batch_seeds = np.random.randint(low=0, high=1_000_000, size=(len(batches),)) + simulation_batch_size = min(simulation_batch_size, num_simulations) - # define seeded simulator. - def simulator_seeded(theta: ndarray, seed: int) -> Tensor: - seed_all_backends(seed) - return simulator(theta) - - try: # catch TypeError to give more informative error message - simulation_outputs: list[Tensor] = [ # pyright: ignore - xx - for xx in tqdm( - Parallel(return_as="generator", n_jobs=num_workers)( - delayed(simulator_seeded)(batch, seed) - for batch, seed in zip(batches, batch_seeds, strict=False) - ), - total=num_simulations, - disable=not show_progress_bar, - ) - ] - except TypeError as err: + # Handle parallel context + context = parallel_config(n_jobs=num_workers) + + with context: + # We enforce simulator_is_batched=True because simulate_for_sbi semantics + # implies that the simulator receives batches (even if size 1). + try: + theta, x = simulate_from_theta( + simulator, + theta_numpy, + simulation_batch_size=simulation_batch_size, + simulator_is_batched=True, + show_progress_bar=show_progress_bar, + seed=seed, + ) + except TypeError as err: + if num_workers > 1: raise TypeError( "There is a TypeError error in your simulator function. Note: For" " multiprocessing, we switch to numpy arrays. Besides confirming" " your simulator works correctly, make sure to preprocess your" " simulator with `process_simulator` to handle numpy arrays." ) from err + else: + raise err - else: - simulation_outputs: list[Tensor] = [] - batches = torch.split(theta, simulation_batch_size) - for batch in tqdm(batches, disable=not show_progress_bar): - simulation_outputs.append(simulator(batch)) + # Correctly format the output to Tensor + theta = torch.as_tensor(theta, dtype=float32) - # Correctly format the output - x = torch.cat(simulation_outputs, dim=0) - theta = torch.as_tensor(theta, dtype=float32) + if isinstance(x, np.ndarray): + x = torch.from_numpy(x) return theta, x + + +@overload +def parallelize_simulator( + simulator: Callable[[Theta], X], + simulator_is_batched: bool = ..., + simulation_batch_size: int = ..., + show_progress_bar: bool = ..., + seed: Optional[int] = ..., +) -> Callable[[Theta], X]: ... + + +@overload +def parallelize_simulator( + simulator: None = None, + simulator_is_batched: bool = ..., + simulation_batch_size: int = ..., + show_progress_bar: bool = ..., + seed: Optional[int] = ..., +) -> Callable[[Callable[[Theta], X]], Callable[[Theta], X]]: ... + + +def parallelize_simulator( + simulator: Callable[[Theta], X] | None = None, + simulator_is_batched: bool = False, + simulation_batch_size: int = 10, + show_progress_bar: bool = True, + seed: Optional[int] = None, +) -> Union[ + Callable[[Theta], X], + Callable[[Callable[[Theta], X]], Callable[[Theta], X]], +]: + r""" + Returns a function that executes simulations in parallel for a given set of + parameters. Can be used as a function or a decorator. + + Args: + simulator: Function to run simulations. + simulator_is_batched: Whether the simulator can handle batches directly. + simulation_batch_size: Number of simulations to run in each batch. + show_progress_bar: Whether to show tqdm progress bar. + seed: Random seed. + + Returns: + Callable that takes a set of :math:`\theta` and returns simulation outputs. + """ + + def decorator(simulator_func: Callable[[Theta], X]) -> Callable[[Theta], X]: + warnings.warn( + "Joblib is used for parallelization. It is recommended to use numpy arrays " + "for the simulator input and output to avoid serialization overhead with " + "torch tensors.", + UserWarning, + stacklevel=2, + ) + + def parallel_simulator(thetas: Theta) -> X: + seed_all_backends(seed) + + num_simulations = len(thetas) + + if num_simulations == 0: + return torch.tensor([], dtype=float32) + + # Create batches + if simulator_is_batched: + num_batches = ( + num_simulations + simulation_batch_size - 1 + ) // simulation_batch_size + batches = [ + thetas[i * simulation_batch_size : (i + 1) * simulation_batch_size] + for i in range(num_batches) + ] + elif simulation_batch_size > 1: + warnings.warn( + "Simulation batch size is greater than 1, but simulator_is_batched " + "is False. Simulations will be run sequentially (batch size 1).", + UserWarning, + stacklevel=2, + ) + batches = [theta for theta in thetas] + else: + batches = [theta for theta in thetas] + + # Run in parallel + # Generate seeds + batch_seeds = np.random.randint(low=0, high=1_000_000, size=(len(batches),)) + + def run_simulation(batch, seed): + seed_all_backends(seed) + return simulator_func(batch) + + # Execute in parallel with joblib + results = Parallel(return_as="generator")( + delayed(run_simulation)(batch, seed) + for batch, seed in zip(batches, batch_seeds, strict=False) + ) + + # Progress bar + simulation_outputs = [] + if show_progress_bar: + pbar = tqdm(total=num_simulations) + + for i, res in enumerate(results): + simulation_outputs.append(res) + if show_progress_bar: + pbar.update(len(batches[i])) + + if show_progress_bar: + pbar.close() + + # Flatten results + output_data = [] + if simulator_is_batched: + for batch_out in simulation_outputs: + if isinstance(batch_out, (list, tuple)): + output_data.extend(batch_out) + elif isinstance(batch_out, (torch.Tensor, np.ndarray)): + output_data.extend([x for x in batch_out]) + else: + output_data.append(batch_out) + else: + output_data = simulation_outputs + + if not output_data: + return torch.tensor([], dtype=float32) + + # Handle file paths (strings) + if isinstance(output_data[0], (str, np.str_, np.bytes_)): + output_data = [str(f) for f in output_data] + return output_data + + if isinstance(output_data[0], torch.Tensor): + return torch.stack(output_data) + elif isinstance(output_data[0], np.ndarray): + return np.stack(output_data) + + return output_data + + return parallel_simulator + + if simulator is None: + return decorator + + return decorator(simulator) + + +def simulate_from_theta( + simulator: Callable[[Theta], X], + thetas: Theta, + simulator_is_batched: bool = False, + simulation_batch_size: int = 10, + show_progress_bar: bool = True, + seed: Optional[int] = None, +) -> Tuple[Theta, X]: + r""" + Execute simulations for a given set of parameters. + + Args: + simulator: Function to run simulations. + thetas: Parameters to simulate (Tensor, Numpy array, or list). + simulator_is_batched: Whether the simulator can handle batches directly. + simulation_batch_size: Number of simulations to run in each batch. + show_progress_bar: Whether to show tqdm progress bar. + seed: Random seed. + + Returns: + Tuple of (:math:`\theta`, simulation_outputs). + """ + parallel_sim = parallelize_simulator( + simulator, + simulator_is_batched=simulator_is_batched, + simulation_batch_size=simulation_batch_size, + show_progress_bar=show_progress_bar, + seed=seed, + ) + + return thetas, parallel_sim(thetas) diff --git a/sbi/utils/user_input_checks.py b/sbi/utils/user_input_checks.py index 56660fe73..557a78966 100644 --- a/sbi/utils/user_input_checks.py +++ b/sbi/utils/user_input_checks.py @@ -471,6 +471,8 @@ def process_simulator( is_numpy_simulator: bool, ) -> Callable: """Returns a simulator that meets the requirements for usage in sbi. + NOTE: This method is deprecated and will be removed in a future release. Please use + `sbi.utils.simulation_utils.parallelize_simulator` instead. Args: user_simulator (Callable): @@ -502,6 +504,13 @@ def process_simulator( simulator = process_simulator(simulator, prior, prior_returns_numpy) """ + warnings.warn( + "process_simulator is deprecated and will be removed in a future release. " + "Please use `sbi.utils.simulation_utils.parallelize_simulator` instead.", + FutureWarning, + stacklevel=2, + ) + assert isinstance(user_simulator, Callable), "Simulator must be a function." joblib_simulator = wrap_as_joblib_efficient_simulator( @@ -518,6 +527,7 @@ def process_simulator( # of the simulator. This is not efficient (~3 times slowdown), but is compatible # with the new joblib and, importantly, does not break previous code. It should # be removed with a future restructuring of the simulation pipeline. +# TODO: delete when removing process_simulator and related functions. def wrap_as_joblib_efficient_simulator( simulator: Callable, prior, is_numpy_simulator ) -> Callable: @@ -555,6 +565,7 @@ def pytorch_simulator(theta: Tensor) -> Tensor: return pytorch_simulator +# TODO: delete when removing process_simulator and related functions. def ensure_batched_simulator(simulator: Callable, prior) -> Callable: """Return a simulator with batched output. diff --git a/tests/simulation_utils_test.py b/tests/simulation_utils_test.py new file mode 100644 index 000000000..2bddc7870 --- /dev/null +++ b/tests/simulation_utils_test.py @@ -0,0 +1,195 @@ +# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed +# under the Apache License Version 2.0, see + +import warnings + +import numpy as np +import pytest +import torch + +from sbi.utils.simulation_utils import parallelize_simulator, simulate_from_theta + +# --- Simulators for testing --- + + +def simple_simulator_torch(theta): + """Simple simulator utilizing Torch operations. + Input: Theta (N, D) or (D,) + Output: (N, D) or (D,) + """ + return theta * 2.0 + + +def simple_simulator_numpy(theta): + """Simple simulator utilizing Numpy operations. + Input: Theta (Tensor or Array) + Output: Array + """ + if isinstance(theta, torch.Tensor): + theta = theta.numpy() + return theta * 2.0 + + +def list_simulator(theta): + """Simulator that returns a list (e.g. file paths).""" + if isinstance(theta, torch.Tensor): + # If batched, return list of strings + if theta.ndim > 1: + return [f"file_{t.sum()}" for t in theta] + else: + return f"file_{theta.sum()}" + return f"file_{theta}" + + +# --- Tests for parallelize_simulator --- + + +def test_parallelize_simulator_decorator(): + """Test parallelize_simulator used as a decorator.""" + + # 1. No arguments + @parallelize_simulator + def decorated_sim(theta): + return theta * 2.0 + + thetas = torch.ones(10, 2) + # By default, decorator assumes simulator_is_batched=False, batch_size=10? + # Checking default args in implementation: + # simulator_is_batched = False + # simulation_batch_size = 10 + # But if simulator_is_batched is False, and batch_size > 1, + # it warns and runs sequentially. + + with pytest.warns(UserWarning, match="Simulation batch size is greater than 1"): + result = decorated_sim(thetas) + + assert result.shape == (10, 2) + assert torch.allclose(result, thetas * 2.0) + + # 2. With arguments + @parallelize_simulator(simulation_batch_size=2, simulator_is_batched=True) + def decorated_sim_batched(theta): + return theta * 2.0 + + result_batched = decorated_sim_batched(thetas) + assert result_batched.shape == (10, 2) + assert torch.allclose(result_batched, thetas * 2.0) + + +def test_parallelize_simulator_function(): + """Test parallelize_simulator used as a function wrapper.""" + sim = parallelize_simulator( + simple_simulator_torch, simulation_batch_size=5, simulator_is_batched=True + ) + thetas = torch.ones(10, 2) + result = sim(thetas) + assert result.shape == (10, 2) + assert torch.allclose(result, thetas * 2.0) + + +def test_parallelize_simulator_numpy_recommendation_warning(): + """Test that warning suggests using Numpy.""" + thetas = torch.ones(2, 2) + with pytest.warns(UserWarning, match="Joblib is used for parallelization"): + wrapper = parallelize_simulator(simple_simulator_torch) + # Warning is emitted at wrapper creation time + + # Run to ensure no error + wrapper(thetas) + + +# --- Tests for simulate_from_thetas --- + + +@pytest.mark.parametrize("simulator", [simple_simulator_torch, simple_simulator_numpy]) +@pytest.mark.parametrize("simulator_is_batched", [True, False]) +@pytest.mark.parametrize("simulation_batch_size", [1, 5]) +def test_simulate_from_theta_correctness( + simulator, simulator_is_batched, simulation_batch_size +): + """Test correctness of simulate_from_theta with various configs.""" + num_sims = 10 + dim = 2 + thetas = torch.rand(num_sims, dim) + + # Compute expected result using the simulator directly + # ensuring we handle both numpy and torch outputs + raw_expected = simulator(thetas) + if isinstance(raw_expected, np.ndarray): + expected_result_torch = torch.from_numpy(raw_expected) + else: + expected_result_torch = raw_expected + + # Run simulation + # We use catch_warnings to handle the expected warnings + # without failing or printing too much + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + + theta_out, x_out = simulate_from_theta( + simulator, + thetas, + simulator_is_batched=simulator_is_batched, + simulation_batch_size=simulation_batch_size, + show_progress_bar=False, + ) + + # Check Joblib warning presence + assert any("Joblib is used" in str(warn.message) for warn in w) + + # Check configuration warning: if batch_size > 1 and not batched + if simulation_batch_size > 1 and not simulator_is_batched: + assert any( + "Simulation batch size is greater than 1" in str(warn.message) + for warn in w + ) + + # Check shapes + assert theta_out.shape == (num_sims, dim) + # output shape depends on simulator return type logic in parallelize_simulator + # It stacks output. + assert x_out.shape == (num_sims, dim) + + # Check values + if isinstance(x_out, torch.Tensor): + assert torch.allclose(x_out, expected_result_torch) + elif isinstance(x_out, np.ndarray): + assert np.allclose(x_out, expected_result_torch.numpy()) + + +def test_simulate_from_thetas_list_return(): + """Test that list returns (e.g. file paths) are handled correctly.""" + thetas = torch.arange(5).float().unsqueeze(1) # [[0], [1], [2], [3], [4]] + + # 1. Batched Simulator + # Simulator receives batch, returns list of strings. + # parallelize_simulator should flatten these lists. + + def list_sim_batched(t): + return [f"path_{v.item()}" for v in t] + + theta_out, x_out = simulate_from_theta( + list_sim_batched, + thetas, + simulator_is_batched=True, + simulation_batch_size=2, + show_progress_bar=False, + ) + + assert len(x_out) == 5 + assert x_out[0] == "path_0.0" + assert x_out[-1] == "path_4.0" + + def list_sim_unbatched(t): + return f"path_{t.item()}" + + theta_out2, x_out2 = simulate_from_theta( + list_sim_unbatched, + thetas, + simulator_is_batched=False, + simulation_batch_size=1, + show_progress_bar=False, + ) + + assert len(x_out2) == 5 + assert x_out2[0] == "path_0.0"